Handle bits selection in results for circuits with non-default regist…

…ers.

pytket/extensions/qiskit/backends/ibm.py

@@ -14,7 +14,7 @@
 
 import itertools
 from ast import literal_eval
-from collections import Counter
+from collections import Counter, OrderedDict
 import json
 from time import sleep
 from typing import (
@@ -28,7 +28,7 @@
 
 import numpy as np
 
-from qiskit.primitives import PrimitiveResult, SamplerResult  # type: ignore
+from qiskit.primitives import PrimitiveResult, SamplerPubResult, DataBin, BitArray  # type: ignore
 
 
 # RuntimeJob has no queue_position attribute, which is referenced
@@ -97,14 +97,10 @@
 _DEBUG_HANDLE_PREFIX = "_MACHINE_DEBUG_"
 
 
-def _gen_debug_results(n_qubits: int, shots: int, index: int) -> SamplerResult:
-    debug_dist = {n: 0.0 for n in range(pow(2, n_qubits))}
-    debug_dist[0] = 1.0
-    qd = QuasiDistribution(debug_dist)
-    return SamplerResult(
-        quasi_dists=[qd] * (index + 1),
-        metadata=[{"header_metadata": {}, "shots": shots}] * (index + 1),
-    )
+def _gen_debug_results(n_qubits: int, shots: int) -> PrimitiveResult:
+    n_u8s = (n_qubits - 1) // 8 + 1
+    arr = np.array([[0] * n_u8s for _ in range(shots)], dtype=np.uint8)
+    return PrimitiveResult([SamplerPubResult(DataBin(c=BitArray(arr, n_qubits)))])
 
 
 class NoIBMQCredentialsError(Exception):
@@ -203,7 +199,9 @@ def __init__(
         self._monitor = monitor
 
         # cache of results keyed by job id and circuit index
-        self._ibm_res_cache: dict[tuple[str, int], Counter] = dict()
+        self._ibm_res_cache: dict[
+            tuple[str, int], tuple[Counter, Optional[list[Bit]]]
+        ] = dict()
 
         if sampler_options is None:
             sampler_options = SamplerOptions()
@@ -519,14 +517,11 @@ def process_circuits(
 
                 qcs, ppcirc_strs = [], []
                 for tkc in batch_chunk:
-                    tkc1 = tkc.copy()
-                    # Flatten bits to default register in lexicographic order:
-                    tkc1.rename_units({bit: Bit(i) for i, bit in enumerate(tkc1.bits)})
                     if postprocess:
-                        c0, ppcirc = prepare_circuit(tkc1, allow_classical=False)
+                        c0, ppcirc = prepare_circuit(tkc, allow_classical=False)
                         ppcirc_rep = ppcirc.to_dict()
                     else:
-                        c0, ppcirc_rep = tkc1, None
+                        c0, ppcirc_rep = tkc, None
                     if simplify_initial:
                         SimplifyInitial(
                             allow_classical=False, create_all_qubits=True
@@ -592,7 +587,7 @@ def get_result(self, handle: ResultHandle, **kwargs: KwargTypes) -> BackendResul
             if self._MACHINE_DEBUG or jobid.startswith(_DEBUG_HANDLE_PREFIX):
                 shots: int
                 shots, _ = literal_eval(jobid[len(_DEBUG_HANDLE_PREFIX) :])
-                res = _gen_debug_results(n_meas, shots, index)
+                res = _gen_debug_results(n_meas, shots)
             else:
                 try:
                     job = self._retrieve_job(jobid)
@@ -611,21 +606,27 @@ def get_result(self, handle: ResultHandle, **kwargs: KwargTypes) -> BackendResul
                         sleep(10)
 
                 res = job.result(timeout=kwargs.get("timeout", None))
-            if isinstance(res, SamplerResult):
-                # TODO Is this code still reachable?
-                for circ_index, (r, d) in enumerate(zip(res.quasi_dists, res.metadata)):
-                    self._ibm_res_cache[(jobid, circ_index)] = Counter(
-                        {n: int(0.5 + d["shots"] * p) for n, p in r.items()}
-                    )
-            else:
-                assert isinstance(res, PrimitiveResult)
-                for circ_index, pub_result in enumerate(res._pub_results):
-                    readouts = pub_result.data.c.array
-                    self._ibm_res_cache[(jobid, circ_index)] = Counter(
-                        _int_from_readout(readout) for readout in readouts
-                    )
-
-        counts = self._ibm_res_cache[cache_key]  # Counter[int]
+            assert isinstance(res, PrimitiveResult)
+            for circ_index, pub_result in enumerate(res._pub_results):
+                data = pub_result.data
+                c_regs = OrderedDict(
+                    (reg_name, data.__getattribute__(reg_name).num_bits)
+                    for reg_name in sorted(data.keys())
+                )
+                readouts = BitArray.concatenate_bits(
+                    [data.__getattribute__(reg_name) for reg_name in c_regs]
+                ).array
+                self._ibm_res_cache[(jobid, circ_index)] = (
+                    Counter(_int_from_readout(readout) for readout in readouts),
+                    list(
+                        itertools.chain.from_iterable(
+                            [Bit(reg_name, i) for i in range(reg_size)]
+                            for reg_name, reg_size in c_regs.items()
+                        )
+                    ),
+                )
+
+        counts, c_bits = self._ibm_res_cache[cache_key]  # Counter[int], list[Bit]
         # Convert to `OutcomeArray`:
         tket_counts: Counter = Counter()
         for outcome_key, sample_count in counts.items():
@@ -636,7 +637,7 @@ def get_result(self, handle: ResultHandle, **kwargs: KwargTypes) -> BackendResul
             )
             tket_counts[array] = sample_count
         # Convert to `BackendResult`:
-        result = BackendResult(counts=tket_counts, ppcirc=ppcirc)
+        result = BackendResult(c_bits=c_bits, counts=tket_counts, ppcirc=ppcirc)
 
         self._cache[handle] = {"result": result}
         return result