Refactoring QutipEmulator class

pulser-simulation/pulser_simulation/hamiltonian.py

@@ -25,10 +25,11 @@
 import numpy as np
 import qutip
 
+from pulser.backend.noise_model import NoiseModel
 from pulser.devices._device_datacls import BaseDevice
 from pulser.register.base_register import QubitId
 from pulser.sampler.samples import SequenceSamples, _PulseTargetSlot
-from pulser_simulation.simconfig import SimConfig
+from pulser_simulation.simconfig import SUPPORTED_NOISES, doppler_sigma
 
 
 class Hamiltonian:
@@ -50,7 +51,7 @@ def __init__(
         qdict: dict[QubitId, np.ndarray],
         device: BaseDevice,
         sampling_rate: float,
-        config: SimConfig,
+        config: NoiseModel,
     ) -> None:
         """Instantiates a Hamiltonian object."""
         self.samples_obj = samples_obj
@@ -60,7 +61,7 @@ def __init__(
 
         # Type hints for attributes defined outside of __init__
         self.basis_name: str
-        self._config: SimConfig
+        self._config: NoiseModel
         self.op_matrix: dict[str, qutip.Qobj]
         self.basis: dict[str, qutip.Qobj]
         self.dim: int
@@ -75,10 +76,10 @@ def __init__(
         self._qid_index = {qid: i for i, qid in enumerate(self._qdict)}
 
         # Compute sampling times
-        self.duration = self.samples_obj.max_duration
+        self._duration = self.samples_obj.max_duration
         self.sampling_times = self._adapt_to_sampling_rate(
             # Include extra time step for final instruction from samples:
-            np.arange(self.duration, dtype=np.double)
+            np.arange(self._duration, dtype=np.double)
             / 1000
         )
 
@@ -95,29 +96,28 @@ def _adapt_to_sampling_rate(self, full_array: np.ndarray) -> np.ndarray:
         indices = np.linspace(
             0,
             len(full_array) - 1,
-            int(self._sampling_rate * self.duration),
+            int(self._sampling_rate * self._duration),
             dtype=int,
         )
         return cast(np.ndarray, full_array[indices])
 
     @property
-    def config(self) -> SimConfig:
-        """The current configuration, as a SimConfig instance."""
+    def config(self) -> NoiseModel:
+        """The current configuration, as a NoiseModel instance."""
         return self._config
 
-    def _build_collapse_operators(self, prev_config: SimConfig) -> None:
+    def _build_collapse_operators(self, config: NoiseModel) -> None:
         kraus_ops = []
         self._collapse_ops = []
-        if "dephasing" in self.config.noise:
+        if "dephasing" in config.noise_types:
             if self.basis_name == "digital" or self.basis_name == "all":
                 # Go back to previous config
-                self.set_config(prev_config)
                 raise NotImplementedError(
                     "Cannot include dephasing noise in digital- or all-basis."
                 )
             # Probability of phase (Z) flip:
             # First order in prob
-            prob = self.config.dephasing_prob / 2
+            prob = config.dephasing_prob / 2
             n = self._size
             if prob > 0.1 and n > 1:
                 warnings.warn(
@@ -134,17 +134,16 @@ def _build_collapse_operators(self, prev_config: SimConfig) -> None:
             ]
             kraus_ops.append(k * qutip.sigmaz())
 
-        if "depolarizing" in self.config.noise:
+        if "depolarizing" in config.noise_types:
             if self.basis_name == "digital" or self.basis_name == "all":
                 # Go back to previous config
-                self.set_config(prev_config)
                 raise NotImplementedError(
                     "Cannot include depolarizing "
                     + "noise in digital- or all-basis."
                 )
             # Probability of error occurrence
 
-            prob = self.config.depolarizing_prob / 4
+            prob = config.depolarizing_prob / 4
             n = self._size
             if prob > 0.1 and n > 1:
                 warnings.warn(
@@ -163,20 +162,19 @@ def _build_collapse_operators(self, prev_config: SimConfig) -> None:
             kraus_ops.append(k * qutip.sigmay())
             kraus_ops.append(k * qutip.sigmaz())
 
-        if "eff_noise" in self.config.noise:
+        if "eff_noise" in config.noise_types:
             if self.basis_name == "digital" or self.basis_name == "all":
                 # Go back to previous config
-                self.set_config(prev_config)
                 raise NotImplementedError(
                     "Cannot include general "
                     + "noise in digital- or all-basis."
                 )
             # Probability distribution of error occurences
             n = self._size
-            m = len(self.config.eff_noise_opers)
+            m = len(config.eff_noise_opers)
             if n > 1:
                 for i in range(1, m):
-                    prob_i = self.config.eff_noise_probs[i]
+                    prob_i = config.eff_noise_probs[i]
                     if prob_i > 0.1:
                         warnings.warn(
                             "The effective noise model is a first-order"
@@ -186,16 +184,14 @@ def _build_collapse_operators(self, prev_config: SimConfig) -> None:
                         )
                         break
             # Deriving Kraus operators
-            prob_id = self.config.eff_noise_probs[0]
+            prob_id = config.eff_noise_probs[0]
             self._collapse_ops += [
                 np.sqrt(prob_id**n)
                 * qutip.tensor([self.op_matrix["I"] for _ in range(n)])
             ]
             for i in range(1, m):
-                k = np.sqrt(
-                    self.config.eff_noise_probs[i] * prob_id ** (n - 1)
-                )
-                k_op = k * self.config.eff_noise_opers[i]
+                k = np.sqrt(config.eff_noise_probs[i] * prob_id ** (n - 1))
+                k_op = k * config.eff_noise_opers[i]
                 kraus_ops.append(k_op)
 
         # Building collapse operators
@@ -205,67 +201,72 @@ def _build_collapse_operators(self, prev_config: SimConfig) -> None:
                 for qid in self._qid_index
             ]
 
-    def set_config(self, cfg: SimConfig) -> None:
+    def set_config(self, cfg: NoiseModel) -> None:
         """Sets current config to cfg and updates simulation parameters.
 
         Args:
             cfg: New configuration.
         """
-        if not isinstance(cfg, SimConfig):
-            raise ValueError(f"Object {cfg} is not a valid `SimConfig`.")
+        if not isinstance(cfg, NoiseModel):
+            raise ValueError(f"Object {cfg} is not a valid `NoiseModel`.")
         not_supported = (
-            set(cfg.noise) - cfg.supported_noises[self._interaction]
+            set(cfg.noise_types) - SUPPORTED_NOISES[self._interaction]
         )
         if not_supported:
             raise NotImplementedError(
                 f"Interaction mode '{self._interaction}' does not support "
                 f"simulation of noise types: {', '.join(not_supported)}."
             )
-        prev_config = self.config if hasattr(self, "_config") else SimConfig()
+        if not hasattr(self, "basis_name"):
+            self._build_basis_and_op_matrices()
+        self._build_collapse_operators(cfg)
         self._config = cfg
-        if not ("SPAM" in self.config.noise and self.config.eta > 0):
+        if not (
+            "SPAM" in self.config.noise_types
+            and self.config.state_prep_error > 0
+        ):
             self._bad_atoms = {qid: False for qid in self._qid_index}
-        if "doppler" not in self.config.noise:
+        if "doppler" not in self.config.noise_types:
             self._doppler_detune = {qid: 0.0 for qid in self._qid_index}
         # Noise, samples and Hamiltonian update routine
         self._construct_hamiltonian()
-        self._build_collapse_operators(prev_config)
 
-    def add_config(self, config: SimConfig) -> None:
+    def add_config(self, config: NoiseModel) -> None:
         """Updates the current configuration with parameters of another one.
 
         Mostly useful when dealing with multiple noise types in different
         configurations and wanting to merge these configurations together.
         Adds simulation parameters to noises that weren't available in the
-        former SimConfig. Noises specified in both SimConfigs will keep
+        former NoiseModel. Noises specified in both NoiseModels will keep
         former noise parameters.
 
         Args:
-            config: SimConfig to retrieve parameters from.
+            config: NoiseModel to retrieve parameters from.
         """
-        if not isinstance(config, SimConfig):
-            raise ValueError(f"Object {config} is not a valid `SimConfig`")
+        if not isinstance(config, NoiseModel):
+            raise ValueError(f"Object {config} is not a valid `NoiseModel`")
 
         not_supported = (
-            set(config.noise) - config.supported_noises[self._interaction]
+            set(config.noise_types) - SUPPORTED_NOISES[self._interaction]
         )
         if not_supported:
             raise NotImplementedError(
                 f"Interaction mode '{self._interaction}' does not support "
                 f"simulation of noise types: {', '.join(not_supported)}."
             )
 
-        old_noise_set = set(self.config.noise)
-        new_noise_set = old_noise_set.union(config.noise)
+        old_noise_set = set(self.config.noise_types)
+        new_noise_set = old_noise_set.union(config.noise_types)
         diff_noise_set = new_noise_set - old_noise_set
+        print(diff_noise_set)
         # Create temporary param_dict to add noise parameters:
         param_dict: dict[str, Any] = asdict(self._config)
         # Begin populating with added noise parameters:
-        param_dict["noise"] = tuple(new_noise_set)
+        param_dict["noise_types"] = tuple(new_noise_set)
         if "SPAM" in diff_noise_set:
-            param_dict["eta"] = config.eta
-            param_dict["epsilon"] = config.epsilon
-            param_dict["epsilon_prime"] = config.epsilon_prime
+            param_dict["state_prep_error"] = config.state_prep_error
+            param_dict["p_false_pos"] = config.p_false_pos
+            param_dict["p_false_neg"] = config.p_false_neg
         if "doppler" in diff_noise_set:
             param_dict["temperature"] = config.temperature
         if "amplitude" in diff_noise_set:
@@ -277,29 +278,23 @@ def add_config(self, config: SimConfig) -> None:
         if "eff_noise" in diff_noise_set:
             param_dict["eff_noise_opers"] = config.eff_noise_opers
             param_dict["eff_noise_probs"] = config.eff_noise_probs
-        param_dict["temperature"] *= 1.0e6
         # update runs:
         param_dict["runs"] = config.runs
         param_dict["samples_per_run"] = config.samples_per_run
 
         # set config with the new parameters:
-        self.set_config(SimConfig(**param_dict))
-
-    def show_config(self, solver_options: bool = False) -> None:
-        """Shows current configuration."""
-        print(self._config.__str__(solver_options))
-
-    def reset_config(self) -> None:
-        """Resets configuration to default."""
-        self.set_config(SimConfig())
+        self.set_config(NoiseModel(**param_dict))
 
     def _extract_samples(self) -> None:
         """Populates samples dictionary with every pulse in the sequence."""
         local_noises = True
-        if set(self.config.noise).issubset(
+        if set(self.config.noise_types).issubset(
             {"dephasing", "SPAM", "depolarizing", "eff_noise"}
         ):
-            local_noises = "SPAM" in self.config.noise and self.config.eta > 0
+            local_noises = (
+                "SPAM" in self.config.noise_types
+                and self.config.state_prep_error > 0
+            )
         samples = self.samples_obj.to_nested_dict(all_local=local_noises)
 
         def add_noise(
@@ -316,12 +311,12 @@ def add_noise(
                 0, np.random.normal(1.0, self.config.amp_sigma)
             )
             for qid in slot.targets:
-                if "doppler" in self.config.noise:
+                if "doppler" in self.config.noise_types:
                     noise_det = self._doppler_detune[qid]
                     samples_dict[qid]["det"][slot.ti : slot.tf] += noise_det
                 # Gaussian beam loss in amplitude for global pulses only
                 # Noise is drawn at random for each pulse
-                if "amplitude" in self.config.noise and is_global_pulse:
+                if "amplitude" in self.config.noise_types and is_global_pulse:
                     position = self._qdict[qid]
                     r = np.linalg.norm(position)
                     w0 = self.config.laser_waist
@@ -393,23 +388,28 @@ def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
                 for qubit in qubits:
                     k = self._qid_index[qubit]
                     op_list[k] = operator
-        return qutip.tensor(op_list)
+        return qutip.tensor(list(map(qutip.Qobj, op_list)))
 
     def _update_noise(self) -> None:
         """Updates noise random parameters.
 
         Used at the start of each run. If SPAM isn't in chosen noises, all
         atoms are set to be correctly prepared.
         """
-        if "SPAM" in self.config.noise and self.config.eta > 0:
+        if (
+            "SPAM" in self.config.noise_types
+            and self.config.state_prep_error > 0
+        ):
             dist = (
                 np.random.uniform(size=len(self._qid_index))
-                < self.config.spam_dict["eta"]
+                < self.config.state_prep_error
             )
             self._bad_atoms = dict(zip(self._qid_index, dist))
-        if "doppler" in self.config.noise:
+        if "doppler" in self.config.noise_types:
             detune = np.random.normal(
-                0, self.config.doppler_sigma, size=len(self._qid_index)
+                0,
+                doppler_sigma(self.config.temperature / 1e6),
+                size=len(self._qid_index),
             )
             self._doppler_detune = dict(zip(self._qid_index, detune))
 
@@ -457,8 +457,6 @@ def _construct_hamiltonian(self, update: bool = True) -> None:
         if update:
             self._update_noise()
         self._extract_samples()
-        if not hasattr(self, "basis_name"):
-            self._build_basis_and_op_matrices()
 
         def make_vdw_term(q1: QubitId, q2: QubitId) -> qutip.Qobj:
             """Construct the Van der Waals interaction Term.
@@ -604,7 +602,7 @@ def build_coeffs_ops(basis: str, addr: str) -> list[list]:
             ):
                 # Build an array of binary coefficients for the interaction
                 # term of unmasked qubits
-                coeff = np.ones(self.duration - 1)
+                coeff = np.ones(self._duration - 1)
                 coeff[0 : self.samples_obj._slm_mask.end] = 0
                 # Build the interaction term for unmasked qubits
                 qobj_list = [

pulser-simulation/pulser_simulation/simconfig.py

@@ -32,6 +32,27 @@
 
 T = TypeVar("T", bound="SimConfig")
 
+SUPPORTED_NOISES: dict = {
+    "ising": {
+        "dephasing",
+        "doppler",
+        "amplitude",
+        "SPAM",
+        "depolarizing",
+        "eff_noise",
+    },
+    "XY": {"SPAM"},
+}
+
+
+def doppler_sigma(temperature: float) -> float:
+    """Standard deviation for Doppler shifting due to thermal motion.
+
+    Arg:
+        temperature: The temperature in K.
+    """
+    return KEFF * sqrt(KB * temperature / MASS)
+
 
 @dataclass(frozen=True)
 class SimConfig:
@@ -160,7 +181,7 @@ def spam_dict(self) -> dict[str, float]:
     @property
     def doppler_sigma(self) -> float:
         """Standard deviation for Doppler shifting due to thermal motion."""
-        return KEFF * sqrt(KB * self.temperature / MASS)
+        return doppler_sigma(self.temperature)
 
     def __str__(self, solver_options: bool = False) -> str:
         lines = [
@@ -220,14 +241,4 @@ def _check_eff_noise_opers_type(self) -> None:
     @property
     def supported_noises(self) -> dict:
         """Return the noises implemented on pulser."""
-        return {
-            "ising": {
-                "dephasing",
-                "doppler",
-                "amplitude",
-                "SPAM",
-                "depolarizing",
-                "eff_noise",
-            },
-            "XY": {"SPAM"},
-        }
+        return SUPPORTED_NOISES