Fixes to detuning modulation and EOM mode buffer addition (#461)

* Keep detuning ends fixed in modulation

* Fix y-axes range in Sequence.draw()

* Change criterion for adding the EOM start buffer

* Complement UTs

* Bump to v0.9.1

* Improve code documentation

VERSION.txt

@@ -1 +1 @@
-0.9.0
+0.9.1

pulser-core/pulser/sampler/samples.py

@@ -214,7 +214,9 @@ def masked(samples: np.ndarray, mask: np.ndarray) -> np.ndarray:
                 # 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])
+                modulated_std = channel_obj.modulate(
+                    std_samples[key], keep_ends=key == "det"
+                )
                 std = masked(modulated_std, ~eom_mask)
 
                 # At the end of an EOM block, the EOM(s) are switched back
@@ -268,7 +270,7 @@ def masked(samples: np.ndarray, mask: np.ndarray) -> np.ndarray:
 
         else:
             new_samples["amp"] = channel_obj.modulate(self.amp)
-            new_samples["det"] = channel_obj.modulate(self.det)
+            new_samples["det"] = channel_obj.modulate(self.det, keep_ends=True)
 
         new_samples["phase"] = channel_obj.modulate(self.phase, keep_ends=True)
         for key in new_samples:

pulser-core/pulser/sequence/_schedule.py

@@ -240,9 +240,8 @@ def enable_eom(
         _skip_buffer: bool = False,
     ) -> None:
         channel_obj = self[channel_id].channel_obj
-        if not _skip_buffer and any(
-            isinstance(op.type, Pulse) for op in self[channel_id]
-        ):
+        # 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)
             # Account for time needed to ramp to desired amplitude

pulser-core/pulser/sequence/_seq_drawer.py

@@ -16,7 +16,7 @@
 
 from collections import defaultdict
 from dataclasses import dataclass, field
-from itertools import combinations
+from itertools import chain, combinations
 from typing import Any, Optional, Union, cast
 
 import matplotlib.pyplot as plt
@@ -313,6 +313,8 @@ def phase_str(phi: float) -> str:
         ch_data = data[ch]
         basis = ch_obj.basis
         ys = ch_data.get_input_curves()
+        ys_mod = [()] * 3
+        yseff = [()] * 3
         draw_output = draw_modulation and (
             ch_obj.mod_bandwidth or not draw_input
         )
@@ -322,7 +324,10 @@ def phase_str(phi: float) -> str:
         if sampling_rate:
             curves = ys_mod if draw_output else ys
             yseff = ch_data.interpolate_curves(curves, sampling_rate)
-        ref_ys = yseff if sampling_rate else ys
+        ref_ys = [
+            list(chain.from_iterable(all_ys))
+            for all_ys in zip(ys, ys_mod, yseff)
+        ]
         max_amp = np.max(ref_ys[0])
         max_amp = 1 if max_amp == 0 else max_amp
         amp_top = max_amp * 1.2

pulser-core/pulser/sequence/sequence.py

@@ -717,9 +717,12 @@ def enable_eom_mode(
         (through `Sequence.add_eom_pulse()`) or delays.
 
         Note:
-            Enabling the EOM mode will automatically enforce a buffer time from
-            the last pulse on the chosen channel. The detuning will go to the
-            `detuning_off` value during this buffer.
+            Enabling the EOM mode will automatically enforce a buffer unless
+            the channel is empty. 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()` before enabling the EOM mode is necessary to
+            avoid eventual conflicts.
 
         Args:
             channel: The name of the channel to put in EOM mode.

tests/test_sequence.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import dataclasses
+import itertools
 import json
 from typing import Any
 from unittest.mock import patch
@@ -1327,3 +1328,53 @@ def test_eom_mode(mod_device):
     assert last_slot.type == Pulse.ConstantPulse(
         duration, 0.0, new_eom_block.detuning_off, last_pulse_slot.type.phase
     )
+
+
+@pytest.mark.parametrize(
+    "initial_instruction, non_zero_detuning_off",
+    list(itertools.product([None, "delay", "add"], [True, False])),
+)
+def test_eom_buffer(mod_device, initial_instruction, non_zero_detuning_off):
+    seq = Sequence(reg, mod_device)
+    seq.declare_channel("ch0", "rydberg_local", initial_target="q0")
+    seq.declare_channel("other", "rydberg_global")
+    if initial_instruction == "delay":
+        seq.delay(16, "ch0")
+        phase = 0
+    elif initial_instruction == "add":
+        phase = np.pi
+        seq.add(Pulse.ConstantPulse(16, 1, 0, np.pi), "ch0")
+    eom_block_starts = seq.get_duration(include_fall_time=True)
+    # Adjust the moment the EOM block starts to the clock period
+    eom_block_starts = seq._schedule["ch0"].adjust_duration(eom_block_starts)
+
+    eom_config = seq.declared_channels["ch0"].eom_config
+    limit_rabi_freq = eom_config.max_limiting_amp**2 / (
+        2 * eom_config.intermediate_detuning
+    )
+    amp_on = limit_rabi_freq * (1.1 if non_zero_detuning_off else 0.5)
+
+    # Show that EOM mode ignores other channels and uses "no-delay" by default
+    seq.add(Pulse.ConstantPulse(100, 1, -1, 0), "other")
+    seq.enable_eom_mode("ch0", amp_on, 0)
+    assert len(seq._schedule["ch0"].eom_blocks) == 1
+    eom_block = seq._schedule["ch0"].eom_blocks[0]
+    if non_zero_detuning_off:
+        assert eom_block.detuning_off != 0
+    else:
+        assert eom_block.detuning_off == 0
+    if not initial_instruction:
+        assert seq.get_duration(channel="ch0") == 0  # Channel remains empty
+    else:
+        last_slot = seq._schedule["ch0"][-1]
+        assert last_slot.ti == eom_block_starts  # Nothing else was added
+        duration = last_slot.tf - last_slot.ti
+        # The buffer is a Pulse at 'detuning_off' and zero amplitude
+        assert (
+            last_slot.type
+            == Pulse.ConstantPulse(
+                duration, 0.0, eom_block.detuning_off, phase
+            )
+            if non_zero_detuning_off
+            else "delay"
+        )

tests/test_sequence_sampler.py

@@ -126,7 +126,7 @@ def test_modulation(mod_seq: pulser.Sequence) -> None:
     got_amp = mod_samples.to_nested_dict()["Global"]["ground-rydberg"]["amp"]
     np.testing.assert_array_equal(got_amp, want_amp)
 
-    want_det = chan.modulate(np.ones(N))
+    want_det = chan.modulate(np.ones(N), keep_ends=True)
     got_det = mod_samples.to_nested_dict()["Global"]["ground-rydberg"]["det"]
     np.testing.assert_array_equal(got_det, want_det)
 
@@ -219,7 +219,9 @@ def test_eom_modulation(mod_device, disable_eom):
         samples = getattr(input_samples, qty)
         aom_input = samples.copy()
         aom_input[eom_mask] = det_off if qty == "det" else 0.0
-        aom_output = chan.modulate(aom_input, eom=False)[:full_duration]
+        aom_output = chan.modulate(
+            aom_input, eom=False, keep_ends=(qty == "det")
+        )[:full_duration]
 
         eom_input = samples.copy()
         eom_input[ext_eom_mask] = aom_output[ext_eom_mask]