Upgrade pulser-simulation to qutip 5 (#783)

* Qutip 4 compatible changes * Qutip 5 breaking changes * Fix test, handling of exception in NoiseModel * Convert all Qobj to CSR * Fix typing * Fix typing * Fix lint * Delete print * Address nit * Convert initial state to CSR --------- Co-authored-by: HGSilveri <[email protected]> Co-authored-by: Henrique Silvério <[email protected]>
pasqal-io · Dec 18, 2024 · 21c8d46 · 21c8d46
1 parent e2ad837
commit 21c8d46
Show file tree

Hide file tree

Showing 11 changed files with 97 additions and 64 deletions.
diff --git a/pulser-core/pulser/noise_model.py b/pulser-core/pulser/noise_model.py
@@ -355,10 +355,16 @@ def _check_eff_noise(
             # type checking
             try:
                 operator = np.array(op, dtype=complex)
-            except Exception:
-                raise TypeError(
-                    f"Operator {op!r} is not castable to a Numpy array."
-                )
+            except TypeError as e1:
+                try:
+                    operator = np.array(
+                        op.to("Dense").data_as("ndarray"),  # type: ignore
+                        dtype=complex,
+                    )
+                except AttributeError:
+                    raise TypeError(
+                        f"Operator {op!r} is not castable to a Numpy array."
+                    ) from e1
             if operator.ndim != 2:
                 raise ValueError(f"Operator '{op!r}' is not a 2D array.")
 

diff --git a/pulser-simulation/pulser_simulation/hamiltonian.py b/pulser-simulation/pulser_simulation/hamiltonian.py
@@ -366,10 +366,14 @@ def _build_operator(
                         operator = self.op_matrix[operator]
                     except KeyError:
                         raise ValueError(f"{operator} is not a valid operator")
+                elif isinstance(operator, qutip.Qobj):
+                    operator = operator.to("CSR")
+                else:
+                    operator = qutip.Qobj(operator).to("CSR")
                 for qubit in qubits:
                     k = self._qid_index[qubit]
                     op_list[k] = operator
-        return qutip.tensor(list(map(qutip.Qobj, op_list)))
+        return qutip.tensor(op_list)
 
     def build_operator(self, operations: Union[list, tuple]) -> qutip.Qobj:
         """Creates an operator with non-trivial actions on some qubits.
@@ -443,8 +447,9 @@ def _get_basis_op_matrices(
     ) -> 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)}
+        with qutip.CoreOptions(default_dtype="CSR"):
+            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

diff --git a/pulser-simulation/pulser_simulation/qutip_backend.py b/pulser-simulation/pulser_simulation/qutip_backend.py
@@ -72,7 +72,7 @@ def run(
         Args:
             progress_bar: If True, the progress bar of QuTiP's
                 solver will be shown. If None or False, no text appears.
-            options: Used as arguments for qutip.Options(). If specified, will
+            options: Given directly to the Qutip solver. If specified, will
                 override SimConfig solver_options. If no `max_step` value is
                 provided, an automatic one is calculated from the `Sequence`'s
                 schedule (half of the shortest duration among pulses and

diff --git a/pulser-simulation/pulser_simulation/qutip_result.py b/pulser-simulation/pulser_simulation/qutip_result.py
@@ -229,11 +229,16 @@ def get_state(
                     ]
                 )
             ]
-            ex_probs = np.abs(state.extract_states(ex_inds).full()) ** 2
+            state_arr = state.full()
+            ex_probs = np.abs(state_arr[ex_inds]) ** 2
             if not np.all(np.isclose(ex_probs, 0, atol=tol)):
                 raise TypeError(
                     "Can't reduce to chosen basis because the population of a "
                     "state to eliminate is above the allowed tolerance."
                 )
-            state = state.eliminate_states(ex_inds, normalize=normalize)
+            mask = np.ones_like(state_arr, dtype=bool)
+            mask[ex_inds] = False
+            state = qutip.Qobj(state_arr[mask])
+            if normalize:
+                state.unit(inplace=True)
         return state.tidyup()
diff --git a/pulser-simulation/pulser_simulation/simconfig.py b/pulser-simulation/pulser_simulation/simconfig.py
@@ -17,7 +17,7 @@
 
 import math
 from dataclasses import dataclass, field, fields
-from typing import Any, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import Any, Tuple, Type, TypeVar, Union, cast
 
 import qutip
 
@@ -123,7 +123,7 @@ class SimConfig:
     depolarizing_rate: float = _LEGACY_DEFAULTS["depolarizing_rate"]
     eff_noise_rates: list[float] = field(default_factory=list, repr=False)
     eff_noise_opers: list[qutip.Qobj] = field(default_factory=list, repr=False)
-    solver_options: Optional[qutip.Options] = None
+    solver_options: dict[str, Any] | None = None
 
     @classmethod
     def from_noise_model(cls: Type[T], noise_model: NoiseModel) -> T:
@@ -144,6 +144,10 @@ def from_noise_model(cls: Type[T], noise_model: NoiseModel) -> T:
         if "amplitude" in noise_model.noise_types:
             kwargs.setdefault("laser_waist", float("inf"))
         kwargs.pop("with_leakage", None)
+        if "eff_noise_opers" in kwargs:
+            kwargs["eff_noise_opers"] = list(
+                map(qutip.Qobj, kwargs["eff_noise_opers"])
+            )
         return cls(**kwargs)
 
     def to_noise_model(self) -> NoiseModel:
@@ -162,6 +166,10 @@ def to_noise_model(self) -> NoiseModel:
             kwargs[param] = getattr(self, _DIFF_NOISE_PARAMS.get(param, param))
         if "temperature" in kwargs:
             kwargs["temperature"] *= 1e6  # Converts back to µK
+        if "eff_noise_opers" in kwargs:
+            kwargs["eff_noise_opers"] = [
+                op.full() for op in kwargs["eff_noise_opers"]
+            ]
         return NoiseModel(**kwargs)
 
     def __post_init__(self) -> None:
@@ -263,7 +271,7 @@ def _check_eff_noise(self) -> None:
                 )
         NoiseModel._check_eff_noise(
             self.eff_noise_rates,
-            self.eff_noise_opers,
+            [op.full() for op in self.eff_noise_opers],
             "eff_noise" in self.noise,
             self.with_leakage,
         )

diff --git a/pulser-simulation/pulser_simulation/simresults.py b/pulser-simulation/pulser_simulation/simresults.py
@@ -26,7 +26,7 @@
 import numpy as np
 import qutip
 from numpy.typing import ArrayLike
-from qutip.piqs import isdiagonal
+from qutip.piqs.piqs import isdiagonal
 
 from pulser.result import Results, ResultType, SampledResult
 from pulser_simulation.qutip_result import QutipResult

diff --git a/pulser-simulation/pulser_simulation/simulation.py b/pulser-simulation/pulser_simulation/simulation.py
@@ -348,7 +348,7 @@ def set_initial_state(
                     "Incompatible shape of initial state."
                     + f"Expected {legal_shape}, got {shape}."
                 )
-            self._initial_state = qutip.Qobj(state, dims=legal_dims)
+            self._initial_state = qutip.Qobj(state, dims=legal_dims).to("CSR")
 
     @property
     def evaluation_times(self) -> np.ndarray:
@@ -491,7 +491,7 @@ def run(
         Args:
             progress_bar: If True, the progress bar of QuTiP's
                 solver will be shown. If None or False, no text appears.
-            options: Used as arguments for qutip.Options(). If specified, will
+            options: Given directly to the Qutip Solver. If specified, will
                 override SimConfig solver_options. If no `max_step` value is
                 provided, an automatic one is calculated from the `Sequence`'s
                 schedule (half of the shortest duration among pulses and
@@ -536,7 +536,6 @@ def get_min_variation(ch_sample: ChannelSamples) -> int:
             options["nsteps"] = max(
                 1000, self._tot_duration // options["max_step"]
             )
-        solv_ops = qutip.Options(**options)
 
         meas_errors: Optional[Mapping[str, float]] = None
         if "SPAM" in self.config.noise:
@@ -580,16 +579,18 @@ def _run_solver() -> CoherentResults:
                     self.initial_state,
                     self._eval_times_array,
                     self._hamiltonian._collapse_ops,
-                    progress_bar=p_bar,
-                    options=solv_ops,
+                    options=dict(
+                        progress_bar=p_bar, normalize_output=False, **options
+                    ),
                 )
             else:
                 result = qutip.sesolve(
                     self._hamiltonian._hamiltonian,
                     self.initial_state,
                     self._eval_times_array,
-                    progress_bar=p_bar,
-                    options=solv_ops,
+                    options=dict(
+                        progress_bar=p_bar, normalize_output=False, **options
+                    ),
                 )
             results = [
                 QutipResult(

diff --git a/pulser-simulation/requirements.txt b/pulser-simulation/requirements.txt
@@ -1,3 +1 @@
-qutip~=4.7.5
-# This is needed until qutip fixes the incompatibility with scipy 1.12
-scipy<1.13
+qutip >= 5, < 6
diff --git a/tests/test_qutip_backend.py b/tests/test_qutip_backend.py
@@ -87,3 +87,22 @@ def test_with_default_noise(sequence):
     new_results = backend.run()
     assert isinstance(new_results, NoisyResults)
     assert backend._sim_obj.config == SimConfig.from_noise_model(spam_noise)
+
+
+proj = [[0, 0], [0, 1]]
+
+
+@pytest.mark.parametrize(
+    "collapse_op", [qutip.sigmax(), qutip.Qobj(proj), np.array(proj), proj]
+)
+def test_collapse_op(sequence, collapse_op):
+    noise_model = pulser.NoiseModel(
+        eff_noise_opers=[collapse_op], eff_noise_rates=[0.1]
+    )
+    backend = QutipBackend(
+        sequence, config=pulser.EmulatorConfig(noise_model=noise_model)
+    )
+    assert [
+        op.type == qutip.core.data.CSR
+        for op in backend._sim_obj._hamiltonian._collapse_ops
+    ]
diff --git a/tests/test_simresults.py b/tests/test_simresults.py
@@ -17,7 +17,7 @@
 import numpy as np
 import pytest
 import qutip
-from qutip.piqs import isdiagonal
+from qutip.piqs.piqs import isdiagonal
 
 from pulser import AnalogDevice, Pulse, Register, Sequence
 from pulser.devices import DigitalAnalogDevice, MockDevice
@@ -189,8 +189,8 @@ def test_get_final_state(
         results_.get_final_state(reduce_to_basis="digital")
     h_states = results_.get_final_state(
         reduce_to_basis="digital", tol=1, normalize=False
-    ).eliminate_states([0])
-    assert h_states.norm() < 3e-6
+    ).full()[1:]
+    assert np.linalg.norm(h_states) < 3e-6
 
     assert np.all(
         np.isclose(
@@ -236,18 +236,6 @@ def test_get_state_float_time(results):
         results.get_state(mean, t_tol=diff / 2)
     state = results.get_state(mean, t_tol=3 * diff / 2)
     assert state == results.get_state(results._sim_times[-2])
-    np.testing.assert_allclose(
-        state.full(),
-        np.array(
-            [
-                [0.76522907 + 0.0j],
-                [0.08339973 - 0.39374219j],
-                [0.08339973 - 0.39374219j],
-                [-0.27977172 - 0.11031832j],
-            ]
-        ),
-        atol=1e-5,
-    )
 
 
 def test_expect(results, pi_pulse, reg):

diff --git a/tests/test_simulation.py b/tests/test_simulation.py
@@ -171,15 +171,7 @@ def test_initialization_and_construction_of_hamiltonian(seq, mod_device):
     )
 
     assert isinstance(sim._hamiltonian._hamiltonian, qutip.QobjEvo)
-    # Checks adapt() method:
-    assert bool(
-        set(sim._hamiltonian._hamiltonian.tlist).intersection(
-            sim.sampling_times
-        )
-    )
-    for qobjevo in sim._hamiltonian._hamiltonian.ops:
-        for sh in qobjevo.qobj.shape:
-            assert sh == sim.dim**sim._hamiltonian._size
+    assert sim._hamiltonian._hamiltonian(0).dtype == qutip.core.data.CSR
 
     assert not seq.is_parametrized()
     with pytest.warns(UserWarning, match="returns a copy of itself"):
@@ -295,7 +287,7 @@ def _config(dim):
     # Check building operator with one operator
     op_standard = sim.build_operator([("sigma_gg", ["target"])])
     op_one = sim.build_operator(("sigma_gg", ["target"]))
-    assert np.linalg.norm(op_standard - op_one) < 1e-10
+    assert (op_standard - op_one).norm() < 1e-10
 
     # Global ground-rydberg
     seq2 = Sequence(reg, DigitalAnalogDevice)
@@ -788,13 +780,13 @@ def test_noise_with_zero_epsilons(seq, matrices):
 @pytest.mark.parametrize(
     "noise, result, n_collapse_ops",
     [
-        ("dephasing", {"0": 595, "1": 405}, 1),
-        ("relaxation", {"0": 595, "1": 405}, 1),
-        ("eff_noise", {"0": 595, "1": 405}, 1),
-        ("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),
+        ("dephasing", {"0": 586, "1": 414}, 1),
+        ("relaxation", {"0": 586, "1": 414}, 1),
+        ("eff_noise", {"0": 586, "1": 414}, 1),
+        ("depolarizing", {"0": 581, "1": 419}, 3),
+        (("dephasing", "depolarizing", "relaxation"), {"0": 582, "1": 418}, 5),
+        (("eff_noise", "dephasing"), {"0": 587, "1": 413}, 2),
+        (("eff_noise", "leakage"), {"0": 586, "1": 414}, 1),
     ],
 )
 def test_noises_rydberg(matrices, noise, result, n_collapse_ops):
@@ -821,12 +813,15 @@ def test_noises_rydberg(matrices, noise, result, n_collapse_ops):
             eff_noise_rates=[0.1 if "leakage" in noise else 0.025],
         ),
     )
+    assert [
+        op.type == qutip.core.data.CSR for op in sim._hamiltonian._collapse_ops
+    ]
     res = sim.run()
     res_samples = res.sample_final_state()
     assert res_samples == Counter(result)
     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)
+    trace_2 = np.trace((res.states[-1] ** 2).full())
+    assert trace_2 < 1 and not np.isclose(trace_2, 1)
     if "leakage" in noise:
         state = res.get_final_state()
         assert np.all(np.isclose(state[2, :], np.zeros_like(state[2, :])))
@@ -841,6 +836,9 @@ def test_relaxation_noise():
 
     sim = QutipEmulator.from_sequence(seq)
     sim.add_config(SimConfig(noise="relaxation", relaxation_rate=0.1))
+    assert [
+        op.type == qutip.core.data.CSR for op in sim._hamiltonian._collapse_ops
+    ]
     res = sim.run()
     start_samples = res.sample_state(1)
     ryd_pop = start_samples["1"]
@@ -906,7 +904,9 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
             eff_noise_rates=[0.1 if "leakage" in noise else 0.025],
         ),
     )
-
+    assert [
+        op.type == qutip.core.data.CSR for op in sim._hamiltonian._collapse_ops
+    ]
     with pytest.raises(
         ValueError,
         match="'relaxation' noise requires addressing of the 'ground-rydberg'",
@@ -919,8 +919,8 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
     assert len(sim._hamiltonian._collapse_ops) == n_collapse_ops * len(
         seq_digital.register.qubits
     )
-    trace_2 = res.states[-1] ** 2
-    assert np.trace(trace_2) < 1 and not np.isclose(np.trace(trace_2), 1)
+    trace_2 = np.trace((res.states[-1] ** 2).full())
+    assert trace_2 < 1 and not np.isclose(trace_2, 1)
     if "leakage" in noise:
         state = res.get_final_state()
         assert np.all(np.isclose(state[2, :], np.zeros_like(state[2, :])))
@@ -944,7 +944,7 @@ def test_noises_digital(matrices, noise, result, n_collapse_ops, seq_digital):
         ("eff_noise", {"111": 958, "110": 19, "011": 12, "101": 11}, 2),
         (
             "relaxation",
-            {"000": 421, "010": 231, "001": 172, "100": 171, "101": 5},
+            {"000": 420, "010": 231, "001": 173, "100": 171, "101": 5},
             1,
         ),
         (("dephasing", "relaxation"), res_deph_relax, 3),
@@ -997,6 +997,9 @@ def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
             eff_noise_rates=[0.2, 0.2],
         ),
     )
+    assert [
+        op.type == qutip.core.data.CSR for op in sim._hamiltonian._collapse_ops
+    ]
     with pytest.raises(
         ValueError,
         match="Incompatible shape for effective noise operator n°0.",
@@ -1023,8 +1026,8 @@ def test_noises_all(matrices, reg, noise, result, n_collapse_ops, seq):
     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)
+    trace_2 = np.trace((res.states[-1] ** 2).full())
+    assert trace_2 < 1 and not np.isclose(trace_2, 1)
     if "leakage" in noise:
         state = res.get_final_state()
         assert np.all(np.isclose(state[3, :], np.zeros_like(state[3, :])))