Return more info about circuit separation during partitioning

circuit_knitting/cutting/cutting_decomposition.py

@@ -28,17 +28,24 @@
 from ..utils.observable_grouping import observables_restricted_to_subsystem
 from ..utils.transforms import separate_circuit
 from .qpd.qpd_basis import QPDBasis
-from .qpd.instructions import TwoQubitQPDGate
+from .qpd.instructions import BaseQPDGate, SingleQubitQPDGate, TwoQubitQPDGate
 
 
 class PartitionedCuttingProblem(NamedTuple):
     """The result of decomposing and separating a circuit and observable(s)."""
 
     subcircuits: dict[str | int, QuantumCircuit]
-    bases: list[QPDBasis]
+    cuts: list[CutInfo]
     subobservables: dict[str | int, QuantumCircuit] | None = None
 
 
+class CutInfo(NamedTuple):
+    """The decomposition and location information associated with one cut."""
+
+    basis: QPDBasis
+    gates: list[tuple[Hashable, int]] | list[int] | int
+
+
 def partition_circuit_qubits(
     circuit: QuantumCircuit, partition_labels: Sequence[Hashable], inplace: bool = False
 ) -> QuantumCircuit:
@@ -195,6 +202,7 @@ def partition_problem(
         ValueError: An input observable acts on a different number of qubits than the input circuit.
         ValueError: An input observable has a phase not equal to 1.
         ValueError: The input circuit should contain no classical bits or registers.
+        ValueError: The input circuit should contain no SingleQubitQPDGate instances.
     """
     if len(partition_labels) != circuit.num_qubits:
         raise ValueError(
@@ -221,14 +229,28 @@ def partition_problem(
     bases = []
     i = 0
     for inst in qpd_circuit.data:
+        if isinstance(inst.operation, SingleQubitQPDGate):
+            raise ValueError(
+                "Input circuit may not contain SingleQubitQPDGate instances."
+            )
         if isinstance(inst.operation, TwoQubitQPDGate):
             bases.append(inst.operation.basis)
             inst.operation.label = inst.operation.label + f"_{i}"
             i += 1
 
     # Separate the decomposed circuit into its subcircuits
     qpd_circuit_dx = qpd_circuit.decompose(TwoQubitQPDGate)
-    separated_circs = separate_circuit(qpd_circuit_dx, partition_labels)
+    subcircuits = separate_circuit(qpd_circuit_dx, partition_labels).subcircuits
+
+    # Gather the basis and location information for the cuts
+    cuts_dict = defaultdict(list)
+    for label in subcircuits.keys():
+        circuit = subcircuits[label]
+        for i, inst in enumerate(circuit.data):
+            if isinstance(inst.operation, BaseQPDGate):
+                cut_num = int(inst.operation.label.split("_")[-1])
+                cuts_dict[cut_num].append((label, i))
+    cuts = [CutInfo(basis, cuts_dict[cut_num]) for cut_num, basis in enumerate(bases)]
 
     # Decompose the observables, if provided
     subobservables_by_subsystem = None
@@ -238,8 +260,8 @@ def partition_problem(
         )
 
     return PartitionedCuttingProblem(
-        separated_circs.subcircuits,  # type: ignore
-        bases,
+        subcircuits,  # type: ignore
+        cuts,
         subobservables=subobservables_by_subsystem,
     )