Refactoring _run_analysis in composite_analysis #1268

qiskit_experiments/framework/composite/composite_analysis.py

@@ -124,13 +124,27 @@ def run(
     def _run_analysis(self, experiment_data: ExperimentData):
         # Return list of experiment data containers for each component experiment
         # containing the marginalized data from the composite experiment
-        component_expdata = self._component_experiment_data(experiment_data)
+        component_expdata = []
+        if not self._flatten_results:
+            # Retrieve child data for component experiments for updating
+            component_index = experiment_data.metadata.get("component_child_index", [])
+            if not component_index:
+                raise AnalysisError("Unable to extract component child experiment data")
+            component_expdata = [experiment_data.child_data(i) for i in component_index]
+        else:
+            # Initialize temporary ExperimentData containers for
+            # each component experiment to analysis on. These will
+            # not be saved but results and figures will be collected
+            # from them
+            component_expdata = self._initialize_component_experiment_data(experiment_data)
+
+        experiment_data._add_data(experiment_data.child_data(),experiment_data.data())
 
         # Run the component analysis on each component data
         for i, sub_expdata in enumerate(component_expdata):
             # Since copy for replace result is handled at the parent level
             # we always run with replace result on component analysis
-            self._analyses[i].run(sub_expdata, replace_results=True)
+            sub_expdata._result_data = self._analyses[i].run(sub_expdata, replace_results=True)._result_data
 
         # Analysis is running in parallel so we add loop to wait
         # for all component analysis to finish before returning
@@ -147,48 +161,6 @@ def _run_analysis(self, experiment_data: ExperimentData):
             return analysis_results, figures
         return [], []
 
-    def _component_experiment_data(self, experiment_data: ExperimentData) -> List[ExperimentData]:
-        """Return a list of marginalized experiment data for component experiments.
-
-        Args:
-            experiment_data: a composite experiment data container.
-
-        Returns:
-            The list of analysis-ready marginalized experiment data for each
-            component experiment.
-
-        Raises:
-            AnalysisError: If the component experiment data cannot be extracted.
-        """
-        if not self._flatten_results:
-            # Retrieve child data for component experiments for updating
-            component_index = experiment_data.metadata.get("component_child_index", [])
-            if not component_index:
-                raise AnalysisError("Unable to extract component child experiment data")
-            component_expdata = [experiment_data.child_data(i) for i in component_index]
-        else:
-            # Initialize temporary ExperimentData containers for
-            # each component experiment to analysis on. These will
-            # not be saved but results and figures will be collected
-            # from them
-            component_expdata = self._initialize_component_experiment_data(experiment_data)
-
-        # Compute marginalize data for each component experiment
-        marginalized_data = self._marginalized_component_data(experiment_data.data())
-
-        # Add the marginalized component data and component job metadata
-        # to each component child experiment. Note that this will clear
-        # any currently stored data in the experiment. Since copying of
-        # child data is handled by the `replace_results` kwarg of the
-        # parent container it is safe to always clear and replace the
-        # results of child containers in this step
-        for sub_expdata, sub_data in zip(component_expdata, marginalized_data):
-            # Clear any previously stored data and add marginalized data
-            sub_expdata._result_data.clear()
-            sub_expdata.add_data(sub_data)
-
-        return component_expdata
-
     def _marginalized_component_data(self, composite_data: List[Dict]) -> List[List[Dict]]:
         """Return marginalized data for component experiments.
 

qiskit_experiments/framework/experiment_data.py

@@ -38,6 +38,9 @@
 from matplotlib import pyplot
 from matplotlib.figure import Figure as MatplotlibFigure
 from qiskit.result import Result
+from qiskit.result import marginal_distribution
+from qiskit.result.postprocess import format_counts_memory
+from qiskit.result.utils import marginal_memory
 from qiskit.providers.jobstatus import JobStatus, JOB_FINAL_STATES
 from qiskit.exceptions import QiskitError
 from qiskit.providers import Job, Backend, Provider
@@ -767,6 +770,49 @@ def add_data(
     ) -> None:
         """Add experiment data.
 
+        Args:
+            data: Experiment data to add. Several types are accepted for convenience:
+
+                * Result: Add data from this ``Result`` object.
+                * List[Result]: Add data from the ``Result`` objects.
+                * Dict: Add this data.
+                * List[Dict]: Add this list of data.
+
+        Raises:
+            TypeError: If the input data type is invalid.
+        """
+
+        if any(not future.done() for future in self._analysis_futures.values()):
+            LOG.warning(
+                "Not all analysis has finished running. Adding new data may "
+                "create unexpected analysis results."
+            )
+        if not isinstance(data, list):
+            data = [data]
+
+        # Directly add non-job data
+        with self._result_data.lock:
+            for datum in data:
+                if isinstance(datum, dict):
+                    if datum["metadata"].get("composite_metadata"):
+                        tmp_exp_data = ExperimentData()
+                        marginalized_data = self._marginalized_component_data([datum])
+                        for inner_datum in marginalized_data:
+                            tmp_exp_data.__add_data(inner_datum)
+                        self._set_child_data([tmp_exp_data])
+                    else:
+                        self._result_data.append(datum)
+                elif isinstance(datum, Result):
+                    self.__add_data(datum)
+                else:
+                    raise TypeError(f"Invalid data type {type(datum)}.")
+
+    def __add_data(
+        self,
+        data: Union[Result, List[Result], Dict, List[Dict]],
+    ) -> None:
+        """Add experiment data.
+
         Args:
             data: Experiment data to add. Several types are accepted for convenience:
 
@@ -792,9 +838,128 @@ def add_data(
                 if isinstance(datum, dict):
                     self._result_data.append(datum)
                 elif isinstance(datum, Result):
-                    self._add_result_data(datum)
+                    if datum["metadata"]:
+                        self._set_child_data(datum["metadata"]._metadata())
+                    else:
+                        self._add_result_data(datum)
                 else:
                     raise TypeError(f"Invalid data type {type(datum)}.")
+
+    def _add_data(
+            self,
+            component_expdata: List[ExperimentData],
+            data: Union[Result, List[Result], Dict, List[Dict]],
+        ) -> None:
+            """Add experiment data.
+
+            Args:
+                data: Experiment data to add. Several types are accepted for convenience:
+
+                    * Result: Add data from this ``Result`` object.
+                    * List[Result]: Add data from the ``Result`` objects.
+                    * Dict: Add this data.
+                    * List[Dict]: Add this list of data.
+
+            Raises:
+                TypeError: If the input data type is invalid.
+            """
+            #TODO: Continue from here
+
+            if not isinstance(data, list):
+                data = [data]
+
+            # self._marginalized_component_data()
+            # Directly add non-job data
+            marginalized_data = self._marginalized_component_data(data)
+
+            with self._result_data.lock:
+                for sub_expdata, sub_data in zip(component_expdata, marginalized_data):
+                    # Clear any previously stored data and add marginalized data
+                    sub_expdata._result_data.clear()
+                    sub_expdata.__add_data(sub_data)
+
+    def _marginalized_component_data(self, composite_data: List[Dict]) -> List[List[Dict]]:
+        """Return marginalized data for component experiments.
+
+        Args:
+            composite_data: a list of composite experiment circuit data.
+
+        Returns:
+            A List of lists of marginalized circuit data for each component
+            experiment in the composite experiment.
+        """
+        # Marginalize data
+        marginalized_data = {}
+        for datum in composite_data:
+            metadata = datum.get("metadata", {})
+
+            # Add marginalized data to sub experiments
+            if "composite_clbits" in metadata:
+                composite_clbits = metadata["composite_clbits"]
+            else:
+                composite_clbits = None
+
+            # Pre-process the memory if any to avoid redundant calls to format_counts_memory
+            f_memory = None
+            if (
+                "memory" in datum
+                and composite_clbits is not None
+                and isinstance(datum["memory"][0], str)
+            ):
+                f_memory = self._format_memory(datum, composite_clbits)
+
+            for i, index in enumerate(metadata["composite_index"]):
+                if index not in marginalized_data:
+                    # Initialize data list for marginalized
+                    marginalized_data[index] = []
+                sub_data = {"metadata": metadata["composite_metadata"][i]}
+                if "counts" in datum:
+                    if composite_clbits is not None:
+                        sub_data["counts"] = marginal_distribution(
+                            counts=datum["counts"],
+                            indices=composite_clbits[i],
+                        )
+                    else:
+                        sub_data["counts"] = datum["counts"]
+                if "memory" in datum:
+                    if composite_clbits is not None:
+                        # level 2
+                        if f_memory is not None:
+                            idx = slice(
+                                -1 - composite_clbits[i][-1], -composite_clbits[i][0] or None
+                            )
+                            sub_data["memory"] = [shot[idx] for shot in f_memory]
+                        # level 1
+                        else:
+                            mem = np.array(datum["memory"])
+
+                            # Averaged level 1 data
+                            if len(mem.shape) == 2:
+                                sub_data["memory"] = mem[composite_clbits[i]].tolist()
+                            # Single-shot level 1 data
+                            if len(mem.shape) == 3:
+                                sub_data["memory"] = mem[:, composite_clbits[i]].tolist()
+                    else:
+                        sub_data["memory"] = datum["memory"]
+                marginalized_data[index].append(sub_data)
+
+        # Sort by index
+        return [marginalized_data[i] for i in sorted(marginalized_data.keys())]
+
+    @staticmethod
+    def _format_memory(datum: Dict, composite_clbits: List):
+        """A helper method to convert level 2 memory (if it exists) to bit-string format."""
+        f_memory = None
+        if (
+            "memory" in datum
+            and composite_clbits is not None
+            and isinstance(datum["memory"][0], str)
+        ):
+            num_cbits = 1 + max(cbit for cbit_list in composite_clbits for cbit in cbit_list)
+            header = {"memory_slots": num_cbits}
+            f_memory = list(format_counts_memory(shot, header) for shot in datum["memory"])
+
+        return f_memory
 
     def add_jobs(
         self,

test/framework/test_composite.py

@@ -651,10 +651,10 @@ def test_composite_count_memory_marginalization(self, memory):
         }
 
         test_data.add_data(datum)
-
         sub_data = CompositeAnalysis([], flatten_results=False)._marginalized_component_data(
             test_data.data()
         )
+        # print(sub_data)
         expected = [
             [
                 {