Provide solution for indexing in general for batchable lists.

qua-platform · Jan 23, 2025 · 78e4e2c · 78e4e2c
1 parent 00609a8
commit 78e4e2c
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 26 deletions.
diff --git a/Quantum-Control-Applications-QuAM/Superconducting/calibration_graph/06_Ramsey.py b/Quantum-Control-Applications-QuAM/Superconducting/calibration_graph/06_Ramsey.py
@@ -82,38 +82,43 @@
         state = [declare(int) for _ in range(num_qubits)]
         state_st = [declare_stream() for _ in range(num_qubits)]
 
-    for i, qubit in enumerate(qubits):
-        machine.set_all_fluxes(flux_point, target=qubit)
+    for multiplexed_qubits in qubits.batch():
+        # todo: is this the right behaviour?
+        for qubit in  multiplexed_qubits:
+            machine.set_all_fluxes(flux_point, target=qubit)
 
         with for_(n, 0, n < n_avg, n + 1):
             save(n, n_st)
+
             with for_each_(idle_time, idle_times):
                 with for_(*from_array(detuning_sign, detuning_signs)):
 
-                    with if_(detuning_sign == 1):
-                        assign(virtual_detuning_phase, Cast.mul_fixed_by_int(detuning * 1e-9, 4 * idle_time))
-                    with else_():
-                        assign(virtual_detuning_phase, Cast.mul_fixed_by_int(-detuning * 1e-9, 4 * idle_time))
+                    for i, qubit in multiplexed_qubits.items():
+
+                        with if_(detuning_sign == 1):
+                            assign(virtual_detuning_phase, Cast.mul_fixed_by_int(detuning * 1e-9, 4 * idle_time))
+                        with else_():
+                            assign(virtual_detuning_phase, Cast.mul_fixed_by_int(-detuning * 1e-9, 4 * idle_time))
 
-                    qubit.align()
+                        qubit.align()
 
-                    qubit.xy.play("x90")
-                    qubit.xy.wait(idle_time)
-                    qubit.xy.frame_rotation_2pi(virtual_detuning_phase)
-                    qubit.xy.play("x90")
+                        qubit.xy.play("x90")
+                        qubit.xy.wait(idle_time)
+                        qubit.xy.frame_rotation_2pi(virtual_detuning_phase)
+                        qubit.xy.play("x90")
 
-                    qubit.align()
+                        qubit.align()
 
-                    if node.parameters.use_state_discrimination:
-                        readout_state(qubit, state[i])
-                        save(state[i], state_st[i])
-                    else:
-                        qubit.resonator.measure("readout", qua_vars=(I[i], Q[i]))
-                        save(I[i], I_st[i])
-                        save(Q[i], Q_st[i])
+                        if node.parameters.use_state_discrimination:
+                            readout_state(qubit, state[i])
+                            save(state[i], state_st[i])
+                        else:
+                            qubit.resonator.measure("readout", qua_vars=(I[i], Q[i]))
+                            save(I[i], I_st[i])
+                            save(Q[i], Q_st[i])
 
-                    qubit.resonator.wait(qubit.thermalization_time * u.ns)
-                    reset_frame(qubit.xy.name)
+                        qubit.resonator.wait(qubit.thermalization_time * u.ns)
+                        reset_frame(qubit.xy.name)
 
         # Measure sequentially
         if not node.parameters.multiplexed:

diff --git a/Quantum-Control-Applications-QuAM/Superconducting/quam_libs/components/batchable_list.py b/Quantum-Control-Applications-QuAM/Superconducting/quam_libs/components/batchable_list.py
@@ -1,4 +1,4 @@
-from typing import List, Any
+from typing import List, Any, Dict
 from collections.abc import MutableSequence
 from quam_libs.components import Transmon
 
@@ -27,8 +27,8 @@ def insert(self, index: int, value: Transmon) -> None:
     def __repr__(self) -> str:
         return repr(self._items)
 
-    def batch(self):
+    def batch(self) -> List[Dict[int, Any]]:
         if self.multiplexed:
-            return [self._items]
+            return [{i: item for i, item in enumerate(self._items)}]
         else:
-            return [[item] for item in self._items]
+            return [{i: item} for i, item in enumerate(self._items)]
diff --git a/Quantum-Control-Applications-QuAM/Superconducting/quam_libs/components/quam_root.py b/Quantum-Control-Applications-QuAM/Superconducting/quam_libs/components/quam_root.py
@@ -29,7 +29,7 @@
 
 __all__ = ["QuAM", "FEMQuAM", "OPXPlusQuAM"]
 
-from ..experiments.node import QubitExperimentNodeParameters
+from ..experiments.node_parameters import QubitsExperimentNodeParameters
 
 
 @quam_dataclass