Coverage for qml_essentials / coefficients.py: 97%
684 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-08-21 08:21 +0000
1from __future__ import annotations
2import sys
3import math
4import warnings
5import itertools
6from collections import defaultdict
7import jax.numpy as jnp
8from jax import random
9import numpy as np
10from scipy.stats import rankdata
11from functools import reduce, lru_cache
12from typing import List, Tuple, Optional, Any, Dict, Union
14from qml_essentials.model import Model
15from qml_essentials.pauli import PauliCircuit
16from qml_essentials.operations import PauliWord
18import logging
20log = logging.getLogger(__name__)
23class Coefficients:
24 @classmethod
25 def get_spectrum(
26 cls,
27 model: Model,
28 mfs: int = 1,
29 mts: int = 1,
30 shift=False,
31 trim=False,
32 numerical_cap: Optional[float] = -1,
33 **kwargs,
34 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
35 """
36 Extracts the coefficients of a given model using a FFT (jnp-fft).
38 Note that the coefficients are complex numbers, but the imaginary part
39 of the coefficients should be very close to zero, since the expectation
40 values of the Pauli operators are real numbers.
42 It can perform oversampling in both the frequency and time domain
43 using the `mfs` and `mts` arguments.
45 Args:
46 model (Model): The model to sample.
47 mfs (int): Multiplicator for the highest frequency. Default is 2.
48 mts (int): Multiplicator for the number of time samples. Default is 1.
49 shift (bool): Whether to apply jnp-fftshift. Default is False.
50 trim (bool): Whether to remove the Nyquist frequency if spectrum is even.
51 Default is False.
52 numerical_cap (Optional[float]): Numerical cap for the coefficients.
53 If positive, coefficients with magnitude below the cap are
54 zeroed and, for a single input feature, frequencies that
55 vanish entirely are removed from both `coeffs` and `freqs`.
56 kwargs (Any): Additional keyword arguments for the model function.
58 Returns:
59 Tuple[jnp.ndarray, jnp.ndarray]: Tuple containing the coefficients
60 and frequencies.
61 """
62 kwargs.setdefault("force_mean", True)
63 kwargs.setdefault("execution_type", "expval")
65 coeffs, freqs = cls._fourier_transform(model, mfs=mfs, mts=mts, **kwargs)
67 if not jnp.isclose(jnp.sum(coeffs).imag, 0.0, atol=1.0e-6):
68 raise ValueError(
69 f"Spectrum is not real. Imaginary part of coefficients is:\
70 {jnp.sum(coeffs).imag}"
71 )
73 if trim:
74 for ax in range(model.n_input_feat):
75 if coeffs.shape[ax] % 2 == 0:
76 coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
77 freqs = [np.delete(freq, len(freq) // 2, axis=ax) for freq in freqs]
79 if shift:
80 coeffs = jnp.fft.fftshift(coeffs, axes=list(range(model.n_input_feat)))
81 freqs = np.fft.fftshift(freqs)
83 if numerical_cap > 0:
84 # set coeffs below threshold to zero
85 coeffs = jnp.where(
86 jnp.abs(coeffs) < numerical_cap,
87 jnp.zeros_like(coeffs),
88 coeffs,
89 )
91 # Drop frequencies whose coefficients vanish entirely after
92 # capping, so the returned spectrum reflects only the surviving
93 # frequencies. Well-defined only for a single (1-D) frequency
94 # axis; for multi-dim input the rectangular grid is left intact.
95 if model.n_input_feat == 1:
96 if coeffs.ndim == 1:
97 surviving = coeffs != 0
98 else:
99 surviving = jnp.any(coeffs != 0, axis=tuple(range(1, coeffs.ndim)))
100 coeffs = coeffs[surviving]
101 freqs = [freqs[0][surviving]]
103 if len(freqs) == 1:
104 freqs = freqs[0]
106 return coeffs, freqs
108 @classmethod
109 def _fourier_transform(
110 cls, model: Model, mfs: int, mts: int, **kwargs: Any
111 ) -> jnp.ndarray:
112 # Create a frequency vector with as many frequencies as model degrees,
113 # oversampled by mfs
114 n_freqs: jnp.ndarray = jnp.array(
115 [mfs * model.degree[i] for i in range(model.n_input_feat)]
116 )
118 start, stop, step = 0, 2 * mts * jnp.pi, 2 * jnp.pi / n_freqs
119 # Stretch according to the number of frequencies
120 inputs: List = [
121 jnp.arange(start, stop, step[i]) for i in range(model.n_input_feat)
122 ]
124 # permute with input dimensionality
125 nd_inputs = jnp.array(
126 jnp.meshgrid(*[inputs[i] for i in range(model.n_input_feat)])
127 ).T.reshape(-1, model.n_input_feat)
129 # Output vector is not necessarily the same length as input
130 outputs = model(inputs=nd_inputs, **kwargs)
131 outputs = outputs.reshape(
132 *[inputs[i].shape[0] for i in range(model.n_input_feat)], -1
133 ).squeeze()
135 coeffs = jnp.fft.fftn(outputs, axes=list(range(model.n_input_feat)))
137 freqs = [
138 jnp.fft.fftfreq(int(mts * n_freqs[i]), 1 / n_freqs[i])
139 for i in range(model.n_input_feat)
140 ]
141 # freqs = jnp.fft.fftfreq(mts * n_freqs, 1 / n_freqs)
143 # TODO: this could cause issues with multidim input
144 # FIXME: account for different frequencies in multidim input scenarios
145 # Run the fft and rearrange +
146 # normalize the output (using product if multidim)
147 return (
148 coeffs / math.prod(outputs.shape[0 : model.n_input_feat]),
149 freqs,
150 )
152 @classmethod
153 def get_psd(cls, coeffs: jnp.ndarray) -> jnp.ndarray:
154 """
155 Calculates the power spectral density (PSD) from given Fourier coefficients.
157 Args:
158 coeffs (jnp.ndarray): The Fourier coefficients.
160 Returns:
161 jnp.ndarray: The power spectral density.
162 """
163 # TODO: if we apply trim=True in advance, this will be slightly wrong..
165 def abs2(x):
166 return x.real**2 + x.imag**2
168 scale = 2.0 / (len(coeffs) ** 2)
169 return scale * abs2(coeffs)
171 @classmethod
172 def evaluate_Fourier_series(
173 cls,
174 coefficients: jnp.ndarray,
175 frequencies: jnp.ndarray,
176 inputs: Union[jnp.ndarray, list, float],
177 ) -> float:
178 """
179 Evaluate the function value of a Fourier series at one point.
181 Args:
182 coefficients (jnp.ndarray): Coefficients of the Fourier series.
183 frequencies (jnp.ndarray): Corresponding frequencies.
184 inputs (jnp.ndarray): Point at which to evaluate the function.
185 Returns:
186 float: The function value at the input point.
187 """
188 coefficients = jnp.asarray(coefficients)
190 def flatten_grid(freq_axes):
191 freq_axes = [jnp.asarray(freq) for freq in freq_axes]
192 freq_grid = jnp.stack(jnp.meshgrid(*freq_axes, indexing="ij"), axis=-1)
193 flat_frequencies = freq_grid.reshape(-1, len(freq_axes))
194 flat_coefficients = coefficients.reshape(
195 flat_frequencies.shape[0], *coefficients.shape[len(freq_axes) :]
196 )
197 return flat_coefficients, flat_frequencies
199 if isinstance(frequencies, list):
200 flat_coefficients, flat_frequencies = flatten_grid(frequencies)
201 else:
202 frequencies = jnp.asarray(frequencies)
203 if frequencies.ndim == 1:
204 flat_frequencies = frequencies[:, jnp.newaxis]
205 flat_coefficients = coefficients.reshape(
206 flat_frequencies.shape[0], *coefficients.shape[1:]
207 )
208 else:
209 n_features, n_axis_freqs = frequencies.shape
210 is_axis_frequencies = (
211 coefficients.shape[:n_features] == (n_axis_freqs,) * n_features
212 )
214 if is_axis_frequencies:
215 flat_coefficients, flat_frequencies = flatten_grid(frequencies)
216 else:
217 flat_frequencies = frequencies
218 flat_coefficients = coefficients.reshape(
219 flat_frequencies.shape[0], *coefficients.shape[1:]
220 )
222 inputs = jnp.asarray(inputs)
223 if inputs.ndim == 0:
224 inputs = inputs.reshape(1, 1)
225 elif inputs.ndim == 1:
226 if flat_frequencies.shape[1] == 1:
227 inputs = inputs[:, jnp.newaxis]
228 elif inputs.shape[0] == flat_frequencies.shape[1]:
229 inputs = inputs[jnp.newaxis, :]
230 else:
231 inputs = jnp.repeat(
232 inputs[:, jnp.newaxis], flat_frequencies.shape[1], axis=1
233 )
234 exponents = jnp.exp(1j * (inputs @ flat_frequencies.T))
235 exp = jnp.tensordot(exponents, flat_coefficients, axes=([1], [0]))
237 return jnp.squeeze(jnp.real(exp))
240class FourierTree:
241 """
242 Sine-cosine tree representation for the algorithm by Nemkov et al.
244 Computes the analytical Fourier coefficients/frequencies of a Pauli-Clifford
245 circuit. The symbolic structure of the tree (which Pauli rotations
246 contribute sine/cosine factors to which leaf, and the leaf observables) is
247 built once in NumPy; the parameter-dependent coefficients are then obtained
248 with a small number of vectorised JAX operations, so the result remains
249 jittable / differentiable with respect to the model parameters.
251 The resulting spectrum is the d-dimensional set of frequency vectors,
252 where $d$ is the input dimensionality.
254 **Usage**:
255 ```
256 model = Model(...)
257 tree = FourierTree(model)
258 exp = tree() # expectation value
259 coeff_list, freq_list = tree.get_spectrum()
260 ```
261 """
263 def __init__(self, model: Model):
264 """
265 Tree initialisation, based on the Pauli-Clifford representation of a
266 model.
268 Args:
269 model (Model): The Model, for which to build the tree.
270 """
271 self.model = model
272 self.n_qubits = model.n_qubits
274 # A single (de-batched) parameter set drives the whole tree.
275 self._params = self._single_param_set(model.params)
277 # Canonical Pauli-Clifford structure, recorded once at a fixed base
278 # input. The base value is irrelevant to the structure (it only sets
279 # the rotation angles, not which Pauli words appear).
280 base_inputs = np.ones(model.n_input_feat)
281 operations, observables = self._build_canonical_tape(self._params, base_inputs)
283 self.parameters = [
284 jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
285 ]
286 self.n_params = len(self.parameters)
288 # Pauli generators of the (canonical) rotations, as symbolic words.
289 self.pauli_words: List[PauliWord] = [
290 PauliWord.from_operation(op, self.n_qubits) for op in operations
291 ]
293 # Cumulative X/Y support of the rotations[0..k] (for light-cone early
294 # stopping). cumulative_xy[k] is True on every qubit touched by an X/Y
295 # generator in any rotation up to index k.
296 self.cumulative_xy: List[np.ndarray] = []
297 running = np.zeros(self.n_qubits, dtype=bool)
298 for pw in self.pauli_words:
299 running = np.logical_or(running, pw.xy_mask)
300 self.cumulative_xy.append(running.copy())
302 # Observable Pauli words (one tree root each).
303 self.observable_words: List[PauliWord] = [
304 PauliWord.from_operation(obs, self.n_qubits) for obs in observables
305 ]
307 # Identify the input-encoding columns, their feature, and integer
308 # frequency scaling directly from the tape (no per-gate tagging). Sets
309 # ``input_indices``, ``all_input_indices``, ``input_scaling``,
310 # ``var_positions`` and ``features``.
311 self._detect_inputs(base_inputs)
313 # The explicit leaf structure is built lazily: for deep circuits the
314 # number of tree paths explodes combinatorially, while the canonical
315 # form above (and the merged-state support DP) remain cheap.
316 self._structure_built = False
318 def _ensure_structure(self) -> None:
319 """Build the explicit leaf/spectrum structure on first use."""
320 if not self._structure_built:
321 # Symbolic structure: per root (S, C, terms) leaf arrays ...
322 self._build_leaf_arrays()
323 # ... and the parameter-independent frequency/weight structure.
324 self._build_spectrum_structure()
325 self._structure_built = True
327 def _single_param_set(self, params) -> jnp.ndarray:
328 """De-batch the model parameters to the single set the tree describes.
330 Models can carry batched parameters (e.g. after FCC sampling); the tree
331 is defined for one set, so fall back to the first and warn.
332 """
333 params = jnp.asarray(params)
334 if params.ndim > 2 and params.shape[0] > 1:
335 warnings.warn(
336 f"FourierTree supports a single parameter set; using the first "
337 f"of {params.shape[0]} batched parameter sets.",
338 UserWarning,
339 )
340 params = params[0]
341 return params
343 def _build_canonical_tape(self, params, inputs):
344 """Record the circuit and transform it to Pauli-Clifford normal form.
346 Returns the ``(operations, observables)`` of the canonical circuit
347 (see :meth:`PauliCircuit.from_parameterised_circuit`).
348 """
349 params = self._single_param_set(params)
350 inputs = self.model._inputs_validation(inputs)
351 raw_tape = self.model.script._record(params=params, inputs=inputs)
352 _, obs_list = self.model._build_obs()
353 return PauliCircuit.from_parameterised_circuit(
354 raw_tape, observables=obs_list, n_qubits=self.n_qubits
355 )
357 def _canonical_parameters(self, inputs) -> np.ndarray:
358 """Recorded canonical rotation angles (1-D float array) for ``inputs``."""
359 operations, _ = self._build_canonical_tape(self._params, inputs)
360 return np.array(
361 [float(jnp.squeeze(p)) for p in PauliCircuit.get_parameters(operations)]
362 )
364 def _detect_inputs(self, base_inputs: np.ndarray) -> None:
365 r"""Infer the input-encoding columns directly from the tape (tag-free).
367 Each encoding rotation applies an angle :math:`\omega_k\,x_f` that is
368 linear in a single input feature :math:`x_f`, and Clifford commutation
369 only multiplies a rotation generator by :math:`\pm 1`. Every canonical
370 rotation angle is therefore an affine function of the inputs, so
371 perturbing one feature at a time and differencing the recorded angles
372 isolates exactly the columns that depend on it, together with the
373 signed scaling :math:`\omega_k` (rational for diagonal-Hamiltonian
374 encodings such as Golomb).
376 Sets :attr:`input_indices` (``{feature: [columns]}``),
377 :attr:`all_input_indices`, :attr:`input_scaling` (per column, ``1`` for
378 variational columns), :attr:`var_positions`, and :attr:`features`.
380 Raises:
381 NotImplementedError: If a rotation depends on more than one feature
382 (the tree requires single-feature encodings).
383 """
384 tol = 1e-6
385 d = self.model.n_input_feat
386 base = np.asarray(base_inputs, dtype=float)
387 p_base = np.array([float(p) for p in self.parameters])
389 # response[f, k] = d(angle_k) / d(x_f), the linear response of column k.
390 response = np.zeros((d, self.n_params))
391 for f in range(d):
392 step = base.copy()
393 step[f] += 1.0
394 response[f] = self._canonical_parameters(step) - p_base
396 input_indices: Dict[int, list] = defaultdict(list)
397 all_input_indices: List[int] = []
398 scaling = np.ones(self.n_params, dtype=np.float64)
399 for k in range(self.n_params):
400 feats = np.flatnonzero(np.abs(response[:, k]) > tol)
401 if feats.size == 0:
402 continue # variational column
403 if feats.size > 1:
404 raise NotImplementedError(
405 f"Rotation {k} depends on multiple input features "
406 f"{feats.tolist()}; the Fourier tree requires each encoding "
407 "rotation to be linear in a single feature."
408 )
409 f = int(feats[0])
410 omega = float(response[f, k])
411 # Scalings may be rational (e.g. diagonal-Hamiltonian / Golomb
412 # encodings, whose per-Pauli-string scalings are dyadic rationals);
413 # the integer-frequency model spectrum is recovered downstream.
414 input_indices[f].append(k)
415 all_input_indices.append(k)
416 scaling[k] = omega
418 self.input_indices = input_indices
419 self.all_input_indices = all_input_indices
420 self.input_scaling = scaling
421 input_set = set(all_input_indices)
422 self.var_positions = np.array(
423 [i for i in range(self.n_params) if i not in input_set], dtype=np.int64
424 )
425 # Ordered list of input feature keys (d-dimensional spectrum).
426 self.features = sorted(input_indices.keys())
428 # Symbolic tree construction (NumPy)
429 def _build_leaf_arrays(self) -> None:
430 """Collect the tree leaves for every root into integer count matrices.
432 For each root (observable) this produces:
433 - ``S``: (n_leaves, n_params) sine-factor counts per parameter,
434 - ``C``: (n_leaves, n_params) cosine-factor counts per parameter,
435 - ``terms``: (n_leaves,) complex leaf constants ``<0|O_leaf|0>``.
436 """
437 self.leaf_arrays: List[Tuple[np.ndarray, np.ndarray, np.ndarray]] = []
438 for obs_word in self.observable_words:
439 leaves: List[Tuple[np.ndarray, np.ndarray, complex]] = []
440 zeros = np.zeros(self.n_params, dtype=np.int64)
441 self._collect_leaves(
442 obs_word, self.n_params - 1, zeros.copy(), zeros.copy(), leaves
443 )
444 if leaves:
445 S = np.stack([leaf[0] for leaf in leaves])
446 C = np.stack([leaf[1] for leaf in leaves])
447 terms = np.array([leaf[2] for leaf in leaves], dtype=np.complex128)
448 else:
449 S = np.zeros((0, self.n_params), dtype=np.int64)
450 C = np.zeros((0, self.n_params), dtype=np.int64)
451 terms = np.zeros(0, dtype=np.complex128)
452 self.leaf_arrays.append((S, C, terms))
454 def _collect_leaves(
455 self,
456 observable: PauliWord,
457 pauli_idx: int,
458 sin_counts: np.ndarray,
459 cos_counts: np.ndarray,
460 leaves: List[Tuple[np.ndarray, np.ndarray, complex]],
461 ) -> None:
462 """Recursively enumerate the leaves of the coefficient tree.
464 The incoming sine/cosine factor (from the parent edge) is already
465 accumulated into ``sin_counts``/``cos_counts``. This fuses the tree
466 construction and leaf traversal of the original implementation into a
467 single NumPy pass (no per-node JAX scatter updates).
468 """
469 if self._early_stopping_possible(pauli_idx, observable):
470 return
472 # Skip trailing Pauli rotations that commute with the observable.
473 while pauli_idx >= 0:
474 last = self.pauli_words[pauli_idx]
475 if not observable.commutes_with(last):
476 break
477 pauli_idx -= 1
478 else: # leaf reached
479 term = observable.zero_expectation()
480 if term != 0:
481 leaves.append((sin_counts, cos_counts, term))
482 return
484 last = self.pauli_words[pauli_idx]
486 # Left child: cosine factor for this parameter, same observable.
487 cos_left = cos_counts.copy()
488 cos_left[pauli_idx] += 1
489 self._collect_leaves(
490 observable, pauli_idx - 1, sin_counts.copy(), cos_left, leaves
491 )
493 # Right child: sine factor, observable becomes P . O.
494 sin_right = sin_counts.copy()
495 sin_right[pauli_idx] += 1
496 self._collect_leaves(
497 last.compose(observable),
498 pauli_idx - 1,
499 sin_right,
500 cos_counts.copy(),
501 leaves,
502 )
504 def _early_stopping_possible(self, pauli_idx: int, observable: PauliWord) -> bool:
505 """Whether a node can be discarded (all reachable expectations vanish).
507 Mirrors the criterion of Nemkov et al. (light cone): a qubit on which
508 the observable carries an X/Y must be covered by an X/Y generator of
509 some remaining rotation (rotations[0..pauli_idx]); otherwise that X/Y can
510 never be rotated into a diagonal term and the whole node contributes
511 zero. Equivalently, the node survives iff every qubit is either I/Z in
512 the observable or covered by the cumulative rotation X/Y support.
513 """
514 obs_iz = np.logical_not(observable.xy_mask)
515 combined = np.logical_or(obs_iz, self.cumulative_xy[pauli_idx]).all()
516 return not bool(combined)
518 # Frequency / weight structure (NumPy, parameter independent)
519 def _build_spectrum_structure(self) -> None:
520 """Build, per root, the frequency vectors and the (n_freq, n_leaves)
521 weight matrix ``W`` such that ``coeffs = W @ (terms * variational)``.
522 """
523 self.freqs_per_root: List[np.ndarray] = []
524 self.weights_per_root: List[np.ndarray] = []
525 d = len(self.features)
527 for S, C, _ in self.leaf_arrays:
528 n_leaves = S.shape[0]
529 freq_to_col: Dict[tuple, np.ndarray] = defaultdict(
530 lambda: np.zeros(n_leaves, dtype=np.complex128)
531 )
532 for leaf in range(n_leaves):
533 # One expansion factor per *active* input column, each carrying
534 # its feature axis and integer frequency scaling. Per leaf a
535 # column contributes at most one sin/cos factor (square-free),
536 # but different columns of the same feature may carry different
537 # scalings, so they are expanded individually and convolved
538 # rather than aggregating counts (which would assume a common
539 # unit scaling).
540 col_factors: List[List[Tuple[int, float, float]]] = []
541 half_exp = 0
542 for axis, feat in enumerate(self.features):
543 for k in self.input_indices[feat]:
544 s = int(S[leaf, k])
545 c = int(C[leaf, k])
546 if s == 0 and c == 0:
547 continue
548 half_exp += s + c
549 w_k = float(self.input_scaling[k])
550 col_factors.append(
551 [
552 (axis, o * w_k, wt)
553 for o, wt in self._binomial_terms(s, c)
554 ]
555 )
556 half = 0.5**half_exp
558 if d == 0:
559 freq_to_col[(0,)][leaf] += half
560 continue
562 if not col_factors:
563 freq_to_col[(0,) * d][leaf] += half
564 continue
566 for combo in itertools.product(*col_factors):
567 omega = [0.0] * d
568 weight = half
569 for axis, o, wt in combo:
570 omega[axis] += o
571 weight *= wt
572 # Snap to a tolerance so dyadic-rational contributions that
573 # are numerically equal share a single frequency key.
574 key = tuple(round(v, 9) for v in omega)
575 freq_to_col[key][leaf] += weight
577 if freq_to_col:
578 omegas = sorted(freq_to_col.keys())
579 W = np.stack([freq_to_col[o] for o in omegas]) # (n_freq, n_leaves)
580 freqs = np.array(omegas, dtype=float) # (n_freq, d)
581 # Diagonal-Hamiltonian (Golomb) encodings produce rational
582 # per-rotation scalings but integer model frequencies; snap to
583 # int64 when every entry is integral so integer-encoding
584 # consumers are unchanged, otherwise keep the rational floats.
585 rounded = np.rint(freqs)
586 if np.all(np.abs(freqs - rounded) < 1e-6):
587 freqs = rounded.astype(np.int64)
588 else:
589 freqs = np.zeros((1, max(d, 1)), dtype=np.int64)
590 W = np.zeros((1, n_leaves), dtype=np.complex128)
592 # Collapse to 1-D frequency array for the single-feature case.
593 if freqs.shape[1] == 1:
594 freqs = freqs[:, 0]
595 self.freqs_per_root.append(freqs)
596 # Keep W in NumPy complex128: its entries are dyadic rationals
597 # (binomial weights x 0.5^k x i^m), which are exact in float64 --
598 # this allows exact symbolic zero-tests in get_exact_support.
599 self.weights_per_root.append(W)
601 @staticmethod
602 def _binomial_terms(s: int, c: int) -> List[Tuple[int, float]]:
603 """Expand ``cos^c (i sin)^s`` in ``e^{i omega x}`` (without the 0.5 factor).
605 Returns a list of ``(omega, weight)`` with
606 ``omega = 2a + 2b - s - c`` and ``weight = C(s,a) C(c,b) (-1)^{s-a}``.
607 """
608 terms = []
609 for a in range(s + 1):
610 for b in range(c + 1):
611 weight = math.comb(s, a) * math.comb(c, b) * (-1) ** (s - a)
612 terms.append((2 * a + 2 * b - s - c, float(weight)))
613 return terms
615 # Vectorised numeric evaluation (JAX)
616 @staticmethod
617 def _safe_pow(base: jnp.ndarray, exp: jnp.ndarray) -> jnp.ndarray:
618 """Elementwise ``base ** exp`` for real base and non-negative integer
619 exponents, correct for negative bases (avoids ``log`` of negatives).
621 Args:
622 base: real array of shape ``(n,)``.
623 exp: integer array of shape ``(n_leaves, n)``.
624 """
625 mag = jnp.abs(base)[None, :] ** exp
626 sign = jnp.where(exp % 2 == 0, 1.0, jnp.sign(base)[None, :])
627 return sign * mag
629 _I_POW = None # set lazily to jnp.array([1, 1j, -1, -1j])
631 def _leaf_factors(
632 self, S: np.ndarray, C: np.ndarray, columns: np.ndarray
633 ) -> jnp.ndarray:
634 """Per-leaf product ``prod_i cos(theta_i)^{C} (i sin(theta_i))^{S}`` over
635 the given parameter ``columns`` (vectorised over leaves).
636 """
637 if FourierTree._I_POW is None:
638 FourierTree._I_POW = jnp.array([1, 1j, -1, -1j])
640 if S.shape[0] == 0:
641 return jnp.zeros(0, dtype=jnp.complex64)
643 theta = jnp.stack([self.parameters[i] for i in columns])
644 S_sub = jnp.asarray(S[:, columns])
645 C_sub = jnp.asarray(C[:, columns])
647 cos_part = self._safe_pow(jnp.cos(theta), C_sub)
648 sin_mag = self._safe_pow(jnp.sin(theta), S_sub)
649 i_part = FourierTree._I_POW[S_sub % 4]
650 return jnp.prod(cos_part * sin_mag * i_part, axis=1)
652 def __call__(
653 self,
654 params: Optional[jnp.ndarray] = None,
655 inputs: Optional[jnp.ndarray] = None,
656 **kwargs,
657 ) -> jnp.ndarray:
658 """
659 Evaluate the expectation value(s) of the model's observables via the
660 sine-cosine tree (equivalent to the circuit expectation).
662 Args:
663 params (Optional[jnp.ndarray]): Model parameters. Defaults to the
664 model's parameters.
665 inputs (Optional[jnp.ndarray]): Inputs to the circuit. Defaults to 1.
667 Returns:
668 jnp.ndarray: Expectation value per observable (or their mean if
669 ``force_mean`` is set).
671 Raises:
672 NotImplementedError: For execution types other than "expval" or when
673 noise is requested.
674 """
675 params = (
676 self.model._params_validation(params)
677 if params is not None
678 else self.model.params
679 )
680 inputs = (
681 self.model._inputs_validation(inputs)
682 if inputs is not None
683 else self.model._inputs_validation(1.0)
684 )
686 if kwargs.get("execution_type", "expval") != "expval":
687 raise NotImplementedError(
688 f'Currently, only "expval" execution type is supported when '
689 f"building FourierTree. Got {kwargs.get('execution_type', 'expval')}."
690 )
691 if kwargs.get("noise_params", None) is not None:
692 raise NotImplementedError(
693 "Currently, noise is not supported when building FourierTree."
694 )
696 # Re-derive the (canonical) parameter values for the requested inputs;
697 # the tree structure (leaf arrays) is unchanged.
698 operations, _ = self._build_canonical_tape(params, inputs)
699 self.parameters = [
700 jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
701 ]
703 self._ensure_structure()
704 all_columns = np.arange(self.n_params, dtype=np.int64)
705 results = []
706 for S, C, terms in self.leaf_arrays:
707 factors = self._leaf_factors(S, C, all_columns)
708 results.append(jnp.real(jnp.sum(jnp.asarray(terms) * factors)))
709 results = jnp.array(results)
711 if kwargs.get("force_mean", False):
712 return jnp.mean(results)
713 return results
715 def get_spectrum(
716 self, force_mean: bool = False
717 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
718 """
719 Compute the Fourier spectrum (coefficients and frequencies) of the tree.
721 Args:
722 force_mean (bool, optional): Average the coefficients over all
723 observables (roots). Defaults to False.
725 Returns:
726 Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
727 - List of coefficients, one entry per observable (root).
728 - List of corresponding frequencies, one entry per root.
729 When ``force_mean`` is set, both lists have a single entry.
730 """
731 self._ensure_structure()
732 per_root_coeffs: List[jnp.ndarray] = []
733 for (S, C, terms), W in zip(self.leaf_arrays, self.weights_per_root):
734 leaf_const = jnp.asarray(terms) * self._leaf_factors(
735 S, C, self.var_positions
736 )
737 per_root_coeffs.append(jnp.asarray(W) @ leaf_const)
739 return self._combine_roots(per_root_coeffs, self.freqs_per_root, force_mean)
741 def get_exact_support(self, method: str = "tree") -> List[np.ndarray]:
742 r"""Symbolically derive the exact frequency support (no sampling).
744 A frequency :math:`\omega` belongs to the exact spectrum iff its
745 coefficient :math:`c_\omega(\theta) = \sum_l W_{\omega l}\,
746 \text{term}_l\, v_l(\theta)` is not identically zero in the
747 variational parameters :math:`\theta`.
749 Two methods are available:
751 - ``"tree"`` (default, fully exact): enumerates the explicit tree
752 leaves. Because the branch index strictly decreases along every tree
753 path, each parameter contributes **at most one** sine *or* cosine
754 factor per leaf (:math:`S_{li}, C_{li} \in \{0, 1\}`). Every
755 variational leaf factor :math:`v_l` is therefore a *square-free*
756 monomial over :math:`\{1, \cos\theta_i, i\sin\theta_i\}`, and
757 monomials with distinct signatures are linearly independent functions
758 (no :math:`\cos^2 + \sin^2` identities can arise without squares).
759 Hence
761 .. math::
762 c_\omega \equiv 0 \iff \sum_{l \in g} W_{\omega l}\,\text{term}_l
763 = 0 \quad \text{for every signature group } g.
765 Since all involved quantities are dyadic rationals times
766 :math:`\{\pm 1, \pm i\}`, the group sums are exact in float64 and the
767 zero-test is exact. The number of leaves can however grow
768 exponentially with circuit depth.
770 - ``"dp"`` (scalable): merges tree nodes with identical
771 ``(rotation index, observable)`` — at most ``n_params * 4^n_qubits``
772 states — and tracks, per state, the achievable per-feature sine/cosine
773 count vectors ``(s_f, c_f)`` as a mixed-radix bitmask. Each feature's
774 support is the union of the (exact) expansion supports of
775 :math:`\cos^{c_f} x_f\, (i \sin x_f)^{s_f}`, and the model support is
776 their Cartesian product across features. This is exact per tree path
777 (including interior zero coefficients of the expansions), but unlike
778 ``"tree"`` it cannot detect coefficients that cancel identically
779 *across* paths with identical variational signatures (e.g. directly
780 repeated encodings). It therefore yields a tight superset in such
781 corner cases. Supports any number of input features, but requires
782 unit-magnitude input scaling: per-gate :math:`|\omega| \neq 1`
783 scalings (e.g. Golomb encodings) are rejected — use ``"tree"``.
785 Args:
786 method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).
788 Returns:
789 List[np.ndarray]: For each observable (root), the frequency vectors
790 with not-identically-zero coefficient — shape ``(n_freq,)`` for a
791 single input feature, ``(n_freq, n_features)`` otherwise.
792 """
793 if method == "dp":
794 return self._support_dp()
795 if method != "tree":
796 raise ValueError(f"Unknown method '{method}'. Use 'tree' or 'dp'.")
798 self._ensure_structure()
799 supports = []
800 for (S, C, terms), W, freqs in zip(
801 self.leaf_arrays, self.weights_per_root, self.freqs_per_root
802 ):
803 freqs = np.asarray(freqs)
804 n_leaves = S.shape[0]
805 if n_leaves == 0:
806 supports.append(freqs[:0])
807 continue
809 # Group leaves by their variational sine/cosine signature.
810 signature = np.hstack([S[:, self.var_positions], C[:, self.var_positions]])
811 _, groups = np.unique(signature, axis=0, return_inverse=True)
812 n_groups = int(groups.max()) + 1
814 # Per-group sums of W[omega, l] * term_l, accumulated exactly.
815 contrib = (W * terms[None, :]).T # (n_leaves, n_freq)
816 group_sums = np.zeros((n_groups, W.shape[0]), dtype=np.complex128)
817 np.add.at(group_sums, groups, contrib)
819 mask = (np.abs(group_sums) > 1e-12).any(axis=0) # (n_freq,)
820 supports.append(freqs[mask])
821 return supports
823 def _support_dp(self) -> List[np.ndarray]:
824 """Merged-state dynamic program for the frequency support.
826 Instead of enumerating all (worst-case exponentially many) tree paths,
827 nodes are merged on ``(rotation index, bare observable)``. Each state
828 stores the set of achievable per-axis count vectors
829 ``(c_0, s_0, ..., c_{d-1}, s_{d-1})`` as a mixed-radix bitmask (one
830 digit per feature sine/cosine count), so transitions are O(1) big-int
831 operations. See :meth:`get_exact_support` for semantics and
832 limitations.
833 """
834 # Count aggregation is valid as long as every input rotation has
835 # unit-magnitude scaling. A Clifford-commutation sign flip (scaling -1)
836 # leaves the frequency support unchanged -- cos is even and sin odd, so
837 # the sign only flips the coefficient, which the support ignores -- but a
838 # genuine per-gate scaling (|omega| != 1, e.g. Golomb / heterogeneous
839 # frequencies) cannot be represented by sin/cos counts.
840 if self.all_input_indices and np.any(
841 np.abs(self.input_scaling[self.all_input_indices]) != 1
842 ):
843 raise NotImplementedError(
844 "The 'dp' support method does not support non-unit input "
845 "frequency scaling (it aggregates sin/cos counts and cannot "
846 "represent per-gate scalings); use method='tree'."
847 )
849 n = self.n_qubits
850 d = len(self.features)
851 # Pack the achievable per-axis sine/cosine counts into one big-int as a
852 # mixed-radix bitmask over (c_0, s_0, ..., c_{d-1}, s_{d-1}). Axis a's
853 # counts range 0..n_a (n_a input rotations encode feature a), so each
854 # digit has radix n_a + 1; a left-shift by a digit's place value
855 # increments that count, OR is set-union. (For d == 1 this reduces to
856 # shift_c = 1, shift_s = n_inp + 1, i.e. the flat (s, c) layout.)
857 ranges = [len(self.input_indices[self.features[a]]) + 1 for a in range(d)]
858 shift_c = [0] * d
859 shift_s = [0] * d
860 place = 1
861 for a in range(d):
862 shift_c[a] = place
863 place *= ranges[a]
864 shift_s[a] = place
865 place *= ranges[a]
866 # Feature axis of each input rotation (-1 if variational), in the same
867 # enumerate(self.features) order the tree method uses for its axes.
868 axis_of_col = np.full(self.n_params, -1, dtype=np.int64)
869 for a in range(d):
870 for k in self.input_indices[self.features[a]]:
871 axis_of_col[k] = a
873 def encode(word: PauliWord) -> Tuple[int, int]:
874 x = z = 0
875 for q in range(n):
876 x |= int(word.x[q]) << q
877 z |= int(word.z[q]) << q
878 return x, z
880 paulis = [encode(w) for w in self.pauli_words]
881 cum_xy = []
882 running = 0
883 for xp, _ in paulis:
884 running |= xp
885 cum_xy.append(running)
887 def parity(v: int) -> int:
888 return bin(v).count("1") & 1
890 def dp(idx: int, xo: int, zo: int, memo: dict) -> int:
891 # Light-cone early stopping (cf. _early_stopping_possible).
892 if idx >= 0 and (xo & ~cum_xy[idx]):
893 return 0
894 # Skip trailing rotations that commute with the observable.
895 while idx >= 0:
896 xp, zp = paulis[idx]
897 if parity(xo & zp) ^ parity(zo & xp):
898 break
899 idx -= 1
900 else: # leaf: counts (s=0, c=0) iff the observable is diagonal
901 return 1 if xo == 0 else 0
902 key = (idx, xo, zo)
903 hit = memo.get(key)
904 if hit is not None:
905 return hit
906 xp, zp = paulis[idx]
907 cos_child = dp(idx - 1, xo, zo, memo)
908 sin_child = dp(idx - 1, xo ^ xp, zo ^ zp, memo)
909 a = int(axis_of_col[idx])
910 if a >= 0:
911 # Active input gate: cosine increments c_a, sine increments s_a.
912 val = (cos_child << shift_c[a]) | (sin_child << shift_s[a])
913 else:
914 val = cos_child | sin_child
915 memo[key] = val
916 return val
918 # Recursion depth is bounded by the number of rotations.
919 old_limit = sys.getrecursionlimit()
920 sys.setrecursionlimit(max(old_limit, self.n_params + 1000))
921 try:
922 supports = []
923 for obs in self.observable_words:
924 memo: dict = {}
925 xo, zo = encode(obs)
926 mask = dp(self.n_params - 1, xo, zo, memo)
927 supports.append(self._dp_mask_to_support(mask, d, ranges))
928 finally:
929 sys.setrecursionlimit(old_limit)
930 return supports
932 def _dp_mask_to_support(self, mask: int, d: int, ranges: List[int]) -> np.ndarray:
933 """Decode a count bitmask (see :meth:`_support_dp`) into a frequency
934 support. Each set bit is a per-axis count vector ``(c_a, s_a)``; the
935 per-axis expansion supports are combined across axes (Cartesian
936 product) and unioned over all bits.
938 Returns an ``(n_freq, d)`` array, collapsed to 1-D for ``d <= 1`` to
939 match :meth:`get_exact_support` with ``method="tree"`` (``d == 0`` keeps
940 only the DC term).
941 """
942 tupleset: set = set()
943 while mask:
944 bit = mask & -mask
945 i = bit.bit_length() - 1
946 rem = i
947 axis_freqs = []
948 for a in range(d):
949 c_a = rem % ranges[a]
950 rem //= ranges[a]
951 s_a = rem % ranges[a]
952 rem //= ranges[a]
953 axis_freqs.append(sorted(self._expansion_support(s_a, c_a)))
954 tupleset.update(itertools.product(*axis_freqs))
955 mask ^= bit
957 if d >= 2:
958 if not tupleset:
959 return np.empty((0, d), dtype=np.int64)
960 return np.array(sorted(tupleset), dtype=np.int64)
961 # d in {0, 1}: collapse to a 1-D frequency array.
962 flat = sorted(t[0] for t in tupleset) if d == 1 else ([0] if tupleset else [])
963 return np.array(flat, dtype=np.int64)
965 @staticmethod
966 @lru_cache(maxsize=None)
967 def _expansion_support(s: int, c: int) -> frozenset:
968 r"""Frequencies with non-zero coefficient in :math:`\cos^c x (i\sin x)^s`.
970 Computed exactly with integer arithmetic via the polynomial
971 :math:`(t - 1)^s (t + 1)^c` (with :math:`t = e^{2ix}` up to a shift);
972 interior coefficients can vanish, e.g. :math:`\cos x \sin x` only
973 contains :math:`\pm 2`.
974 """
975 coeffs = [1]
976 for _ in range(s): # multiply by (t - 1)
977 new = [0] * (len(coeffs) + 1)
978 for i, a in enumerate(coeffs):
979 new[i + 1] += a
980 new[i] -= a
981 coeffs = new
982 for _ in range(c): # multiply by (t + 1)
983 new = [0] * (len(coeffs) + 1)
984 for i, a in enumerate(coeffs):
985 new[i + 1] += a
986 new[i] += a
987 coeffs = new
988 m = s + c
989 return frozenset(2 * k - m for k, a in enumerate(coeffs) if a != 0)
991 def _combine_roots(
992 self,
993 per_root_coeffs: List[jnp.ndarray],
994 per_root_freqs: List[np.ndarray],
995 force_mean: bool,
996 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
997 """Assemble the per-root spectra, optionally averaging over roots."""
998 if not force_mean:
999 coefficients = [jnp.asarray(c) for c in per_root_coeffs]
1000 frequencies = [jnp.asarray(f) for f in per_root_freqs]
1001 return coefficients, frequencies
1003 # Average over roots on the union of all frequency vectors. Keys are
1004 # snapped to a tolerance so rational frequencies dedup robustly.
1005 accum: Dict[tuple, complex] = defaultdict(complex)
1006 for coeffs, freqs in zip(per_root_coeffs, per_root_freqs):
1007 freqs_np = np.asarray(freqs)
1008 for k in range(freqs_np.shape[0]):
1009 key = (
1010 (round(float(freqs_np[k]), 9),)
1011 if freqs_np.ndim == 1
1012 else tuple(round(float(v), 9) for v in freqs_np[k])
1013 )
1014 accum[key] += complex(coeffs[k])
1015 n_roots = max(len(per_root_coeffs), 1)
1016 keys = sorted(accum.keys())
1017 mean_coeffs = jnp.array([accum[k] / n_roots for k in keys])
1018 freq_arr = np.array(keys, dtype=float)
1019 rounded = np.rint(freq_arr)
1020 if np.all(np.abs(freq_arr - rounded) < 1e-6):
1021 freq_arr = rounded.astype(np.int64)
1022 if freq_arr.shape[1] == 1:
1023 freq_arr = freq_arr[:, 0]
1024 return [mean_coeffs], [jnp.asarray(freq_arr)]
1027class FCC:
1028 @classmethod
1029 def get_fcc(
1030 cls,
1031 model: Model,
1032 n_samples: int,
1033 random_key: Optional[random.PRNGKey] = None,
1034 method: Optional[str] = "pearson",
1035 scale: Optional[bool] = False,
1036 weight: Optional[bool] = False,
1037 trim_redundant: Optional[bool] = True,
1038 **kwargs,
1039 ) -> float:
1040 """
1041 Shortcut method to get just the FCC.
1042 This includes
1043 1. What is done in `get_fourier_fingerprint`:
1044 1. Calculating the coefficients (using `n_samples`)
1045 2. Correlating the result from 1) using `method`
1046 3. Weighting the correlation matrix (if `weight` is True)
1047 4. Remove redundancies
1048 2. What is done in `calculate_fcc`:
1049 1. Absolute of the fingerprint
1050 2. Average
1052 Args:
1053 model (Model): The QFM model
1054 n_samples (int): Number of samples to calculate average of coefficients
1055 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1056 initialization. If None, uses the model's internal random key.
1057 method (Optional[str], optional): Correlation method. Supported values are
1058 "pearson", "complex_pearson", "spearman", and "covariance".
1059 Defaults to "pearson".
1060 scale (Optional[bool], optional): Whether to scale the number of samples.
1061 Defaults to False.
1062 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1063 Defaults to False.
1064 trim_redundant (Optional[bool], optional): Whether to remove redundant
1065 correlations. Defaults to False.
1066 **kwargs (Any): Additional keyword arguments for the model function.
1068 Returns:
1069 float: The FCC
1070 """
1072 # Memory-efficient fast path
1073 if trim_redundant and not weight:
1074 _, coeffs, freqs = cls._calculate_coefficients(
1075 model, n_samples, random_key, scale, **kwargs
1076 )
1077 pos_idx = cls._calculate_mask(freqs)
1078 coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
1079 coeffs_sub = coeffs_flat[pos_idx]
1081 fp = cls._correlate(coeffs_sub.transpose(), method=method)
1082 abs_fp = jnp.abs(fp)
1083 diag = jnp.abs(jnp.diagonal(fp))
1085 total_sum = jnp.nansum(abs_fp)
1086 total_count = jnp.sum(jnp.isfinite(abs_fp))
1087 diag_sum = jnp.nansum(diag)
1088 diag_count = jnp.sum(jnp.isfinite(diag))
1090 lower_sum = (total_sum - diag_sum) / 2.0
1091 lower_count = (total_count - diag_count) / 2.0
1092 return lower_sum / lower_count
1094 fourier_fingerprint, _, _ = cls.get_fourier_fingerprint(
1095 model,
1096 n_samples,
1097 random_key,
1098 method,
1099 scale,
1100 weight,
1101 trim_redundant=trim_redundant,
1102 **kwargs,
1103 )
1105 return cls.calculate_fcc(fourier_fingerprint)
1107 @classmethod
1108 def get_fourier_fingerprint(
1109 cls,
1110 model: Model,
1111 n_samples: int,
1112 random_key: Optional[random.PRNGKey] = None,
1113 method: Optional[str] = "pearson",
1114 scale: Optional[bool] = False,
1115 weight: Optional[bool] = False,
1116 trim_redundant: Optional[bool] = True,
1117 nan_to_one: Optional[bool] = False,
1118 **kwargs: Any,
1119 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
1120 """
1121 Shortcut method to get just the fourier fingerprint.
1122 This includes
1123 1. Calculating the coefficients (using `n_samples`)
1124 2. Correlating the result from 1) using `method`
1125 3. Weighting the correlation matrix (if `weight` is True)
1126 4. Remove redundancies (if `trim_redundant` is True)
1128 Args:
1129 model (Model): The QFM model
1130 n_samples (int): Number of samples to calculate average of coefficients
1131 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1132 initialization. If None, uses the model's internal random key.
1133 method (Optional[str], optional): Correlation method. Supported values are
1134 "pearson", "complex_pearson", "spearman", and "covariance".
1135 Defaults to "pearson".
1136 scale (Optional[bool], optional): Whether to scale the number of samples.
1137 Defaults to False.
1138 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1139 Defaults to False.
1140 trim_redundant (Optional[bool], optional): Whether to remove redundant
1141 correlations. Defaults to True.
1142 nan_to_one (Optional[bool], optional): Whether to set nan to 1.
1143 Defaults to False.
1144 **kwargs: Additional keyword arguments for the model function.
1146 Returns:
1147 Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: The fourier
1148 fingerprint, the corresponding frequency indices and the
1149 corresponding coefficients. If `trim_redundant` is True the
1150 frequencies are returned as a `(row_freqs, col_freqs)` tuple that
1151 labels the two (redundancy-trimmed) matrix axes and the
1152 coefficients as a matching `(row_coeffs, col_coeffs)` tuple whose
1153 rows align with those frequencies; otherwise the full frequency
1154 vector and full coefficient array are returned.
1155 """
1156 _, coeffs, freqs = cls._calculate_coefficients(
1157 model, n_samples, random_key, scale, **kwargs
1158 )
1160 # Memory-efficient fast path
1161 if trim_redundant and not weight:
1162 pos_idx = cls._calculate_mask(freqs)
1163 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1165 # Flatten all frequency axes; the last axis is the sample
1166 # axis. `_calculate_mask` returns flat indices in C order,
1167 # matching this reshape.
1168 coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
1169 coeffs_sub = coeffs_flat[pos_idx]
1171 fourier_fingerprint = cls._correlate(coeffs_sub.transpose(), method=method)
1173 if nan_to_one:
1174 fourier_fingerprint = jnp.where(
1175 jnp.isnan(fourier_fingerprint), 1.0, fourier_fingerprint
1176 )
1178 M = fourier_fingerprint.shape[0]
1179 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1180 fourier_fingerprint = jnp.where(
1181 lower_tri_mask, fourier_fingerprint, jnp.nan
1182 )
1184 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1185 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1186 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1188 return (
1189 fourier_fingerprint,
1190 (pos_freqs[row_mask], pos_freqs[col_mask]),
1191 (coeffs_sub[row_mask], coeffs_sub[col_mask]),
1192 )
1194 fourier_fingerprint = cls._correlate(coeffs.transpose(), method=method)
1196 if nan_to_one:
1197 # set nan to 1
1198 fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0
1200 # perform weighting if requested
1201 fourier_fingerprint = (
1202 cls._weighting_mean(fourier_fingerprint, coeffs)
1203 if weight
1204 else fourier_fingerprint
1205 )
1207 if trim_redundant:
1208 pos_idx = cls._calculate_mask(freqs)
1209 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1210 coeffs_sub = coeffs.reshape(-1, coeffs.shape[-1])[pos_idx]
1212 # restrict to the positive-frequency sub-block (M x M with
1213 # M = number of non-negative flat-frequencies) instead of
1214 # building a full N x N mask. This avoids the O(N^2) float
1215 fourier_fingerprint = fourier_fingerprint[pos_idx][:, pos_idx]
1217 # keep only the strict lower triangle; the rest -> nan
1218 M = fourier_fingerprint.shape[0]
1219 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1220 fourier_fingerprint = jnp.where(
1221 lower_tri_mask, fourier_fingerprint, jnp.nan
1222 )
1224 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1225 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1227 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1229 return (
1230 fourier_fingerprint,
1231 (pos_freqs[row_mask], pos_freqs[col_mask]),
1232 (coeffs_sub[row_mask], coeffs_sub[col_mask]),
1233 )
1235 return fourier_fingerprint, freqs, coeffs
1237 @classmethod
1238 def calculate_fcc(
1239 cls,
1240 fourier_fingerprint: jnp.ndarray,
1241 ) -> float:
1242 """
1243 Method to calculate the FCC based on an existing correlation matrix.
1244 Calculate absolute and then the average over this matrix.
1245 The Fingerprint can be obtained via `get_fourier_fingerprint`
1247 Args:
1248 fourier_fingerprint (jnp.ndarray): Correlation matrix of coefficients
1249 Returns:
1250 float: The FCC
1251 """
1252 # apply the mask on the fingerprint
1253 return jnp.nanmean(jnp.abs(fourier_fingerprint))
1255 @classmethod
1256 def _calculate_mask(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1257 """
1258 Determine the flat indices of the Fourier correlation matrix
1259 that lie on a non-negative-frequency row/column. Together with
1260 the strict-lower-triangle condition (handled by the caller),
1261 these indices select the entries of the correlation matrix
1262 that survive the redundancy filter applied in
1263 `get_fourier_fingerprint`:
1265 - rows/columns whose flat frequency component is negative are
1266 discarded (they are the complex-conjugate redundancies of
1267 their positive counterparts);
1268 - of the remaining positive-frequency sub-block, only the
1269 strict lower triangle is kept (the upper triangle, including
1270 the diagonal, contains either duplicates from symmetry or
1271 self-correlations).
1273 Args:
1274 freqs (jnp.ndarray): Array of frequencies. Either a 1-D
1275 vector (single input feature) or a 2-D array of shape
1276 ``(n_input_feat, K)`` whose rows are the per-axis
1277 frequency vectors.
1279 Returns:
1280 jnp.ndarray: 1-D int array of flat indices selecting the
1281 non-negative-frequency rows/cols of the fingerprint.
1282 """
1283 freqs_arr = jnp.asarray(freqs)
1285 if freqs_arr.ndim == 1:
1286 pos_flat = freqs_arr >= 0
1287 else:
1288 # N-D case: build the per-axis non-negativity masks and
1289 # combine them via broadcasting (no float `jnp.outer`!),
1290 # then flatten to match the row-major flattening used by
1291 # the upstream coefficient/correlation pipeline.
1292 axes_pos = [freqs_arr[i] >= 0 for i in range(freqs_arr.shape[0])]
1293 expanded = []
1294 n_axes = len(axes_pos)
1295 for i, p in enumerate(axes_pos):
1296 shape = [1] * n_axes
1297 shape[i] = p.shape[0]
1298 expanded.append(p.reshape(shape))
1299 nd_pos = reduce(jnp.logical_and, expanded)
1300 pos_flat = nd_pos.flatten()
1302 return jnp.where(pos_flat)[0]
1304 @classmethod
1305 def _flat_frequencies(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1306 """
1307 Build the per-coefficient flat frequency labels in the same
1308 C-order used to flatten the coefficient/correlation pipeline, so
1309 they can be indexed by the flat indices from `_calculate_mask`.
1311 Args:
1312 freqs (jnp.ndarray): Either a 1-D vector (single input feature)
1313 or a ``(n_input_feat, K)`` stack / list of per-axis frequency
1314 vectors (multi-dim input).
1316 Returns:
1317 jnp.ndarray: 1-D frequency vector (single input feature) or a
1318 ``(N, n_input_feat)`` array of per-coefficient frequency
1319 tuples (multi-dim input).
1320 """
1321 fa = jnp.asarray(freqs)
1322 if fa.ndim == 1:
1323 return fa
1324 # Multi-dim: per-axis vectors -> flat grid of frequency tuples in the
1325 # same C-order used by `_calculate_mask` and the coefficient reshape.
1326 grids = jnp.meshgrid(*[fa[i] for i in range(fa.shape[0])], indexing="ij")
1327 return jnp.stack(grids, axis=-1).reshape(-1, fa.shape[0])
1329 @classmethod
1330 def _calculate_coefficients(
1331 cls,
1332 model: Model,
1333 n_samples: int,
1334 random_key: Optional[random.PRNGKey] = None,
1335 scale: bool = False,
1336 **kwargs: Any,
1337 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
1338 """
1339 Calculates the Fourier coefficients of a given model
1340 using `n_samples`.
1341 Optionally, `noise_params` can be passed to perform noisy simulation.
1343 Args:
1344 model (Model): The QFM model
1345 n_samples (int): Number of samples to calculate average of coefficients
1346 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1347 initialization. If None, uses the model's internal random key.
1348 scale (bool, optional): Whether to scale the number of samples.
1349 Defaults to False.
1350 **kwargs: Additional keyword arguments for the model function.
1352 Returns:
1353 Tuple[jnp.ndarray, jnp.ndarray]: Parameters and Coefficients of size NxK
1354 """
1355 if n_samples > 0:
1356 if scale:
1357 total_samples = int(
1358 jnp.power(2, model.n_qubits) * n_samples * model.n_input_feat
1359 )
1360 log.info(f"Using {total_samples} samples.")
1361 else:
1362 total_samples = n_samples
1363 model.initialize_params(random_key, repeat=total_samples)
1364 else:
1365 total_samples = 1
1367 coeffs, freqs = Coefficients.get_spectrum(
1368 model, shift=True, trim=True, **kwargs
1369 )
1371 return model.params, coeffs, freqs
1373 @classmethod
1374 def _correlate(cls, mat: jnp.ndarray, method: str = "pearson") -> jnp.ndarray:
1375 """
1376 Correlates two arrays using `method`.
1377 Currently, `pearson`, `complex_pearson`, `spearman`, and `covariance`
1378 are supported.
1380 Args:
1381 mat (jnp.ndarray): Array of shape (N, K)
1382 method (str, optional): Correlation method. Defaults to "pearson".
1384 Raises:
1385 ValueError: If the method is not supported.
1387 Returns:
1388 jnp.ndarray: Correlation matrix of `a` and `b`.
1389 """
1390 assert len(mat.shape) >= 2, "Input matrix must have at least 2 dimensions"
1392 # Note that for the general n-D case, we have to flatten along
1393 # the first axis (last one is batch).
1394 # Note that the order here is important so we can easily filter out
1395 # negative coefficients later.
1396 # Consider the following example: [[1,2,3],[4,5,6],[7,8,9]]
1397 # we want to get [1, 4, 7, 2, 5, 8, 3, 6, 9]
1398 # such that after correlation, all positive indexed coefficients
1399 # will be in the bottom right quadrant
1400 if method == "pearson":
1401 result = cls._pearson(mat.reshape(mat.shape[0], -1))
1402 # result = cls._pearson(mat.reshape(mat.shape[-1], -1, order="F"))
1403 elif method == "complex_pearson":
1404 result = cls._complex_pearson(mat.reshape(mat.shape[0], -1))
1405 elif method == "spearman":
1406 result = cls._spearman(mat.reshape(mat.shape[0], -1))
1407 # result = cls._spearman(mat.reshape(mat.shape[-1], -1, order="F"))
1408 elif method == "covariance":
1409 result = cls._covariance(mat.reshape(mat.shape[0], -1))
1410 else:
1411 raise ValueError(
1412 f"Unknown correlation method: {method}. Must be 'pearson', \
1413 'complex_pearson', 'spearman' or 'covariance'."
1414 )
1416 return result
1418 @classmethod
1419 def _covariance(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1420 """
1421 Compute the Hermitian sample covariance between columns of `mat`,
1422 permitting missing values (NaN or ±Inf).
1424 For each pair (i, j) the covariance is computed over the rows that are
1425 finite in both columns, as
1426 sum(conj(x_i - mean_i) * (x_j - mean_j)) / (nobs - 1),
1427 so it computes `X.conj().T @ X`.
1428 Real input collapses to the ordinary real sample covariance; complex
1429 input yields a complex matrix whose magnitude and angle carry the
1430 covariance strength and relative phase.
1433 Args:
1434 mat : array_like, shape (N, K)
1435 Input data.
1436 minp : int, optional
1437 Minimum number of paired observations required to form a
1438 covariance. If the number of valid pairs for (i, j) is < minp,
1439 the result is NaN.
1441 Returns:
1442 cov : ndarray, shape (K, K)
1443 Sample covariance matrix.
1444 """
1445 mat = jnp.asarray(mat)
1446 real_dtype = jnp.asarray(mat.real).dtype
1448 mask = jnp.isfinite(mat)
1449 fmask = mask.astype(real_dtype)
1450 safe = jnp.where(mask, mat, 0.0)
1452 nobs = fmask.T @ fmask
1453 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1455 sum_x = safe.T @ fmask
1456 sum_y = fmask.T @ safe
1458 masked = safe * fmask
1459 sum_conj_xy = jnp.conj(masked).T @ masked
1461 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1463 denom = jnp.where(nobs > 1, nobs - 1, jnp.nan)
1464 result = sxy / denom
1466 result = jnp.where(nobs < minp, jnp.nan, result)
1468 return result
1470 @classmethod
1471 def _complex_pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1472 """
1473 Compute the complex Pearson correlation between columns of `mat`,
1474 permitting missing values (NaN or ±Inf).
1476 This uses the Hermitian normalized covariance
1477 sum(conj(x_i - mean_i) * (x_j - mean_j)) /
1478 sqrt(sum(abs(x_i - mean_i)**2) * sum(abs(x_j - mean_j)**2)).
1479 Consequently, if column j is exp(1j * phi) times column i, then
1480 abs(corr[i, j]) is 1 and angle(corr[i, j]) is phi.
1482 Args:
1483 mat : array_like, shape (N, K)
1484 Input data.
1485 minp : int, optional
1486 Minimum number of paired observations required to form a correlation.
1487 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1489 Returns:
1490 corr : ndarray, shape (K, K)
1491 Complex Pearson correlation matrix.
1492 """
1493 mat = jnp.asarray(mat)
1494 real_dtype = jnp.asarray(mat.real).dtype
1496 mask = jnp.isfinite(mat)
1497 fmask = mask.astype(real_dtype)
1498 safe = jnp.where(mask, mat, 0.0)
1500 nobs = fmask.T @ fmask
1501 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1503 sum_x = safe.T @ fmask
1504 sum_y = fmask.T @ safe
1506 masked = safe * fmask
1507 sum_conj_xy = jnp.conj(masked).T @ masked
1509 safe_abs_sq = jnp.abs(safe) ** 2
1510 sum_abs_x2 = safe_abs_sq.T @ fmask
1511 sum_abs_y2 = fmask.T @ safe_abs_sq
1513 ssx = sum_abs_x2 - jnp.abs(sum_x) ** 2 / nobs_safe
1514 ssy = sum_abs_y2 - jnp.abs(sum_y) ** 2 / nobs_safe
1515 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1517 denom = jnp.sqrt(ssx * ssy)
1518 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1519 magnitude = jnp.abs(result)
1520 result = jnp.where(magnitude > 1.0, result / magnitude, result)
1522 result = jnp.where(nobs < minp, jnp.nan, result)
1524 return result
1526 @classmethod
1527 def _pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1528 """
1529 Compute Pearson correlation between columns of `mat`,
1530 permitting missing values (NaN or ±Inf).
1532 The Pearson correlation is the normalized covariance,
1533 corr[i, j] = cov[i, j] / sqrt(cov[i, i] * cov[j, j]),
1534 so it is obtained by normalizing `_covariance` by the per-column
1535 standard deviations.
1537 If the input is complex, real and imaginary parts are stacked along
1538 the sample axis so that both components contribute to the correlation
1539 without discarding information.
1541 Args:
1542 mat : array_like, shape (N, K)
1543 Input data.
1544 minp : int, optional
1545 Minimum number of paired observations required to form a correlation.
1546 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1548 Returns:
1549 corr : ndarray, shape (K, K)
1550 Pearson correlation matrix.
1551 """
1552 # Preserve complex information by splitting into real / imag samples.
1553 # After stacking the data is real, so the Hermitian `_covariance`
1554 # reduces to the ordinary real sample covariance.
1555 if jnp.iscomplexobj(mat):
1556 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1558 cov = cls._covariance(mat, minp=minp)
1560 # corr[i, j] = cov[i, j] / (std_i * std_j) with std_i = sqrt(cov[i, i])
1561 std = jnp.sqrt(jnp.diagonal(cov))
1562 denom = std[:, None] * std[None, :]
1563 result = jnp.where(denom > 0, cov / denom, jnp.nan)
1565 # clip numerical drift to [-1, 1]
1566 result = jnp.clip(jnp.real(result), -1.0, 1.0)
1568 return result
1570 @classmethod
1571 def _spearman(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1572 """
1573 Based on Pandas correlation method as implemented here:
1574 https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx
1576 Compute Spearman correlation between columns of `mat`,
1577 permitting missing values (NaN or ±Inf).
1579 If the input is complex, real and imaginary parts are stacked along
1580 the sample axis so that both components contribute to the correlation
1581 without discarding information.
1583 Args:
1584 mat : array_like, shape (N, K)
1585 Input data.
1586 minp : int, optional
1587 Minimum number of paired observations required to form a correlation.
1588 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1590 Returns:
1591 corr : ndarray, shape (K, K)
1592 Spearman correlation matrix.
1593 """
1594 # Preserve complex information by splitting into real / imag samples
1595 if jnp.iscomplexobj(mat):
1596 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1598 mat = jnp.asarray(mat)
1599 N, K = mat.shape
1601 # trivial all-NaN answer if too few rows
1602 if N < minp:
1603 return jnp.full((K, K), jnp.nan)
1605 # mask of finite entries
1606 mask = jnp.isfinite(mat) # shape (N, K), dtype=bool
1608 # precompute ranks column-wise ignoring NaNs
1609 ranks = np.full((N, K), np.nan)
1610 for j in range(K):
1611 valid = mask[:, j]
1612 if valid.any():
1613 ranks[valid, j] = rankdata(mat[valid, j], method="average")
1615 ranks = jnp.asarray(ranks)
1617 # Vectorised Pearson on the ranks
1618 # Replace NaN ranks with 0; use mask to track validity.
1619 rank_mask = jnp.isfinite(ranks)
1620 safe_ranks = jnp.where(rank_mask, ranks, 0.0)
1622 # Pairwise valid-observation counts (K, K)
1623 fmask = rank_mask.astype(ranks.dtype)
1624 nobs = fmask.T @ fmask
1626 # Pairwise sums over mutually-valid rows
1627 sum_x = safe_ranks.T @ fmask # (K, K)
1628 sum_y = fmask.T @ safe_ranks # (K, K)
1630 # Pairwise products
1631 masked_ranks = safe_ranks * fmask # same as safe_ranks
1632 sum_xy = masked_ranks.T @ masked_ranks # (K, K)
1634 safe_sq = safe_ranks**2
1635 sum_x2 = safe_sq.T @ fmask # (K, K)
1636 sum_y2 = fmask.T @ safe_sq # (K, K)
1638 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1639 ssx = sum_x2 - sum_x**2 / nobs_safe
1640 ssy = sum_y2 - sum_y**2 / nobs_safe
1641 sxy = sum_xy - (sum_x * sum_y) / nobs_safe
1643 denom = jnp.sqrt(ssx * ssy)
1644 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1645 result = jnp.clip(result, -1.0, 1.0)
1647 # Enforce minp
1648 result = jnp.where(nobs < minp, jnp.nan, result)
1650 return result
1652 @classmethod
1653 def _weighting_linear(cls, fourier_fingerprint: jnp.ndarray) -> jnp.ndarray:
1654 """
1655 Performs weighting on the given correlation matrix.
1656 Here, low-frequent coefficients are weighted more heavily.
1658 Args:
1659 fourier_fingerprint (jnp.ndarray): Correlation matrix
1660 """
1661 assert (
1662 fourier_fingerprint.shape[0] % 2 != 0
1663 and fourier_fingerprint.shape[1] % 2 != 0
1664 ), (
1665 "Correlation matrix must have odd dimensions. \
1666 Hint: use `trim` argument when calling `get_spectrum`."
1667 )
1668 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1669 "Correlation matrix must be square."
1670 )
1672 # The weight matrix produced by the previous quadrant-mirror
1673 # construction has a closed form: it is a "tent" sum along the
1674 # two axes. Concretely, with N = fourier_fingerprint.shape[0]
1675 # (odd) and center = N // 2,
1676 # W[i, j] = u[i] + u[j]
1677 # where u[k] = (center - |k - center|) / (2 * center)
1678 # is a triangular weighting peaking at the centre (the zero
1679 # frequency) and decaying linearly to 0 at the spectrum edges.
1680 N = fourier_fingerprint.shape[0]
1681 center = N // 2
1682 k = jnp.arange(N)
1683 u = (center - jnp.abs(k - center)) / (2 * center)
1685 return fourier_fingerprint * (u[:, None] + u[None, :])
1687 @classmethod
1688 def _weighting_mean(
1689 cls, fourier_fingerprint: jnp.ndarray, coeffs: jnp.ndarray
1690 ) -> jnp.ndarray:
1691 """
1692 Performs weighting on the given correlation matrix.
1693 Here, we use the product of the mean of the coefficients as weights.
1694 This suppresses correlations where the mean of the coefficients is near zero.
1696 Args:
1697 fourier_fingerprint (jnp.ndarray): Correlation matrix
1698 coeffs (jnp.ndarray): Fourier coefficients
1699 """
1700 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1701 "Correlation matrix must be square."
1702 )
1703 assert len(coeffs.shape) >= 2, (
1704 "Coefficient matrix must contain coefficient axes and a sample axis."
1705 )
1707 coefficient_means = jnp.abs(jnp.mean(coeffs, axis=-1))
1708 coefficient_means = coefficient_means.T.reshape(-1)
1710 assert fourier_fingerprint.shape[0] == coefficient_means.shape[0], (
1711 "Correlation matrix size must match the number of Fourier coefficients."
1712 )
1714 # Apply the rank-1 weight w[i] * w[j] via broadcasting instead
1715 # of materialising an explicit `jnp.outer` N x N intermediate.
1716 return (
1717 fourier_fingerprint
1718 * coefficient_means[:, None]
1719 * coefficient_means[None, :]
1720 )
1723class Datasets:
1724 @classmethod
1725 def generate_fourier_series(
1726 cls,
1727 random_key: random.PRNGKey,
1728 model: Model,
1729 coefficients_min: float = 0.0,
1730 coefficients_max: float = 1.0,
1731 zero_centered: bool = False,
1732 ) -> jnp.ndarray:
1733 """
1734 Generates the Fourier series representation of a function.
1735 It uses the `model.frequencies` property to retrieve the frequency
1736 information. This ensures that the resulting Fourier series is
1737 compatible with the model.
1739 This function is capable of generating $D$-dimensional Fourier series
1740 (again defined by `model.n_input_feat`).
1741 The highest frequency $N$ is retrieved per dimension.
1743 Samples of the Fourier coefficients are drawn from a uniform circle.
1745 Args:
1746 random_key (random.PRNGKey): Random number key for JAX.
1747 model (Model): The quantum circuit model.
1748 coefficients_min (float, optional): Minimum value for the coefficients.
1749 Defaults to 0.0.
1750 coefficients_max (float, optional): Maximum value for the coefficients.
1751 Defaults to 1.0.
1752 zero_centered (bool, optional): Whether to zero-center the coefficients.
1753 Defaults to False.
1755 Returns:
1756 jnp.ndarray: Input domain samples with shape ((N,)*D, D)
1757 jnp.ndarray: Fourier series values with shape ((N,)*D)
1758 jnp.ndarray: Fourier coefficients with shape ((N,)*D)
1760 """
1761 # TODO: the following code can be considered to
1762 # capturing a truly random spectrum.
1763 # add some constraints on the spectrum, i.e. not fully
1765 # Note: one key observation for understanding the following code is,
1766 # that instead of wrapping your head around symmetries in multi-
1767 # dimensional coefficient matrices, one can simply look at the flattened
1768 # version of such a matrix and reshape later. It just works out.
1769 domain_samples_per_input_dim = cls.construct_domain_samples(model)
1771 frequencies = cls.construct_frequencies(model)
1773 coefficients = cls.construct_coefficients(
1774 random_key, model, coefficients_min, coefficients_max, zero_centered
1775 )
1777 values = cls.calculate_values(
1778 domain_samples_per_input_dim, frequencies, coefficients
1779 )
1781 # return all the information we have
1782 return [
1783 domain_samples_per_input_dim.reshape(*model.degree, -1),
1784 values.reshape(model.degree),
1785 coefficients.reshape(model.degree),
1786 ]
1788 @classmethod
1789 def construct_domain_samples(cls, model: Model) -> jnp.ndarray:
1790 """
1791 Builds the input-domain sample grid for the model spectrum.
1793 Going from $[0, 2\\pi]$ with the resolution required for the highest
1794 frequency, permuted with the input dimensionality to get an n-d grid
1795 of domain samples (a "coordinate system").
1797 Args:
1798 model (Model): The quantum circuit model.
1800 Returns:
1801 jnp.ndarray: Domain samples with shape
1802 ($\\prod$ degree, n_input_feat).
1803 """
1804 return jnp.stack(
1805 jnp.meshgrid(
1806 *[jnp.arange(0, 2 * jnp.pi, 2 * jnp.pi / d) for d in model.degree]
1807 )
1808 ).T.reshape(-1, model.n_input_feat)
1810 @classmethod
1811 def construct_frequencies(cls, model: Model) -> jnp.ndarray:
1812 """
1813 Builds the frequency-index grid for the model spectrum.
1815 This has the same shape as the domain samples returned by
1816 `construct_domain_samples`.
1818 Args:
1819 model (Model): The quantum circuit model.
1821 Returns:
1822 jnp.ndarray: Frequency indices with shape
1823 ($\\prod$ degree, n_input_feat).
1824 """
1825 return jnp.stack(jnp.meshgrid(*model.frequencies)).T.reshape(
1826 -1, model.n_input_feat
1827 )
1829 @classmethod
1830 def construct_coefficients(
1831 cls,
1832 random_key: random.PRNGKey,
1833 model: Model,
1834 coefficients_min: float = 0.0,
1835 coefficients_max: float = 1.0,
1836 zero_centered: bool = False,
1837 ) -> jnp.ndarray:
1838 """
1839 Samples the conjugate-symmetric Fourier coefficient vector.
1841 Coefficients are drawn from a uniform circle (see `uniform_circle`).
1842 The offset coefficient (first entry) is either zeroed or made real,
1843 then the spectrum is mirrored to enforce conjugate symmetry.
1845 Args:
1846 random_key (random.PRNGKey): Random number key for JAX.
1847 model (Model): The quantum circuit model.
1848 coefficients_min (float, optional): Minimum value for the
1849 coefficients. Defaults to 0.0.
1850 coefficients_max (float, optional): Maximum value for the
1851 coefficients. Defaults to 1.0.
1852 zero_centered (bool, optional): Whether to zero-center the
1853 coefficients. Defaults to False.
1855 Returns:
1856 jnp.ndarray: Conjugate-symmetric coefficient vector of size
1857 $\\prod$ degree.
1858 """
1859 coefficients = cls.uniform_circle(
1860 random_key,
1861 low=coefficients_min,
1862 high=coefficients_max,
1863 size=math.prod(model.degree) // 2 + 1,
1864 )
1866 # zero center (first coeff = 0)
1867 # we can assume the first coeff is the offset, because we're dealing
1868 # with a non-symmetric spectrum here
1869 if zero_centered:
1870 coefficients = coefficients.at[0].set(0.0)
1871 else:
1872 coefficients = coefficients.at[0].set(coefficients[0].real)
1874 # ensure symmetry (here, non_negative_ is removed!),
1875 # giving us the full coefficients vector
1876 return jnp.concat(
1877 [
1878 jnp.flip(coefficients[..., 1:]).conjugate(),
1879 coefficients,
1880 ],
1881 axis=-1,
1882 )
1884 @classmethod
1885 def calculate_values(
1886 cls,
1887 domain_samples: jnp.ndarray,
1888 frequencies: jnp.ndarray,
1889 coefficients: jnp.ndarray,
1890 ) -> jnp.ndarray:
1891 """
1892 Evaluates the real-valued Fourier series on the domain grid.
1894 Vectorized version of
1895 $f(x) = \\sum_{n=0}^{N-1} c_n e^{i \\omega_n x}$ that takes the input
1896 dimension into account, normalized by the number of coefficients.
1898 Args:
1899 domain_samples (jnp.ndarray): Domain samples with shape
1900 (n_points, n_input_feat).
1901 frequencies (jnp.ndarray): Frequency indices with shape
1902 (n_freqs, n_input_feat).
1903 coefficients (jnp.ndarray): Fourier coefficients with shape
1904 (n_freqs,).
1906 Returns:
1907 jnp.ndarray: Real-valued Fourier series samples with shape
1908 (n_points,).
1909 """
1910 return jnp.real(
1911 (jnp.exp(1j * (domain_samples @ frequencies.T)) * coefficients).sum(axis=1)
1912 / coefficients.size
1913 )
1915 @classmethod
1916 def uniform_circle(
1917 cls,
1918 random_key: random.PRNGKey,
1919 size: Union[jnp.ndarray, List, int],
1920 low=0.0,
1921 high=1.0,
1922 ):
1923 """
1924 Random number generator for complex numbers sampled inside the unit circle
1926 Args:
1927 random_key (random.PRNGKey): Random number key for JAX.
1928 size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
1929 the first dimension will be the number of dimensions.
1930 low (float, optional): Minimum Radius. Defaults to 0.0.
1931 high (float, optional): Maximum Radius. Defaults to 1.0.
1933 Returns
1934 jnp.ndarray: Array of complex numbers with shape of `size`
1935 """
1937 if isinstance(size, int):
1938 size = jnp.array([size])
1940 random_key, random_key1 = random.split(random_key)
1941 return jnp.sqrt(
1942 random.uniform(random_key, size, minval=low, maxval=high)
1943 ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))