Coverage for qml_essentials / model.py: 91%
529 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-08-18 14:58 +0000
1from typing import Any, Dict, Optional, Tuple, Callable, Union, List
3import warnings
4import jax.numpy as jnp
5import numpy as np
6from jax import random
8from qml_essentials import jaqsi as js
9from qml_essentials import operations as op
10from qml_essentials.tape import recording
11from qml_essentials.operations import KrausChannel
12from qml_essentials.ansaetze import Ansaetze, Circuit, Encoding
13from qml_essentials.gates import Gates, PulseInformation as pinfo
14from qml_essentials.utils import safe_random_split
16import logging
18log = logging.getLogger(__name__)
21class Model:
22 """
23 A quantum circuit model.
24 """
26 def __init__(
27 self,
28 n_qubits: int,
29 n_layers: int,
30 circuit_type: Union[str, Circuit] = "No_Ansatz",
31 data_reupload: Union[bool, List[List[bool]], List[List[List[bool]]]] = True,
32 state_preparation: Union[
33 str, Callable, List[Union[str, Callable]], None
34 ] = None,
35 encoding: Union[Encoding, str, Callable, List[Union[str, Callable]]] = Gates.RX,
36 trainable_frequencies: bool = False,
37 initialization: str = "random",
38 initialization_domain: List[float] = [0, 2 * jnp.pi],
39 output_qubit: Union[List[int], int] = -1,
40 shots: Optional[int] = None,
41 random_seed: int = 1000,
42 remove_zero_encoding: bool = True,
43 repeat_batch_axis: List[bool] = [True, True, True],
44 pulse_shape: str = "gaussian",
45 ) -> None:
46 """
47 Initialize the quantum circuit model.
48 Parameters will have the shape [impl_n_layers, parameters_per_layer]
49 where impl_n_layers is the number of layers provided and added by one
50 depending if data_reupload is True and parameters_per_layer is given by
51 the chosen ansatz.
53 The model is initialized with the following parameters as defaults:
54 - noise_params: None
55 - execution_type: "expval"
56 - shots: None
58 Args:
59 n_qubits (int): The number of qubits in the circuit.
60 n_layers (int): The number of layers in the circuit.
61 circuit_type (str, Circuit): The type of quantum circuit to use.
62 If None, defaults to "no_ansatz".
63 data_reupload (Union[bool, List[bool], List[List[bool]]], optional):
64 Whether to reupload data to the quantum device on each
65 layer and qubit. Detailed re-uploading instructions can be given
66 as a list/array of 0/False and 1/True with shape (n_qubits,
67 n_layers) to specify where to upload the data. Defaults to True
68 for applying data re-uploading to the full circuit.
69 encoding (Union[str, Callable, List[str], List[Callable]], optional):
70 The unitary to use for encoding the input data. Can be a string
71 (e.g. "RX") or a callable (e.g. op.RX). Defaults to op.RX.
72 If input is multidimensional it is assumed to be a list of
73 unitaries or a list of strings.
74 trainable_frequencies (bool, optional):
75 Sets trainable encoding parameters for trainable frequencies.
76 Defaults to False.
77 initialization (str, optional): The strategy to initialize the parameters.
78 Can be "random", "zeros", "zero-controlled", "pi", or "pi-controlled".
79 Defaults to "random".
80 output_qubit (List[int], int, optional): The index of the output
81 qubit (or qubits). When set to -1 all qubits are measured, or a
82 global measurement is conducted, depending on the execution
83 type.
84 shots (Optional[int], optional): The number of shots to use for
85 the quantum device. Defaults to None.
86 random_seed (int, optional): seed for the random number generator
87 in initialization is "random" and for random noise parameters.
88 Defaults to 1000.
89 remove_zero_encoding (bool, optional): whether to
90 remove the zero encoding from the circuit. Defaults to True.
91 repeat_batch_axis (List[bool], optional): Each boolean in the array
92 determines over which axes to parallelise computation. The axes
93 correspond to [inputs, params, pulse_params]. Defaults to
94 [True, True, True], meaning that batching is enabled over all
95 axes.
96 pulse_shape (str, optional): Pulse envelope shape for pulse-level
97 simulation. One of ``PulseEnvelope.available()``.
98 Defaults to ``"gaussian"``.
100 Returns:
101 None
102 """
103 # Initialize default parameters needed for circuit evaluation
104 self.n_qubits: int = n_qubits
105 self.output_qubit: Union[List[int], int] = output_qubit
106 self.n_layers: int = n_layers
107 self.noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None
108 self.shots = shots
109 self.remove_zero_encoding = remove_zero_encoding
110 self.trainable_frequencies: bool = trainable_frequencies
111 self.execution_type: str = "expval"
112 self.repeat_batch_axis: List[bool] = repeat_batch_axis
114 # --- Pulse envelope ---
115 pinfo.set_envelope(pulse_shape)
117 # --- State Preparation ---
118 try:
119 self._sp = Gates.parse_gates(state_preparation, Gates)
120 except ValueError as e:
121 raise ValueError(f"Error parsing encodings: {e}")
123 # prepare corresponding pulse parameters (always optimized pulses)
124 self.sp_pulse_params = []
125 for sp in self._sp:
126 sp_name = sp.__name__ if hasattr(sp, "__name__") else str(sp)
128 if pinfo.gate_by_name(sp_name) is not None:
129 self.sp_pulse_params.append(pinfo.gate_by_name(sp_name).params)
130 else:
131 # gate has no pulse parametrization
132 self.sp_pulse_params.append(None)
134 # --- Encoding ---
135 if isinstance(encoding, Encoding):
136 # user wants custom strategy? do it!
137 self._enc = encoding
138 else:
139 # use hammming encoding by default
140 self._enc = Encoding("hamming", encoding)
142 if self._enc.is_golomb:
143 self._enc._n_qubits = n_qubits
145 # Number of possible inputs
146 self.n_input_feat = len(self._enc)
147 log.debug(f"Number of input features: {self.n_input_feat}")
149 # Trainable frequencies, default initialization as in arXiv:2309.03279v2
150 self.enc_params = jnp.ones((self.n_layers, self.n_qubits, self.n_input_feat))
152 self._zero_inputs = False
154 # --- Data-Reuploading ---
156 # Keep as NumPy array (not JAX) so that ``if data_reupload[q, idx]``
157 # in _iec remains a concrete Python bool even under jax.jit tracing.
158 # note that setting this will also update self.degree and self.frequencies
159 # and in consequence also self.has_dru
160 self.data_reupload = data_reupload
162 # check for the highest degree among all input dimensions
163 if self.has_dru:
164 impl_n_layers: int = n_layers + 1 # we need L+1 according to Schuld et al.
165 else:
166 impl_n_layers = n_layers
167 log.info(f"Number of implicit layers: {impl_n_layers}.")
169 # --- Ansatz ---
170 # only weak check for str. We trust the user to provide sth useful
171 if isinstance(circuit_type, str):
172 self.pqc: Callable[[Optional[jnp.ndarray], int], int] = getattr(
173 Ansaetze, circuit_type or "No_Ansatz"
174 )()
175 else:
176 self.pqc = circuit_type()
177 log.info(f"Using Ansatz {circuit_type}.")
179 # calculate the shape of the parameter vector here, we will re-use this in init.
180 params_per_layer = self.pqc.n_params_per_layer(self.n_qubits)
181 self._params_shape: Tuple[int, int] = (impl_n_layers, params_per_layer)
182 log.info(f"Parameters per layer: {params_per_layer}")
184 pulse_params_per_layer = self.pqc.n_pulse_params_per_layer(self.n_qubits)
185 self._pulse_params_shape: Tuple[int, int] = (
186 impl_n_layers,
187 pulse_params_per_layer,
188 )
190 # intialize to None as we can't know this yet
191 self._batch_shape = None
193 # this will also be re-used in the init method,
194 # however, only if nothing is provided
195 self._inialization_strategy = initialization
196 self._initialization_domain = initialization_domain
198 # ..here! where we only require a JAX random key
199 self.random_key = self.initialize_params(random.key(random_seed))
201 # Initializing pulse params
202 self.pulse_params: jnp.ndarray = jnp.ones((1, *self._pulse_params_shape))
204 log.info(f"Initialized pulse parameters with shape {self.pulse_params.shape}.")
206 # Initialise the jaqsi Script that wraps _variational.
207 # No device selection needed - jaqsi auto-routes between statevector
208 # and density-matrix simulation based on whether noise channels are
209 # present on the tape.
210 self.script = js.Script(f=self._variational, n_qubits=self.n_qubits)
212 @property
213 def noise_params(self) -> Optional[Dict[str, Union[float, Dict[str, float]]]]:
214 """
215 Gets the noise parameters of the model.
217 Returns:
218 Optional[Dict[str, float]]: A dictionary of
219 noise parameters or None if not set.
220 """
221 return self._noise_params
223 @noise_params.setter
224 def noise_params(
225 self, kvs: Optional[Dict[str, Union[float, Dict[str, float]]]]
226 ) -> None:
227 """
228 Sets the noise parameters of the model.
230 Typically a "noise parameter" refers to the error probability.
231 ThermalRelaxation is a special case, and supports a dict as value with
232 structure:
233 "ThermalRelaxation":
234 {
235 "t1": 2000, # relative t1 time.
236 "t2": 1000, # relative t2 time
237 "t_factor" 1: # relative gate time factor
238 },
240 Args:
241 kvs (Optional[Dict[str, Union[float, Dict[str, float]]]]): A
242 dictionary of noise parameters. If all values are 0.0, the noise
243 parameters are set to None.
245 Returns:
246 None
247 """
248 # set to None if only zero values provided
249 if kvs is not None and all(v == 0.0 for v in kvs.values()):
250 kvs = None
252 # set default values
253 if kvs is not None:
254 defaults = {
255 "BitFlip": 0.0,
256 "PhaseFlip": 0.0,
257 "Depolarizing": 0.0,
258 "MultiQubitDepolarizing": 0.0,
259 "AmplitudeDamping": 0.0,
260 "PhaseDamping": 0.0,
261 "GateError": 0.0,
262 "ThermalRelaxation": None,
263 "StatePreparation": 0.0,
264 "Measurement": 0.0,
265 }
266 for key, default_val in defaults.items():
267 kvs.setdefault(key, default_val)
269 # check if there are any keys not supported
270 for key in kvs.keys():
271 if key not in defaults:
272 warnings.warn(
273 f"Noise type {key} is not supported by this package",
274 UserWarning,
275 )
277 # check valid params for thermal relaxation noise channel
278 tr_params = kvs["ThermalRelaxation"]
279 if isinstance(tr_params, dict):
280 tr_params.setdefault("t1", 0.0)
281 tr_params.setdefault("t2", 0.0)
282 tr_params.setdefault("t_factor", 0.0)
283 valid_tr_keys = {"t1", "t2", "t_factor"}
284 for k in tr_params.keys():
285 if k not in valid_tr_keys:
286 warnings.warn(
287 f"Thermal Relaxation parameter {k} is not supported "
288 f"by this package",
289 UserWarning,
290 )
291 if not all(tr_params.values()) or tr_params["t2"] > 2 * tr_params["t1"]:
292 warnings.warn(
293 "Received invalid values for Thermal Relaxation noise "
294 "parameter. Thermal relaxation is not applied!",
295 UserWarning,
296 )
297 kvs["ThermalRelaxation"] = 0.0
299 self._noise_params = kvs
301 @property
302 def output_qubit(self) -> List[int]:
303 """Get the output qubit indices for measurement."""
304 return self._output_qubit
306 @output_qubit.setter
307 def output_qubit(self, value: Union[int, List[int]]) -> None:
308 """
309 Set the output qubit(s) for measurement.
311 Args:
312 value: Qubit index or list of indices. Use -1 for all qubits.
313 """
314 if isinstance(value, list):
315 assert len(value) <= self.n_qubits, (
316 f"Size of output_qubit {len(value)} cannot be\
317 larger than number of qubits {self.n_qubits}."
318 )
319 elif isinstance(value, int):
320 if value == -1:
321 value = list(range(self.n_qubits))
322 else:
323 assert value < self.n_qubits, (
324 f"Output qubit {value} cannot be larger than {self.n_qubits}."
325 )
326 value = [value]
328 self._output_qubit = value
330 @property
331 def execution_type(self) -> str:
332 """
333 Gets the execution type of the model.
335 Returns:
336 str: The execution type, one of 'density', 'expval', or 'probs'.
337 """
338 return self._execution_type
340 @execution_type.setter
341 def execution_type(self, value: str) -> None:
342 if value == "density":
343 self._result_shape = (
344 2 ** len(self.output_qubit),
345 2 ** len(self.output_qubit),
346 )
347 elif value == "expval":
348 # check if all qubits are used
349 if len(self.output_qubit) == self.n_qubits:
350 self._result_shape = (len(self.output_qubit),)
351 # if not -> parity measurement with only 1D output per pair
352 # or n_local measurement
353 else:
354 self._result_shape = (len(self.output_qubit),)
355 elif value == "probs":
356 # in case this is a list of parities,
357 # each pair has 2^len(qubits) probabilities
358 n_parity = (
359 (2,) * len(self.output_qubit)
360 if isinstance(self.output_qubit, (Tuple, List))
361 else (2,)
362 )
363 self._result_shape = n_parity
364 elif value == "state":
365 self._result_shape = (2 ** len(self.output_qubit),)
366 else:
367 raise ValueError(f"Invalid execution type: {value}.")
369 if value == "state" and not self.all_qubit_measurement:
370 warnings.warn(
371 f"{value} measurement does ignore output_qubit, which is "
372 f"{self.output_qubit}.",
373 UserWarning,
374 )
376 if value == "probs" and self.shots is None:
377 warnings.warn(
378 "Setting execution_type to probs without specifying shots.",
379 UserWarning,
380 )
382 if value == "density" and self.shots is not None:
383 raise ValueError("Setting execution_type to density with shots not None.")
385 self._execution_type = value
387 @property
388 def shots(self) -> Optional[int]:
389 """
390 Gets the number of shots to use for the quantum device.
392 Returns:
393 Optional[int]: The number of shots.
394 """
395 return self._shots
397 @shots.setter
398 def shots(self, value: Optional[int]) -> None:
399 """
400 Sets the number of shots to use for the quantum device.
402 Args:
403 value (Optional[int]): The number of shots.
404 If an integer less than or equal to 0 is provided, it is set to None.
406 Returns:
407 None
408 """
409 if type(value) is int and value <= 0:
410 value = None
411 self._shots = value
413 @property
414 def params(self) -> jnp.ndarray:
415 """Get the variational parameters of the model."""
416 return self._params
418 @params.setter
419 def params(self, value: jnp.ndarray) -> None:
420 """Set the variational parameters, ensuring batch dimension exists."""
421 if len(value.shape) == 2:
422 value = value.reshape(1, *value.shape)
424 self._params = value
426 @property
427 def enc_params(self) -> jnp.ndarray:
428 """Get the encoding parameters used for input transformation."""
429 return self._enc_params
431 @enc_params.setter
432 def enc_params(self, value: jnp.ndarray) -> None:
433 """Set the encoding parameters."""
434 self._enc_params = value
436 @property
437 def pulse_params(self) -> jnp.ndarray:
438 """Get the pulse parameters for pulse-mode gate execution."""
439 return self._pulse_params
441 @pulse_params.setter
442 def pulse_params(self, value: jnp.ndarray) -> None:
443 """Set the pulse parameters."""
444 self._pulse_params = value
446 @property
447 def data_reupload(self) -> jnp.ndarray:
448 """Get the data reupload mask."""
449 return self._data_reupload
451 @data_reupload.setter
452 def data_reupload(self, value: jnp.ndarray) -> None:
453 """Set the data reupload mask.
455 Always converts to a concrete NumPy boolean array so that
456 ``if data_reupload[q, idx]`` in :meth:`_iec` remains a plain
457 Python ``bool`` even inside JAX-traced functions (jit / grad / vmap).
458 """
459 # Process data reuploading strategy and set degree
460 if not isinstance(value, bool):
461 if not isinstance(value, np.ndarray):
462 value = np.array(value)
464 if len(value.shape) == 2:
465 assert value.shape == (
466 self.n_layers,
467 self.n_qubits,
468 ), (
469 f"Data reuploading array has wrong shape. \
470 Expected {(self.n_layers, self.n_qubits)} or\
471 {(self.n_layers, self.n_qubits, self.n_input_feat)},\
472 got {value.shape}."
473 )
474 value = value.reshape(*value.shape, 1)
475 value = np.repeat(value, self.n_input_feat, axis=2)
477 assert value.shape == (
478 self.n_layers,
479 self.n_qubits,
480 self.n_input_feat,
481 ), (
482 f"Data reuploading array has wrong shape. \
483 Expected {(self.n_layers, self.n_qubits, self.n_input_feat)},\
484 got {value.shape}."
485 )
487 log.debug(f"Data reuploading array:\n{value}")
488 else:
489 if value:
490 value = np.ones((self.n_layers, self.n_qubits, self.n_input_feat))
491 log.debug("Full data reuploading.")
492 else:
493 value = np.zeros((self.n_layers, self.n_qubits, self.n_input_feat))
494 value[0][0] = 1
495 log.debug("No data reuploading.")
497 # convert to boolean values
498 self._data_reupload = np.asarray(value).astype(bool)
500 self.degree: Tuple = tuple(
501 self._enc.get_n_freqs(np.count_nonzero(self.data_reupload[..., i]))
502 for i in range(self.n_input_feat)
503 )
505 self.frequencies: Tuple = tuple(
506 self._enc.get_spectrum(np.count_nonzero(self.data_reupload[..., i]))
507 for i in range(self.n_input_feat)
508 )
510 # Cache has_dru as a plain Python bool so that it can be used in
511 # Python ``if`` statements even inside JAX-traced functions.
512 self._has_dru: bool = bool(max(int(np.max(f)) for f in self._frequencies) > 1)
514 @property
515 def degree(self) -> Tuple:
516 """Get the degree of the model."""
517 return self._degree
519 @degree.setter
520 def degree(self, value: Tuple):
521 self._degree = value
523 @property
524 def frequencies(self) -> Tuple:
525 """Get the frequencies of the model."""
526 return self._frequencies
528 @frequencies.setter
529 def frequencies(self, value: Tuple):
530 self._frequencies = value
532 def exact_spectrum(self, method: str = "tree") -> Tuple[np.ndarray, ...]:
533 """Compute the exact per-feature Fourier spectrum via the FourierTree.
535 Unlike :attr:`frequencies` -- a naive per-feature estimate derived purely
536 from the encoding, which can *overestimate* the spectrum (some
537 coefficients are constrained to zero for all parameters) -- this builds
538 the analytical Fourier tree (Nemkov et al.) and returns, for each input
539 feature, the integer frequencies whose Fourier coefficient is not
540 identically zero. The result is always a subset of :attr:`frequencies`.
542 The support is derived purely symbolically (no parameter sampling): see
543 :meth:`~qml_essentials.coefficients.FourierTree.get_exact_support`.
544 With ``method="tree"`` (default), frequencies whose contributions cancel
545 identically across tree paths (e.g. two consecutive encodings combining
546 into a single rotation) are excluded exactly; this enumerates the
547 explicit tree, which can be infeasible for deep entangling circuits.
548 With ``method="dp"``, a merged-state dynamic program derives the support
549 without enumerating paths, which scales to deep circuits at the cost of
550 not detecting identical cross-path cancellations.
552 Requires a Clifford + Pauli-rotation ansatz (see
553 :class:`~qml_essentials.pauli.PauliCircuit`); other gate sets raise
554 ``NotImplementedError`` during tree construction.
556 Args:
557 method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).
559 Returns:
560 Tuple[np.ndarray, ...]: One sorted integer frequency array per input
561 feature (same layout as :attr:`frequencies`).
562 """
563 from qml_essentials.coefficients import FourierTree # avoid circular imp.
565 tree = FourierTree(self)
567 # Position of each model feature within the tree's frequency vectors.
568 feature_pos = {feat: i for i, feat in enumerate(tree.features)}
570 # Union of the symbolic supports over all observables (roots).
571 support = set()
572 for freqs in tree.get_exact_support(method=method):
573 farr = np.asarray(freqs)
574 for k in range(farr.shape[0]):
575 key = (
576 (int(farr[k]),)
577 if farr.ndim == 1
578 else tuple(int(v) for v in farr[k])
579 )
580 support.add(key)
582 spectrum = []
583 for feat in range(self.n_input_feat):
584 if support and feat in feature_pos:
585 pos = feature_pos[feat]
586 vals = sorted({k[pos] for k in support})
587 else:
588 vals = [0]
589 spectrum.append(np.array(vals, dtype=int))
590 return tuple(spectrum)
592 @property
593 def has_dru(self) -> bool:
594 """Check if the model has data reupload."""
595 return self._has_dru
597 @property
598 def all_qubit_measurement(self) -> bool:
599 """Check if measurement is performed on all qubits."""
600 return self.output_qubit == list(range(self.n_qubits))
602 @property
603 def batch_shape(self) -> Tuple[int, ...]:
604 """
605 Get the batch shape (B_I, B_P, B_R).
606 If the model was not called before,
607 it returns (1, 1, 1).
609 Returns:
610 Tuple[int, ...]: Tuple of (input_batch, param_batch, pulse_batch).
611 Returns (1, 1, 1) if model has not been called yet.
612 """
613 if self._batch_shape is None:
614 log.debug("Model was not called yet. Returning (1,1,1) as batch shape.")
615 return (1, 1, 1)
616 return self._batch_shape
618 @property
619 def eff_batch_shape(self) -> Tuple[int, ...]:
620 """
621 Get the effective batch shape after applying repeat_batch_axis mask.
623 Returns:
624 Tuple[int, ...]: Effective batch dimensions, excluding zeros.
625 """
626 batch_shape = np.array(self.batch_shape) * self.repeat_batch_axis
627 batch_shape = batch_shape[batch_shape != 0]
628 return batch_shape
630 def initialize_params(
631 self,
632 random_key: Optional[random.PRNGKey] = None,
633 repeat: int = 1,
634 initialization: Optional[str] = None,
635 initialization_domain: Optional[List[float]] = None,
636 ) -> random.PRNGKey:
637 """
638 Initialize the variational parameters of the model.
640 Args:
641 random_key (Optional[random.PRNGKey]): JAX random key for initialization.
642 If None, uses the model's internal random key.
643 repeat (int): Number of parameter sets to create (batch dimension).
644 Defaults to 1.
645 initialization (Optional[str]): Strategy for parameter initialization.
646 Options: "random", "zeros", "pi", "zero-controlled", "pi-controlled".
647 If None, uses the strategy specified in the constructor.
648 initialization_domain (Optional[List[float]]): Domain [min, max] for
649 random initialization. If None, uses the domain from constructor.
651 Returns:
652 random.PRNGKey: Updated random key after initialization.
654 Raises:
655 Exception: If an invalid initialization method is specified.
656 """
657 # Initializing params
658 params_shape = (repeat, *self._params_shape)
660 # use existing strategy if not specified
661 initialization = initialization or self._inialization_strategy
662 initialization_domain = initialization_domain or self._initialization_domain
664 random_key, sub_key = safe_random_split(
665 random_key if random_key is not None else self.random_key
666 )
668 def set_control_params(params: jnp.ndarray, value: float) -> jnp.ndarray:
669 indices = self.pqc.get_control_indices(self.n_qubits)
670 if indices is None:
671 warnings.warn(
672 f"Specified {initialization} but circuit\
673 does not contain controlled rotation gates.\
674 Parameters are intialized randomly.",
675 UserWarning,
676 )
677 else:
678 np_params = np.array(params)
679 np_params[:, :, indices[0] : indices[1] : indices[2]] = (
680 np.ones_like(params[:, :, indices[0] : indices[1] : indices[2]])
681 * value
682 )
683 params = jnp.array(np_params)
684 return params
686 if initialization == "random":
687 self.params: jnp.ndarray = random.uniform(
688 sub_key,
689 params_shape,
690 minval=initialization_domain[0],
691 maxval=initialization_domain[1],
692 )
693 elif initialization == "zeros":
694 self.params: jnp.ndarray = jnp.zeros(params_shape)
695 elif initialization == "pi":
696 self.params: jnp.ndarray = jnp.ones(params_shape) * jnp.pi
697 elif initialization == "zero-controlled":
698 self.params: jnp.ndarray = random.uniform(
699 sub_key,
700 params_shape,
701 minval=initialization_domain[0],
702 maxval=initialization_domain[1],
703 )
704 self.params = set_control_params(self.params, 0)
705 elif initialization == "pi-controlled":
706 self.params: jnp.ndarray = random.uniform(
707 sub_key,
708 params_shape,
709 minval=initialization_domain[0],
710 maxval=initialization_domain[1],
711 )
712 self.params = set_control_params(self.params, jnp.pi)
713 else:
714 raise Exception("Invalid initialization method")
716 log.info(
717 f"Initialized parameters with shape {self.params.shape}\
718 using strategy {initialization}."
719 )
721 return random_key
723 def transform_input(
724 self, inputs: jnp.ndarray, enc_params: jnp.ndarray
725 ) -> jnp.ndarray:
726 """
727 Transform input data by scaling with encoding parameters.
729 Implements the input transformation as described in arXiv:2309.03279v2,
730 where inputs are linearly scaled by encoding parameters before being
731 used in the quantum circuit.
733 Args:
734 inputs (jnp.ndarray): Input data point of shape (n_input_feat,) or
735 (batch_size, n_input_feat).
736 enc_params (jnp.ndarray): Encoding weight scalar or vector used to
737 scale the input.
739 Returns:
740 jnp.ndarray: Transformed input, element-wise product of inputs
741 and enc_params.
742 """
743 return inputs * enc_params
745 def _iec(
746 self,
747 inputs: jnp.ndarray,
748 data_reupload: jnp.ndarray,
749 enc: Encoding,
750 enc_params: jnp.ndarray,
751 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
752 random_key: Optional[random.PRNGKey] = None,
753 ) -> None:
754 """
755 Apply Input Encoding Circuit (IEC) with angle encoding.
757 Encodes classical input data into the quantum circuit using rotation
758 gates (e.g., RX, RY, RZ). Supports data re-uploading at specified
759 positions in the circuit.
761 For Golomb encoding, a single multi-qubit diagonal unitary is applied
762 to all qubits simultaneously instead of per-qubit rotation gates.
764 Args:
765 inputs (jnp.ndarray): Input data of shape (n_input_feat,) or
766 (batch_size, n_input_feat).
767 data_reupload (jnp.ndarray): Boolean array of shape (n_qubits, n_input_feat)
768 indicating where to apply encoding gates.
769 enc (Encoding): Encoding strategy containing the encoding gate functions.
770 enc_params (jnp.ndarray): Encoding parameters of shape
771 (n_qubits, n_input_feat) used to scale inputs.
772 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
773 Noise parameters for gate-level noise simulation. Defaults to None.
774 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
775 noise. Defaults to None.
777 Returns:
778 None: Gates are applied in-place to the quantum circuit.
779 """
780 # check for zero, because due to input validation, input cannot be none
781 if self.remove_zero_encoding and self._zero_inputs and self.batch_shape[0] == 1:
782 return
784 # --- Golomb encoding: single multi-qubit gate on all qubits --------
785 if enc.is_golomb:
786 idx = 0 # Golomb encoding supports a single input feature
787 # Check if any qubit has re-uploading enabled for this layer
788 if data_reupload[:, idx].any():
789 random_key, sub_key = safe_random_split(random_key)
790 # Use the mean of enc_params across qubits as scalar scaling
791 # (Golomb acts on all qubits jointly)
792 mean_enc_param = jnp.mean(enc_params[:, idx])
793 all_wires = list(range(self.n_qubits))
794 enc[idx](
795 self.transform_input(inputs[..., idx], mean_enc_param),
796 wires=all_wires,
797 noise_params=noise_params,
798 random_key=sub_key,
799 )
800 return
802 # --- Standard per-qubit encoding -----------------------------------
803 for q in range(self.n_qubits):
804 # use the last dimension of the inputs (feature dimension)
805 for idx in range(inputs.shape[-1]):
806 if data_reupload[q, idx]:
807 # use elipsis to indiex only the last dimension
808 # as inputs are generally *not* qubit dependent
809 random_key, sub_key = safe_random_split(random_key)
810 enc[idx](
811 self.transform_input(inputs[..., idx], enc_params[q, idx]),
812 wires=q,
813 noise_params=noise_params,
814 random_key=sub_key,
815 )
817 def _variational(
818 self,
819 params: jnp.ndarray,
820 inputs: jnp.ndarray,
821 pulse_params: Optional[jnp.ndarray] = None,
822 random_key: Optional[random.PRNGKey] = None,
823 enc_params: Optional[jnp.ndarray] = None,
824 gate_mode: str = "unitary",
825 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
826 ) -> None:
827 """
828 Build the variational quantum circuit structure.
830 Constructs the circuit by applying state preparation, alternating
831 variational ansatz layers with input encoding layers, and optional
832 noise channels.
834 The first five parameters (after ``self``) - ``params``, ``inputs``,
835 ``pulse_params``, ``random_key``, ``enc_params`` - are the batchable
836 positional arguments.
837 The remaining keyword arguments are broadcast across the batch.
839 Args:
840 params (jnp.ndarray): Variational parameters of shape
841 (n_layers, n_params_per_layer).
842 inputs (jnp.ndarray): Input data of shape (n_input_feat,).
843 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers of shape
844 (n_layers, n_pulse_params_per_layer) for pulse-mode execution.
845 Defaults to None (uses model's pulse_params).
846 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
847 operations. Defaults to None.
848 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
849 (n_qubits, n_input_feat). Defaults to None (uses model's enc_params).
850 gate_mode (str): Gate execution mode, either "unitary" or "pulse".
851 Defaults to "unitary".
852 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
853 Noise parameters for simulation. Defaults to None.
855 Returns:
856 None: Gates are applied in-place to the quantum circuit.
858 Note:
859 Issues RuntimeWarning if called directly without providing parameters
860 that would normally be passed through the forward method.
861 """
862 # TODO: rework and double check params shape
863 if len(params.shape) > 2 and params.shape[0] == 1:
864 params = params[0]
866 if len(inputs.shape) > 1 and inputs.shape[0] == 1:
867 inputs = inputs[0]
869 if enc_params is None:
870 # TODO: Raise warning if trainable frequencies is True, or similar. I.e., no
871 # warning if user does not care for frequencies or enc_params
872 if self.trainable_frequencies:
873 warnings.warn(
874 "Explicit call to `_circuit` or `_variational` detected: "
875 "`enc_params` is None, using `self.enc_params` instead.",
876 RuntimeWarning,
877 )
878 enc_params = self.enc_params
880 if pulse_params is None:
881 if gate_mode == "pulse":
882 warnings.warn(
883 "Explicit call to `_circuit` or `_variational` detected: "
884 "`pulse_params` is None, using `self.pulse_params` instead.",
885 RuntimeWarning,
886 )
887 pulse_params = self.pulse_params
889 # Squeeze batch dimension for pulse_params (batch-first convention)
890 if len(pulse_params.shape) > 2 and pulse_params.shape[0] == 1:
891 pulse_params = pulse_params[0]
893 if noise_params is None:
894 if self.noise_params is not None:
895 warnings.warn(
896 "Explicit call to `_circuit` or `_variational` detected: "
897 "`noise_params` is None, using `self.noise_params` instead.",
898 RuntimeWarning,
899 )
900 noise_params = self.noise_params
902 if noise_params is not None:
903 if random_key is None:
904 warnings.warn(
905 "Explicit call to `_circuit` or `_variational` detected: "
906 "`random_key` is None, using `random.PRNGKey(0)` instead.",
907 RuntimeWarning,
908 )
909 random_key = self.random_key
910 self._apply_state_prep_noise(noise_params=noise_params)
912 # state preparation
913 for q in range(self.n_qubits):
914 for _sp, sp_pulse_params in zip(self._sp, self.sp_pulse_params):
915 random_key, sub_key = safe_random_split(random_key)
916 _sp(
917 wires=q,
918 pulse_params=sp_pulse_params,
919 noise_params=noise_params,
920 random_key=sub_key,
921 gate_mode=gate_mode,
922 )
924 # circuit building
925 for layer in range(0, self.n_layers):
926 random_key, sub_key = safe_random_split(random_key)
927 # ansatz layers
928 self.pqc(
929 params[layer],
930 self.n_qubits,
931 pulse_params=pulse_params[layer],
932 noise_params=noise_params,
933 random_key=sub_key,
934 gate_mode=gate_mode,
935 )
937 random_key, sub_key = safe_random_split(random_key)
938 # encoding layers
939 self._iec(
940 inputs,
941 data_reupload=self.data_reupload[layer],
942 enc=self._enc,
943 enc_params=enc_params[layer],
944 noise_params=noise_params,
945 random_key=sub_key,
946 )
948 # final ansatz layer
949 if self.has_dru: # same check as in init
950 random_key, sub_key = safe_random_split(random_key)
951 self.pqc(
952 params[self.n_layers],
953 self.n_qubits,
954 pulse_params=pulse_params[-1],
955 noise_params=noise_params,
956 random_key=sub_key,
957 gate_mode=gate_mode,
958 )
960 # channel noise
961 if noise_params is not None:
962 self._apply_general_noise(noise_params=noise_params)
964 def _build_obs(self) -> Tuple[str, List[op.Operation]]:
965 """Build the jaqsi measurement type and observable list.
967 Translates the model's ``execution_type`` and ``output_qubit``
968 settings into parameters suitable for
969 :meth:`~qml_essentials.jaqsi.Script.execute`.
971 Returns:
972 Tuple ``(meas_type, obs)`` where *meas_type* is one of
973 ``"expval"``, ``"probs"``, ``"density"``, ``"state"`` and *obs*
974 is a (possibly empty) list of :class:`Operation` observables.
975 """
976 if self.execution_type == "density":
977 return "density", []
979 if self.execution_type == "state":
980 return "state", []
982 if self.execution_type == "expval":
983 obs: List[op.Operation] = []
984 for qubit_spec in self.output_qubit:
985 if isinstance(qubit_spec, int):
986 obs.append(op.PauliZ(wires=qubit_spec))
987 else:
988 # parity: Z \\otimes Z \\otimes …
989 obs.append(js.build_parity_observable(list(qubit_spec)))
990 return "expval", obs
992 if self.execution_type == "probs":
993 # probs are computed on the full system; subsystem
994 # marginalisation is handled in _postprocess_res
995 return "probs", []
997 raise ValueError(f"Invalid execution_type: {self.execution_type}.")
999 def _apply_state_prep_noise(
1000 self, noise_params: Dict[str, Union[float, Dict[str, float]]]
1001 ) -> None:
1002 """
1003 Apply state preparation noise to all qubits.
1005 Simulates imperfect state preparation by applying BitFlip errors
1006 to each qubit with the specified probability.
1008 Args:
1009 noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
1010 containing noise parameters. Uses the "StatePreparation" key
1011 for the BitFlip probability.
1013 Returns:
1014 None: Noise channels are applied in-place to the circuit.
1015 """
1016 p = noise_params.get("StatePreparation", 0.0)
1017 if p > 0:
1018 for q in range(self.n_qubits):
1019 op.BitFlip(p, wires=q)
1021 def _apply_general_noise(
1022 self, noise_params: Dict[str, Union[float, Dict[str, float]]]
1023 ) -> None:
1024 """
1025 Apply general noise channels to all qubits.
1027 Applies various decoherence and error channels after the circuit
1028 execution, simulating environmental noise effects.
1030 Args:
1031 noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
1032 containing noise parameters with the following supported keys:
1033 - "AmplitudeDamping" (float): Probability for amplitude damping.
1034 - "PhaseDamping" (float): Probability for phase damping.
1035 - "Measurement" (float): Probability for measurement error (BitFlip).
1036 - "ThermalRelaxation" (Dict): Dictionary with keys "t1", "t2",
1037 "t_factor" for thermal relaxation simulation.
1039 Returns:
1040 None: Noise channels are applied in-place to the circuit.
1042 Note:
1043 Gate-level noise (e.g., GateError) is handled separately in the
1044 Gates.Noise module and applied at the individual gate level.
1045 """
1046 amp_damp = noise_params.get("AmplitudeDamping", 0.0)
1047 phase_damp = noise_params.get("PhaseDamping", 0.0)
1048 thermal_relax = noise_params.get("ThermalRelaxation", 0.0)
1049 meas = noise_params.get("Measurement", 0.0)
1050 for q in range(self.n_qubits):
1051 if amp_damp > 0:
1052 op.AmplitudeDamping(amp_damp, wires=q)
1053 if phase_damp > 0:
1054 op.PhaseDamping(phase_damp, wires=q)
1055 if meas > 0:
1056 op.BitFlip(meas, wires=q)
1057 if isinstance(thermal_relax, dict):
1058 t1 = thermal_relax["t1"]
1059 t2 = thermal_relax["t2"]
1060 t_factor = thermal_relax["t_factor"]
1061 circuit_depth = self._get_circuit_depth()
1062 tg = circuit_depth * t_factor
1063 op.ThermalRelaxationError(1.0, t1, t2, tg, q)
1065 def _get_circuit_depth(self, inputs: Optional[jnp.ndarray] = None) -> int:
1066 """
1067 Calculate the depth of the quantum circuit.
1069 Records the circuit onto a tape (without noise) and computes the
1070 depth as the length of the critical path: each gate is scheduled
1071 at the earliest time step after all of its qubits are free.
1073 Args:
1074 inputs (Optional[jnp.ndarray]): Input data for circuit evaluation.
1075 If None, default zero inputs are used.
1077 Returns:
1078 int: The circuit depth (longest path of gates in the circuit).
1079 """
1080 # Return cached value if available
1081 if hasattr(self, "_cached_circuit_depth"):
1082 return self._cached_circuit_depth
1084 inputs = self._inputs_validation(inputs)
1086 # Temporarily clear noise_params to prevent _variational from
1087 # picking them up (which would call _apply_general_noise ->
1088 # _get_circuit_depth again, causing infinite recursion).
1089 saved_noise = self._noise_params
1090 self._noise_params = None
1092 with recording() as tape:
1093 self._variational(
1094 self.params[0] if self.params.ndim == 3 else self.params,
1095 inputs[0] if inputs.ndim == 2 else inputs,
1096 noise_params=None,
1097 )
1099 self._noise_params = saved_noise
1101 # Filter out noise channels - only count unitary gates
1102 ops = [o for o in tape if not isinstance(o, KrausChannel)]
1104 if not ops:
1105 self._cached_circuit_depth = 0
1106 return 0
1108 # Schedule each gate at the earliest time step where all its wires
1109 # are free. ``wire_busy[q]`` tracks the next free time step for
1110 # qubit ``q``.
1111 wire_busy: Dict[int, int] = {}
1112 depth = 0
1113 for gate in ops:
1114 start = max((wire_busy.get(w, 0) for w in gate.wires), default=0)
1115 end = start + 1
1116 for w in gate.wires:
1117 wire_busy[w] = end
1118 depth = max(depth, end)
1120 self._cached_circuit_depth = depth
1121 return depth
1123 def draw(
1124 self,
1125 inputs: Optional[jnp.ndarray] = None,
1126 figure: str = "text",
1127 **kwargs: Any,
1128 ) -> Union[str, Any]:
1129 """Visualize the quantum circuit.
1131 Records the circuit tape (without noise) and renders the gate
1132 sequence using the requested backend.
1134 Args:
1135 inputs (Optional[jnp.ndarray]): Input data for the circuit.
1136 If ``None``, default zero inputs are used.
1137 figure (str): Rendering backend. One of:
1139 * ``"text"`` - ASCII art (returned as a ``str``).
1140 * ``"mpl"`` - Matplotlib figure (returns ``(fig, ax)``).
1141 * ``"tikz"`` - LaTeX/TikZ ``quantikz`` code (returns a
1142 :class:`TikzFigure`).
1143 * ``"pulse"`` - Pulse schedule (returns ``(fig, axes)``).
1144 Only meaningful for pulse-mode models.
1146 **kwargs: Extra options forwarded to the drawing backend
1147 (e.g. ``gate_values=True``).
1149 Returns:
1150 Depends on figure:
1152 * ``"text"`` -> ``str``
1153 * ``"mpl"`` -> ``(matplotlib.figure.Figure, matplotlib.axes.Axes)``
1154 * ``"tikz"`` -> :class:`TikzFigure`
1156 Raises:
1157 ValueError: If figure is not one of the supported modes.
1158 """
1159 inputs = self._inputs_validation(inputs)
1160 params = self.params[0] if self.params.ndim == 3 else self.params
1161 inp = inputs[0] if inputs.ndim == 2 else inputs
1163 if figure == "pulse":
1164 return self.draw_pulse(inputs=inputs, **kwargs)
1166 # Record without noise to get a clean circuit
1167 saved_noise = self._noise_params
1168 self._noise_params = None
1170 draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
1171 result = draw_script.draw(
1172 figure=figure,
1173 args=(params, inp),
1174 kwargs={"noise_params": None},
1175 **kwargs,
1176 )
1178 self._noise_params = saved_noise
1179 return result
1181 def draw_pulse(
1182 self,
1183 inputs: Optional[jnp.ndarray] = None,
1184 **kwargs: Any,
1185 ) -> Any:
1186 """Visualize the pulse schedule for the circuit.
1188 Records the circuit in pulse mode and collects PulseEvents
1189 automatically via the pulse-event tape, then renders them.
1191 Args:
1192 inputs: Input data. If ``None``, default zero inputs are used.
1193 **kwargs: Forwarded to
1194 :func:`~qml_essentials.drawing.draw_pulse_schedule`
1195 (e.g. ``show_carrier=True``, ``n_samples=300``).
1197 Returns:
1198 ``(fig, axes)`` — Matplotlib Figure and array of Axes.
1199 """
1200 inputs = self._inputs_validation(inputs)
1201 params = self.params[0] if self.params.ndim == 3 else self.params
1202 inp = inputs[0] if inputs.ndim == 2 else inputs
1204 draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
1205 return draw_script.draw(
1206 figure="pulse",
1207 args=(params, inp),
1208 kwargs={
1209 "gate_mode": "pulse",
1210 "noise_params": None,
1211 },
1212 **kwargs,
1213 )
1215 def __repr__(self) -> str:
1216 """Return text representation of the quantum circuit model."""
1217 return self.draw(figure="text")
1219 def __str__(self) -> str:
1220 """Return string representation of the quantum circuit model."""
1221 return self.draw(figure="text")
1223 def _params_validation(self, params: Optional[jnp.ndarray]) -> jnp.ndarray:
1224 """
1225 Validate and normalize variational parameters.
1227 Ensures parameters have the correct shape with a batch dimension,
1228 and updates the model's internal parameters if new ones are provided.
1230 Args:
1231 params (Optional[jnp.ndarray]): Variational parameters to validate.
1232 If None, returns the model's current parameters.
1234 Returns:
1235 jnp.ndarray: Validated parameters with shape
1236 (batch_size, n_layers, n_params_per_layer).
1237 """
1238 # append batch axis if not provided
1239 if params is not None:
1240 if len(params.shape) == 2:
1241 params = np.expand_dims(params, axis=0)
1243 # Avoid stashing JAX tracers on ``self``: under an outer
1244 # transform (e.g. ``jacrev``) the tracer becomes invalid once
1245 # the transform returns, and a subsequent read of
1246 # ``self.params`` would feed a leaked tracer into the next
1247 # call (raising ``UnexpectedTracerError``).
1248 # if not isinstance(params, jax.core.Tracer):
1249 # self.params = params
1250 self.params = params
1251 else:
1252 params = self.params
1254 return params
1256 def _pulse_params_validation(
1257 self, pulse_params: Optional[jnp.ndarray]
1258 ) -> jnp.ndarray:
1259 """
1260 Validate and normalize pulse parameters.
1262 Ensures pulse parameters are set, using model defaults if not provided.
1264 Args:
1265 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers.
1266 If None, returns the model's current pulse parameters.
1268 Returns:
1269 jnp.ndarray: Validated pulse parameters with shape
1270 (batch_size, n_layers, n_pulse_params_per_layer).
1271 """
1272 if pulse_params is None:
1273 pulse_params = self.pulse_params
1274 else:
1275 # ensure batch dimension exists (batch-first convention)
1276 if len(pulse_params.shape) == 2:
1277 pulse_params = jnp.expand_dims(pulse_params, axis=0)
1278 # See note in _params_validation: never stash JAX tracers on
1279 # ``self``.
1280 # if not isinstance(pulse_params, jax.core.Tracer):
1281 # self.pulse_params = pulse_params
1282 self.pulse_params = pulse_params
1284 return pulse_params
1286 def _enc_params_validation(self, enc_params: Optional[jnp.ndarray]) -> jnp.ndarray:
1287 """
1288 Validate and normalize encoding parameters.
1290 Ensures encoding parameters have the correct shape for the model's
1291 input feature dimensions.
1293 Args:
1294 enc_params (Optional[jnp.ndarray]): Encoding parameters to validate.
1295 If None, returns the model's current encoding parameters.
1297 Returns:
1298 jnp.ndarray: Validated encoding parameters with shape
1299 (n_qubits, n_input_feat).
1301 Raises:
1302 ValueError: If enc_params shape is incompatible with n_input_feat > 1.
1303 """
1304 if enc_params is None:
1305 enc_params = self.enc_params
1306 else:
1307 # See note in _params_validation: never stash JAX tracers on
1308 # ``self``.
1309 # if not isinstance(enc_params, jax.core.Tracer):
1310 # if self.trainable_frequencies:
1311 # self.enc_params = enc_params
1312 # else:
1313 # self.enc_params = jnp.array(enc_params)
1314 if self.trainable_frequencies:
1315 self.enc_params = enc_params
1316 else:
1317 self.enc_params = jnp.array(enc_params)
1319 if len(enc_params.shape) == 1 and self.n_input_feat == 1:
1320 enc_params = enc_params.reshape(-1, 1)
1321 elif len(enc_params.shape) == 1 and self.n_input_feat > 1:
1322 raise ValueError(
1323 f"Input dimension {self.n_input_feat} >1 but \
1324 `enc_params` has shape {enc_params.shape}"
1325 )
1327 return enc_params
1329 def _inputs_validation(
1330 self, inputs: Union[None, List, float, int, jnp.ndarray]
1331 ) -> jnp.ndarray:
1332 """
1333 Validate and normalize input data.
1335 Converts various input formats to a standardized 2D array shape
1336 suitable for batch processing in the quantum circuit.
1338 Args:
1339 inputs (Union[None, List, float, int, jnp.ndarray]): Input data in
1340 various formats:
1341 - None: Returns zeros with shape (1, n_input_feat)
1342 - float/int: Single scalar value
1343 - List: List of values or batched inputs
1344 - jnp.ndarray: NumPy/JAX array
1346 Returns:
1347 jnp.ndarray: Validated inputs with shape (batch_size, n_input_feat).
1349 Raises:
1350 ValueError: If input shape is incompatible with expected n_input_feat.
1352 Warns:
1353 UserWarning: If input is replicated to match n_input_feat.
1354 """
1355 self._zero_inputs = False
1356 if isinstance(inputs, List):
1357 inputs = jnp.array(np.stack(inputs))
1358 elif isinstance(inputs, float) or isinstance(inputs, int):
1359 inputs = jnp.array([inputs])
1360 elif inputs is None:
1361 inputs = jnp.array([[0] * self.n_input_feat])
1363 if not inputs.any():
1364 self._zero_inputs = True
1366 if len(inputs.shape) <= 1:
1367 if self.n_input_feat == 1:
1368 # add a batch dimension
1369 inputs = inputs.reshape(-1, 1)
1370 else:
1371 if inputs.shape[0] == self.n_input_feat:
1372 inputs = inputs.reshape(1, -1)
1373 else:
1374 inputs = inputs.reshape(-1, 1)
1375 inputs = inputs.repeat(self.n_input_feat, axis=1)
1376 warnings.warn(
1377 f"Expected {self.n_input_feat} inputs, but {inputs.shape[0]} "
1378 "was provided, replicating input for all input features.",
1379 UserWarning,
1380 )
1381 else:
1382 if inputs.shape[1] != self.n_input_feat:
1383 raise ValueError(
1384 f"Wrong number of inputs provided. Expected {self.n_input_feat} "
1385 f"inputs, but input has shape {inputs.shape}."
1386 )
1388 return inputs
1390 def _postprocess_res(self, result: Union[List, jnp.ndarray]) -> jnp.ndarray:
1391 """
1392 Post-process circuit execution results for uniform shape.
1394 Converts list outputs (from multiple measurements) to stacked arrays
1395 and reorders axes for consistent batch dimension placement.
1397 Args:
1398 result (Union[List, jnp.ndarray]): Raw circuit output, either a
1399 list of measurement results or a single array.
1401 Returns:
1402 jnp.ndarray: Uniformly shaped result array with batch dimension first.
1403 """
1404 if isinstance(result, list):
1405 # we use moveaxis here because in case of parity measure,
1406 # there is another dimension appended to the end and
1407 # simply transposing would result in a wrong shape
1408 result = jnp.stack(result)
1409 if len(result.shape) > 1:
1410 result = jnp.moveaxis(result, 0, 1)
1411 return result
1413 def _assimilate_batch(
1414 self,
1415 inputs: jnp.ndarray,
1416 params: jnp.ndarray,
1417 pulse_params: jnp.ndarray,
1418 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
1419 """
1420 Align batch dimensions across inputs, parameters, and pulse parameters.
1422 Broadcasts and reshapes arrays to have compatible batch dimensions
1423 for vectorized circuit execution. Sets the internal batch_shape.
1425 Args:
1426 inputs (jnp.ndarray): Input data of shape (B_I, n_input_feat).
1427 params (jnp.ndarray): Parameters of shape (B_P, n_layers, n_params).
1428 pulse_params (jnp.ndarray): Pulse params of shape (B_R, n_layers, n_pulse).
1430 Returns:
1431 Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: Tuple containing:
1432 - inputs: Reshaped to (B, n_input_feat) where B = B_I * B_P * B_R
1433 - params: Reshaped to (B, n_layers, n_params)
1434 - pulse_params: Reshaped to (B, n_layers, n_pulse)
1436 Note:
1437 The effective batch shape depends on repeat_batch_axis configuration.
1438 This is the only method that sets self._batch_shape.
1439 """
1440 B_I = inputs.shape[0]
1441 # we check for the product because there is a chance that
1442 # there are no params. In this case we want B_P to be 1
1443 B_P = 1 if 0 in params.shape else params.shape[0]
1444 B_R = pulse_params.shape[0]
1446 # THIS is the only place where we set the batch shape
1447 self._batch_shape = (B_I, B_P, B_R)
1448 B = np.prod(self.eff_batch_shape)
1450 # [B_I, ...] -> [B_I, B_P, B_R, ...] -> [B, ...]
1451 if B_I > 1 and self.repeat_batch_axis[0]:
1452 if self.repeat_batch_axis[1]:
1453 inputs = jnp.repeat(inputs[:, None, None, ...], B_P, axis=1)
1454 if self.repeat_batch_axis[2]:
1455 inputs = jnp.repeat(inputs, B_R, axis=2)
1456 inputs = inputs.reshape(B, *inputs.shape[3:])
1458 # [B_P, ..., ...] -> [B_I, B_P, B_R, ..., ...] -> [B, ..., ...]
1459 if B_P > 1 and self.repeat_batch_axis[1]:
1460 # add B_I axis before first, and B_R axis after first batch dim
1461 params = params[None, :, None, ...] # [B_I(=1), B_P, B_R(=1), ...]
1462 if self.repeat_batch_axis[0]:
1463 params = jnp.repeat(params, B_I, axis=0) # [B_I, B_P, 1, ...]
1464 if self.repeat_batch_axis[2]:
1465 params = jnp.repeat(params, B_R, axis=2) # [B_I, B_P, B_R, ...]
1466 params = params.reshape(B, *params.shape[3:])
1468 # [B_R, ..., ...] -> [B_I, B_P, B_R, ..., ...] -> [B, ..., ...]
1469 if B_R > 1 and self.repeat_batch_axis[2]:
1470 # add B_I axis and B_P axis before B_R
1471 pulse_params = pulse_params[None, None, ...] # [B_I(=1), B_P(=1), B_R, ...]
1472 if self.repeat_batch_axis[0]:
1473 pulse_params = jnp.repeat(
1474 pulse_params, B_I, axis=0
1475 ) # [B_I, 1, B_R, ...]
1476 if self.repeat_batch_axis[1]:
1477 pulse_params = jnp.repeat(
1478 pulse_params, B_P, axis=1
1479 ) # [B_I, B_P, B_R, ...]
1480 pulse_params = pulse_params.reshape(B, *pulse_params.shape[3:])
1482 return inputs, params, pulse_params
1484 def _requires_density(self) -> bool:
1485 """
1486 Check if density matrix simulation is required.
1488 Determines whether the circuit must be executed with the mixed-state
1489 simulator based on execution type and noise configuration.
1491 Returns:
1492 bool: True if density matrix simulation is required, False otherwise.
1493 Returns True if:
1494 - execution_type is "density", or
1495 - Any non-coherent noise channel has non-zero probability
1496 """
1497 if self.execution_type == "density":
1498 return True
1500 if self.noise_params is None:
1501 return False
1503 coherent_noise = {"GateError"}
1504 for k, v in self.noise_params.items():
1505 if k in coherent_noise:
1506 continue
1507 if v is not None and v > 0:
1508 return True
1509 return False
1511 def __call__(
1512 self,
1513 params: Optional[jnp.ndarray] = None,
1514 inputs: Optional[jnp.ndarray] = None,
1515 pulse_params: Optional[jnp.ndarray] = None,
1516 enc_params: Optional[jnp.ndarray] = None,
1517 data_reupload: Union[bool, List[List[bool]], List[List[List[bool]]]] = None,
1518 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
1519 execution_type: Optional[str] = None,
1520 force_mean: bool = False,
1521 gate_mode: str = "unitary",
1522 ) -> jnp.ndarray:
1523 """
1524 Execute the quantum circuit (callable interface).
1526 Provides a convenient callable interface for circuit execution,
1527 delegating to the _forward method.
1529 Args:
1530 params (Optional[jnp.ndarray]): Variational parameters of shape
1531 (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer).
1532 If None, uses model's internal parameters.
1533 inputs (Optional[jnp.ndarray]): Input data of shape
1534 (batch_size, n_input_feat). If None, uses zero inputs.
1535 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
1536 pulse-mode gate execution.
1537 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
1538 (n_qubits, n_input_feat). If None, uses model's encoding parameters.
1539 data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]]]):
1540 Data reupload configuration. If None, uses previously set reupload
1541 configuration.
1542 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
1543 Noise configuration. If None, uses previously set noise parameters.
1544 execution_type (Optional[str]): Measurement type: "expval", "density",
1545 "probs", or "state". If None, uses current execution_type setting.
1546 force_mean (bool): If True, averages results over measurement qubits.
1547 Defaults to False.
1548 gate_mode (str): Gate execution backend, "unitary" or "pulse".
1549 Defaults to "unitary".
1551 Returns:
1552 jnp.ndarray: Circuit output with shape depending on execution_type:
1553 - "expval": (n_output_qubits,) or scalar
1554 - "density": (2^n_output, 2^n_output)
1555 - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
1556 - "state": (2^n_qubits,)
1557 """
1558 # Call forward method which handles the actual caching etc.
1559 return self._forward(
1560 params=params,
1561 inputs=inputs,
1562 pulse_params=pulse_params,
1563 enc_params=enc_params,
1564 data_reupload=data_reupload,
1565 noise_params=noise_params,
1566 execution_type=execution_type,
1567 force_mean=force_mean,
1568 gate_mode=gate_mode,
1569 )
1571 def _forward(
1572 self,
1573 params: Optional[jnp.ndarray] = None,
1574 inputs: Optional[jnp.ndarray] = None,
1575 pulse_params: Optional[jnp.ndarray] = None,
1576 enc_params: Optional[jnp.ndarray] = None,
1577 data_reupload: Union[bool, List[List[bool]], List[List[List[bool]]]] = None,
1578 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
1579 execution_type: Optional[str] = None,
1580 force_mean: bool = False,
1581 gate_mode: str = "unitary",
1582 ) -> jnp.ndarray:
1583 """
1584 Execute the quantum circuit forward pass.
1586 Internal implementation of the forward pass that handles parameter
1587 validation, batch alignment, and circuit execution routing.
1589 Args:
1590 params (Optional[jnp.ndarray]): Variational parameters of shape
1591 (n_layers, n_params_per_layer) or
1592 (batch, n_layers, n_params_per_layer).
1593 If None, uses model's internal parameters.
1594 inputs (Optional[jnp.ndarray]): Input data of shape
1595 (batch_size, n_input_feat).
1596 If None, uses zero inputs.
1597 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
1598 pulse-mode gate execution.
1599 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
1600 (n_qubits, n_input_feat). If None, uses model's encoding parameters.
1601 data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]]]):
1602 Data reupload configuration. If None, uses previously set reupload
1603 configuration.
1604 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
1605 Noise configuration. If None, uses previously set noise parameters.
1606 execution_type (Optional[str]): Measurement type: "expval", "density",
1607 "probs", or "state". If None, uses current execution_type setting.
1608 force_mean (bool): If True, averages results over measurement qubits.
1609 Defaults to False.
1610 gate_mode (str): Gate execution backend, "unitary" or "pulse".
1611 Defaults to "unitary".
1613 Returns:
1614 jnp.ndarray: Circuit output with shape depending on execution_type:
1615 - "expval": (n_output_qubits,) or scalar
1616 - "density": (2^n_output, 2^n_output)
1617 - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
1618 - "state": (2^n_qubits,)
1620 Raises:
1621 ValueError: If pulse_params provided without pulse gate_mode, or
1622 if noise_params provided with pulse gate_mode.
1623 """
1624 # set the parameters as object attributes
1625 if noise_params is not None:
1626 self.noise_params = noise_params
1627 if execution_type is not None:
1628 self.execution_type = execution_type
1629 self.gate_mode = gate_mode
1631 # consistency checks
1632 if pulse_params is not None and gate_mode != "pulse":
1633 raise ValueError(
1634 "pulse_params were provided but gate_mode is not 'pulse'. "
1635 "Either switch gate_mode='pulse' or do not pass pulse_params."
1636 )
1638 # TODO: add testing
1639 if data_reupload is not None:
1640 self.data_reupload = data_reupload
1642 params = self._params_validation(params)
1643 pulse_params = self._pulse_params_validation(pulse_params)
1644 inputs = self._inputs_validation(inputs)
1645 enc_params = self._enc_params_validation(enc_params)
1647 inputs, params, pulse_params = self._assimilate_batch(
1648 inputs,
1649 params,
1650 pulse_params,
1651 )
1653 # split to generate a sub_key, required for actual execution
1654 self.random_key, sub_key = safe_random_split(self.random_key)
1656 # Build measurement type & observables from execution_type / output_qubit
1657 meas_type, obs = self._build_obs()
1659 # Jaqsi auto-routes between statevector and density-matrix simulation
1660 # based on whether noise channels appear on the tape, so a single
1661 B = np.prod(self.eff_batch_shape)
1663 # kwargs are broadcast (not vmapped over)
1664 exec_kwargs = dict(
1665 noise_params=self.noise_params,
1666 gate_mode=self.gate_mode,
1667 )
1669 # Build a shot key from the random_key if shots are requested
1670 shot_key = None
1671 if self.shots is not None:
1672 # overwrite subkey and split shot_key
1673 sub_key, shot_key = safe_random_split(sub_key)
1675 if B > 1:
1676 # use random keys, derived from the subkey
1677 random_keys = safe_random_split(sub_key, num=B)
1679 in_axes = (
1680 0 if self.batch_shape[1] > 1 else None, # params
1681 0 if self.batch_shape[0] > 1 else None, # inputs
1682 0 if self.batch_shape[2] > 1 else None, # pulse_params
1683 0, # random_keys
1684 None, # enc_params (broadcast, not batched)
1685 )
1687 result = self.script.execute(
1688 type=meas_type,
1689 obs=obs,
1690 args=(params, inputs, pulse_params, random_keys, enc_params),
1691 kwargs=exec_kwargs,
1692 in_axes=in_axes,
1693 shots=self.shots,
1694 key=shot_key,
1695 )
1696 else:
1697 # use the subkey directly
1698 result = self.script.execute(
1699 type=meas_type,
1700 obs=obs,
1701 args=(params, inputs, pulse_params, sub_key, enc_params),
1702 kwargs=exec_kwargs,
1703 shots=self.shots,
1704 key=shot_key,
1705 )
1707 result = self._postprocess_res(result)
1709 # --- Post-processing for partial-qubit measurements ---------------
1710 if self.execution_type == "density" and not self.all_qubit_measurement:
1711 result = js.partial_trace(result, self.n_qubits, self.output_qubit)
1713 if self.execution_type == "probs" and not self.all_qubit_measurement:
1714 if isinstance(self.output_qubit[0], (list, tuple)):
1715 # list of qubit groups - marginalize each independently
1716 result = jnp.stack(
1717 [
1718 js.marginalize_probs(result, self.n_qubits, list(group))
1719 for group in self.output_qubit
1720 ]
1721 )
1722 else:
1723 result = js.marginalize_probs(result, self.n_qubits, self.output_qubit)
1725 result = jnp.asarray(result)
1726 result = result.reshape((*self.eff_batch_shape, *self._result_shape)).squeeze()
1728 if (
1729 self.execution_type in ("expval", "probs")
1730 and force_mean
1731 and len(result.shape) > 0
1732 and self._result_shape[0] > 1
1733 ):
1734 result = result.mean(axis=-1)
1736 return result