Coverage for qml_essentials / coefficients.py: 97%
747 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-09-02 14:50 +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 jaqsi.paulis 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 1.
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.
62 Note:
63 The FFT grid is built from the nominal `model.degree`, which does
64 not account for frequency scaling via `enc_params` or
65 `enc_pulse_params`. For a model whose effective frequencies exceed
66 the nominal degree by a factor $s > 1$, choose `mfs` at least
67 $\\lceil s \\rceil$ (e.g. `mfs=2` covers scalings up to 2),
68 otherwise the scaled components alias.
69 """
70 kwargs.setdefault("force_mean", True)
71 kwargs.setdefault("execution_type", "expval")
73 coeffs, freqs = cls._fourier_transform(model, mfs=mfs, mts=mts, **kwargs)
75 if not jnp.isclose(jnp.sum(coeffs).imag, 0.0, atol=1.0e-6):
76 raise ValueError(
77 f"Spectrum is not real. Imaginary part of coefficients is:\
78 {jnp.sum(coeffs).imag}"
79 )
81 if trim:
82 for ax in range(model.n_input_feat):
83 if coeffs.shape[ax] % 2 == 0:
84 coeffs = np.delete(coeffs, len(coeffs) // 2, axis=ax)
85 freqs = [np.delete(freq, len(freq) // 2, axis=ax) for freq in freqs]
87 if shift:
88 coeffs = jnp.fft.fftshift(coeffs, axes=list(range(model.n_input_feat)))
89 freqs = np.fft.fftshift(freqs)
91 if numerical_cap > 0:
92 # set coeffs below threshold to zero
93 coeffs = jnp.where(
94 jnp.abs(coeffs) < numerical_cap,
95 jnp.zeros_like(coeffs),
96 coeffs,
97 )
99 # Drop frequencies whose coefficients vanish entirely after
100 # capping, so the returned spectrum reflects only the surviving
101 # frequencies. Well-defined only for a single (1-D) frequency
102 # axis; for multi-dim input the rectangular grid is left intact.
103 if model.n_input_feat == 1:
104 if coeffs.ndim == 1:
105 surviving = coeffs != 0
106 else:
107 surviving = jnp.any(coeffs != 0, axis=tuple(range(1, coeffs.ndim)))
108 coeffs = coeffs[surviving]
109 freqs = [freqs[0][surviving]]
111 if len(freqs) == 1:
112 freqs = freqs[0]
114 return coeffs, freqs
116 @classmethod
117 def _fourier_transform(
118 cls, model: Model, mfs: int, mts: int, **kwargs: Any
119 ) -> jnp.ndarray:
120 # Create a frequency vector with as many frequencies as model degrees,
121 # oversampled by mfs
122 n_freqs: jnp.ndarray = jnp.array(
123 [mfs * model.degree[i] for i in range(model.n_input_feat)]
124 )
126 start, stop, step = 0, 2 * mts * jnp.pi, 2 * jnp.pi / n_freqs
127 # Stretch according to the number of frequencies
128 inputs: List = [
129 jnp.arange(start, stop, step[i]) for i in range(model.n_input_feat)
130 ]
132 # permute with input dimensionality
133 nd_inputs = jnp.array(
134 jnp.meshgrid(*[inputs[i] for i in range(model.n_input_feat)])
135 ).T.reshape(-1, model.n_input_feat)
137 # Output vector is not necessarily the same length as input
138 outputs = model(inputs=nd_inputs, **kwargs)
139 outputs = outputs.reshape(
140 *[inputs[i].shape[0] for i in range(model.n_input_feat)], -1
141 ).squeeze()
143 coeffs = jnp.fft.fftn(outputs, axes=list(range(model.n_input_feat)))
145 freqs = [
146 jnp.fft.fftfreq(int(mts * n_freqs[i]), 1 / n_freqs[i])
147 for i in range(model.n_input_feat)
148 ]
149 # freqs = jnp.fft.fftfreq(mts * n_freqs, 1 / n_freqs)
151 # TODO: this could cause issues with multidim input
152 # FIXME: account for different frequencies in multidim input scenarios
153 # Run the fft and rearrange +
154 # normalize the output (using product if multidim)
155 return (
156 coeffs / math.prod(outputs.shape[0 : model.n_input_feat]),
157 freqs,
158 )
160 @classmethod
161 def get_psd(cls, coeffs: jnp.ndarray) -> jnp.ndarray:
162 """
163 Calculates the power spectral density (PSD) from given Fourier coefficients.
165 Args:
166 coeffs (jnp.ndarray): The Fourier coefficients.
168 Returns:
169 jnp.ndarray: The power spectral density.
170 """
171 # TODO: if we apply trim=True in advance, this will be slightly wrong..
173 def abs2(x):
174 return x.real**2 + x.imag**2
176 scale = 2.0 / (len(coeffs) ** 2)
177 return scale * abs2(coeffs)
179 @classmethod
180 def evaluate_Fourier_series(
181 cls,
182 coefficients: jnp.ndarray,
183 frequencies: jnp.ndarray,
184 inputs: Union[jnp.ndarray, list, float],
185 ) -> float:
186 """
187 Evaluate the function value of a Fourier series at one point.
189 Args:
190 coefficients (jnp.ndarray): Coefficients of the Fourier series.
191 frequencies (jnp.ndarray): Corresponding frequencies.
192 inputs (jnp.ndarray): Point at which to evaluate the function.
193 Returns:
194 float: The function value at the input point.
195 """
196 coefficients = jnp.asarray(coefficients)
198 def flatten_grid(freq_axes):
199 freq_axes = [jnp.asarray(freq) for freq in freq_axes]
200 freq_grid = jnp.stack(jnp.meshgrid(*freq_axes, indexing="ij"), axis=-1)
201 flat_frequencies = freq_grid.reshape(-1, len(freq_axes))
202 flat_coefficients = coefficients.reshape(
203 flat_frequencies.shape[0], *coefficients.shape[len(freq_axes) :]
204 )
205 return flat_coefficients, flat_frequencies
207 if isinstance(frequencies, list):
208 flat_coefficients, flat_frequencies = flatten_grid(frequencies)
209 else:
210 frequencies = jnp.asarray(frequencies)
211 if frequencies.ndim == 1:
212 flat_frequencies = frequencies[:, jnp.newaxis]
213 flat_coefficients = coefficients.reshape(
214 flat_frequencies.shape[0], *coefficients.shape[1:]
215 )
216 else:
217 n_features, n_axis_freqs = frequencies.shape
218 is_axis_frequencies = (
219 coefficients.shape[:n_features] == (n_axis_freqs,) * n_features
220 )
222 if is_axis_frequencies:
223 flat_coefficients, flat_frequencies = flatten_grid(frequencies)
224 else:
225 flat_frequencies = frequencies
226 flat_coefficients = coefficients.reshape(
227 flat_frequencies.shape[0], *coefficients.shape[1:]
228 )
230 inputs = jnp.asarray(inputs)
231 if inputs.ndim == 0:
232 inputs = inputs.reshape(1, 1)
233 elif inputs.ndim == 1:
234 if flat_frequencies.shape[1] == 1:
235 inputs = inputs[:, jnp.newaxis]
236 elif inputs.shape[0] == flat_frequencies.shape[1]:
237 inputs = inputs[jnp.newaxis, :]
238 else:
239 inputs = jnp.repeat(
240 inputs[:, jnp.newaxis], flat_frequencies.shape[1], axis=1
241 )
242 exponents = jnp.exp(1j * (inputs @ flat_frequencies.T))
243 exp = jnp.tensordot(exponents, flat_coefficients, axes=([1], [0]))
245 return jnp.squeeze(jnp.real(exp))
248class FourierTree:
249 """
250 Sine-cosine tree representation for the algorithm by Nemkov et al.
252 Computes the analytical Fourier coefficients/frequencies of a Pauli-Clifford
253 circuit. The symbolic structure of the tree (which Pauli rotations
254 contribute sine/cosine factors to which leaf, and the leaf observables) is
255 built once in NumPy; the parameter-dependent coefficients are then obtained
256 with a small number of vectorised JAX operations, so the result remains
257 jittable / differentiable with respect to the model parameters.
259 The resulting spectrum is the d-dimensional set of frequency vectors,
260 where $d$ is the input dimensionality.
262 **Usage**:
263 ```
264 model = Model(...)
265 tree = FourierTree(model)
266 exp = tree() # expectation value
267 coeff_list, freq_list = tree.get_spectrum()
268 ```
269 """
271 def __init__(self, model: Model):
272 """
273 Tree initialisation, based on the Pauli-Clifford representation of a
274 model.
276 Args:
277 model (Model): The Model, for which to build the tree.
278 """
279 self.model = model
280 self.n_qubits = model.n_qubits
282 # A single (de-batched) parameter set drives the whole tree.
283 self._params = self._single_param_set(model.params)
285 # Canonical Pauli-Clifford structure, recorded once at a fixed base
286 # input. The base value is irrelevant to the structure (it only sets
287 # the rotation angles, not which Pauli words appear).
288 base_inputs = np.ones(model.n_input_feat)
289 operations, observables = self._build_canonical_tape(self._params, base_inputs)
291 self.parameters = [
292 jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
293 ]
294 self.n_params = len(self.parameters)
296 # Pauli generators of the (canonical) rotations, as symbolic words.
297 self.pauli_words: List[PauliWord] = [
298 PauliWord.from_operation(op, self.n_qubits) for op in operations
299 ]
301 # Cumulative X/Y support of the rotations[0..k] (for light-cone early
302 # stopping). cumulative_xy[k] is True on every qubit touched by an X/Y
303 # generator in any rotation up to index k.
304 self.cumulative_xy: List[np.ndarray] = []
305 running = np.zeros(self.n_qubits, dtype=bool)
306 for pw in self.pauli_words:
307 running = np.logical_or(running, pw.xy_mask)
308 self.cumulative_xy.append(running.copy())
310 # Observable Pauli words (one tree root each).
311 self.observable_words: List[PauliWord] = [
312 PauliWord.from_operation(obs, self.n_qubits) for obs in observables
313 ]
315 # Identify the input-encoding columns, their feature, and integer
316 # frequency scaling directly from the tape (no per-gate tagging). Sets
317 # ``input_indices``, ``all_input_indices``, ``input_scaling``,
318 # ``var_positions`` and ``features``.
319 self._detect_inputs(base_inputs)
321 # The explicit leaf structure is built lazily: for deep circuits the
322 # number of tree paths explodes combinatorially, while the canonical
323 # form above (and the merged-state support DP) remain cheap.
324 self._structure_built = False
326 def _ensure_structure(self) -> None:
327 """Build the explicit leaf/spectrum structure on first use."""
328 if not self._structure_built:
329 # Symbolic structure: per root (S, C, terms) leaf arrays ...
330 self._build_leaf_arrays()
331 # ... and the parameter-independent frequency/weight structure.
332 self._build_spectrum_structure()
333 self._structure_built = True
335 def _single_param_set(self, params) -> jnp.ndarray:
336 """De-batch the model parameters to the single set the tree describes.
338 Models can carry batched parameters (e.g. after FCC sampling); the tree
339 is defined for one set, so fall back to the first and warn.
340 """
341 params = jnp.asarray(params)
342 if params.ndim > 2 and params.shape[0] > 1:
343 warnings.warn(
344 f"FourierTree supports a single parameter set; using the first "
345 f"of {params.shape[0]} batched parameter sets.",
346 UserWarning,
347 )
348 params = params[0]
349 return params
351 def _build_canonical_tape(self, params, inputs):
352 """Record the circuit and transform it to Pauli-Clifford normal form.
354 Returns the ``(operations, observables)`` of the canonical circuit
355 (see :meth:`PauliCircuit.from_parameterised_circuit`).
356 """
357 params = self._single_param_set(params)
358 inputs = self.model._inputs_validation(inputs)
359 raw_tape = self.model.script.record(params=params, inputs=inputs)
360 _, obs_list = self.model._build_obs()
361 return PauliCircuit.from_parameterised_circuit(
362 raw_tape, observables=obs_list, n_qubits=self.n_qubits
363 )
365 def _canonical_parameters(self, inputs) -> np.ndarray:
366 """Recorded canonical rotation angles (1-D float array) for ``inputs``."""
367 operations, _ = self._build_canonical_tape(self._params, inputs)
368 return np.array(
369 [float(jnp.squeeze(p)) for p in PauliCircuit.get_parameters(operations)]
370 )
372 def _detect_inputs(self, base_inputs: np.ndarray) -> None:
373 r"""Infer the input-encoding columns directly from the tape (tag-free).
375 Each encoding rotation applies an angle :math:`\omega_k\,x_f` that is
376 linear in a single input feature :math:`x_f`, and Clifford commutation
377 only multiplies a rotation generator by :math:`\pm 1`. Every canonical
378 rotation angle is therefore an affine function of the inputs, so
379 perturbing one feature at a time and differencing the recorded angles
380 isolates exactly the columns that depend on it, together with the
381 signed scaling :math:`\omega_k` (rational for diagonal-Hamiltonian
382 encodings such as Golomb).
384 Sets :attr:`input_indices` (``{feature: [columns]}``),
385 :attr:`all_input_indices`, :attr:`input_scaling` (per column, ``1`` for
386 variational columns), :attr:`var_positions`, and :attr:`features`.
388 Raises:
389 NotImplementedError: If a rotation depends on more than one feature
390 (the tree requires single-feature encodings).
391 """
392 tol = 1e-6
393 d = self.model.n_input_feat
394 base = np.asarray(base_inputs, dtype=float)
395 p_base = np.array([float(p) for p in self.parameters])
397 # response[f, k] = d(angle_k) / d(x_f), the linear response of column k.
398 response = np.zeros((d, self.n_params))
399 for f in range(d):
400 step = base.copy()
401 step[f] += 1.0
402 response[f] = self._canonical_parameters(step) - p_base
404 input_indices: Dict[int, list] = defaultdict(list)
405 all_input_indices: List[int] = []
406 scaling = np.ones(self.n_params, dtype=np.float64)
407 for k in range(self.n_params):
408 feats = np.flatnonzero(np.abs(response[:, k]) > tol)
409 if feats.size == 0:
410 continue # variational column
411 if feats.size > 1:
412 raise NotImplementedError(
413 f"Rotation {k} depends on multiple input features "
414 f"{feats.tolist()}; the Fourier tree requires each encoding "
415 "rotation to be linear in a single feature."
416 )
417 f = int(feats[0])
418 omega = float(response[f, k])
419 # Scalings may be rational (e.g. diagonal-Hamiltonian / Golomb
420 # encodings, whose per-Pauli-string scalings are dyadic rationals);
421 # the integer-frequency model spectrum is recovered downstream.
422 input_indices[f].append(k)
423 all_input_indices.append(k)
424 scaling[k] = omega
426 self.input_indices = input_indices
427 self.all_input_indices = all_input_indices
428 self.input_scaling = scaling
429 input_set = set(all_input_indices)
430 self.var_positions = np.array(
431 [i for i in range(self.n_params) if i not in input_set], dtype=np.int64
432 )
433 # Ordered list of input feature keys (d-dimensional spectrum).
434 self.features = sorted(input_indices.keys())
436 # Symbolic tree construction (NumPy)
437 def _build_leaf_arrays(self) -> None:
438 """Collect the tree leaves for every root into integer count matrices.
440 For each root (observable) this produces:
441 - ``S``: (n_leaves, n_params) sine-factor counts per parameter,
442 - ``C``: (n_leaves, n_params) cosine-factor counts per parameter,
443 - ``terms``: (n_leaves,) complex leaf constants ``<0|O_leaf|0>``.
444 """
445 self.leaf_arrays: List[Tuple[np.ndarray, np.ndarray, np.ndarray]] = []
446 for obs_word in self.observable_words:
447 leaves: List[Tuple[np.ndarray, np.ndarray, complex]] = []
448 zeros = np.zeros(self.n_params, dtype=np.int64)
449 self._collect_leaves(
450 obs_word, self.n_params - 1, zeros.copy(), zeros.copy(), leaves
451 )
452 if leaves:
453 S = np.stack([leaf[0] for leaf in leaves])
454 C = np.stack([leaf[1] for leaf in leaves])
455 terms = np.array([leaf[2] for leaf in leaves], dtype=np.complex128)
456 else:
457 S = np.zeros((0, self.n_params), dtype=np.int64)
458 C = np.zeros((0, self.n_params), dtype=np.int64)
459 terms = np.zeros(0, dtype=np.complex128)
460 self.leaf_arrays.append((S, C, terms))
462 def _collect_leaves(
463 self,
464 observable: PauliWord,
465 pauli_idx: int,
466 sin_counts: np.ndarray,
467 cos_counts: np.ndarray,
468 leaves: List[Tuple[np.ndarray, np.ndarray, complex]],
469 ) -> None:
470 """Recursively enumerate the leaves of the coefficient tree.
472 The incoming sine/cosine factor (from the parent edge) is already
473 accumulated into ``sin_counts``/``cos_counts``. This fuses the tree
474 construction and leaf traversal of the original implementation into a
475 single NumPy pass (no per-node JAX scatter updates).
476 """
477 if self._early_stopping_possible(pauli_idx, observable):
478 return
480 # Skip trailing Pauli rotations that commute with the observable.
481 while pauli_idx >= 0:
482 last = self.pauli_words[pauli_idx]
483 if not observable.commutes_with(last):
484 break
485 pauli_idx -= 1
486 else: # leaf reached
487 term = observable.zero_expectation()
488 if term != 0:
489 leaves.append((sin_counts, cos_counts, term))
490 return
492 last = self.pauli_words[pauli_idx]
494 # Left child: cosine factor for this parameter, same observable.
495 cos_left = cos_counts.copy()
496 cos_left[pauli_idx] += 1
497 self._collect_leaves(
498 observable, pauli_idx - 1, sin_counts.copy(), cos_left, leaves
499 )
501 # Right child: sine factor, observable becomes P . O.
502 sin_right = sin_counts.copy()
503 sin_right[pauli_idx] += 1
504 self._collect_leaves(
505 last.compose(observable),
506 pauli_idx - 1,
507 sin_right,
508 cos_counts.copy(),
509 leaves,
510 )
512 def _early_stopping_possible(self, pauli_idx: int, observable: PauliWord) -> bool:
513 """Whether a node can be discarded (all reachable expectations vanish).
515 Mirrors the criterion of Nemkov et al. (light cone): a qubit on which
516 the observable carries an X/Y must be covered by an X/Y generator of
517 some remaining rotation (rotations[0..pauli_idx]); otherwise that X/Y can
518 never be rotated into a diagonal term and the whole node contributes
519 zero. Equivalently, the node survives iff every qubit is either I/Z in
520 the observable or covered by the cumulative rotation X/Y support.
521 """
522 obs_iz = np.logical_not(observable.xy_mask)
523 combined = np.logical_or(obs_iz, self.cumulative_xy[pauli_idx]).all()
524 return not bool(combined)
526 # Frequency / weight structure (NumPy, parameter independent)
527 def _build_spectrum_structure(self) -> None:
528 """Build, per root, the frequency vectors and the (n_freq, n_leaves)
529 weight matrix ``W`` such that ``coeffs = W @ (terms * variational)``.
530 """
531 self.freqs_per_root: List[np.ndarray] = []
532 self.weights_per_root: List[np.ndarray] = []
533 d = len(self.features)
535 for S, C, _ in self.leaf_arrays:
536 n_leaves = S.shape[0]
537 freq_to_col: Dict[tuple, np.ndarray] = defaultdict(
538 lambda: np.zeros(n_leaves, dtype=np.complex128)
539 )
540 for leaf in range(n_leaves):
541 # One expansion factor per *active* input column, each carrying
542 # its feature axis and integer frequency scaling. Per leaf a
543 # column contributes at most one sin/cos factor (square-free),
544 # but different columns of the same feature may carry different
545 # scalings, so they are expanded individually and convolved
546 # rather than aggregating counts (which would assume a common
547 # unit scaling).
548 col_factors: List[List[Tuple[int, float, float]]] = []
549 half_exp = 0
550 for axis, feat in enumerate(self.features):
551 for k in self.input_indices[feat]:
552 s = int(S[leaf, k])
553 c = int(C[leaf, k])
554 if s == 0 and c == 0:
555 continue
556 half_exp += s + c
557 w_k = float(self.input_scaling[k])
558 col_factors.append(
559 [
560 (axis, o * w_k, wt)
561 for o, wt in self._binomial_terms(s, c)
562 ]
563 )
564 half = 0.5**half_exp
566 if d == 0:
567 freq_to_col[(0,)][leaf] += half
568 continue
570 if not col_factors:
571 freq_to_col[(0,) * d][leaf] += half
572 continue
574 for combo in itertools.product(*col_factors):
575 omega = [0.0] * d
576 weight = half
577 for axis, o, wt in combo:
578 omega[axis] += o
579 weight *= wt
580 # Snap to a tolerance so dyadic-rational contributions that
581 # are numerically equal share a single frequency key.
582 key = tuple(round(v, 9) for v in omega)
583 freq_to_col[key][leaf] += weight
585 if freq_to_col:
586 omegas = sorted(freq_to_col.keys())
587 W = np.stack([freq_to_col[o] for o in omegas]) # (n_freq, n_leaves)
588 freqs = np.array(omegas, dtype=float) # (n_freq, d)
589 # Diagonal-Hamiltonian (Golomb) encodings produce rational
590 # per-rotation scalings but integer model frequencies; snap to
591 # int64 when every entry is integral so integer-encoding
592 # consumers are unchanged, otherwise keep the rational floats.
593 rounded = np.rint(freqs)
594 if np.all(np.abs(freqs - rounded) < 1e-6):
595 freqs = rounded.astype(np.int64)
596 else:
597 freqs = np.zeros((1, max(d, 1)), dtype=np.int64)
598 W = np.zeros((1, n_leaves), dtype=np.complex128)
600 # Collapse to 1-D frequency array for the single-feature case.
601 if freqs.shape[1] == 1:
602 freqs = freqs[:, 0]
603 self.freqs_per_root.append(freqs)
604 # Keep W in NumPy complex128: its entries are dyadic rationals
605 # (binomial weights x 0.5^k x i^m), which are exact in float64 --
606 # this allows exact symbolic zero-tests in get_exact_support.
607 self.weights_per_root.append(W)
609 @staticmethod
610 def _binomial_terms(s: int, c: int) -> List[Tuple[int, float]]:
611 """Expand ``cos^c (i sin)^s`` in ``e^{i omega x}`` (without the 0.5 factor).
613 Returns a list of ``(omega, weight)`` with
614 ``omega = 2a + 2b - s - c`` and ``weight = C(s,a) C(c,b) (-1)^{s-a}``.
615 """
616 terms = []
617 for a in range(s + 1):
618 for b in range(c + 1):
619 weight = math.comb(s, a) * math.comb(c, b) * (-1) ** (s - a)
620 terms.append((2 * a + 2 * b - s - c, float(weight)))
621 return terms
623 # Vectorised numeric evaluation (JAX)
624 @staticmethod
625 def _safe_pow(base: jnp.ndarray, exp: jnp.ndarray) -> jnp.ndarray:
626 """Elementwise ``base ** exp`` for real base and non-negative integer
627 exponents, correct for negative bases (avoids ``log`` of negatives).
629 Args:
630 base: real array of shape ``(n,)``.
631 exp: integer array of shape ``(n_leaves, n)``.
632 """
633 mag = jnp.abs(base)[None, :] ** exp
634 sign = jnp.where(exp % 2 == 0, 1.0, jnp.sign(base)[None, :])
635 return sign * mag
637 _I_POW = None # set lazily to jnp.array([1, 1j, -1, -1j])
639 def _leaf_factors(
640 self, S: np.ndarray, C: np.ndarray, columns: np.ndarray
641 ) -> jnp.ndarray:
642 """Per-leaf product ``prod_i cos(theta_i)^{C} (i sin(theta_i))^{S}`` over
643 the given parameter ``columns`` (vectorised over leaves).
644 """
645 if FourierTree._I_POW is None:
646 FourierTree._I_POW = jnp.array([1, 1j, -1, -1j])
648 if S.shape[0] == 0:
649 return jnp.zeros(0, dtype=jnp.complex64)
651 theta = jnp.stack([self.parameters[i] for i in columns])
652 S_sub = jnp.asarray(S[:, columns])
653 C_sub = jnp.asarray(C[:, columns])
655 cos_part = self._safe_pow(jnp.cos(theta), C_sub)
656 sin_mag = self._safe_pow(jnp.sin(theta), S_sub)
657 i_part = FourierTree._I_POW[S_sub % 4]
658 return jnp.prod(cos_part * sin_mag * i_part, axis=1)
660 def __call__(
661 self,
662 params: Optional[jnp.ndarray] = None,
663 inputs: Optional[jnp.ndarray] = None,
664 **kwargs,
665 ) -> jnp.ndarray:
666 """
667 Evaluate the expectation value(s) of the model's observables via the
668 sine-cosine tree (equivalent to the circuit expectation).
670 Args:
671 params (Optional[jnp.ndarray]): Model parameters. Defaults to the
672 model's parameters.
673 inputs (Optional[jnp.ndarray]): Inputs to the circuit. Defaults to 1.
675 Returns:
676 jnp.ndarray: Expectation value per observable (or their mean if
677 ``force_mean`` is set).
679 Raises:
680 NotImplementedError: For execution types other than "expval" or when
681 noise is requested.
682 """
683 params = (
684 self.model._params_validation(params)
685 if params is not None
686 else self.model.params
687 )
688 inputs = (
689 self.model._inputs_validation(inputs)
690 if inputs is not None
691 else self.model._inputs_validation(1.0)
692 )
694 if kwargs.get("execution_type", "expval") != "expval":
695 raise NotImplementedError(
696 f'Currently, only "expval" execution type is supported when '
697 f"building FourierTree. Got {kwargs.get('execution_type', 'expval')}."
698 )
699 if kwargs.get("noise_params", None) is not None:
700 raise NotImplementedError(
701 "Currently, noise is not supported when building FourierTree."
702 )
704 # Re-derive the (canonical) parameter values for the requested inputs;
705 # the tree structure (leaf arrays) is unchanged.
706 operations, _ = self._build_canonical_tape(params, inputs)
707 self.parameters = [
708 jnp.squeeze(p) for p in PauliCircuit.get_parameters(operations)
709 ]
711 self._ensure_structure()
712 all_columns = np.arange(self.n_params, dtype=np.int64)
713 results = []
714 for S, C, terms in self.leaf_arrays:
715 factors = self._leaf_factors(S, C, all_columns)
716 results.append(jnp.real(jnp.sum(jnp.asarray(terms) * factors)))
717 results = jnp.array(results)
719 if kwargs.get("force_mean", False):
720 return jnp.mean(results)
721 return results
723 def get_spectrum(
724 self, force_mean: bool = False
725 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
726 """
727 Compute the Fourier spectrum (coefficients and frequencies) of the tree.
729 Args:
730 force_mean (bool, optional): Average the coefficients over all
731 observables (roots). Defaults to False.
733 Returns:
734 Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
735 - List of coefficients, one entry per observable (root).
736 - List of corresponding frequencies, one entry per root.
737 When ``force_mean`` is set, both lists have a single entry.
738 """
739 self._ensure_structure()
740 per_root_coeffs: List[jnp.ndarray] = []
741 for (S, C, terms), W in zip(self.leaf_arrays, self.weights_per_root):
742 leaf_const = jnp.asarray(terms) * self._leaf_factors(
743 S, C, self.var_positions
744 )
745 per_root_coeffs.append(jnp.asarray(W) @ leaf_const)
747 return self._combine_roots(per_root_coeffs, self.freqs_per_root, force_mean)
749 def get_exact_support(self, method: str = "tree") -> List[np.ndarray]:
750 r"""Symbolically derive the exact frequency support (no sampling).
752 A frequency :math:`\omega` belongs to the exact spectrum iff its
753 coefficient :math:`c_\omega(\theta) = \sum_l W_{\omega l}\,
754 \text{term}_l\, v_l(\theta)` is not identically zero in the
755 variational parameters :math:`\theta`.
757 Two methods are available:
759 - ``"tree"`` (default, fully exact): enumerates the explicit tree
760 leaves. Because the branch index strictly decreases along every tree
761 path, each parameter contributes **at most one** sine *or* cosine
762 factor per leaf (:math:`S_{li}, C_{li} \in \{0, 1\}`). Every
763 variational leaf factor :math:`v_l` is therefore a *square-free*
764 monomial over :math:`\{1, \cos\theta_i, i\sin\theta_i\}`, and
765 monomials with distinct signatures are linearly independent functions
766 (no :math:`\cos^2 + \sin^2` identities can arise without squares).
767 Hence
769 .. math::
770 c_\omega \equiv 0 \iff \sum_{l \in g} W_{\omega l}\,\text{term}_l
771 = 0 \quad \text{for every signature group } g.
773 Since all involved quantities are dyadic rationals times
774 :math:`\{\pm 1, \pm i\}`, the group sums are exact in float64 and the
775 zero-test is exact. The number of leaves can however grow
776 exponentially with circuit depth.
778 - ``"dp"`` (scalable): merges tree nodes with identical
779 ``(rotation index, observable)`` — at most ``n_params * 4^n_qubits``
780 states — and tracks, per state, the achievable per-feature sine/cosine
781 count vectors ``(s_f, c_f)`` as a mixed-radix bitmask. Each feature's
782 support is the union of the (exact) expansion supports of
783 :math:`\cos^{c_f} x_f\, (i \sin x_f)^{s_f}`, and the model support is
784 their Cartesian product across features. This is exact per tree path
785 (including interior zero coefficients of the expansions), but unlike
786 ``"tree"`` it cannot detect coefficients that cancel identically
787 *across* paths with identical variational signatures (e.g. directly
788 repeated encodings). It therefore yields a tight superset in such
789 corner cases. Supports any number of input features, but requires
790 unit-magnitude input scaling: per-gate :math:`|\omega| \neq 1`
791 scalings (e.g. Golomb encodings) are rejected — use ``"tree"``.
793 Args:
794 method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).
796 Returns:
797 List[np.ndarray]: For each observable (root), the frequency vectors
798 with not-identically-zero coefficient — shape ``(n_freq,)`` for a
799 single input feature, ``(n_freq, n_features)`` otherwise.
800 """
801 if method == "dp":
802 return self._support_dp()
803 if method != "tree":
804 raise ValueError(f"Unknown method '{method}'. Use 'tree' or 'dp'.")
806 self._ensure_structure()
807 supports = []
808 for (S, C, terms), W, freqs in zip(
809 self.leaf_arrays, self.weights_per_root, self.freqs_per_root
810 ):
811 freqs = np.asarray(freqs)
812 n_leaves = S.shape[0]
813 if n_leaves == 0:
814 supports.append(freqs[:0])
815 continue
817 # Group leaves by their variational sine/cosine signature.
818 signature = np.hstack([S[:, self.var_positions], C[:, self.var_positions]])
819 _, groups = np.unique(signature, axis=0, return_inverse=True)
820 n_groups = int(groups.max()) + 1
822 # Per-group sums of W[omega, l] * term_l, accumulated exactly.
823 contrib = (W * terms[None, :]).T # (n_leaves, n_freq)
824 group_sums = np.zeros((n_groups, W.shape[0]), dtype=np.complex128)
825 np.add.at(group_sums, groups, contrib)
827 mask = (np.abs(group_sums) > 1e-12).any(axis=0) # (n_freq,)
828 supports.append(freqs[mask])
829 return supports
831 def _support_dp(self) -> List[np.ndarray]:
832 """Merged-state dynamic program for the frequency support.
834 Instead of enumerating all (worst-case exponentially many) tree paths,
835 nodes are merged on ``(rotation index, bare observable)``. Each state
836 stores the set of achievable per-axis count vectors
837 ``(c_0, s_0, ..., c_{d-1}, s_{d-1})`` as a mixed-radix bitmask (one
838 digit per feature sine/cosine count), so transitions are O(1) big-int
839 operations. See :meth:`get_exact_support` for semantics and
840 limitations.
841 """
842 # Count aggregation is valid as long as every input rotation has
843 # unit-magnitude scaling. A Clifford-commutation sign flip (scaling -1)
844 # leaves the frequency support unchanged -- cos is even and sin odd, so
845 # the sign only flips the coefficient, which the support ignores -- but a
846 # genuine per-gate scaling (|omega| != 1, e.g. Golomb / heterogeneous
847 # frequencies) cannot be represented by sin/cos counts.
848 if self.all_input_indices and np.any(
849 np.abs(self.input_scaling[self.all_input_indices]) != 1
850 ):
851 raise NotImplementedError(
852 "The 'dp' support method does not support non-unit input "
853 "frequency scaling (it aggregates sin/cos counts and cannot "
854 "represent per-gate scalings); use method='tree'."
855 )
857 n = self.n_qubits
858 d = len(self.features)
859 # Pack the achievable per-axis sine/cosine counts into one big-int as a
860 # mixed-radix bitmask over (c_0, s_0, ..., c_{d-1}, s_{d-1}). Axis a's
861 # counts range 0..n_a (n_a input rotations encode feature a), so each
862 # digit has radix n_a + 1; a left-shift by a digit's place value
863 # increments that count, OR is set-union. (For d == 1 this reduces to
864 # shift_c = 1, shift_s = n_inp + 1, i.e. the flat (s, c) layout.)
865 ranges = [len(self.input_indices[self.features[a]]) + 1 for a in range(d)]
866 shift_c = [0] * d
867 shift_s = [0] * d
868 place = 1
869 for a in range(d):
870 shift_c[a] = place
871 place *= ranges[a]
872 shift_s[a] = place
873 place *= ranges[a]
874 # Feature axis of each input rotation (-1 if variational), in the same
875 # enumerate(self.features) order the tree method uses for its axes.
876 axis_of_col = np.full(self.n_params, -1, dtype=np.int64)
877 for a in range(d):
878 for k in self.input_indices[self.features[a]]:
879 axis_of_col[k] = a
881 def encode(word: PauliWord) -> Tuple[int, int]:
882 x = z = 0
883 for q in range(n):
884 x |= int(word.x[q]) << q
885 z |= int(word.z[q]) << q
886 return x, z
888 paulis = [encode(w) for w in self.pauli_words]
889 cum_xy = []
890 running = 0
891 for xp, _ in paulis:
892 running |= xp
893 cum_xy.append(running)
895 def parity(v: int) -> int:
896 return bin(v).count("1") & 1
898 def dp(idx: int, xo: int, zo: int, memo: dict) -> int:
899 # Light-cone early stopping (cf. _early_stopping_possible).
900 if idx >= 0 and (xo & ~cum_xy[idx]):
901 return 0
902 # Skip trailing rotations that commute with the observable.
903 while idx >= 0:
904 xp, zp = paulis[idx]
905 if parity(xo & zp) ^ parity(zo & xp):
906 break
907 idx -= 1
908 else: # leaf: counts (s=0, c=0) iff the observable is diagonal
909 return 1 if xo == 0 else 0
910 key = (idx, xo, zo)
911 hit = memo.get(key)
912 if hit is not None:
913 return hit
914 xp, zp = paulis[idx]
915 cos_child = dp(idx - 1, xo, zo, memo)
916 sin_child = dp(idx - 1, xo ^ xp, zo ^ zp, memo)
917 a = int(axis_of_col[idx])
918 if a >= 0:
919 # Active input gate: cosine increments c_a, sine increments s_a.
920 val = (cos_child << shift_c[a]) | (sin_child << shift_s[a])
921 else:
922 val = cos_child | sin_child
923 memo[key] = val
924 return val
926 # Recursion depth is bounded by the number of rotations.
927 old_limit = sys.getrecursionlimit()
928 sys.setrecursionlimit(max(old_limit, self.n_params + 1000))
929 try:
930 supports = []
931 for obs in self.observable_words:
932 memo: dict = {}
933 xo, zo = encode(obs)
934 mask = dp(self.n_params - 1, xo, zo, memo)
935 supports.append(self._dp_mask_to_support(mask, d, ranges))
936 finally:
937 sys.setrecursionlimit(old_limit)
938 return supports
940 def _dp_mask_to_support(self, mask: int, d: int, ranges: List[int]) -> np.ndarray:
941 """Decode a count bitmask (see :meth:`_support_dp`) into a frequency
942 support. Each set bit is a per-axis count vector ``(c_a, s_a)``; the
943 per-axis expansion supports are combined across axes (Cartesian
944 product) and unioned over all bits.
946 Returns an ``(n_freq, d)`` array, collapsed to 1-D for ``d <= 1`` to
947 match :meth:`get_exact_support` with ``method="tree"`` (``d == 0`` keeps
948 only the DC term).
949 """
950 tupleset: set = set()
951 while mask:
952 bit = mask & -mask
953 i = bit.bit_length() - 1
954 rem = i
955 axis_freqs = []
956 for a in range(d):
957 c_a = rem % ranges[a]
958 rem //= ranges[a]
959 s_a = rem % ranges[a]
960 rem //= ranges[a]
961 axis_freqs.append(sorted(self._expansion_support(s_a, c_a)))
962 tupleset.update(itertools.product(*axis_freqs))
963 mask ^= bit
965 if d >= 2:
966 if not tupleset:
967 return np.empty((0, d), dtype=np.int64)
968 return np.array(sorted(tupleset), dtype=np.int64)
969 # d in {0, 1}: collapse to a 1-D frequency array.
970 flat = sorted(t[0] for t in tupleset) if d == 1 else ([0] if tupleset else [])
971 return np.array(flat, dtype=np.int64)
973 @staticmethod
974 @lru_cache(maxsize=None)
975 def _expansion_support(s: int, c: int) -> frozenset:
976 r"""Frequencies with non-zero coefficient in :math:`\cos^c x (i\sin x)^s`.
978 Computed exactly with integer arithmetic via the polynomial
979 :math:`(t - 1)^s (t + 1)^c` (with :math:`t = e^{2ix}` up to a shift);
980 interior coefficients can vanish, e.g. :math:`\cos x \sin x` only
981 contains :math:`\pm 2`.
982 """
983 coeffs = [1]
984 for _ in range(s): # multiply by (t - 1)
985 new = [0] * (len(coeffs) + 1)
986 for i, a in enumerate(coeffs):
987 new[i + 1] += a
988 new[i] -= a
989 coeffs = new
990 for _ in range(c): # multiply by (t + 1)
991 new = [0] * (len(coeffs) + 1)
992 for i, a in enumerate(coeffs):
993 new[i + 1] += a
994 new[i] += a
995 coeffs = new
996 m = s + c
997 return frozenset(2 * k - m for k, a in enumerate(coeffs) if a != 0)
999 def _combine_roots(
1000 self,
1001 per_root_coeffs: List[jnp.ndarray],
1002 per_root_freqs: List[np.ndarray],
1003 force_mean: bool,
1004 ) -> Tuple[List[jnp.ndarray], List[jnp.ndarray]]:
1005 """Assemble the per-root spectra, optionally averaging over roots."""
1006 if not force_mean:
1007 coefficients = [jnp.asarray(c) for c in per_root_coeffs]
1008 frequencies = [jnp.asarray(f) for f in per_root_freqs]
1009 return coefficients, frequencies
1011 # Average over roots on the union of all frequency vectors. Keys are
1012 # snapped to a tolerance so rational frequencies dedup robustly.
1013 accum: Dict[tuple, complex] = defaultdict(complex)
1014 for coeffs, freqs in zip(per_root_coeffs, per_root_freqs):
1015 freqs_np = np.asarray(freqs)
1016 for k in range(freqs_np.shape[0]):
1017 key = (
1018 (round(float(freqs_np[k]), 9),)
1019 if freqs_np.ndim == 1
1020 else tuple(round(float(v), 9) for v in freqs_np[k])
1021 )
1022 accum[key] += complex(coeffs[k])
1023 n_roots = max(len(per_root_coeffs), 1)
1024 keys = sorted(accum.keys())
1025 mean_coeffs = jnp.array([accum[k] / n_roots for k in keys])
1026 freq_arr = np.array(keys, dtype=float)
1027 rounded = np.rint(freq_arr)
1028 if np.all(np.abs(freq_arr - rounded) < 1e-6):
1029 freq_arr = rounded.astype(np.int64)
1030 if freq_arr.shape[1] == 1:
1031 freq_arr = freq_arr[:, 0]
1032 return [mean_coeffs], [jnp.asarray(freq_arr)]
1035class FCC:
1036 @classmethod
1037 def get_fcc(
1038 cls,
1039 model: Model,
1040 n_samples: int,
1041 random_key: Optional[random.PRNGKey] = None,
1042 method: Optional[str] = "pearson",
1043 scale: Optional[bool] = False,
1044 weight: Optional[bool] = False,
1045 trim_redundant: Optional[bool] = True,
1046 **kwargs,
1047 ) -> float:
1048 """
1049 Shortcut method to get just the FCC.
1050 This includes
1051 1. What is done in `get_fourier_fingerprint`:
1052 1. Calculating the coefficients (using `n_samples`)
1053 2. Correlating the result from 1) using `method`
1054 3. Weighting the correlation matrix (if `weight` is True)
1055 4. Remove redundancies
1056 2. What is done in `calculate_fcc`:
1057 1. Absolute of the fingerprint
1058 2. Average
1060 Args:
1061 model (Model): The QFM model
1062 n_samples (int): Number of samples to calculate average of coefficients
1063 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1064 initialization. If None, uses the model's internal random key.
1065 method (Optional[str], optional): Correlation method. Supported values are
1066 "pearson", "complex_pearson", "spearman", and "covariance".
1067 Defaults to "pearson".
1068 scale (Optional[bool], optional): Whether to scale the number of samples.
1069 Defaults to False.
1070 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1071 Defaults to False.
1072 trim_redundant (Optional[bool], optional): Whether to remove redundant
1073 correlations. Defaults to False.
1074 **kwargs (Any): Additional keyword arguments for the model function.
1076 Returns:
1077 float: The FCC
1078 """
1080 # Memory-efficient fast path
1081 if trim_redundant and not weight:
1082 _, coeffs, freqs = cls._calculate_coefficients(
1083 model, n_samples, random_key, scale, **kwargs
1084 )
1085 pos_idx = cls._calculate_mask(freqs)
1086 coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
1087 coeffs_sub = coeffs_flat[pos_idx]
1089 fp = cls._correlate(coeffs_sub.transpose(), method=method)
1090 abs_fp = jnp.abs(fp)
1091 diag = jnp.abs(jnp.diagonal(fp))
1093 total_sum = jnp.nansum(abs_fp)
1094 total_count = jnp.sum(jnp.isfinite(abs_fp))
1095 diag_sum = jnp.nansum(diag)
1096 diag_count = jnp.sum(jnp.isfinite(diag))
1098 lower_sum = (total_sum - diag_sum) / 2.0
1099 lower_count = (total_count - diag_count) / 2.0
1100 return lower_sum / lower_count
1102 fourier_fingerprint, _, _ = cls.get_fourier_fingerprint(
1103 model,
1104 n_samples,
1105 random_key,
1106 method,
1107 scale,
1108 weight,
1109 trim_redundant=trim_redundant,
1110 **kwargs,
1111 )
1113 return cls.calculate_fcc(fourier_fingerprint)
1115 @classmethod
1116 def get_fourier_fingerprint(
1117 cls,
1118 model: Model,
1119 n_samples: int,
1120 random_key: Optional[random.PRNGKey] = None,
1121 method: Optional[str] = "pearson",
1122 scale: Optional[bool] = False,
1123 weight: Optional[bool] = False,
1124 trim_redundant: Optional[bool] = True,
1125 nan_to_one: Optional[bool] = False,
1126 **kwargs: Any,
1127 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
1128 """
1129 Shortcut method to get just the fourier fingerprint.
1130 This includes
1131 1. Calculating the coefficients (using `n_samples`)
1132 2. Correlating the result from 1) using `method`
1133 3. Weighting the correlation matrix (if `weight` is True)
1134 4. Remove redundancies (if `trim_redundant` is True)
1136 Args:
1137 model (Model): The QFM model
1138 n_samples (int): Number of samples to calculate average of coefficients
1139 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1140 initialization. If None, uses the model's internal random key.
1141 method (Optional[str], optional): Correlation method. Supported values are
1142 "pearson", "complex_pearson", "spearman", and "covariance".
1143 Defaults to "pearson".
1144 scale (Optional[bool], optional): Whether to scale the number of samples.
1145 Defaults to False.
1146 weight (Optional[bool], optional): Whether to weight the correlation matrix.
1147 Defaults to False.
1148 trim_redundant (Optional[bool], optional): Whether to remove redundant
1149 correlations. Defaults to True.
1150 nan_to_one (Optional[bool], optional): Whether to set nan to 1.
1151 Defaults to False.
1152 **kwargs: Additional keyword arguments for the model function.
1154 Returns:
1155 Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: The fourier
1156 fingerprint, the corresponding frequency indices and the
1157 corresponding coefficients. If `trim_redundant` is True the
1158 frequencies are returned as a `(row_freqs, col_freqs)` tuple that
1159 labels the two (redundancy-trimmed) matrix axes and the
1160 coefficients as a matching `(row_coeffs, col_coeffs)` tuple whose
1161 rows align with those frequencies; otherwise the full frequency
1162 vector and full coefficient array are returned.
1163 """
1164 _, coeffs, freqs = cls._calculate_coefficients(
1165 model, n_samples, random_key, scale, **kwargs
1166 )
1168 # Memory-efficient fast path
1169 if trim_redundant and not weight:
1170 pos_idx = cls._calculate_mask(freqs)
1171 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1173 # Flatten all frequency axes; the last axis is the sample
1174 # axis. `_calculate_mask` returns flat indices in C order,
1175 # matching this reshape.
1176 coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
1177 coeffs_sub = coeffs_flat[pos_idx]
1179 fourier_fingerprint = cls._correlate(coeffs_sub.transpose(), method=method)
1181 if nan_to_one:
1182 fourier_fingerprint = jnp.where(
1183 jnp.isnan(fourier_fingerprint), 1.0, fourier_fingerprint
1184 )
1186 M = fourier_fingerprint.shape[0]
1187 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1188 fourier_fingerprint = jnp.where(
1189 lower_tri_mask, fourier_fingerprint, jnp.nan
1190 )
1192 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1193 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1194 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1196 return (
1197 fourier_fingerprint,
1198 (pos_freqs[row_mask], pos_freqs[col_mask]),
1199 (coeffs_sub[row_mask], coeffs_sub[col_mask]),
1200 )
1202 fourier_fingerprint = cls._correlate(coeffs.transpose(), method=method)
1204 if nan_to_one:
1205 # set nan to 1
1206 fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0
1208 # perform weighting if requested
1209 fourier_fingerprint = (
1210 cls._weighting_mean(fourier_fingerprint, coeffs)
1211 if weight
1212 else fourier_fingerprint
1213 )
1215 if trim_redundant:
1216 pos_idx = cls._calculate_mask(freqs)
1217 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1218 coeffs_sub = coeffs.reshape(-1, coeffs.shape[-1])[pos_idx]
1220 # restrict to the positive-frequency sub-block (M x M with
1221 # M = number of non-negative flat-frequencies) instead of
1222 # building a full N x N mask. This avoids the O(N^2) float
1223 fourier_fingerprint = fourier_fingerprint[pos_idx][:, pos_idx]
1225 # keep only the strict lower triangle; the rest -> nan
1226 M = fourier_fingerprint.shape[0]
1227 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1228 fourier_fingerprint = jnp.where(
1229 lower_tri_mask, fourier_fingerprint, jnp.nan
1230 )
1232 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1233 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1235 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1237 return (
1238 fourier_fingerprint,
1239 (pos_freqs[row_mask], pos_freqs[col_mask]),
1240 (coeffs_sub[row_mask], coeffs_sub[col_mask]),
1241 )
1243 return fourier_fingerprint, freqs, coeffs
1245 @classmethod
1246 def calculate_fcc(
1247 cls,
1248 fourier_fingerprint: jnp.ndarray,
1249 ) -> float:
1250 """
1251 Method to calculate the FCC based on an existing correlation matrix.
1252 Calculate absolute and then the average over this matrix.
1253 The Fingerprint can be obtained via `get_fourier_fingerprint`
1255 Args:
1256 fourier_fingerprint (jnp.ndarray): Correlation matrix of coefficients
1257 Returns:
1258 float: The FCC
1259 """
1260 # apply the mask on the fingerprint
1261 return jnp.nanmean(jnp.abs(fourier_fingerprint))
1263 @classmethod
1264 def _calculate_mask(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1265 """
1266 Determine the flat indices of the Fourier correlation matrix
1267 that lie on a non-negative-frequency row/column. Together with
1268 the strict-lower-triangle condition (handled by the caller),
1269 these indices select the entries of the correlation matrix
1270 that survive the redundancy filter applied in
1271 `get_fourier_fingerprint`:
1273 - rows/columns whose flat frequency component is negative are
1274 discarded (they are the complex-conjugate redundancies of
1275 their positive counterparts);
1276 - of the remaining positive-frequency sub-block, only the
1277 strict lower triangle is kept (the upper triangle, including
1278 the diagonal, contains either duplicates from symmetry or
1279 self-correlations).
1281 Args:
1282 freqs (jnp.ndarray): Array of frequencies. Either a 1-D
1283 vector (single input feature) or a 2-D array of shape
1284 ``(n_input_feat, K)`` whose rows are the per-axis
1285 frequency vectors.
1287 Returns:
1288 jnp.ndarray: 1-D int array of flat indices selecting the
1289 non-negative-frequency rows/cols of the fingerprint.
1290 """
1291 freqs_arr = jnp.asarray(freqs)
1293 if freqs_arr.ndim == 1:
1294 pos_flat = freqs_arr >= 0
1295 else:
1296 # N-D case: build the per-axis non-negativity masks and
1297 # combine them via broadcasting (no float `jnp.outer`!),
1298 # then flatten to match the row-major flattening used by
1299 # the upstream coefficient/correlation pipeline.
1300 axes_pos = [freqs_arr[i] >= 0 for i in range(freqs_arr.shape[0])]
1301 expanded = []
1302 n_axes = len(axes_pos)
1303 for i, p in enumerate(axes_pos):
1304 shape = [1] * n_axes
1305 shape[i] = p.shape[0]
1306 expanded.append(p.reshape(shape))
1307 nd_pos = reduce(jnp.logical_and, expanded)
1308 pos_flat = nd_pos.flatten()
1310 return jnp.where(pos_flat)[0]
1312 @classmethod
1313 def _flat_frequencies(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1314 """
1315 Build the per-coefficient flat frequency labels in the same
1316 C-order used to flatten the coefficient/correlation pipeline, so
1317 they can be indexed by the flat indices from `_calculate_mask`.
1319 Args:
1320 freqs (jnp.ndarray): Either a 1-D vector (single input feature)
1321 or a ``(n_input_feat, K)`` stack / list of per-axis frequency
1322 vectors (multi-dim input).
1324 Returns:
1325 jnp.ndarray: 1-D frequency vector (single input feature) or a
1326 ``(N, n_input_feat)`` array of per-coefficient frequency
1327 tuples (multi-dim input).
1328 """
1329 fa = jnp.asarray(freqs)
1330 if fa.ndim == 1:
1331 return fa
1332 # Multi-dim: per-axis vectors -> flat grid of frequency tuples in the
1333 # same C-order used by `_calculate_mask` and the coefficient reshape.
1334 grids = jnp.meshgrid(*[fa[i] for i in range(fa.shape[0])], indexing="ij")
1335 return jnp.stack(grids, axis=-1).reshape(-1, fa.shape[0])
1337 @classmethod
1338 def _calculate_coefficients(
1339 cls,
1340 model: Model,
1341 n_samples: int,
1342 random_key: Optional[random.PRNGKey] = None,
1343 scale: bool = False,
1344 **kwargs: Any,
1345 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
1346 """
1347 Calculates the Fourier coefficients of a given model
1348 using `n_samples`.
1349 Optionally, `noise_params` can be passed to perform noisy simulation.
1351 Args:
1352 model (Model): The QFM model
1353 n_samples (int): Number of samples to calculate average of coefficients
1354 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1355 initialization. If None, uses the model's internal random key.
1356 scale (bool, optional): Whether to scale the number of samples.
1357 Defaults to False.
1358 **kwargs: Additional keyword arguments for the model function.
1360 Returns:
1361 Tuple[jnp.ndarray, jnp.ndarray]: Parameters and Coefficients of size NxK
1362 """
1363 if n_samples > 0:
1364 if scale:
1365 total_samples = int(
1366 jnp.power(2, model.n_qubits) * n_samples * model.n_input_feat
1367 )
1368 log.info(f"Using {total_samples} samples.")
1369 else:
1370 total_samples = n_samples
1371 model.initialize_params(random_key, repeat=total_samples)
1372 else:
1373 total_samples = 1
1375 coeffs, freqs = Coefficients.get_spectrum(
1376 model, shift=True, trim=True, **kwargs
1377 )
1379 return model.params, coeffs, freqs
1381 @classmethod
1382 def _correlate(cls, mat: jnp.ndarray, method: str = "pearson") -> jnp.ndarray:
1383 """
1384 Correlates two arrays using `method`.
1385 Currently, `pearson`, `complex_pearson`, `spearman`, and `covariance`
1386 are supported.
1388 Args:
1389 mat (jnp.ndarray): Array of shape (N, K)
1390 method (str, optional): Correlation method. Defaults to "pearson".
1392 Raises:
1393 ValueError: If the method is not supported.
1395 Returns:
1396 jnp.ndarray: Correlation matrix of `a` and `b`.
1397 """
1398 assert len(mat.shape) >= 2, "Input matrix must have at least 2 dimensions"
1400 # Note that for the general n-D case, we have to flatten along
1401 # the first axis (last one is batch).
1402 # Note that the order here is important so we can easily filter out
1403 # negative coefficients later.
1404 # Consider the following example: [[1,2,3],[4,5,6],[7,8,9]]
1405 # we want to get [1, 4, 7, 2, 5, 8, 3, 6, 9]
1406 # such that after correlation, all positive indexed coefficients
1407 # will be in the bottom right quadrant
1408 if method == "pearson":
1409 result = cls._pearson(mat.reshape(mat.shape[0], -1))
1410 # result = cls._pearson(mat.reshape(mat.shape[-1], -1, order="F"))
1411 elif method == "complex_pearson":
1412 result = cls._complex_pearson(mat.reshape(mat.shape[0], -1))
1413 elif method == "spearman":
1414 result = cls._spearman(mat.reshape(mat.shape[0], -1))
1415 # result = cls._spearman(mat.reshape(mat.shape[-1], -1, order="F"))
1416 elif method == "covariance":
1417 result = cls._covariance(mat.reshape(mat.shape[0], -1))
1418 else:
1419 raise ValueError(
1420 f"Unknown correlation method: {method}. Must be 'pearson', \
1421 'complex_pearson', 'spearman' or 'covariance'."
1422 )
1424 return result
1426 @classmethod
1427 def _covariance(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1428 """
1429 Compute the Hermitian sample covariance between columns of `mat`,
1430 permitting missing values (NaN or ±Inf).
1432 For each pair (i, j) the covariance is computed over the rows that are
1433 finite in both columns, as
1434 sum(conj(x_i - mean_i) * (x_j - mean_j)) / (nobs - 1),
1435 so it computes `X.conj().T @ X`.
1436 Real input collapses to the ordinary real sample covariance; complex
1437 input yields a complex matrix whose magnitude and angle carry the
1438 covariance strength and relative phase.
1441 Args:
1442 mat : array_like, shape (N, K)
1443 Input data.
1444 minp : int, optional
1445 Minimum number of paired observations required to form a
1446 covariance. If the number of valid pairs for (i, j) is < minp,
1447 the result is NaN.
1449 Returns:
1450 cov : ndarray, shape (K, K)
1451 Sample covariance matrix.
1452 """
1453 mat = jnp.asarray(mat)
1454 real_dtype = jnp.asarray(mat.real).dtype
1456 mask = jnp.isfinite(mat)
1457 fmask = mask.astype(real_dtype)
1458 safe = jnp.where(mask, mat, 0.0)
1460 nobs = fmask.T @ fmask
1461 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1463 sum_x = safe.T @ fmask
1464 sum_y = fmask.T @ safe
1466 masked = safe * fmask
1467 sum_conj_xy = jnp.conj(masked).T @ masked
1469 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1471 denom = jnp.where(nobs > 1, nobs - 1, jnp.nan)
1472 result = sxy / denom
1474 result = jnp.where(nobs < minp, jnp.nan, result)
1476 return result
1478 @classmethod
1479 def _complex_pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1480 """
1481 Compute the complex Pearson correlation between columns of `mat`,
1482 permitting missing values (NaN or ±Inf).
1484 This uses the Hermitian normalized covariance
1485 sum(conj(x_i - mean_i) * (x_j - mean_j)) /
1486 sqrt(sum(abs(x_i - mean_i)**2) * sum(abs(x_j - mean_j)**2)).
1487 Consequently, if column j is exp(1j * phi) times column i, then
1488 abs(corr[i, j]) is 1 and angle(corr[i, j]) is phi.
1490 Args:
1491 mat : array_like, shape (N, K)
1492 Input data.
1493 minp : int, optional
1494 Minimum number of paired observations required to form a correlation.
1495 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1497 Returns:
1498 corr : ndarray, shape (K, K)
1499 Complex Pearson correlation matrix.
1500 """
1501 mat = jnp.asarray(mat)
1502 real_dtype = jnp.asarray(mat.real).dtype
1504 mask = jnp.isfinite(mat)
1505 fmask = mask.astype(real_dtype)
1506 safe = jnp.where(mask, mat, 0.0)
1508 nobs = fmask.T @ fmask
1509 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1511 sum_x = safe.T @ fmask
1512 sum_y = fmask.T @ safe
1514 masked = safe * fmask
1515 sum_conj_xy = jnp.conj(masked).T @ masked
1517 safe_abs_sq = jnp.abs(safe) ** 2
1518 sum_abs_x2 = safe_abs_sq.T @ fmask
1519 sum_abs_y2 = fmask.T @ safe_abs_sq
1521 ssx = sum_abs_x2 - jnp.abs(sum_x) ** 2 / nobs_safe
1522 ssy = sum_abs_y2 - jnp.abs(sum_y) ** 2 / nobs_safe
1523 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1525 denom = jnp.sqrt(ssx * ssy)
1526 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1527 magnitude = jnp.abs(result)
1528 result = jnp.where(magnitude > 1.0, result / magnitude, result)
1530 result = jnp.where(nobs < minp, jnp.nan, result)
1532 return result
1534 @classmethod
1535 def _pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1536 """
1537 Compute Pearson correlation between columns of `mat`,
1538 permitting missing values (NaN or ±Inf).
1540 The Pearson correlation is the normalized covariance,
1541 corr[i, j] = cov[i, j] / sqrt(cov[i, i] * cov[j, j]),
1542 so it is obtained by normalizing `_covariance` by the per-column
1543 standard deviations.
1545 If the input is complex, real and imaginary parts are stacked along
1546 the sample axis so that both components contribute to the correlation
1547 without discarding information.
1549 Args:
1550 mat : array_like, shape (N, K)
1551 Input data.
1552 minp : int, optional
1553 Minimum number of paired observations required to form a correlation.
1554 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1556 Returns:
1557 corr : ndarray, shape (K, K)
1558 Pearson correlation matrix.
1559 """
1560 # Preserve complex information by splitting into real / imag samples.
1561 # After stacking the data is real, so the Hermitian `_covariance`
1562 # reduces to the ordinary real sample covariance.
1563 if jnp.iscomplexobj(mat):
1564 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1566 cov = cls._covariance(mat, minp=minp)
1568 # corr[i, j] = cov[i, j] / (std_i * std_j) with std_i = sqrt(cov[i, i])
1569 std = jnp.sqrt(jnp.diagonal(cov))
1570 denom = std[:, None] * std[None, :]
1571 result = jnp.where(denom > 0, cov / denom, jnp.nan)
1573 # clip numerical drift to [-1, 1]
1574 result = jnp.clip(jnp.real(result), -1.0, 1.0)
1576 return result
1578 @classmethod
1579 def _spearman(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1580 """
1581 Based on Pandas correlation method as implemented here:
1582 https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx
1584 Compute Spearman correlation between columns of `mat`,
1585 permitting missing values (NaN or ±Inf).
1587 If the input is complex, real and imaginary parts are stacked along
1588 the sample axis so that both components contribute to the correlation
1589 without discarding information.
1591 Args:
1592 mat : array_like, shape (N, K)
1593 Input data.
1594 minp : int, optional
1595 Minimum number of paired observations required to form a correlation.
1596 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1598 Returns:
1599 corr : ndarray, shape (K, K)
1600 Spearman correlation matrix.
1601 """
1602 # Preserve complex information by splitting into real / imag samples
1603 if jnp.iscomplexobj(mat):
1604 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1606 mat = jnp.asarray(mat)
1607 N, K = mat.shape
1609 # trivial all-NaN answer if too few rows
1610 if N < minp:
1611 return jnp.full((K, K), jnp.nan)
1613 # mask of finite entries
1614 mask = jnp.isfinite(mat) # shape (N, K), dtype=bool
1616 # precompute ranks column-wise ignoring NaNs
1617 ranks = np.full((N, K), np.nan)
1618 for j in range(K):
1619 valid = mask[:, j]
1620 if valid.any():
1621 ranks[valid, j] = rankdata(mat[valid, j], method="average")
1623 ranks = jnp.asarray(ranks)
1625 # Vectorised Pearson on the ranks
1626 # Replace NaN ranks with 0; use mask to track validity.
1627 rank_mask = jnp.isfinite(ranks)
1628 safe_ranks = jnp.where(rank_mask, ranks, 0.0)
1630 # Pairwise valid-observation counts (K, K)
1631 fmask = rank_mask.astype(ranks.dtype)
1632 nobs = fmask.T @ fmask
1634 # Pairwise sums over mutually-valid rows
1635 sum_x = safe_ranks.T @ fmask # (K, K)
1636 sum_y = fmask.T @ safe_ranks # (K, K)
1638 # Pairwise products
1639 masked_ranks = safe_ranks * fmask # same as safe_ranks
1640 sum_xy = masked_ranks.T @ masked_ranks # (K, K)
1642 safe_sq = safe_ranks**2
1643 sum_x2 = safe_sq.T @ fmask # (K, K)
1644 sum_y2 = fmask.T @ safe_sq # (K, K)
1646 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1647 ssx = sum_x2 - sum_x**2 / nobs_safe
1648 ssy = sum_y2 - sum_y**2 / nobs_safe
1649 sxy = sum_xy - (sum_x * sum_y) / nobs_safe
1651 denom = jnp.sqrt(ssx * ssy)
1652 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1653 result = jnp.clip(result, -1.0, 1.0)
1655 # Enforce minp
1656 result = jnp.where(nobs < minp, jnp.nan, result)
1658 return result
1660 @classmethod
1661 def _weighting_linear(cls, fourier_fingerprint: jnp.ndarray) -> jnp.ndarray:
1662 """
1663 Performs weighting on the given correlation matrix.
1664 Here, low-frequent coefficients are weighted more heavily.
1666 Args:
1667 fourier_fingerprint (jnp.ndarray): Correlation matrix
1668 """
1669 assert (
1670 fourier_fingerprint.shape[0] % 2 != 0
1671 and fourier_fingerprint.shape[1] % 2 != 0
1672 ), (
1673 "Correlation matrix must have odd dimensions. \
1674 Hint: use `trim` argument when calling `get_spectrum`."
1675 )
1676 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1677 "Correlation matrix must be square."
1678 )
1680 # The weight matrix produced by the previous quadrant-mirror
1681 # construction has a closed form: it is a "tent" sum along the
1682 # two axes. Concretely, with N = fourier_fingerprint.shape[0]
1683 # (odd) and center = N // 2,
1684 # W[i, j] = u[i] + u[j]
1685 # where u[k] = (center - |k - center|) / (2 * center)
1686 # is a triangular weighting peaking at the centre (the zero
1687 # frequency) and decaying linearly to 0 at the spectrum edges.
1688 N = fourier_fingerprint.shape[0]
1689 center = N // 2
1690 k = jnp.arange(N)
1691 u = (center - jnp.abs(k - center)) / (2 * center)
1693 return fourier_fingerprint * (u[:, None] + u[None, :])
1695 @classmethod
1696 def _weighting_mean(
1697 cls, fourier_fingerprint: jnp.ndarray, coeffs: jnp.ndarray
1698 ) -> jnp.ndarray:
1699 """
1700 Performs weighting on the given correlation matrix.
1701 Here, we use the product of the mean of the coefficients as weights.
1702 This suppresses correlations where the mean of the coefficients is near zero.
1704 Args:
1705 fourier_fingerprint (jnp.ndarray): Correlation matrix
1706 coeffs (jnp.ndarray): Fourier coefficients
1707 """
1708 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1709 "Correlation matrix must be square."
1710 )
1711 assert len(coeffs.shape) >= 2, (
1712 "Coefficient matrix must contain coefficient axes and a sample axis."
1713 )
1715 coefficient_means = jnp.abs(jnp.mean(coeffs, axis=-1))
1716 coefficient_means = coefficient_means.T.reshape(-1)
1718 assert fourier_fingerprint.shape[0] == coefficient_means.shape[0], (
1719 "Correlation matrix size must match the number of Fourier coefficients."
1720 )
1722 # Apply the rank-1 weight w[i] * w[j] via broadcasting instead
1723 # of materialising an explicit `jnp.outer` N x N intermediate.
1724 return (
1725 fourier_fingerprint
1726 * coefficient_means[:, None]
1727 * coefficient_means[None, :]
1728 )
1731class Datasets:
1732 @classmethod
1733 def generate_fourier_series(
1734 cls,
1735 random_key: random.PRNGKey,
1736 model: Model,
1737 coefficients_min: float = 0.0,
1738 coefficients_max: float = 1.0,
1739 zero_centered: bool = False,
1740 ) -> jnp.ndarray:
1741 """
1742 Generates the Fourier series representation of a function.
1743 It uses the `model.frequencies` property to retrieve the frequency
1744 information. This ensures that the resulting Fourier series is
1745 compatible with the model.
1747 This function is capable of generating $D$-dimensional Fourier series
1748 (again defined by `model.n_input_feat`).
1749 The highest frequency $N$ is retrieved per dimension.
1751 Samples of the Fourier coefficients are drawn from a uniform circle.
1753 Args:
1754 random_key (random.PRNGKey): Random number key for JAX.
1755 model (Model): The quantum circuit model.
1756 coefficients_min (float, optional): Minimum value for the coefficients.
1757 Defaults to 0.0.
1758 coefficients_max (float, optional): Maximum value for the coefficients.
1759 Defaults to 1.0.
1760 zero_centered (bool, optional): Whether to zero-center the coefficients.
1761 Defaults to False.
1763 Returns:
1764 jnp.ndarray: Input domain samples with shape ((N,)*D, D)
1765 jnp.ndarray: Fourier series values with shape ((N,)*D)
1766 jnp.ndarray: Fourier coefficients with shape ((N,)*D)
1768 """
1769 # TODO: the following code can be considered to
1770 # capturing a truly random spectrum.
1771 # add some constraints on the spectrum, i.e. not fully
1773 # Note: one key observation for understanding the following code is,
1774 # that instead of wrapping your head around symmetries in multi-
1775 # dimensional coefficient matrices, one can simply look at the flattened
1776 # version of such a matrix and reshape later. It just works out.
1777 domain_samples_per_input_dim = cls.construct_domain_samples(model)
1779 frequencies = cls.construct_frequencies(model)
1781 coefficients = cls.construct_coefficients(
1782 random_key, model, coefficients_min, coefficients_max, zero_centered
1783 )
1785 values = cls.calculate_values(
1786 domain_samples_per_input_dim, frequencies, coefficients
1787 )
1789 # return all the information we have
1790 return [
1791 domain_samples_per_input_dim.reshape(*model.degree, -1),
1792 values.reshape(model.degree),
1793 coefficients.reshape(model.degree),
1794 ]
1796 @classmethod
1797 def construct_domain_samples(
1798 cls, model: Model, mts: int = 1, mfs: int = 1
1799 ) -> jnp.ndarray:
1800 """
1801 Builds the input-domain sample grid for the model spectrum.
1803 Going from $[0, 2 \\pi \\, \\mathrm{mts}]$ with the resolution required
1804 for the highest frequency, permuted with the input dimensionality to get
1805 an n-d grid of domain samples (a "coordinate system").
1807 The grid follows the same convention as
1808 `Coefficients._fourier_transform`, so a dataset built here lands on the
1809 bins that `Coefficients.get_spectrum` analyses with the same `mts` and
1810 `mfs`. A target component at $k + j/r$ has period $2 \\pi r$, hence
1811 `mts` should be at least $r$ to cover a full period.
1813 Args:
1814 model (Model): The quantum circuit model.
1815 mts (int, optional): Domain oversampling, i.e. the number of
1816 periods covered. Defaults to 1.
1817 mfs (int, optional): Frequency oversampling, i.e. the sample
1818 density per period. Defaults to 1.
1820 Returns:
1821 jnp.ndarray: Domain samples with shape
1822 (mts $\\cdot$ mfs $\\cdot$ $\\prod$ degree, n_input_feat).
1823 """
1824 return jnp.stack(
1825 jnp.meshgrid(
1826 *[
1827 jnp.arange(0, 2 * mts * jnp.pi, 2 * jnp.pi / (mfs * d))
1828 for d in model.degree
1829 ]
1830 )
1831 ).T.reshape(-1, model.n_input_feat)
1833 @classmethod
1834 def construct_frequencies(
1835 cls,
1836 model: Model,
1837 random_key: Optional[random.PRNGKey] = None,
1838 offgrid_mode: str = "none",
1839 offgrid_prob: float = 0.0,
1840 offgrid_resolution: int = 2,
1841 ) -> jnp.ndarray:
1842 """
1843 Builds the frequency-index grid for the model spectrum.
1845 This has the same shape as the domain samples returned by
1846 `construct_domain_samples`.
1848 By default the grid is the model's own comb, so the dataset is exactly
1849 representable. The off-grid modes move a controllable fraction of the
1850 components off that comb.
1851 Offsets are always multiples of $1/r$ for the given resolution $r$.
1853 Args:
1854 model (Model): The quantum circuit model.
1855 random_key (Optional[random.PRNGKey]): Random number key for JAX.
1856 Required unless `offgrid_mode` is "none".
1857 offgrid_mode (str, optional): How to displace components off the
1858 model comb. "none" keeps the model comb. "index" perturbs each
1859 frequency independently, which spans arbitrary combs that are
1860 in general not exactly reachable. "generator" perturbs the
1861 per-gate generator frequencies and rebuilds the comb as their
1862 Minkowski sum, which stays exactly reachable by an encoding
1863 pulse configuration. Defaults to "none".
1864 offgrid_prob (float, optional): Probability that a single component
1865 ("index") or generator ("generator") is displaced. Defaults to
1866 0.0, which reproduces the model comb in every mode. Note that
1867 this is the fraction of components that end up off the comb
1868 only in "index" mode: a sum of displaced generators can land
1869 back on an integer, so "generator" mode displaces noticeably
1870 fewer components than asked for and saturates well below one.
1871 offgrid_resolution (int, optional): Denominator $r$ of the offset
1872 grid, i.e. offsets are drawn from $\\{\\pm j/r\\}$ with
1873 $j = 1 \\dots r-1$. Defaults to 2, giving half-integer offsets.
1875 Returns:
1876 jnp.ndarray: Frequency indices with shape
1877 ($\\prod$ degree, n_input_feat).
1878 """
1879 if offgrid_mode == "none":
1880 frequencies = model.frequencies
1881 else:
1882 if random_key is None:
1883 raise ValueError(f"offgrid_mode={offgrid_mode!r} requires a random_key")
1884 if offgrid_resolution < 2:
1885 raise ValueError(
1886 f"offgrid_resolution must be at least 2, "
1887 f"got {offgrid_resolution}. There is no non-integer offset "
1888 "on a grid of resolution 1."
1889 )
1890 if offgrid_mode == "index":
1891 displace = cls._displace_indices
1892 elif offgrid_mode == "generator":
1893 displace = cls._displace_generators
1894 else:
1895 raise ValueError(
1896 f"Unknown offgrid_mode: {offgrid_mode!r}. Use one of "
1897 "'none', 'index', 'generator'."
1898 )
1900 frequencies = []
1901 for i in range(model.n_input_feat):
1902 random_key, sub_key = random.split(random_key)
1903 frequencies.append(
1904 displace(model, i, sub_key, offgrid_prob, offgrid_resolution)
1905 )
1907 return jnp.stack(jnp.meshgrid(*frequencies)).T.reshape(-1, model.n_input_feat)
1909 @classmethod
1910 def _offsets(
1911 cls,
1912 random_key: random.PRNGKey,
1913 shape: Tuple[int, ...],
1914 prob: float,
1915 resolution: int,
1916 ) -> jnp.ndarray:
1917 """
1918 Draws signed offsets on the $1/r$ grid, zero where not displaced.
1920 Args:
1921 random_key (random.PRNGKey): Random number key for JAX.
1922 shape (Tuple[int, ...]): Shape of the offset array.
1923 prob (float): Probability that an entry is displaced.
1924 resolution (int): Denominator $r$ of the offset grid.
1926 Returns:
1927 jnp.ndarray: Offsets drawn from $\\{0\\} \\cup \\{\\pm j/r\\}$ with
1928 $j = 1 \\dots r-1$.
1929 """
1930 move_key, magnitude_key, sign_key = random.split(random_key, 3)
1931 magnitude = random.randint(magnitude_key, shape, 1, resolution) / resolution
1932 return (
1933 random.bernoulli(move_key, prob, shape)
1934 * random.rademacher(sign_key, shape)
1935 * magnitude
1936 )
1938 @classmethod
1939 def _displace_indices(
1940 cls,
1941 model: Model,
1942 feature: int,
1943 random_key: random.PRNGKey,
1944 prob: float,
1945 resolution: int,
1946 ) -> jnp.ndarray:
1947 """
1948 Displaces individual frequencies of one input feature off the comb.
1950 Each positive frequency is displaced independently, the result is
1951 re-sorted and mirrored so that the comb stays antisymmetric. This is
1952 what `construct_coefficients` relies on to enforce conjugate symmetry,
1953 and in turn what keeps the series real-valued. The comb never leaves
1954 the model's frequency range: an offset that would push a component past
1955 the highest frequency has its sign flipped rather than being clipped,
1956 which would put the component back on the comb.
1958 Args:
1959 model (Model): The quantum circuit model.
1960 feature (int): Index of the input feature.
1961 random_key (random.PRNGKey): Random number key for JAX.
1962 prob (float): Probability that a component is displaced.
1963 resolution (int): Denominator $r$ of the offset grid.
1965 Returns:
1966 jnp.ndarray: Displaced comb, same size as the model comb.
1967 """
1968 nominal = jnp.asarray(model.frequencies[feature])
1969 positive = nominal[nominal > 0]
1970 limit = positive[-1]
1972 offsets = cls._offsets(random_key, positive.shape, prob, resolution)
1973 # the smallest positive frequency is 1 and offsets are below 1, so only
1974 # the upper end of the range can be overshot
1975 offsets = jnp.where(positive + offsets > limit, -offsets, offsets)
1977 # ponytail: two components can collide (1 + 0.5 and 2 - 0.5), in which
1978 # case their coefficients simply add. Deduplicating would change the
1979 # number of components, which is the one thing the study holds fixed.
1980 positive = jnp.sort(positive + offsets)
1982 return jnp.concatenate([-jnp.flip(positive), jnp.zeros(1), positive])
1984 @classmethod
1985 def _displace_generators(
1986 cls,
1987 model: Model,
1988 feature: int,
1989 random_key: random.PRNGKey,
1990 prob: float,
1991 resolution: int,
1992 ) -> jnp.ndarray:
1993 """
1994 Displaces the generator frequencies of one input feature off the comb.
1996 Mirrors `Encoding.get_spectrum`, but scales each encoding gate's
1997 generator by a displaced $\\eta$ before taking the Minkowski sum, which
1998 is exactly what an encoding pulse scaler does to the gate it drives.
1999 The reachable comb then grows past the model degree in both count and
2000 range, so each model frequency claims the closest reachable one that is
2001 still inside the model's range. Staying in range matters: a component
2002 beyond the highest model frequency would be unreachable.
2003 Two model frequencies may end up claiming the same
2004 reachable one, as in `_displace_indices`.
2006 Args:
2007 model (Model): The quantum circuit model.
2008 feature (int): Index of the input feature.
2009 random_key (random.PRNGKey): Random number key for JAX.
2010 prob (float): Probability that a generator is displaced.
2011 resolution (int): Denominator $r$ of the offset grid.
2013 Returns:
2014 jnp.ndarray: Displaced comb, same size as the model comb.
2015 """
2016 # the offset draw and in-range flip live in _generator_etas, so the
2017 # scalers exposed by generator_etas cannot desync from this comb
2018 eta = cls._generator_etas(model, feature, random_key, prob, resolution)
2020 base = {"hamming": 1, "binary": 2, "ternary": 3}[model._enc._strategy]
2021 nominal = np.asarray(model.frequencies[feature])
2022 limit = nominal.max()
2023 mask = np.asarray(model.data_reupload[..., feature], dtype=bool)
2024 scale = base ** np.arange(mask.shape[1])
2026 # Minkowski sum over the displaced per-gate generators. Rounded before
2027 # deduplication, which is exact for a power-of-two resolution.
2028 reach = {0.0}
2029 for layer, qubit in zip(*np.nonzero(mask)):
2030 generator = scale[qubit] * eta[layer, qubit]
2031 reach = {
2032 round(a + s * generator, 9) for a in reach for s in (-1.0, 0.0, 1.0)
2033 }
2034 reachable = sorted(v for v in reach if 0 < v <= limit)
2036 # claim the closest reachable frequency, so the displaced comb tracks
2037 # the original one
2038 positive = jnp.sort(
2039 jnp.asarray(
2040 [
2041 min(reachable, key=lambda v: abs(v - frequency))
2042 for frequency in nominal[nominal > 0]
2043 ],
2044 dtype=float,
2045 )
2046 )
2048 return jnp.concatenate([-jnp.flip(positive), jnp.zeros(1), positive])
2050 @classmethod
2051 def _generator_etas(
2052 cls,
2053 model: Model,
2054 feature: int,
2055 random_key: random.PRNGKey,
2056 prob: float,
2057 resolution: int,
2058 ) -> np.ndarray:
2059 """
2060 The per-gate scalers $\\eta = 1 + \\text{offset}$ applied to the
2061 encoding generators of one input feature in `offgrid_mode='generator'`.
2063 This is the offset draw and in-range flip shared with
2064 `_displace_generators`; the returned array has the shape of the
2065 data-reupload mask ($n_\\text{layers}, n_\\text{qubits}$) and holds
2066 exactly the encoding pulse amplitude scalers that make the generator
2067 comb reachable.
2069 Args:
2070 model (Model): The quantum circuit model.
2071 feature (int): Index of the input feature.
2072 random_key (random.PRNGKey): Random number key for JAX.
2073 prob (float): Probability that a generator is displaced.
2074 resolution (int): Denominator $r$ of the offset grid.
2076 Returns:
2077 np.ndarray: Amplitude scalers $\\eta$, shape
2078 ($n_\\text{layers}, n_\\text{qubits}$).
2079 """
2080 base = {"hamming": 1, "binary": 2, "ternary": 3}.get(model._enc._strategy)
2081 if base is None:
2082 raise ValueError(
2083 f"offgrid_mode='generator' does not support the "
2084 f"{model._enc._strategy!r} encoding strategy, which has no "
2085 "per-gate pulse parametrization to displace."
2086 )
2088 nominal = np.asarray(model.frequencies[feature])
2089 limit = nominal.max()
2090 mask = np.asarray(model.data_reupload[..., feature], dtype=bool)
2091 scale = base ** np.arange(mask.shape[1])
2092 offsets = np.asarray(cls._offsets(random_key, mask.shape, prob, resolution))
2093 offsets = np.where(scale * (1.0 + offsets) > limit, -offsets, offsets)
2094 return 1.0 + offsets
2096 @classmethod
2097 def generator_etas(
2098 cls,
2099 model: Model,
2100 random_key: random.PRNGKey,
2101 offgrid_prob: float,
2102 offgrid_resolution: int,
2103 ) -> List[np.ndarray]:
2104 """
2105 The encoding pulse amplitude scalers `construct_frequencies` applies in
2106 `offgrid_mode='generator'`, one ($n_\\text{layers}, n_\\text{qubits}$)
2107 array per input feature.
2109 Call with the same `random_key` passed to `construct_frequencies` to
2110 recover the encoding pulse configuration that makes the off-grid target
2111 reachable, e.g. to oracle-initialize or score trained scalers against
2112 it. The per-feature key split mirrors `construct_frequencies`.
2114 Args:
2115 model (Model): The quantum circuit model.
2116 random_key (random.PRNGKey): The key passed to
2117 `construct_frequencies`.
2118 offgrid_prob (float): Probability that a generator is displaced.
2119 offgrid_resolution (int): Denominator $r$ of the offset grid.
2121 Returns:
2122 List[np.ndarray]: Amplitude scalers $\\eta$ per input feature.
2123 """
2124 etas = []
2125 for i in range(model.n_input_feat):
2126 random_key, sub_key = random.split(random_key)
2127 etas.append(
2128 cls._generator_etas(model, i, sub_key, offgrid_prob, offgrid_resolution)
2129 )
2130 return etas
2132 @classmethod
2133 def construct_coefficients(
2134 cls,
2135 random_key: random.PRNGKey,
2136 model: Model,
2137 coefficients_min: float = 0.0,
2138 coefficients_max: float = 1.0,
2139 zero_centered: bool = False,
2140 ) -> jnp.ndarray:
2141 """
2142 Samples the conjugate-symmetric Fourier coefficient vector.
2144 Coefficients are drawn from a uniform circle (see `uniform_circle`).
2145 The offset coefficient (first entry) is either zeroed or made real,
2146 then the spectrum is mirrored to enforce conjugate symmetry.
2148 Args:
2149 random_key (random.PRNGKey): Random number key for JAX.
2150 model (Model): The quantum circuit model.
2151 coefficients_min (float, optional): Minimum value for the
2152 coefficients. Defaults to 0.0.
2153 coefficients_max (float, optional): Maximum value for the
2154 coefficients. Defaults to 1.0.
2155 zero_centered (bool, optional): Whether to zero-center the
2156 coefficients. Defaults to False.
2158 Returns:
2159 jnp.ndarray: Conjugate-symmetric coefficient vector of size
2160 $\\prod$ degree.
2161 """
2162 coefficients = cls.uniform_circle(
2163 random_key,
2164 low=coefficients_min,
2165 high=coefficients_max,
2166 size=math.prod(model.degree) // 2 + 1,
2167 )
2169 # zero center (first coeff = 0)
2170 # we can assume the first coeff is the offset, because we're dealing
2171 # with a non-symmetric spectrum here
2172 if zero_centered:
2173 coefficients = coefficients.at[0].set(0.0)
2174 else:
2175 coefficients = coefficients.at[0].set(coefficients[0].real)
2177 # ensure symmetry (here, non_negative_ is removed!),
2178 # giving us the full coefficients vector
2179 return jnp.concat(
2180 [
2181 jnp.flip(coefficients[..., 1:]).conjugate(),
2182 coefficients,
2183 ],
2184 axis=-1,
2185 )
2187 @classmethod
2188 def calculate_values(
2189 cls,
2190 domain_samples: jnp.ndarray,
2191 frequencies: jnp.ndarray,
2192 coefficients: jnp.ndarray,
2193 ) -> jnp.ndarray:
2194 """
2195 Evaluates the real-valued Fourier series on the domain grid.
2197 Vectorized version of
2198 $f(x) = \\sum_{n=0}^{N-1} c_n e^{i \\omega_n x}$ that takes the input
2199 dimension into account, normalized by the number of coefficients.
2201 Args:
2202 domain_samples (jnp.ndarray): Domain samples with shape
2203 (n_points, n_input_feat).
2204 frequencies (jnp.ndarray): Frequency indices with shape
2205 (n_freqs, n_input_feat).
2206 coefficients (jnp.ndarray): Fourier coefficients with shape
2207 (n_freqs,).
2209 Returns:
2210 jnp.ndarray: Real-valued Fourier series samples with shape
2211 (n_points,).
2212 """
2213 return jnp.real(
2214 (jnp.exp(1j * (domain_samples @ frequencies.T)) * coefficients).sum(axis=1)
2215 / coefficients.size
2216 )
2218 @classmethod
2219 def uniform_circle(
2220 cls,
2221 random_key: random.PRNGKey,
2222 size: Union[jnp.ndarray, List, int],
2223 low=0.0,
2224 high=1.0,
2225 ):
2226 """
2227 Random number generator for complex numbers sampled inside the unit circle
2229 Args:
2230 random_key (random.PRNGKey): Random number key for JAX.
2231 size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
2232 the first dimension will be the number of dimensions.
2233 low (float, optional): Minimum Radius. Defaults to 0.0.
2234 high (float, optional): Maximum Radius. Defaults to 1.0.
2236 Returns
2237 jnp.ndarray: Array of complex numbers with shape of `size`
2238 """
2240 if isinstance(size, int):
2241 size = jnp.array([size])
2243 random_key, random_key1 = random.split(random_key)
2244 return jnp.sqrt(
2245 random.uniform(random_key, size, minval=low, maxval=high)
2246 ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))