WIP: simplify handling of sweepers and sweep parameters

src/qibolab/instruments/zhinst.py

@@ -132,6 +132,23 @@ def select_pulse(pulse, pulse_type):
         )
 
 
+def classify_sweepers(sweepers):
+    """"""
+    nt_sweepers = []
+    rt_sweepers = []
+    for sweeper in sweepers:
+        if (
+            sweeper.parameter is Parameter.amplitude
+            and sweeper.pulses[0].type is PulseType.READOUT
+        ):
+            nt_sweepers.append(sweeper)
+        elif sweeper.parameter.name in SWEEPER_BIAS:
+            nt_sweepers.append(sweeper)
+        else:
+            rt_sweepers.append(sweeper)
+    return nt_sweepers, rt_sweepers
+
+
 @dataclass
 class ZhPort(Port):
     name: Tuple[str, str]
@@ -310,7 +327,7 @@ def __init__(self, name, device_setup, time_of_flight=0.0, smearing=0.0):
         "Sub sequences between each measurement"
 
         self.sweepers = []
-        self.nt_sweeps = None
+        self.nt_sweeps = []
         "Storing sweepers"
         # Improve the storing of multiple sweeps
 
@@ -625,14 +642,7 @@ def sequence_zh(self, sequence, qubits, couplers):
             else:
                 zhsequence[pulse.channel].append(ZhPulse(pulse))
 
-        def nt_loop(sweeper):
-            if not self.nt_sweeps:
-                self.nt_sweeps = [sweeper]
-            else:
-                self.nt_sweeps.append(sweeper)
-            self.sweepers.remove(sweeper)
-
-        for sweeper in self.sweepers.copy():
+        for sweeper in self.sweepers:
             if sweeper.parameter.name in SWEEPER_SET:
                 for pulse in sweeper.pulses:
                     if pulse.type == PulseType.READOUT:
@@ -649,7 +659,7 @@ def nt_loop(sweeper):
                         and pulse.type is PulseType.READOUT
                     ):
                         self.acquisition_type = lo.AcquisitionType.SPECTROSCOPY
-                        nt_loop(sweeper)
+
                     for element in aux_list:
                         if pulse == element.pulse:
                             if isinstance(aux_list[aux_list.index(element)], ZhPulse):
@@ -673,9 +683,6 @@ def nt_loop(sweeper):
                                         sweeper, qubits[pulse.qubit]
                                     )
 
-            if sweeper.parameter.name in SWEEPER_BIAS:
-                nt_loop(sweeper)
-
             # This may not place the Zhsweeper when the start occurs among different sections or lines
             if sweeper.parameter.name in SWEEPER_START:
                 pulse = sweeper.pulses[0]
@@ -696,6 +703,7 @@ def nt_loop(sweeper):
                             break
 
         self.sequence = zhsequence
+        self.nt_sweeps, self.sweepers = classify_sweepers(self.sweepers)
 
     def create_exp(self, qubits, couplers, options):
         """Zurich experiment initialization using their Experiment class."""
@@ -720,7 +728,7 @@ def create_exp(self, qubits, couplers, options):
 
         exp_calib = lo.Calibration()
         # Near Time recursion loop or directly to Real Time recursion loop
-        if self.nt_sweeps is not None:
+        if self.nt_sweeps:
             self.sweep_recursion_nt(qubits, couplers, exp_options, exp, exp_calib)
         else:
             self.define_exp(qubits, couplers, exp_options, exp, exp_calib)
@@ -1138,7 +1146,7 @@ def sweep(self, qubits, couplers, sequence: PulseSequence, options, *sweepers):
         """Play pulse and sweepers sequence."""
 
         self.signal_map = {}
-        self.nt_sweeps = None
+        self.nt_sweeps = []
         sweepers = list(sweepers)
         rearranging_axes, sweepers = self.rearrange_sweepers(sweepers)
         self.sweepers = sweepers
@@ -1186,9 +1194,9 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
         if sweeper.parameter is Parameter.frequency:
             for pulse in sweeper.pulses:
                 line = "drive" if pulse.type is PulseType.DRIVE else "readout"
-                zhsweeper = ZhSweeper(
-                    pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                ).zhsweeper
+                zhsweeper = ZhSweeper.select_sweeper(
+                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
+                )
                 zhsweeper.uid = "frequency"  # Changing the name from "frequency" breaks it f"frequency_{i}
                 if line == "readout":
                     channel_name = measure_channel_name(qubits[pulse.qubit])
@@ -1208,30 +1216,21 @@ def sweep_recursion(self, qubits, couplers, exp, exp_calib, exp_options):
                 aux_max = max(abs(sweeper.values))
 
                 sweeper.values /= aux_max
-                parameter = ZhSweeper(
-                    pulse, sweeper, qubits[sweeper.pulses[0].qubit]
-                ).zhsweeper
+                parameter = ZhSweeper.select_sweeper(
+                    pulse.type, sweeper, qubits[sweeper.pulses[0].qubit]
+                )
                 sweeper.values *= aux_max
 
         if sweeper.parameter is Parameter.bias:
-            if sweeper.qubits:
-                for qubit in sweeper.qubits:
-                    parameter = ZhSweeperLine(
-                        sweeper, qubit, self.sequence_qibo
-                    ).zhsweeper
-            if sweeper.couplers:
-                for qubit in sweeper.couplers:
-                    parameter = ZhSweeperLine(
-                        sweeper, qubit, self.sequence_qibo
-                    ).zhsweeper
+            parameter = ZhSweeperLine.select_sweeper(sweeper)
 
         elif sweeper.parameter is Parameter.start:
-            parameter = ZhSweeperLine(sweeper).zhsweeper
+            parameter = ZhSweeperLine.select_sweeper(sweeper)
 
         elif parameter is None:
-            parameter = ZhSweeper(
-                sweeper.pulses[0], sweeper, qubits[sweeper.pulses[0].qubit]
-            ).zhsweeper
+            parameter = ZhSweeper.select_sweeper(
+                sweeper.pulses[0].type, sweeper, qubits[sweeper.pulses[0].qubit]
+            )
 
         with exp.sweep(
             uid=f"sweep_{sweeper.parameter.name.lower()}_{i}",