Coverage for qml_essentials / coefficients.py: 95%
427 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-30 11:43 +0000
1from __future__ import annotations
2import math
3from collections import defaultdict
4from dataclasses import dataclass
5import jax.numpy as jnp
6from jax import random
7import numpy as np
8from scipy.stats import rankdata
9from functools import reduce
10from typing import List, Tuple, Optional, Any, Dict, Union
12from qml_essentials.model import Model
13from qml_essentials.utils import PauliCircuit
14from qml_essentials.operations import (
15 Operation,
16 PauliX,
17 PauliY,
18 PauliZ,
19)
21import logging
23log = logging.getLogger(__name__)
26class Coefficients:
27 @staticmethod
28 def get_spectrum(
29 model: Model,
30 mfs: int = 1,
31 mts: int = 1,
32 shift=False,
33 trim=False,
34 numerical_cap: Optional[float] = -1,
35 **kwargs,
36 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
37 """
38 Extracts the coefficients of a given model using a FFT (jnp-fft).
40 Note that the coefficients are complex numbers, but the imaginary part
41 of the coefficients should be very close to zero, since the expectation
42 values of the Pauli operators are real numbers.
44 It can perform oversampling in both the frequency and time domain
45 using the `mfs` and `mts` arguments.
47 Args:
48 model (Model): The model to sample.
49 mfs (int): Multiplicator for the highest frequency. Default is 2.
50 mts (int): Multiplicator for the number of time samples. Default is 1.
51 shift (bool): Whether to apply jnp-fftshift. Default is False.
52 trim (bool): Whether to remove the Nyquist frequency if spectrum is even.
53 Default is False.
54 numerical_cap (Optional[float]): Numerical cap for the coefficients.
55 kwargs (Any): Additional keyword arguments for the model function.
57 Returns:
58 Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the coefficients
59 and frequencies.
60 """
61 kwargs.setdefault("force_mean", True)
62 kwargs.setdefault("execution_type", "expval")
64 coeffs, freqs = Coefficients._fourier_transform(
65 model, mfs=mfs, mts=mts, **kwargs
66 )
68 if not jnp.isclose(jnp.sum(coeffs).imag, 0.0, rtol=1.0e-5):
69 raise ValueError(
70 f"Spectrum is not real. Imaginary part of coefficients is:\
71 {jnp.sum(coeffs).imag}"
72 )
74 if trim:
75 for ax in range(model.n_input_feat):
76 if coeffs.shape[ax] % 2 == 0:
77 coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
78 freqs = [np.delete(freq, len(freq) // 2, axis=ax) for freq in freqs]
80 if shift:
81 coeffs = jnp.fft.fftshift(coeffs, axes=list(range(model.n_input_feat)))
82 freqs = np.fft.fftshift(freqs)
84 if numerical_cap > 0:
85 # set coeffs below threshold to zero
86 coeffs = jnp.where(
87 jnp.abs(coeffs) < numerical_cap,
88 jnp.zeros_like(coeffs),
89 coeffs,
90 )
92 if len(freqs) == 1:
93 freqs = freqs[0]
95 return coeffs, freqs
97 @staticmethod
98 def _fourier_transform(
99 model: Model, mfs: int, mts: int, **kwargs: Any
100 ) -> jnp.ndarray:
101 # Create a frequency vector with as many frequencies as model degrees,
102 # oversampled by mfs
103 n_freqs: jnp.ndarray = jnp.array(
104 [mfs * model.degree[i] for i in range(model.n_input_feat)]
105 )
107 start, stop, step = 0, 2 * mts * jnp.pi, 2 * jnp.pi / n_freqs
108 # Stretch according to the number of frequencies
109 inputs: List = [
110 jnp.arange(start, stop, step[i]) for i in range(model.n_input_feat)
111 ]
113 # permute with input dimensionality
114 nd_inputs = jnp.array(
115 jnp.meshgrid(*[inputs[i] for i in range(model.n_input_feat)])
116 ).T.reshape(-1, model.n_input_feat)
118 # Output vector is not necessarily the same length as input
119 outputs = model(inputs=nd_inputs, **kwargs)
120 outputs = outputs.reshape(
121 *[inputs[i].shape[0] for i in range(model.n_input_feat)], -1
122 ).squeeze()
124 coeffs = jnp.fft.fftn(outputs, axes=list(range(model.n_input_feat)))
126 freqs = [
127 jnp.fft.fftfreq(int(mts * n_freqs[i]), 1 / n_freqs[i])
128 for i in range(model.n_input_feat)
129 ]
130 # freqs = jnp.fft.fftfreq(mts * n_freqs, 1 / n_freqs)
132 # TODO: this could cause issues with multidim input
133 # FIXME: account for different frequencies in multidim input scenarios
134 # Run the fft and rearrange +
135 # normalize the output (using product if multidim)
136 return (
137 coeffs / math.prod(outputs.shape[0 : model.n_input_feat]),
138 freqs,
139 )
141 @staticmethod
142 def get_psd(coeffs: jnp.ndarray) -> jnp.ndarray:
143 """
144 Calculates the power spectral density (PSD) from given Fourier coefficients.
146 Args:
147 coeffs (jnp.ndarray): The Fourier coefficients.
149 Returns:
150 jnp.ndarray: The power spectral density.
151 """
152 # TODO: if we apply trim=True in advance, this will be slightly wrong..
154 def abs2(x):
155 return x.real**2 + x.imag**2
157 scale = 2.0 / (len(coeffs) ** 2)
158 return scale * abs2(coeffs)
160 @staticmethod
161 def evaluate_Fourier_series(
162 coefficients: jnp.ndarray,
163 frequencies: jnp.ndarray,
164 inputs: Union[jnp.ndarray, list, float],
165 ) -> float:
166 """
167 Evaluate the function value of a Fourier series at one point.
169 Args:
170 coefficients (jnp.ndarray): Coefficients of the Fourier series.
171 frequencies (jnp.ndarray): Corresponding frequencies.
172 inputs (jnp.ndarray): Point at which to evaluate the function.
173 Returns:
174 float: The function value at the input point.
175 """
176 if isinstance(frequencies, list):
177 if len(coefficients.shape) <= len(frequencies):
178 coefficients = coefficients[..., jnp.newaxis]
179 else:
180 if len(coefficients.shape) == 1:
181 coefficients = coefficients[..., jnp.newaxis]
183 if isinstance(inputs, list):
184 inputs = jnp.array(inputs)
185 if len(inputs.shape) < 1:
186 inputs = inputs[jnp.newaxis, ...]
188 if isinstance(frequencies, list):
189 input_dim = len(frequencies)
190 frequencies = jnp.stack(jnp.meshgrid(*frequencies))
191 if input_dim != len(inputs):
192 frequencies = jnp.repeat(
193 frequencies[jnp.newaxis, ...], inputs.shape[0], axis=0
194 )
195 freq_inputs = jnp.einsum("bi...,b->b...", frequencies, inputs)
196 exponents = jnp.exp(1j * freq_inputs).T
197 exp = jnp.einsum("jl...k,jl...b->b...k", coefficients, exponents)
198 else:
199 freq_inputs = jnp.einsum("i...,i->...", frequencies, inputs)
200 exponents = jnp.exp(1j * freq_inputs).T
201 exp = jnp.einsum("jl...k,jl...->k...", coefficients, exponents)
202 else:
203 frequencies = jnp.repeat(
204 frequencies[jnp.newaxis, ...], inputs.shape[0], axis=0
205 )
206 freq_inputs = jnp.einsum("i...,i->i...", frequencies, inputs)
207 exponents = jnp.exp(1j * freq_inputs)
208 exp = jnp.einsum("j...k,ij...->ik...", coefficients, exponents)
210 return jnp.squeeze(jnp.real(exp))
213class FourierTree:
214 """
215 Sine-cosine tree representation for the algorithm by Nemkov et al.
216 This tree can be used to obtain analytical Fourier coefficients for a given
217 Pauli-Clifford circuit.
218 """
220 class CoefficientsTreeNode:
221 """
222 Representation of a node in the coefficients tree for the algorithm by
223 Nemkov et al.
224 """
226 def __init__(
227 self,
228 parameter_idx: Optional[int],
229 observable: FourierTree.PauliOperator,
230 is_sine_factor: bool,
231 is_cosine_factor: bool,
232 left: Optional[FourierTree.CoefficientsTreeNode] = None,
233 right: Optional[FourierTree.CoefficientsTreeNode] = None,
234 ):
235 """
236 Coefficient tree node initialisation. Each node has information about
237 its creation context and it's children, i.e.:
239 Args:
240 parameter_idx (Optional[int]): Index of the corresp. param. index i.
241 observable (FourierTree.PauliOperator): The nodes observable to
242 obtain the expectation value that contributes to the constant
243 term.
244 is_sine_factor (bool): If this node belongs to a sine coefficient.
245 is_cosine_factor (bool): If this node belongs to a cosine coefficient.
246 left (Optional[CoefficientsTreeNode]): left child (if any).
247 right (Optional[CoefficientsTreeNode]): right child (if any).
248 """
249 self.parameter_idx = parameter_idx
251 assert not (
252 is_sine_factor and is_cosine_factor
253 ), "Cannot be sine and cosine at the same time"
254 self.is_sine_factor = is_sine_factor
255 self.is_cosine_factor = is_cosine_factor
257 # If the observable does not constist of only Z and I, the
258 # expectation (and therefore the constant node term) is zero
259 if jnp.logical_or(
260 observable.list_repr == 0, observable.list_repr == 1
261 ).any():
262 self.term = 0.0
263 else:
264 self.term = observable.phase
266 self.left = left
267 self.right = right
269 def evaluate(self, parameters: list[float]) -> float:
270 """
271 Recursive function to evaluate the expectation of the coefficient tree,
272 starting from the current node.
274 Args:
275 parameters (list[float]): The parameters, by which the circuit (and
276 therefore the tree) is parametrised.
278 Returns:
279 float: The expectation for the current node and it's children.
280 """
281 factor = (
282 parameters[self.parameter_idx]
283 if self.parameter_idx is not None
284 else 1.0
285 )
286 if self.is_sine_factor:
287 factor = 1j * jnp.sin(factor)
288 elif self.is_cosine_factor:
289 factor = jnp.cos(factor)
290 if not (self.left or self.right): # leaf
291 return factor * self.term
293 sum_children = 0.0
294 if self.left:
295 left = self.left.evaluate(parameters)
296 sum_children = sum_children + left
297 if self.right:
298 right = self.right.evaluate(parameters)
299 sum_children = sum_children + right
301 return factor * sum_children
303 def get_leafs(
304 self,
305 sin_list: List[int],
306 cos_list: List[int],
307 existing_leafs: List[FourierTree.TreeLeaf] = [],
308 ) -> List[FourierTree.TreeLeaf]:
309 """
310 Traverse the tree starting from the current node, to obtain the tree
311 leafs only.
312 The leafs correspond to the terms in the sine-cosine tree
313 representation that eventually are used to obtain coefficients and
314 frequencies.
315 Sine and cosine lists are recursively passed to the children until a
316 leaf is reached (top to bottom).
317 Leafs are then passed bottom to top to the caller.
319 Args:
320 sin_list (List[int]): Current number of sine contributions for each
321 parameter. Has the same length as the parameters, as each
322 position corresponds to one parameter.
323 cos_list (List[int]): Current number of cosine contributions for
324 each parameter. Has the same length as the parameters, as each
325 position corresponds to one parameter.
326 existing_leafs (List[TreeLeaf]): Current list of leaf nodes from
327 parents.
329 Returns:
330 List[TreeLeaf]: Updated list of leaf nodes.
331 """
333 if self.is_sine_factor:
334 sin_list = sin_list.at[self.parameter_idx].set(
335 sin_list[self.parameter_idx] + 1
336 )
337 if self.is_cosine_factor:
338 cos_list = cos_list.at[self.parameter_idx].set(
339 cos_list[self.parameter_idx] + 1
340 )
342 if not (self.left or self.right): # leaf
343 if self.term != 0.0:
344 return [FourierTree.TreeLeaf(sin_list, cos_list, self.term)]
345 else:
346 return []
348 if self.left:
349 leafs_left = self.left.get_leafs(
350 sin_list.copy(), cos_list.copy(), existing_leafs.copy()
351 )
352 else:
353 leafs_left = []
355 if self.right:
356 leafs_right = self.right.get_leafs(
357 sin_list.copy(), cos_list.copy(), existing_leafs.copy()
358 )
359 else:
360 leafs_right = []
362 existing_leafs.extend(leafs_left)
363 existing_leafs.extend(leafs_right)
364 return existing_leafs
366 @dataclass
367 class TreeLeaf:
368 """
369 Coefficient tree leafs according to the algorithm by Nemkov et al., which
370 correspond to the terms in the sine-cosine tree representation that
371 eventually are used to obtain coefficients and frequencies.
373 Args:
374 sin_indices (List[int]): Current number of sine contributions for each
375 parameter. Has the same length as the parameters, as each
376 position corresponds to one parameter.
377 cos_list (List[int]): Current number of cosine contributions for
378 each parameter. Has the same length as the parameters, as each
379 position corresponds to one parameter.
380 term (jnp.complex): Constant factor of the leaf, depending on the
381 expectation value of the observable, and a phase.
382 """
384 sin_indices: List[int]
385 cos_indices: List[int]
386 term: complex
388 class PauliOperator:
389 """
390 Utility class for storing Pauli Rotations, the corresponding indices
391 in the XY-Space (whether there is a gate with X or Y generator at a
392 certain qubit) and the phase.
394 Args:
395 pauli (Union[Operator, jnp.ndarray[int]]): Pauli Rotation Operation
396 or list representation
397 n_qubits (int): Number of qubits in the circuit
398 prev_xy_indices (Optional[jnp.ndarray[bool]]): X/Y indices of the
399 previous Pauli sequence. Defaults to None.
400 is_observable (bool): If the operator is an observable. Defaults to
401 False.
402 is_init (bool): If this Pauli operator is initialised the first
403 time. Defaults to True.
404 phase (float): Phase of the operator. Defaults to 1.0
405 """
407 def __init__(
408 self,
409 pauli: Union[Operation, jnp.ndarray[int]],
410 n_qubits: int,
411 prev_xy_indices: Optional[jnp.ndarray[bool]] = None,
412 is_observable: bool = False,
413 is_init: bool = True,
414 phase: float = 1.0,
415 ):
416 self.is_observable = is_observable
417 self.phase = phase
419 if is_init:
420 if not is_observable:
421 pauli = pauli.generator()
422 self.list_repr = self._create_list_representation(pauli, n_qubits)
423 else:
424 assert isinstance(pauli, jnp.ndarray)
425 self.list_repr = pauli
427 if prev_xy_indices is None:
428 prev_xy_indices = jnp.zeros(n_qubits, dtype=bool)
429 self.xy_indices = jnp.logical_or(
430 prev_xy_indices,
431 self._compute_xy_indices(self.list_repr, rev=is_observable),
432 )
434 @staticmethod
435 def _compute_xy_indices(
436 op: jnp.ndarray[int], rev: bool = False
437 ) -> jnp.ndarray[bool]:
438 """
439 Computes the positions of X or Y gates to an one-hot encoded boolen
440 array.
442 Args:
443 op (jnp.ndarray[int]): Pauli-Operation list representation.
444 rev (bool): Whether to negate the array.
446 Returns:
447 jnp.ndarray[bool]: One hot encoded boolean array.
448 """
449 xy_indices = (op == 0) + (op == 1)
450 if rev:
451 xy_indices = ~xy_indices
452 return xy_indices
454 @staticmethod
455 def _create_list_representation(
456 op: Operation, n_qubits: int
457 ) -> jnp.ndarray[int]:
458 """
459 Create list representation of an Operation.
460 Generally, the list representation is a list of length n_qubits,
461 where at each position a Pauli Operator is encoded as such:
462 I: -1
463 X: 0
464 Y: 1
465 Z: 2
467 Args:
468 op (Operation): Gate operation (PauliX, PauliY, PauliZ, or
469 Hermitian wrapping a multi-qubit Pauli tensor product).
470 n_qubits (int): number of qubits in the circuit
472 Returns:
473 jnp.ndarray[int]: List representation
474 """
475 pauli_repr = -jnp.ones(n_qubits, dtype=int)
477 _NAME_TO_IDX = {"PauliX": 0, "PauliY": 1, "PauliZ": 2}
479 if op.name in _NAME_TO_IDX:
480 pauli_repr = pauli_repr.at[op.wires[0]].set(_NAME_TO_IDX[op.name])
481 elif isinstance(op, PauliX):
482 pauli_repr = pauli_repr.at[op.wires[0]].set(0)
483 elif isinstance(op, PauliY):
484 pauli_repr = pauli_repr.at[op.wires[0]].set(1)
485 elif isinstance(op, PauliZ):
486 pauli_repr = pauli_repr.at[op.wires[0]].set(2)
487 else:
488 # Multi-qubit case: decompose via pauli_string_from_operation
489 from qml_essentials.operations import pauli_string_from_operation
491 pauli_str = pauli_string_from_operation(op)
492 char_to_idx = {"X": 0, "Y": 1, "Z": 2, "I": -1}
493 for i, (wire, ch) in enumerate(zip(op.wires, pauli_str)):
494 idx = char_to_idx.get(ch, -1)
495 if idx >= 0:
496 pauli_repr = pauli_repr.at[wire].set(idx)
498 return pauli_repr
500 def is_commuting(self, pauli: jnp.ndarray[int]) -> bool:
501 """
502 Computes if this Pauli commutes with another Pauli operator.
503 This computation is based on the fact that The commutator is zero
504 if and only if the number of anticommuting single-qubit Paulis is
505 even.
507 Args:
508 pauli (jnp.ndarray[int]): List representation of another Pauli
510 Returns:
511 bool: If the current and other Pauli are commuting.
512 """
513 anticommutator = jnp.where(
514 pauli < 0,
515 False,
516 jnp.where(
517 self.list_repr < 0,
518 False,
519 jnp.where(self.list_repr == pauli, False, True),
520 ),
521 )
522 return not (jnp.sum(anticommutator) % 2)
524 def tensor(self, pauli: jnp.ndarray[int]) -> FourierTree.PauliOperator:
525 """
526 Compute tensor product between the current Pauli and a given list
527 representation of another Pauli operator.
529 Args:
530 pauli (jnp.ndarray[int]): List representation of Pauli
532 Returns:
533 FourierTree.PauliOperator: New Pauli operator object, which
534 contains the tensor product
535 """
536 diff = (pauli - self.list_repr + 3) % 3
537 phase = jnp.where(
538 self.list_repr < 0,
539 1.0,
540 jnp.where(
541 pauli < 0,
542 1.0,
543 jnp.where(
544 diff == 2,
545 1.0j,
546 jnp.where(diff == 1, -1.0j, 1.0),
547 ),
548 ),
549 )
551 obs = jnp.where(
552 self.list_repr < 0,
553 pauli,
554 jnp.where(
555 pauli < 0,
556 self.list_repr,
557 jnp.where(
558 diff == 2,
559 (self.list_repr + 1) % 3,
560 jnp.where(diff == 1, (self.list_repr + 2) % 3, -1),
561 ),
562 ),
563 )
564 phase = self.phase * jnp.prod(phase)
565 return FourierTree.PauliOperator(
566 obs, phase=phase, n_qubits=obs.size, is_init=False, is_observable=True
567 )
569 def __init__(self, model: Model):
570 """
571 Tree initialisation, based on the Pauli-Clifford representation of a model.
572 Currently, only one input feature is supported.
574 **Usage**:
575 ```
576 # initialise a model
577 model = Model(...)
579 # initialise and build FourierTree
580 tree = FourierTree(model)
582 # get expectaion value
583 exp = tree()
585 # Get spectrum (for each observable, we have one list element)
586 coeff_list, freq_list = tree.spectrum()
587 ```
589 Args:
590 model (Model): The Model, for which to build the tree
591 """
592 self.model = model
593 self.tree_roots = None
595 inputs = self.model._inputs_validation([1.0])
597 # Record the circuit tape using yaqsi's tape recording
598 raw_tape = self.model.script._record(params=model.params, inputs=inputs)
600 # Build observables from the model's output_qubit configuration
601 _, obs_list = self.model._build_obs()
603 quantum_tape = PauliCircuit.from_parameterised_circuit(
604 raw_tape, observables=obs_list
605 )
607 self.parameters = [jnp.squeeze(p) for p in quantum_tape.get_parameters()]
609 self.input_indices, self.all_input_indices = quantum_tape.get_input_indices()
611 self.observables = self._encode_observables(quantum_tape.observables)
613 pauli_rot = FourierTree.PauliOperator(
614 quantum_tape.operations[0],
615 self.model.n_qubits,
616 )
617 self.pauli_rotations = [pauli_rot]
618 for op in quantum_tape.operations[1:]:
619 pauli_rot = FourierTree.PauliOperator(
620 op, self.model.n_qubits, pauli_rot.xy_indices
621 )
622 self.pauli_rotations.append(pauli_rot)
624 self.tree_roots = self.build()
625 self.leafs: List[List[FourierTree.TreeLeaf]] = self._get_tree_leafs()
627 def __call__(
628 self,
629 params: Optional[jnp.ndarray] = None,
630 inputs: Optional[jnp.ndarray] = None,
631 **kwargs,
632 ) -> jnp.ndarray:
633 """
634 Evaluates the Fourier tree via sine-cosine terms sum. This is
635 equivalent to computing the expectation value of the observables with
636 respect to the corresponding circuit.
638 Args:
639 params (Optional[jnp.ndarray], optional): Parameters of the model.
640 Defaults to None.
641 inputs (Optional[jnp.ndarray], optional): Inputs to the circuit.
642 Defaults to None.
644 Returns:
645 jnp.ndarray: Expectation value of the tree.
647 Raises:
648 NotImplementedError: When using other "execution_type" as expval.
649 NotImplementedError: When using "noise_params"
652 """
653 params = (
654 self.model._params_validation(params)
655 if params is not None
656 else self.model.params
657 )
658 inputs = (
659 self.model._inputs_validation(inputs)
660 if inputs is not None
661 else self.model._inputs_validation(1.0)
662 )
664 if kwargs.get("execution_type", "expval") != "expval":
665 raise NotImplementedError(
666 f'Currently, only "expval" execution type is supported when '
667 f"building FourierTree. Got {kwargs.get('execution_type', 'expval')}."
668 )
669 if kwargs.get("noise_params", None) is not None:
670 raise NotImplementedError(
671 "Currently, noise is not supported when building FourierTree."
672 )
674 # Record the circuit tape using yaqsi's tape recording
675 raw_tape = self.model.script._record(params=self.model.params, inputs=inputs)
677 # Build observables from the model's output_qubit configuration
678 _, obs_list = self.model._build_obs()
680 quantum_tape = PauliCircuit.from_parameterised_circuit(
681 raw_tape, observables=obs_list
682 )
684 self.parameters = [jnp.squeeze(p) for p in quantum_tape.get_parameters()]
686 results = jnp.zeros(len(self.tree_roots))
687 for i, root in enumerate(self.tree_roots):
688 results = results.at[i].set(jnp.real(root.evaluate(self.parameters)))
690 if kwargs.get("force_mean", False):
691 return jnp.mean(results)
692 else:
693 return results
695 def build(self) -> List[CoefficientsTreeNode]:
696 """
697 Creates the coefficient tree, i.e. it creates and initialises the tree
698 nodes.
699 Leafs can be obtained separately in _get_tree_leafs, once the tree is
700 set up.
702 Returns:
703 List[CoefficientsTreeNode]: The list of root nodes (one root for
704 each observable).
705 """
706 tree_roots = []
707 pauli_rotation_idx = len(self.pauli_rotations) - 1
708 for obs in self.observables:
709 root = self._create_tree_node(obs, pauli_rotation_idx)
710 tree_roots.append(root)
711 return tree_roots
713 def _encode_observables(
714 self, tape_obs: List[Operation]
715 ) -> List[FourierTree.PauliOperator]:
716 """
717 Encodes observables from tape as FourierTree.PauliOperator
718 utility objects.
720 Args:
721 tape_obs (List[Operation]): Observable operations
723 Returns:
724 List[FourierTree.PauliOperator]: List of Pauli operators
725 """
726 observables = []
727 for obs in tape_obs:
728 observables.append(
729 FourierTree.PauliOperator(obs, self.model.n_qubits, is_observable=True)
730 )
731 return observables
733 def _get_tree_leafs(self) -> List[List[TreeLeaf]]:
734 """
735 Obtain all Leaf Nodes with its sine- and cosine terms.
737 Returns:
738 List[List[TreeLeaf]]: For each observable (root), the list of leaf
739 nodes.
740 """
741 leafs = []
742 for root in self.tree_roots:
743 sin_list = jnp.zeros(len(self.parameters), dtype=jnp.int32)
744 cos_list = jnp.zeros(len(self.parameters), dtype=jnp.int32)
745 leafs.append(root.get_leafs(sin_list, cos_list, []))
746 return leafs
748 def get_spectrum(
749 self, force_mean: bool = False
750 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
751 """
752 Computes the Fourier spectrum for the tree, consisting of the
753 frequencies and its corresponding coefficinets.
754 If the frag force_mean was set in the constructor, the mean coefficient
755 over all observables (roots) are computed.
757 Args:
758 force_mean (bool, optional): Whether to average over multiple
759 observables. Defaults to False.
761 Returns:
762 Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
763 - List of frequencies, one list for each observable (root).
764 - List of corresponding coefficents, one list for each
765 observable (root).
766 """
767 parameter_indices = [
768 i for i in range(len(self.parameters)) if i not in self.all_input_indices
769 ]
771 coeffs = []
772 for leafs in self.leafs:
773 freq_terms = defaultdict(np.complex128)
774 for input_idx in self.input_indices:
775 for leaf in leafs:
776 leaf_factor, s, c = self._compute_leaf_factors(
777 leaf, parameter_indices, input_idx
778 )
780 for a in range(s + 1):
781 for b in range(c + 1):
782 comb = math.comb(s, a) * math.comb(c, b) * (-1) ** (s - a)
783 freq_terms[2 * a + 2 * b - s - c] += comb * leaf_factor
785 coeffs.append(freq_terms)
787 frequencies, coefficients = self._freq_terms_to_coeffs(coeffs, force_mean)
788 return coefficients, frequencies
790 def _freq_terms_to_coeffs(
791 self, coeffs: List[Dict[int, jnp.ndarray]], force_mean: bool
792 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
793 """
794 Given a list of dictionaries of the form:
795 [
796 {
797 freq_obs1_1: coeff1,
798 freq_obs1_2: coeff2,
799 ...
800 },
801 {
802 freq_obs2_1: coeff3,
803 freq_obs2_2: coeff4,
804 ...
805 }
806 ...
807 ],
808 Compute two separate lists of frequencies and coefficients.
809 such that:
810 freqs: [
811 [freq_obs1_1, freq_obs1_1, ...],
812 [freq_obs2_1, freq_obs2_1, ...],
813 ...
814 ]
815 coeffs: [
816 [coeff1, coeff2, ...],
817 [coeff3, coeff4, ...],
818 ...
819 ]
821 If force_mean is set length of the resulting frequency and coefficent
822 list is 1.
824 Args:
825 coeffs (List[Dict[int, jnp.ndarray]]): Frequency->Coefficients
826 dictionary list, one dict for each observable (root).
827 force_mean (bool): Whether to average coefficients over multiple
828 observables.
830 Returns:
831 Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
832 - List of frequencies, one list for each observable (root).
833 - List of corresponding coefficents, one list for each
834 observable (root).
835 """
836 frequencies = []
837 coefficients = []
838 if force_mean:
839 all_freqs = sorted(set([f for c in coeffs for f in c.keys()]))
840 coefficients.append(
841 jnp.array(
842 [
843 jnp.mean(jnp.array([c.get(f, 0.0) for c in coeffs]))
844 for f in all_freqs
845 ]
846 )
847 )
848 frequencies.append(jnp.array(all_freqs))
849 else:
850 for freq_terms in coeffs:
851 freq_terms = dict(sorted(freq_terms.items()))
852 frequencies.append(jnp.array(list(freq_terms.keys())))
853 coefficients.append(jnp.array(list(freq_terms.values())))
854 return frequencies, coefficients
856 def _compute_leaf_factors(
857 self,
858 leaf: TreeLeaf,
859 parameter_indices: List[int],
860 input_idx: int,
861 ) -> Tuple[float, int, int]:
862 """
863 Computes the constant coefficient factor for each leaf.
864 Additionally sine and cosine contributions of the input parameters for
865 this leaf are returned, which are required to obtain the corresponding
866 frequencies.
868 Args:
869 leaf (TreeLeaf): The leaf for which to compute the factor.
870 parameter_indices (List[int]): Variational parameter indices.
872 Returns:
873 Tuple[float, int, int]:
874 - float: the constant factor for the leaf
875 - int: number of sine contributions of the input
876 - int: number of cosine contributions of the input
877 """
878 leaf_factor = 1.0
879 for i in parameter_indices:
880 interm_factor = (
881 jnp.cos(self.parameters[i]) ** leaf.cos_indices[i]
882 * (1j * jnp.sin(self.parameters[i])) ** leaf.sin_indices[i]
883 )
884 leaf_factor = leaf_factor * interm_factor
886 # Get number of sine and cosine factors to which the input contributes
887 c = jnp.sum(
888 jnp.array([leaf.cos_indices[k] for k in self.input_indices[input_idx]])
889 )
890 s = jnp.sum(
891 jnp.array([leaf.sin_indices[k] for k in self.input_indices[input_idx]])
892 )
894 leaf_factor = leaf.term * leaf_factor * 0.5 ** (s + c)
896 return leaf_factor, int(s), int(c)
898 def _early_stopping_possible(
899 self, pauli_rotation_idx: int, observable: FourierTree.PauliOperator
900 ):
901 """
902 Checks if a node for an observable can be discarded as all expecation
903 values that can result through further branching are zero.
904 The method is mentioned in the paper by Nemkov et al.: If the one-hot
905 encoded indices for X/Y operations in the Pauli-rotation generators are
906 a basis for that of the observable, the node must be processed further.
907 If not, it can be discarded.
909 Args:
910 pauli_rotation_idx (int): Index of remaining Pauli rotation gates.
911 Gates itself are attributes of the class.
912 observable (FourierTree.PauliOperator): Current observable
913 """
914 xy_indices_obs = jnp.logical_or(
915 observable.xy_indices, self.pauli_rotations[pauli_rotation_idx].xy_indices
916 ).all()
918 return not xy_indices_obs
920 def _create_tree_node(
921 self,
922 observable: FourierTree.PauliOperator,
923 pauli_rotation_idx: int,
924 parameter_idx: Optional[int] = None,
925 is_sine: bool = False,
926 is_cosine: bool = False,
927 ) -> Optional[CoefficientsTreeNode]:
928 """
929 Builds the Fourier-Tree according to the algorithm by Nemkov et al.
931 Args:
932 observable (FourierTree.PauliOperator): Current observable
933 pauli_rotation_idx (int): Index of remaining Pauli rotation gates.
934 Gates itself are attributes of the class.
935 parameter_idx (Optional[int]): Index of the current parameter.
936 Parameters itself are attributes of the class.
937 is_sine (bool): If the current node is a sine (left) node.
938 is_cosine (bool): If the current node is a cosine (right) node.
940 Returns:
941 Optional[CoefficientsTreeNode]: The resulting node. Children are set
942 recursively. The top level receives the tree root.
943 """
944 if self._early_stopping_possible(pauli_rotation_idx, observable):
945 return None
947 # remove commuting paulis
948 while pauli_rotation_idx >= 0:
949 last_pauli = self.pauli_rotations[pauli_rotation_idx]
950 if not observable.is_commuting(last_pauli.list_repr):
951 break
952 pauli_rotation_idx -= 1
953 else: # leaf
954 return FourierTree.CoefficientsTreeNode(
955 parameter_idx, observable, is_sine, is_cosine
956 )
958 last_pauli = self.pauli_rotations[pauli_rotation_idx]
960 left = self._create_tree_node(
961 observable,
962 pauli_rotation_idx - 1,
963 pauli_rotation_idx,
964 is_cosine=True,
965 )
967 next_observable = self._create_new_observable(last_pauli.list_repr, observable)
968 right = self._create_tree_node(
969 next_observable,
970 pauli_rotation_idx - 1,
971 pauli_rotation_idx,
972 is_sine=True,
973 )
975 return FourierTree.CoefficientsTreeNode(
976 parameter_idx,
977 observable,
978 is_sine,
979 is_cosine,
980 left,
981 right,
982 )
984 def _create_new_observable(
985 self, pauli: jnp.ndarray[int], observable: FourierTree.PauliOperator
986 ) -> FourierTree.PauliOperator:
987 """
988 Utility function to obtain the new observable for a tree node, if the
989 last Pauli and the observable do not commute.
991 Args:
992 pauli (jnp.ndarray[int]): The int array representation of the last
993 Pauli rotation in the operation sequence.
994 observable (FourierTree.PauliOperator): The current observable.
996 Returns:
997 FourierTree.PauliOperator: The new observable.
998 """
999 observable = observable.tensor(pauli)
1000 return observable
1003class FCC:
1004 @staticmethod
1005 def get_fcc(
1006 model: Model,
1007 n_samples: int,
1008 random_key: Optional[random.PRNGKey] = None,
1009 method: Optional[str] = "pearson",
1010 scale: Optional[bool] = False,
1011 weight: Optional[bool] = False,
1012 trim_redundant: Optional[bool] = True,
1013 **kwargs,
1014 ) -> float:
1015 """
1016 Shortcut method to get just the FCC.
1017 This includes
1018 1. What is done in `get_fourier_fingerprint`:
1019 1. Calculating the coefficients (using `n_samples`)
1020 2. Correlating the result from 1) using `method`
1021 3. Weighting the correlation matrix (if `weight` is True)
1022 4. Remove redundancies
1023 2. What is done in `calculate_fcc`:
1024 1. Absolute of the fingerprint
1025 2. Average
1027 Args:
1028 model (Model): The QFM model
1029 n_samples (int): Number of samples to calculate average of coefficients
1030 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1031 initialization. If None, uses the model's internal random key.
1032 method (Optional[str], optional): Correlation method. Defaults to "pearson".
1033 scale (Optional[bool], optional): Whether to scale the number of samples.
1034 Defaults to False.
1035 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1036 Defaults to False.
1037 trim_redundant (Optional[bool], optional): Whether to remove redundant
1038 correlations. Defaults to False.
1039 **kwargs: Additional keyword arguments for the model function.
1041 Returns:
1042 float: The FCC
1043 """
1044 fourier_fingerprint, _ = FCC.get_fourier_fingerprint(
1045 model,
1046 n_samples,
1047 random_key,
1048 method,
1049 scale,
1050 weight,
1051 trim_redundant=trim_redundant,
1052 **kwargs,
1053 )
1055 return FCC.calculate_fcc(fourier_fingerprint)
1057 def get_fourier_fingerprint(
1058 model: Model,
1059 n_samples: int,
1060 random_key: Optional[random.PRNGKey] = None,
1061 method: Optional[str] = "pearson",
1062 scale: Optional[bool] = False,
1063 weight: Optional[bool] = False,
1064 trim_redundant: Optional[bool] = True,
1065 nan_to_one: Optional[bool] = False,
1066 **kwargs,
1067 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
1068 """
1069 Shortcut method to get just the fourier fingerprint.
1070 This includes
1071 1. Calculating the coefficients (using `n_samples`)
1072 2. Correlating the result from 1) using `method`
1073 3. Weighting the correlation matrix (if `weight` is True)
1074 4. Remove redundancies (if `trim_redundant` is True)
1076 Args:
1077 model (Model): The QFM model
1078 n_samples (int): Number of samples to calculate average of coefficients
1079 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1080 initialization. If None, uses the model's internal random key.
1081 method (Optional[str], optional): Correlation method. Defaults to "pearson".
1082 scale (Optional[bool], optional): Whether to scale the number of samples.
1083 Defaults to False.
1084 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1085 Defaults to False.
1086 trim_redundant (Optional[bool], optional): Whether to remove redundant
1087 correlations. Defaults to True.
1088 nan_to_one (Optional[bool], optional): Whether to set nan to 1.
1089 Defaults to False.
1090 **kwargs: Additional keyword arguments for the model function.
1092 Returns:
1093 Tuple[jnp.ndarray, jnp.ndarray]: The fourier fingerprint
1094 and the frequency indices
1095 """
1096 _, coeffs, freqs = FCC._calculate_coefficients(
1097 model, n_samples, random_key, scale, **kwargs
1098 )
1100 fourier_fingerprint = FCC._correlate(coeffs.transpose(), method=method)
1102 if nan_to_one:
1103 # set nan to 1
1104 fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0
1106 # perform weighting if requested
1107 fourier_fingerprint = (
1108 FCC._weighting(fourier_fingerprint) if weight else fourier_fingerprint
1109 )
1111 if trim_redundant:
1112 mask = FCC._calculate_mask(freqs)
1114 # apply the mask on the fingerprint
1115 fourier_fingerprint = mask * fourier_fingerprint
1117 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1118 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1120 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1122 return fourier_fingerprint, freqs
1124 @staticmethod
1125 def calculate_fcc(
1126 fourier_fingerprint: jnp.ndarray,
1127 ) -> float:
1128 """
1129 Method to calculate the FCC based on an existing correlation matrix.
1130 Calculate absolute and then the average over this matrix.
1131 The Fingerprint can be obtained via `get_fourier_fingerprint`
1133 Args:
1134 coeff_coeff_correlation (jnp.ndarray): Correlation matrix of coefficients
1135 Returns:
1136 float: The FCC
1137 """
1138 # apply the mask on the fingerprint
1139 return jnp.nanmean(jnp.abs(fourier_fingerprint))
1141 def _calculate_mask(freqs: jnp.ndarray) -> jnp.ndarray:
1142 """
1143 Method to calculate a mask filtering out redundant elements
1144 of the Fourier correlation matrix, based on the provided frequency vector.
1145 It does so by 'simulating' the operations that would be performed
1146 by `_correlate`.
1148 Args:
1149 freqs (jnp.ndarray): Array of frequencies
1151 Returns:
1152 jnp.ndarray: The mask
1153 """
1154 # TODO: this part can be heavily optimized, by e.g. using a "positive_only"
1155 # flag when calculating the coefficients.
1156 # However this would change the numerical values
1157 # (while the order should be still the same).
1159 # disregard all the negativ frequencies
1160 freqs[freqs < 0] = jnp.nan
1161 # compute the outer product of the frequency vectors for arbitrary dimensions
1162 # or just use the existing frequency vector if it is 1D
1163 nd_freqs = (
1164 reduce(jnp.multiply, jnp.ix_(*freqs)) if len(freqs.shape) > 1 else freqs
1165 )
1166 # TODO: could prevent this if we're not using .squeeze()..
1168 # "simulate" what would happen on correlating the coefficients
1169 corr_freqs = jnp.outer(nd_freqs, nd_freqs)
1170 # mask all frequencies that are nan now
1171 # (i.e. all correlations with a negative frequency component)
1172 corr_mask = jnp.where(jnp.isnan(corr_freqs), corr_freqs, 1)
1173 # from this, disregard all the other redundant correlations (i.e. c_0_1 = c_1_0)
1174 corr_mask = corr_mask.at[jnp.triu_indices(corr_mask.shape[0], 0)].set(jnp.nan)
1176 return corr_mask
1178 @staticmethod
1179 def _calculate_coefficients(
1180 model: Model,
1181 n_samples: int,
1182 random_key: Optional[random.PRNGKey] = None,
1183 scale: bool = False,
1184 **kwargs,
1185 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
1186 """
1187 Calculates the Fourier coefficients of a given model
1188 using `n_samples`.
1189 Optionally, `noise_params` can be passed to perform noisy simulation.
1191 Args:
1192 model (Model): The QFM model
1193 n_samples (int): Number of samples to calculate average of coefficients
1194 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1195 initialization. If None, uses the model's internal random key.
1196 scale (bool, optional): Whether to scale the number of samples.
1197 Defaults to False.
1198 **kwargs: Additional keyword arguments for the model function.
1200 Returns:
1201 Tuple[jnp.ndarray, jnp.ndarray]: Parameters and Coefficients of size NxK
1202 """
1203 if n_samples > 0:
1204 if scale:
1205 total_samples = int(
1206 jnp.power(2, model.n_qubits) * n_samples * model.n_input_feat
1207 )
1208 log.info(f"Using {total_samples} samples.")
1209 else:
1210 total_samples = n_samples
1211 model.initialize_params(random_key, repeat=total_samples)
1212 else:
1213 total_samples = 1
1215 coeffs, freqs = Coefficients.get_spectrum(
1216 model, shift=True, trim=True, **kwargs
1217 )
1219 return model.params, coeffs, freqs
1221 @staticmethod
1222 def _correlate(mat: jnp.ndarray, method: str = "pearson") -> jnp.ndarray:
1223 """
1224 Correlates two arrays using `method`.
1225 Currently, `pearson` and `spearman` are supported.
1227 Args:
1228 mat (jnp.ndarray): Array of shape (N, K)
1229 method (str, optional): Correlation method. Defaults to "pearson".
1231 Raises:
1232 ValueError: If the method is not supported.
1234 Returns:
1235 jnp.ndarray: Correlation matrix of `a` and `b`.
1236 """
1237 assert len(mat.shape) >= 2, "Input matrix must have at least 2 dimensions"
1239 # Note that for the general n-D case, we have to flatten along
1240 # the first axis (last one is batch).
1241 # Note that the order here is important so we can easily filter out
1242 # negative coefficients later.
1243 # Consider the following example: [[1,2,3],[4,5,6],[7,8,9]]
1244 # we want to get [1, 4, 7, 2, 5, 8, 3, 6, 9]
1245 # such that after correlation, all positive indexed coefficients
1246 # will be in the bottom right quadrant
1247 if method == "pearson":
1248 result = FCC._pearson(mat.reshape(mat.shape[0], -1))
1249 # result = FCC._pearson(mat.reshape(mat.shape[-1], -1, order="F"))
1250 elif method == "spearman":
1251 result = FCC._spearman(mat.reshape(mat.shape[0], -1))
1252 # result = FCC._spearman(mat.reshape(mat.shape[-1], -1, order="F"))
1253 else:
1254 raise ValueError(
1255 f"Unknown correlation method: {method}. \
1256 Must be 'pearson' or 'spearman'."
1257 )
1259 return result
1261 @staticmethod
1262 def _pearson(
1263 mat: jnp.ndarray, cov: Optional[bool] = False, minp: Optional[int] = 1
1264 ) -> jnp.ndarray:
1265 """
1266 Based on Pandas correlation method as implemented here:
1267 https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx
1269 Compute Pearson correlation between columns of `mat`,
1270 permitting missing values (NaN or ±Inf).
1272 Args:
1273 mat : array_like, shape (N, K)
1274 Input data.
1275 minp : int, optional
1276 Minimum number of paired observations required to form a correlation.
1277 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1279 Returns:
1280 corr : ndarray, shape (K, K)
1281 Pearson correlation matrix.
1282 """
1284 mat = jnp.asarray(mat, dtype=jnp.float64)
1285 N, K = mat.shape
1287 # pre‐compute finite‐mask
1288 mask = jnp.isfinite(mat)
1290 # output
1291 result = np.empty((K, K), dtype=jnp.float64)
1293 # TODO: optimize in future iterations
1294 # loop over column‐pairs
1295 for i in range(K):
1296 for j in range(i + 1):
1297 # find rows where both columns are finite
1298 m = mask[:, i] & mask[:, j]
1299 n = jnp.count_nonzero(m)
1300 if n < minp:
1301 # too few pairs
1302 value = jnp.nan
1303 else:
1304 x = mat[m, i]
1305 y = mat[m, j]
1307 # compute means
1308 mean_x = x.mean()
1309 mean_y = y.mean()
1311 # demeaned data
1312 dx = x - mean_x
1313 dy = y - mean_y
1315 # sum of squares and cross‐prod
1316 ssx = jnp.dot(dx, dx)
1317 ssy = jnp.dot(dy, dy)
1318 cxy = jnp.dot(dx, dy)
1320 if cov:
1321 # sample covariance (denominator n−1)
1322 value = cxy / (n - 1) if n > 1 else jnp.nan
1323 else:
1324 # Pearson r = cov / (σx σy)
1325 denom = jnp.sqrt(ssx * ssy)
1326 if denom == 0.0:
1327 value = jnp.nan
1328 else:
1329 value = cxy / denom
1330 # clip numerical drift
1331 if value > 1.0:
1332 value = 1.0
1333 elif value < -1.0:
1334 value = -1.0
1336 result[i, j] = result[j, i] = value
1338 return result
1340 def _spearman(mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1341 """
1342 Based on Pandas correlation method as implemented here:
1343 https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx
1345 Compute Spearman correlation between columns of `mat`,
1346 permitting missing values (NaN or ±Inf).
1348 Args:
1349 mat : array_like, shape (N, K)
1350 Input data.
1351 minp : int, optional
1352 Minimum number of paired observations required to form a correlation.
1353 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1355 Returns:
1356 corr : ndarray, shape (K, K)
1357 Spearman correlation matrix.
1358 """
1359 N, K = mat.shape
1360 # trivial all-NaN answer if too few rows
1361 if N < minp:
1362 return jnp.full((K, K), jnp.nan, dtype=float)
1364 # mask of finite entries
1365 mask = jnp.isfinite(mat) # shape (N, K), dtype=bool
1367 # precompute ranks column-wise ignoring NaNs
1368 ranks = np.full((N, K), jnp.nan, dtype=float)
1369 for j in range(K):
1370 valid = mask[:, j]
1371 if valid.any():
1372 # rankdata by default gives average ranks for ties
1373 ranks[valid, j] = rankdata(mat[valid, j], method="average")
1375 # allocate result
1376 result = np.empty((K, K), dtype=float)
1378 # TODO: optimize in future iterations
1379 # loop lower triangle (including diagonal)
1380 for i in range(K):
1381 for j in range(i + 1):
1382 # find rows where both columns are finite
1383 valid = mask[:, i] & mask[:, j]
1384 nobs = valid.sum()
1386 if nobs < minp:
1387 rho = jnp.nan
1388 else:
1389 xi = ranks[valid, i]
1390 yj = ranks[valid, j]
1391 # subtract means
1392 xi = xi - xi.mean()
1393 yj = yj - yj.mean()
1394 num = jnp.dot(xi, yj)
1395 den = jnp.sqrt(jnp.dot(xi, xi) * jnp.dot(yj, yj))
1396 rho = num / den if den > 0 else jnp.nan
1398 result[i, j] = rho
1399 result[j, i] = rho
1401 return result
1403 @staticmethod
1404 def _weighting(fourier_fingerprint: jnp.ndarray) -> jnp.ndarray:
1405 """
1406 Performs weighting on the given correlation matrix.
1407 Here, low-frequent coefficients are weighted more heavily.
1409 Args:
1410 correlation (jnp.ndarray): Correlation matrix
1411 """
1412 # TODO: in Future iterations, this can be optimized by computing
1413 # on the trimmed matrix instead.
1415 assert (
1416 fourier_fingerprint.shape[0] % 2 != 0
1417 and fourier_fingerprint.shape[1] % 2 != 0
1418 ), "Correlation matrix must have odd dimensions. \
1419 Hint: use `trim` argument when calling `get_spectrum`."
1420 assert (
1421 fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1]
1422 ), "Correlation matrix must be square."
1424 def quadrant_to_matrix(a: jnp.ndarray) -> jnp.ndarray:
1425 """
1426 Transforms [[1,2],[3,4]] to
1427 [[1,2,1],[3,4,3],[1,2,1]]
1429 Args:
1430 a (jnp.ndarray): _description_
1432 Returns:
1433 jnp.ndarray: _description_
1434 """
1435 # rotates a from [[1,2],[3,4]] to [[3,4],[1,2]]
1436 a_rot = jnp.rot90(a)
1437 # merge the two matrices
1438 left = jnp.concat([a, a_rot])
1439 # merges left and right (left flipped)
1440 b = jnp.concat(
1441 [left, jnp.flip(left)],
1442 axis=1,
1443 )
1444 # remove the middle column and row
1445 return jnp.delete(
1446 jnp.delete(b, (b.shape[0] // 2), axis=0), (b.shape[1] // 2), axis=1
1447 )
1449 nc = fourier_fingerprint.shape[0] // 2 + 1
1450 weights = jnp.mgrid[0:nc:1, 0:nc:1].sum(axis=0) / ((nc - 1) * 2)
1451 weights_matrix = quadrant_to_matrix(weights)
1453 return fourier_fingerprint * weights_matrix
1454 raise NotImplementedError("Weighting method is not implemented")
1457class Datasets:
1458 @staticmethod
1459 def generate_fourier_series(
1460 random_key: random.PRNGKey,
1461 model: Model,
1462 coefficients_min: float = 0.0,
1463 coefficients_max: float = 1.0,
1464 zero_centered: bool = False,
1465 ) -> jnp.ndarray:
1466 """
1467 Generates the Fourier series representation of a function.
1468 It uses the `model.frequencies` property to retrieve the frequency
1469 information. This ensures that the resulting Fourier series is
1470 compatible with the model.
1472 This function is capable of generating $D$-dimensional Fourier series
1473 (again defined by `model.n_input_feat`).
1474 The highest frequency $N$ is retrieved per dimension.
1476 Samples of the Fourier coefficients are drawn from a uniform circle.
1478 Args:
1479 random_key (random.PRNGKey): Random number key for JAX.
1480 model (Model): The quantum circuit model.
1481 coefficients_min (float, optional): Minimum value for the coefficients.
1482 Defaults to 0.0.
1483 coefficients_max (float, optional): Maximum value for the coefficients.
1484 Defaults to 1.0.
1485 zero_centered (bool, optional): Whether to zero-center the coefficients.
1486 Defaults to False.
1488 Returns:
1489 jnp.ndarray: Input domain samples with shape ((N,)*D, D)
1490 jnp.ndarray: Fourier series values with shape ((N,)*D)
1491 jnp.ndarray: Fourier coefficients with shape ((N,)*D)
1493 """
1494 # TODO: the following code can be considered to
1495 # capturing a truly random spectrum.
1496 # add some constraints on the spectrum, i.e. not fully
1498 # Note: one key observation for understanding the following code is,
1499 # that instead of wrapping your head around symmetries in multi-
1500 # dimensional coefficient matrices, one can simply look at the flattened
1501 # version of such a matrix and reshape later. It just works out.
1503 # going from [0, 2pi] with the resolution required for highest frequency
1504 # permute with input dimensionality to get an n-d grid of domain samples
1505 # the output shape comes from the fact that want to create a "coordinate system"
1506 domain_samples_per_input_dim = jnp.stack(
1507 jnp.meshgrid(
1508 *[jnp.arange(0, 2 * jnp.pi, 2 * jnp.pi / d) for d in model.degree]
1509 )
1510 ).T.reshape(-1, model.n_input_feat)
1512 # generate the frequency indices for each dimension.
1513 # this will have the same shape as the domain samples
1514 frequencies = jnp.stack(jnp.meshgrid(*model.frequencies)).T.reshape(
1515 -1, model.n_input_feat
1516 )
1518 # using the frequency information, sample coefficients for each dimension
1519 # shape: (input_dims, n_freqs_per_input_dim // 2 + 1)
1520 coefficients = Datasets.uniform_circle(
1521 random_key,
1522 low=coefficients_min,
1523 high=coefficients_max,
1524 size=math.prod(model.degree) // 2 + 1,
1525 )
1527 # zero center (first coeff = 0)
1528 # we can assume the first coeff is the offset, because we're dealing
1529 # with a non-symmetric spectrum here
1530 if zero_centered:
1531 coefficients = coefficients.at[0].set(0.0)
1532 else:
1533 coefficients = coefficients.at[0].set(coefficients[0].real)
1535 # ensure symmetry (here, non_negative_ is removed!),
1536 # giving us the full coefficients vector
1537 coefficients = jnp.concat(
1538 [
1539 jnp.flip(coefficients[..., 1:]).conjugate(),
1540 coefficients,
1541 ],
1542 axis=-1,
1543 )
1545 # Vectorized version of $f(x) = \sum_{n=0}^{N-1} c_n * e^{i * \omega_n * x}$
1546 # it takes into account the input dimension, i.e. the output is a matrix
1547 # normalization uses the n_freqs component of the coefficients
1548 values = jnp.real(
1549 (
1550 jnp.exp(1j * (domain_samples_per_input_dim @ frequencies.T))
1551 * coefficients
1552 ).sum(axis=1)
1553 / coefficients.size
1554 )
1556 # return all the information we have
1557 return [
1558 domain_samples_per_input_dim.reshape(*model.degree, -1),
1559 values.reshape(model.degree),
1560 coefficients.reshape(model.degree),
1561 ]
1563 @staticmethod
1564 def uniform_circle(
1565 random_key: random.PRNGKey,
1566 size: Union[jnp.ndarray, List, int],
1567 low=0.0,
1568 high=1.0,
1569 ):
1570 """
1571 Random number generator for complex numbers sampled inside the unit circle
1573 Args:
1574 random_key (random.PRNGKey): Random number key for JAX.
1575 size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
1576 the first dimension will be the number of dimensions.
1577 low (float, optional): Minimum Radius. Defaults to 0.0.
1578 high (float, optional): Maximum Radius. Defaults to 1.0.
1580 Returns
1581 jnp.ndarray: Array of complex numbers with shape of `size`
1582 """
1584 if isinstance(size, int):
1585 size = jnp.array([size])
1587 random_key, random_key1 = random.split(random_key)
1588 return jnp.sqrt(
1589 random.uniform(random_key, size, minval=low, maxval=high)
1590 ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))