Add option to switch the Sequence register (#684)

* Add option to switch the Sequence register

* Addressing review comments

* Fix typo

pulser-core/pulser/sequence/sequence.py

@@ -684,6 +684,45 @@ def _config_detuning_map(
         # DMM has Global addressing
         self._add_to_schedule(dmm_name, _TimeSlot("target", -1, 0, self._qids))
 
+    def switch_register(
+        self, new_register: BaseRegister | MappableRegister
+    ) -> Sequence:
+        """Replicate the sequence with a different register.
+
+        The new sequence is reconstructed with the provided register by
+        replicating all the instructions used to build the original sequence.
+        This means that operations referecing specific qubits IDs
+        (eg. `Sequence.target()`) expect to find the same qubit IDs in the new
+        register. By the same token, switching from a register to a mappable
+        register might fail if one of the instructions does not work with
+        mappable registers (e.g. `Sequence.configure_slm_mask()`).
+
+        Warns:
+            UserWarning: If the sequence is configuring a detuning map, a
+            warning is raised to remind the user that the detuning map is
+            unchanged and might no longer be aligned with the qubits in
+            the new register.
+
+        Args:
+            new_register: The new register to give the sequence.
+
+        Returns:
+            The sequence with the new register.
+        """
+        new_seq = type(self)(register=new_register, device=self._device)
+        # Copy the variables to the new sequence
+        new_seq._variables = self.declared_variables
+        for call in self._calls[1:] + self._to_build_calls:
+            if call.name == "config_detuning_map":
+                warnings.warn(
+                    "Switching the register of a sequence that configures"
+                    " a detuning map. Please ensure that the new qubit"
+                    " positions are still aligned.",
+                    stacklevel=2,
+                )
+            getattr(new_seq, call.name)(*call.args, **call.kwargs)
+        return new_seq
+
     def switch_device(
         self, new_device: DeviceType, strict: bool = False
     ) -> Sequence:

tests/test_sequence.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import itertools
 import json
@@ -483,6 +484,7 @@ def init_seq(
     parametrized=False,
     mappable_reg=False,
     config_det_map=False,
+    prefer_slm_mask=True,
 ) -> Sequence:
     register = (
         reg.layout.make_mappable_register(len(reg.qubits))
@@ -506,7 +508,7 @@ def init_seq(
                 for i in range(10)
             }
         )
-        if mappable_reg:
+        if mappable_reg or not prefer_slm_mask:
             seq.config_detuning_map(detuning_map=det_map, dmm_id="dmm_0")
         else:
             seq.config_slm_mask(["q0"], "dmm_0")
@@ -533,6 +535,100 @@ def test_ising_mode(
         seq2._in_ising = True
 
 
+@pytest.mark.parametrize("config_det_map", [False, True])
+@pytest.mark.parametrize("starts_mappable", [False, True])
+@pytest.mark.parametrize("mappable_reg", [False, True])
+@pytest.mark.parametrize("parametrized", [False, True])
+def test_switch_register(
+    reg, mappable_reg, parametrized, starts_mappable, config_det_map
+):
+    pulse = Pulse.ConstantPulse(1000, 1, -1, 2)
+    with_slm_mask = not starts_mappable and not mappable_reg
+    seq = init_seq(
+        reg,
+        DigitalAnalogDevice,
+        "raman",
+        "raman_local",
+        [pulse],
+        initial_target="q0",
+        parametrized=parametrized,
+        mappable_reg=starts_mappable,
+        config_det_map=config_det_map,
+        prefer_slm_mask=with_slm_mask,
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="given ids have to be qubit ids declared in this sequence's"
+        " register",
+    ):
+        seq.switch_register(Register(dict(q1=(0, 0), qN=(10, 10))))
+
+    seq.declare_channel("ryd", "rydberg_global")
+    seq.add(pulse, "ryd", protocol="no-delay")
+
+    if mappable_reg:
+        new_reg = TriangularLatticeLayout(10, 5).make_mappable_register(2)
+    else:
+        new_reg = Register(dict(q0=(0, 0), foo=(10, 10)))
+
+    if config_det_map and not with_slm_mask:
+        context_manager = pytest.warns(
+            UserWarning, match="configures a detuning map"
+        )
+    else:
+        context_manager = contextlib.nullcontext()
+
+    with context_manager:
+        new_seq = seq.switch_register(new_reg)
+    assert seq.declared_variables or not parametrized
+    assert seq.declared_variables == new_seq.declared_variables
+    assert new_seq.is_parametrized() == parametrized
+    assert new_seq.is_register_mappable() == mappable_reg
+    assert new_seq._calls[1:] == seq._calls[1:]  # Excludes __init__
+    assert new_seq._to_build_calls == seq._to_build_calls
+
+    build_kwargs = {}
+    if parametrized:
+        build_kwargs["delay"] = 120
+    if mappable_reg:
+        build_kwargs["qubits"] = {"q0": 1, "q1": 4}
+    if build_kwargs:
+        new_seq = new_seq.build(**build_kwargs)
+
+    assert isinstance(
+        (raman_pulse_slot := new_seq._schedule["raman"][1]).type, Pulse
+    )
+    assert raman_pulse_slot.type == pulse
+    assert raman_pulse_slot.targets == {"q0"}
+
+    assert isinstance(
+        (rydberg_pulse_slot := new_seq._schedule["ryd"][1]).type, Pulse
+    )
+    assert rydberg_pulse_slot.type == pulse
+    assert rydberg_pulse_slot.targets == set(new_reg.qubit_ids)
+
+    if config_det_map:
+        if with_slm_mask:
+            if parametrized:
+                seq = seq.build(**build_kwargs)
+            assert np.any(reg.qubits["q0"] != new_reg.qubits["q0"])
+            assert "dmm_0" in seq.declared_channels
+            prev_qubit_wmap = seq._schedule[
+                "dmm_0"
+            ].detuning_map.get_qubit_weight_map(reg.qubits)
+            new_qubit_wmap = new_seq._schedule[
+                "dmm_0"
+            ].detuning_map.get_qubit_weight_map(new_reg.qubits)
+            assert prev_qubit_wmap["q0"] == 1.0
+            assert new_qubit_wmap == dict(q0=1.0, foo=0.0)
+        elif not parametrized:
+            assert (
+                seq._schedule["dmm_0"].detuning_map
+                == new_seq._schedule["dmm_0"].detuning_map
+            )
+
+
 @pytest.mark.parametrize("mappable_reg", [False, True])
 @pytest.mark.parametrize("parametrized", [False, True])
 def test_switch_device_down(