Add leakage noise in NoiseModel

pulser-core/pulser/json/abstract_repr/deserializer.py

@@ -435,10 +435,12 @@ def convert_complex(obj: Any) -> Any:
         eff_noise_opers.append(convert_complex(oper))
 
     noise_types = noise_model_obj.pop("noise_types")
+    with_leakage = "leakage" in noise_types
     noise_model = pulser.NoiseModel(
         **noise_model_obj,
         eff_noise_rates=tuple(eff_noise_rates),
         eff_noise_opers=tuple(eff_noise_opers),
+        with_leakage=with_leakage,
     )
     assert set(noise_model.noise_types) == set(noise_types)
     return noise_model

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,6 +78,8 @@
     "amp_sigma",
 }
 
+_BOOLEAN = {"with_leakage"}
+
 _LEGACY_DEFAULTS = {
     "runs": 15,
     "samples_per_run": 5,
@@ -98,6 +102,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 +165,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 +184,7 @@ class NoiseModel:
     depolarizing_rate: float
     eff_noise_rates: tuple[float, ...]
     eff_noise_opers: tuple[ArrayLike, ...]
+    with_leakage: bool
 
     def __init__(
         self,
@@ -191,6 +203,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 +227,8 @@ 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,
         )
-
         if noise_types is not None:
             with warnings.catch_warnings():
                 warnings.simplefilter("always")
@@ -231,21 +244,26 @@ def to_tuple(obj: tuple) -> tuple:
                 )
             self._check_noise_types(noise_types)
             for nt_ in noise_types:
+                if nt_ == "leakage":
+                    raise ValueError(
+                        "'leakage' cannot be explicitely defined in the noise"
+                        " types. Set 'with_leakage' to True instead."
+                    )
                 for p_ in _NOISE_TYPE_PARAMS[nt_]:
                     # Replace undefined relevant params by the legacy default
                     if param_vals[p_] is None:
                         param_vals[p_] = _LEGACY_DEFAULTS[p_]
-
         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
         }
-
+        self._check_leakage_noise(true_noise_types)
         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=cast(bool, param_vals["with_leakage"]),
         )
 
         # Get rid of unnecessary None's
@@ -277,7 +295,7 @@ def to_tuple(obj: tuple) -> tuple:
         relevant_param_vals = {
             p: param_vals[p]
             for p in param_vals
-            if param_vals[p] is not None or (p in relevant_params)
+            if param_vals[p] is not None or p in relevant_params
         }
         self._validate_parameters(relevant_param_vals)
 
@@ -314,6 +332,17 @@ def _find_relevant_params(
             relevant_params.discard("laser_waist")
         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
+        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 +358,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(
@@ -355,6 +385,11 @@ def _check_eff_noise(
             raise ValueError("The provided rates must be greater than 0.")
 
         # Check the validity of operators
+        min_shape = 2 if not with_leakage else 3
+        possible_shapes = [
+            (min_shape, min_shape),
+            (min_shape + 1, min_shape + 1),
+        ]
         for op in eff_noise_opers:
             # type checking
             try:
@@ -366,9 +401,17 @@ 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}."
+            # TODO: Modify when effective noise can be provided for qutrit
+            if operator.shape != possible_shapes[0]:
+                err_type = (
+                    NotImplementedError
+                    if operator.shape in possible_shapes
+                    else ValueError
+                )
+                raise err_type(
+                    f"With{'' if with_leakage else 'out'} leakage, operator's "
+                    f"shape must be {possible_shapes[0]}, "
+                    f"not {operator.shape}."
                 )
 
     @staticmethod
@@ -388,11 +431,15 @@ 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}.")
 
     def _to_abstract_repr(self) -> dict[str, Any]:
         all_fields = asdict(self)
+        all_fields.pop("with_leakage")
         eff_noise_rates = all_fields.pop("eff_noise_rates")
         eff_noise_opers = all_fields.pop("eff_noise_opers")
         all_fields["eff_noise"] = list(zip(eff_noise_rates, eff_noise_opers))

pulser-simulation/pulser_simulation/simconfig.py

@@ -78,6 +78,9 @@ class SimConfig:
             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 (which must be defined).
             - "relaxation": Relaxation from the Rydberg to the ground state.
             - "dephasing": Random phase (Z) flip.
             - "depolarizing": Quantum noise where the state (rho) is
@@ -140,6 +143,7 @@ def from_noise_model(cls: Type[T], noise_model: NoiseModel) -> T:
             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:
@@ -176,6 +180,11 @@ def __post_init__(self) -> None:
             {f.name: getattr(self, f.name) for f in fields(self)}
         )
 
+    @property
+    def with_leakage(self) -> bool:
+        """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 +262,7 @@ def _check_eff_noise(self) -> None:
             self.eff_noise_rates,
             self.eff_noise_opers,
             "eff_noise" in self.noise,
+            self.with_leakage,
         )
 
     @property

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):

tests/test_noise_model.py

@@ -176,8 +176,29 @@ def matrices(self):
         matrices["Zh"] = 0.5 * np.array([[1, 0], [0, -1]])
         matrices["ket"] = np.array([[1.0], [2.0]])
         matrices["I3"] = np.eye(3)
+        matrices["I4"] = np.eye(4)
         return matrices
 
+    @pytest.mark.parametrize("value", [False, True])
+    def test_init_bool_like(self, value, matrices):
+        noise_model = NoiseModel(
+            eff_noise_rates=[0.1],
+            eff_noise_opers=[matrices["I3"] if value else matrices["I"]],
+            with_leakage=value,
+        )
+        assert noise_model.with_leakage == value
+
+    @pytest.mark.parametrize("value", [0, 1, 0.1])
+    def test_wrong_init_bool_like(self, value, matrices):
+        with pytest.raises(
+            ValueError, match=f"'with_leakage' must be a boolean, not {value}"
+        ):
+            NoiseModel(
+                eff_noise_rates=[0.1],
+                eff_noise_opers=[matrices["I3"] if value else matrices["I"]],
+                with_leakage=value,
+            )
+
     def test_eff_noise_rates(self, matrices):
         with pytest.raises(
             ValueError, match="The provided rates must be greater than 0."
@@ -207,11 +228,41 @@ def test_eff_noise_opers(self, matrices):
                 eff_noise_opers=[2.0],
                 eff_noise_rates=[1.0],
             )
-        with pytest.raises(NotImplementedError, match="Operator's shape"):
+        with pytest.raises(ValueError, match="With leakage, operator's shape"):
+            NoiseModel(
+                eff_noise_opers=[matrices["I"]],
+                eff_noise_rates=[1.0],
+                with_leakage=True,
+            )
+        with pytest.raises(
+            NotImplementedError, match="With leakage, operator's shape"
+        ):
+            NoiseModel(
+                eff_noise_opers=[matrices["I4"]],
+                eff_noise_rates=[1.0],
+                with_leakage=True,
+            )
+        with pytest.raises(
+            NotImplementedError, match="Without leakage, operator's shape"
+        ):
             NoiseModel(
                 eff_noise_opers=[matrices["I3"]],
                 eff_noise_rates=[1.0],
             )
+        with pytest.raises(
+            ValueError, match="Without leakage, operator's shape"
+        ):
+            NoiseModel(
+                eff_noise_opers=[matrices["I4"]],
+                eff_noise_rates=[1.0],
+            )
+
+    @pytest.mark.parametrize("param", ["dephasing_rate", "depolarizing_rate"])
+    def test_leakage(self, param):
+        with pytest.raises(
+            ValueError, match="At least one effective noise operator"
+        ):
+            NoiseModel(with_leakage=True)
 
     def test_eq(self, matrices):
         final_fields = dict(
@@ -258,17 +309,17 @@ def test_relevant_params(self):
         ) == {"amp_sigma", "laser_waist", "runs", "samples_per_run"}
 
         assert NoiseModel._find_relevant_params(
-            {"dephasing"}, 0.0, 0.0, None
-        ) == {"dephasing_rate", "hyperfine_dephasing_rate"}
+            {"dephasing", "leakage"}, 0.0, 0.0, None
+        ) == {"dephasing_rate", "hyperfine_dephasing_rate", "with_leakage"}
         assert NoiseModel._find_relevant_params(
-            {"relaxation"}, 0.0, 0.0, None
-        ) == {"relaxation_rate"}
+            {"relaxation", "leakage"}, 0.0, 0.0, None
+        ) == {"relaxation_rate", "with_leakage"}
         assert NoiseModel._find_relevant_params(
-            {"depolarizing"}, 0.0, 0.0, None
-        ) == {"depolarizing_rate"}
+            {"depolarizing", "leakage"}, 0.0, 0.0, None
+        ) == {"depolarizing_rate", "with_leakage"}
         assert NoiseModel._find_relevant_params(
-            {"eff_noise"}, 0.0, 0.0, None
-        ) == {"eff_noise_rates", "eff_noise_opers"}
+            {"eff_noise", "leakage"}, 0.0, 0.0, None
+        ) == {"eff_noise_rates", "eff_noise_opers", "with_leakage"}
 
     def test_repr(self):
         assert repr(NoiseModel()) == "NoiseModel(noise_types=())"
@@ -293,6 +344,20 @@ def test_repr(self):
             == "NoiseModel(noise_types=('amplitude',), "
             "laser_waist=100.0, amp_sigma=0.0)"
         )
+        assert (
+            repr(
+                NoiseModel(
+                    hyperfine_dephasing_rate=0.2,
+                    eff_noise_opers=[[[1, 0, 0], [0, 1, 0], [0, 0, 1]]],
+                    eff_noise_rates=[0.1],
+                    with_leakage=True,
+                )
+            )
+            == "NoiseModel(noise_types=('dephasing', 'eff_noise', 'leakage'), "
+            "dephasing_rate=0.0, hyperfine_dephasing_rate=0.2, "
+            "eff_noise_rates=(0.1,), eff_noise_opers=(((1, 0, 0), (0, 1, 0), "
+            "(0, 0, 1)),), with_leakage=True)"
+        )
 
 
 class TestLegacyNoiseModel:
@@ -350,6 +415,16 @@ def test_legacy_init(self, noise_type):
         # Check that the parameter is not overwritten by the default
         assert getattr(noise_model, non_zero_param) == 1
 
+        with pytest.raises(
+            ValueError,
+            match="'leakage' cannot be explicitely defined in the noise",
+        ):
+            with pytest.warns(
+                DeprecationWarning,
+                match="The explicit definition of noise types is deprecated",
+            ):
+                NoiseModel(noise_types=("leakage",))
+
         relevant_params = NoiseModel._find_relevant_params(
             {noise_type},
             # These values don't matter, they just have to be > 0