redefining OperatorEvolution class and fixing typehints

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,33 +240,31 @@ 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
         )
         measurement_times[measurement_times < SAMPLING_INTERVAL] = SAMPLING_INTERVAL
         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:
     """

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."""

src/qibolab/_core/instruments/emulator/engine/dynamiqs.py

@@ -148,30 +148,29 @@ 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:
         """Evolve the system."""
 
         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,

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

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(