Refactor reconstruct_expectation_values

circuit_knitting/cutting/cutting_reconstruction.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 from qiskit.quantum_info import PauliList
-from qiskit.result import QuasiDistribution
+from qiskit.primitives import SamplerResult
 
 from ..utils.observable_grouping import CommutingObservableGroup, ObservableCollection
 from ..utils.bitwise import bit_count
@@ -26,49 +26,57 @@
 
 
 def reconstruct_expectation_values(
-    quasi_dists: Sequence[Sequence[Sequence[tuple[QuasiDistribution, int]]]],
+    results: SamplerResult | dict[str | int, SamplerResult],
     coefficients: Sequence[tuple[float, WeightType]],
     observables: PauliList | dict[str | int, PauliList],
 ) -> list[float]:
     r"""
     Reconstruct an expectation value from the results of the sub-experiments.
 
     Args:
-        quasi_dists: A 3D sequence of length-2 tuples containing the quasi distributions and
-            QPD bit information from each sub-experiment. Its expected shape is
-            (num_unique_samples, num_partitions, num_commuting_observ_groups)
-        coefficients: A sequence of coefficients, such that each coefficient is associated
-            with one unique sample. The length of ``coefficients`` should equal
-            the length of ``quasi_dists``. Each coefficient is a tuple containing the numerical
-            value and the ``WeightType`` denoting how the value was generated.
+        results: The results from running the cutting subexperiments. If the cut circuit
+            was not partitioned between qubits and run separately, the input should be
+            a :class:`~qiskit.primitives.SamplerResult` instance or a dictionary mapping
+            a single partition to the results. If the circuit was partitioned and its
+            pieces run separately, the input should be a dictionary mapping partition labels
+            to the results from each partition's subexperiments.
+        coefficients: The weights associated with each unique subexperiment. Each weight is a tuple
+            containing the scalar value as well as the ``WeightType``, which denotes
+            how the value was generated.
         observables: The observable(s) for which the expectation values will be calculated.
-            This should be a :class:`~qiskit.quantum_info.PauliList` if the decomposed circuit
-            was not separated into subcircuits. If the decomposed circuit was separated, this
-            should be a dictionary mapping from partition label to subobservables.
+            This should be a :class:`~qiskit.quantum_info.PauliList` if ``results`` is a
+            :class:`~qiskit.primitives.SamplerResult` instance. Otherwise, it should be a
+            dictionary mapping partition labels to the observables associated with that partition.
 
     Returns:
-        A ``list`` of ``float``\ s, such that each float is a simulated expectation
+        A ``list`` of ``float``\ s, such that each float is an expectation
         value corresponding to the input observable in the same position
 
     Raises:
-        ValueError: The number of unique samples in quasi_dists does not equal the number of coefficients.
+        ValueError: ``observables`` and ``results`` are of incompatible types.
         ValueError: An input observable has a phase not equal to 1.
     """
-    if len(coefficients) != len(quasi_dists):
+    if isinstance(observables, PauliList) and not isinstance(results, SamplerResult):
         raise ValueError(
-            f"The number of unique samples in the quasi_dists list ({len(quasi_dists)}) does "
-            f"not equal the number of coefficients ({len(coefficients)})."
+            "If observables is a PauliList, results must be a SamplerResult instance."
         )
-    # Create the commuting observable groups
+    if isinstance(observables, dict) and not isinstance(results, dict):
+        raise ValueError(
+            "If observables is a dictionary, results must also be a dictionary."
+        )
+
+    # If circuit was not separated, transform input data structures to dictionary format
     if isinstance(observables, PauliList):
         if any(obs.phase != 0 for obs in observables):
             raise ValueError("An input observable has a phase not equal to 1.")
         subobservables_by_subsystem = decompose_observables(
             observables, "A" * len(observables[0])
         )
+        results_dict: dict[str | int, SamplerResult] = {"A": results}
         expvals = np.zeros(len(observables))
 
     else:
+        results_dict = results
         for label, subobservable in observables.items():
             if any(obs.phase != 0 for obs in subobservable):
                 raise ValueError("An input observable has a phase not equal to 1.")
@@ -79,21 +87,26 @@ def reconstruct_expectation_values(
         label: ObservableCollection(subobservables)
         for label, subobservables in subobservables_by_subsystem.items()
     }
+    sorted_subsystems = sorted(subsystem_observables.keys())  # type: ignore
 
-    # Assign each weight's sign and calculate the expectation values for each observable
-    for i, coeff in enumerate(coefficients):
-        sorted_subsystems = sorted(subsystem_observables.keys())  # type: ignore
+    # Reconstruct the expectation values
+    for i in range(len(coefficients)):
         current_expvals = np.ones((len(expvals),))
-        for j, label in enumerate(sorted_subsystems):
+        for label in sorted_subsystems:
             so = subsystem_observables[label]
+            coeff = coefficients[i]
             subsystem_expvals = [
                 np.zeros(len(cog.commuting_observables)) for cog in so.groups
             ]
             for k, cog in enumerate(so.groups):
-                quasi_probs = quasi_dists[i][j][k][0]
-                for outcome, quasi_prob in quasi_probs.items():
+                quasi_probs = results_dict[label].quasi_dists[i * len(so.groups) + k]  # type: ignore
+                for outcome, quasi_prob in quasi_probs.items():  # type: ignore
                     subsystem_expvals[k] += quasi_prob * _process_outcome(
-                        quasi_dists[i][j][k][1], cog, outcome
+                        results_dict[label].metadata[i * len(so.groups) + k][
+                            "num_qpd_bits"
+                        ],
+                        cog,
+                        outcome,
                     )
 
             for k, subobservable in enumerate(subobservables_by_subsystem[label]):

test/cutting/test_cutting_reconstruction.py

@@ -15,6 +15,7 @@
 import pytest
 import numpy as np
 from qiskit.result import QuasiDistribution
+from qiskit.primitives import SamplerResult
 from qiskit.quantum_info import Pauli, PauliList
 from qiskit.circuit import QuantumCircuit, ClassicalRegister
 
@@ -44,36 +45,64 @@ def setUp(self):
 
     def test_cutting_reconstruction(self):
         with self.subTest("Test PauliList observable"):
-            quasi_dists = [[[(QuasiDistribution({"0": 1.0}), 0)]]]
-            coefficients = [(1.0, WeightType.EXACT)]
-            observables = PauliList(["ZZ"])
-            expvals = reconstruct_expectation_values(
-                quasi_dists, coefficients, observables
+            results = SamplerResult(
+                quasi_dists=[QuasiDistribution({"0": 1.0})], metadata=[{}]
             )
+            results.metadata[0]["num_qpd_bits"] = 1
+            weights = [(1.0, WeightType.EXACT)]
+            subexperiments = [QuantumCircuit(2)]
+            creg1 = ClassicalRegister(1, name="qpd_measurements")
+            creg2 = ClassicalRegister(2, name="observable_measurements")
+            subexperiments[0].add_register(creg1)
+            subexperiments[0].add_register(creg2)
+            observables = PauliList(["ZZ"])
+            expvals = reconstruct_expectation_values(results, weights, observables)
             self.assertEqual([1.0], expvals)
         with self.subTest("Test mismatching inputs"):
-            quasi_dists = [[[(QuasiDistribution({"0": 1.0}), 0)]]]
-            coefficients = [(0.5, WeightType.EXACT), (0.5, WeightType.EXACT)]
+            results = SamplerResult(
+                quasi_dists=[QuasiDistribution({"0": 1.0})], metadata=[{}]
+            )
+            results.metadata[0]["num_qpd_bits"] = 1
+            weights = [(0.5, WeightType.EXACT), (0.5, WeightType.EXACT)]
+            subexperiments = {"A": QuantumCircuit(2)}
+            observables = {"A": PauliList(["Z"]), "B": PauliList(["Z"])}
+            with pytest.raises(ValueError) as e_info:
+                reconstruct_expectation_values(results, weights, observables)
+            assert (
+                e_info.value.args[0]
+                == "If observables is a dictionary, results must also be a dictionary."
+            )
+            results2 = {"A": results}
             observables = PauliList(["ZZ"])
             with pytest.raises(ValueError) as e_info:
-                reconstruct_expectation_values(quasi_dists, coefficients, observables)
+                reconstruct_expectation_values(results2, weights, observables)
             assert (
                 e_info.value.args[0]
-                == "The number of unique samples in the quasi_dists list (1) does not equal the number of coefficients (2)."
+                == "If observables is a PauliList, results must be a SamplerResult instance."
             )
         with self.subTest("Test unsupported phase"):
-            quasi_dists = [[[(QuasiDistribution({"0": 1.0}), 0)]]]
-            coefficients = [(0.5, WeightType.EXACT)]
+            results = SamplerResult(
+                quasi_dists=[QuasiDistribution({"0": 1.0})], metadata=[{}]
+            )
+            results.metadata[0]["num_qpd_bits"] = 1
+            weights = [(0.5, WeightType.EXACT)]
+            subexperiments = [QuantumCircuit(2)]
+            creg1 = ClassicalRegister(1, name="qpd_measurements")
+            creg2 = ClassicalRegister(2, name="observable_measurements")
+            subexperiments[0].add_register(creg1)
+            subexperiments[0].add_register(creg2)
             observables = PauliList(["iZZ"])
             with pytest.raises(ValueError) as e_info:
-                reconstruct_expectation_values(quasi_dists, coefficients, observables)
+                reconstruct_expectation_values(results, weights, observables)
             assert (
                 e_info.value.args[0]
                 == "An input observable has a phase not equal to 1."
             )
-            observables = {"A": PauliList(["iZZ"])}
+            results = {"A": results}
+            subexperiments = {"A": subexperiments}
+            observables = {"A": observables}
             with pytest.raises(ValueError) as e_info:
-                reconstruct_expectation_values(quasi_dists, coefficients, observables)
+                reconstruct_expectation_values(results, weights, observables)
             assert (
                 e_info.value.args[0]
                 == "An input observable has a phase not equal to 1."