Fix output array shape of SVD data processing node

qiskit_experiments/data_processing/nodes.py

@@ -99,20 +99,30 @@ class SVD(TrainableDataAction):
     """Singular Value Decomposition of averaged IQ data."""
 
     def __init__(self, validate: bool = True):
-        """
+        """Create new action.
+
         Args:
             validate: If set to False the DataAction will not validate its input.
         """
         super().__init__(validate=validate)
         self._main_axes = None
         self._means = None
         self._scales = None
+        self._n_circs = 0
+        self._n_shots = 0
+        self._n_slots = 0
+        self._n_iq = 0
 
     def _format_data(self, datum: Any, error: Optional[Any] = None) -> Tuple[Any, Any]:
         """Check that the IQ data is 2D and convert it to a numpy array.
 
         Args:
-            datum: A single item of data which corresponds to single-shot IQ data.
+            datum: All IQ data. This data has different dimensions depending on whether
+                single-shot or averaged data is being processed.
+                Single-shot data is four dimensional, i.e., ``[#circuits, #shots, #slots, 2]``,
+                while averaged IQ data is three dimensional, i.e., ``[#circuits, #slots, 2]``.
+                Here, ``#slots`` is the number of classical registers used in the circuit.
+            error: Optional, accompanied error.
 
         Returns:
             datum and any error estimate as a numpy array.
@@ -125,14 +135,32 @@ def _format_data(self, datum: Any, error: Optional[Any] = None) -> Tuple[Any, An
         if error is not None:
             error = np.asarray(error, dtype=float)
 
+        self._n_circs = 0
+        self._n_shots = 0
+        self._n_slots = 0
+        self._n_iq = 0
+
+        # identify shape
+        try:
+            # level1 single-shot data
+            self._n_circs, self._n_shots, self._n_slots, self._n_iq = datum.shape
+        except ValueError:
+            try:
+                # level1 data averaged over shots
+                self._n_circs, self._n_slots, self._n_iq = datum.shape
+            except ValueError as ex:
+                raise DataProcessorError(
+                    f"Data given to {self.__class__.__name__} is not likely level1 data."
+                ) from ex
+
         if self._validate:
-            if len(datum.shape) not in {2, 3}:
+            if self._n_iq != 2:
                 raise DataProcessorError(
-                    f"IQ data given to {self.__class__.__name__} must be a 2D array. "
-                    f"Instead, a {len(datum.shape)}D array was given."
+                    f"IQ data given to {self.__class__.__name__} does not have two-dimensions "
+                    f"(I and Q). Instead, {self._n_iq} dimensions were found."
                 )
 
-            if error is not None and len(error.shape) not in {2, 3}:
+            if error is not None and error.shape != datum.shape:
                 raise DataProcessorError(
                     f"IQ data error given to {self.__class__.__name__} must be a 2D array."
                     f"Instead, a {len(error.shape)}D array was given."
@@ -192,45 +220,46 @@ def _process(
         Raises:
             DataProcessorError: If the SVD has not been previously trained on data.
         """
-
         if not self.is_trained:
             raise DataProcessorError("SVD must be trained on data before it can be used.")
 
-        n_qubits = datum.shape[0] if len(datum.shape) == 2 else datum.shape[1]
-        processed_data = []
-
-        if error is not None:
-            processed_error = []
+        # IQ axis is reduced by projection
+        if self._n_shots == 0:
+            # level1 average mode
+            dims = self._n_circs, self._n_slots
         else:
-            processed_error = None
+            # level1 single mode
+            dims = self._n_circs, self._n_shots, self._n_slots
 
-        # process each averaged IQ point with its own axis.
-        for idx in range(n_qubits):
+        processed_data = np.zeros(dims, dtype=float)
+        error_vals = np.zeros(dims, dtype=float)
 
+        for idx in range(self._n_slots):
+            scale = self.scales[idx]
             centered = np.array(
                 [datum[..., idx, iq] - self.means(qubit=idx, iq_index=iq) for iq in [0, 1]]
             )
-
-            processed_data.append((self._main_axes[idx] @ centered) / self.scales[idx])
+            processed_data[..., idx] = (self._main_axes[idx] @ centered) / scale
 
             if error is not None:
                 angle = np.arctan(self._main_axes[idx][1] / self._main_axes[idx][0])
-                error_value = np.sqrt(
-                    (error[..., idx, 0] * np.cos(angle)) ** 2
-                    + (error[..., idx, 1] * np.sin(angle)) ** 2
+                error_vals[..., idx] = (
+                    np.sqrt(
+                        (error[..., idx, 0] * np.cos(angle)) ** 2
+                        + (error[..., idx, 1] * np.sin(angle)) ** 2
+                    )
+                    / scale
                 )
-                processed_error.append(error_value / self.scales[idx])
 
-        if len(processed_data) == 1:
+        if self._n_circs == 1:
             if error is None:
                 return processed_data[0], None
             else:
-                return processed_data[0], processed_error[0]
+                return processed_data[0], error_vals[0]
 
         if error is None:
-            return np.array(processed_data), None
-        else:
-            return np.array(processed_data), np.array(processed_error)
+            return processed_data, None
+        return processed_data, error_vals
 
     def train(self, data: List[Any]):
         """Train the SVD on the given data.
@@ -248,14 +277,14 @@ def train(self, data: List[Any]):
         if data is None:
             return
 
-        n_qubits = self._format_data(data[0])[0].shape[0]
+        data, _ = self._format_data(data)
 
         self._main_axes = []
         self._scales = []
         self._means = []
 
-        for qubit_idx in range(n_qubits):
-            datums = np.vstack([self._format_data(datum)[0][qubit_idx] for datum in data]).T
+        for qubit_idx in range(self._n_slots):
+            datums = np.vstack([datum[qubit_idx] for datum in data]).T
 
             # Calculate the mean of the data to recenter it in the IQ plane.
             mean_i = np.average(datums[0, :])

test/data_processing/test_data_processing.py

@@ -322,7 +322,7 @@ def setUp(self):
                 [[0.9, 0.9], [-1.1, 1.0]],
             ]
         )
-        self._sig_gs = np.array([[1.0], [-1.0]]) / np.sqrt(2.0)
+        self._sig_gs = np.array([1.0, -1.0]) / np.sqrt(2.0)
 
         circ_gs = ExperimentResultData(
             memory=[
@@ -332,7 +332,7 @@ def setUp(self):
                 [[-0.9, -0.9], [1.1, -1.0]],
             ]
         )
-        self._sig_es = np.array([[-1.0], [1.0]]) / np.sqrt(2.0)
+        self._sig_es = np.array([-1.0, 1.0]) / np.sqrt(2.0)
 
         circ_x90p = ExperimentResultData(
             memory=[
@@ -342,7 +342,7 @@ def setUp(self):
                 [[1.0, 1.0], [-1.0, 1.0]],
             ]
         )
-        self._sig_x90 = np.array([[0], [0]])
+        self._sig_x90 = np.array([0, 0])
 
         circ_x45p = ExperimentResultData(
             memory=[
@@ -352,7 +352,7 @@ def setUp(self):
                 [[1.0, 1.0], [-1.0, 1.0]],
             ]
         )
-        self._sig_x45 = np.array([[0.5], [-0.5]]) / np.sqrt(2.0)
+        self._sig_x45 = np.array([0.5, -0.5]) / np.sqrt(2.0)
 
         res_es = ExperimentResult(
             shots=4,
@@ -460,7 +460,7 @@ def test_process_all_data(self):
                 self._sig_x90.reshape(1, 2),
                 self._sig_x45.reshape(1, 2),
             )
-        ).T
+        )
 
         # Test processing of all data
         processed = processor(self.data.data())[0]
@@ -480,7 +480,7 @@ def test_normalize(self):
         processor.train([self.data.data(idx) for idx in [0, 1]])
         self.assertTrue(processor.is_trained)
 
-        all_expected = np.array([[0.0, 1.0, 0.5, 0.75], [1.0, 0.0, 0.5, 0.25]])
+        all_expected = np.array([[0.0, 1.0], [1.0, 0.0], [0.5, 0.5], [0.75, 0.25]])
 
         # Test processing of all data
         processed = processor(self.data.data())[0]
@@ -559,7 +559,7 @@ def test_normalize(self):
         processor.train([self.data.data(idx) for idx in [0, 1]])
         self.assertTrue(processor.is_trained)
 
-        all_expected = np.array([[0.0, 1.0, 0.5, 0.75], [1.0, 0.0, 0.5, 0.25]])
+        all_expected = np.array([[0.0, 1.0], [1.0, 0.0], [0.5, 0.5], [0.75, 0.25]])
 
         # Test processing of all data
         processed = processor(self.data.data())[0]

test/data_processing/test_nodes.py

@@ -118,15 +118,17 @@ def test_simple_data(self):
         # qubit 1 IQ data is oriented along (1, -1)
         self.assertTrue(np.allclose(iq_svd._main_axes[1], np.array([-1, 1]) / np.sqrt(2)))
 
-        processed, _ = iq_svd(np.array([[1, 1], [1, -1]]))
+        # Note: input data shape [n_circs, n_slots, n_iq] for avg mode simulation
+
+        processed, _ = iq_svd(np.array([[[1, 1], [1, -1]]]))
         expected = np.array([-1, -1]) / np.sqrt(2)
         self.assertTrue(np.allclose(processed, expected))
 
-        processed, _ = iq_svd(np.array([[2, 2], [2, -2]]))
+        processed, _ = iq_svd(np.array([[[2, 2], [2, -2]]]))
         self.assertTrue(np.allclose(processed, expected * 2))
 
         # Check that orthogonal data gives 0.
-        processed, _ = iq_svd(np.array([[1, -1], [1, 1]]))
+        processed, _ = iq_svd(np.array([[[1, -1], [1, 1]]]))
         expected = np.array([0, 0])
         self.assertTrue(np.allclose(processed, expected))
 
@@ -166,18 +168,18 @@ def test_svd_error(self):
         iq_svd._means = [[0.0, 0.0]]
 
         # Since the axis is along the real part the imaginary error is irrelevant.
-        processed, error = iq_svd([[1.0, 0.2]], [[0.2, 0.1]])
+        processed, error = iq_svd([[[1.0, 0.2]]], [[[0.2, 0.1]]])
         self.assertEqual(processed, np.array([1.0]))
         self.assertEqual(error, np.array([0.2]))
 
         # Since the axis is along the real part the imaginary error is irrelevant.
-        processed, error = iq_svd([[1.0, 0.2]], [[0.2, 0.3]])
+        processed, error = iq_svd([[[1.0, 0.2]]], [[[0.2, 0.3]]])
         self.assertEqual(processed, np.array([1.0]))
         self.assertEqual(error, np.array([0.2]))
 
         # Tilt the axis to an angle of 36.9... degrees
         iq_svd._main_axes = np.array([[0.8, 0.6]])
-        processed, error = iq_svd([[1.0, 0.0]], [[0.2, 0.3]])
+        processed, error = iq_svd([[[1.0, 0.0]]], [[[0.2, 0.3]]])
         cos_ = np.cos(np.arctan(0.6 / 0.8))
         sin_ = np.sin(np.arctan(0.6 / 0.8))
         self.assertEqual(processed, np.array([cos_]))