Add leakage noise in NoiseModel

pulser-core/pulser/json/abstract_repr/schemas/noise-schema.json

@@ -103,6 +103,9 @@
         },
         "temperature": {
           "type": "number"
+        },
+        "with_leakage": {
+          "type": "boolean"
         }
       },
       "required": [
@@ -125,13 +128,13 @@
     },
     "NoiseType": {
       "enum": [
+        "leakage",
         "doppler",
         "amplitude",
         "SPAM",
         "relaxation",
         "dephasing",
         "depolarizing",
-        "leakage",
         "eff_noise"
       ],
       "type": "string"

pulser-core/pulser/noise_model.py

@@ -30,6 +30,7 @@
 __all__ = ["NoiseModel"]
 
 NoiseTypes = Literal[
+    "leakage",
     "doppler",
     "amplitude",
     "SPAM",
@@ -40,6 +41,7 @@
 ]
 
 _NOISE_TYPE_PARAMS: dict[NoiseTypes, tuple[str, ...]] = {
+    "leakage": ("with_leakage",),
     "doppler": ("temperature",),
     "amplitude": ("laser_waist", "amp_sigma"),
     "SPAM": ("p_false_pos", "p_false_neg", "state_prep_error"),
@@ -76,7 +78,10 @@
     "amp_sigma",
 }
 
+_BOOLEAN = {"with_leakage"}
+
 _LEGACY_DEFAULTS = {
+    "with_leakage": True,
     "runs": 15,
     "samples_per_run": 5,
     "state_prep_error": 0.005,
@@ -98,6 +103,11 @@ class NoiseModel:
 
     Supported noise types:
 
+    - "leakage": Adds an error state 'x' to the computational
+        basis, that can interact with the other states via an
+        effective noise channel. Must be defined with an effective
+        noise channel, but is incompatible with dephasing and
+        depolarizing noise channels.
     - **relaxation**: Noise due to a decay from the Rydberg to
       the ground state (parametrized by ``relaxation_rate``),
       commonly characterized experimentally by the T1 time.
@@ -156,6 +166,8 @@ class NoiseModel:
         eff_noise_rates: The rate associated to each effective noise operator
             (in 1/µs).
         eff_noise_opers: The operators for the effective noise model.
+        with_leakage: Whether or not to include an error state in the
+            computations (default to False).
     """
 
     noise_types: tuple[NoiseTypes, ...]
@@ -173,6 +185,7 @@ class NoiseModel:
     depolarizing_rate: float
     eff_noise_rates: tuple[float, ...]
     eff_noise_opers: tuple[ArrayLike, ...]
+    with_leakage: bool
 
     def __init__(
         self,
@@ -191,6 +204,7 @@ def __init__(
         depolarizing_rate: float | None = None,
         eff_noise_rates: tuple[float, ...] = (),
         eff_noise_opers: tuple[ArrayLike, ...] = (),
+        with_leakage: bool = False,
     ) -> None:
         """Initializes a noise model."""
 
@@ -214,8 +228,9 @@ def to_tuple(obj: tuple) -> tuple:
             depolarizing_rate=depolarizing_rate,
             eff_noise_rates=to_tuple(eff_noise_rates),
             eff_noise_opers=to_tuple(eff_noise_opers),
+            with_leakage=with_leakage,
         )
-
+        print("Initial param_vals", param_vals)
         if noise_types is not None:
             with warnings.catch_warnings():
                 warnings.simplefilter("always")
@@ -235,17 +250,22 @@ def to_tuple(obj: tuple) -> tuple:
                     # Replace undefined relevant params by the legacy default
                     if param_vals[p_] is None:
                         param_vals[p_] = _LEGACY_DEFAULTS[p_]
-
+        print("param_vals post noise_types", param_vals)
         true_noise_types: set[NoiseTypes] = {
             _PARAM_TO_NOISE_TYPE[p_]
             for p_ in param_vals
             if param_vals[p_] and p_ in _PARAM_TO_NOISE_TYPE
         }
-
+        print("With leakage", param_vals["with_leakage"])
+        self._check_leakage_noise(
+            true_noise_types if noise_types is None else noise_types
+        )
+        print("With leakage", param_vals["with_leakage"])
         self._check_eff_noise(
             cast(tuple, param_vals["eff_noise_rates"]),
             cast(tuple, param_vals["eff_noise_opers"]),
             "eff_noise" in (noise_types or true_noise_types),
+            with_leakage=False,
         )
 
         # Get rid of unnecessary None's
@@ -257,6 +277,7 @@ def to_tuple(obj: tuple) -> tuple:
             cast(float, param_vals["state_prep_error"]),
             cast(float, param_vals["amp_sigma"]),
             cast(Union[float, None], param_vals["laser_waist"]),
+            cast(bool, param_vals["with_leakage"]),
         )
 
         if noise_types is not None:
@@ -299,6 +320,7 @@ def _find_relevant_params(
         state_prep_error: float,
         amp_sigma: float,
         laser_waist: float | None,
+        with_leakage: bool,
     ) -> set[str]:
         relevant_params: set[str] = set()
         for nt_ in noise_types:
@@ -312,8 +334,29 @@ def _find_relevant_params(
         # Disregard laser_waist when not defined
         if laser_waist is None:
             relevant_params.discard("laser_waist")
+        # Diregard leakage if with_leakage is False
+        if not with_leakage:
+            relevant_params.discard("with_leakage")
         return relevant_params
 
+    @staticmethod
+    def _check_leakage_noise(noise_types: Collection[NoiseTypes]) -> None:
+        # Can't define "dephasing", "depolarizing" with "leakage"
+        if "leakage" not in noise_types:
+            return
+        # TODO: Implement the depolarizing and dephasing operations with
+        # projectors will stop raising an error.
+        if "dephasing" in noise_types or "depolarizing" in noise_types:
+            raise NotImplementedError(
+                "Dephasing and depolarizing channels can't be defined "
+                "with a leakage noise."
+            )
+        if "eff_noise" not in noise_types:
+            raise ValueError(
+                "At least one effective noise operator must be defined to"
+                " simulate leakage."
+            )
+
     @staticmethod
     def _check_noise_types(noise_types: Sequence[NoiseTypes]) -> None:
         for noise_type in noise_types:
@@ -329,6 +372,7 @@ def _check_eff_noise(
         eff_noise_rates: Sequence[float],
         eff_noise_opers: Sequence[ArrayLike],
         check_contents: bool,
+        with_leakage: bool,
     ) -> None:
         if len(eff_noise_opers) != len(eff_noise_rates):
             raise ValueError(
@@ -353,6 +397,12 @@ def _check_eff_noise(
 
         if np.any(np.array(eff_noise_rates) < 0):
             raise ValueError("The provided rates must be greater than 0.")
+        print(with_leakage)
+        min_shape = 2 if not with_leakage else 3
+        possible_shapes = [
+            (min_shape, min_shape),
+            (min_shape + 1, min_shape + 1),
+        ]
 
         # Check the validity of operators
         for op in eff_noise_opers:
@@ -366,9 +416,11 @@ def _check_eff_noise(
             if operator.ndim != 2:
                 raise ValueError(f"Operator '{op!r}' is not a 2D array.")
 
-            if operator.shape != (2, 2):
-                raise NotImplementedError(
-                    f"Operator's shape must be (2,2) not {operator.shape}."
+            if operator.shape not in possible_shapes:
+                raise ValueError(
+                    f"With{'' if with_leakage else 'out'} leakage, operator's "
+                    f"shape must be {possible_shapes[0]} or "
+                    f"{possible_shapes[1]}, not {operator.shape}."
                 )
 
     @staticmethod
@@ -388,6 +440,9 @@ def _validate_parameters(param_vals: dict[str, Any]) -> None:
                     "greater than or equal to zero and smaller than "
                     "or equal to one"
                 )
+            elif param in _BOOLEAN:
+                is_valid = isinstance(value, bool)
+                comp = "a boolean"
             if not is_valid:
                 raise ValueError(f"'{param}' must be {comp}, not {value}.")
 
@@ -404,6 +459,7 @@ def __repr__(self) -> str:
             self.state_prep_error,
             self.amp_sigma,
             self.laser_waist,
+            self.with_leakage,
         )
         relevant_params.add("noise_types")
         params_list = []

pulser-simulation/pulser_simulation/simconfig.py

@@ -77,7 +77,9 @@ class SimConfig:
         noise: Types of noises to be used in the
             simulation. You may specify just one, or a tuple of the allowed
             noise types:
-
+            - "leakage": Adds an error state 'x' to the computational
+                basis, that can interact with the other states via an
+                effective noise channel.
             - "relaxation": Relaxation from the Rydberg to the ground state.
             - "dephasing": Random phase (Z) flip.
             - "depolarizing": Quantum noise where the state (rho) is
@@ -135,11 +137,13 @@ def from_noise_model(cls: Type[T], noise_model: NoiseModel) -> T:
             noise_model.state_prep_error,
             noise_model.amp_sigma,
             noise_model.laser_waist,
+            noise_model.with_leakage,
         )
         for param in relevant_params:
             kwargs[_DIFF_NOISE_PARAMS.get(param, param)] = getattr(
                 noise_model, param
             )
+        kwargs.pop("with_leakage", None)
         return cls(**kwargs)
 
     def to_noise_model(self) -> NoiseModel:
@@ -149,6 +153,7 @@ def to_noise_model(self) -> NoiseModel:
             self.eta,
             self.amp_sigma,
             self.laser_waist,
+            self.with_leakage,
         )
         kwargs = {}
         for param in relevant_params:
@@ -176,6 +181,11 @@ def __post_init__(self) -> None:
             {f.name: getattr(self, f.name) for f in fields(self)}
         )
 
+    @property
+    def with_leakage(self):
+        """Whether or not 'leakage' is included in the noise types."""
+        return "leakage" in self.noise
+
     @property
     def spam_dict(self) -> dict[str, float]:
         """A dictionary combining the SPAM error parameters."""
@@ -253,6 +263,7 @@ def _check_eff_noise(self) -> None:
             self.eff_noise_rates,
             self.eff_noise_opers,
             "eff_noise" in self.noise,
+            self.with_leakage,
         )
 
     @property

pulser-simulation/pulser_simulation/simulation.py

@@ -255,6 +255,7 @@ def add_config(self, config: SimConfig) -> None:
             noise_model.state_prep_error,
             noise_model.amp_sigma,
             noise_model.laser_waist,
+            noise_model.with_leakage,
         )
         for param in relevant_params:
             param_dict[param] = getattr(noise_model, param)

tests/test_abstract_repr.py

@@ -181,6 +181,11 @@ def test_register(reg: Register | Register3D):
             eff_noise_rates=(0.1,),
             eff_noise_opers=(((0, -1j), (1j, 0)),),
         ),
+        NoiseModel(
+            eff_noise_rates=(0.1,),
+            eff_noise_opers=(((0, -1j, 0), (1j, 0, 0), (0, 0, 1)),),
+            with_leakage=True,
+        ),
     ],
 )
 def test_noise_model(noise_model: NoiseModel):