From f4dec604669fbe3d1103256221cb6d10ea1b32c7 Mon Sep 17 00:00:00 2001
From: Jim Garrison <garrison@ibm.com>
Date: Mon, 23 Oct 2023 16:26:28 -0400
Subject: [PATCH] Fix reconstruction when subsystem contains no observable

---
 .../cutting/cutting_reconstruction.py         |  3 +-
 test/cutting/test_cutting_roundtrip.py        | 46 +++++++++++++++----
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/circuit_knitting/cutting/cutting_reconstruction.py b/circuit_knitting/cutting/cutting_reconstruction.py
index 90c87257f..40094a1aa 100644
--- a/circuit_knitting/cutting/cutting_reconstruction.py
+++ b/circuit_knitting/cutting/cutting_reconstruction.py
@@ -22,6 +22,7 @@
 from ..utils.observable_grouping import CommutingObservableGroup, ObservableCollection
 from ..utils.bitwise import bit_count
 from .cutting_decomposition import decompose_observables
+from .cutting_experiments import _get_pauli_indices
 from .qpd import WeightType
 
 
@@ -136,7 +137,7 @@ def _process_outcome(
         this vector correspond to the elements of ``cog.commuting_observables``,
         and each result will be either +1 or -1.
     """
-    num_meas_bits = len(cog.pauli_indices)
+    num_meas_bits = len(_get_pauli_indices(cog))
 
     outcome = _outcome_to_int(outcome)
     meas_outcomes = outcome & ((1 << num_meas_bits) - 1)
diff --git a/test/cutting/test_cutting_roundtrip.py b/test/cutting/test_cutting_roundtrip.py
index 9faa796be..ea4944054 100644
--- a/test/cutting/test_cutting_roundtrip.py
+++ b/test/cutting/test_cutting_roundtrip.py
@@ -179,28 +179,56 @@ def test_cutting_exact_reconstruction(example_circuit):
     assert np.allclose(exact_expvals, simulated_expvals, atol=1e-8)
 
 
-def test_sampler_with_identity_subobservable(example_circuit):
-    """This test ensures that the sampler does not throw an error if you pass it a subcircuit with no observable measurements.
+@pytest.mark.parametrize(
+    "sampler,is_exact_sampler", [(Sampler(), False), (ExactSampler(), True)]
+)
+def test_sampler_with_identity_subobservable(sampler, is_exact_sampler):
+    """This test ensures that the sampler works for a subcircuit with no observable measurements.
 
-    Tests temporary workaround to Issue #422.
+    Specifically, that
 
-    This test passes if no exceptions are raised.
+    - ``Sampler`` does not blow up (Issue #422); and
+    - ``ExactSampler`` returns correct results
 
+    This is related to https://github.com/Qiskit-Extensions/circuit-knitting-toolbox/issues/422.
     """
+    # Create a circuit to cut
+    qc = QuantumCircuit(3)
+    append_random_unitary(qc, [0, 1])
+    append_random_unitary(qc, [2])
+    qc.rxx(np.pi / 3, 1, 2)
+    append_random_unitary(qc, [0, 1])
+    append_random_unitary(qc, [2])
 
-    qc = example_circuit
-    observable_to_test = PauliList(
+    # Determine expectation value using cutting
+    observables = PauliList(
         ["IIZ"]
     )  # Without the workaround to Issue #422, this observable causes a Sampler error.
     subcircuits, bases, subobservables = partition_problem(
-        qc, "AAB", observables=observable_to_test
+        qc, "AAB", observables=observables
     )
     subexperiments, coefficients = generate_cutting_experiments(
         subcircuits, subobservables, num_samples=np.inf
     )
-    samplers = {label: Sampler() for label in subexperiments.keys()}
+    samplers = {label: sampler for label in subexperiments.keys()}
     results = {
         label: sampler.run(subexperiments[label]).result()
         for label, sampler in samplers.items()
     }
-    _ = results
+    reconstructed_expvals = reconstruct_expectation_values(
+        results, coefficients, subobservables
+    )
+
+    if is_exact_sampler:
+        # Determine exact expectation values
+        estimator = Estimator()
+        exact_expvals = (
+            estimator.run([qc] * len(observables), list(observables)).result().values
+        )
+
+        logger.info(
+            "Max error: %f", np.max(np.abs(exact_expvals - reconstructed_expvals))
+        )
+
+        # Ensure both methods yielded equivalent expectation values
+        assert np.allclose(exact_expvals, reconstructed_expvals, atol=1e-8)