Allow modification of the EOM setpoint without disabling EOM mode (#708)

* Allow modification of the EOM setpoint without disabling EOM mode

* Adding abstract repr support

pulser-core/pulser/json/abstract_repr/deserializer.py

@@ -291,6 +291,16 @@ def _deserialize_operation(seq: Sequence, op: dict, vars: dict) -> None:
             ),
             correct_phase_drift=op.get("correct_phase_drift", False),
         )
+    elif op["op"] == "modify_eom_setpoint":
+        seq.modify_eom_setpoint(
+            channel=op["channel"],
+            amp_on=_deserialize_parameter(op["amp_on"], vars),
+            detuning_on=_deserialize_parameter(op["detuning_on"], vars),
+            optimal_detuning_off=_deserialize_parameter(
+                op["optimal_detuning_off"], vars
+            ),
+            correct_phase_drift=op["correct_phase_drift"],
+        )
     elif op["op"] == "add_eom_pulse":
         seq.add_eom_pulse(
             channel=op["channel"],

pulser-core/pulser/json/abstract_repr/schemas/sequence-schema.json

@@ -687,6 +687,44 @@
       ],
       "type": "object"
     },
+    "OpModifyEOM": {
+      "additionalProperties": false,
+      "properties": {
+        "amp_on": {
+          "$ref": "#/definitions/ParametrizedNum",
+          "description": "The new amplitude of the EOM pulses (in rad/µs)."
+        },
+        "channel": {
+          "$ref": "#/definitions/ChannelName",
+          "description": "The name of the channel currently in EOM mode."
+        },
+        "correct_phase_drift": {
+          "description": "Performs a phase shift to correct for the phase drift incurred while modifying the EOM setpoint.",
+          "type": "boolean"
+        },
+        "detuning_on": {
+          "$ref": "#/definitions/ParametrizedNum",
+          "description": "The new detuning of the EOM pulses (in rad/µs)."
+        },
+        "op": {
+          "const": "modify_eom_setpoint",
+          "type": "string"
+        },
+        "optimal_detuning_off": {
+          "$ref": "#/definitions/ParametrizedNum",
+          "description": "The new optimal value of detuning (in rad/µs) when there is no pulse being  played. It will choose the closest value among the existing options."
+        }
+      },
+      "required": [
+        "op",
+        "channel",
+        "amp_on",
+        "detuning_on",
+        "optimal_detuning_off",
+        "correct_phase_drift"
+      ],
+      "type": "object"
+    },
     "OpPhaseShift": {
       "additionalProperties": false,
       "description": "Adds a separate phase shift to atoms. If possible, OpPulse phase and post_phase_shift are preferred.",
@@ -865,6 +903,9 @@
         {
           "$ref": "#/definitions/OpEnableEOM"
         },
+        {
+          "$ref": "#/definitions/OpModifyEOM"
+        },
         {
           "$ref": "#/definitions/OpDisableEOM"
         },

pulser-core/pulser/json/abstract_repr/serializer.py

@@ -358,6 +358,18 @@ def remove_kwarg_if_default(
                 data, call.name, "correct_phase_drift"
             )
             operations.append({"op": "enable_eom_mode", **data})
+        elif call.name == "modify_eom_setpoint":
+            data = get_all_args(
+                (
+                    "channel",
+                    "amp_on",
+                    "detuning_on",
+                    "optimal_detuning_off",
+                    "correct_phase_drift",
+                ),
+                call,
+            )
+            operations.append({"op": "modify_eom_setpoint", **data})
         elif call.name == "add_eom_pulse":
             data = get_all_args(
                 (

pulser-core/pulser/sequence/_schedule.py

@@ -341,12 +341,14 @@ def enable_eom(
         detuning_off: float,
         switching_beams: tuple[RydbergBeam, ...] = (),
         _skip_buffer: bool = False,
+        _skip_wait_for_fall: bool = False,
     ) -> None:
         channel_obj = self[channel_id].channel_obj
         # Adds a buffer unless the channel is empty or _skip_buffer = True
         if not _skip_buffer and self.get_duration(channel_id):
-            # Wait for the last pulse to ramp down (if needed)
-            self.wait_for_fall(channel_id)
+            if not _skip_wait_for_fall:
+                # Wait for the last pulse to ramp down (if needed)
+                self.wait_for_fall(channel_id)
             eom_buffer_time = self[channel_id].adjust_duration(
                 channel_obj._eom_buffer_time
             )

pulser-core/pulser/sequence/sequence.py

@@ -43,7 +43,7 @@
 import pulser.sequence._decorators as seq_decorators
 from pulser.channels.base_channel import Channel, States, get_states_from_bases
 from pulser.channels.dmm import DMM, _dmm_id_from_name, _get_dmm_name
-from pulser.channels.eom import RydbergEOM
+from pulser.channels.eom import RydbergBeam, RydbergEOM
 from pulser.devices._device_datacls import BaseDevice
 from pulser.json.abstract_repr.deserializer import (
     deserialize_abstract_sequence,
@@ -1139,54 +1139,35 @@ def enable_eom_mode(
             raise RuntimeError(
                 f"The '{channel}' channel is already in EOM mode."
             )
+
         channel_obj = self.declared_channels[channel]
         if not channel_obj.supports_eom():
             raise TypeError(f"Channel '{channel}' does not have an EOM.")
 
-        on_pulse = Pulse.ConstantPulse(
-            channel_obj.min_duration, amp_on, detuning_on, 0.0
+        detuning_off, switching_beams = self._process_eom_parameters(
+            channel_obj, amp_on, detuning_on, optimal_detuning_off
         )
-        stored_opt_detuning_off = optimal_detuning_off
-        if not isinstance(on_pulse, Parametrized):
-            channel_obj.validate_pulse(on_pulse)
-            amp_on = cast(float, amp_on)
-            detuning_on = cast(float, detuning_on)
-            eom_config = cast(RydbergEOM, channel_obj.eom_config)
-            if not isinstance(optimal_detuning_off, Parametrized):
-                (
-                    detuning_off,
-                    switching_beams,
-                ) = eom_config.calculate_detuning_off(
-                    amp_on,
-                    detuning_on,
-                    optimal_detuning_off,
-                    return_switching_beams=True,
-                )
-                off_pulse = Pulse.ConstantPulse(
-                    channel_obj.min_duration, 0.0, detuning_off, 0.0
-                )
-                channel_obj.validate_pulse(off_pulse)
-                # Update optimal_detuning_off to match the chosen detuning_off
-                # This minimizes the changes to the sequence when the device
-                # is switched
-                stored_opt_detuning_off = detuning_off
-
-            if not self.is_parametrized():
-                phase_drift_params = _PhaseDriftParams(
-                    drift_rate=-detuning_off,
-                    # enable_eom() calls wait for fall, so the block only
-                    # starts after fall time
-                    ti=self.get_duration(channel, include_fall_time=True),
-                )
-                self._schedule.enable_eom(
-                    channel, amp_on, detuning_on, detuning_off, switching_beams
+        if not self.is_parametrized():
+            detuning_off = cast(float, detuning_off)
+            phase_drift_params = _PhaseDriftParams(
+                drift_rate=-detuning_off,
+                # enable_eom() calls wait for fall, so the block only
+                # starts after fall time
+                ti=self.get_duration(channel, include_fall_time=True),
+            )
+            self._schedule.enable_eom(
+                channel,
+                cast(float, amp_on),
+                cast(float, detuning_on),
+                detuning_off,
+                switching_beams,
+            )
+            if correct_phase_drift:
+                buffer_slot = self._last(channel)
+                drift = phase_drift_params.calc_phase_drift(buffer_slot.tf)
+                self._phase_shift(
+                    -drift, *buffer_slot.targets, basis=channel_obj.basis
                 )
-                if correct_phase_drift:
-                    buffer_slot = self._last(channel)
-                    drift = phase_drift_params.calc_phase_drift(buffer_slot.tf)
-                    self._phase_shift(
-                        -drift, *buffer_slot.targets, basis=channel_obj.basis
-                    )
 
         # Manually store the call to "enable_eom_mode" so that the updated
         # 'optimal_detuning_off' is stored
@@ -1201,7 +1182,7 @@ def enable_eom_mode(
                     channel=channel,
                     amp_on=amp_on,
                     detuning_on=detuning_on,
-                    optimal_detuning_off=stored_opt_detuning_off,
+                    optimal_detuning_off=detuning_off,
                     correct_phase_drift=correct_phase_drift,
                 ),
             )
@@ -1253,6 +1234,90 @@ def disable_eom_mode(
                     basis=ch_schedule.channel_obj.basis,
                 )
 
+    @seq_decorators.verify_parametrization
+    @seq_decorators.block_if_measured
+    def modify_eom_setpoint(
+        self,
+        channel: str,
+        amp_on: Union[float, Parametrized],
+        detuning_on: Union[float, Parametrized],
+        optimal_detuning_off: Union[float, Parametrized] = 0.0,
+        correct_phase_drift: bool = False,
+    ) -> None:
+        """Modifies the setpoint of an ongoing EOM mode operation.
+
+        Note:
+            Modifying the EOM setpoint will automatically enforce a buffer.
+            The detuning will go to the `detuning_off` value during
+            this buffer. This buffer will not wait for pulses on other
+            channels to finish, so calling `Sequence.align()` or
+            `Sequence.delay()` beforehand is necessary to avoid eventual
+            conflicts.
+
+        Args:
+            channel: The name of the channel currently in EOM mode.
+            amp_on: The new amplitude of the EOM pulses (in rad/µs).
+            detuning_on: The new detuning of the EOM pulses (in rad/µs).
+            optimal_detuning_off: The new optimal value of detuning (in rad/µs)
+                when there is no pulse being played. It will choose the closest
+                value among the existing options.
+            correct_phase_drift: Performs a phase shift to correct for the
+                phase drift incurred while modifying the EOM setpoint.
+        """
+        if not self.is_in_eom_mode(channel):
+            raise RuntimeError(f"The '{channel}' channel is not in EOM mode.")
+
+        channel_obj = self.declared_channels[channel]
+        detuning_off, switching_beams = self._process_eom_parameters(
+            channel_obj, amp_on, detuning_on, optimal_detuning_off
+        )
+
+        if not self.is_parametrized():
+            detuning_off = cast(float, detuning_off)
+            self._schedule.disable_eom(channel, _skip_buffer=True)
+            old_phase_drift_params = self._get_last_eom_pulse_phase_drift(
+                channel
+            )
+            new_phase_drift_params = _PhaseDriftParams(
+                drift_rate=-detuning_off,
+                ti=self.get_duration(channel, include_fall_time=False),
+            )
+            self._schedule.enable_eom(
+                channel,
+                cast(float, amp_on),
+                cast(float, detuning_on),
+                detuning_off,
+                switching_beams,
+                _skip_wait_for_fall=True,
+            )
+            if correct_phase_drift:
+                buffer_slot = self._last(channel)
+                drift = old_phase_drift_params.calc_phase_drift(
+                    buffer_slot.ti
+                ) + new_phase_drift_params.calc_phase_drift(buffer_slot.tf)
+                self._phase_shift(
+                    -drift, *buffer_slot.targets, basis=channel_obj.basis
+                )
+
+        # Manually store the call to "modify_eom_setpoint" so that the updated
+        # 'optimal_detuning_off' is stored
+        call_container = (
+            self._to_build_calls if self.is_parametrized() else self._calls
+        )
+        call_container.append(
+            _Call(
+                "modify_eom_setpoint",
+                (),
+                dict(
+                    channel=channel,
+                    amp_on=amp_on,
+                    detuning_on=detuning_on,
+                    optimal_detuning_off=detuning_off,
+                    correct_phase_drift=correct_phase_drift,
+                ),
+            )
+        )
+
     @seq_decorators.store
     @seq_decorators.mark_non_empty
     @seq_decorators.block_if_measured
@@ -2389,6 +2454,43 @@ def _validate_add_protocol(self, protocol: str) -> None:
                 + ", ".join(valid_protocols)
             )
 
+    def _process_eom_parameters(
+        self,
+        channel_obj: Channel,
+        amp_on: Union[float, Parametrized],
+        detuning_on: Union[float, Parametrized],
+        optimal_detuning_off: Union[float, Parametrized],
+    ) -> tuple[float | Parametrized, tuple[RydbergBeam, ...]]:
+        on_pulse = Pulse.ConstantPulse(
+            channel_obj.min_duration, amp_on, detuning_on, 0.0
+        )
+        stored_opt_detuning_off = optimal_detuning_off
+        switching_beams: tuple[RydbergBeam, ...] = ()
+        if not isinstance(on_pulse, Parametrized):
+            channel_obj.validate_pulse(on_pulse)
+            amp_on = cast(float, amp_on)
+            detuning_on = cast(float, detuning_on)
+            eom_config = cast(RydbergEOM, channel_obj.eom_config)
+            if not isinstance(optimal_detuning_off, Parametrized):
+                (
+                    detuning_off,
+                    switching_beams,
+                ) = eom_config.calculate_detuning_off(
+                    amp_on,
+                    detuning_on,
+                    optimal_detuning_off,
+                    return_switching_beams=True,
+                )
+                off_pulse = Pulse.ConstantPulse(
+                    channel_obj.min_duration, 0.0, detuning_off, 0.0
+                )
+                channel_obj.validate_pulse(off_pulse)
+                # Update optimal_detuning_off to match the chosen detuning_off
+                # This minimizes the changes to the sequence when the device
+                # is switched
+                stored_opt_detuning_off = detuning_off
+        return stored_opt_detuning_off, switching_beams
+
     def _reset_parametrized(self) -> None:
         """Resets all attributes related to parametrization."""
         # Signals the sequence as actively "building" ie not parametrized