Add leakage (#720)

* Enable definition of leakage with all basis

* Validate test results

* Test simulation with leakage

* fixing set_config

* Fix failing tests

* Fix set_config, add_config, coverage of qutip_result

* Fix NoisyResults, Improve test_simresults

* Fixing nits

* Testing projections

* Fix set_config, add_config documentation

* Extend set of noise to fix tests

* print deletion

pulser-core/pulser/channels/base_channel.py

@@ -37,7 +37,7 @@
 OPTIONAL_ABSTR_CH_FIELDS = ("min_avg_amp",)
 
 # States ranked in decreasing order of their associated eigenenergy
-States = Literal["u", "d", "r", "g", "h"]  # TODO: add "x" for leakage
+States = Literal["u", "d", "r", "g", "h", "x"]
 
 STATES_RANK = get_args(States)
 
@@ -138,6 +138,9 @@ def eigenstates(self) -> list[States]:
             * - Hyperfine state
               - :math:`|h\rangle`
               - ``"h"``
+            * - Error state
+              - :math:`|x\rangle`
+              - ``"x"``
         """
         return EIGENSTATES[self.basis]
 

pulser-core/pulser/noise_model.py

@@ -401,16 +401,8 @@ def _check_eff_noise(
             if operator.ndim != 2:
                 raise ValueError(f"Operator '{op!r}' is not a 2D array.")
 
-            # TODO: Modify when effective noise can be provided for leakage
-            if operator.shape != possible_shapes[0] and (
-                with_leakage or operator.shape != possible_shapes[1]
-            ):
-                err_type = (
-                    NotImplementedError
-                    if operator.shape in possible_shapes
-                    else ValueError
-                )
-                raise err_type(
+            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]}, "
                     f"not {operator.shape}."

pulser-simulation/pulser_simulation/hamiltonian.py

@@ -23,7 +23,7 @@
 import numpy as np
 import qutip
 
-from pulser.channels.base_channel import STATES_RANK
+from pulser.channels.base_channel import STATES_RANK, States
 from pulser.devices._device_datacls import BaseDevice
 from pulser.noise_model import NoiseModel
 from pulser.register.base_register import QubitId
@@ -62,7 +62,7 @@ def __init__(
         self.basis_name: str
         self._config: NoiseModel
         self.op_matrix: dict[str, qutip.Qobj]
-        self.basis: dict[str, qutip.Qobj]
+        self.basis: dict[States, qutip.Qobj]
         self.dim: int
         self._bad_atoms: dict[Union[str, int], bool] = {}
         self._doppler_detune: dict[Union[str, int], float] = {}
@@ -83,9 +83,6 @@ def __init__(
         )
 
         # Stores the qutip operators used in building the Hamiltonian
-        self.operators: dict[str, defaultdict[str, dict]] = {
-            addr: defaultdict(dict) for addr in ["Global", "Local"]
-        }
         self._collapse_ops: list[qutip.Qobj] = []
 
         self.set_config(config)
@@ -105,14 +102,13 @@ def config(self) -> NoiseModel:
         """The current configuration, as a NoiseModel instance."""
         return self._config
 
-    def _build_collapse_operators(self, config: NoiseModel) -> None:
-        def basis_check(noise_type: str) -> None:
-            """Checks if the basis allows for the use of noise."""
-            if self.basis_name == "all":
-                # Go back to previous config
-                raise NotImplementedError(
-                    f"Cannot include {noise_type} noise in all-basis."
-                )
+    def _build_collapse_operators(
+        self,
+        config: NoiseModel,
+        basis_name: str,
+        eigenbasis: list[States],
+        op_matrix: dict[str, qutip.Qobj],
+    ) -> None:
 
         local_collapse_ops = []
         if "dephasing" in config.noise_types:
@@ -121,33 +117,35 @@ def basis_check(noise_type: str) -> None:
                 "r": config.dephasing_rate,
                 "h": config.hyperfine_dephasing_rate,
             }
-            for state in self.eigenbasis:
+            for state in eigenbasis:
                 if state in dephasing_rates:
                     coeff = np.sqrt(2 * dephasing_rates[state])
-                    op = self.op_matrix[f"sigma_{state}{state}"]
+                    op = op_matrix[f"sigma_{state}{state}"]
                     local_collapse_ops.append(coeff * op)
 
         if "relaxation" in config.noise_types:
             coeff = np.sqrt(config.relaxation_rate)
             try:
-                local_collapse_ops.append(coeff * self.op_matrix["sigma_gr"])
+                local_collapse_ops.append(coeff * op_matrix["sigma_gr"])
             except KeyError:
                 raise ValueError(
                     "'relaxation' noise requires addressing of the"
                     " 'ground-rydberg' basis."
                 )
 
         if "depolarizing" in config.noise_types:
-            basis_check("depolarizing")
+            if "all" in basis_name:
+                # Go back to previous config
+                raise NotImplementedError(
+                    "Cannot include depolarizing noise in all-basis."
+                )
             # NOTE: These operators only make sense when basis != "all"
-            b, a = self.eigenbasis[:2]
+            b, a = eigenbasis[:2]
             pauli_2d = {
-                "x": self.op_matrix[f"sigma_{a}{b}"]
-                + self.op_matrix[f"sigma_{b}{a}"],
-                "y": 1j * self.op_matrix[f"sigma_{a}{b}"]
-                - 1j * self.op_matrix[f"sigma_{b}{a}"],
-                "z": self.op_matrix[f"sigma_{b}{b}"]
-                - self.op_matrix[f"sigma_{a}{a}"],
+                "x": op_matrix[f"sigma_{a}{b}"] + op_matrix[f"sigma_{b}{a}"],
+                "y": 1j * op_matrix[f"sigma_{a}{b}"]
+                - 1j * op_matrix[f"sigma_{b}{a}"],
+                "z": op_matrix[f"sigma_{b}{b}"] - op_matrix[f"sigma_{a}{a}"],
             }
             coeff = np.sqrt(config.depolarizing_rate / 4)
             local_collapse_ops.append(coeff * pauli_2d["x"])
@@ -157,7 +155,7 @@ def basis_check(noise_type: str) -> None:
         if "eff_noise" in config.noise_types:
             for id, rate in enumerate(config.eff_noise_rates):
                 op = np.array(config.eff_noise_opers[id])
-                basis_dim = len(self.eigenbasis)
+                basis_dim = len(eigenbasis)
                 op_shape = (basis_dim, basis_dim)
                 if op.shape != op_shape:
                     raise ValueError(
@@ -169,7 +167,7 @@ def basis_check(noise_type: str) -> None:
         self._collapse_ops = []
         for operator in local_collapse_ops:
             self._collapse_ops += [
-                self.build_operator([(operator, [qid])])
+                self._build_operator([(operator, [qid])], op_matrix)
                 for qid in self._qid_index
             ]
 
@@ -189,9 +187,28 @@ def set_config(self, cfg: NoiseModel) -> None:
                 f"Interaction mode '{self._interaction}' does not support "
                 f"simulation of noise types: {', '.join(not_supported)}."
             )
-        if not hasattr(self, "basis_name"):
-            self._build_basis_and_op_matrices()
-        self._build_collapse_operators(cfg)
+        if not hasattr(self, "_config") or (
+            hasattr(self, "_config")
+            and self.config.with_leakage != cfg.with_leakage
+        ):
+            basis_name = self._get_basis_name(cfg.with_leakage)
+            eigenbasis = self._get_eigenbasis(cfg.with_leakage)
+            basis, op_matrix = self._get_basis_op_matrices(eigenbasis)
+            self._build_collapse_operators(
+                cfg, basis_name, eigenbasis, op_matrix
+            )
+            self.basis_name = basis_name
+            self.eigenbasis = eigenbasis
+            self.basis = basis
+            self.op_matrix = op_matrix
+            self.dim = len(eigenbasis)
+            self.operators: dict[str, defaultdict[str, dict]] = {
+                addr: defaultdict(dict) for addr in ["Global", "Local"]
+            }
+        else:
+            self._build_collapse_operators(
+                cfg, self.basis_name, self.eigenbasis, self.op_matrix
+            )
         self._config = cfg
         if not (
             "SPAM" in self.config.noise_types
@@ -207,7 +224,14 @@ def _extract_samples(self) -> None:
         """Populates samples dictionary with every pulse in the sequence."""
         local_noises = True
         if set(self.config.noise_types).issubset(
-            {"dephasing", "relaxation", "SPAM", "depolarizing", "eff_noise"}
+            {
+                "dephasing",
+                "relaxation",
+                "SPAM",
+                "depolarizing",
+                "eff_noise",
+                "leakage",
+            }
         ):
             local_noises = (
                 "SPAM" in self.config.noise_types
@@ -259,7 +283,9 @@ def add_noise(
                             samples["Local"][basis][qid][qty] = 0.0
         self.samples = samples
 
-    def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
+    def _build_operator(
+        self, operations: Union[list, tuple], op_matrix: dict[str, qutip.Qobj]
+    ) -> qutip.Qobj:
         """Creates an operator with non-trivial actions on some qubits.
 
         Takes as argument a list of tuples ``[(operator_1, qubits_1),
@@ -281,15 +307,15 @@ def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
         Returns:
             The final operator.
         """
-        op_list = [self.op_matrix["I"] for j in range(self._size)]
+        op_list = [op_matrix["I"] for j in range(self._size)]
 
         if not isinstance(operations, list):
             operations = [operations]
 
         for operator, qubits in operations:
             if qubits == "global":
                 return sum(
-                    self.build_operator([(operator, [q_id])])
+                    self._build_operator([(operator, [q_id])], op_matrix)
                     for q_id in self._qdict
                 )
             else:
@@ -311,6 +337,30 @@ def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
                     op_list[k] = operator
         return qutip.tensor(list(map(qutip.Qobj, op_list)))
 
+    def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
+        """Creates an operator with non-trivial actions on some qubits.
+
+        Takes as argument a list of tuples ``[(operator_1, qubits_1),
+        (operator_2, qubits_2)...]``. Returns the operator given by the tensor
+        product of {``operator_i`` applied on ``qubits_i``} and Id on the rest.
+        ``(operator, 'global')`` returns the sum for all ``j`` of operator
+        applied at ``qubit_j`` and identity elsewhere.
+
+        Example for 4 qubits: ``[(Z, [1, 2]), (Y, [3])]`` returns `ZZYI`
+        and ``[(X, 'global')]`` returns `XIII + IXII + IIXI + IIIX`
+
+        Args:
+            operations: List of tuples `(operator, qubits)`.
+                `operator` can be a ``qutip.Quobj`` or a string key for
+                ``self.op_matrix``. `qubits` is the list on which operator
+                will be applied. The qubits can be passed as their
+                index or their label in the register.
+
+        Returns:
+            The final operator.
+        """
+        return self._build_operator(operations, self.op_matrix)
+
     def _update_noise(self) -> None:
         """Updates noise random parameters.
 
@@ -333,36 +383,39 @@ def _update_noise(self) -> None:
             )
             self._doppler_detune = dict(zip(self._qid_index, detune))
 
-    def _build_basis_and_op_matrices(self) -> None:
-        """Determine dimension, basis and projector operators."""
+    def _get_basis_name(self, with_leakage: bool) -> str:
         if len(self.samples_obj.used_bases) == 0:
             if self.samples_obj._in_xy:
-                self.basis_name = "XY"
+                basis_name = "XY"
             else:
-                self.basis_name = "ground-rydberg"
+                basis_name = "ground-rydberg"
         elif len(self.samples_obj.used_bases) == 1:
-            self.basis_name = list(self.samples_obj.used_bases)[0]
+            basis_name = list(self.samples_obj.used_bases)[0]
         else:
-            self.basis_name = "all"  # All three rydberg states
-        eigenbasis = self.samples_obj.eigenbasis
-
-        # TODO: Add leakage
-
-        self.eigenbasis = [
-            state for state in STATES_RANK if state in eigenbasis
-        ]
+            basis_name = "all"  # All three rydberg states
+        if with_leakage:
+            basis_name += "_with_error"
+        return basis_name
 
-        self.dim = len(self.eigenbasis)
-        self.basis = {
-            b: qutip.basis(self.dim, i) for i, b in enumerate(self.eigenbasis)
-        }
-        self.op_matrix = {"I": qutip.qeye(self.dim)}
-        for proj0 in self.eigenbasis:
-            for proj1 in self.eigenbasis:
+    def _get_eigenbasis(self, with_leakage: bool) -> list[States]:
+        eigenbasis = self.samples_obj.eigenbasis
+        if with_leakage:
+            eigenbasis.append("x")
+        return [state for state in STATES_RANK if state in eigenbasis]
+
+    @staticmethod
+    def _get_basis_op_matrices(
+        eigenbasis: list[States],
+    ) -> tuple[dict[States, qutip.Qobj], dict[str, qutip.Qobj]]:
+        """Determine basis and projector operators."""
+        dim = len(eigenbasis)
+        basis = {b: qutip.basis(dim, i) for i, b in enumerate(eigenbasis)}
+        op_matrix = {"I": qutip.qeye(dim)}
+        for proj0 in eigenbasis:
+            for proj1 in eigenbasis:
                 proj_name = "sigma_" + proj0 + proj1
-                self.op_matrix[proj_name] = (
-                    self.basis[proj0] * self.basis[proj1].dag()
-                )
+                op_matrix[proj_name] = basis[proj0] * basis[proj1].dag()
+        return basis, op_matrix
 
     def _construct_hamiltonian(self, update: bool = True) -> None:
         """Constructs the hamiltonian from the sampled Sequence and noise.
@@ -518,7 +571,7 @@ def build_coeffs_ops(basis: str, addr: str) -> list[list]:
         qobj_list = []
         # Time independent term:
         effective_size = self._size - sum(self._bad_atoms.values())
-        if self.basis_name != "digital" and effective_size > 1:
+        if "digital" not in self.basis_name and effective_size > 1:
             # Build time-dependent or time-independent interaction term based
             # on whether an SLM mask was defined or not
             if (