Soften switching device with strict conditions

pulser-core/pulser/sequence/helpers/_switch_device.py

@@ -16,6 +16,7 @@
 
 import itertools
 import warnings
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
@@ -52,7 +53,7 @@ def switch_device(
     """
     # Check if the device is new or not
 
-    if seq._device == new_device:
+    if seq.device == new_device:
         warnings.warn(
             "Switching a sequence to the same device"
             + " returns the sequence unchanged.",
@@ -98,9 +99,10 @@ def check_channels_match(
 
         Returns a tuple that contains a non-strict error message and a
         strict error message. If the channel matches, the two error
-        messages are empty strings. If strict=False (True), the strict
-        (non-strict) error message - second (first) component of the
-        tuple - is always empty.
+        messages are empty strings. If strict=False, only non-strict
+        conditions are checked, and only the non-strict error message
+        will eventually be filled. If strict=True, all the conditions are
+        checked - the returned error can either be non-strict or strict.
         """
         old_ch_obj = seq.declared_channels[old_ch_name]
         # We verify the channel class then
@@ -117,20 +119,41 @@ def check_channels_match(
             if new_ch_obj.eom_config is None:
                 return (" with an EOM configuration.", "")
             if strict:
-                if (
-                    not seq.is_parametrized()
-                    and new_ch_obj.eom_config.mod_bandwidth
-                    != cast(RydbergEOM, old_ch_obj.eom_config).mod_bandwidth
-                ):
-                    return (
-                        "",
-                        " with the same mod_bandwidth for the EOM.",
-                    )
-                if (
-                    seq.is_parametrized()
-                    and new_ch_obj.eom_config != old_ch_obj.eom_config
-                ):
-                    return ("", " with the same EOM configuration.")
+                if not seq.is_parametrized():
+                    if (
+                        new_ch_obj.eom_config.mod_bandwidth
+                        != cast(
+                            RydbergEOM, old_ch_obj.eom_config
+                        ).mod_bandwidth
+                    ):
+                        return (
+                            "",
+                            " with the same mod_bandwidth for the EOM.",
+                        )
+                else:
+                    # Eom configs have to match is Sequence is parametrized
+                    new_eom_config = asdict(new_ch_obj.eom_config)
+                    old_eom_config = asdict(old_ch_obj.eom_config)
+                    # However, multiple_beam_control only matters when
+                    # the two beams are controlled
+                    if len(old_ch_obj.eom_config.controlled_beams) == 1:
+                        new_eom_config.pop("multiple_beam_control")
+                        old_eom_config.pop("multiple_beam_control")
+                        # Controlled beams only matter when only one beam
+                        # is controlled by the new eom
+                        if len(new_ch_obj.eom_config.controlled_beams) > 1:
+                            new_eom_config.pop("controlled_beams")
+                            old_eom_config.pop("controlled_beams")
+                    # And custom_buffer_time doesn't have to match as long
+                    # as `Channel_eom_buffer_time`` does
+                    if (
+                        new_ch_obj._eom_buffer_time
+                        == new_ch_obj._eom_buffer_time
+                    ):
+                        new_eom_config.pop("custom_buffer_time")
+                        old_eom_config.pop("custom_buffer_time")
+                    if new_eom_config != old_eom_config:
+                        return ("", " with the same EOM configuration.")
         if not strict:
             return ("", "")
 
@@ -346,6 +369,6 @@ def build_sequence_from_matching(
     raise ValueError(
         "No matching found between declared channels and channels in the "
         "new device that does not modify the samples of the Sequence. "
-        "Here is a list of matching tested and their associated errors: "
+        "Here is a list of matchings tested and their associated errors: "
         f"{err_channel_match}"
     )

tests/test_sequence.py

@@ -28,6 +28,7 @@
 from pulser import Pulse, Register, Register3D, Sequence
 from pulser.channels import Raman, Rydberg
 from pulser.channels.dmm import DMM
+from pulser.channels.eom import RydbergBeam
 from pulser.devices import AnalogDevice, DigitalAnalogDevice, MockDevice
 from pulser.devices._device_datacls import Device, VirtualDevice
 from pulser.register.base_register import BaseRegister
@@ -790,7 +791,7 @@ def test_switch_device_down(
     error_msg = (
         "No matching found between declared channels and channels in the "
         "new device that does not modify the samples of the Sequence. "
-        "Here is a list of matching tested and their associated errors: "
+        "Here is a list of matchings tested and their associated errors: "
         "{(('global', 'rydberg_global'), ('dmm_0', 'dmm_0'), ('dmm_0_1', "
         "'dmm_1')): ('The detunings on some atoms go below the local bottom "
         "detuning of the DMM (-10 rad/µs).',), (('global', 'rydberg_global'), "
@@ -1075,7 +1076,10 @@ def test_switch_device_up(
 
 @pytest.mark.parametrize("mappable_reg", [False, True])
 @pytest.mark.parametrize("parametrized", [False, True])
-def test_switch_device_eom(reg, mappable_reg, parametrized, patch_plt_show):
+@pytest.mark.parametrize("extension_arg", ["amp", "control", "buffer_time"])
+def test_switch_device_eom(
+    reg, mappable_reg, parametrized, extension_arg, patch_plt_show
+):
     # Sequence with EOM blocks
     seq = init_seq(
         reg,
@@ -1131,15 +1135,26 @@ def test_switch_device_eom(reg, mappable_reg, parametrized, patch_plt_show):
     )
     assert new_seq.declared_channels == {"rydberg": ch_obj}
     # Can if eom extends current eom
-    up_eom_config = dataclasses.replace(
-        ch_obj.eom_config, max_limiting_amp=40 * 2 * np.pi
-    )
+    up_eom_configs = {
+        "amp": dataclasses.replace(
+            ch_obj.eom_config, max_limiting_amp=40 * 2 * np.pi
+        ),
+        "control": dataclasses.replace(
+            ch_obj.eom_config,
+            controlled_beams=tuple(RydbergBeam),
+            multiple_beam_control=False,
+        ),
+        "buffer_time": dataclasses.replace(
+            ch_obj.eom_config,
+            custom_buffer_time=None,
+        ),
+    }
+    up_eom_config = up_eom_configs[extension_arg]
     up_ch_obj = dataclasses.replace(ch_obj, eom_config=up_eom_config)
     up_analog = dataclasses.replace(
         AnalogDevice, channel_objects=(up_ch_obj,), max_atom_num=28
     )
-    if parametrized:
-        # Can't switch to eom if the modulation bandwidth doesn't match
+    if parametrized and extension_arg == "amp":
         with pytest.raises(
             ValueError,
             match=err_base + "with the same EOM configuration.",
@@ -1177,11 +1192,18 @@ def test_switch_device_eom(reg, mappable_reg, parametrized, patch_plt_show):
     err_msg = (
         "No matching found between declared channels and channels in "
         "the new device that does not modify the samples of the "
-        "Sequence. Here is a list of matching tested and their "
+        "Sequence. Here is a list of matchings tested and their "
         "associated errors: {(('rydberg', 'rydberg_global'),): ('No "
         "match for channel rydberg with an EOM configuration that "
         "does not change the samples."
     )
+    if parametrized:
+        with pytest.raises(
+            ValueError,
+            match=err_base + "with the same EOM configuration.",
+        ):
+            seq.switch_device(mod_analog, strict=True)
+        return
     with pytest.raises(ValueError, match=re.escape(err_msg)):
         seq.switch_device(mod_analog, strict=True)
     mod_seq = seq.switch_device(mod_analog, strict=False)