additional emulator fixes

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

@@ -1,5 +1,7 @@
 """Emulator controller."""
 
+import json
+import os
 from collections import defaultdict
 from collections.abc import Iterable
 from functools import reduce
@@ -9,7 +11,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from scipy.interpolate import BSpline
+from pydantic import model_validator
 
 from qibolab._core.components import Config
 from qibolab._core.components.configs import AcquisitionConfig
@@ -25,16 +27,25 @@
 from qibolab._core.sequence import PulseSequence
 from qibolab._core.sweeper import ParallelSweepers
 
-from .engine import Operator, OperatorEvolution, QutipEngine, SimulationEngine
+from .engine import (
+    Operator,
+    OperatorEvolution,
+    QutipEngine,
+    SimulationEngine,
+    TimeDependentOperator,
+)
+from .engine.abstract import (
+    HAMILTONIAN_FILENAME,
+    SIMULATOR_CONFIG,
+    SWEEP_SIMULATION_FILENAME,
+)
 from .hamiltonians import (
     HamiltonianConfig,
     Modulated,
     waveform,
 )
 from .results import acquisitions, index, results
 
-SPLINE_INTERP_ORDER = 3
-"""Polynomial order used for interpolating the pulses with a spline function."""
 NYQUIST_FREQUENCY = 20
 """GHz, Nyquist frequency used for computing the solution and resolve qubit oscillations."""
 SAMPLING_INTERVAL = 1 / (2 * NYQUIST_FREQUENCY)
@@ -51,16 +62,26 @@ class EmulatorController(Controller):
     """Sampling rate used during simulation."""
     engine: SimulationEngine = QutipEngine()
     """SimulationEngine. Default is QutipEngine."""
-    save_dir: Path | None = None
+    save_dir: os.PathLike | str | None = None
     """Flag for saving the full system evolution computed from the simulation
     backend. In order to set it True modify `platform.py` file in the platform folder."""
 
+    @model_validator(mode="after")
+    def validate_save_dir(self):
+        if self.save_dir is not None:
+            # converting every possible output as a pathlib.Path object
+            save_dir = Path(self.save_dir)
+            if save_dir.exists():
+                raise FileExistsError("The given data folder already exists.")
+            object.__setattr__(self, "save_dir", save_dir)
+        return self
+
     @property
     def sampling_rate(self) -> float:
         return self.sampling_rate_
 
     @sampling_rate.setter
-    def sampling_rate(self, value: float):
+    def sampling_rate(self, value: float) -> float:
         self.sampling_rate_ = value
 
     def connect(self):
@@ -69,22 +90,55 @@ def connect(self):
     def disconnect(self):
         """Dummy disconnect method."""
 
-    def _dump_simulation(self, sequence, configs, states, coefficients):
+    def _dump_simulation(
+        self,
+        sequence_idx,
+        static_ham: Operator,
+        evolution: OperatorEvolution,
+        states: NDArray,
+        simulation_config: dict,
+    ) -> None:
         """Write operators (once), time coefficients (n-d), density matrices (n-d)."""
-        self.save_dir.mkdir(parents=True, exist_ok=True)
-
-        # operators: run-invariant, so build them once here (outside the sweep loop)
-        config = cast(HamiltonianConfig, configs["hamiltonian"])
-        static = config.hamiltonian(config=configs, engine=self.engine)
-        evolution = self._pulse_hamiltonian(sequence, configs)
-        time_ops = (
-            [pair[0] for pair in evolution.operators] if evolution is not None else []
+
+        if self.save_dir is None:
+            return
+
+        sequence_dir = self.save_dir / f"sequence_{sequence_idx}"
+        sequence_dir.mkdir(parents=True, exist_ok=True)
+
+        # list of file coefficients; NOTE: the first element is always the simulation timesteps array
+        time_coefficients: list[NDArray] = np.stack(
+            [evolution.times] + [c for _, c in evolution.operators]
         )
-        operators = [static] + time_ops
 
-        self.engine.save_operators(operators, self.save_dir)
-        np.save(self.save_dir / "time_coefficients.npy", coefficients)
-        np.save(self.save_dir / "density_matrices.npy", states)
+        # solver configuration file path
+        json_filename = sequence_dir / (SIMULATOR_CONFIG + ".json")
+        static_hamiltonian_filename = sequence_dir / (HAMILTONIAN_FILENAME + ".npy")
+
+        # check if the the sweeper-independent data for the current sequence have already been dumped
+        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]
+            )
+            np.save(static_hamiltonian_filename, operators)
+
+        if not json_filename.exists():
+            with open(json_filename, "w") as f:
+                json.dump(simulation_config, f)
+
+        # NOTE: this might crash if the process is parallelized, to be review once we enable parallelization
+        sweep_idx = sum(
+            1
+            for file in sequence_dir.iterdir()
+            if file.is_file() and SWEEP_SIMULATION_FILENAME in file.name
+        )
+        np.savez(
+            sequence_dir / (SWEEP_SIMULATION_FILENAME + f"_{sweep_idx}.npz"),
+            time_coeffs=time_coefficients,
+            results=states,
+            sim_config=simulation_config,
+        )
 
     def play(
         self,
@@ -105,15 +159,15 @@ def play(
 
         results_to_process = (
             self._play_sequence(configs, sequence, options, sweepers)
-            for sequence in sequences_
+            for sequence in enumerate(sequences_)
         )
 
         return reduce(or_, results_to_process)
 
     def _play_sequence(
         self,
         configs: dict[str, Config],
-        sequence: PulseSequence,
+        sequence: tuple[int, PulseSequence],
         options: ExecutionParameters,
         sweepers: list[ParallelSweepers],
     ):
@@ -122,21 +176,19 @@ def _play_sequence(
         Executes a sweep of the quantum sequence and processes the results
         into a structured results object containing quantum states and measurement data.
         """
-        sweep_states, sweep_coefficients = self._sweep(sequence, configs, sweepers)
-        if self.save_dir is not None:
-            self._dump_simulation(sequence, configs, sweep_states, sweep_coefficients)
+        sweep_states = self._sweep(sequence, configs, sweepers)
         hamiltonian = cast(HamiltonianConfig, configs["hamiltonian"])
         return results(
             # states in computational basis
             states=sweep_states,
-            sequence=sequence,
+            sequence=sequence[1],
             hamiltonian=hamiltonian,
             options=options,
         )
 
     def _sweep(
         self,
-        sequence: PulseSequence,
+        sequence: tuple[int, PulseSequence],
         configs: dict[str, Config],
         sweepers: list[ParallelSweepers],
         updates: dict | None = None,
@@ -155,7 +207,6 @@ def _sweep(
             return self._evolve(sequence, configs, updates)
 
         state_slices: list[NDArray] = []
-        coeff_slices = []
         parsweep = sweepers[0]
         # execute once for each parallel value
         for values in zip(*(s.values for s in parsweep)):
@@ -167,22 +218,25 @@ def _sweep(
                 if sweeper.channels is not None:
                     for channel in sweeper.channels:
                         updates[channel].update({sweeper.parameter.name: value})
-            states, coeffs = self._sweep(sequence, configs, sweepers[1:], updates)
-            state_slices.append(states)
-            coeff_slices.append(coeffs)
+            state_slices.append(self._sweep(sequence, configs, sweepers[1:], updates))
 
         # stack all slices in a single array, along the current outermost dimension
-        return np.stack(state_slices), np.stack(coeff_slices)
+        return np.stack(state_slices)
 
     def _evolve(
-        self, sequence: PulseSequence, configs: dict[str, Config], updates: dict
+        self,
+        sequence_tuple: tuple[int, PulseSequence],
+        configs: dict[str, Config],
+        updates: dict,
     ) -> NDArray:
         """Evolve a pulse sequence on the quantum emulator.
 
         This method updates the sequence parameters, generates the time grid, constructs
         the time-dependent Hamiltonian, evolves the initial state with optional collapse
         operators, and returns the resulting measurement data.
         """
+        sequence_identifier, sequence = sequence_tuple
+
         sequence_ = update_sequence(sequence, updates)
         configs_ = update_configs(configs, updates)
         config = cast(HamiltonianConfig, configs_["hamiltonian"])
@@ -194,46 +248,49 @@ def _evolve(
         measurement_times[measurement_times < SAMPLING_INTERVAL] = SAMPLING_INTERVAL
         tlist_, index = np.unique(measurement_times, return_inverse=True)
 
-        results = self.engine.evolve(
+        results, simulation_configs = self.engine.evolve(
             hamiltonian=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]
-        coefficients = (
-            time_hamiltonian.coefficients if time_hamiltonian is not None else None
+
+        self._dump_simulation(
+            sequence_identifier,
+            hamiltonian,
+            time_hamiltonian,
+            states,
+            simulation_configs,
         )
-        return states, coefficients
+        return states
 
     def _pulse_hamiltonian(
         self, sequence: PulseSequence, configs: dict[str, Config]
-    ) -> OperatorEvolution | None:
+    ) -> OperatorEvolution:
         """Construct Hamiltonian time dependent term for qutip simulation."""
 
         # processed sampling rate; field `sampling_rate` of the `EmulatorController`
         # mimic a real hardware sampling rate, but it is insufficient for us to resolve
         # the oscillation and correctly solve the system evolution, hence we
         # set a nyquist frequency to define the timesteps in order to compute the solution
         times = tlist(sequence)
-        channels, raw_coefficients = [], []
-        for operator, waveforms in hamiltonians(
-            sequence, configs, self.engine, self.sampling_rate
-        ):
-            raw = channel_coefficients(
-                waveforms, sampling_rate=self.sampling_rate, times=times
+        channels: list[TimeDependentOperator] = [
+            TimeDependentOperator(
+                (
+                    operator,
+                    channel_coefficients(
+                        waveforms, sampling_rate=self.sampling_rate, times=times
+                    ),
+                )
             )
-            channels.append(operator)
-            raw_coefficients.append(raw)
-
-        return (
-            OperatorEvolution(
-                operators=channels, coefficients=np.stack(raw_coefficients), times=times
+            for operator, waveforms in hamiltonians(
+                sequence, configs, self.engine, self.sampling_rate
             )
-            if len(channels) > 0
-            else None
-        )
+        ]
+
+        return OperatorEvolution(operators=channels, times=times)
 
 
 def update_sequence(sequence: PulseSequence, updates: dict) -> PulseSequence:
@@ -256,6 +313,7 @@ def tlist(sequence: PulseSequence) -> NDArray:
     or Readout operation, it is excluded from the duration calculation.
     """
 
+    # TODO: maybe this can be a fragility in the case of 0 duration pulses.
     end = max(sequence.duration, SAMPLING_INTERVAL)
     return np.arange(0, end, SAMPLING_INTERVAL)
 
@@ -270,7 +328,9 @@ def hamiltonian(
 ) -> tuple[Operator, list[Modulated]]:
     n = hamiltonian.transmon_levels
     op = engine.expand(
-        config.operator(n=n, engine=engine), hamiltonian.dims, hilbert_space_index
+        op=config.operator(n=n, engine=engine),
+        targets=hilbert_space_index,
+        dims=hamiltonian.dims,
     )
     waveforms = (
         waveform(pulse, config, hamiltonian.qubits[hilbert_space_index], sampling_rate)
@@ -306,7 +366,7 @@ def channel_coefficients(
     waveforms: Iterable[Modulated],
     sampling_rate: int,
     times: NDArray,
-) -> BSpline:
+) -> NDArray:
     """
     Generate a B-spline interpolation of waveforms over a time evolution.
     This function processes a sequence of pulses, accumulating their waveforms
@@ -319,11 +379,16 @@ def channel_coefficients(
     cumulative_phase = 0
     cumulative_time = 0
     for pulse in waveforms:
-        next_pulse_time = cumulative_time + pulse.duration
-        pulse_times_idx = (times >= cumulative_time) & (times < next_pulse_time)
+        times_ = times - cumulative_time  # local times
+        # in this mask we take into account finite sampling rate with might mismatch with timestep of the emulator
+        # (i.e. float time durations and int sampling_rate)
+        pulse_times_idx = (times_ >= 0) & (
+            times_ < int(pulse.duration * sampling_rate) / sampling_rate
+        )
         times_samples = np.floor(
             (times[pulse_times_idx] - cumulative_time) * sampling_rate
         ).astype(int)
+
         # in case of virtual operations (such as VirtualZ or in general
         # zero-duration pulses), we apply the phase jump without
         # affecting the waveform
@@ -333,7 +398,6 @@ def channel_coefficients(
             )
 
         cumulative_phase += pulse.phase
-        cumulative_time = next_pulse_time
+        cumulative_time += pulse.duration
 
-    # return pulse_waveforms
     return pulse_waveforms