Separate draw_samples and draw_sequence

pulser-core/pulser/sampler/samples.py

@@ -92,6 +92,8 @@ class ChannelSamples:
     phase: np.ndarray
     slots: list[_PulseTargetSlot] = field(default_factory=list)
     eom_blocks: list[_EOMSettings] = field(default_factory=list)
+    eom_start_buffers: list[tuple[int, int]] = field(default_factory=list)
+    eom_end_buffers: list[tuple[int, int]] = field(default_factory=list)
     target_time_slots: list[_TimeSlot] = field(default_factory=list)
 
     def __post_init__(self) -> None:

pulser-core/pulser/sequence/_schedule.py

@@ -141,6 +141,15 @@ def get_samples(
         target_time_slots: list[_TimeSlot] = [
             s for s in self.slots if s.type == "target"
         ]
+        # Extracting the EOM Buffers
+        eom_intervals_ti = [
+            eom_interval[0] for eom_interval in self.get_eom_mode_intervals()
+        ]
+        nb_eom_intervals = len(eom_intervals_ti)
+        eom_start_buffers = [(0, 0) for _ in range(nb_eom_intervals)]
+        eom_end_buffers = [(0, 0) for _ in range(nb_eom_intervals)]
+        in_eom_mode = False
+        eom_block_n = -1
 
         for ind, s in enumerate(channel_slots):
             pulse = cast(Pulse, s.type)
@@ -184,8 +193,40 @@ def get_samples(
             # the same, so the last phase is automatically kept till the end
             phase[t_start:] = pulse.phase
 
+        # Create EOM start and end buffers
+        for s in self.slots:
+            if s.ti == -1:
+                continue
+
+            # If slot is not the first element in schedule
+            if self.in_eom_mode(s):
+                # EOM mode starts
+                if not in_eom_mode:
+                    in_eom_mode = True
+                    eom_block_n += 1
+            elif in_eom_mode:
+                # Buffer when EOM mode is disabled and next slot has 0 amp
+                in_eom_mode = False
+                if amp[s.ti] == 0:
+                    eom_end_buffers[eom_block_n] = (s.ti, s.tf)
+            if (
+                eom_block_n + 1 < nb_eom_intervals
+                and s.tf == eom_intervals_ti[eom_block_n + 1]
+                and det[s.tf - 1]
+                == self.eom_blocks[eom_block_n + 1].detuning_off
+            ):
+                # Buffer if next is eom and final det matches det_off
+                eom_start_buffers[eom_block_n + 1] = (s.ti, s.tf)
+
         return ChannelSamples(
-            amp, det, phase, slots, self.eom_blocks, target_time_slots
+            amp,
+            det,
+            phase,
+            slots,
+            self.eom_blocks,
+            eom_start_buffers,
+            eom_end_buffers,
+            target_time_slots,
         )
 
     @overload

pulser-core/pulser/sequence/_seq_drawer.py

@@ -28,7 +28,6 @@
 import pulser
 from pulser import Register, Register3D
 from pulser.channels.base_channel import Channel
-from pulser.channels.eom import BaseEOM
 from pulser.pulse import Pulse
 from pulser.register.base_register import BaseRegister
 from pulser.sampler.sampler import sample
@@ -187,58 +186,23 @@ def gather_data(
             EOMSegment(eom_interval[0], eom_interval[1])
             for eom_interval in ch_samples.get_eom_mode_intervals()
         ]
-        nb_eom_intervals = len(eom_intervals)
-        eom_start_buffers = [EOMSegment() for _ in range(nb_eom_intervals)]
-        eom_end_buffers = [EOMSegment() for _ in range(nb_eom_intervals)]
-        in_eom_mode = False
-        eom_block_n = -1
         # Last eom interval is extended if eom mode not disabled at the end
         if (
-            nb_eom_intervals > 0
+            len(eom_intervals) > 0
             and ch_samples.duration == eom_intervals[-1].tf
         ):
             eom_intervals[-1].tf = total_duration
         # sampling the channel schedule
         extended_samples = ch_samples.extend_duration(total_duration)
 
-        for slot in ch_samples.slots:
-            # If slot is not the first element in schedule
-            if ch_samples.in_eom_mode(slot):
-                # EOM mode starts
-                if not in_eom_mode:
-                    in_eom_mode = True
-                    eom_block_n += 1
-            elif in_eom_mode:
-                # Buffer when EOM mode is disabled and next slot has 0 amp
-                in_eom_mode = False
-
-                tf = eom_intervals[eom_block_n].tf
-                if tf is not None and slot.ti != eom_intervals[eom_block_n].tf:
-                    if np.isclose(
-                        extended_samples.amp[
-                            tf
-                            + 2
-                            * cast(
-                                BaseEOM, sampled_seq._ch_objs[ch].eom_config
-                            ).rise_time
-                        ],
-                        0,
-                    ):
-                        eom_end_buffers[eom_block_n] = EOMSegment(tf, slot.ti)
-
-            if (
-                eom_block_n + 1 < nb_eom_intervals
-                and slot.tf == eom_intervals[eom_block_n + 1].ti
-                and np.isclose(
-                    extended_samples.det[slot.tf - 1],
-                    ch_samples.eom_blocks[eom_block_n + 1].detuning_off,
-                    0.006,
-                )
-            ):
-                # Buffer if next is eom and final det matches det_off
-                eom_start_buffers[eom_block_n + 1] = EOMSegment(
-                    slot.ti, slot.tf
-                )
+        eom_start_buffers = [
+            EOMSegment(eom_interval[0], eom_interval[1])
+            for eom_interval in ch_samples.eom_start_buffers
+        ]
+        eom_end_buffers = [
+            EOMSegment(eom_interval[0], eom_interval[1])
+            for eom_interval in ch_samples.eom_end_buffers
+        ]
 
         for time_slot in ch_samples.target_time_slots:
             if time_slot.ti == -1:

tests/test_sequence_sampler.py

@@ -240,6 +240,9 @@ def test_eom_modulation(mod_device, disable_eom):
     input_samples = sample(
         seq, extended_duration=full_duration
     ).channel_samples["ch0"]
+    assert input_samples.in_eom_mode(input_samples.slots[-1]) == (
+        not disable_eom
+    )
     mod_samples = sample(seq, modulation=True, extended_duration=full_duration)
     chan = seq.declared_channels["ch0"]
     for qty in ("amp", "det"):