From d5b8b45738b829e76f26d798d91dbf75a51bc6d5 Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Mon, 19 Feb 2024 18:09:19 +0400
Subject: [PATCH] refactor: simplify fetching

---
 src/qibolab/instruments/qm/acquisition.py | 35 ++++++-----------------
 src/qibolab/result.py                     |  2 +-
 2 files changed, 9 insertions(+), 28 deletions(-)

diff --git a/src/qibolab/instruments/qm/acquisition.py b/src/qibolab/instruments/qm/acquisition.py
index e9be966b08..c79e5a31c9 100644
--- a/src/qibolab/instruments/qm/acquisition.py
+++ b/src/qibolab/instruments/qm/acquisition.py
@@ -131,8 +131,8 @@ def measure(self, operation, element):
             qua.dual_demod.full("cos", "out1", "sin", "out2", self.i),
             qua.dual_demod.full("minus_sin", "out1", "cos", "out2", self.q),
         )
-        qua.save(self.I, self.istream)
-        qua.save(self.Q, self.qstream)
+        qua.save(self.i, self.istream)
+        qua.save(self.q, self.qstream)
 
     def download(self, *dimensions):
         istream = self.istream
@@ -153,19 +153,10 @@ def fetch(self, handles):
         ires = handles.get(f"{self.name}_I").fetch_all()
         qres = handles.get(f"{self.name}_Q").fetch_all()
         signal = ires + 1j * qres
+        res_cls = AveragedIntegratedResults if self.average else IntegratedResults
         if self.npulses > 1:
-            if self.average:
-                # TODO: calculate std
-                return [
-                    AveragedIntegratedResults(signal[..., i])
-                    for i in range(self.npulses)
-                ]
-            return [IntegratedResults(signal[..., i]) for i in range(self.npulses)]
-        else:
-            if self.average:
-                # TODO: calculate std
-                return [AveragedIntegratedResults(signal)]
-            return [IntegratedResults(signal)]
+            return [res_cls(signal[..., i]) for i in range(self.npulses)]
+        return [res_cls(signal)]
 
 
 @dataclass
@@ -221,20 +212,10 @@ def download(self, *dimensions):
 
     def fetch(self, handles):
         shots = handles.get(f"{self.name}_shots").fetch_all()
+        res_cls = AveragedSampleResults if self.average else SampleResults
         if self.npulses > 1:
-            if self.average:
-                # TODO: calculate std
-                return [
-                    AveragedSampleResults(shots[..., i]) for i in range(self.npulses)
-                ]
-            return [
-                SampleResults(shots[..., i].astype(int)) for i in range(self.npulses)
-            ]
-        else:
-            if self.average:
-                # TODO: calculate std
-                return [AveragedSampleResults(shots)]
-            return [SampleResults(shots.astype(int))]
+            return [res_cls(shots[..., i]) for i in range(self.npulses)]
+        return [res_cls(shots.astype(int))]
 
 
 ACQUISITION_TYPES = {
diff --git a/src/qibolab/result.py b/src/qibolab/result.py
index 1ee07539a6..a727578984 100644
--- a/src/qibolab/result.py
+++ b/src/qibolab/result.py
@@ -108,7 +108,7 @@ class SampleResults:
     """
 
     def __init__(self, data: np.ndarray):
-        self.samples: npt.NDArray[np.uint32] = data
+        self.samples: npt.NDArray[np.uint32] = np.array(data).astype(int)
 
     def __add__(self, data):
         return self.__class__(np.append(self.samples, data.samples))