Coverage for qml_essentials / coefficients.py: 97%
671 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-08-18 14:58 +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]:
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]: The fourier fingerprint and the
1148 corresponding frequency indices. If `trim_redundant` is True the
1149 frequencies are returned as a `(row_freqs, col_freqs)` tuple that
1150 labels the two (redundancy-trimmed) matrix axes; otherwise the
1151 full frequency vector is returned.
1152 """
1153 _, coeffs, freqs = cls._calculate_coefficients(
1154 model, n_samples, random_key, scale, **kwargs
1155 )
1157 # Memory-efficient fast path
1158 if trim_redundant and not weight:
1159 pos_idx = cls._calculate_mask(freqs)
1160 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1162 # Flatten all frequency axes; the last axis is the sample
1163 # axis. `_calculate_mask` returns flat indices in C order,
1164 # matching this reshape.
1165 coeffs_flat = coeffs.reshape(-1, coeffs.shape[-1])
1166 coeffs_sub = coeffs_flat[pos_idx]
1168 fourier_fingerprint = cls._correlate(coeffs_sub.transpose(), method=method)
1170 if nan_to_one:
1171 fourier_fingerprint = jnp.where(
1172 jnp.isnan(fourier_fingerprint), 1.0, fourier_fingerprint
1173 )
1175 M = fourier_fingerprint.shape[0]
1176 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1177 fourier_fingerprint = jnp.where(
1178 lower_tri_mask, fourier_fingerprint, jnp.nan
1179 )
1181 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1182 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1183 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1185 return fourier_fingerprint, (pos_freqs[row_mask], pos_freqs[col_mask])
1187 fourier_fingerprint = cls._correlate(coeffs.transpose(), method=method)
1189 if nan_to_one:
1190 # set nan to 1
1191 fourier_fingerprint[jnp.isnan(fourier_fingerprint)] = 1.0
1193 # perform weighting if requested
1194 fourier_fingerprint = (
1195 cls._weighting_mean(fourier_fingerprint, coeffs)
1196 if weight
1197 else fourier_fingerprint
1198 )
1200 if trim_redundant:
1201 pos_idx = cls._calculate_mask(freqs)
1202 pos_freqs = cls._flat_frequencies(freqs)[pos_idx]
1204 # restrict to the positive-frequency sub-block (M x M with
1205 # M = number of non-negative flat-frequencies) instead of
1206 # building a full N x N mask. This avoids the O(N^2) float
1207 fourier_fingerprint = fourier_fingerprint[pos_idx][:, pos_idx]
1209 # keep only the strict lower triangle; the rest -> nan
1210 M = fourier_fingerprint.shape[0]
1211 lower_tri_mask = jnp.tri(M, k=-1, dtype=bool)
1212 fourier_fingerprint = jnp.where(
1213 lower_tri_mask, fourier_fingerprint, jnp.nan
1214 )
1216 row_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=1)
1217 col_mask = jnp.any(jnp.isfinite(fourier_fingerprint), axis=0)
1219 fourier_fingerprint = fourier_fingerprint[row_mask][:, col_mask]
1221 return fourier_fingerprint, (pos_freqs[row_mask], pos_freqs[col_mask])
1223 return fourier_fingerprint, freqs
1225 @classmethod
1226 def calculate_fcc(
1227 cls,
1228 fourier_fingerprint: jnp.ndarray,
1229 ) -> float:
1230 """
1231 Method to calculate the FCC based on an existing correlation matrix.
1232 Calculate absolute and then the average over this matrix.
1233 The Fingerprint can be obtained via `get_fourier_fingerprint`
1235 Args:
1236 fourier_fingerprint (jnp.ndarray): Correlation matrix of coefficients
1237 Returns:
1238 float: The FCC
1239 """
1240 # apply the mask on the fingerprint
1241 return jnp.nanmean(jnp.abs(fourier_fingerprint))
1243 @classmethod
1244 def _calculate_mask(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1245 """
1246 Determine the flat indices of the Fourier correlation matrix
1247 that lie on a non-negative-frequency row/column. Together with
1248 the strict-lower-triangle condition (handled by the caller),
1249 these indices select the entries of the correlation matrix
1250 that survive the redundancy filter applied in
1251 `get_fourier_fingerprint`:
1253 - rows/columns whose flat frequency component is negative are
1254 discarded (they are the complex-conjugate redundancies of
1255 their positive counterparts);
1256 - of the remaining positive-frequency sub-block, only the
1257 strict lower triangle is kept (the upper triangle, including
1258 the diagonal, contains either duplicates from symmetry or
1259 self-correlations).
1261 Args:
1262 freqs (jnp.ndarray): Array of frequencies. Either a 1-D
1263 vector (single input feature) or a 2-D array of shape
1264 ``(n_input_feat, K)`` whose rows are the per-axis
1265 frequency vectors.
1267 Returns:
1268 jnp.ndarray: 1-D int array of flat indices selecting the
1269 non-negative-frequency rows/cols of the fingerprint.
1270 """
1271 freqs_arr = jnp.asarray(freqs)
1273 if freqs_arr.ndim == 1:
1274 pos_flat = freqs_arr >= 0
1275 else:
1276 # N-D case: build the per-axis non-negativity masks and
1277 # combine them via broadcasting (no float `jnp.outer`!),
1278 # then flatten to match the row-major flattening used by
1279 # the upstream coefficient/correlation pipeline.
1280 axes_pos = [freqs_arr[i] >= 0 for i in range(freqs_arr.shape[0])]
1281 expanded = []
1282 n_axes = len(axes_pos)
1283 for i, p in enumerate(axes_pos):
1284 shape = [1] * n_axes
1285 shape[i] = p.shape[0]
1286 expanded.append(p.reshape(shape))
1287 nd_pos = reduce(jnp.logical_and, expanded)
1288 pos_flat = nd_pos.flatten()
1290 return jnp.where(pos_flat)[0]
1292 @classmethod
1293 def _flat_frequencies(cls, freqs: jnp.ndarray) -> jnp.ndarray:
1294 """
1295 Build the per-coefficient flat frequency labels in the same
1296 C-order used to flatten the coefficient/correlation pipeline, so
1297 they can be indexed by the flat indices from `_calculate_mask`.
1299 Args:
1300 freqs (jnp.ndarray): Either a 1-D vector (single input feature)
1301 or a ``(n_input_feat, K)`` stack / list of per-axis frequency
1302 vectors (multi-dim input).
1304 Returns:
1305 jnp.ndarray: 1-D frequency vector (single input feature) or a
1306 ``(N, n_input_feat)`` array of per-coefficient frequency
1307 tuples (multi-dim input).
1308 """
1309 fa = jnp.asarray(freqs)
1310 if fa.ndim == 1:
1311 return fa
1312 # Multi-dim: per-axis vectors -> flat grid of frequency tuples in the
1313 # same C-order used by `_calculate_mask` and the coefficient reshape.
1314 grids = jnp.meshgrid(*[fa[i] for i in range(fa.shape[0])], indexing="ij")
1315 return jnp.stack(grids, axis=-1).reshape(-1, fa.shape[0])
1317 @classmethod
1318 def _calculate_coefficients(
1319 cls,
1320 model: Model,
1321 n_samples: int,
1322 random_key: Optional[random.PRNGKey] = None,
1323 scale: bool = False,
1324 **kwargs: Any,
1325 ) -> Tuple[jnp.ndarray, jnp.ndarray]:
1326 """
1327 Calculates the Fourier coefficients of a given model
1328 using `n_samples`.
1329 Optionally, `noise_params` can be passed to perform noisy simulation.
1331 Args:
1332 model (Model): The QFM model
1333 n_samples (int): Number of samples to calculate average of coefficients
1334 random_key (Optional[random.PRNGKey]): JAX random key for parameter
1335 initialization. If None, uses the model's internal random key.
1336 scale (bool, optional): Whether to scale the number of samples.
1337 Defaults to False.
1338 **kwargs: Additional keyword arguments for the model function.
1340 Returns:
1341 Tuple[jnp.ndarray, jnp.ndarray]: Parameters and Coefficients of size NxK
1342 """
1343 if n_samples > 0:
1344 if scale:
1345 total_samples = int(
1346 jnp.power(2, model.n_qubits) * n_samples * model.n_input_feat
1347 )
1348 log.info(f"Using {total_samples} samples.")
1349 else:
1350 total_samples = n_samples
1351 model.initialize_params(random_key, repeat=total_samples)
1352 else:
1353 total_samples = 1
1355 coeffs, freqs = Coefficients.get_spectrum(
1356 model, shift=True, trim=True, **kwargs
1357 )
1359 return model.params, coeffs, freqs
1361 @classmethod
1362 def _correlate(cls, mat: jnp.ndarray, method: str = "pearson") -> jnp.ndarray:
1363 """
1364 Correlates two arrays using `method`.
1365 Currently, `pearson`, `complex_pearson`, `spearman`, and `covariance`
1366 are supported.
1368 Args:
1369 mat (jnp.ndarray): Array of shape (N, K)
1370 method (str, optional): Correlation method. Defaults to "pearson".
1372 Raises:
1373 ValueError: If the method is not supported.
1375 Returns:
1376 jnp.ndarray: Correlation matrix of `a` and `b`.
1377 """
1378 assert len(mat.shape) >= 2, "Input matrix must have at least 2 dimensions"
1380 # Note that for the general n-D case, we have to flatten along
1381 # the first axis (last one is batch).
1382 # Note that the order here is important so we can easily filter out
1383 # negative coefficients later.
1384 # Consider the following example: [[1,2,3],[4,5,6],[7,8,9]]
1385 # we want to get [1, 4, 7, 2, 5, 8, 3, 6, 9]
1386 # such that after correlation, all positive indexed coefficients
1387 # will be in the bottom right quadrant
1388 if method == "pearson":
1389 result = cls._pearson(mat.reshape(mat.shape[0], -1))
1390 # result = cls._pearson(mat.reshape(mat.shape[-1], -1, order="F"))
1391 elif method == "complex_pearson":
1392 result = cls._complex_pearson(mat.reshape(mat.shape[0], -1))
1393 elif method == "spearman":
1394 result = cls._spearman(mat.reshape(mat.shape[0], -1))
1395 # result = cls._spearman(mat.reshape(mat.shape[-1], -1, order="F"))
1396 elif method == "covariance":
1397 result = cls._covariance(mat.reshape(mat.shape[0], -1))
1398 else:
1399 raise ValueError(
1400 f"Unknown correlation method: {method}. Must be 'pearson', \
1401 'complex_pearson', 'spearman' or 'covariance'."
1402 )
1404 return result
1406 @classmethod
1407 def _covariance(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1408 """
1409 Compute the Hermitian sample covariance between columns of `mat`,
1410 permitting missing values (NaN or ±Inf).
1412 For each pair (i, j) the covariance is computed over the rows that are
1413 finite in both columns, as
1414 sum(conj(x_i - mean_i) * (x_j - mean_j)) / (nobs - 1),
1415 so it computes `X.conj().T @ X`.
1416 Real input collapses to the ordinary real sample covariance; complex
1417 input yields a complex matrix whose magnitude and angle carry the
1418 covariance strength and relative phase.
1421 Args:
1422 mat : array_like, shape (N, K)
1423 Input data.
1424 minp : int, optional
1425 Minimum number of paired observations required to form a
1426 covariance. If the number of valid pairs for (i, j) is < minp,
1427 the result is NaN.
1429 Returns:
1430 cov : ndarray, shape (K, K)
1431 Sample covariance matrix.
1432 """
1433 mat = jnp.asarray(mat)
1434 real_dtype = jnp.asarray(mat.real).dtype
1436 mask = jnp.isfinite(mat)
1437 fmask = mask.astype(real_dtype)
1438 safe = jnp.where(mask, mat, 0.0)
1440 nobs = fmask.T @ fmask
1441 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1443 sum_x = safe.T @ fmask
1444 sum_y = fmask.T @ safe
1446 masked = safe * fmask
1447 sum_conj_xy = jnp.conj(masked).T @ masked
1449 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1451 denom = jnp.where(nobs > 1, nobs - 1, jnp.nan)
1452 result = sxy / denom
1454 result = jnp.where(nobs < minp, jnp.nan, result)
1456 return result
1458 @classmethod
1459 def _complex_pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1460 """
1461 Compute the complex Pearson correlation between columns of `mat`,
1462 permitting missing values (NaN or ±Inf).
1464 This uses the Hermitian normalized covariance
1465 sum(conj(x_i - mean_i) * (x_j - mean_j)) /
1466 sqrt(sum(abs(x_i - mean_i)**2) * sum(abs(x_j - mean_j)**2)).
1467 Consequently, if column j is exp(1j * phi) times column i, then
1468 abs(corr[i, j]) is 1 and angle(corr[i, j]) is phi.
1470 Args:
1471 mat : array_like, shape (N, K)
1472 Input data.
1473 minp : int, optional
1474 Minimum number of paired observations required to form a correlation.
1475 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1477 Returns:
1478 corr : ndarray, shape (K, K)
1479 Complex Pearson correlation matrix.
1480 """
1481 mat = jnp.asarray(mat)
1482 real_dtype = jnp.asarray(mat.real).dtype
1484 mask = jnp.isfinite(mat)
1485 fmask = mask.astype(real_dtype)
1486 safe = jnp.where(mask, mat, 0.0)
1488 nobs = fmask.T @ fmask
1489 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1491 sum_x = safe.T @ fmask
1492 sum_y = fmask.T @ safe
1494 masked = safe * fmask
1495 sum_conj_xy = jnp.conj(masked).T @ masked
1497 safe_abs_sq = jnp.abs(safe) ** 2
1498 sum_abs_x2 = safe_abs_sq.T @ fmask
1499 sum_abs_y2 = fmask.T @ safe_abs_sq
1501 ssx = sum_abs_x2 - jnp.abs(sum_x) ** 2 / nobs_safe
1502 ssy = sum_abs_y2 - jnp.abs(sum_y) ** 2 / nobs_safe
1503 sxy = sum_conj_xy - (jnp.conj(sum_x) * sum_y) / nobs_safe
1505 denom = jnp.sqrt(ssx * ssy)
1506 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1507 magnitude = jnp.abs(result)
1508 result = jnp.where(magnitude > 1.0, result / magnitude, result)
1510 result = jnp.where(nobs < minp, jnp.nan, result)
1512 return result
1514 @classmethod
1515 def _pearson(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1516 """
1517 Compute Pearson correlation between columns of `mat`,
1518 permitting missing values (NaN or ±Inf).
1520 The Pearson correlation is the normalized covariance,
1521 corr[i, j] = cov[i, j] / sqrt(cov[i, i] * cov[j, j]),
1522 so it is obtained by normalizing `_covariance` by the per-column
1523 standard deviations.
1525 If the input is complex, real and imaginary parts are stacked along
1526 the sample axis so that both components contribute to the correlation
1527 without discarding information.
1529 Args:
1530 mat : array_like, shape (N, K)
1531 Input data.
1532 minp : int, optional
1533 Minimum number of paired observations required to form a correlation.
1534 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1536 Returns:
1537 corr : ndarray, shape (K, K)
1538 Pearson correlation matrix.
1539 """
1540 # Preserve complex information by splitting into real / imag samples.
1541 # After stacking the data is real, so the Hermitian `_covariance`
1542 # reduces to the ordinary real sample covariance.
1543 if jnp.iscomplexobj(mat):
1544 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1546 cov = cls._covariance(mat, minp=minp)
1548 # corr[i, j] = cov[i, j] / (std_i * std_j) with std_i = sqrt(cov[i, i])
1549 std = jnp.sqrt(jnp.diagonal(cov))
1550 denom = std[:, None] * std[None, :]
1551 result = jnp.where(denom > 0, cov / denom, jnp.nan)
1553 # clip numerical drift to [-1, 1]
1554 result = jnp.clip(jnp.real(result), -1.0, 1.0)
1556 return result
1558 @classmethod
1559 def _spearman(cls, mat: jnp.ndarray, minp: Optional[int] = 1) -> jnp.ndarray:
1560 """
1561 Based on Pandas correlation method as implemented here:
1562 https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/algos.pyx
1564 Compute Spearman correlation between columns of `mat`,
1565 permitting missing values (NaN or ±Inf).
1567 If the input is complex, real and imaginary parts are stacked along
1568 the sample axis so that both components contribute to the correlation
1569 without discarding information.
1571 Args:
1572 mat : array_like, shape (N, K)
1573 Input data.
1574 minp : int, optional
1575 Minimum number of paired observations required to form a correlation.
1576 If the number of valid pairs for (i, j) is < minp, the result is NaN.
1578 Returns:
1579 corr : ndarray, shape (K, K)
1580 Spearman correlation matrix.
1581 """
1582 # Preserve complex information by splitting into real / imag samples
1583 if jnp.iscomplexobj(mat):
1584 mat = jnp.concatenate([mat.real, mat.imag], axis=0)
1586 mat = jnp.asarray(mat)
1587 N, K = mat.shape
1589 # trivial all-NaN answer if too few rows
1590 if N < minp:
1591 return jnp.full((K, K), jnp.nan)
1593 # mask of finite entries
1594 mask = jnp.isfinite(mat) # shape (N, K), dtype=bool
1596 # precompute ranks column-wise ignoring NaNs
1597 ranks = np.full((N, K), np.nan)
1598 for j in range(K):
1599 valid = mask[:, j]
1600 if valid.any():
1601 ranks[valid, j] = rankdata(mat[valid, j], method="average")
1603 ranks = jnp.asarray(ranks)
1605 # Vectorised Pearson on the ranks
1606 # Replace NaN ranks with 0; use mask to track validity.
1607 rank_mask = jnp.isfinite(ranks)
1608 safe_ranks = jnp.where(rank_mask, ranks, 0.0)
1610 # Pairwise valid-observation counts (K, K)
1611 fmask = rank_mask.astype(ranks.dtype)
1612 nobs = fmask.T @ fmask
1614 # Pairwise sums over mutually-valid rows
1615 sum_x = safe_ranks.T @ fmask # (K, K)
1616 sum_y = fmask.T @ safe_ranks # (K, K)
1618 # Pairwise products
1619 masked_ranks = safe_ranks * fmask # same as safe_ranks
1620 sum_xy = masked_ranks.T @ masked_ranks # (K, K)
1622 safe_sq = safe_ranks**2
1623 sum_x2 = safe_sq.T @ fmask # (K, K)
1624 sum_y2 = fmask.T @ safe_sq # (K, K)
1626 nobs_safe = jnp.where(nobs > 0, nobs, 1.0)
1627 ssx = sum_x2 - sum_x**2 / nobs_safe
1628 ssy = sum_y2 - sum_y**2 / nobs_safe
1629 sxy = sum_xy - (sum_x * sum_y) / nobs_safe
1631 denom = jnp.sqrt(ssx * ssy)
1632 result = jnp.where(denom > 0, sxy / denom, jnp.nan)
1633 result = jnp.clip(result, -1.0, 1.0)
1635 # Enforce minp
1636 result = jnp.where(nobs < minp, jnp.nan, result)
1638 return result
1640 @classmethod
1641 def _weighting_linear(cls, fourier_fingerprint: jnp.ndarray) -> jnp.ndarray:
1642 """
1643 Performs weighting on the given correlation matrix.
1644 Here, low-frequent coefficients are weighted more heavily.
1646 Args:
1647 fourier_fingerprint (jnp.ndarray): Correlation matrix
1648 """
1649 assert (
1650 fourier_fingerprint.shape[0] % 2 != 0
1651 and fourier_fingerprint.shape[1] % 2 != 0
1652 ), (
1653 "Correlation matrix must have odd dimensions. \
1654 Hint: use `trim` argument when calling `get_spectrum`."
1655 )
1656 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1657 "Correlation matrix must be square."
1658 )
1660 # The weight matrix produced by the previous quadrant-mirror
1661 # construction has a closed form: it is a "tent" sum along the
1662 # two axes. Concretely, with N = fourier_fingerprint.shape[0]
1663 # (odd) and center = N // 2,
1664 # W[i, j] = u[i] + u[j]
1665 # where u[k] = (center - |k - center|) / (2 * center)
1666 # is a triangular weighting peaking at the centre (the zero
1667 # frequency) and decaying linearly to 0 at the spectrum edges.
1668 N = fourier_fingerprint.shape[0]
1669 center = N // 2
1670 k = jnp.arange(N)
1671 u = (center - jnp.abs(k - center)) / (2 * center)
1673 return fourier_fingerprint * (u[:, None] + u[None, :])
1675 @classmethod
1676 def _weighting_mean(
1677 cls, fourier_fingerprint: jnp.ndarray, coeffs: jnp.ndarray
1678 ) -> jnp.ndarray:
1679 """
1680 Performs weighting on the given correlation matrix.
1681 Here, we use the product of the mean of the coefficients as weights.
1682 This suppresses correlations where the mean of the coefficients is near zero.
1684 Args:
1685 fourier_fingerprint (jnp.ndarray): Correlation matrix
1686 coeffs (jnp.ndarray): Fourier coefficients
1687 """
1688 assert fourier_fingerprint.shape[0] == fourier_fingerprint.shape[1], (
1689 "Correlation matrix must be square."
1690 )
1691 assert len(coeffs.shape) >= 2, (
1692 "Coefficient matrix must contain coefficient axes and a sample axis."
1693 )
1695 coefficient_means = jnp.abs(jnp.mean(coeffs, axis=-1))
1696 coefficient_means = coefficient_means.T.reshape(-1)
1698 assert fourier_fingerprint.shape[0] == coefficient_means.shape[0], (
1699 "Correlation matrix size must match the number of Fourier coefficients."
1700 )
1702 # Apply the rank-1 weight w[i] * w[j] via broadcasting instead
1703 # of materialising an explicit `jnp.outer` N x N intermediate.
1704 return (
1705 fourier_fingerprint
1706 * coefficient_means[:, None]
1707 * coefficient_means[None, :]
1708 )
1711class Datasets:
1712 @classmethod
1713 def generate_fourier_series(
1714 cls,
1715 random_key: random.PRNGKey,
1716 model: Model,
1717 coefficients_min: float = 0.0,
1718 coefficients_max: float = 1.0,
1719 zero_centered: bool = False,
1720 ) -> jnp.ndarray:
1721 """
1722 Generates the Fourier series representation of a function.
1723 It uses the `model.frequencies` property to retrieve the frequency
1724 information. This ensures that the resulting Fourier series is
1725 compatible with the model.
1727 This function is capable of generating $D$-dimensional Fourier series
1728 (again defined by `model.n_input_feat`).
1729 The highest frequency $N$ is retrieved per dimension.
1731 Samples of the Fourier coefficients are drawn from a uniform circle.
1733 Args:
1734 random_key (random.PRNGKey): Random number key for JAX.
1735 model (Model): The quantum circuit model.
1736 coefficients_min (float, optional): Minimum value for the coefficients.
1737 Defaults to 0.0.
1738 coefficients_max (float, optional): Maximum value for the coefficients.
1739 Defaults to 1.0.
1740 zero_centered (bool, optional): Whether to zero-center the coefficients.
1741 Defaults to False.
1743 Returns:
1744 jnp.ndarray: Input domain samples with shape ((N,)*D, D)
1745 jnp.ndarray: Fourier series values with shape ((N,)*D)
1746 jnp.ndarray: Fourier coefficients with shape ((N,)*D)
1748 """
1749 # TODO: the following code can be considered to
1750 # capturing a truly random spectrum.
1751 # add some constraints on the spectrum, i.e. not fully
1753 # Note: one key observation for understanding the following code is,
1754 # that instead of wrapping your head around symmetries in multi-
1755 # dimensional coefficient matrices, one can simply look at the flattened
1756 # version of such a matrix and reshape later. It just works out.
1758 # going from [0, 2pi] with the resolution required for highest frequency
1759 # permute with input dimensionality to get an n-d grid of domain samples
1760 # the output shape comes from the fact that want to create a "coordinate system"
1761 domain_samples_per_input_dim = jnp.stack(
1762 jnp.meshgrid(
1763 *[jnp.arange(0, 2 * jnp.pi, 2 * jnp.pi / d) for d in model.degree]
1764 )
1765 ).T.reshape(-1, model.n_input_feat)
1767 # generate the frequency indices for each dimension.
1768 # this will have the same shape as the domain samples
1769 frequencies = jnp.stack(jnp.meshgrid(*model.frequencies)).T.reshape(
1770 -1, model.n_input_feat
1771 )
1773 # using the frequency information, sample coefficients for each dimension
1774 # shape: (input_dims, n_freqs_per_input_dim // 2 + 1)
1776 coefficients = cls.uniform_circle(
1777 random_key,
1778 low=coefficients_min,
1779 high=coefficients_max,
1780 size=math.prod(model.degree) // 2 + 1,
1781 )
1783 # zero center (first coeff = 0)
1784 # we can assume the first coeff is the offset, because we're dealing
1785 # with a non-symmetric spectrum here
1786 if zero_centered:
1787 coefficients = coefficients.at[0].set(0.0)
1788 else:
1789 coefficients = coefficients.at[0].set(coefficients[0].real)
1791 # ensure symmetry (here, non_negative_ is removed!),
1792 # giving us the full coefficients vector
1793 coefficients = jnp.concat(
1794 [
1795 jnp.flip(coefficients[..., 1:]).conjugate(),
1796 coefficients,
1797 ],
1798 axis=-1,
1799 )
1801 # Vectorized version of $f(x) = \sum_{n=0}^{N-1} c_n * e^{i * \omega_n * x}$
1802 # it takes into account the input dimension, i.e. the output is a matrix
1803 # normalization uses the n_freqs component of the coefficients
1804 values = jnp.real(
1805 (
1806 jnp.exp(1j * (domain_samples_per_input_dim @ frequencies.T))
1807 * coefficients
1808 ).sum(axis=1)
1809 / coefficients.size
1810 )
1812 # return all the information we have
1813 return [
1814 domain_samples_per_input_dim.reshape(*model.degree, -1),
1815 values.reshape(model.degree),
1816 coefficients.reshape(model.degree),
1817 ]
1819 @classmethod
1820 def uniform_circle(
1821 cls,
1822 random_key: random.PRNGKey,
1823 size: Union[jnp.ndarray, List, int],
1824 low=0.0,
1825 high=1.0,
1826 ):
1827 """
1828 Random number generator for complex numbers sampled inside the unit circle
1830 Args:
1831 random_key (random.PRNGKey): Random number key for JAX.
1832 size (Union[jnp.ndarray, int]): Number of samples. If a 2D array is passed,
1833 the first dimension will be the number of dimensions.
1834 low (float, optional): Minimum Radius. Defaults to 0.0.
1835 high (float, optional): Maximum Radius. Defaults to 1.0.
1837 Returns
1838 jnp.ndarray: Array of complex numbers with shape of `size`
1839 """
1841 if isinstance(size, int):
1842 size = jnp.array([size])
1844 random_key, random_key1 = random.split(random_key)
1845 return jnp.sqrt(
1846 random.uniform(random_key, size, minval=low, maxval=high)
1847 ) * jnp.exp(2j * jnp.pi * random.uniform(random_key1, size))