diff --git a/pulser-core/pulser/channels/eom.py b/pulser-core/pulser/channels/eom.py index 9409a259..56bdd9f2 100644 --- a/pulser-core/pulser/channels/eom.py +++ b/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,42 @@ def __post_init__(self) -> None: f" enumeration, not {self.limiting_beam}." ) + @property + def _switching_beams_combos(self) -> list[tuple[RydbergBeam, ...]]: + switching_beams: list[tuple[RydbergBeam, ...]] = [ + (beam,) for beam in self.controlled_beams + ] + if len(self.controlled_beams) > 1 and self.multiple_beam_control: + switching_beams.append(tuple(RydbergBeam)) + return switching_beams + + @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,13 +240,20 @@ 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) 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, self._switching_beams_combos[closest_option] def detuning_off_options( - self, rabi_frequency: float, detuning_on: float + self, + rabi_frequency: float, + detuning_on: float, ) -> np.ndarray: """Calculates the possible detuning values when the amplitude is off. @@ -214,24 +270,12 @@ 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) - 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 - # switched off) on. - lightshifts = [ - self._lightshift(rabi_frequency, ~self.controlled_beams[0]) - ] - - else: - # 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) - 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) + all_beams: set[RydbergBeam] = set(RydbergBeam) + lightshifts = [] + for beams_off in self._switching_beams_combos: + # The beams that don't switch off contribute to the lightshift + beams_on: set[RydbergBeam] = all_beams - set(beams_off) + lightshifts.append(self._lightshift(rabi_frequency, *beams_on)) # We sum the offset to all lightshifts to get the effective detuning return np.array(lightshifts) + offset @@ -241,7 +285,10 @@ def _lightshift( ) -> 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 +297,25 @@ 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_squared = 2 * rabi_frequency * self.intermediate_detuning + return { + self.limiting_beam: np.sqrt(base_amp_squared / shift_factor), + ~self.limiting_beam: np.sqrt(base_amp_squared * shift_factor), + } # The limiting beam is at its maximum amplitude while the other # has the necessary amplitude to reach the desired effective rabi freq diff --git a/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json b/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json index 8c8370da..88349bd8 100644 --- a/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json +++ b/pulser-core/pulser/json/abstract_repr/schemas/device-schema.json @@ -377,6 +377,10 @@ { "additionalProperties": false, "properties": { + "blue_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" + }, "controlled_beams": { "description": "The beams that can be switched on/off with an EOM.", "items": { @@ -415,6 +419,10 @@ "multiple_beam_control": { "description": "Whether both EOMs can be used simultaneously or not.", "type": "boolean" + }, + "red_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" } }, "required": [ @@ -703,6 +711,10 @@ { "additionalProperties": false, "properties": { + "blue_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" + }, "controlled_beams": { "description": "The beams that can be switched on/off with an EOM.", "items": { @@ -741,6 +753,10 @@ "multiple_beam_control": { "description": "Whether both EOMs can be used simultaneously or not.", "type": "boolean" + }, + "red_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" } }, "required": [ @@ -1043,6 +1059,10 @@ { "additionalProperties": false, "properties": { + "blue_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" + }, "controlled_beams": { "description": "The beams that can be switched on/off with an EOM.", "items": { @@ -1081,6 +1101,10 @@ "multiple_beam_control": { "description": "Whether both EOMs can be used simultaneously or not.", "type": "boolean" + }, + "red_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" } }, "required": [ @@ -1342,6 +1366,10 @@ { "additionalProperties": false, "properties": { + "blue_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" + }, "controlled_beams": { "description": "The beams that can be switched on/off with an EOM.", "items": { @@ -1380,6 +1408,10 @@ "multiple_beam_control": { "description": "Whether both EOMs can be used simultaneously or not.", "type": "boolean" + }, + "red_shift_coeff": { + "description": "The weight coefficient of the blue beam's contribution to the lightshift.", + "type": "number" } }, "required": [ diff --git a/pulser-core/pulser/sequence/_schedule.py b/pulser-core/pulser/sequence/_schedule.py index 31edf0b1..8f2847ba 100644 --- a/pulser-core/pulser/sequence/_schedule.py +++ b/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) diff --git a/pulser-core/pulser/sequence/sequence.py b/pulser-core/pulser/sequence/sequence.py index e9efa2ab..dc12b640 100644 --- a/pulser-core/pulser/sequence/sequence.py +++ b/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) diff --git a/tests/test_abstract_repr.py b/tests/test_abstract_repr.py index 64075c8f..01439435 100644 --- a/tests/test_abstract_repr.py +++ b/tests/test_abstract_repr.py @@ -352,6 +352,32 @@ def test_optional_device_fields(self, og_device, field, value): custom_buffer_time=500, ), ), + 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, + ), + ), + 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, + ), + ), ], ) def test_optional_channel_fields(self, ch_obj): diff --git a/tests/test_eom.py b/tests/test_eom.py index 7f428b05..58f61833 100644 --- a/tests/test_eom.py +++ b/tests/test_eom.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytest from pulser.channels.eom import MODBW_TO_TR, RydbergBeam, RydbergEOM @@ -39,6 +40,10 @@ def params(): ("intermediate_detuning", 0), ("custom_buffer_time", 0.1), ("custom_buffer_time", 0), + ("blue_shift_coeff", -1e-3), + ("blue_shift_coeff", 0), + ("red_shift_coeff", -1.1), + ("red_shift_coeff", 0), ], ) def test_bad_value_init_eom(bad_param, bad_value, params): @@ -93,13 +98,32 @@ def test_bad_controlled_beam(params): assert RydbergEOM(**params).controlled_beams == tuple(RydbergBeam) +@pytest.mark.parametrize("limiting_beam", list(RydbergBeam)) +@pytest.mark.parametrize("blue_shift_coeff", [0.5, 1.0, 2.0]) +@pytest.mark.parametrize("red_shift_coeff", [0.5, 1.0, 1.8]) @pytest.mark.parametrize("multiple_beam_control", [True, False]) @pytest.mark.parametrize("limit_amp_fraction", [0.5, 2]) -def test_detuning_off(multiple_beam_control, limit_amp_fraction, params): +def test_detuning_off( + limiting_beam, + blue_shift_coeff, + red_shift_coeff, + multiple_beam_control, + limit_amp_fraction, + params, +): params["multiple_beam_control"] = multiple_beam_control + params["blue_shift_coeff"] = blue_shift_coeff + params["red_shift_coeff"] = red_shift_coeff + params["limiting_beam"] = limiting_beam eom = RydbergEOM(**params) - limit_amp = params["max_limiting_amp"] ** 2 / ( - 2 * params["intermediate_detuning"] + limit_amp = ( + params["max_limiting_amp"] ** 2 + / (2 * params["intermediate_detuning"]) + * np.sqrt( + red_shift_coeff / blue_shift_coeff + if limiting_beam == RydbergBeam.RED + else blue_shift_coeff / red_shift_coeff + ) ) amp = limit_amp_fraction * limit_amp @@ -107,30 +131,55 @@ def calc_offset(amp): # Manually calculates the offset needed to correct the lightshift # coming from a difference in power between the beams if amp <= limit_amp: - # Below limit_amp, red_amp=blue_amp so there is no lightshift + # Below limit_amp, there is no lightshift return 0.0 - assert params["limiting_beam"] == RydbergBeam.RED - red_amp = params["max_limiting_amp"] - blue_amp = 2 * params["intermediate_detuning"] * amp / red_amp - # The offset to have resonance when the pulse is on is -lightshift - return -(blue_amp**2 - red_amp**2) / ( - 4 * params["intermediate_detuning"] + limit_amp_ = params["max_limiting_amp"] + non_limit_amp = 2 * params["intermediate_detuning"] * amp / limit_amp_ + red_amp = ( + limit_amp_ if limiting_beam == RydbergBeam.RED else non_limit_amp ) + blue_amp = ( + limit_amp_ if limiting_beam == RydbergBeam.BLUE else non_limit_amp + ) + # The offset to have resonance when the pulse is on is -lightshift + return -( + blue_shift_coeff * blue_amp**2 - red_shift_coeff * red_amp**2 + ) / (4 * params["intermediate_detuning"]) # Case where the EOM pulses are resonant detuning_on = 0.0 zero_det = calc_offset(amp) # detuning when both beams are off = offset - assert eom._lightshift(amp, *RydbergBeam) == -zero_det + assert np.isclose(eom._lightshift(amp, *RydbergBeam), -zero_det) assert eom._lightshift(amp) == 0.0 det_off_options = eom.detuning_off_options(amp, detuning_on) + switching_beams_opts = eom._switching_beams_combos + assert len(det_off_options) == len(switching_beams_opts) assert len(det_off_options) == 2 + multiple_beam_control - det_off_options.sort() + order = np.argsort(det_off_options) + det_off_options = det_off_options[order] + switching_beams_opts = [switching_beams_opts[ind] for ind in order] assert det_off_options[0] < zero_det # RED on + assert switching_beams_opts[0] == (RydbergBeam.BLUE,) next_ = 1 if multiple_beam_control: - assert det_off_options[next_] == zero_det # All off + assert np.isclose(det_off_options[next_], zero_det) # All off + assert switching_beams_opts[1] == tuple(RydbergBeam) next_ += 1 assert det_off_options[next_] > zero_det # BLUE on + assert switching_beams_opts[next_] == (RydbergBeam.RED,) + calculated_det_off, switching_beams = eom.calculate_detuning_off( + amp, + detuning_on, + optimal_detuning_off=0, + return_switching_beams=True, + ) + assert ( + switching_beams + == switching_beams_opts[ + det_off_options.tolist().index(calculated_det_off) + ] + ) + assert calculated_det_off == min(det_off_options, key=abs) # Case where the EOM pulses are off-resonant detuning_on = 1.0 @@ -143,4 +192,7 @@ def calc_offset(amp): assert len(off_options) == 1 # The new detuning_off is shifted by the new detuning_on, # since that changes the offset compared the resonant case - assert off_options[0] == det_off_options[ind] + detuning_on + assert np.isclose(off_options[0], det_off_options[ind] + detuning_on) + assert off_options[0] == eom_.calculate_detuning_off( + amp, detuning_on, optimal_detuning_off=0.0 + )