Separate draw_samples and draw_sequence

pulser-core/pulser/sampler/sampler.py

@@ -56,5 +56,6 @@ def sample(
         list(seq.declared_channels.keys()),
         samples_list,
         seq.declared_channels,
+        seq._basis_ref,
         **optionals,
     )

pulser-core/pulser/sampler/samples.py

@@ -10,9 +10,10 @@
 from pulser.channels.base_channel import Channel
 from pulser.channels.eom import BaseEOM
 from pulser.register import QubitId
+from pulser.sequence._basis_ref import _QubitRef
 
 if TYPE_CHECKING:
-    from pulser.sequence._schedule import _EOMSettings
+    from pulser.sequence._schedule import _EOMSettings, _TimeSlot
 
 """Literal constants for addressing."""
 _GLOBAL = "Global"
@@ -58,7 +59,7 @@ def _default_to_regular(d: dict | defaultdict) -> dict:
 
 
 @dataclass
-class _TargetSlot:
+class _PulseTargetSlot:
     """Auxiliary class to store target information.
 
     Recopy of the sequence._TimeSlot but without the unrelevant `type` field,
@@ -89,9 +90,11 @@ class ChannelSamples:
     amp: np.ndarray
     det: np.ndarray
     phase: np.ndarray
-    slots: list[_TargetSlot] = field(default_factory=list)
+    slots: list[_PulseTargetSlot] = field(default_factory=list)
     eom_blocks: list[_EOMSettings] = field(default_factory=list)
-    initial_targets: set[QubitId] = field(default_factory=set)
+    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:
         assert len(self.amp) == len(self.det) == len(self.phase)
@@ -102,6 +105,15 @@ def __post_init__(self) -> None:
         for t1, t2 in zip(self.slots, self.slots[1:]):
             assert t1.tf <= t2.ti  # no overlaps on a given channel
 
+    @property
+    def initial_targets(self) -> set[QubitId]:
+        """Returns the initial targets."""
+        return (
+            self.target_time_slots[0].targets
+            if self.target_time_slots
+            else set()
+        )
+
     def extend_duration(self, new_duration: int) -> ChannelSamples:
         """Extends the duration of the samples.
 
@@ -160,6 +172,23 @@ def _generate_std_samples(self) -> ChannelSamples:
 
         return replace(self, **new_samples)
 
+    def get_eom_mode_intervals(self) -> list[tuple[int, int]]:
+        """Returns EOM mode intervals."""
+        return [
+            (
+                block.ti,
+                block.tf if block.tf is not None else self.duration,
+            )
+            for block in self.eom_blocks
+        ]
+
+    def in_eom_mode(self, slot: _TimeSlot | _PulseTargetSlot) -> bool:
+        """States if a time slot is inside an EOM mode block."""
+        return any(
+            start <= slot.ti < end
+            for start, end in self.get_eom_mode_intervals()
+        )
+
     def modulate(
         self, channel_obj: Channel, max_duration: Optional[int] = None
     ) -> ChannelSamples:
@@ -292,6 +321,9 @@ class SequenceSamples:
     channels: list[str]
     samples_list: list[ChannelSamples]
     _ch_objs: dict[str, Channel]
+    _basis_ref: dict[str, dict[QubitId, _QubitRef]] = field(
+        default_factory=dict
+    )
     _slm_mask: _SlmMask = field(default_factory=_SlmMask)
     _magnetic_field: np.ndarray | None = None
     _measurement: str | None = None
@@ -396,3 +428,8 @@ def __repr__(self) -> str:
             for chname, cs in zip(self.channels, self.samples_list)
         ]
         return "\n\n".join(blocks)
+
+
+# This is just to preserve backwards compatibility after the renaming of
+# _TargetSlot to _PulseTarget slot
+_TargetSlot = _PulseTargetSlot

pulser-core/pulser/sequence/_schedule.py

@@ -24,7 +24,7 @@
 from pulser.channels.base_channel import Channel
 from pulser.pulse import Pulse
 from pulser.register.base_register import QubitId
-from pulser.sampler.samples import ChannelSamples, _TargetSlot
+from pulser.sampler.samples import ChannelSamples, _PulseTargetSlot
 from pulser.waveforms import ConstantWaveform
 
 
@@ -137,8 +137,17 @@ def get_samples(
         channel_slots = [s for s in self.slots if isinstance(s.type, Pulse)]
         dt = self.get_duration()
         amp, det, phase = np.zeros(dt), np.zeros(dt), np.zeros(dt)
-        slots: list[_TargetSlot] = []
-        initial_targets = self.slots[0].targets if self.slots else set()
+        slots: list[_PulseTargetSlot] = []
+        target_time_slots: list[_TimeSlot] = [
+            s for s in self.slots if s.type == "target"
+        ]
+        # Extracting the EOM Buffers
+        eom_intervals_ti = [block.ti for block in self.eom_blocks]
+        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)
@@ -156,7 +165,7 @@ def get_samples(
                 if ind < len(channel_slots) - 1
                 else fall_time
             )
-            slots.append(_TargetSlot(s.ti, tf, s.targets))
+            slots.append(_PulseTargetSlot(s.ti, tf, s.targets))
 
             if ignore_detuned_delay_phase and self.is_detuned_delay(pulse):
                 # The phase of detuned delays is not considered
@@ -182,8 +191,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, initial_targets
+            amp,
+            det,
+            phase,
+            slots,
+            self.eom_blocks,
+            eom_start_buffers,
+            eom_end_buffers,
+            target_time_slots,
         )
 
     @overload