Enable definition of effective noise operators in all basis (#716)

* Enable definition of effective noise operators in all basis

* Modifying definition of collapse operators in dephasing

* Fix tests

* Fixing notebook

* Revert abs

---------

Co-authored-by: Henrique Silvério <[email protected]>

pulser-core/pulser/noise_model.py

@@ -401,8 +401,10 @@ 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 qutrit
-            if operator.shape != possible_shapes[0]:
+            # 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

pulser-simulation/pulser_simulation/hamiltonian.py

@@ -114,26 +114,18 @@ def basis_check(noise_type: str) -> None:
                     f"Cannot include {noise_type} noise in all-basis."
                 )
 
-        # NOTE: These operators only make sense when basis != "all"
-        b, a = self.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}"],
-        }
-
         local_collapse_ops = []
         if "dephasing" in config.noise_types:
-            basis_check("dephasing")
-            rate = (
-                config.hyperfine_dephasing_rate
-                if self.basis_name == "digital"
-                else config.dephasing_rate
-            )
-            local_collapse_ops.append(np.sqrt(rate / 2) * pauli_2d["z"])
+            dephasing_rates = {
+                "d": config.dephasing_rate,
+                "r": config.dephasing_rate,
+                "h": config.hyperfine_dephasing_rate,
+            }
+            for state in self.eigenbasis:
+                if state in dephasing_rates:
+                    coeff = np.sqrt(2 * dephasing_rates[state])
+                    op = self.op_matrix[f"sigma_{state}{state}"]
+                    local_collapse_ops.append(coeff * op)
 
         if "relaxation" in config.noise_types:
             coeff = np.sqrt(config.relaxation_rate)
@@ -147,18 +139,32 @@ def basis_check(noise_type: str) -> None:
 
         if "depolarizing" in config.noise_types:
             basis_check("depolarizing")
+            # NOTE: These operators only make sense when basis != "all"
+            b, a = self.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}"],
+            }
             coeff = np.sqrt(config.depolarizing_rate / 4)
             local_collapse_ops.append(coeff * pauli_2d["x"])
             local_collapse_ops.append(coeff * pauli_2d["y"])
             local_collapse_ops.append(coeff * pauli_2d["z"])
 
         if "eff_noise" in config.noise_types:
-            basis_check("effective")
             for id, rate in enumerate(config.eff_noise_rates):
-                local_collapse_ops.append(
-                    np.sqrt(rate) * np.array(config.eff_noise_opers[id])
-                )
-
+                op = np.array(config.eff_noise_opers[id])
+                basis_dim = len(self.eigenbasis)
+                op_shape = (basis_dim, basis_dim)
+                if op.shape != op_shape:
+                    raise ValueError(
+                        "Incompatible shape for effective noise operator n°"
+                        f"{id}. Operator {op} should be of shape {op_shape}."
+                    )
+                local_collapse_ops.append(np.sqrt(rate) * op)
         # Building collapse operators
         self._collapse_ops = []
         for operator in local_collapse_ops:

pulser-simulation/pulser_simulation/qutip_result.py

@@ -171,8 +171,8 @@ def get_state(
                     + f" to the {reduce_to_basis} basis."
                 )
         elif reduce_to_basis is not None:
-            if is_density_matrix:  # pragma: no cover
-                # Not tested as noise in digital or all basis not implemented
+            if is_density_matrix:
+                # TODO
                 raise NotImplementedError(
                     "Reduce to basis not implemented for density matrix"
                     " states."

tests/test_noise_model.py

@@ -242,13 +242,6 @@ def test_eff_noise_opers(self, matrices):
                 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"
         ):

tests/test_result.py

@@ -112,6 +112,21 @@ def test_qutip_result():
     ):
         result.sampling_dist
 
+    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,
+    )
+    assert result._basis_name == "all"
+
+    with pytest.raises(
+        NotImplementedError,
+        match="Reduce to basis not implemented for density matrix states.",
+    ):
+        result.get_state(reduce_to_basis="ground-rydberg")
+
     density_matrix = qutip.Qobj(np.eye(4) / 4)
     result = QutipResult(
         atom_order=("a", "b"),

tests/test_simconfig.py

@@ -124,14 +124,6 @@ def test_eff_noise_opers(matrices):
             eff_noise_opers=[matrices["I4"]],
             eff_noise_rates=[1.0],
         )
-    with pytest.raises(
-        NotImplementedError, match="Without leakage, operator's shape"
-    ):
-        SimConfig(
-            noise=("eff_noise",),
-            eff_noise_opers=[matrices["I3"]],
-            eff_noise_rates=[1.0],
-        )
     SimConfig(
         noise=("eff_noise"),
         eff_noise_opers=[matrices["X"], matrices["I"]],

tests/test_simulation.py

@@ -694,18 +694,8 @@ def test_noise(seq, matrices):
     assert sim2.run().sample_final_state() == Counter(
         {"000": 857, "110": 73, "100": 70}
     )
-    with pytest.raises(NotImplementedError, match="Cannot include"):
-        sim2.set_config(SimConfig(noise="dephasing"))
     with pytest.raises(NotImplementedError, match="Cannot include"):
         sim2.set_config(SimConfig(noise="depolarizing"))
-    with pytest.raises(NotImplementedError, match="Cannot include"):
-        sim2.set_config(
-            SimConfig(
-                noise="eff_noise",
-                eff_noise_opers=[matrices["I"]],
-                eff_noise_rates=[1.0],
-            )
-        )
     with pytest.raises(
         NotImplementedError,
         match="mode 'ising' does not support simulation of",
@@ -868,6 +858,71 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
     assert np.trace(trace_2) < 1 and not np.isclose(np.trace(trace_2), 1)
 
 
+@pytest.mark.parametrize(
+    "noise, result, n_collapse_ops",
+    [
+        ("dephasing", {"111": 958, "110": 19, "011": 12, "101": 11}, 2),
+        ("eff_noise", {"111": 958, "110": 19, "011": 12, "101": 11}, 2),
+        ("relaxation", {"111": 1000}, 1),
+        (
+            ("dephasing", "relaxation"),
+            {"111": 958, "110": 19, "011": 12, "101": 11},
+            3,
+        ),
+        (
+            ("eff_noise", "dephasing"),
+            {"111": 922, "110": 33, "011": 23, "101": 21, "100": 1},
+            4,
+        ),
+    ],
+)
+def test_noises_all(matrices, noise, result, n_collapse_ops, seq):
+    # Test with Digital Sequence
+    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,
+        config=SimConfig(
+            noise=noise,
+            dephasing_rate=0.1,
+            hyperfine_dephasing_rate=0.1,
+            relaxation_rate=1000,
+            eff_noise_opers=[deph_op, hyp_deph_op],
+            eff_noise_rates=[0.2, 0.2],
+        ),
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Incompatible shape for effective noise operator n°0.",
+    ):
+        # Only raised if 'eff_noise' in noise
+        sim.set_config(
+            SimConfig(
+                noise=("eff_noise",),
+                eff_noise_opers=[matrices["Z"]],
+                eff_noise_rates=[1.0],
+            )
+        )
+
+    with pytest.raises(
+        NotImplementedError,
+        match="Cannot include depolarizing noise in all-basis.",
+    ):
+        sim.set_config(SimConfig(noise="depolarizing"))
+
+    assert len(sim._hamiltonian._collapse_ops) == n_collapse_ops * len(
+        seq.register.qubits
+    )
+    np.random.seed(123)
+    res = sim.run()
+    res_samples = res.sample_final_state()
+    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)
+
+
 def test_add_config(matrices):
     reg = Register.from_coordinates([(0, 0)], prefix="q")
     seq = Sequence(reg, DigitalAnalogDevice)