Merge pull request #936 from qiboteam/diagonal-sweepers

Diagonal sweepers

doc/source/getting-started/experiment.rst

@@ -208,7 +208,7 @@ We leave to the dedicated tutorial a full explanation of the experiment, but her
         acquisition_type=AcquisitionType.INTEGRATION,
     )
 
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
     # plot the results
     amplitudes = results[ro_pulse.id][0].magnitude

doc/source/main-documentation/qibolab.rst

@@ -496,7 +496,7 @@ A tipical resonator spectroscopy experiment could be defined with:
         type=SweeperType.OFFSET,
     )
 
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
 .. note::
 
@@ -543,7 +543,7 @@ For example:
         type=SweeperType.FACTOR,
     )
 
-    results = platform.execute([sequence], options, sweeper_freq, sweeper_amp)
+    results = platform.execute([sequence], options, [[sweeper_freq], [sweeper_amp]])
 
 Let's say that the RX pulse has, from the runcard, a frequency of 4.5 GHz and an amplitude of 0.3, the parameter space probed will be:
 

doc/source/tutorials/calibration.rst

@@ -65,7 +65,7 @@ We then define the execution parameters and launch the experiment.
         acquisition_type=AcquisitionType.INTEGRATION,
     )
 
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
 In few seconds, the experiment will be finished and we can proceed to plot it.
 
@@ -153,7 +153,7 @@ We can now proceed to launch on hardware:
         acquisition_type=AcquisitionType.INTEGRATION,
     )
 
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
     amplitudes = results[readout_pulse.id][0].magnitude
     frequencies = np.arange(-2e8, +2e8, 1e6) + drive_pulse.frequency

src/qibolab/instruments/dummy.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Dict, List, Optional
+from typing import Dict, Optional
 
 import numpy as np
 from qibo.config import log
@@ -12,7 +12,7 @@
 )
 from qibolab.pulses import PulseSequence
 from qibolab.qubits import Qubit, QubitId
-from qibolab.sweeper import Sweeper
+from qibolab.sweeper import ParallelSweepers
 from qibolab.unrolling import Bounds
 
 from .abstract import Controller
@@ -120,16 +120,18 @@ def play(
         couplers: Dict[QubitId, Coupler],
         sequence: PulseSequence,
         options: ExecutionParameters,
-        *sweepers: List[Sweeper],
+        sweepers: list[ParallelSweepers],
     ):
         results = {}
 
         if options.averaging_mode is not AveragingMode.CYCLIC:
             shape = (options.nshots,) + tuple(
-                len(sweeper.values) for sweeper in sweepers
+                min(len(sweep.values) for sweep in parsweeps) for parsweeps in sweepers
             )
         else:
-            shape = tuple(len(sweeper.values) for sweeper in sweepers)
+            shape = tuple(
+                min(len(sweep.values) for sweep in parsweeps) for parsweeps in sweepers
+            )
 
         for ro_pulse in sequence.ro_pulses:
             values = self.get_values(options, ro_pulse, shape)

src/qibolab/platform/platform.py

@@ -2,7 +2,8 @@
 
 from collections import defaultdict
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict, List, Optional, Tuple
+from math import prod
+from typing import Any, Dict, List, Optional, Tuple, TypeVar
 
 import networkx as nx
 from qibo.config import log, raise_error
@@ -13,7 +14,7 @@
 from qibolab.pulses import Delay, Drag, PulseSequence, PulseType
 from qibolab.qubits import Qubit, QubitId, QubitPair, QubitPairId
 from qibolab.serialize_ import replace
-from qibolab.sweeper import Sweeper
+from qibolab.sweeper import ParallelSweepers
 from qibolab.unrolling import batch
 
 InstrumentMap = Dict[InstrumentId, Instrument]
@@ -23,6 +24,15 @@
 
 NS_TO_SEC = 1e-9
 
+# TODO: replace with https://docs.python.org/3/reference/compound_stmts.html#type-params
+T = TypeVar("T")
+
+
+# TODO: lift for general usage in Qibolab
+def default(value: Optional[T], default: T) -> T:
+    """None replacement shortcut."""
+    return value if value is not None else default
+
 
 def unroll_sequences(
     sequences: List[PulseSequence], relaxation_time: int
@@ -60,6 +70,23 @@ def unroll_sequences(
     return total_sequence, readout_map
 
 
+def estimate_duration(
+    sequences: list[PulseSequence],
+    options: ExecutionParameters,
+    sweepers: list[ParallelSweepers],
+) -> float:
+    """Estimate experiment duration."""
+    duration = sum(seq.duration for seq in sequences)
+    relaxation = default(options.relaxation_time, 0)
+    nshots = default(options.nshots, 0)
+    return (
+        (duration + len(sequences) * relaxation)
+        * nshots
+        * NS_TO_SEC
+        * prod(len(s[0].values) for s in sweepers)
+    )
+
+
 @dataclass
 class Settings:
     """Default execution settings read from the runcard."""
@@ -231,14 +258,14 @@ def _controller(self):
         assert len(controllers) == 1
         return controllers[0]
 
-    def _execute(self, sequence, options, *sweepers):
+    def _execute(self, sequence, options, sweepers):
         """Executes sequence on the controllers."""
         result = {}
 
         for instrument in self.instruments.values():
             if isinstance(instrument, Controller):
                 new_result = instrument.play(
-                    self.qubits, self.couplers, sequence, options, *sweepers
+                    self.qubits, self.couplers, sequence, options, sweepers
                 )
                 if isinstance(new_result, dict):
                     result.update(new_result)
@@ -249,11 +276,14 @@ def execute(
         self,
         sequences: List[PulseSequence],
         options: ExecutionParameters,
-        *sweepers: Sweeper,
+        sweepers: Optional[list[ParallelSweepers]] = None,
     ) -> dict[Any, list]:
         """Execute a pulse sequences.
 
-        If any sweeper is passed, the execution is performed for the different values of sweeped parameters.
+        If any sweeper is passed, the execution is performed for the different values
+        of sweeped parameters.
+
+        Returns readout results acquired by after execution.
 
         Example:
             .. testcode::
@@ -271,26 +301,15 @@ def execute(
                 pulse = platform.create_qubit_readout_pulse(qubit=0)
                 sequence.append(pulse)
                 parameter_range = np.random.randint(10, size=10)
-                sweeper = Sweeper(parameter, parameter_range, [pulse])
-                platform.execute([sequence], ExecutionParameters(), sweeper)
-
-        Args:
-            sequence (List[:class:`qibolab.pulses.PulseSequence`]): Pulse sequences to execute.
-            options (:class:`qibolab.platforms.platform.ExecutionParameters`): Object holding the execution options.
-            **kwargs: May need them for something
-        Returns:
-            Readout results acquired by after execution.
+                sweeper = [Sweeper(parameter, parameter_range, [pulse])]
+                platform.execute([sequence], ExecutionParameters(), [sweeper])
         """
+        if sweepers is None:
+            sweepers = []
+
         options = self.settings.fill(options)
 
-        duration = sum(seq.duration for seq in sequences)
-        time = (
-            (duration + len(sequences) * options.relaxation_time)
-            * options.nshots
-            * NS_TO_SEC
-        )
-        for sweep in sweepers:
-            time *= len(sweep.values)
+        time = estimate_duration(sequences, options, sweepers)
         log.info(f"Minimal execution time: {time}")
 
         # find readout pulses
@@ -303,7 +322,7 @@ def execute(
         results = defaultdict(list)
         for b in batch(sequences, self._controller.bounds):
             sequence, readouts = unroll_sequences(b, options.relaxation_time)
-            result = self._execute(sequence, options, *sweepers)
+            result = self._execute(sequence, options, sweepers)
             for serial, new_serials in readouts.items():
                 results[serial].extend(result[ser] for ser in new_serials)
 

src/qibolab/sweeper.py

@@ -45,9 +45,9 @@ class SweeperType(Enum):
 class Sweeper:
     """Data structure for Sweeper object.
 
-    This object is passed as an argument to the method :func:`qibolab.platforms.abstract.Platform.sweep`
+    This object is passed as an argument to the method :func:`qibolab.platforms.platform.Platform.execute`
     which enables the user to sweep a specific parameter for one or more pulses. For information on how to
-    perform sweeps see :func:`qibolab.platforms.abstract.Platform.sweep`.
+    perform sweeps see :func:`qibolab.platforms.platform.Platform.execute`.
 
     Example:
         .. testcode::
@@ -66,7 +66,7 @@ class Sweeper:
             sequence.append(pulse)
             parameter_range = np.random.randint(10, size=10)
             sweeper = Sweeper(parameter, parameter_range, [pulse])
-            platform.sweep(sequence, ExecutionParameters(), sweeper)
+            platform.execute([sequence], ExecutionParameters(), [[sweeper]])
 
     Args:
         parameter (`qibolab.sweeper.Parameter`): parameter to be swept, possible choices are frequency, attenuation, amplitude, current and gain.
@@ -111,3 +111,7 @@ def __post_init__(self):
     def get_values(self, base_value):
         """Convert sweeper values depending on the sweeper type."""
         return self.type.value(self.values, base_value)
+
+
+ParallelSweepers = list[Sweeper]
+"""Sweepers that should be iterated in parallel."""

tests/test_dummy.py

@@ -117,7 +117,7 @@ def test_dummy_single_sweep_raw(name):
         averaging_mode=AveragingMode.CYCLIC,
         acquisition_type=AcquisitionType.RAW,
     )
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
     assert pulse.id and pulse.qubit in results
     shape = results[pulse.qubit][0].magnitude.shape
     assert shape == (pulse.duration * SWEPT_POINTS,)
@@ -162,7 +162,7 @@ def test_dummy_single_sweep_coupler(
         fast_reset=fast_reset,
     )
     average = not options.averaging_mode is AveragingMode.SINGLESHOT
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
     assert ro_pulse.id and ro_pulse.qubit in results
     if average:
@@ -208,7 +208,7 @@ def test_dummy_single_sweep(name, fast_reset, parameter, average, acquisition, n
         fast_reset=fast_reset,
     )
     average = not options.averaging_mode is AveragingMode.SINGLESHOT
-    results = platform.execute([sequence], options, sweeper)
+    results = platform.execute([sequence], options, [[sweeper]])
 
     assert pulse.id and pulse.qubit in results
     if average:
@@ -270,7 +270,7 @@ def test_dummy_double_sweep(name, parameter1, parameter2, average, acquisition,
         acquisition_type=acquisition,
     )
     average = not options.averaging_mode is AveragingMode.SINGLESHOT
-    results = platform.execute([sequence], options, sweeper1, sweeper2)
+    results = platform.execute([sequence], options, [[sweeper1], [sweeper2]])
 
     assert ro_pulse.id and ro_pulse.qubit in results
 
@@ -333,7 +333,7 @@ def test_dummy_single_sweep_multiplex(name, parameter, average, acquisition, nsh
         acquisition_type=acquisition,
     )
     average = not options.averaging_mode is AveragingMode.SINGLESHOT
-    results = platform.execute([sequence], options, sweeper1)
+    results = platform.execute([sequence], options, [[sweeper1]])
 
     for ro_pulse in ro_pulses.values():
         assert ro_pulse.id and ro_pulse.qubit in results

tests/test_result_shapes.py

@@ -35,7 +35,7 @@ def execute(platform: Platform, acquisition_type, averaging_mode, sweep=False):
         sweeper1 = Sweeper(Parameter.bias, amp_values, qubits=[platform.qubits[qubit]])
         # sweeper1 = Sweeper(Parameter.amplitude, amp_values, pulses=[qd_pulse])
         sweeper2 = Sweeper(Parameter.frequency, freq_values, pulses=[ro_pulse])
-        results = platform.execute([sequence], options, sweeper1, sweeper2)
+        results = platform.execute([sequence], options, [[sweeper1], [sweeper2]])
     else:
         results = platform.execute([sequence], options)
     return results[qubit][0]