Update composite experiment and analysis initialization

qiskit_experiments/framework/composite/batch_experiment.py

@@ -18,7 +18,9 @@
 
 from qiskit import QuantumCircuit
 from qiskit.providers.backend import Backend
+
 from .composite_experiment import CompositeExperiment, BaseExperiment
+from .composite_analysis import CompositeAnalysis
 
 
 class BatchExperiment(CompositeExperiment):
@@ -40,12 +42,20 @@ class BatchExperiment(CompositeExperiment):
     documentation for additional information.
     """
 
-    def __init__(self, experiments: List[BaseExperiment], backend: Optional[Backend] = None):
+    def __init__(
+        self,
+        experiments: List[BaseExperiment],
+        backend: Optional[Backend] = None,
+        analysis: Optional[CompositeAnalysis] = None,
+    ):
         """Initialize a batch experiment.
 
         Args:
             experiments: a list of experiments.
             backend: Optional, the backend to run the experiment on.
+            analysis: Optional, the composite analysis class to use. If not
+                      provided this will be initialized automatically from the
+                      supplied experiments.
         """
 
         # Generate qubit map
@@ -57,7 +67,7 @@ def __init__(self, experiments: List[BaseExperiment], backend: Optional[Backend]
                     self._qubit_map[physical_qubit] = logical_qubit
                     logical_qubit += 1
         qubits = tuple(self._qubit_map.keys())
-        super().__init__(experiments, qubits, backend=backend)
+        super().__init__(experiments, qubits, backend=backend, analysis=analysis)
 
     def circuits(self):
         return self._batch_circuits(to_transpile=False)

qiskit_experiments/framework/composite/composite_analysis.py

@@ -13,10 +13,11 @@
 Composite Experiment Analysis class.
 """
 
-from typing import List, Dict, Union
+from typing import List, Dict, Union, Optional
 import numpy as np
 from qiskit.result import marginal_counts
 from qiskit_experiments.framework import BaseAnalysis, ExperimentData
+from qiskit_experiments.framework.base_analysis import _requires_copy
 from qiskit_experiments.exceptions import AnalysisError
 
 
@@ -26,23 +27,26 @@ class CompositeAnalysis(BaseAnalysis):
     Composite experiments consist of several component experiments
     run together in a single execution, the results of which are returned
     as a single list of circuit result data in the :class:`ExperimentData`
-    container. Analysis of this composite circuit data involves constructing
-    a child experiment data container for each component experiment containing
-    the marginalized circuit result data for that experiment. Each component
-    child data is then analyzed using the analysis class from the corresponding
-    component experiment.
+    container.
+
+    Analysis of this composite circuit data involves constructing
+    a list of experiment data containers for each component experiment containing
+    the marginalized circuit result data for that experiment. These are saved as
+    :meth:.~ExperimentData.child_data` in the main :class:`.ExperimentData` container.
+    Each component experiment data is then analyzed using the analysis class from
+    the corresponding component experiment.
 
     .. note::
 
-        If the child :class:`ExperimentData` for each component experiment
-        does not already exist in the experiment data they will be initialized
-        and added to the experiment data when :meth:`run` is called on the
-        composite :class:`ExperimentData`.
+        If the composite :class:`ExperimentData` does not already contain
+        child experiment data containers for the component experiments
+        they will be initialized and added to the experiment data when :meth:`run`
+        is called on the composite data.
 
         When calling :meth:`run` on experiment data already containing
-        initialized component experiment child data, any previously stored
+        initialized component experiment data, any previously stored
         circuit data will be cleared and replaced with the marginalized data
-        reconstructed from the parent composite experiment data.
+        from the composite experiment data.
     """
 
     def __init__(self, analyses: List[BaseAnalysis]):
@@ -54,8 +58,19 @@ def __init__(self, analyses: List[BaseAnalysis]):
         super().__init__()
         self._analyses = analyses
 
-    def component_analysis(self, index=None) -> Union[BaseAnalysis, List[BaseAnalysis]]:
-        """Return the component experiment Analysis object"""
+    def component_analysis(
+        self, index: Optional[int] = None
+    ) -> Union[BaseAnalysis, List[BaseAnalysis]]:
+        """Return the component experiment Analysis instance.
+
+        Args:
+            index: Optional, the component index to return analysis for.
+                   If None return a list of all component analysis instances.
+
+        Returns:
+            The analysis instance for the specified index, or a list of all
+            analysis instances if index is None.
+        """
         if index is None:
             return self._analyses
         return self._analyses[index]
@@ -66,68 +81,87 @@ def copy(self):
         ret._analyses = [analysis.copy() for analysis in ret._analyses]
         return ret
 
+    def run(
+        self,
+        experiment_data: ExperimentData,
+        replace_results: bool = False,
+        **options,
+    ) -> ExperimentData:
+        # Make a new copy of experiment data if not updating results
+        if not replace_results and _requires_copy(experiment_data):
+            experiment_data = experiment_data.copy()
+
+        # Initialize child components if they are not initalized.
+        self._add_child_data(experiment_data)
+
+        # Run analysis with replace_results = True since we have already
+        # created the copy if it was required
+        return super().run(experiment_data, replace_results=True, **options)
+
     def _run_analysis(self, experiment_data: ExperimentData):
         # Return list of experiment data containers for each component experiment
         # containing the marginalied data from the composite experiment
-        component_exp_data = self._component_experiment_data(experiment_data)
+        component_expdata = self._component_experiment_data(experiment_data)
 
         # Run the component analysis on each component data
-        for sub_exp_data, sub_analysis in zip(component_exp_data, self._analyses):
+        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
-            sub_analysis.run(sub_exp_data, replace_results=True)
+            self._analyses[i].run(sub_expdata, replace_results=True)
 
-        # Wait for all component analysis to finish before returning
+        # Analysis is running in parallel so we add loop to wait
+        # for all component analysis to finish before returning
         # the parent experiment analysis results
-        for sub_exp_data in component_exp_data:
-            sub_exp_data.block_for_results()
+        for sub_expdata in component_expdata:
+            sub_expdata.block_for_results()
 
         return [], []
 
     def _component_experiment_data(self, experiment_data: ExperimentData) -> List[ExperimentData]:
-        """Return a list of component child experiment data"""
-        # Initialize component data for updating and get the experiment IDs for
-        # the component child experiments in case there are other child experiments
-        # in the experiment data
-        component_ids = self._initialize_components(experiment_data)
-        if len(component_ids) != len(self._analyses):
-            raise AnalysisError(
-                "Number of experiment components does not match number of"
-                " component analysis classes"
-            )
-
-        # Extract job metadata for the component experiments so it can be added
-        # to the child experiment data in case it is required by the child experiments
-        # analysis classes
-        component_metadata = experiment_data.metadata.get(
-            "component_metadata", [{}] * len(component_ids)
-        )
+        """Return a list of marginalized experiment data for component experiments.
+
+        Args:
+            experiment_data: a composite experiment 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.
+        """
+        # Retrieve or initialize the component data for updating
+        component_index = experiment_data.metadata.get("component_child_index", [])
+        if not component_index:
+            raise AnalysisError("Unable to extract component experiment data")
+        component_expdata = [experiment_data.child_data(i) for i in component_index]
 
         # Compute marginalize data for each component experiment
-        marginalized_data = self._component_data(experiment_data.data())
+        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
-        component_data = []
-        for i, sub_data in enumerate(marginalized_data):
-            sub_exp_data = experiment_data.child_data(component_ids[i])
-
+        for sub_expdata, sub_data in zip(component_expdata, marginalized_data):
             # Clear any previously stored data and add marginalized data
-            sub_exp_data._data.clear()
-            sub_exp_data.add_data(sub_data)
+            sub_expdata._data.clear()
+            sub_expdata.add_data(sub_data)
 
-            # Add component job metadata
-            sub_exp_data.metadata.update(component_metadata[i])
-            component_data.append(sub_exp_data)
+        return component_expdata
 
-        return component_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.
 
-    def _component_data(self, composite_data: List[Dict]) -> List[List[Dict]]:
-        """Return marginalized data for component experiments"""
+        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:
@@ -160,32 +194,62 @@ def _component_data(self, composite_data: List[Dict]) -> List[List[Dict]]:
         # Sort by index
         return [marginalized_data[i] for i in sorted(marginalized_data.keys())]
 
-    def _initialize_components(self, experiment_data: ExperimentData) -> List[str]:
-        """Initialize child data components and return list of child experiment IDs"""
-        # Check if component child experiment data containers have already
-        # been created. If so the list of indices for their positions in the
-        # ordered dict should exist. Index is used to extract the experiment
-        # IDs for each child experiment which can change when re-running analysis
-        # if replace_results=False, so that we update the correct child data
-        # for each component experiment
+    def _add_child_data(self, experiment_data: ExperimentData):
+        """Save empty component experiment data as child data.
+
+        This will initialize empty ExperimentData objects for each component
+        experiment and add them as child data to the main composite experiment
+        ExperimentData container container for saving.
+
+        Args:
+            experiment_data: a composite experiment experiment data container.
+        """
         component_index = experiment_data.metadata.get("component_child_index", [])
-        if not component_index:
-            experiment = experiment_data.experiment
-            if experiment is None:
-                raise AnalysisError(
-                    "Cannot run composite analysis on an experiment data without either "
-                    "a composite experiment, or composite experiment metadata."
-                )
-            # If the experiment Construct component data and update indices
-            start_index = len(experiment_data.child_data())
-            component_index = []
-            for i, sub_exp in enumerate(experiment.component_experiment()):
-                sub_data = sub_exp._initialize_experiment_data()
-                experiment_data.add_child_data(sub_data)
-                component_index.append(start_index + i)
-            experiment_data.metadata["component_child_index"] = component_index
-
-        # Child components exist so we can get their ID for accessing them
-        child_ids = experiment_data._child_data.keys()
-        component_ids = [child_ids[idx] for idx in component_index]
-        return component_ids
+        if component_index:
+            # Child components are already initialized
+            return
+
+        # Initialize the component experiment data containers and add them
+        # as child data to the current experiment data
+        child_components = self._initialize_component_experiment_data(experiment_data)
+        start_index = len(experiment_data.child_data())
+        for i, subdata in enumerate(child_components):
+            experiment_data.add_child_data(subdata)
+            component_index.append(start_index + i)
+
+        # Store the indices of the added child data in metadata
+        experiment_data.metadata["component_child_index"] = component_index
+
+    def _initialize_component_experiment_data(
+        self, experiment_data: ExperimentData
+    ) -> List[ExperimentData]:
+        """Initialize empty experiment data containers for component experiments.
+
+        Args:
+            experiment_data: a composite experiment experiment data container.
+
+        Returns:
+            The list of experiment data containers for each component experiment
+            containing the component metadata, and tags, share level, and
+            auto save settings of the composite experiment.
+        """
+        # Extract component experiment types and metadata so they can be
+        # added to the component experiment data containers
+        metadata = experiment_data.metadata
+        num_components = len(self._analyses)
+        experiment_types = metadata.get("component_types", [None] * num_components)
+        component_metadata = metadata.get("component_metadata", [{}] * num_components)
+
+        # Create component experiments and copy backend, tags, share level
+        # and auto save from the parent experiment data
+        component_expdata = []
+        for i, _ in enumerate(self._analyses):
+            subdata = ExperimentData(backend=experiment_data.backend)
+            subdata._type = experiment_types[i]
+            subdata.metadata.update(component_metadata[i])
+            subdata.tags = experiment_data.tags
+            subdata.share_level = experiment_data.share_level
+            subdata.auto_save = experiment_data.auto_save
+            component_expdata.append(subdata)
+
+        return component_expdata

qiskit_experiments/framework/composite/composite_experiment.py

@@ -16,8 +16,9 @@
 from typing import List, Sequence, Optional, Union
 from abc import abstractmethod
 import warnings
+from qiskit import QiskitError
 from qiskit.providers.backend import Backend
-from qiskit_experiments.framework import BaseExperiment, ExperimentData
+from qiskit_experiments.framework import BaseExperiment
 from qiskit_experiments.framework.base_analysis import BaseAnalysis
 from .composite_analysis import CompositeAnalysis
 
@@ -31,6 +32,7 @@ def __init__(
         qubits: Sequence[int],
         backend: Optional[Backend] = None,
         experiment_type: Optional[str] = None,
+        analysis: Optional[CompositeAnalysis] = None,
     ):
         """Initialize the composite experiment object.
 
@@ -39,10 +41,23 @@ def __init__(
             qubits: list of physical qubits for the experiment.
             backend: Optional, the backend to run the experiment on.
             experiment_type: Optional, composite experiment subclass name.
+            analysis: Optional, the composite analysis class to use. If not
+                      provided this will be initialized automatically from the
+                      supplied experiments.
+
+        Raises:
+            QiskitError: if the provided analysis class is not a CompositeAnalysis
+                         instance.
         """
         self._experiments = experiments
         self._num_experiments = len(experiments)
-        analysis = CompositeAnalysis([exp.analysis for exp in self._experiments])
+        if analysis is None:
+            analysis = CompositeAnalysis([exp.analysis for exp in self._experiments])
+        elif not isinstance(analysis, CompositeAnalysis):
+            raise QiskitError(
+                f"{type(analysis)} is not a CompositeAnalysis instance. CompositeExperiments"
+                " require a CompositeAnalysis class or subclass for analysis."
+            )
         super().__init__(
             qubits,
             analysis=analysis,
@@ -142,20 +157,11 @@ def _finalize(self):
             # Call sub-experiments finalize method
             subexp._finalize()
 
-    def _initialize_experiment_data(self):
-        """Initialize the return data container for the experiment run"""
-        experiment_data = ExperimentData(experiment=self)
-        # Initialize child experiment data
-        for sub_exp in self._experiments:
-            sub_data = sub_exp._initialize_experiment_data()
-            experiment_data.add_child_data(sub_data)
-        experiment_data.metadata["component_child_index"] = list(range(self.num_experiments))
-        return experiment_data
-
     def _additional_metadata(self):
         """Add component experiment metadata"""
         return {
-            "component_metadata": [sub_exp._metadata() for sub_exp in self.component_experiment()]
+            "component_types": [sub_exp.experiment_type for sub_exp in self.component_experiment()],
+            "component_metadata": [sub_exp._metadata() for sub_exp in self.component_experiment()],
         }
 
     def _add_job_metadata(self, metadata, jobs, **run_options):