diff --git a/src/qibolab/_core/instruments/emulator/emulator.py b/src/qibolab/_core/instruments/emulator/emulator.py index a1c4d7355d..c5d95cd1b2 100644 --- a/src/qibolab/_core/instruments/emulator/emulator.py +++ b/src/qibolab/_core/instruments/emulator/emulator.py @@ -14,7 +14,6 @@ from pydantic import model_validator from qibolab._core.components import Config -from qibolab._core.components.configs import AcquisitionConfig from qibolab._core.execution_parameters import AveragingMode, ExecutionParameters from qibolab._core.identifier import Result from qibolab._core.instruments.abstract import Controller @@ -40,8 +39,10 @@ SWEEP_SIMULATION_FILENAME, ) from .hamiltonians import ( + ControlLine, + DriveEmulatorConfig, + FluxEmulatorConfig, HamiltonianConfig, - Modulated, waveform, ) from .results import acquisitions, index, results @@ -81,7 +82,7 @@ def sampling_rate(self) -> float: return self.sampling_rate_ @sampling_rate.setter - def sampling_rate(self, value: float) -> float: + def sampling_rate(self, value: float) -> None: self.sampling_rate_ = value def connect(self): @@ -93,7 +94,6 @@ def disconnect(self): def _dump_simulation( self, sequence_idx, - static_ham: Operator, evolution: OperatorEvolution, states: NDArray, simulation_config: dict, @@ -119,7 +119,7 @@ def _dump_simulation( if not static_hamiltonian_filename.exists(): # list of file operators of the pulse sequence; NOTE: the first element is always the time independent hamiltonian operators = np.stack( - [static_ham.full()] + [op.full() for op, _ in evolution.operators] + [evolution.static.full()] + [op.full() for op, _ in evolution.operators] ) np.save(static_hamiltonian_filename, operators) @@ -137,7 +137,6 @@ def _dump_simulation( sequence_dir / (SWEEP_SIMULATION_FILENAME + f"_{sweep_idx}.npz"), time_coeffs=time_coefficients, results=states, - sim_config=simulation_config, ) def play( @@ -241,7 +240,7 @@ def _evolve( configs_ = update_configs(configs, updates) config = cast(HamiltonianConfig, configs_["hamiltonian"]) hamiltonian = config.hamiltonian(config=configs_, engine=self.engine) - time_hamiltonian = self._pulse_hamiltonian(sequence_, configs_) + complete_hamiltonian = self._pulse_hamiltonian(sequence_, hamiltonian, configs_) measurement_times = np.array( list(acquisitions(sequence_).values()), dtype=float ) @@ -249,25 +248,23 @@ def _evolve( tlist_, index = np.unique(measurement_times, return_inverse=True) results, simulation_configs = self.engine.evolve( - hamiltonian=hamiltonian, + hamiltonian=complete_hamiltonian, initial_state=config.initial_state(self.engine), time=np.concatenate(([0], tlist_)), collapse_operators=config.dissipation(self.engine), - time_hamiltonian=time_hamiltonian, ) states = np.stack([s.full() for s in results.states[1:]])[index] self._dump_simulation( sequence_identifier, - hamiltonian, - time_hamiltonian, + complete_hamiltonian, states, simulation_configs, ) return states def _pulse_hamiltonian( - self, sequence: PulseSequence, configs: dict[str, Config] + self, sequence: PulseSequence, static_ham: Operator, configs: dict[str, Config] ) -> OperatorEvolution: """Construct Hamiltonian time dependent term for qutip simulation.""" @@ -290,7 +287,7 @@ def _pulse_hamiltonian( ) ] - return OperatorEvolution(operators=channels, times=times) + return OperatorEvolution(static=static_ham, operators=channels, times=times) def update_sequence(sequence: PulseSequence, updates: dict) -> PulseSequence: @@ -320,12 +317,12 @@ def tlist(sequence: PulseSequence) -> NDArray: def hamiltonian( pulses: Iterable[PulseLike], - config: Config, + config: FluxEmulatorConfig | DriveEmulatorConfig, hamiltonian: HamiltonianConfig, hilbert_space_index: int, engine: SimulationEngine, sampling_rate: float, -) -> tuple[Operator, list[Modulated]]: +) -> tuple[Operator, list[ControlLine]]: n = hamiltonian.transmon_levels op = engine.expand( op=config.operator(n=n, engine=engine), @@ -345,7 +342,7 @@ def hamiltonians( configs: dict[str, Config], engine: SimulationEngine, sampling_rate: float, -) -> Iterable[tuple[Operator, list[Modulated]]]: +) -> Iterable[tuple[Operator, list[ControlLine]]]: hconfig = cast(HamiltonianConfig, configs["hamiltonian"]) return ( hamiltonian( @@ -358,13 +355,13 @@ def hamiltonians( ) for ch in sequence.channels # TODO: drop the following, and treat acquisitions just as empty channels - if not isinstance(configs[ch], AcquisitionConfig) + if isinstance(configs[ch], FluxEmulatorConfig | DriveEmulatorConfig) ) def channel_coefficients( - waveforms: Iterable[Modulated], - sampling_rate: int, + waveforms: Iterable[ControlLine], + sampling_rate: float, times: NDArray, ) -> NDArray: """ diff --git a/src/qibolab/_core/instruments/emulator/engine/abstract.py b/src/qibolab/_core/instruments/emulator/engine/abstract.py index efef751407..4636daebe2 100644 --- a/src/qibolab/_core/instruments/emulator/engine/abstract.py +++ b/src/qibolab/_core/instruments/emulator/engine/abstract.py @@ -5,7 +5,7 @@ from typing import Protocol import numpy as np -from numpy.typing import NDArray +from numpy.typing import ArrayLike, NDArray from qibolab._core.serialize import Model @@ -43,9 +43,27 @@ class Operator(Protocol): def dag(self) -> "Operator": """Return the adjoint of the operator.""" + def full(self) -> "Operator": + """Return the matrix form of the operator.""" + def __add__(self, other: "Operator") -> "Operator": """Add two operators.""" + def __sub__(self, other: "Operator") -> "Operator": + """Subtract two operators.""" + + def __mul__(self, other: float | int | complex) -> "Operator": + """Scalar multiplication.""" + + def __rmul__(self, other: float | int | complex) -> "Operator": + """Right-hand scalar multiplication.""" + + def __truediv__(self, other: float | int | complex) -> "Operator": + """Scalar division.""" + + def __matmul__(self, other: "Operator") -> "Operator": + """Multiply two operators.""" + TimeDependentOperator = tuple[Operator, NDArray] """Abstract time dependent operator type.""" @@ -62,7 +80,9 @@ class EvolutionResult(Protocol): class OperatorEvolution: """Abstract operator evolution interface.""" - operators: list[Operator | TimeDependentOperator] = field(default_factory=list) + static: Operator + """static term in the system evolution""" + operators: list[TimeDependentOperator] = field(default_factory=list) """List of static or time-dependent operators for evolution.""" times: NDArray = field(default_factory=lambda: np.array([], dtype=float)) """Evolution times with time step equal to the waveforms resolution.""" @@ -79,10 +99,10 @@ def engine(self): @abstractmethod def evolve( self, - hamiltonian: Operator, + hamiltonian: OperatorEvolution, initial_state: Operator, - time: list[float], - collapse_operators: list[Operator] = None, + time: ArrayLike, + collapse_operators: list[Operator] | None = None, **kwargs, ) -> tuple[EvolutionResult, dict]: """Evolve the system.""" @@ -110,5 +130,5 @@ def expand( """Expand operator in larger Hilbert space.""" @abstractmethod - def basis(self, n: int, state: int) -> Operator: + def basis(self, dim: int | list[int], state: int | list[int]) -> Operator: """Basis operator for n levels system.""" diff --git a/src/qibolab/_core/instruments/emulator/engine/dynamiqs.py b/src/qibolab/_core/instruments/emulator/engine/dynamiqs.py index 007f7a35b4..a780b285b0 100644 --- a/src/qibolab/_core/instruments/emulator/engine/dynamiqs.py +++ b/src/qibolab/_core/instruments/emulator/engine/dynamiqs.py @@ -148,10 +148,9 @@ def _method(self, time: np.ndarray): def evolve( self, - hamiltonian: Operator, + hamiltonian: OperatorEvolution, initial_state: Operator, time: Iterable[float], - time_hamiltonian: OperatorEvolution = None, collapse_operators: list[Operator] = None, **kwargs, ) -> DynamiqsEvolutionResult: @@ -159,19 +158,19 @@ def evolve( time = np.asarray(list(time), dtype=float) - hamiltonian = self.engine.constant(_unwrap(hamiltonian)) - if time_hamiltonian is not None: - times = time_hamiltonian.times - for operator, coefficient in time_hamiltonian.operators: - spline = make_interp_spline(times, coefficient, k=SPLINE_INTERP_ORDER) - hamiltonian += self.engine.modulated( - jax_interpolation(spline), _unwrap(operator) - ) + complete_hamiltonian = self.engine.constant(_unwrap(hamiltonian.static)) + for operator, coefficient in hamiltonian.operators: + spline = make_interp_spline( + hamiltonian.times, coefficient, k=SPLINE_INTERP_ORDER + ) + complete_hamiltonian += self.engine.modulated( + jax_interpolation(spline), _unwrap(operator) + ) method = kwargs.pop("method", self._method(time)) options = kwargs.pop("options", self.engine.Options(progress_meter=False)) result = self.engine.mesolve( - hamiltonian, + complete_hamiltonian, [_unwrap(op) for op in collapse_operators or []], _unwrap(initial_state), time, diff --git a/src/qibolab/_core/instruments/emulator/engine/qutip.py b/src/qibolab/_core/instruments/emulator/engine/qutip.py index 5df9ab1d69..71774dd38b 100644 --- a/src/qibolab/_core/instruments/emulator/engine/qutip.py +++ b/src/qibolab/_core/instruments/emulator/engine/qutip.py @@ -1,12 +1,12 @@ import json import os -from collections.abc import Callable, Iterable +from collections.abc import Callable from functools import cached_property from pathlib import Path from typing import Literal import numpy as np -from numpy.typing import NDArray +from numpy.typing import ArrayLike, NDArray from scipy.interpolate import make_interp_spline from .abstract import ( @@ -44,6 +44,7 @@ def engine(self): if self.device == "gpu" and any( device.platform == "gpu" for device in devices ): + # must be always imported import qutip_jax # noqa: F401 else: object.__setattr__(self, "device", "cpu") @@ -56,7 +57,7 @@ def _to_device(self, op: Operator) -> Operator: def interpolate_coeffs( self, x: NDArray, y: NDArray - ) -> Callable[[NDArray], Iterable[float]]: + ) -> Callable[[NDArray], ArrayLike]: spline = make_interp_spline(x, y, k=SPLINE_INTERP_ORDER) if self.device == "gpu": @@ -71,10 +72,9 @@ def interpolate_coeffs( def evolve( self, - hamiltonian: Operator, + hamiltonian: OperatorEvolution, initial_state: Operator, - time: Iterable[float], - time_hamiltonian: OperatorEvolution, + time: ArrayLike, collapse_operators: list[Operator] | None = None, **kwargs, ): @@ -108,16 +108,18 @@ def evolve( # define nsteps instead options = {"max_step": INTEGRATION_MAX_TIME_STEP, "nsteps": nsteps} - hamiltonian = [self._to_device(hamiltonian)] + [ + qutip_hamiltonian: list[Operator | list[Operator | Callable]] = [ + self._to_device(hamiltonian.static) + ] + [ [ self._to_device(operator), - self.interpolate_coeffs(time_hamiltonian.times, coefficient), + self.interpolate_coeffs(hamiltonian.times, coefficient), ] - for operator, coefficient in time_hamiltonian.operators + for operator, coefficient in hamiltonian.operators ] sim_results = self.engine.mesolve( - hamiltonian, + qutip_hamiltonian, self._to_device(initial_state), time, [self._to_device(op) for op in collapse_operators or []], @@ -147,7 +149,7 @@ def expand(self, op: Operator, dims: list[int], targets: int | list[int]): """Expand operator in larger Hilbert space.""" return self._to_device(self.engine.expand_operator(op, dims, targets)) - def basis(self, dim: int, state: int) -> Operator: + def basis(self, dim: int | list[int], state: int | list[int]) -> Operator: """Basis operator for n levels system.""" return self._to_device(self.engine.basis(dimensions=dim, n=state)) @@ -187,11 +189,13 @@ def load_simulation( with open(simulated_sequence_path / (SIMULATOR_CONFIG + ".json")) as f: sim_configs = json.load(f) - system = [Operator(hamiltonians[0])] + [ + system = [ TimeDependentOperator([ham, coeffs]) for ham, coeffs in zip(hamiltonians[1:], time_coeffs[1:]) ] timesteps = time_coeffs[0] - system_evo = OperatorEvolution(operators=system, times=timesteps) + system_evo = OperatorEvolution( + static=hamiltonians[0], operators=system, times=timesteps + ) return system_evo, result_states, sim_configs diff --git a/src/qibolab/_core/instruments/emulator/hamiltonians.py b/src/qibolab/_core/instruments/emulator/hamiltonians.py index a3ee24fced..f53fc227ba 100644 --- a/src/qibolab/_core/instruments/emulator/hamiltonians.py +++ b/src/qibolab/_core/instruments/emulator/hamiltonians.py @@ -68,11 +68,11 @@ class Qubit(Config): t2: dict[TransitionId, float] = Field(default_factory=dict) """Dictionary with dephasing time per transition.""" - def omega(self, flux: float = 0) -> float: + def omega(self, flux: float | NDArray = 0) -> float | NDArray: """Angular velocity.""" return 2 * np.pi * self.detuned_frequency(flux) - def detuned_frequency(self, flux: float) -> float: + def detuned_frequency(self, flux: float | NDArray) -> float | NDArray: """Return frequency of the qubit modified by the flux.""" return (self.frequency - self.anharmonicity) * ( self.asymmetry**2 @@ -183,7 +183,7 @@ def __call__(self, times: NDArray, samples: NDArray, phase: float) -> IqWaveform # we are passing the relative frequency because the term with the offset # is already included in the time-independent part of the Hamiltonian # and it corresponds to changing the static bias - return ( + return np.asarray( 2 * np.pi * ( @@ -240,7 +240,7 @@ class ModulatedDelay(Model): phase: float = 0 """Delay has 0 virtual z phase.""" - def __call__(self, t: float, sample: int, phase: float) -> float: + def __call__(self, t: NDArray, sample: NDArray, phase: float) -> float: """Delay waveform.""" return 0 @@ -253,7 +253,7 @@ class ModulatedVirtualZ(Model): duration: float = 0 """Duration is 0 for virtual Z.""" - def __call__(self, t: float, sample: int, phase: float) -> float: + def __call__(self, t: NDArray, sample: NDArray, phase: float) -> None: """Delay waveform.""" raise_error(ValueError, "VirtualZ doesn't have waveform.") @@ -322,11 +322,11 @@ def hamiltonian(self, config: dict, engine: SimulationEngine) -> Operator: ) return qubit_terms + coupling - def dissipation(self, engine: SimulationEngine) -> Operator: + def dissipation(self, engine: SimulationEngine) -> list[Operator]: """Dissipation operators for the hamiltonian. They are going to be passed to mesolve as collapse operators.""" - collapse_operators = [] + collapse_operators: list[Operator] = [] for i, qubit in self.qubits.items(): if len(qubit.t1) > 0: collapse_operators.append( diff --git a/src/qibolab/_core/instruments/emulator/results.py b/src/qibolab/_core/instruments/emulator/results.py index d55e226021..ffbd6f2a25 100644 --- a/src/qibolab/_core/instruments/emulator/results.py +++ b/src/qibolab/_core/instruments/emulator/results.py @@ -205,6 +205,10 @@ def _singleshot_results( corresponding measurement results. """ + assert options.nshots is not None, ( + "Number of shots not inserted for SINGLESHOT simulation." + ) + # select only unique times of measurements _, direct_map, inverse_map = np.unique( list(acquisitions(sequence).values()), @@ -252,7 +256,7 @@ def results( sequence: PulseSequence, hamiltonian: HamiltonianConfig, options: ExecutionParameters, -) -> dict[int, Result]: +) -> dict[PulseId, Result]: """Collect results for a single pulse sequence. The dictionary returned is already compliant with the expected diff --git a/tests/instruments/emulator/test_dynamiqs.py b/tests/instruments/emulator/test_dynamiqs.py index 2bc8364995..d3bd97e5b6 100644 --- a/tests/instruments/emulator/test_dynamiqs.py +++ b/tests/instruments/emulator/test_dynamiqs.py @@ -89,16 +89,15 @@ def test_dynamiqs_evolution_matches_qutip(): times = np.linspace(0, 1, 5) dynamiqs_result, _ = dynamiqs.evolve( - hamiltonian=dynamiqs.identity(2) * 0, + hamiltonian=OperatorEvolution(static=dynamiqs.identity(2) * 0), initial_state=dynamiqs.basis(2, 1), time=times, collapse_operators=[0.2 * dynamiqs.destroy(2)], ) qutip_result, _ = qutip.evolve( - hamiltonian=qutip.identity(2) * 0, + hamiltonian=OperatorEvolution(static=qutip.identity(2) * 0), initial_state=qutip.basis(2, 1), time=times, - time_hamiltonian=OperatorEvolution(), collapse_operators=[0.2 * qutip.destroy(2)], ) @@ -125,14 +124,14 @@ def _driven_evolution(engine, **kwargs): * np.cos(2 * np.pi * 0.2 * times) ) evolution = OperatorEvolution( + static=hamiltonian, operators=[(a + a.dag(), coefficient)], times=times, ) result, _ = engine.evolve( - hamiltonian=hamiltonian, + hamiltonian=evolution, initial_state=engine.basis(3, 0), time=np.linspace(0.0, 20.0, 5), - time_hamiltonian=evolution, collapse_operators=[0.005 * a], **kwargs, )