Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add lightshift coefficients to RydbergEOM #687

Merged
merged 6 commits into from
Jun 5, 2024
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 21 additions & 55 deletions pulser-core/pulser/channels/eom.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,15 @@ def __post_init__(self) -> None:
f" enumeration, not {self.limiting_beam}."
)

@property
def _switching_beams_combos(self) -> list[tuple[RydbergBeam, ...]]:
HGSilveri marked this conversation as resolved.
Show resolved Hide resolved
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,
Expand Down Expand Up @@ -234,47 +243,24 @@ def calculate_detuning_off(
return_switching_beams: Whether to return the beams that switch on
on and off.
"""
off_options, switching_options = self.detuning_off_options(
amp_on, detuning_on, return_switching_beams=True
)
off_options = self.detuning_off_options(amp_on, detuning_on)
closest_option = np.abs(off_options - optimal_detuning_off).argmin()
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]
return best_det_off, self._switching_beams_combos[closest_option]

@overload
def detuning_off_options(
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.
Expand All @@ -284,34 +270,15 @@ 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
# 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)
switching_beams.append(tuple(RydbergBeam))
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
det_offs = np.array(lightshifts) + offset
if not return_switching_beams:
return det_offs
return det_offs, switching_beams
return np.array(lightshifts) + offset

def _lightshift(
self, rabi_frequency: float, *beams_on: RydbergBeam
Expand Down Expand Up @@ -344,11 +311,10 @@ def _rabi_freq_per_beam(
# limit_rabi_freq is the maximum effective rabi frequency value
# below which the lightshift can be zero
if rabi_frequency <= limit_rabi_freq:
base_amp = np.sqrt(2 * rabi_frequency * self.intermediate_detuning)
sqrt_shift_factor = np.sqrt(shift_factor)
base_amp_squared = 2 * rabi_frequency * self.intermediate_detuning
return {
self.limiting_beam: base_amp / sqrt_shift_factor,
~self.limiting_beam: base_amp * sqrt_shift_factor,
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
Expand Down
32 changes: 32 additions & 0 deletions pulser-core/pulser/json/abstract_repr/schemas/device-schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": [
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": [
Expand Down
50 changes: 22 additions & 28 deletions tests/test_abstract_repr.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,37 +352,31 @@ 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,
),
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,
),
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"),
),
],
)
Expand Down
5 changes: 2 additions & 3 deletions tests/test_eom.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,9 +151,8 @@ def calc_offset(amp):
zero_det = calc_offset(amp) # detuning when both beams are off = offset
assert np.isclose(eom._lightshift(amp, *RydbergBeam), -zero_det)
assert eom._lightshift(amp) == 0.0
det_off_options, switching_beams_opts = eom.detuning_off_options(
amp, detuning_on, return_switching_beams=True
)
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
order = np.argsort(det_off_options)
Expand Down