Including detuning in the EOM buffer (#444)

* Include detuning on the EOM mode buffer

* Update the samples' modulation

* Mypy + isort

* UTs

* Bump to v0.8.1

* Fix detuning modulation when starting in EOM mode

* Updating docs and tutorials

* Replacing `is_eom_delay` with `is_detuned_delay`

* Correct phase sampling

* Improving documentation

VERSION.txt

@@ -1 +1 @@
-0.8.0
+0.8.1

pulser-core/pulser/sampler/samples.py

@@ -3,13 +3,17 @@
 
 from collections import defaultdict
 from dataclasses import dataclass, field, replace
-from typing import Optional
+from typing import TYPE_CHECKING, Optional, cast
 
 import numpy as np
 
 from pulser.channels.base_channel import Channel
+from pulser.channels.eom import BaseEOM
 from pulser.register import QubitId
 
+if TYPE_CHECKING:
+    from pulser.sequence._schedule import _EOMSettings
+
 """Literal constants for addressing."""
 _GLOBAL = "Global"
 _LOCAL = "Local"
@@ -86,8 +90,7 @@ class ChannelSamples:
     det: np.ndarray
     phase: np.ndarray
     slots: list[_TargetSlot] = field(default_factory=list)
-    # (t_start, t_end) of each EOM mode block
-    eom_intervals: list[tuple[int, int]] = field(default_factory=list)
+    eom_blocks: list[_EOMSettings] = field(default_factory=list)
 
     def __post_init__(self) -> None:
         assert len(self.amp) == len(self.det) == len(self.phase)
@@ -116,7 +119,17 @@ def extend_duration(self, new_duration: int) -> ChannelSamples:
             raise ValueError("Can't extend samples to a lower duration.")
 
         new_amp = np.pad(self.amp, (0, extension))
-        new_detuning = np.pad(self.det, (0, extension))
+        # When in EOM mode, we need to keep the detuning at detuning_off
+        if self.eom_blocks and self.eom_blocks[-1].tf is None:
+            final_detuning = self.eom_blocks[-1].detuning_off
+        else:
+            final_detuning = 0.0
+        new_detuning = np.pad(
+            self.det,
+            (0, extension),
+            constant_values=(final_detuning,),
+            mode="constant",
+        )
         new_phase = np.pad(
             self.phase,
             (0, extension),
@@ -153,29 +166,103 @@ def modulate(
 
         def masked(samples: np.ndarray, mask: np.ndarray) -> np.ndarray:
             new_samples = samples.copy()
+            # Extend the mask to fit the size of the samples
+            mask = np.pad(mask, (0, len(new_samples) - len(mask)), mode="edge")
             new_samples[~mask] = 0
             return new_samples
 
         new_samples: dict[str, np.ndarray] = {}
 
-        if self.eom_intervals:
+        std_samples = {
+            key: getattr(self, key).copy() for key in ("amp", "det")
+        }
+        eom_samples = {
+            key: getattr(self, key).copy() for key in ("amp", "det")
+        }
+
+        if self.eom_blocks:
+            # Note: self.duration already includes the fall time
             eom_mask = np.zeros(self.duration, dtype=bool)
-            for start, end in self.eom_intervals:
-                end = min(end, self.duration)  # This is defensive
-                eom_mask[np.arange(start, end)] = True
+            # Extension of the EOM mask outside of the EOM interval
+            eom_mask_ext = eom_mask.copy()
+            eom_fall_time = 2 * cast(BaseEOM, channel_obj.eom_config).rise_time
+            for block in self.eom_blocks:
+                # If block.tf is None, uses the full duration as the tf
+                end = block.tf or self.duration
+                eom_mask[block.ti : end] = True
+                std_samples["amp"][block.ti : end] = 0
+                # For modulation purposes, the detuning on the standard
+                # samples is kept at 'detuning_off', which permits a smooth
+                # transition to/from the EOM modulated samples
+                std_samples["det"][block.ti : end] = block.detuning_off
+                # Extends EOM masks to include fall time
+                ext_end = end + eom_fall_time
+                eom_mask_ext[end:ext_end] = True
+
+            # We need 'eom_mask_ext' on its own, but we can already add it
+            # to the 'eom_mask'
+            eom_mask = eom_mask + eom_mask_ext
+
+            if block.tf is None:
+                # The sequence finishes in EOM mode, so 'end' was already
+                # including the fall time (unlike when it is disabled).
+                # For modulation, we make the detuning during the last
+                # fall time to be kept at 'detuning_off'
+                eom_samples["det"][-eom_fall_time:] = block.detuning_off
 
             for key in ("amp", "det"):
-                samples = getattr(self, key)
-                std = channel_obj.modulate(masked(samples, ~eom_mask))
-                eom = channel_obj.modulate(masked(samples, eom_mask), eom=True)
+                # First, we modulated the pre-filtered standard samples, then
+                # we mask them to include only the parts outside the EOM mask
+                # This ensures smooth transitions between EOM and STD samples
+                modulated_std = channel_obj.modulate(std_samples[key])
+                std = masked(modulated_std, ~eom_mask)
+
+                # At the end of an EOM block, the EOM(s) are switched back
+                # to the OFF configuration, so the detuning should go quickly
+                # back to `detuning_off`.
+                # However, the applied detuning and the lightshift are
+                # simultaneously being ramped to zero, so the fast ramp doesn't
+                # reach `detuning_off` but rather a modified detuning value
+                # (closer to zero). Then, the detuning goes slowly
+                # to zero (as dictacted by the standard modulation bandwidth).
+                # To mimick this effect, we substitute the detuning at the end
+                # of each block by the standard modulated detuning during the
+                # transition period, so the EOM modulation is superimposed on
+                # the standard modulation
+                if key == "det":
+                    samples_ = eom_samples[key]
+                    samples_[eom_mask_ext] = modulated_std[
+                        : len(eom_mask_ext)
+                    ][eom_mask_ext]
+                    # Starts out in EOM mode, so we prepend 'detuning_off'
+                    # such that the modulation starts off from that value
+                    # We then remove the extra value after modulation
+                    if eom_mask[0]:
+                        samples_ = np.insert(
+                            samples_,
+                            0,
+                            self.eom_blocks[0].detuning_off,
+                        )
+                    # Finally, the modified EOM samples are modulated
+                    modulated_eom = channel_obj.modulate(
+                        samples_, eom=True, keep_ends=True
+                    )[(1 if eom_mask[0] else 0) :]
+                else:
+                    modulated_eom = channel_obj.modulate(
+                        eom_samples[key], eom=True
+                    )
+
+                # filtered to include only the parts inside the EOM mask
+                eom = masked(modulated_eom, eom_mask)
+
+                # 'std' and 'eom' are then summed, but before the shortest
+                # array is extended so that they are of the same length
                 sample_arrs = [std, eom]
                 sample_arrs.sort(key=len)
                 # Extend shortest array to match the longest
-                sample_arrs[0] = np.concatenate(
-                    (
-                        sample_arrs[0],
-                        np.zeros(sample_arrs[1].size - sample_arrs[0].size),
-                    )
+                sample_arrs[0] = np.pad(
+                    sample_arrs[0],
+                    (0, sample_arrs[1].size - sample_arrs[0].size),
                 )
                 new_samples[key] = sample_arrs[0] + sample_arrs[1]
 

pulser-core/pulser/sequence/_schedule.py

@@ -62,10 +62,12 @@ def last_target(self) -> int:
                 return slot.tf
         return 0  # pragma: no cover
 
-    def last_pulse_slot(self) -> _TimeSlot:
+    def last_pulse_slot(self, ignore_detuned_delay: bool = False) -> _TimeSlot:
         """The last slot with a Pulse."""
         for slot in self.slots[::-1]:
-            if isinstance(slot.type, Pulse) and not self.is_eom_delay(slot):
+            if isinstance(slot.type, Pulse) and not (
+                ignore_detuned_delay and self.is_detuned_delay(slot.type)
+            ):
                 return slot
         raise RuntimeError("There is no slot with a pulse.")
 
@@ -78,13 +80,14 @@ def in_eom_mode(self, time_slot: Optional[_TimeSlot] = None) -> bool:
             for start, end in self.get_eom_mode_intervals()
         )
 
-    def is_eom_delay(self, slot: _TimeSlot) -> bool:
-        """Tells if a pulse slot is actually an EOM delay."""
+    @staticmethod
+    def is_detuned_delay(pulse: Pulse) -> bool:
+        """Tells if a pulse is actually a delay with a constant detuning."""
         return (
-            self.in_eom_mode(time_slot=slot)
-            and isinstance(slot.type, Pulse)
-            and isinstance(slot.type.amplitude, ConstantWaveform)
-            and slot.type.amplitude[0] == 0.0
+            isinstance(pulse, Pulse)
+            and isinstance(pulse.amplitude, ConstantWaveform)
+            and pulse.amplitude[0] == 0.0
+            and isinstance(pulse.detuning, ConstantWaveform)
         )
 
     def get_eom_mode_intervals(self) -> list[tuple[int, int]]:
@@ -138,14 +141,7 @@ def get_samples(self) -> ChannelSamples:
             pulse = cast(Pulse, s.type)
             amp[s.ti : s.tf] += pulse.amplitude.samples
             det[s.ti : s.tf] += pulse.detuning.samples
-            ph_jump_t = self.channel_obj.phase_jump_time
-            t_start = s.ti - ph_jump_t if ind > 0 else 0
-            t_end = (
-                channel_slots[ind + 1].ti - ph_jump_t
-                if ind < len(channel_slots) - 1
-                else dt
-            )
-            phase[t_start:t_end] += pulse.phase
+
             tf = s.tf
             # Account for the extended duration of the pulses
             # after modulation, which is at most fall_time
@@ -157,14 +153,30 @@ def get_samples(self) -> ChannelSamples:
                 if ind < len(channel_slots) - 1
                 else fall_time
             )
-
             slots.append(_TargetSlot(s.ti, tf, s.targets))
 
-        ch_samples = ChannelSamples(
-            amp, det, phase, slots, self.get_eom_mode_intervals()
-        )
+            # The phase of detuned delays is not considered
+            if self.is_detuned_delay(pulse):
+                continue
 
-        return ch_samples
+            ph_jump_t = self.channel_obj.phase_jump_time
+            for last_pulse_ind in range(ind - 1, -1, -1):  # From ind-1 to 0
+                last_pulse_slot = channel_slots[last_pulse_ind]
+                # Skips over detuned delay pulses
+                if not self.is_detuned_delay(
+                    cast(Pulse, last_pulse_slot.type)
+                ):
+                    # Accounts for when pulse is added with 'no-delay'
+                    # i.e. there is no phase_jump_time in between a phase jump
+                    t_start = max(s.ti - ph_jump_t, last_pulse_slot.tf)
+                    break
+            else:
+                t_start = 0
+            # Overrides all values from t_start on. The next pulses will do
+            # the same, so the last phase is automatically kept till the endm
+            phase[t_start:] = pulse.phase
+
+        return ChannelSamples(amp, det, phase, slots, self.eom_blocks)
 
     @overload
     def __getitem__(self, key: int) -> _TimeSlot:
@@ -233,7 +245,20 @@ def enable_eom(
             # Account for time needed to ramp to desired amplitude
             # By definition, rise_time goes from 10% to 90%
             # Roughly 2*rise_time is enough to go from 0% to 100%
-            self.add_delay(2 * channel_obj.rise_time, channel_id)
+            if detuning_off != 0:
+                self.add_pulse(
+                    Pulse.ConstantPulse(
+                        2 * channel_obj.rise_time,
+                        0.0,
+                        detuning_off,
+                        self._get_last_pulse_phase(channel_id),
+                    ),
+                    channel_id,
+                    phase_barrier_ts=[0],
+                    protocol="no-delay",
+                )
+            else:
+                self.add_delay(2 * channel_obj.rise_time, channel_id)
 
         # Set up the EOM
         eom_settings = _EOMSettings(
@@ -268,7 +293,9 @@ def add_pulse(
             )
             try:
                 # Gets the last pulse on the channel
-                last_pulse_slot = self[channel].last_pulse_slot()
+                last_pulse_slot = self[channel].last_pulse_slot(
+                    ignore_detuned_delay=True
+                )
                 last_pulse = cast(Pulse, last_pulse_slot.type)
                 # Checks if the current pulse changes the phase
                 if last_pulse.phase != pulse.phase:
@@ -304,11 +331,7 @@ def add_delay(self, duration: int, channel: str) -> None:
             self[channel].in_eom_mode()
             and self[channel].eom_blocks[-1].detuning_off != 0
         ):
-            try:
-                last_pulse = cast(Pulse, self[channel].last_pulse_slot().type)
-                phase = last_pulse.phase
-            except RuntimeError:
-                phase = 0.0
+            phase = self._get_last_pulse_phase(channel)
             delay_pulse = Pulse.ConstantPulse(
                 tf - ti, 0.0, self[channel].eom_blocks[-1].detuning_off, phase
             )
@@ -385,3 +408,11 @@ def _find_add_delay(self, t0: int, channel: str, protocol: str) -> int:
                     break
 
         return current_max_t
+
+    def _get_last_pulse_phase(self, channel: str) -> float:
+        try:
+            last_pulse = cast(Pulse, self[channel].last_pulse_slot().type)
+            phase = last_pulse.phase
+        except RuntimeError:
+            phase = 0.0
+        return phase

pulser-core/pulser/sequence/_seq_str.py

@@ -28,7 +28,7 @@ def seq_to_str(sequence: Sequence) -> str:
     pulse_line = "t: {}->{} | {} | Targets: {}\n"
     target_line = "t: {}->{} | Target: {} | Phase Reference: {}\n"
     delay_line = "t: {}->{} | Delay \n"
-    eom_delay_line = "t: {}->{} | EOM Delay | Detuning: {:.3g} rad/µs\n"
+    det_delay_line = "t: {}->{} | Detuned Delay | Detuning: {:.3g} rad/µs\n"
     for ch, seq in sequence._schedule.items():
         basis = sequence.declared_channels[ch].basis
         full += f"Channel: {ch}\n"
@@ -46,9 +46,9 @@ def seq_to_str(sequence: Sequence) -> str:
                 )
             tgt_txt = ", ".join(map(str, tgts))
             if isinstance(ts.type, Pulse):
-                if seq.is_eom_delay(ts):
+                if seq.is_detuned_delay(ts.type):
                     det = ts.type.detuning[0]
-                    full += eom_delay_line.format(ts.ti, ts.tf, det)
+                    full += det_delay_line.format(ts.ti, ts.tf, det)
                 else:
                     full += pulse_line.format(ts.ti, ts.tf, ts.type, tgt_txt)
             elif ts.type == "target":

pulser-core/pulser/sequence/sequence.py

@@ -716,7 +716,8 @@ def enable_eom_mode(
 
         Note:
             Enabling the EOM mode will automatically enforce a buffer time from
-            the last pulse on the chose channel.
+            the last pulse on the chosen channel. The detuning will go to the
+            `detuning_off` value during this buffer.
 
         Args:
             channel: The name of the channel to put in EOM mode.
@@ -781,8 +782,8 @@ def disable_eom_mode(self, channel: str) -> None:
         (through `Sequence.add_eom_pulse()`) or delays.
 
         Note:
-            Disable the EOM mode will automatically enforce a buffer time from
-            the moment it is turned off.
+            Disabling the EOM mode will automatically enforce a buffer time
+            from the moment it is turned off.
 
         Args:
             channel: The name of the channel to take out of EOM mode.