Fix marginalize problems

qiskit_experiments/framework/base_analysis.py

@@ -169,7 +169,7 @@ def run_analysis(expdata: ExperimentData):
             # Clearing previous analysis data
             experiment_data._clear_results()
 
-            if not expdata.data():
+            if not expdata.data() and not expdata.child_data():
                 warnings.warn("ExperimentData object data is empty.\n")
 
             # Making new analysis

qiskit_experiments/framework/composite/composite_analysis.py

@@ -118,6 +118,7 @@ def copy(self):
     def _run_analysis(self, experiment_data: ExperimentData):
         child_data = experiment_data.child_data()
 
+<<<<<<< HEAD
 <<<<<<< HEAD
 <<<<<<< HEAD
         if len(child_data) == 0:  
@@ -152,6 +153,21 @@ def _run_analysis(self, experiment_data: ExperimentData):
 =======
             self._analyses[i].run(sub_expdata, replace_results=True)
 >>>>>>> 0bd3a186 (Updated add_data and deprecated _add_data #1268)
+=======
+        if len(self._analyses) != len(child_data):
+            # Child data is automatically created when composite result data is added.
+            # Validate that child data size matches with number of analysis entries.
+            raise RuntimeError(
+                "Number of sub-analysis and child data don't match: "
+                f"{len(self._analyses)} != {len(child_data)}. "
+                "Please check if the composite experiment and analysis are properly instantiated."
+            )
+
+        for sub_analysis, sub_data in zip(self._analyses, child_data):
+            # Since copy for replace result is handled at the parent level
+            # we always run with replace result on component analysis
+            sub_analysis.run(sub_data, replace_results=True)
+>>>>>>> a3abf4d2 (Fix marginalize problems)
 
         # Analysis is running in parallel so we add loop to wait
         # for all component analysis to finish before returning
@@ -163,7 +179,10 @@ def _run_analysis(self, experiment_data: ExperimentData):
         # for adding to the main experiment data container
         if self._flatten_results:
             analysis_results, figures = self._combine_results(child_data)
+<<<<<<< HEAD
 
+=======
+>>>>>>> a3abf4d2 (Fix marginalize problems)
             for res in analysis_results:
                 # Override experiment  ID because entries are flattened
                 res.experiment_id = experiment_data.experiment_id

qiskit_experiments/framework/experiment_data.py

@@ -23,9 +23,14 @@
 from functools import wraps, partial
 =======
 from threading import Event
+<<<<<<< HEAD
 from functools import wraps, singledispatch
 >>>>>>> 0bd3a186 (Updated add_data and deprecated _add_data #1268)
 from collections import deque, defaultdict
+=======
+from functools import wraps, singledispatch, partial
+from collections import deque
+>>>>>>> a3abf4d2 (Fix marginalize problems)
 import contextlib
 import copy
 import uuid
@@ -41,9 +46,12 @@
 from qiskit.result import Result
 from qiskit.result import marginal_distribution
 from qiskit.result.postprocess import format_counts_memory
+<<<<<<< HEAD
 from qiskit.result import marginal_distribution
 from qiskit.result.postprocess import format_counts_memory
 from qiskit.result.utils import marginal_memory
+=======
+>>>>>>> a3abf4d2 (Fix marginalize problems)
 from qiskit.providers.jobstatus import JobStatus, JOB_FINAL_STATES
 from qiskit.exceptions import QiskitError
 from qiskit.providers import Job, Backend, Provider
@@ -866,7 +874,6 @@ def __add_data(
     def add_data(
         self,
         data: Union[Result, List[Result], Dict, List[Dict]],
-        **kwargs
     ) -> None:
         """Add experiment data.
 
@@ -891,75 +898,82 @@ def add_data(
             data = [data]
 
         # Directly add non-job data
-        with self._result_data.lock and self._child_data.lock:
-
-            for datum in data:
-                if isinstance(datum, dict):
-                    if "metadata" in datum and "composite_metadata" in datum["metadata"]:
-
-                        marginalized_datum = self._marginalized_component_data([datum])
-                        composite_index = datum["metadata"]["composite_index"]
-                        max_index = max(composite_index)
-                        while max_index > len(self._child_data) -1:
-                            self.add_child_data(ExperimentData())
-                        composite_expdata = [self.child_data(i) for i in composite_index]
-                        for sub_expdata, sub_data in zip(composite_expdata, marginalized_datum):
-                            sub_expdata.add_data(sub_data)
-                        for inner_datum in datum["metadata"]["composite_metadata"]:
-                            if "composite_index" in inner_datum:
-                                for sub_expdata in composite_expdata:
-                                    self.add_data(inner_datum,inner_comoposite_flag=False)
-
-                        self._result_data.append(datum)
-
-                    elif "composite_metadata" in datum and "metadata" not in datum:
-
-                        marginalized_datum = self._marginalized_component_data([datum])
-                        composite_index = datum["composite_index"]
-                        max_index = max(composite_index)
-                        while max(composite_index) > len(self._child_data) -1:
-                            self.add_child_data(ExperimentData())
-                        composite_expdata = [self.child_data(i) for i in composite_index]
-                        for sub_expdata, sub_data in zip(composite_expdata, marginalized_datum):
-                            sub_expdata.add_data(sub_data)
-                        for inner_datum in datum["composite_metadata"]:
-                            if "composite_index" in inner_datum:
-                                for sub_expdata in composite_expdata:
-                                    self.add_data(inner_datum,inner_comoposite_flag=False)
-                    else:
-                        try:
-                            if kwargs["inner_comoposite_flag"]:
-                                self._result_data.append(datum)
-                        except KeyError:
-                            self._result_data.append(datum)
+        for datum in data:
+            if isinstance(datum, dict):
+                self._add_canonical_dict_data(datum)
+            elif isinstance(datum, Result):
+                self._add_result_data(datum)
+            else:
+                raise TypeError(f"Invalid data type {type(datum)}.")
 
-                elif isinstance(datum, Result):
-                    self._add_result_data(datum)
-                else:
-                    raise TypeError(f"Invalid data type {type(datum)}.")
+    def _add_canonical_dict_data(self, data: dict):
+        """A common subroutine to store result dictionary in canonical format.
 
+        Args:
+            data: A single formatted entry of experiment results.
+                ExperimentData expects this data dictionary to include keys such as
+                metadata, counts, memory and so forth.
+        """
+        if "metadata" in data and "composite_metadata" in data["metadata"]:
+            composite_index = data["metadata"]["composite_index"]
+            max_index = max(composite_index)
+            with self._child_data.lock:
+                while (new_idx := len(self._child_data)) <= max_index:
+                    child_data = ExperimentData()
+                    # Add automatically generated component experiment metadata
+                    try:
+                        component_metadata = self.metadata["component_metadata"][new_idx].copy()
+                        child_data.metadata.update(component_metadata)
+                    except (KeyError, IndexError):
+                        pass
+                    try:
+                        component_type = self.metadata["component_types"][new_idx]
+                        child_data.experiment_type = component_type
+                    except (KeyError, IndexError):
+                        pass
+                    self.add_child_data(child_data)
+            for idx, sub_data in self._decompose_component_data(data):
+                self.child_data(idx).add_data(sub_data)
+        else:
+            with self._result_data.lock:
+                self._result_data.append(data)
 
-    def _marginalized_component_data(self, composite_data: List[Dict]) -> List[List[Dict]]:
+    @staticmethod
+    def _decompose_component_data(
+        composite_data: dict,
+    ) -> Iterator[tuple[int, dict]]:
         """Return marginalized data for component experiments.
 
         Args:
-            composite_data: a list of composite experiment circuit data.
+            composite_data: a composite experiment result dictionary.
 
-        Returns:
-            A List of lists of marginalized circuit data for each component
-            experiment in the composite experiment.
+        Yields:
+            Tuple of composite index and result dictionary for each component experiment.
         """
-        # Marginalize data
-        marginalized_data = {}
-        for datum in composite_data:
-            metadata = datum.get("metadata", {})
+        metadata = composite_data.get("metadata", {})
 
-            # Add marginalized data to sub experiments
-            if "composite_clbits" in metadata:
-                composite_clbits = metadata["composite_clbits"]
+        tmp_sub_data = {
+            k: v for k, v in composite_data.items() if k not in ("metadata", "counts", "memory")
+        }
+        composite_clbits = metadata.get("composite_clbits", None)
+
+        if composite_clbits is not None and "memory" in composite_data:
+            # TODO use qiskit.result.utils.marginal_memory function implemented in Rust.
+            #  This function expects a complex data-type ndarray for IQ data,
+            #  while Qiskit Experiments stores IQ data in list format, i.e. [Re, Im].
+            #  This format is tied to the data processor module and we cannot easily switch.
+            #  We need to overhaul the data processor and related unit tests first.
+            memory = composite_data["memory"]
+            if isinstance(memory[0], str):
+                n_clbits = max(sum(composite_clbits, [])) + 1
+                formatter = partial(format_counts_memory, header={"memory_slots": n_clbits})
+                formatted_mem = list(map(formatter, memory))
             else:
-                composite_clbits = None
+                formatted_mem = np.array(memory, dtype=float)
+        else:
+            formatted_mem = None
 
+<<<<<<< HEAD
             # Pre-process the memory if any to avoid redundant calls to format_counts_memory
             f_memory = None
             if (
@@ -1010,6 +1024,44 @@ def _marginalized_component_data(self, composite_data: List[Dict]) -> List[List[
         # Sort by index
         return [marginalized_data[i] for i in sorted(marginalized_data.keys())]
 
+=======
+        for i, exp_idx in enumerate(metadata["composite_index"]):
+            sub_data = tmp_sub_data.copy()
+            try:
+                sub_data["metadata"] = metadata["composite_metadata"][i]
+            except (KeyError, IndexError):
+                sub_data["metadata"] = {}
+            if "counts" in composite_data:
+                if composite_clbits is not None:
+                    sub_data["counts"] = marginal_distribution(
+                        counts=composite_data["counts"],
+                        indices=composite_clbits[i],
+                    )
+                else:
+                    sub_data["counts"] = composite_data["counts"]
+            if "memory" in composite_data:
+                if isinstance(formatted_mem, list):
+                    # level 2
+                    idx = slice(-1 - composite_clbits[i][-1], -composite_clbits[i][0] or None)
+                    sub_data["memory"] = [shot[idx] for shot in formatted_mem]
+                elif isinstance(formatted_mem, np.ndarray):
+                    # level 1
+                    if len(formatted_mem.shape) == 2:
+                        # Averaged
+                        sub_data["memory"] = formatted_mem[composite_clbits[i]].tolist()
+                    elif len(formatted_mem.shape) == 3:
+                        # Single shot
+                        sub_data["memory"] = formatted_mem[:, composite_clbits[i]].tolist()
+                    else:
+                        raise ValueError(
+                            f"Invalid memory shape of {formatted_mem.shape}. "
+                            "This data cannot be marginalized."
+                        )
+                else:
+                    sub_data["memory"] = composite_data["memory"]
+            yield exp_idx, sub_data
+
+>>>>>>> a3abf4d2 (Fix marginalize problems)
     def add_jobs(
         self,
         jobs: Union[Job, List[Job]],
@@ -1265,6 +1317,10 @@ def _add_result_data(self, result: Result, job_id: Optional[str] = None) -> None
             self._jobs[job_id] = None
             self.job_ids.append(job_id)
 <<<<<<< HEAD
+<<<<<<< HEAD
+=======
+
+>>>>>>> a3abf4d2 (Fix marginalize problems)
         for i, _ in enumerate(result.results):
             data = result.data(i)
             data["job_id"] = job_id
@@ -1278,6 +1334,7 @@ def _add_result_data(self, result: Result, job_id: Optional[str] = None) -> None
             data["meas_level"] = expr_result.meas_level
             if hasattr(expr_result, "meas_return"):
                 data["meas_return"] = expr_result.meas_return
+<<<<<<< HEAD
             self.add_data(data)
 =======
         with self._result_data.lock:
@@ -1300,6 +1357,9 @@ def _add_result_data(self, result: Result, job_id: Optional[str] = None) -> None
 
             self.add_data(results)
 >>>>>>> 5e4b9d2d (Updated add_data and _add_result_data, deprecated _add_data #1268)
+=======
+            self._add_canonical_dict_data(data)
+>>>>>>> a3abf4d2 (Fix marginalize problems)
 
     def _retrieve_data(self):
         """Retrieve job data if missing experiment data."""