Add lightshift coefficients to RydbergEOM

pulser-core/pulser/channels/eom.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass, fields
 from enum import Flag
 from itertools import chain
-from typing import Any, cast
+from typing import Any, Literal, cast, overload
 
 import numpy as np
 
@@ -26,7 +26,12 @@
 # Conversion factor from modulation bandwith to rise time
 # For more info, see https://tinyurl.com/bdeumc8k
 MODBW_TO_TR = 0.48
-OPTIONAL_ABSTR_EOM_FIELDS = ("multiple_beam_control", "custom_buffer_time")
+OPTIONAL_ABSTR_EOM_FIELDS = (
+    "multiple_beam_control",
+    "custom_buffer_time",
+    "blue_shift_coeff",
+    "red_shift_coeff",
+)
 
 
 class RydbergBeam(Flag):
@@ -132,6 +137,8 @@ class _RydbergEOM:
 @dataclass(frozen=True)
 class _RydbergEOMDefaults:
     multiple_beam_control: bool = True
+    blue_shift_coeff: float = 1.0
+    red_shift_coeff: float = 1.0
 
 
 @dataclass(frozen=True)
@@ -149,11 +156,20 @@ class RydbergEOM(_RydbergEOMDefaults, BaseEOM, _RydbergEOM):
         custom_buffer_time: A custom wait time to enforce during EOM buffers.
         multiple_beam_control: Whether both EOMs can be used simultaneously.
             Ignored when only one beam can be controlled.
+        blue_shift_coeff: The weight coefficient of the blue beam's
+            contribution to the lightshift.
+        red_shift_coeff: The weight coefficient of the red beam's contribution
+            to the lightshift.
     """
 
     def __post_init__(self) -> None:
         super().__post_init__()
-        for param in ["max_limiting_amp", "intermediate_detuning"]:
+        for param in [
+            "max_limiting_amp",
+            "intermediate_detuning",
+            "blue_shift_coeff",
+            "red_shift_coeff",
+        ]:
             value = getattr(self, param)
             if value <= 0.0:
                 raise ValueError(
@@ -180,9 +196,33 @@ def __post_init__(self) -> None:
                     f" enumeration, not {self.limiting_beam}."
                 )
 
+    @overload
     def calculate_detuning_off(
-        self, amp_on: float, detuning_on: float, optimal_detuning_off: float
+        self,
+        amp_on: float,
+        detuning_on: float,
+        optimal_detuning_off: float,
+        return_switching_beams: Literal[False],
     ) -> float:
+        pass
+
+    @overload
+    def calculate_detuning_off(
+        self,
+        amp_on: float,
+        detuning_on: float,
+        optimal_detuning_off: float,
+        return_switching_beams: Literal[True],
+    ) -> tuple[float, tuple[RydbergBeam, ...]]:
+        pass
+
+    def calculate_detuning_off(
+        self,
+        amp_on: float,
+        detuning_on: float,
+        optimal_detuning_off: float,
+        return_switching_beams: bool = False,
+    ) -> float | tuple[float, tuple[RydbergBeam, ...]]:
         """Calculates the detuning when the amplitude is off in EOM mode.
 
         Args:
@@ -191,20 +231,50 @@ def calculate_detuning_off(
             optimal_detuning_off: The optimal value of detuning (in rad/µs)
                 when there is no pulse being played. It will choose the closest
                 value among the existing options.
+            return_switching_beams: Whether to return the beams that switch on
+                on and off.
         """
-        off_options = self.detuning_off_options(amp_on, detuning_on)
+        off_options, switching_options = self.detuning_off_options(
+            amp_on, detuning_on, return_switching_beams=True
+        )
         closest_option = np.abs(off_options - optimal_detuning_off).argmin()
-        return cast(float, off_options[closest_option])
+        best_det_off = cast(float, off_options[closest_option])
+        if not return_switching_beams:
+            return best_det_off
+        return best_det_off, switching_options[closest_option]
 
+    @overload
     def detuning_off_options(
-        self, rabi_frequency: float, detuning_on: float
+        self,
+        rabi_frequency: float,
+        detuning_on: float,
+        return_switching_beams: Literal[False],
     ) -> np.ndarray:
+        pass
+
+    @overload
+    def detuning_off_options(
+        self,
+        rabi_frequency: float,
+        detuning_on: float,
+        return_switching_beams: Literal[True],
+    ) -> tuple[np.ndarray, list[tuple[RydbergBeam, ...]]]:
+        pass
+
+    def detuning_off_options(
+        self,
+        rabi_frequency: float,
+        detuning_on: float,
+        return_switching_beams: bool = False,
+    ) -> np.ndarray | tuple[np.ndarray, list[tuple[RydbergBeam, ...]]]:
         """Calculates the possible detuning values when the amplitude is off.
 
         Args:
             rabi_frequency: The Rabi frequency when executing a pulse,
                 in rad/µs.
             detuning_on: The detuning when executing a pulse, in rad/µs.
+            return_switching_beams: Whether to return the beams that switch for
+                each option.
 
         Returns:
             The possible detuning values when in between pulses.
@@ -214,6 +284,9 @@ def detuning_off_options(
         # offset takes into account the lightshift when both beams are on
         # which is not zero when the Rabi freq of both beams is not equal
         offset = detuning_on - self._lightshift(rabi_frequency, *RydbergBeam)
+        switching_beams: list[tuple[RydbergBeam, ...]] = [
+            (beam,) for beam in self.controlled_beams
+        ]
         if len(self.controlled_beams) == 1:
             # When only one beam is controlled, the lighshift during delays
             # corresponds to having only the other beam (which can't be
@@ -226,22 +299,29 @@ def detuning_off_options(
             # When both beams are controlled, we have three options for the
             # lightshift: (ON, OFF), (OFF, ON) and (OFF, OFF)
             lightshifts = [
-                self._lightshift(rabi_frequency, beam)
+                self._lightshift(rabi_frequency, ~beam)
                 for beam in self.controlled_beams
             ]
             if self.multiple_beam_control:
                 # Case where both beams are off ie (OFF, OFF) -> no lightshift
                 lightshifts.append(0.0)
+                switching_beams.append(tuple(RydbergBeam))
 
         # We sum the offset to all lightshifts to get the effective detuning
-        return np.array(lightshifts) + offset
+        det_offs = np.array(lightshifts) + offset
+        if not return_switching_beams:
+            return det_offs
+        return det_offs, switching_beams
 
     def _lightshift(
         self, rabi_frequency: float, *beams_on: RydbergBeam
     ) -> float:
         # lightshift = (rabi_blue**2 - rabi_red**2) / 4 * int_detuning
         rabi_freqs = self._rabi_freq_per_beam(rabi_frequency)
-        bias = {RydbergBeam.RED: -1, RydbergBeam.BLUE: 1}
+        bias = {
+            RydbergBeam.RED: -self.red_shift_coeff,
+            RydbergBeam.BLUE: self.blue_shift_coeff,
+        }
         # beam off -> beam_rabi_freq = 0
         return sum(bias[beam] * rabi_freqs[beam] ** 2 for beam in beams_on) / (
             4 * self.intermediate_detuning
@@ -250,16 +330,26 @@ def _lightshift(
     def _rabi_freq_per_beam(
         self, rabi_frequency: float
     ) -> dict[RydbergBeam, float]:
+        shift_factor = np.sqrt(
+            self.red_shift_coeff / self.blue_shift_coeff
+            if self.limiting_beam == RydbergBeam.RED
+            else self.blue_shift_coeff / self.red_shift_coeff
+        )
         # rabi_freq = (rabi_red * rabi_blue) / (2 * int_detuning)
-        limit_rabi_freq = self.max_limiting_amp**2 / (
-            2 * self.intermediate_detuning
+        limit_rabi_freq = (
+            shift_factor
+            * self.max_limiting_amp**2
+            / (2 * self.intermediate_detuning)
         )
         # limit_rabi_freq is the maximum effective rabi frequency value
-        # below which the rabi frequency of both beams can be matched
+        # below which the lightshift can be zero
         if rabi_frequency <= limit_rabi_freq:
-            # Both beams the same rabi_freq
-            beam_amp = np.sqrt(2 * rabi_frequency * self.intermediate_detuning)
-            return {beam: beam_amp for beam in RydbergBeam}
+            base_amp = np.sqrt(2 * rabi_frequency * self.intermediate_detuning)
+            sqrt_shift_factor = np.sqrt(shift_factor)
+            return {
+                self.limiting_beam: base_amp / sqrt_shift_factor,
+                ~self.limiting_beam: base_amp * sqrt_shift_factor,
+            }
 
         # The limiting beam is at its maximum amplitude while the other
         # has the necessary amplitude to reach the desired effective rabi freq

pulser-core/pulser/sequence/_schedule.py

@@ -23,6 +23,7 @@
 
 from pulser.channels.base_channel import Channel
 from pulser.channels.dmm import DMM
+from pulser.channels.eom import RydbergBeam
 from pulser.pulse import Pulse
 from pulser.register.base_register import QubitId
 from pulser.register.weight_maps import DetuningMap
@@ -45,7 +46,8 @@ class _EOMSettings:
     detuning_on: float
     detuning_off: float
     ti: int
-    tf: Optional[int] = None
+    tf: int | None = None
+    switching_beams: tuple[RydbergBeam, ...] = ()
 
 
 @dataclass
@@ -337,6 +339,7 @@ def enable_eom(
         amp_on: float,
         detuning_on: float,
         detuning_off: float,
+        switching_beams: tuple[RydbergBeam, ...] = (),
         _skip_buffer: bool = False,
     ) -> None:
         channel_obj = self[channel_id].channel_obj
@@ -368,6 +371,7 @@ def enable_eom(
             detuning_on=detuning_on,
             detuning_off=detuning_off,
             ti=self[channel_id][-1].tf,
+            switching_beams=switching_beams,
         )
 
         self[channel_id].eom_blocks.append(eom_settings)

pulser-core/pulser/sequence/sequence.py

@@ -1149,8 +1149,14 @@ def enable_eom_mode(
             detuning_on = cast(float, detuning_on)
             eom_config = cast(RydbergEOM, channel_obj.eom_config)
             if not isinstance(optimal_detuning_off, Parametrized):
-                detuning_off = eom_config.calculate_detuning_off(
-                    amp_on, detuning_on, optimal_detuning_off
+                (
+                    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
@@ -1166,7 +1172,7 @@ def enable_eom_mode(
                     drift_rate=-detuning_off, ti=self.get_duration(channel)
                 )
                 self._schedule.enable_eom(
-                    channel, amp_on, detuning_on, detuning_off
+                    channel, amp_on, detuning_on, detuning_off, switching_beams
                 )
                 if correct_phase_drift:
                     buffer_slot = self._last(channel)

tests/test_abstract_repr.py

@@ -352,6 +352,38 @@ def test_optional_device_fields(self, og_device, field, value):
                     custom_buffer_time=500,
                 ),
             ),
+            pytest.param(
+                Rydberg.Global(
+                    None,
+                    None,
+                    mod_bandwidth=5,
+                    eom_config=RydbergEOM(
+                        max_limiting_amp=10,
+                        mod_bandwidth=20,
+                        limiting_beam=RydbergBeam.RED,
+                        intermediate_detuning=1000,
+                        controlled_beams=tuple(RydbergBeam),
+                        red_shift_coeff=1.4,
+                    ),
+                ),
+                marks=pytest.mark.xfail(reason="Needs new schema"),
+            ),
+            pytest.param(
+                Rydberg.Global(
+                    None,
+                    None,
+                    mod_bandwidth=5,
+                    eom_config=RydbergEOM(
+                        max_limiting_amp=10,
+                        mod_bandwidth=20,
+                        limiting_beam=RydbergBeam.RED,
+                        intermediate_detuning=1000,
+                        controlled_beams=tuple(RydbergBeam),
+                        blue_shift_coeff=1.4,
+                    ),
+                ),
+                marks=pytest.mark.xfail(reason="Needs new schema"),
+            ),
         ],
     )
     def test_optional_channel_fields(self, ch_obj):