From fd90b5328f35c526e0a99771712b072986d06858 Mon Sep 17 00:00:00 2001
From: Alessandro Candido <candido.ale@gmail.com>
Date: Wed, 19 Feb 2025 14:55:47 +0100
Subject: [PATCH] feat: Introduce stretched waveforms generation

---
 .../instruments/qblox/sequence/sequence.py    | 14 ++++--
 .../instruments/qblox/sequence/waveforms.py   | 48 ++++++++++++++-----
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/src/qibolab/_core/instruments/qblox/sequence/sequence.py b/src/qibolab/_core/instruments/qblox/sequence/sequence.py
index 460fccc82..5529fa964 100644
--- a/src/qibolab/_core/instruments/qblox/sequence/sequence.py
+++ b/src/qibolab/_core/instruments/qblox/sequence/sequence.py
@@ -7,7 +7,7 @@
 from qibolab._core.components.configs import OscillatorConfig
 from qibolab._core.execution_parameters import ExecutionParameters
 from qibolab._core.identifier import ChannelId
-from qibolab._core.pulses import Align, PulseLike
+from qibolab._core.pulses import Align, PulseLike, Pulse, Readout
 from qibolab._core.sequence import PulseSequence
 from qibolab._core.serialize import Model
 from qibolab._core.sweeper import ParallelSweepers, Parameter, swept_pulses
@@ -79,12 +79,16 @@ def from_pulses(
         sampling_rate: float,
         channel: ChannelId,
         lo: Optional[float],
-    ):
+    ) -> "Q1Sequence":
         waveforms_ = waveforms(
             sequence,
             sampling_rate,
-            amplitude_swept=swept_pulses(sweepers, {Parameter.amplitude}),
-            duration_swept=swept_pulses(sweepers, {Parameter.duration}),
+            amplitude_swept=set(swept_pulses(sweepers, {Parameter.amplitude})),
+            duration_swept={
+                k: v
+                for k, v in swept_pulses(sweepers, {Parameter.duration}).items()
+                if isinstance(v, (Pulse, Readout))
+            },
         )
         sequence, sweepers = _apply_sampling_rate(sequence, sweepers, sampling_rate)
         sweepers = _subtract_lo(sweepers, lo) if lo is not None else sweepers
@@ -106,7 +110,7 @@ def from_pulses(
         )
 
     @classmethod
-    def empty(cls):
+    def empty(cls) -> "Q1Sequence":
         return cls(
             waveforms={}, weights={}, acquisitions={}, program=Program(elements=[])
         )
diff --git a/src/qibolab/_core/instruments/qblox/sequence/waveforms.py b/src/qibolab/_core/instruments/qblox/sequence/waveforms.py
index 0d2105b70..ce5bcfc1d 100644
--- a/src/qibolab/_core/instruments/qblox/sequence/waveforms.py
+++ b/src/qibolab/_core/instruments/qblox/sequence/waveforms.py
@@ -1,10 +1,11 @@
 from collections.abc import Iterable
-from typing import Annotated, Union
+from typing import Annotated, Optional, Union
 
 from pydantic import AfterValidator
 
 from qibolab._core.pulses import Pulse, PulseId, PulseLike, Readout
 from qibolab._core.serialize import ArrayList, Model
+from qibolab._core.sweeper import Sweeper
 
 __all__ = []
 
@@ -29,34 +30,59 @@ class WaveformSpec(Model):
 Waveforms = dict[ComponentId, Waveform]
 
 
+def _pulse(event: Union[Pulse, Readout]) -> Pulse:
+    return event.probe if isinstance(event, Readout) else event
+
+
 def waveforms(
     sequence: Iterable[PulseLike],
     sampling_rate: float,
     amplitude_swept: set[PulseId],
-    duration_swept: set[PulseId],
+    duration_swept: dict[PulseId, Sweeper],
 ) -> dict[ComponentId, WaveformSpec]:
-    def waveform(pulse: Pulse, component: str) -> WaveformSpec:
+    def _waveform(
+        pulse: Pulse, component: str, duration: Optional[float] = None
+    ) -> WaveformSpec:
         update = {"amplitude": 1.0} if pulse.id in amplitude_swept else {}
         return WaveformSpec(
             waveform=Waveform(
                 data=getattr(pulse.model_copy(update=update), component)(sampling_rate),
                 index=0,
             ),
-            duration=int(pulse.duration),
+            duration=int(pulse.duration if duration is None else duration),
         )
 
-    def pulse_(event: Union[Pulse, Readout]) -> Pulse:
-        return event.probe if isinstance(event, Readout) else event
-
+    events = {pulse.id: pulse for pulse in sequence}
     indexless = {
         k: v
         for d in (
             {
-                (pulse_uid(pulse_(event)), 0): waveform(pulse_(event), "i"),
-                (pulse_uid(pulse_(event)), 1): waveform(pulse_(event), "q"),
+                (pulse_uid(pulse), 0): _waveform(pulse, "i"),
+                (pulse_uid(pulse), 1): _waveform(pulse, "q"),
+            }
+            for pulse in (
+                _pulse(event)
+                for event in events.values()
+                if isinstance(event, (Pulse, Readout))
+            )
+            if pulse not in duration_swept
+        )
+        for k, v in d.items()
+    } | {
+        k: v
+        for d in (
+            {
+                (pulse_uid(pulse), 2 * i): _waveform(pulse, "i", duration),
+                (pulse_uid(pulse), 2 * i + 1): _waveform(pulse, "q", duration),
             }
-            for event in sequence
-            if isinstance(event, (Pulse, Readout))
+            for pulse, sweep in (
+                (_pulse(event), duration_swept[id_])
+                for id_, event in ((id_, events[id_]) for id_ in duration_swept)
+                if isinstance(event, (Pulse, Readout))
+            )
+            for i, duration in (
+                (i, sweep.irange[0] + sweep.irange[2] * i) for i in range(len(sweep))
+            )
         )
         for k, v in d.items()
     }