From 43cc7ca7bb8b0cd9c77801339e5b5397b234f0f1 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 13 Nov 2023 23:09:36 +0400
Subject: [PATCH] Add split_batches test

---
 src/qibolab/instruments/qblox/controller.py |  5 ++++-
 src/qibolab/instruments/zhinst.py           |  4 +++-
 src/qibolab/platform.py                     |  4 ++--
 tests/test_instruments_zhinst.py            | 18 ++++++++++++++++++
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/qibolab/instruments/qblox/controller.py b/src/qibolab/instruments/qblox/controller.py
index ce0904114..8d889ee92 100644
--- a/src/qibolab/instruments/qblox/controller.py
+++ b/src/qibolab/instruments/qblox/controller.py
@@ -13,6 +13,9 @@
 from qibolab.pulses import PulseSequence, PulseType
 from qibolab.sweeper import Parameter, Sweeper, SweeperType
 
+MAX_BATCH_SIZE = 30
+"""Maximum number of sequences that can be unrolled in a single one (independent of measurements)."""
+
 
 class QbloxController(Controller):
     """A controller to manage qblox devices.
@@ -236,7 +239,7 @@ def play(self, qubits, couplers, sequence, options):
         return self._execute_pulse_sequence(qubits, sequence, options)
 
     def split_batches(self, sequences):
-        return chunked(sequences, 30)
+        return chunked(sequences, MAX_BATCH_SIZE)
 
     def play_sequences(self, *args, **kwargs):
         raise_error(NotImplementedError, "play_sequences is not implemented in qblox driver yet.")
diff --git a/src/qibolab/instruments/zhinst.py b/src/qibolab/instruments/zhinst.py
index eb303f676..4ca78f4f1 100644
--- a/src/qibolab/instruments/zhinst.py
+++ b/src/qibolab/instruments/zhinst.py
@@ -54,6 +54,9 @@
 SWEEPER_BIAS = {"bias"}
 SWEEPER_START = {"start"}
 
+MAX_MEASUREMENTS = 32
+"""Maximum number of readout pulses in a single sequence."""
+
 
 def select_pulse(pulse, pulse_type):
     """Pulse translation"""
@@ -1250,7 +1253,6 @@ def sweep_recursion_nt(self, qubits, couplers, options, exp, exp_calib):
                 self.define_exp(qubits, couplers, options, exp, exp_calib)
 
     def split_batches(self, sequences):
-        MAX_MEASUREMENTS = 34
         batch_measurements, batch = 0, []
         for sequence in sequences:
             nmeasurements = len(sequence.ro_pulses)
diff --git a/src/qibolab/platform.py b/src/qibolab/platform.py
index 74bdbe0d8..70db74975 100644
--- a/src/qibolab/platform.py
+++ b/src/qibolab/platform.py
@@ -204,7 +204,7 @@ def execute_pulse_sequence(self, sequence: PulseSequence, options: ExecutionPara
         options = self.settings.apply(options)
 
         time = (sequence.duration + options.relaxation_time) * options.nshots * 1e-9
-        log.info(f"Minimal execution time (sec): {time}")
+        log.info(f"Minimal execution time (sequence): {time}")
 
         return self._execute("play", sequence, options, **kwargs)
 
@@ -221,7 +221,7 @@ def execute_pulse_sequences(self, sequences: List[PulseSequence], options: Execu
 
         duration = sum(seq.duration for seq in sequences)
         time = (duration + len(sequences) * options.relaxation_time) * options.nshots * 1e-9
-        log.info(f"Minimal execution time (sec): {time}")
+        log.info(f"Minimal execution time (unrolling): {time}")
 
         try:
             return self._execute("play_sequences", sequences, options, **kwargs)
diff --git a/tests/test_instruments_zhinst.py b/tests/test_instruments_zhinst.py
index e08465260..b67ff0606 100644
--- a/tests/test_instruments_zhinst.py
+++ b/tests/test_instruments_zhinst.py
@@ -679,6 +679,24 @@ def test_sim(dummy_qrc):
         sequence.add(qf_pulses[qubit])
 
 
+def test_split_batches(dummy_qrc):
+    platform = create_platform("zurich")
+    platform.setup()
+    instrument = platform.controller
+
+    sequence = PulseSequence()
+    sequence.add(platform.create_RX_pulse(0, start=0))
+    sequence.add(platform.create_RX_pulse(1, start=0))
+    measurement_start = sequence.finish
+    sequence.add(platform.create_MZ_pulse(0, start=measurement_start))
+    sequence.add(platform.create_MZ_pulse(1, start=measurement_start))
+
+    batches = list(instrument.split_batches(20 * [sequence]))
+    assert len(batches) == 2
+    assert len(batches[0]) == 16
+    assert len(batches[1]) == 4
+
+
 @pytest.fixture(scope="module")
 def instrument(connected_platform):
     return get_instrument(connected_platform, Zurich)