Fix reconstruction to work with fragmented partition labels

app/gate_cutting_reconstruct_distribution.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from collections import defaultdict
 from typing import Hashable, Sequence
 
@@ -16,7 +17,11 @@
 from qiskit.primitives import SamplerResult
 from qiskit.quantum_info import PauliList
 
-from app.utils import counts_to_array
+from app.utils import (
+    product_dicts,
+    shift_bits_by_index,
+    find_character_in_string,
+)
 
 
 def _process_outcome_distribution(
@@ -51,7 +56,8 @@ def reconstruct_distribution(
     results: SamplerResult | dict[Hashable, SamplerResult],
     coefficients: Sequence[tuple[float, WeightType]],
     observables: PauliList | dict[Hashable, PauliList],
-) -> list[float]:
+    partition_labels: str = None,
+) -> dict[int, float]:
     r"""
     Reconstruct an expectation value from the results of the sub-experiments.
 
@@ -119,12 +125,14 @@ def reconstruct_distribution(
         for label, subobservables in subobservables_by_subsystem.items()
     }
 
-    qubits = {
-        label: len(subobservables[0])
-        for label, subobservables in subobservables_by_subsystem.items()
-    }
+    result_dict = defaultdict(float)
 
-    result = np.zeros(2 ** sum(qubits.values()))
+    labels = {*partition_labels}
+
+    label_index_lists = {
+        label: list(find_character_in_string(partition_labels, label))
+        for label in labels
+    }
 
     # Reconstruct the expectation values
     for i, coeff in enumerate(coefficients):
@@ -141,10 +149,12 @@ def reconstruct_distribution(
                     )
                     coeff_result_dict[label][meas_outcomes] += qpd_factor * quasi_prob
 
-        temp = np.ones(1)
-        for label, r in coeff_result_dict.items():
-            temp = np.kron(counts_to_array(r, qubits[label]), temp)
-
-        result += coeff[0] * temp
+        for meas_keys, quasi_prob_vals in product_dicts(
+            *list(coeff_result_dict.values())
+        ).items():
+            combined_meas = 0
+            for meas, label in zip(meas_keys, coeff_result_dict.keys()):
+                combined_meas += shift_bits_by_index(meas, label_index_lists[label])
+            result_dict[combined_meas] += coeff[0] * math.prod(quasi_prob_vals)
 
-    return result
+    return result_dict

app/utils.py

@@ -40,3 +40,42 @@ def sampler_result_to_array_dict(sampler_result: SamplerResult):
         )
 
     return result_array_dict
+
+
+def find_character_in_string(string, ch):
+    for i, ltr in enumerate(string):
+        if ltr == ch:
+            yield i
+
+
+def shift_bits_by_index(num, idx_list):
+    n_bits = len(idx_list)
+    new_num = 0
+    sorted_idx_list = sorted(idx_list, reverse=True)
+    prev_index = sorted_idx_list[0]
+    for i, idx in enumerate(sorted_idx_list):
+        new_num = new_num << (prev_index - idx)
+        new_num += (num >> (n_bits - i - 1)) & 1
+        prev_index = idx
+
+    return new_num << prev_index
+
+
+def product_dicts(*dicts):
+    # Base case: If there is only one dictionary, convert its values to lists
+    if len(dicts) == 1:
+        return {(key,): [value] for key, value in dicts[0].items()}
+
+    # Recursive case: Compute the product of n-1 dictionaries
+    dicts_except_last = dicts[:-1]
+    product_except_last = product_dicts(*dicts_except_last)
+
+    # Get the last dictionary
+    last_dict = dicts[-1]
+
+    # Compute the product of the last dictionary with the result from the previous step using dictionary comprehensions
+    return {
+        (*prod_key, key): prod_val + [val]
+        for prod_key, prod_val in product_except_last.items()
+        for key, val in last_dict.items()
+    }

test/test_reconstruction_gate_cutting.py

@@ -65,11 +65,10 @@ def _generate_reconstruction_test(num_qubits=4, partition_label=None, reps=2):
         exact_experiment["metadata"]["simulator_metadata"]["num_qubits"],
     )
 
-    reconstructed_dist = reconstruct_distribution(
-        results,
-        coefficients,
-        subobservables,
+    reconstructed_counts = reconstruct_distribution(
+        results, coefficients, subobservables, partition_label
     )
+    reconstructed_dist = counts_to_array(reconstructed_counts, num_qubits)
     return exact_distribution, reconstructed_dist
 
 
@@ -80,7 +79,8 @@ def test_generate_test_data(self):
             num_qubits, "AABB"
         )
         self.assertTrue(
-            np.allclose(exact_distribution, reconstructed_dist, atol=2 ** (-num_qubits))
+            np.allclose(exact_distribution, reconstructed_dist, atol=0.025),
+            msg=f"\nExact distribution: {exact_distribution}\nReconstruced distribution: {reconstructed_dist}\nDiff: {np.abs(exact_distribution- reconstructed_dist)}",
         )
 
     def test_generate_test_data_2(self):
@@ -89,7 +89,8 @@ def test_generate_test_data_2(self):
             num_qubits, "ABBA", reps=1
         )
         self.assertTrue(
-            np.allclose(exact_distribution, reconstructed_dist, atol=2 ** (-num_qubits))
+            np.allclose(exact_distribution, reconstructed_dist, atol=0.025),
+            msg=f"\nExact distribution: {exact_distribution}\nReconstruced distribution: {reconstructed_dist}\nDiff: {np.abs(exact_distribution- reconstructed_dist)}",
         )
 
     def test_generate_test_data_3(self):
@@ -98,5 +99,16 @@ def test_generate_test_data_3(self):
             num_qubits, "ABBCD", reps=1
         )
         self.assertTrue(
-            np.allclose(exact_distribution, reconstructed_dist, atol=2 ** (-num_qubits))
+            np.allclose(exact_distribution, reconstructed_dist, atol=0.025),
+            msg=f"\nExact distribution: {exact_distribution}\nReconstruced distribution: {reconstructed_dist}\nDiff: {np.abs(exact_distribution- reconstructed_dist)}",
+        )
+
+    def test_generate_test_data_4(self):
+        num_qubits = 8
+        exact_distribution, reconstructed_dist = _generate_reconstruction_test(
+            num_qubits, "AAAABBBB"
+        )
+        self.assertTrue(
+            np.allclose(exact_distribution, reconstructed_dist, atol=0.025),
+            msg=f"\nExact distribution: {exact_distribution}\nReconstruced distribution: {reconstructed_dist}\nDiff: {np.abs(exact_distribution- reconstructed_dist)}",
         )