Coverage for qml_essentials/coefficients.py: 18%
213 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-15 15:48 +0000
1from __future__ import annotations
2import numpy as np
3import math
4from collections import defaultdict
5from dataclasses import dataclass
6import pennylane as qml
7from pennylane.operation import Operator
8import pennylane.ops.op_math as qml_op
9from typing import List, Tuple, Optional, Any, Dict, Union
11from qml_essentials.model import Model
14class Coefficients:
16 @staticmethod
17 def get_spectrum(
18 model: Model,
19 mfs: int = 1,
20 mts: int = 1,
21 shift=False,
22 trim=False,
23 **kwargs,
24 ) -> np.ndarray:
25 """
26 Extracts the coefficients of a given model using a FFT (np-fft).
28 Note that the coefficients are complex numbers, but the imaginary part
29 of the coefficients should be very close to zero, since the expectation
30 values of the Pauli operators are real numbers.
32 It can perform oversampling in both the frequency and time domain
33 using the `mfs` and `mts` arguments.
35 Args:
36 model (Model): The model to sample.
37 mfs (int): Multiplicator for the highest frequency. Default is 2.
38 mts (int): Multiplicator for the number of time samples. Default is 1.
39 shift (bool): Whether to apply np-fftshift. Default is False.
40 trim (bool): Whether to remove the Nyquist frequency if spectrum is even.
41 Default is False.
42 kwargs (Any): Additional keyword arguments for the model function.
44 Returns:
45 np.ndarray: The sampled Fourier coefficients.
46 """
47 kwargs.setdefault("force_mean", True)
48 kwargs.setdefault("execution_type", "expval")
50 coeffs, freqs = Coefficients._fourier_transform(
51 model, mfs=mfs, mts=mts, **kwargs
52 )
54 if not np.isclose(np.sum(coeffs).imag, 0.0, rtol=1.0e-5):
55 raise ValueError(
56 f"Spectrum is not real. Imaginary part of coefficients is:\
57 {np.sum(coeffs).imag}"
58 )
60 if trim:
61 for ax in range(len(coeffs.shape) - 1):
62 if coeffs.shape[ax] % 2 == 0:
63 coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
64 freqs = np.delete(freqs, len(freqs) // 2, axis=ax)
66 if shift:
67 return np.fft.fftshift(
68 coeffs, axes=list(range(model.n_input_feat))
69 ), np.fft.fftshift(freqs)
70 else:
71 return coeffs, freqs
73 @staticmethod
74 def _fourier_transform(
75 model: Model, mfs: int, mts: int, **kwargs: Any
76 ) -> np.ndarray:
77 # Create a frequency vector with as many frequencies as model degrees,
78 # oversampled by nfs
79 n_freqs: int = 2 * mfs * model.degree + 1
81 start, stop, step = 0, 2 * mts * np.pi, 2 * np.pi / n_freqs
82 # Stretch according to the number of frequencies
83 inputs: np.ndarray = np.arange(start, stop, step) % (2 * np.pi)
85 # permute with input dimensionality
86 nd_inputs = np.array(np.meshgrid(*[inputs] * model.n_input_feat)).T.reshape(
87 -1, model.n_input_feat
88 )
90 # Output vector is not necessarily the same length as input
91 outputs = model(inputs=nd_inputs, **kwargs)
92 outputs = outputs.reshape(*(inputs.shape * model.n_input_feat), -1).squeeze()
94 coeffs = np.fft.fftn(outputs, axes=list(range(model.n_input_feat)))
96 # assert (
97 # mts * n_freqs,
98 # ) * model.n_input_feat == coeffs.shape, f"Expected shape\
99 # {(mts * n_freqs,) * model.n_input_feat} but got {coeffs.shape}"
101 freqs = np.fft.fftfreq(mts * n_freqs, 1 / n_freqs)
103 # TODO: this could cause issues with multidim input
104 # FIXME: account for different frequencies in multidim input scenarios
105 # Run the fft and rearrange +
106 # normalize the output (using product if multidim)
107 return (
108 coeffs / np.prod(outputs.shape[0 : model.n_input_feat]),
109 freqs,
110 # np.repeat(freqs[:, np.newaxis], model.n_input_feat, axis=1).squeeze(),
111 )
113 @staticmethod
114 def get_psd(coeffs: np.ndarray) -> np.ndarray:
115 """
116 Calculates the power spectral density (PSD) from given Fourier coefficients.
118 Args:
119 coeffs (np.ndarray): The Fourier coefficients.
121 Returns:
122 np.ndarray: The power spectral density.
123 """
124 # TODO: if we apply trim=True in advance, this will be slightly wrong..
126 def abs2(x):
127 return x.real**2 + x.imag**2
129 scale = 2.0 / (len(coeffs) ** 2)
130 return scale * abs2(coeffs)
132 @staticmethod
133 def evaluate_Fourier_series(
134 coefficients: np.ndarray,
135 frequencies: np.ndarray,
136 inputs: Union[np.ndarray, list, float],
137 ) -> float:
138 """
139 Evaluate the function value of a Fourier series at one point.
141 Args:
142 coefficients (np.ndarray): Coefficients of the Fourier series.
143 frequencies (np.ndarray): Corresponding frequencies.
144 inputs (np.ndarray): Point at which to evaluate the function.
145 Returns:
146 float: The function value at the input point.
147 """
148 dims = len(coefficients.shape)
150 if not isinstance(inputs, (np.ndarray, list)):
151 inputs = [inputs]
153 frequencies = np.stack(np.meshgrid(*[frequencies] * dims)).T.reshape(-1, dims)
154 freq_inputs = np.einsum("...j,j->...", frequencies, inputs)
155 coeffs = coefficients.flatten()
156 freq_inputs = freq_inputs.flatten()
158 exp = 0.0
159 for omega_x, c in zip(freq_inputs, coeffs):
160 exp += c * np.exp(1j * omega_x)
162 return np.real_if_close(exp)
165class FourierTree:
166 """
167 Sine-cosine tree representation for the algorithm by Nemkov et al.
168 This tree can be used to obtain analytical Fourier coefficients for a given
169 Pauli-Clifford circuit.
170 """
172 class CoefficientsTreeNode:
173 """
174 Representation of a node in the coefficients tree for the algorithm by
175 Nemkov et al.
176 """
178 def __init__(
179 self,
180 parameter_idx: Optional[int],
181 observable: Operator,
182 is_sine_factor: bool,
183 is_cosine_factor: bool,
184 left: Optional[FourierTree.CoefficientsTreeNode] = None,
185 right: Optional[FourierTree.CoefficientsTreeNode] = None,
186 ):
187 """
188 Coefficient tree node initialisation. Each node has information about
189 its creation context and it's children, i.e.:
191 Args:
192 parameter_idx (Optional[int]): Index of the corresp. param. index i.
193 observable (Operator): The nodes observable to obtain the
194 expectation value that contributes to the constant term.
195 is_sine_factor (bool): If this node belongs to a sine coefficient.
196 is_cosine_factor (bool): If this node belongs to a cosine coefficient.
197 left (Optional[CoefficientsTreeNode]): left child (if any).
198 right (Optional[CoefficientsTreeNode]): right child (if any).
199 """
200 self.parameter_idx = parameter_idx
202 assert not (
203 is_sine_factor and is_cosine_factor
204 ), "Cannot be sine and cosine at the same time"
205 self.is_sine_factor = is_sine_factor
206 self.is_cosine_factor = is_cosine_factor
208 if isinstance(observable, qml_op.SProd):
209 term = observable.terms()[0][0]
210 observable = observable.terms()[1][0]
211 else:
212 term = 1.0
214 # If the observable does not constist of only Z and I, the
215 # expectation (and therefore the constant node term) is zero
216 if (
217 isinstance(observable, qml_op.Prod)
218 and any([isinstance(p, (qml.X, qml.Y)) for p in observable])
219 or isinstance(observable, (qml.PauliX, qml.PauliY))
220 ):
221 self.term = 0.0
222 else:
223 self.term = term
225 self.observable = observable
227 self.left = left
228 self.right = right
230 def evaluate(self, parameters: list[float]) -> float:
231 """
232 Recursive function to evaluate the expectation of the coefficient tree,
233 starting from the current node.
235 Args:
236 parameters (list[float]): The parameters, by which the circuit (and
237 therefore the tree) is parametrised.
239 Returns:
240 float: The expectation for the current node and it's children.
241 """
242 factor = (
243 parameters[self.parameter_idx]
244 if self.parameter_idx is not None
245 else 1.0
246 )
247 if self.is_sine_factor:
248 factor = 1j * np.sin(factor)
249 elif self.is_cosine_factor:
250 factor = np.cos(factor)
251 if not (self.left or self.right): # leaf
252 return factor * self.term
254 sum_children = 0.0
255 if self.left:
256 left = self.left.evaluate(parameters)
257 sum_children = sum_children + left
258 if self.right:
259 right = self.right.evaluate(parameters)
260 sum_children = sum_children + right
262 return factor * sum_children
264 def get_leafs(
265 self,
266 sin_list: List[int],
267 cos_list: List[int],
268 existing_leafs: List[FourierTree.TreeLeaf] = [],
269 ) -> List[FourierTree.TreeLeaf]:
270 """
271 Traverse the tree starting from the current node, to obtain the tree
272 leafs only.
273 The leafs correspond to the terms in the sine-cosine tree
274 representation that eventually are used to obtain coefficients and
275 frequencies.
276 Sine and cosine lists are recursively passed to the children until a
277 leaf is reached (top to bottom).
278 Leafs are then passed bottom to top to the caller.
280 Args:
281 sin_list (List[int]): Current number of sine contributions for each
282 parameter. Has the same length as the parameters, as each
283 position corresponds to one parameter.
284 cos_list (List[int]): Current number of cosine contributions for
285 each parameter. Has the same length as the parameters, as each
286 position corresponds to one parameter.
287 existing_leafs (List[TreeLeaf]): Current list of leaf nodes from
288 parents.
290 Returns:
291 List[TreeLeaf]: Updated list of leaf nodes.
292 """
294 if self.is_sine_factor:
295 sin_list[self.parameter_idx] += 1
296 if self.is_cosine_factor:
297 cos_list[self.parameter_idx] += 1
299 if not (self.left or self.right): # leaf
300 if self.term != 0.0:
301 return [FourierTree.TreeLeaf(sin_list, cos_list, self.term)]
302 else:
303 return []
305 if self.left:
306 leafs_left = self.left.get_leafs(
307 sin_list.copy(), cos_list.copy(), existing_leafs.copy()
308 )
309 else:
310 leafs_left = []
312 if self.right:
313 leafs_right = self.right.get_leafs(
314 sin_list.copy(), cos_list.copy(), existing_leafs.copy()
315 )
316 else:
317 leafs_right = []
319 existing_leafs.extend(leafs_left)
320 existing_leafs.extend(leafs_right)
321 return existing_leafs
323 @dataclass
324 class TreeLeaf:
325 """
326 Coefficient tree leafs according to the algorithm by Nemkov et al., which
327 correspond to the terms in the sine-cosine tree representation that
328 eventually are used to obtain coefficients and frequencies.
330 Attributes:
331 sin_indices (List[int]): Current number of sine contributions for each
332 parameter. Has the same length as the parameters, as each
333 position corresponds to one parameter.
334 cos_list (List[int]): Current number of cosine contributions for
335 each parameter. Has the same length as the parameters, as each
336 position corresponds to one parameter.
337 term (np.complex): Constant factor of the leaf, depending on the
338 expectation value of the observable, and a phase.
339 """
341 sin_indices: List[int]
342 cos_indices: List[int]
343 term: np.complex128
345 def __init__(self, model: Model, inputs=1.0):
346 """
347 Tree initialisation, based on the Pauli-Clifford representation of a model.
348 Currently, only one input feature is supported.
350 **Usage**:
351 ```
352 # initialise a model
353 model = Model(...)
355 # initialise and build FourierTree
356 tree = FourierTree(model)
358 # get expectaion value
359 exp = tree()
361 # Get spectrum (for each observable, we have one list element)
362 coeff_list, freq_list = tree.spectrum()
363 ```
365 Args:
366 model (Model): The Model, for which to build the tree
367 inputs (bool, optional): Possible default inputs. Defaults to 1.0.
368 """
369 self.model = model
370 self.tree_roots = None
372 if not model.as_pauli_circuit:
373 model.as_pauli_circuit = True
375 inputs = (
376 self.model._inputs_validation(inputs)
377 if inputs is not None
378 else self.model._inputs_validation(1.0)
379 )
380 inputs.requires_grad = False
382 quantum_tape = qml.workflow.construct_tape(self.model.circuit)(
383 params=model.params, inputs=inputs
384 )
385 self.parameters = [np.squeeze(p) for p in quantum_tape.get_parameters()]
386 self.input_indices = [
387 i for (i, p) in enumerate(self.parameters) if not p.requires_grad
388 ]
390 self.observables = quantum_tape.observables
391 self.pauli_rotations = quantum_tape.operations
393 self.tree_roots = self.build()
394 self.leafs: List[List[FourierTree.TreeLeaf]] = self._get_tree_leafs()
396 def __call__(
397 self,
398 params: Optional[np.ndarray] = None,
399 inputs: Optional[np.ndarray] = None,
400 **kwargs,
401 ) -> np.ndarray:
402 """
403 Evaluates the Fourier tree via sine-cosine terms sum. This is
404 equivalent to computing the expectation value of the observables with
405 respect to the corresponding circuit.
407 Args:
408 params (Optional[np.ndarray], optional): Parameters of the model.
409 Defaults to None.
410 inputs (Optional[np.ndarray], optional): Inputs to the circuit.
411 Defaults to None.
413 Returns:
414 np.ndarray: Expectation value of the tree.
416 Raises:
417 NotImplementedError: When using other "execution_type" as expval.
418 NotImplementedError: When using "noise_params"
421 """
422 params = (
423 self.model._params_validation(params)
424 if params is not None
425 else self.model.params
426 )
427 inputs = (
428 self.model._inputs_validation(inputs)
429 if inputs is not None
430 else self.model._inputs_validation(1.0)
431 )
432 inputs.requires_grad = False
434 if kwargs.get("execution_type", "expval") != "expval":
435 raise NotImplementedError(
436 f'Currently, only "expval" execution type is supported when '
437 f'building FourierTree. Got {kwargs.get("execution_type", "expval")}.'
438 )
439 if kwargs.get("noise_params", None) is not None:
440 raise NotImplementedError(
441 "Currently, noise is not supported when building FourierTree."
442 )
444 quantum_tape = qml.workflow.construct_tape(self.model.circuit)(
445 params=self.model.params, inputs=inputs
446 )
447 self.parameters = [np.squeeze(p) for p in quantum_tape.get_parameters()]
448 self.input_indices = [
449 i for (i, p) in enumerate(self.parameters) if not p.requires_grad
450 ]
452 results = np.zeros(len(self.tree_roots))
453 for i, root in enumerate(self.tree_roots):
454 results[i] = np.real_if_close(root.evaluate(self.parameters))
456 if kwargs.get("force_mean", False):
457 return np.mean(results)
458 else:
459 return results
461 def build(self) -> List[CoefficientsTreeNode]:
462 """
463 Creates the coefficient tree, i.e. it creates and initialises the tree
464 nodes.
465 Leafs can be obtained separately in _get_tree_leafs, once the tree is
466 set up.
468 Returns:
469 List[CoefficientsTreeNode]: The list of root nodes (one root for
470 each observable).
471 """
472 tree_roots = []
473 for obs in self.observables:
474 pauli_rotation_indices = np.arange(
475 len(self.pauli_rotations), dtype=np.int16
476 )
477 root = self._create_tree_node(obs, pauli_rotation_indices)
478 tree_roots.append(root)
479 return tree_roots
481 def _get_tree_leafs(self) -> List[List[TreeLeaf]]:
482 """
483 Obtain all Leaf Nodes with its sine- and cosine terms.
485 Returns:
486 List[List[TreeLeaf]]: For each observable (root), the list of leaf
487 nodes.
488 """
489 leafs = []
490 for root in self.tree_roots:
491 sin_list = np.zeros(len(self.parameters), dtype=np.int32)
492 cos_list = np.zeros(len(self.parameters), dtype=np.int32)
493 leafs.append(root.get_leafs(sin_list, cos_list, []))
494 return leafs
496 def get_spectrum(
497 self, force_mean: bool = False
498 ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
499 """
500 Computes the Fourier spectrum for the tree, consisting of the
501 frequencies and its corresponding coefficinets.
502 If the frag force_mean was set in the constructor, the mean coefficient
503 over all observables (roots) are computed.
505 Args:
506 force_mean (bool, optional): Whether to average over multiple
507 observables. Defaults to False.
509 Returns:
510 Tuple[List[np.ndarray], List[np.ndarray]]:
511 - List of frequencies, one list for each observable (root).
512 - List of corresponding coefficents, one list for each
513 observable (root).
514 """
515 parameter_indices = [
516 i for i in range(len(self.parameters)) if i not in self.input_indices
517 ]
519 coeffs = []
520 for leafs in self.leafs:
521 freq_terms = defaultdict(np.complex128)
522 for leaf in leafs:
523 leaf_factor, s, c = self._compute_leaf_factors(leaf, parameter_indices)
525 for a in range(s + 1):
526 for b in range(c + 1):
527 comb = math.comb(s, a) * math.comb(c, b) * (-1) ** (s - a)
528 freq_terms[2 * a + 2 * b - s - c] += comb * leaf_factor
530 coeffs.append(freq_terms)
532 frequencies, coefficients = self._freq_terms_to_coeffs(coeffs, force_mean)
533 return coefficients, frequencies
535 def _freq_terms_to_coeffs(
536 self, coeffs: List[Dict[int, np.ndarray]], force_mean: bool
537 ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
538 """
539 Given a list of dictionaries of the form:
540 [
541 {
542 freq_obs1_1: coeff1,
543 freq_obs1_2: coeff2,
544 ...
545 },
546 {
547 freq_obs2_1: coeff3,
548 freq_obs2_2: coeff4,
549 ...
550 }
551 ...
552 ],
553 Compute two separate lists of frequencies and coefficients.
554 such that:
555 freqs: [
556 [freq_obs1_1, freq_obs1_1, ...],
557 [freq_obs2_1, freq_obs2_1, ...],
558 ...
559 ]
560 coeffs: [
561 [coeff1, coeff2, ...],
562 [coeff3, coeff4, ...],
563 ...
564 ]
566 If force_mean is set length of the resulting frequency and coefficent
567 list is 1.
569 Args:
570 coeffs (List[Dict[int, np.ndarray]]): Frequency->Coefficients
571 dictionary list, one dict for each observable (root).
572 force_mean (bool, optional): Whether to average coefficients over
573 multiple observables. Defaults to False.
575 Returns:
576 Tuple[List[np.ndarray], List[np.ndarray]]:
577 - List of frequencies, one list for each observable (root).
578 - List of corresponding coefficents, one list for each
579 observable (root).
580 """
581 frequencies = []
582 coefficients = []
583 if force_mean:
584 all_freqs = sorted(set([f for c in coeffs for f in c.keys()]))
585 coefficients.append(
586 np.array([np.mean([c.get(f, 0.0) for c in coeffs]) for f in all_freqs])
587 )
588 frequencies.append(np.array(all_freqs))
589 else:
590 for freq_terms in coeffs:
591 freq_terms = dict(sorted(freq_terms.items()))
592 frequencies.append(np.array(list(freq_terms.keys())))
593 coefficients.append(np.array(list(freq_terms.values())))
594 return frequencies, coefficients
596 def _compute_leaf_factors(
597 self, leaf: TreeLeaf, parameter_indices: List[int]
598 ) -> Tuple[float, int, int]:
599 """
600 Computes the constant coefficient factor for each leaf.
601 Additionally sine and cosine contributions of the input parameters for
602 this leaf are returned, which are required to obtain the corresponding
603 frequencies.
605 Args:
606 leaf (TreeLeaf): The leaf for which to compute the factor.
607 parameter_indices (List[int]): Variational parameter indices.
609 Returns:
610 Tuple[float, int, int]:
611 - float: the constant factor for the leaf
612 - int: number of sine contributions of the input
613 - int: number of cosine contributions of the input
614 """
615 leaf_factor = 1.0
616 for i in parameter_indices:
617 interm_factor = (
618 np.cos(self.parameters[i]) ** leaf.cos_indices[i]
619 * (1j * np.sin(self.parameters[i])) ** leaf.sin_indices[i]
620 )
621 leaf_factor = leaf_factor * interm_factor
623 # Get number of sine and cosine factors to which the input contributes
624 c = np.sum([leaf.cos_indices[k] for k in self.input_indices], dtype=np.int32)
625 s = np.sum([leaf.sin_indices[k] for k in self.input_indices], dtype=np.int32)
627 leaf_factor = leaf.term * leaf_factor * 0.5 ** (s + c)
629 return leaf_factor, s, c
631 def _create_tree_node(
632 self,
633 observable: Operator,
634 pauli_rotation_indices: List[int],
635 parameter_idx: Optional[int] = None,
636 is_sine: bool = False,
637 is_cosine: bool = False,
638 ) -> CoefficientsTreeNode:
639 """
640 Builds the Fourier-Tree according to the algorithm by Nemkov et al.
642 Args:
643 observable (Operator): Current observable
644 pauli_rotation_indices (List[int]): Indices of remaining Pauli
645 rotation gates. Gates itself are attributes of the class.
646 parameter_idx (Optional[int]): Index of the current parameter.
647 Parameters itself are attributes of the class.
648 is_sine (bool): If the current node is a sine (left) node.
649 is_cosine (bool): If the current node is a cosine (right) node.
651 Returns:
652 CoefficientsTreeNode: The resulting node. Children are set
653 recursively. The top level receives the tree root.
654 """
656 # remove commuting paulis
657 idx = len(pauli_rotation_indices) - 1
658 while idx >= 0:
659 last_pauli = self.pauli_rotations[pauli_rotation_indices[idx]]
660 if not qml.is_commuting(last_pauli.generator(), observable):
661 break
662 idx -= 1
664 if idx < 0: # leaf
665 return FourierTree.CoefficientsTreeNode(
666 parameter_idx, observable, is_sine, is_cosine
667 )
669 next_pauli_rotation_indices = pauli_rotation_indices[:idx]
670 last_pauli_idx = pauli_rotation_indices[idx]
671 last_pauli = self.pauli_rotations[last_pauli_idx]
673 left = self._create_tree_node(
674 observable,
675 next_pauli_rotation_indices,
676 last_pauli_idx,
677 is_cosine=True,
678 )
680 next_observable = self._create_new_observable(
681 last_pauli.generator(), observable
682 )
683 right = self._create_tree_node(
684 next_observable,
685 next_pauli_rotation_indices,
686 last_pauli_idx,
687 is_sine=True,
688 )
690 return FourierTree.CoefficientsTreeNode(
691 parameter_idx,
692 observable,
693 is_sine,
694 is_cosine,
695 left,
696 right,
697 )
699 def _create_new_observable(self, pauli: Operator, observable: Operator) -> Operator:
700 """
701 Utility function to obtain the new observable for a tree node, if the
702 last Pauli and the observable do not commute.
704 Args:
705 pauli (Operator): The generator of the last Pauli rotation in the
706 operation sequence.
707 observable (Operator): The current observable.
709 Returns:
710 Operator: The new observable.
711 """
713 pauli = pauli[0] / pauli.coeffs[0] # ignore coefficients of generator
714 obs = pauli @ observable
715 obs = qml.simplify(obs)
717 return obs