From 3a48f8338d1933cee58777029cb9b750f6e1739a Mon Sep 17 00:00:00 2001
From: Paul Tan <paultan_1@hotmail.com>
Date: Thu, 24 Oct 2024 11:41:06 +0800
Subject: [PATCH] refactor: clean up result management

---
 src/qibolab/_core/instruments/keysight/qcs.py |  2 +-
 .../_core/instruments/keysight/results.py     | 33 +++++++------------
 2 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/src/qibolab/_core/instruments/keysight/qcs.py b/src/qibolab/_core/instruments/keysight/qcs.py
index 542452387..3267d76ac 100644
--- a/src/qibolab/_core/instruments/keysight/qcs.py
+++ b/src/qibolab/_core/instruments/keysight/qcs.py
@@ -164,7 +164,7 @@ def play(
 
                 for result, input_op in zip(raw.values(), input_ops):
 
-                    ret[input_op.id] = parse_result(result, options, sweepers)
+                    ret[input_op.id] = parse_result(result, options)
 
         return ret
 
diff --git a/src/qibolab/_core/instruments/keysight/results.py b/src/qibolab/_core/instruments/keysight/results.py
index 5408deb17..ff93a6949 100644
--- a/src/qibolab/_core/instruments/keysight/results.py
+++ b/src/qibolab/_core/instruments/keysight/results.py
@@ -8,7 +8,6 @@
     AveragingMode,
     ExecutionParameters,
 )
-from qibolab._core.sweeper import ParallelSweepers
 
 
 def fetch_result(
@@ -40,9 +39,7 @@ def fetch_result(
     return raw
 
 
-def parse_result(
-    result: np.ndarray, options: ExecutionParameters, sweepers: list[ParallelSweepers]
-) -> np.ndarray:
+def parse_result(result: np.ndarray, options: ExecutionParameters) -> np.ndarray:
     """Parses resulting numpy array into Qibolab expected array shape.
 
     Arguments:
@@ -55,22 +52,16 @@ def parse_result(
     """
     # For single shot, qibolab expects result format (nshots, ...)
     # QCS returns (..., nshots), so we need to shuffle the arrays
-    if options.averaging_mode is AveragingMode.SINGLESHOT and len(sweepers) > 0:
-        tmp = np.zeros(options.results_shape(sweepers))
-        # For IQ data, QCS returns complex results
-        if options.acquisition_type is AcquisitionType.INTEGRATION:
-            for k in range(options.nshots):
-                tmp[k, ..., 0] = np.real(result[..., k])
-                tmp[k, ..., 1] = np.imag(result[..., k])
-        else:
-            for k in range(options.nshots):
-                tmp[k, ...] = result[..., k]
-
+    if (
+        options.averaging_mode is not AveragingMode.SINGLESHOT
+        and options.acquisition_type is not AcquisitionType.INTEGRATION
+    ):
+        return result
+    # For single shot, qibolab expects result format (nshots, ...)
+    # QCS returns (..., nshots), so we need to shuffle the arrays
+    if options.averaging_mode is AveragingMode.SINGLESHOT:
+        result = np.moveaxis(result, -1, 0)
     # For IQ data, QCS returns complex results
-    elif options.acquisition_type is AcquisitionType.INTEGRATION:
-        tmp = np.zeros(options.results_shape(sweepers))
-        tmp[..., 0] = np.real(result)
-        tmp[..., 1] = np.imag(result)
-    else:
+    if options.acquisition_type is not AcquisitionType.INTEGRATION:
         return result
-    return tmp
+    return np.moveaxis(np.stack([np.real(result), np.imag(result)]), 0, -1)