ffsim sampler: return bit array directly (#324)

* ffsim sampler: return bit array directly * remove incorrect type annotations
qiskit-community · Oct 9, 2024 · 076d790 · 076d790
1 parent 54f91e6
commit 076d790
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/python/ffsim/qiskit/sampler.py b/python/ffsim/qiskit/sampler.py
@@ -76,23 +76,24 @@ def _run_pub(self, pub: SamplerPub) -> SamplerPubResult:
             for item in meas_info
         }
         for index, bound_circuit in np.ndenumerate(bound_circuits):
-            final_state = final_state_vector(bound_circuit)
             if qargs:
+                final_state = final_state_vector(bound_circuit)
                 norb, nelec = final_state.norb, final_state.nelec
                 if isinstance(nelec, int):
                     orbs = qargs
                 else:
                     orbs_a = [q for q in qargs if q < norb]
                     orbs_b = [q % norb for q in qargs if q >= norb]
                     orbs = (orbs_a, orbs_b)
-                samples = states.sample_state_vector(
-                    final_state, orbs=orbs, shots=pub.shots, seed=self._rng
+                samples_array = states.sample_state_vector(
+                    final_state,
+                    orbs=orbs,
+                    shots=pub.shots,
+                    bitstring_type=states.BitstringType.BIT_ARRAY,
+                    seed=self._rng,
                 )
             else:
-                samples = [""] * pub.shots
-            samples_array = np.array(
-                [np.fromiter(sample, dtype=np.uint8) for sample in samples]
-            )
+                samples_array = np.empty((pub.shots, 0), dtype=bool)
             for item in meas_info:
                 ary = _samples_to_packed_array(
                     samples_array, item.num_bits, item.qreg_indices

diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py
@@ -421,7 +421,7 @@ def sample_state_vector(
     concatenate: bool = True,
     bitstring_type: BitstringType = BitstringType.STRING,
     seed: np.random.Generator | int | None = None,
-) -> list[str] | tuple[list[str], list[str]]:
+):
     """Sample bitstrings from a state vector.
 
     Args:
@@ -488,7 +488,7 @@ def _sample_state_vector_spinless(
     shots: int,
     bitstring_type: BitstringType,
     seed: np.random.Generator | int | None,
-) -> list[str]:
+):
     if orbs is None:
         orbs = range(norb)
     rng = np.random.default_rng(seed)
@@ -510,7 +510,7 @@ def _sample_state_vector_spinful(
     concatenate: bool = True,
     bitstring_type: BitstringType,
     seed: np.random.Generator | int | None,
-) -> list[str] | tuple[list[str], list[str]]:
+):
     if orbs is None:
         orbs = range(norb), range(norb)
     rng = np.random.default_rng(seed)

diff --git a/tests/python/qiskit/sampler_test.py b/tests/python/qiskit/sampler_test.py
@@ -296,3 +296,17 @@ def test_reproducible_with_seed():
     counts_2 = pub_result.data.meas.get_counts()
 
     assert counts_1 == counts_2
+
+
+def test_edge_cases():
+    """Test edge cases."""
+    with pytest.raises(
+        ValueError, match="Circuit must contain at least one instruction."
+    ):
+        qubits = QuantumRegister(1, name="q")
+        circuit = QuantumCircuit(qubits)
+        circuit.measure_all()
+        sampler = ffsim.qiskit.FfsimSampler(default_shots=1)
+        pub = (circuit,)
+        job = sampler.run([pub])
+        _ = job.result()