From 393526f0ebaea062e73f5f6fa597c4ab5bd99502 Mon Sep 17 00:00:00 2001
From: Antoine Cornillot <61453516+a-corni@users.noreply.github.com>
Date: Fri, 13 Sep 2024 10:13:20 +0200
Subject: [PATCH] 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   |   5 +-
 pulser-core/pulser/noise_model.py             |  12 +-
 .../pulser_simulation/hamiltonian.py          | 167 ++++++----
 .../pulser_simulation/qutip_result.py         | 108 ++++---
 .../pulser_simulation/simconfig.py            |   8 +-
 .../pulser_simulation/simresults.py           |  51 ++--
 .../pulser_simulation/simulation.py           |  50 ++-
 tests/test_noise_model.py                     |   6 +-
 tests/test_result.py                          | 112 ++++++-
 tests/test_simconfig.py                       |   6 +-
 tests/test_simresults.py                      |  58 +++-
 tests/test_simulation.py                      | 285 +++++++++++++-----
 12 files changed, 624 insertions(+), 244 deletions(-)

diff --git a/pulser-core/pulser/channels/base_channel.py b/pulser-core/pulser/channels/base_channel.py
index 6fdc1fb81..ea07f6799 100644
--- a/pulser-core/pulser/channels/base_channel.py
+++ b/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]
 
diff --git a/pulser-core/pulser/noise_model.py b/pulser-core/pulser/noise_model.py
index 7690ad684..cfe1d67ae 100644
--- a/pulser-core/pulser/noise_model.py
+++ b/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}."
diff --git a/pulser-simulation/pulser_simulation/hamiltonian.py b/pulser-simulation/pulser_simulation/hamiltonian.py
index 746a9838e..605c0ab76 100644
--- a/pulser-simulation/pulser_simulation/hamiltonian.py
+++ b/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,16 +117,16 @@ 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"
@@ -138,16 +134,18 @@ def basis_check(noise_type: str) -> None:
                 )
 
         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,7 +307,7 @@ 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]
@@ -289,7 +315,7 @@ def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
         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 (
diff --git a/pulser-simulation/pulser_simulation/qutip_result.py b/pulser-simulation/pulser_simulation/qutip_result.py
index b899beb22..b8c7da0f0 100644
--- a/pulser-simulation/pulser_simulation/qutip_result.py
+++ b/pulser-simulation/pulser_simulation/qutip_result.py
@@ -15,11 +15,16 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Union, cast
+from typing import cast
 
 import numpy as np
 import qutip
 
+from pulser.channels.base_channel import (
+    EIGENSTATES,
+    States,
+    get_states_from_bases,
+)
 from pulser.register import QubitId
 from pulser.result import Result
 
@@ -62,10 +67,22 @@ def _dim(self) -> int:
 
     @property
     def _basis_name(self) -> str:
-        if self._dim > 2:
-            return "all"
         if self.meas_basis == "XY":
+            if self._dim == 3:
+                return "XY_with_error"
+            assert (
+                self._dim == 2
+            ), f"In XY, state's dimension can only be 2 or 3, not {self._dim}."
             return "XY"
+        if self._dim == 4:
+            return "all_with_error"
+        if self._dim == 3:
+            if self.matching_meas_basis:
+                return self.meas_basis + "_with_error"
+            return "all"
+        assert (
+            self._dim == 2
+        ), f"In Ising, state's dimension can be 2, 3 or 4, not {self._dim}."
         if not self.matching_meas_basis:
             return (
                 "digital"
@@ -74,8 +91,17 @@ def _basis_name(self) -> str:
             )
         return self.meas_basis
 
+    @property
+    def _eigenbasis(self) -> list[States]:
+        bases = self._basis_name.split("_with_error")
+        states = get_states_from_bases(
+            ["ground-rydberg", "digital"] if bases[0] == "all" else [bases[0]]
+        )
+        states += ["x"] if len(bases) == 2 else []
+        return states
+
     def _weights(self) -> np.ndarray:
-        n = self._size
+        size = self._size
         if not self.state.isket:
             probs = np.abs(self.state.diag())
         else:
@@ -97,36 +123,38 @@ def _weights(self) -> np.ndarray:
                 weights = np.zeros(probs.size)
                 weights[0] = 1.0
 
-        elif self._dim == 3:
-            if self.meas_basis == "ground-rydberg":
-                one_state = 0  # 1 = |r>
-                ex_one = slice(1, 3)
-            elif self.meas_basis == "digital":
-                one_state = 2  # 1 = |h>
-                ex_one = slice(0, 2)
-            else:
+        elif self._dim == 3 or self._dim == 4:
+            one_state_dict: dict[str, States] = {
+                "ground-rydberg": "r",
+                "digital": "h",
+                "XY": "d",
+            }
+            if self.meas_basis not in one_state_dict:
                 raise RuntimeError(
-                    f"Unknown measurement basis '{self.meas_basis}' "
-                    "for a three-level system.'"
+                    f"Unknown measurement basis '{self.meas_basis}'."
                 )
-            probs = probs.reshape([3] * n)
-            weights = np.zeros(2**n)
-            for dec_val in range(2**n):
-                ind: list[Union[int, slice]] = []
-                for v in np.binary_repr(dec_val, width=n):
+            one_state_idx = self._eigenbasis.index(
+                one_state_dict[self.meas_basis]
+            )
+            ex_one = [i for i in range(self._dim) if i != one_state_idx]
+            probs = probs.reshape([self._dim] * size)
+            weights = np.zeros(2**size)
+            for dec_val in range(2**size):
+                ind: list[int | list[int]] = []
+                for v in np.binary_repr(dec_val, width=size):
                     if v == "0":
                         ind.append(ex_one)
                     else:
-                        ind.append(one_state)
+                        ind.append([one_state_idx])
                 # Eg: 'digital' basis : |1> = index2, |0> = index0, 1 = 0:2
                 # p_11010 = sum(probs[2, 2, 0:2, 2, 0:2])
                 # We sum all probabilites that correspond to measuring
                 # 11010, namely hhghg, hhrhg, hhghr, hhrhr
-                weights[dec_val] = np.sum(probs[tuple(ind)])
+                weights[dec_val] = np.sum(probs[np.ix_(*ind)])
         else:
             raise NotImplementedError(
                 "Cannot sample system with single-atom state vectors of "
-                "dimension > 3."
+                "dimension > 4."
             )
         # Takes care of numerical artefacts in case sum(weights) != 1
         return cast(np.ndarray, weights / sum(weights))
@@ -142,9 +170,8 @@ def get_state(
 
         Args:
             reduce_to_basis: Reduces the full state vector
-                to the given basis ("ground-rydberg" or "digital"), if the
-                population of the states to be ignored is negligible. Doesn't
-                apply to XY mode.
+                to the given basis ("ground-rydberg", "digital" or "XY"), if
+                the population of the states to be ignored is negligible.
             ignore_global_phase: If True and if the final state is a vector,
                 changes the final state's global phase such that the largest
                 term (in absolute value) is real.
@@ -164,7 +191,7 @@ def get_state(
             full = state.full()
             global_ph = float(np.angle(full[np.argmax(np.abs(full))])[0])
             state *= np.exp(-1j * global_ph)
-        if self._dim != 3:
+        if self._dim == 2:
             if reduce_to_basis not in [None, self._basis_name]:
                 raise TypeError(
                     f"Can't reduce a system in {self._basis_name}"
@@ -177,19 +204,30 @@ def get_state(
                     "Reduce to basis not implemented for density matrix"
                     " states."
                 )
-            if reduce_to_basis == "ground-rydberg":
-                ex_state = "2"
-            elif reduce_to_basis == "digital":
-                ex_state = "0"
-            else:
+            if reduce_to_basis not in EIGENSTATES:
                 raise ValueError(
-                    "'reduce_to_basis' must be 'ground-rydberg' "
-                    + f"or 'digital', not '{reduce_to_basis}'."
+                    "'reduce_to_basis' must be 'ground-rydberg', "
+                    f"'XY', or 'digital', not '{reduce_to_basis}'."
                 )
+            basis_states = set(self._eigenbasis)
+            target_states = set(EIGENSTATES[reduce_to_basis])
+            if not target_states.issubset(basis_states):
+                raise ValueError(
+                    f"Can't reduce a state expressed in {self._basis_name}"
+                    f" into {reduce_to_basis}"
+                )
+            # Exclude the states that are not in the basis into which to reduce
+            ex_states = basis_states - target_states
             ex_inds = [
                 i
-                for i in range(3**self._size)
-                if ex_state in np.base_repr(i, base=3).zfill(self._size)
+                for i in range(self._dim**self._size)
+                if any(
+                    [
+                        str(self._eigenbasis.index(ex_state))
+                        in np.base_repr(i, base=self._dim).zfill(self._size)
+                        for ex_state in ex_states
+                    ]
+                )
             ]
             ex_probs = np.abs(state.extract_states(ex_inds).full()) ** 2
             if not np.all(np.isclose(ex_probs, 0, atol=tol)):
diff --git a/pulser-simulation/pulser_simulation/simconfig.py b/pulser-simulation/pulser_simulation/simconfig.py
index ec7f6dbd0..5811495a7 100644
--- a/pulser-simulation/pulser_simulation/simconfig.py
+++ b/pulser-simulation/pulser_simulation/simconfig.py
@@ -38,13 +38,9 @@
         "doppler",
         "eff_noise",
         "SPAM",
+        "leakage",
     },
-    "XY": {
-        "dephasing",
-        "depolarizing",
-        "eff_noise",
-        "SPAM",
-    },
+    "XY": {"dephasing", "depolarizing", "eff_noise", "SPAM", "leakage"},
 }
 
 # Maps the noise model parameters with a different name in SimConfig
diff --git a/pulser-simulation/pulser_simulation/simresults.py b/pulser-simulation/pulser_simulation/simresults.py
index ec22169b9..493a28691 100644
--- a/pulser-simulation/pulser_simulation/simresults.py
+++ b/pulser-simulation/pulser_simulation/simresults.py
@@ -51,17 +51,17 @@ def __init__(
         Args:
             size: The number of atoms in the register.
             basis_name: The basis indicating the addressed atoms after
-                the pulse sequence ('ground-rydberg', 'digital' or 'all').
+                the pulse sequence ('ground-rydberg', 'digital' or 'all' or one
+                of these 3 bases with the suffix "_with_error").
             sim_times: Array of times (in µs) when simulation results are
                 returned.
         """
         self._dim = 3 if basis_name == "all" else 2
         self._size = size
-        if basis_name not in {"ground-rydberg", "digital", "all", "XY"}:
-            raise ValueError(
-                "`basis_name` must be 'ground-rydberg', 'digital', 'all' or "
-                "'XY'."
-            )
+        bases = ["ground-rydberg", "digital", "all", "XY"]
+        bases += [basis + "_with_error" for basis in bases]
+        if basis_name not in bases:
+            raise ValueError(f"`basis_name` must be in {bases}")
         self._basis_name = basis_name
         self._sim_times = sim_times
 
@@ -259,16 +259,20 @@ def __init__(
                 represented as a bitstring. There is one Counter for each time
                 the simulation was asked to return a result.
             size: The number of atoms in the register.
-            basis_name: Basis indicating the addressed atoms after
-                the pulse sequence ('ground-rydberg' or 'digital' - 'all' basis
-                makes no sense after projection on bitstrings). Defaults to
-                'digital' if given value 'all'.
+            basis_name: Basis indicating the addressed atoms after the pulse
+                sequence ('ground-rydberg' or 'digital' - 'all' basis or any
+                basis with the suffix "with_error" make no sense after
+                projection on bitstrings). Defaults to 'digital' if given value
+                'all' or 'all_with_error', and to 'ground-rydberg', 'XY',
+                'digital' if given respectively 'ground-rydberg_with_error',
+                'XY_with_error' or 'digital_with_error'.
             sim_times: Times at which Simulation object returned
                 the results.
             n_measures: Number of measurements needed to compute this
                 result when doing the simulation.
         """
-        basis_name_ = "digital" if basis_name == "all" else basis_name
+        basis = basis_name.replace("_with_error", "")
+        basis_name_ = "digital" if basis == "all" else basis
         super().__init__(size, basis_name_, sim_times)
         self.n_measures = n_measures
         self._results = tuple(run_output)
@@ -375,25 +379,28 @@ def __init__(
                 simulated.
             size: The number of atoms in the register.
             basis_name: The basis indicating the addressed atoms after
-                the pulse sequence ('ground-rydberg', 'digital' or 'all').
+                the pulse sequence ('ground-rydberg', 'digital' or 'all' or
+                one of these bases with the suffix "_with_error").
             sim_times: Times at which Simulation object returned the
                 results.
             meas_basis: The basis in which a sampling measurement
-                is desired.
+                is desired (must be in "ground-rydberg" or "digital").
             meas_errors: If measurement errors
                 are involved, give them in a dictionary with "epsilon" and
                 "epsilon_prime".
         """
         super().__init__(size, basis_name, sim_times)
-        if self._basis_name == "all":
+        if "all" in self._basis_name:
             if meas_basis not in {"ground-rydberg", "digital"}:
                 raise ValueError(
                     "`meas_basis` must be 'ground-rydberg' or 'digital'."
                 )
         else:
-            if meas_basis != self._basis_name:
+            expected_meas_basis = self._basis_name.replace("_with_error", "")
+            if meas_basis != expected_meas_basis:
                 raise ValueError(
-                    "`meas_basis` and `basis_name` must have the same value."
+                    f"`meas_basis` associated to basis_name '"
+                    f"{self._basis_name}' must be '{expected_meas_basis}'."
                 )
         self._meas_basis = meas_basis
         self._results = tuple(run_output)
@@ -425,9 +432,8 @@ def get_state(
         Args:
             t: Time (in µs) at which to return the state.
             reduce_to_basis: Reduces the full state vector
-                to the given basis ("ground-rydberg" or "digital"), if the
-                population of the states to be ignored is negligible. Doesn't
-                apply to XY mode.
+                to the given basis ("ground-rydberg", "digital" or "XY"), if
+                the population of the states to be ignored is negligible.
             ignore_global_phase: If True and if the final state is a vector,
                 changes the final state's global phase such that the largest
                 term (in absolute value) is real.
@@ -461,9 +467,8 @@ def get_final_state(
 
         Args:
             reduce_to_basis: Reduces the full state vector
-                to the given basis ("ground-rydberg" or "digital"), if the
-                population of the states to be ignored is negligible. Doesn't
-                apply to XY mode.
+                to the given basis ("ground-rydberg", "digital" or "XY"), if
+                the population of the states to be ignored is negligible.
             ignore_global_phase: If True, changes the
                 final state's global phase such that the largest term (in
                 absolute value) is real.
@@ -499,7 +504,7 @@ def _meas_projector(self, state_n: int) -> qutip.Qobj:
             # ground-rydberg
             good = (
                 1 - state_n
-                if self._basis_name == "ground-rydberg"
+                if "ground-rydberg" in self._basis_name
                 else state_n
             )
             return (
diff --git a/pulser-simulation/pulser_simulation/simulation.py b/pulser-simulation/pulser_simulation/simulation.py
index 4a7185370..77a3d11c1 100644
--- a/pulser-simulation/pulser_simulation/simulation.py
+++ b/pulser-simulation/pulser_simulation/simulation.py
@@ -28,6 +28,7 @@
 
 import pulser.sampler as sampler
 from pulser import Sequence
+from pulser.channels.base_channel import States
 from pulser.devices._device_datacls import BaseDevice
 from pulser.noise_model import NoiseModel
 from pulser.register.base_register import BaseRegister
@@ -163,10 +164,10 @@ def __init__(
         if self.samples_obj._measurement:
             self._meas_basis = self.samples_obj._measurement
         else:
-            if self._hamiltonian.basis_name in {"digital", "all"}:
+            if "all" in self.basis_name:
                 self._meas_basis = "digital"
             else:
-                self._meas_basis = self._hamiltonian.basis_name
+                self._meas_basis = self.basis_name.replace("_with_error", "")
         self.set_initial_state("all-ground")
 
     @property
@@ -190,7 +191,7 @@ def basis_name(self) -> str:
         return self._hamiltonian.basis_name
 
     @property
-    def basis(self) -> dict[str, Any]:
+    def basis(self) -> dict[States, Any]:
         """The basis in which result is expressed."""
         return self._hamiltonian.basis
 
@@ -217,7 +218,25 @@ def set_config(self, cfg: SimConfig) -> None:
                 " support simulation of noise types:"
                 f"{', '.join(not_supported)}."
             )
+        former_dim = self.dim
+        former_basis = self._hamiltonian.basis
         self._hamiltonian.set_config(cfg.to_noise_model())
+        if self.dim == former_dim:
+            self.set_initial_state(self._initial_state)
+            return
+        if self._initial_state != qutip.tensor(
+            [
+                former_basis[
+                    "u" if self._hamiltonian._interaction == "XY" else "g"
+                ]
+                for _ in range(self._hamiltonian._size)
+            ]
+        ):
+            warnings.warn(
+                "Current initial state's dimension does not match new"
+                " dimensions. Setting it to 'all-ground'."
+            )
+        self.set_initial_state("all-ground")
 
     def add_config(self, config: SimConfig) -> None:
         """Updates the current configuration with parameters of another one.
@@ -260,7 +279,25 @@ def add_config(self, config: SimConfig) -> None:
             param_dict[param] = getattr(noise_model, param)
         # set config with the new parameters:
         param_dict.pop("noise_types")
+        former_dim = self.dim
+        former_basis = self._hamiltonian.basis
         self._hamiltonian.set_config(NoiseModel(**param_dict))
+        if self.dim == former_dim:
+            self.set_initial_state(self._initial_state)
+            return
+        if self._initial_state != qutip.tensor(
+            [
+                former_basis[
+                    "u" if self._hamiltonian._interaction == "XY" else "g"
+                ]
+                for _ in range(self._hamiltonian._size)
+            ]
+        ):
+            warnings.warn(
+                "Current initial state's dimension does not match new"
+                " dimensions. Setting initial state to 'all-ground'."
+            )
+        self.set_initial_state("all-ground")
 
     def show_config(self, solver_options: bool = False) -> None:
         """Shows current configuration."""
@@ -556,14 +593,14 @@ def _run_solver() -> CoherentResults:
                     tuple(self._hamiltonian._qdict),
                     self._meas_basis,
                     state,
-                    self._meas_basis == self._hamiltonian.basis_name,
+                    self._meas_basis in self.basis_name,
                 )
                 for state in result.states
             ]
             return CoherentResults(
                 results,
                 self._hamiltonian._size,
-                self._hamiltonian.basis_name,
+                self.basis_name,
                 self._eval_times_array,
                 self._meas_basis,
                 meas_errors,
@@ -578,6 +615,7 @@ def _run_solver() -> CoherentResults:
                 "depolarizing",
                 "eff_noise",
                 "amplitude",
+                "leakage",
             }
         ) and (
             # If amplitude is in noise, not resampling needs amp_sigma=0.
@@ -650,7 +688,7 @@ def _run_solver() -> CoherentResults:
         return NoisyResults(
             results,
             self._hamiltonian._size,
-            self._hamiltonian.basis_name,
+            self.basis_name,
             self._eval_times_array,
             n_measures,
         )
diff --git a/tests/test_noise_model.py b/tests/test_noise_model.py
index dba514bb4..4e052e9ec 100644
--- a/tests/test_noise_model.py
+++ b/tests/test_noise_model.py
@@ -234,11 +234,9 @@ def test_eff_noise_opers(self, matrices):
                 eff_noise_rates=[1.0],
                 with_leakage=True,
             )
-        with pytest.raises(
-            NotImplementedError, match="With leakage, operator's shape"
-        ):
+        with pytest.raises(ValueError, match="With leakage, operator's shape"):
             NoiseModel(
-                eff_noise_opers=[matrices["I4"]],
+                eff_noise_opers=[np.eye(5)],
                 eff_noise_rates=[1.0],
                 with_leakage=True,
             )
diff --git a/tests/test_result.py b/tests/test_result.py
index cc0a22cca..feaee1c20 100644
--- a/tests/test_result.py
+++ b/tests/test_result.py
@@ -56,17 +56,20 @@ def test_sampled_result(patch_plt_show):
     result.plot_histogram()
 
 
-def test_qutip_result():
+def test_qutip_result_state():
     qutrit_state = qutip.tensor(qutip.basis(3, 0), qutip.basis(3, 1))
+
+    # Associated to "all" basis
     result = QutipResult(
         atom_order=("q0", "q1"),
         meas_basis="ground-rydberg",
         state=qutrit_state,
-        matching_meas_basis=True,
+        matching_meas_basis=False,
     )
     assert result.sampling_dist == {"10": 1.0}
     assert result.sampling_errors == {"10": 0.0}
     assert result._basis_name == "all"
+    assert result._eigenbasis == ["r", "g", "h"]
 
     assert result.get_state() == qutrit_state
     qubit_state = qutip.tensor(qutip.basis(2, 0), qutip.basis(2, 1))
@@ -74,13 +77,38 @@ def test_qutip_result():
         result.get_state(reduce_to_basis="ground-rydberg").full(),
         qubit_state.full(),
     )
+    with pytest.raises(
+        ValueError,
+        match="'reduce_to_basis' must be 'ground-rydberg', 'XY', or 'digital'",
+    ):
+        result.get_state("rydberg")
+    with pytest.raises(
+        ValueError, match="Can't reduce a state expressed in all into XY"
+    ):
+        result.get_state("XY")
 
     result.meas_basis = "digital"
     assert result.sampling_dist == {"00": 1.0}
+    assert result._basis_name == "all"
 
+    # Associated to bases with error state
+    # Associated to "digital_with_error"
+    result.matching_meas_basis = True
+    assert result._basis_name == "digital_with_error"
+    assert result._eigenbasis == ["g", "h", "x"]
+    assert result.sampling_dist == {"01": 1.0}
+
+    # Associated to "ground-rydberg_with_error"
+    result.meas_basis = "ground-rydberg"
+    assert result._basis_name == "ground-rydberg_with_error"
+    assert result._eigenbasis == ["r", "g", "x"]
+    assert result.sampling_dist == {"10": 1.0}
+
+    # Associated to "XY_with_error"
     result.meas_basis = "XY"
-    with pytest.raises(RuntimeError, match="Unknown measurement basis 'XY'"):
-        result.sampling_dist
+    assert result._basis_name == "XY_with_error"
+    assert result._eigenbasis == ["u", "d", "x"]
+    assert result.sampling_dist == {"01": 1.0}
 
     new_result = QutipResult(
         atom_order=("q0", "q1"),
@@ -102,22 +130,77 @@ def test_qutip_result():
     ):
         new_result.get_state(reduce_to_basis="ground-rydberg")
 
-    oversized_state = qutip.Qobj(np.eye(16) / 16)
-    result.state = oversized_state
-    assert result._dim == 4
+    # Associated with "all_wih_error_basis"
+    qudit_state = qutip.tensor(qutip.basis(4, 0), qutip.basis(4, 1))
+    qudit_result = QutipResult(
+        atom_order=("q0", "q1"),
+        meas_basis="ground-rydberg",
+        state=qudit_state,
+        matching_meas_basis=False,
+    )
+    assert qudit_result._dim == 4
+    assert qudit_result._basis_name == "all_with_error"
+    assert qudit_result._eigenbasis == ["r", "g", "h", "x"]
+    assert qudit_result.sampling_dist == {"10": 1.0}
+
+    qudit_result.meas_basis = "digital"
+    assert qudit_result.sampling_dist == {"00": 1.0}
+
+    qudit_result.meas_basis = "XY"
+    with pytest.raises(
+        AssertionError,
+        match="In XY, state's dimension can only be 2 or 3, not 4",
+    ):
+        qudit_result._basis_name
+    wrong_result = QutipResult(
+        atom_order=("q0", "q1"),
+        meas_basis="ground-rydberg",
+        state=qutip.tensor(qutip.basis(5, 0), qutip.basis(5, 1)),
+        matching_meas_basis=False,
+    )
+    assert wrong_result._dim == 5
+    with pytest.raises(
+        AssertionError,
+        match="In Ising, state's dimension can be 2, 3 or 4, not 5.",
+    ):
+        wrong_result._basis_name
+
     with pytest.raises(
         NotImplementedError,
         match="Cannot sample system with single-atom state vectors of"
-        " dimension > 3",
+        " dimension > 4",
+    ):
+        wrong_result.sampling_dist
+
+    qudit_result = QutipResult(
+        atom_order=("q0", "q1"),
+        meas_basis="rydberg",
+        state=qudit_state,
+        matching_meas_basis=False,
+    )
+    with pytest.raises(
+        RuntimeError,
+        match="Unknown measurement basis 'rydberg'.",
     ):
-        result.sampling_dist
+        qudit_result.sampling_dist
+
+
+def test_qutip_result_density_matrices():
+    qudit_density_matrix = qutip.Qobj(np.eye(16) / 16)
+    result = QutipResult(
+        atom_order=("a", "b"),
+        meas_basis="ground-rydberg",
+        state=qudit_density_matrix,
+        matching_meas_basis=False,
+    )
+    assert result._basis_name == "all_with_error"
 
     density_matrix = qutip.Qobj(np.eye(8) / 8)
     result = QutipResult(
         atom_order=("a", "b"),
         meas_basis="ground-rydberg",
         state=density_matrix,
-        matching_meas_basis=True,
+        matching_meas_basis=False,
     )
     assert result._basis_name == "all"
 
@@ -127,6 +210,15 @@ def test_qutip_result():
     ):
         result.get_state(reduce_to_basis="ground-rydberg")
 
+    result.matching_meas_basis = True
+    assert result._basis_name == "ground-rydberg_with_error"
+
+    result.meas_basis = "digital"
+    assert result._basis_name == "digital_with_error"
+
+    result.meas_basis = "XY"
+    assert result._basis_name == "XY_with_error"
+
     density_matrix = qutip.Qobj(np.eye(4) / 4)
     result = QutipResult(
         atom_order=("a", "b"),
diff --git a/tests/test_simconfig.py b/tests/test_simconfig.py
index ea9c3999c..5d7fe04d7 100644
--- a/tests/test_simconfig.py
+++ b/tests/test_simconfig.py
@@ -110,12 +110,10 @@ def test_eff_noise_opers(matrices):
             eff_noise_opers=[matrices["I"]],
             eff_noise_rates=[1.0],
         )
-    with pytest.raises(
-        NotImplementedError, match="With leakage, operator's shape"
-    ):
+    with pytest.raises(ValueError, match="With leakage, operator's shape"):
         SimConfig(
             noise=("eff_noise", "leakage"),
-            eff_noise_opers=[matrices["I4"]],
+            eff_noise_opers=[qeye(5)],
             eff_noise_rates=[1.0],
         )
     with pytest.raises(ValueError, match="Without leakage, operator's shape"):
diff --git a/tests/test_simresults.py b/tests/test_simresults.py
index 9943232ed..3941923f8 100644
--- a/tests/test_simresults.py
+++ b/tests/test_simresults.py
@@ -73,32 +73,41 @@ def results(sim):
     return sim.run()
 
 
-def test_initialization(results):
+@pytest.mark.parametrize(
+    ["basis", "exp_basis"],
+    [
+        ("ground-rydberg_with_error", "ground-rydberg"),
+        ("digital_with_error", "digital"),
+        ("all_with_error", "digital"),
+        ("all", "digital"),
+        ("XY_with_error", "XY"),
+    ],
+)
+def test_initialization(results, basis, exp_basis):
     rr_state = qutip.tensor([qutip.basis(2, 0), qutip.basis(2, 0)])
     with pytest.raises(ValueError, match="`basis_name` must be"):
         CoherentResults(rr_state, 2, "bad_basis", None, [0])
-    with pytest.raises(
-        ValueError, match="`meas_basis` must be 'ground-rydberg' or 'digital'."
-    ):
-        CoherentResults(rr_state, 1, "all", None, "XY")
-    with pytest.raises(
-        ValueError,
-        match="`meas_basis` and `basis_name` must have the same value.",
-    ):
-        CoherentResults(
-            rr_state, 1, "ground-rydberg", [0], "wrong_measurement_basis"
-        )
-    with pytest.raises(ValueError, match="`basis_name` must be"):
-        NoisyResults(rr_state, 2, "bad_basis", [0], 123)
+    if "all" in basis:
+        with pytest.raises(
+            ValueError,
+            match="`meas_basis` must be 'ground-rydberg' or 'digital'.",
+        ):
+            CoherentResults(rr_state, 1, basis, None, "XY")
+    else:
+        with pytest.raises(
+            ValueError,
+            match=f"`meas_basis` associated to basis_name '{basis}' must be",
+        ):
+            CoherentResults(rr_state, 1, basis, [0], "wrong_measurement_basis")
     with pytest.raises(
         ValueError, match="only values of 'epsilon' and 'epsilon_prime'"
     ):
         CoherentResults(
             rr_state,
             1,
-            "ground-rydberg",
+            basis,
             [0],
-            "ground-rydberg",
+            exp_basis,
             {"eta": 0.1, "epsilon": 0.0, "epsilon_prime": 0.4},
         )
 
@@ -111,6 +120,23 @@ def test_initialization(results):
     )
 
 
+@pytest.mark.parametrize(
+    ["basis", "exp_basis"],
+    [
+        ("ground-rydberg_with_error", "ground-rydberg"),
+        ("digital_with_error", "digital"),
+        ("all_with_error", "digital"),
+        ("all", "digital"),
+        ("XY_with_error", "XY"),
+    ],
+)
+def test_init_noisy(basis, exp_basis):
+    state = qutip.tensor([qutip.basis(2, 0), qutip.basis(2, 0)])
+    with pytest.raises(ValueError, match="`basis_name` must be"):
+        NoisyResults(state, 2, "bad_basis", [0], 123)
+    assert NoisyResults(state, 2, basis, [0], 100)._basis_name == exp_basis
+
+
 @pytest.mark.parametrize("noisychannel", [True, False])
 def test_get_final_state(
     noisychannel, sim: QutipEmulator, results, reg, pi_pulse
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 4c1427b4a..6c7d9cc0d 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -93,6 +93,7 @@ def matrices():
     pauli["Y"] = qutip.sigmay()
     pauli["Z"] = qutip.sigmaz()
     pauli["I3"] = qutip.qeye(3)
+    pauli["Z3"] = qutip.Qobj([[1, 0, 0], [0, -1, 0], [0, 0, 0]])
     return pauli
 
 
@@ -248,29 +249,51 @@ def test_extraction_of_sequences(seq):
                         ).all()
 
 
-def test_building_basis_and_projection_operators(seq, reg):
+@pytest.mark.parametrize("leakage", [False, True])
+def test_building_basis_and_projection_operators(seq, reg, leakage, matrices):
     # All three levels:
-    sim = QutipEmulator.from_sequence(seq, sampling_rate=0.01)
-    assert sim.basis_name == "all"
-    assert sim.dim == 3
-    assert sim.basis == {
-        "r": qutip.basis(3, 0),
-        "g": qutip.basis(3, 1),
-        "h": qutip.basis(3, 2),
+    def _config(dim):
+        return (
+            SimConfig(
+                ("leakage", "eff_noise"),
+                eff_noise_opers=[qutip.qeye(dim)],
+                eff_noise_rates=[0.0],
+            )
+            if leakage
+            else SimConfig()
+        )
+
+    dim = 3 + leakage
+    sim = QutipEmulator.from_sequence(
+        seq, sampling_rate=0.01, config=_config(dim)
+    )
+    assert sim.basis_name == "all" + ("_with_error" if leakage else "")
+    assert sim.dim == dim
+    basis_dict = {
+        "r": qutip.basis(dim, 0),
+        "g": qutip.basis(dim, 1),
+        "h": qutip.basis(dim, 2),
     }
+    if leakage:
+        basis_dict["x"] = qutip.basis(dim, 3)
+    assert sim.basis == basis_dict
     assert (
         sim._hamiltonian.op_matrix["sigma_rr"]
-        == qutip.basis(3, 0) * qutip.basis(3, 0).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim._hamiltonian.op_matrix["sigma_gr"]
-        == qutip.basis(3, 1) * qutip.basis(3, 0).dag()
+        == qutip.basis(dim, 1) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim._hamiltonian.op_matrix["sigma_hg"]
-        == qutip.basis(3, 2) * qutip.basis(3, 1).dag()
+        == qutip.basis(dim, 2) * qutip.basis(dim, 1).dag()
     )
-
+    if leakage:
+        assert (
+            sim._hamiltonian.op_matrix["sigma_xr"]
+            == qutip.basis(dim, 3) * qutip.basis(dim, 0).dag()
+        )
     # Check local operator building method:
     with pytest.raises(ValueError, match="Duplicate atom"):
         sim.build_operator([("sigma_gg", ["target", "target"])])
@@ -289,53 +312,86 @@ def test_building_basis_and_projection_operators(seq, reg):
     seq2.declare_channel("global", "rydberg_global")
     pi_pls = Pulse.ConstantDetuning(BlackmanWaveform(1000, np.pi), 0.0, 0)
     seq2.add(pi_pls, "global")
-    sim2 = QutipEmulator.from_sequence(seq2, sampling_rate=0.01)
-    assert sim2.basis_name == "ground-rydberg"
-    assert sim2.dim == 2
-    assert sim2.basis == {"r": qutip.basis(2, 0), "g": qutip.basis(2, 1)}
+    dim = 2 + leakage
+    sim2 = QutipEmulator.from_sequence(
+        seq2, sampling_rate=0.01, config=_config(dim)
+    )
+    assert sim2.basis_name == "ground-rydberg" + (
+        "_with_error" if leakage else ""
+    )
+    assert sim2.dim == dim
+    basis_dict = {"r": qutip.basis(dim, 0), "g": qutip.basis(dim, 1)}
+    if leakage:
+        basis_dict["x"] = qutip.basis(dim, 2)
+    assert sim2.basis == basis_dict
     assert (
         sim2._hamiltonian.op_matrix["sigma_rr"]
-        == qutip.basis(2, 0) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim2._hamiltonian.op_matrix["sigma_gr"]
-        == qutip.basis(2, 1) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 1) * qutip.basis(dim, 0).dag()
     )
-
+    if leakage:
+        assert (
+            sim2._hamiltonian.op_matrix["sigma_xr"]
+            == qutip.basis(dim, 2) * qutip.basis(dim, 0).dag()
+        )
     # Digital
     seq2b = Sequence(reg, DigitalAnalogDevice)
     seq2b.declare_channel("local", "raman_local", "target")
     seq2b.add(pi_pls, "local")
-    sim2b = QutipEmulator.from_sequence(seq2b, sampling_rate=0.01)
-    assert sim2b.basis_name == "digital"
-    assert sim2b.dim == 2
-    assert sim2b.basis == {"g": qutip.basis(2, 0), "h": qutip.basis(2, 1)}
+    sim2b = QutipEmulator.from_sequence(
+        seq2b, sampling_rate=0.01, config=_config(dim)
+    )
+    assert sim2b.basis_name == "digital" + ("_with_error" if leakage else "")
+    assert sim2b.dim == dim
+    basis_dict = {"g": qutip.basis(dim, 0), "h": qutip.basis(dim, 1)}
+    if leakage:
+        basis_dict["x"] = qutip.basis(dim, 2)
+    assert sim2b.basis == basis_dict
     assert (
         sim2b._hamiltonian.op_matrix["sigma_gg"]
-        == qutip.basis(2, 0) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim2b._hamiltonian.op_matrix["sigma_hg"]
-        == qutip.basis(2, 1) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 1) * qutip.basis(dim, 0).dag()
     )
+    if leakage:
+        assert (
+            sim2b._hamiltonian.op_matrix["sigma_xh"]
+            == qutip.basis(dim, 2) * qutip.basis(dim, 1).dag()
+        )
 
     # Local ground-rydberg
     seq2c = Sequence(reg, DigitalAnalogDevice)
     seq2c.declare_channel("local_ryd", "rydberg_local", "target")
     seq2c.add(pi_pls, "local_ryd")
-    sim2c = QutipEmulator.from_sequence(seq2c, sampling_rate=0.01)
-    assert sim2c.basis_name == "ground-rydberg"
-    assert sim2c.dim == 2
-    assert sim2c.basis == {"r": qutip.basis(2, 0), "g": qutip.basis(2, 1)}
+    sim2c = QutipEmulator.from_sequence(
+        seq2c, sampling_rate=0.01, config=_config(dim)
+    )
+    assert sim2c.basis_name == "ground-rydberg" + (
+        "_with_error" if leakage else ""
+    )
+    assert sim2c.dim == dim
+    basis_dict = {"r": qutip.basis(dim, 0), "g": qutip.basis(dim, 1)}
+    if leakage:
+        basis_dict["x"] = qutip.basis(dim, 2)
+    assert sim2c.basis == basis_dict
     assert (
         sim2c._hamiltonian.op_matrix["sigma_rr"]
-        == qutip.basis(2, 0) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim2c._hamiltonian.op_matrix["sigma_gr"]
-        == qutip.basis(2, 1) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 1) * qutip.basis(dim, 0).dag()
     )
-
+    if leakage:
+        assert (
+            sim2c._hamiltonian.op_matrix["sigma_xg"]
+            == qutip.basis(dim, 2) * qutip.basis(dim, 1).dag()
+        )
     # Global XY
     seq2 = Sequence(reg, MockDevice)
     seq2.declare_channel("global", "mw_global")
@@ -346,22 +402,32 @@ def test_building_basis_and_projection_operators(seq, reg):
         match="Bases used in samples should be supported by device.",
     ):
         QutipEmulator(sampler.sample(seq2), seq2.register, DigitalAnalogDevice)
-    sim2 = QutipEmulator.from_sequence(seq2, sampling_rate=0.01)
-    assert sim2.basis_name == "XY"
-    assert sim2.dim == 2
-    assert sim2.basis == {"u": qutip.basis(2, 0), "d": qutip.basis(2, 1)}
+    sim2 = QutipEmulator.from_sequence(
+        seq2, sampling_rate=0.01, config=_config(dim)
+    )
+    assert sim2.basis_name == "XY" + ("_with_error" if leakage else "")
+    assert sim2.dim == dim
+    basis_dict = {"u": qutip.basis(dim, 0), "d": qutip.basis(dim, 1)}
+    if leakage:
+        basis_dict["x"] = qutip.basis(dim, 2)
+    assert sim2.basis == basis_dict
     assert (
         sim2._hamiltonian.op_matrix["sigma_uu"]
-        == qutip.basis(2, 0) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim2._hamiltonian.op_matrix["sigma_du"]
-        == qutip.basis(2, 1) * qutip.basis(2, 0).dag()
+        == qutip.basis(dim, 1) * qutip.basis(dim, 0).dag()
     )
     assert (
         sim2._hamiltonian.op_matrix["sigma_ud"]
-        == qutip.basis(2, 0) * qutip.basis(2, 1).dag()
+        == qutip.basis(dim, 0) * qutip.basis(dim, 1).dag()
     )
+    if leakage:
+        assert (
+            sim2._hamiltonian.op_matrix["sigma_ux"]
+            == qutip.basis(dim, 0) * qutip.basis(dim, 2).dag()
+        )
 
 
 def test_empty_sequences(reg):
@@ -655,7 +721,7 @@ def test_eval_times(seq):
     )
 
 
-def test_config():
+def test_config(matrices):
     np.random.seed(123)
     reg = Register.from_coordinates([(0, 0), (0, 5)], prefix="q")
     seq = Sequence(reg, DigitalAnalogDevice)
@@ -684,6 +750,30 @@ def test_config():
         noisy_amp_ham[0, 0] == clean_ham[0, 0]
         and noisy_amp_ham[0, 1] != clean_ham[0, 1]
     )
+    assert sim._initial_state == qutip.tensor(
+        [qutip.basis(2, 1) for _ in range(2)]
+    )
+    # Currently in ground state => initial state is extended without warning
+    sim.set_config(
+        SimConfig(
+            noise=("leakage", "eff_noise"),
+            eff_noise_opers=[matrices["Z3"]],
+            eff_noise_rates=[0.1],
+        )
+    )
+    assert sim._initial_state == qutip.tensor(
+        [qutip.basis(3, 1) for _ in range(2)]
+    )
+    # Otherwise initial state is set to ground-state
+    sim.set_initial_state(qutip.tensor([qutip.basis(3, 0) for _ in range(2)]))
+    with pytest.warns(
+        UserWarning,
+        match="Current initial state's dimension does not match new dim",
+    ):
+        sim.set_config(SimConfig(noise="SPAM", eta=0.5))
+    assert sim._initial_state == qutip.tensor(
+        [qutip.basis(2, 1) for _ in range(2)]
+    )
 
 
 def test_noise(seq, matrices):
@@ -696,17 +786,6 @@ def test_noise(seq, matrices):
     )
     with pytest.raises(NotImplementedError, match="Cannot include"):
         sim2.set_config(SimConfig(noise="depolarizing"))
-    with pytest.raises(
-        NotImplementedError,
-        match="mode 'ising' does not support simulation of",
-    ):
-        sim2.set_config(
-            SimConfig(
-                ("leakage", "eff_noise"),
-                eff_noise_opers=[matrices["I3"]],
-                eff_noise_rates=[0.1],
-            )
-        )
     assert sim2.config.spam_dict == {
         "eta": 0.9,
         "epsilon": 0.01,
@@ -749,6 +828,7 @@ def test_noise_with_zero_epsilons(seq, matrices):
         ("depolarizing", {"0": 587, "1": 413}, 3),
         (("dephasing", "depolarizing", "relaxation"), {"0": 587, "1": 413}, 5),
         (("eff_noise", "dephasing"), {"0": 595, "1": 405}, 2),
+        (("eff_noise", "leakage"), {"0": 595, "1": 405}, 1),
     ],
 )
 def test_noises_rydberg(matrices, noise, result, n_collapse_ops):
@@ -765,8 +845,14 @@ def test_noises_rydberg(matrices, noise, result, n_collapse_ops):
         sampling_rate=0.01,
         config=SimConfig(
             noise=noise,
-            eff_noise_opers=[matrices["Z"]],
-            eff_noise_rates=[0.025],
+            eff_noise_opers=[
+                (
+                    qutip.Qobj([[1, 0, 0], [0, 0, 0], [0, 0, 0]])
+                    if "leakage" in noise
+                    else matrices["Z"]
+                )
+            ],
+            eff_noise_rates=[0.1 if "leakage" in noise else 0.025],
         ),
     )
     res = sim.run()
@@ -775,6 +861,10 @@ def test_noises_rydberg(matrices, noise, result, n_collapse_ops):
     assert len(sim._hamiltonian._collapse_ops) == n_collapse_ops
     trace_2 = res.states[-1] ** 2
     assert np.trace(trace_2) < 1 and not np.isclose(np.trace(trace_2), 1)
+    if "leakage" in noise:
+        state = res.get_final_state()
+        assert np.all(np.isclose(state[2, :], np.zeros_like(state[2, :])))
+        assert np.all(np.isclose(state[:, 2], np.zeros_like(state[:, 2])))
 
 
 def test_relaxation_noise():
@@ -796,6 +886,7 @@ def test_relaxation_noise():
         ryd_pop = new_ryd_pop
 
 
+deph_res = {"111": 978, "110": 11, "011": 6, "101": 5}
 depo_res = {
     "111": 821,
     "110": 61,
@@ -821,11 +912,13 @@ def test_relaxation_noise():
 @pytest.mark.parametrize(
     "noise, result, n_collapse_ops",
     [
-        ("dephasing", {"111": 978, "110": 11, "011": 6, "101": 5}, 1),
-        ("eff_noise", {"111": 978, "110": 11, "011": 6, "101": 5}, 1),
+        ("dephasing", deph_res, 1),
+        ("eff_noise", deph_res, 1),
         ("depolarizing", depo_res, 3),
         (("dephasing", "depolarizing"), deph_depo_res, 4),
         (("eff_noise", "dephasing"), eff_deph_res, 2),
+        (("eff_noise", "leakage"), deph_res, 1),
+        (("eff_noise", "leakage", "dephasing"), eff_deph_res, 2),
     ],
 )
 def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
@@ -837,8 +930,14 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
         config=SimConfig(
             noise=noise,
             hyperfine_dephasing_rate=0.05,
-            eff_noise_opers=[matrices["Z"]],
-            eff_noise_rates=[0.025],
+            eff_noise_opers=[
+                (
+                    qutip.Qobj([[0, 0, 0], [0, 1, 0], [0, 0, 0]])
+                    if "leakage" in noise
+                    else matrices["Z"]
+                )
+            ],
+            eff_noise_rates=[0.1 if "leakage" in noise else 0.025],
         ),
     )
 
@@ -856,6 +955,10 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
     )
     trace_2 = res.states[-1] ** 2
     assert np.trace(trace_2) < 1 and not np.isclose(np.trace(trace_2), 1)
+    if "leakage" in noise:
+        state = res.get_final_state()
+        assert np.all(np.isclose(state[2, :], np.zeros_like(state[2, :])))
+        assert np.all(np.isclose(state[:, 2], np.zeros_like(state[:, 2])))
 
 
 res_deph_relax = {
@@ -884,6 +987,11 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
             {"111": 922, "110": 33, "011": 23, "101": 21, "100": 1},
             4,
         ),
+        (
+            ("eff_noise", "leakage"),
+            {"111": 958, "110": 19, "011": 12, "101": 11},
+            2,
+        ),
     ],
 )
 def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
@@ -901,8 +1009,16 @@ def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
         seq.add(twopi_pulse, "ryd_glob")
         # Measure in the rydberg basis
         seq.measure()
-    deph_op = qutip.Qobj([[1, 0, 0], [0, 0, 0], [0, 0, 0]])
-    hyp_deph_op = qutip.Qobj([[0, 0, 0], [0, 0, 0], [0, 0, 1]])
+    if "leakage" in noise:
+        deph_op = qutip.Qobj(
+            [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
+        )
+        hyp_deph_op = qutip.Qobj(
+            [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]]
+        )
+    else:
+        deph_op = qutip.Qobj([[1, 0, 0], [0, 0, 0], [0, 0, 0]])
+        hyp_deph_op = qutip.Qobj([[0, 0, 0], [0, 0, 0], [0, 0, 1]])
     sim = QutipEmulator.from_sequence(
         seq,  # resulting state should be hhh
         sampling_rate=0.01,
@@ -915,7 +1031,6 @@ def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
             eff_noise_rates=[0.2, 0.2],
         ),
     )
-
     with pytest.raises(
         ValueError,
         match="Incompatible shape for effective noise operator n°0.",
@@ -944,6 +1059,10 @@ def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
     assert res_samples == Counter(result)
     trace_2 = res.states[-1] ** 2
     assert np.trace(trace_2) < 1 and not np.isclose(np.trace(trace_2), 1)
+    if "leakage" in noise:
+        state = res.get_final_state()
+        assert np.all(np.isclose(state[3, :], np.zeros_like(state[3, :])))
+        assert np.all(np.isclose(state[:, 3], np.zeros_like(state[:, 3])))
 
 
 def test_add_config(matrices):
@@ -994,6 +1113,31 @@ def test_add_config(matrices):
     sim.set_config(SimConfig(noise="SPAM", eta=0.5))
     sim.add_config(SimConfig(noise="depolarizing"))
     assert "depolarizing" in sim.config.noise
+    assert sim._initial_state == qutip.basis(2, 1)
+    # Currently in ground state => initial state is extended without warning
+    sim.add_config(
+        SimConfig(
+            noise=("leakage", "eff_noise"),
+            eff_noise_opers=[matrices["Z3"]],
+            eff_noise_rates=[0.1],
+        )
+    )
+    assert sim._initial_state == qutip.basis(3, 1)
+    # Otherwise initial state is set to ground-state
+    sim.set_config(SimConfig(noise="SPAM", eta=0.5))
+    sim.set_initial_state(qutip.basis(2, 0))
+    with pytest.warns(
+        UserWarning,
+        match="Current initial state's dimension does not match new dim",
+    ):
+        sim.add_config(
+            SimConfig(
+                noise=("leakage", "eff_noise"),
+                eff_noise_opers=[matrices["Z3"]],
+                eff_noise_rates=[0.1],
+            )
+        )
+    assert sim._initial_state == qutip.basis(3, 1)
 
 
 def test_concurrent_pulses():
@@ -1103,6 +1247,7 @@ def test_run_xy():
     [
         (None, "dephasing", res1, 1),
         (None, "eff_noise", res1, 1),
+        (None, "leakage", res1, 1),
         (None, "depolarizing", res2, 3),
         ("atom0", "dephasing", res3, 1),
         ("atom1", "dephasing", res4, 1),
@@ -1121,16 +1266,6 @@ def test_noisy_xy(matrices, masked_qubit, noise, result, n_collapse_ops):
     seq.add(rise, "ch0")
 
     sim = QutipEmulator.from_sequence(seq, sampling_rate=0.1)
-    with pytest.raises(
-        NotImplementedError, match="mode 'XY' does not support simulation of"
-    ):
-        sim.set_config(
-            SimConfig(
-                ("leakage", "eff_noise"),
-                eff_noise_opers=[matrices["I3"]],
-                eff_noise_rates=[0.1],
-            )
-        )
     with pytest.raises(
         NotImplementedError, match="mode 'XY' does not support simulation of"
     ):
@@ -1151,9 +1286,15 @@ def test_noisy_xy(matrices, masked_qubit, noise, result, n_collapse_ops):
     # SPAM simulation is implemented:
     sim.set_config(
         SimConfig(
-            ("SPAM", noise),
+            (
+                ("SPAM", noise)
+                if noise != "leakage"
+                else ("SPAM", "leakage", "eff_noise")
+            ),
             eta=0.4,
-            eff_noise_opers=[matrices["Z"]],
+            eff_noise_opers=[
+                matrices["Z"] if noise != "leakage" else matrices["Z3"]
+            ],
             eff_noise_rates=[0.025],
         )
     )