From 1c1e1f5d0178c5c0f7870281c676c1c84e6fb277 Mon Sep 17 00:00:00 2001
From: Alessandro Candido <candido.ale@gmail.com>
Date: Tue, 18 Feb 2025 16:47:09 +0100
Subject: [PATCH] fix: Subtract lo frequency from channel frequency sweepers

---
 .../_core/instruments/qblox/cluster.py        | 19 ++++++---
 .../instruments/qblox/sequence/sequence.py    | 41 ++++++++++---------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/src/qibolab/_core/instruments/qblox/cluster.py b/src/qibolab/_core/instruments/qblox/cluster.py
index 212930614..7f8830f94 100644
--- a/src/qibolab/_core/instruments/qblox/cluster.py
+++ b/src/qibolab/_core/instruments/qblox/cluster.py
@@ -89,10 +89,10 @@ def _los(self) -> dict[ChannelId, str]:
         """
         channels = self.channels
         return {
-            iq: lo
-            for iq, lo in (
-                (iq, cast(IqChannel, channels[iq]).lo)
-                for iq in (channels[ch].iqout(ch) for ch in self.channels)
+            ch: lo
+            for ch, lo in (
+                (ch, cast(IqChannel, channels[iq]).lo)
+                for ch, iq in ((ch, channels[ch].iqout(ch)) for ch in self.channels)
                 if iq is not None
             )
             if lo is not None
@@ -140,7 +140,16 @@ def play(
 
         for ps in sequences:
             assert_channels_exclusion(ps, self._probes)
-            sequences_ = compile(ps, sweepers, options, self.sampling_rate)
+            sequences_ = compile(
+                ps,
+                sweepers,
+                options,
+                self.sampling_rate,
+                {
+                    ch: cast(OscillatorConfig, configs[lo])
+                    for ch, lo in self._los.items()
+                },
+            )
             log.sequences(sequences_)
             sequencers = self._configure(sequences_, configs, options.acquisition_type)
             log.status(self.cluster, sequencers)
diff --git a/src/qibolab/_core/instruments/qblox/sequence/sequence.py b/src/qibolab/_core/instruments/qblox/sequence/sequence.py
index ab77cb4cd..fc3afa84c 100644
--- a/src/qibolab/_core/instruments/qblox/sequence/sequence.py
+++ b/src/qibolab/_core/instruments/qblox/sequence/sequence.py
@@ -4,6 +4,7 @@
 import numpy as np
 from pydantic import PlainSerializer, PlainValidator
 
+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
@@ -41,24 +42,7 @@ def _apply_sampling_rate(
     ]
     sweepers_ = [
         [
-            s.model_copy(
-                update=(
-                    {
-                        "range": tuple(t * sampling_rate for t in s.range),
-                    }
-                    if s.range is not None
-                    else {}
-                )
-                | (
-                    {
-                        "values": sampling_rate * s.values,
-                    }
-                    if s.values is not None
-                    else {}
-                )
-            )
-            if s.parameter is Parameter.duration
-            else s
+            (s * sampling_rate) if s.parameter is Parameter.duration else s
             for s in parsweep
         ]
         for parsweep in sweepers
@@ -66,6 +50,16 @@ def _apply_sampling_rate(
     return (sequence_, sweepers_)
 
 
+def _subtract_lo(sweepers: list[ParallelSweepers], lo: float) -> list[ParallelSweepers]:
+    return [
+        [
+            (sweep - lo) if sweep.parameter is Parameter.frequency else sweep
+            for sweep in parsweep
+        ]
+        for parsweep in sweepers
+    ]
+
+
 class Q1Sequence(Model):
     waveforms: Waveforms
     weights: Weights
@@ -84,6 +78,7 @@ def from_pulses(
         options: ExecutionParameters,
         sampling_rate: float,
         channel: ChannelId,
+        lo: Optional[float],
     ):
         waveforms_ = waveforms(
             sequence,
@@ -97,6 +92,7 @@ def from_pulses(
             },
         )
         sequence, sweepers = _apply_sampling_rate(sequence, sweepers, sampling_rate)
+        sweepers = _subtract_lo(sweepers, lo) if lo is not None else sweepers
         acquisitions_ = acquisitions(
             sequence, np.prod(options.bins(sweepers), dtype=int)
         )
@@ -138,13 +134,20 @@ def integration_lengths(self) -> dict[MeasureId, Optional[int]]:
         return {acq: _weight_len(self.weights.get(acq)) for acq in self.acquisitions}
 
 
+def _lo_frequency(lo: Optional[OscillatorConfig]) -> Optional[float]:
+    return lo.frequency if lo is not None else None
+
+
 def compile(
     sequence: PulseSequence,
     sweepers: list[ParallelSweepers],
     options: ExecutionParameters,
     sampling_rate: float,
+    los: dict[ChannelId, OscillatorConfig],
 ) -> dict[ChannelId, Q1Sequence]:
     return {
-        ch: Q1Sequence.from_pulses(seq, sweepers, options, sampling_rate, ch)
+        ch: Q1Sequence.from_pulses(
+            seq, sweepers, options, sampling_rate, ch, _lo_frequency(los.get(ch))
+        )
         for ch, seq in sequence.by_channel.items()
     }