Accept NumPy bitstrings in fermionic diagonalization

qiskit_addon_sqd/counts.py

@@ -1,6 +1,6 @@
 # This code is a Qiskit project.
 #
-# (C) Copyright IBM 2024.
+# (C) Copyright IBM 2024, 2026.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory
@@ -42,22 +42,29 @@ def counts_to_arrays(counts: Mapping[str, float | int]) -> tuple[np.ndarray, np.
     return bs_mat, freq_arr
 
 
-def bit_array_to_arrays(bit_array: BitArray) -> tuple[np.ndarray, np.ndarray]:
-    """Convert a bit array into a bitstring matrix and a probability array.
+def bit_array_to_arrays(bit_array: BitArray | np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Convert sampled bitstrings into a bitstring matrix and a probability array.
 
     Args:
-        bit_array: The bit array to convert
+        bit_array: The bit array or boolean bitstring matrix to convert.
 
     Returns:
         - A 2D array representing the sampled bitstrings. Each row represents a
           bitstring, and each element is a ``bool`` representation of the
           bit's value
         - A 1D array containing the probability with which each bitstring was sampled
     """
-    # TODO can use bit_array.to_bool_array() when it's available
-    bool_array = np.unpackbits(bit_array.array, axis=-1)[..., -bit_array.num_bits :].astype(bool)
+    if isinstance(bit_array, BitArray):
+        # TODO can use bit_array.to_bool_array() when it's available
+        bool_array = np.unpackbits(bit_array.array, axis=-1)[..., -bit_array.num_bits :].astype(
+            bool
+        )
+        num_samples = bit_array.num_shots
+    else:
+        bool_array = bit_array
+        num_samples = len(bool_array)
     bitstrings, counts = np.unique(bool_array, axis=0, return_counts=True)
-    probs = counts / bit_array.num_shots
+    probs = counts / num_samples
     return bitstrings, probs
 
 

qiskit_addon_sqd/fermion.py

@@ -1,6 +1,6 @@
 # This code is a Qiskit project.
 #
-# (C) Copyright IBM 2024.
+# (C) Copyright IBM 2024, 2026.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory
@@ -204,7 +204,7 @@ class _IterationState:
 def diagonalize_fermionic_hamiltonian(
     one_body_tensor: np.ndarray,
     two_body_tensor: np.ndarray,
-    bit_array: BitArray,
+    bit_array: BitArray | np.ndarray,
     samples_per_batch: int,
     norb: int,
     nelec: tuple[int, int],
@@ -231,10 +231,11 @@ def diagonalize_fermionic_hamiltonian(
     Args:
         one_body_tensor: The one-body tensor of the Hamiltonian.
         two_body_tensor: The two-body tensor of the Hamiltonian.
-        bit_array: Array of sampled bitstrings. Each bitstring should have both the
-            alpha part and beta part concatenated together, with the alpha part
-            concatenated on the right-hand side, like this:
-            ``[b_N, ..., b_0, a_N, ..., a_0]``.
+        bit_array: Array of sampled bitstrings, provided as either a Qiskit
+            :class:`~qiskit.primitives.BitArray` or a two-dimensional NumPy boolean
+            array. Each bitstring should have both the alpha part and beta part
+            concatenated together, with the alpha part concatenated on the right-hand
+            side, like this: ``[b_N, ..., b_0, a_N, ..., a_0]``.
         samples_per_batch: The number of bitstrings to include in each subsampled batch
             of bitstrings.
         norb: The number of spatial orbitals.
@@ -383,7 +384,7 @@ def diagonalize_fermionic_hamiltonian(
     carryover_strings_a = np.array([], dtype=np.int64)
     carryover_strings_b = np.array([], dtype=np.int64)
 
-    # Convert BitArray into bitstring and probability arrays
+    # Convert the samples into bitstring and probability arrays
     raw_bitstrings, raw_probs = bit_array_to_arrays(bit_array)
 
     # Bundle the loop-invariant configuration once, so the per-iteration helper

releasenotes/notes/accept-numpy-bitstrings-d944a5825e0ea36f.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    :func:`qiskit_addon_sqd.fermion.diagonalize_fermionic_hamiltonian` now accepts
+    sampled bitstrings as a two-dimensional NumPy boolean array in addition to a
+    Qiskit :class:`~qiskit.primitives.BitArray`.

test/test_counts.py

@@ -1,6 +1,6 @@
 # This code is a Qiskit project.
 #
-# (C) Copyright IBM 2024.
+# (C) Copyright IBM 2024, 2026.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory
@@ -14,8 +14,11 @@
 
 import unittest
 
+import numpy as np
 import pytest
+from qiskit.primitives import BitArray
 from qiskit_addon_sqd.counts import (
+    bit_array_to_arrays,
     counts_to_arrays,
     generate_counts_bipartite_hamming,
     generate_counts_uniform,
@@ -43,6 +46,19 @@ def test_counts_to_arrays(self):
             self.assertEqual((0,), bitstring_matrix.shape)
             self.assertEqual((0,), probs.shape)
 
+    def test_bit_array_to_arrays(self):
+        samples = np.array(
+            [[False, True, False], [True, False, True], [False, True, False]], dtype=bool
+        )
+        expected_bitstrings = np.array([[False, True, False], [True, False, True]])
+        expected_probs = np.array([2 / 3, 1 / 3])
+
+        for source in (samples, BitArray.from_bool_array(samples)):
+            with self.subTest(source=type(source).__name__):
+                bitstrings, probs = bit_array_to_arrays(source)
+                np.testing.assert_array_equal(bitstrings, expected_bitstrings)
+                np.testing.assert_allclose(probs, expected_probs)
+
     def test_generate_counts_uniform(self):
         with self.subTest("Basic test"):
             num_samples = 10

test/test_fermion.py

@@ -1,6 +1,6 @@
 # This code is a Qiskit project.
 #
-# (C) Copyright IBM 2024.
+# (C) Copyright IBM 2024, 2026.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory
@@ -99,6 +99,7 @@ def test_diagonalize_fermionic_hamiltonian_basic(self):
 
         # Merge bitstrings
         bit_array = BitArray.concatenate_shots([bit_array_ground_state, bit_array_random])
+        bit_array = np.unpackbits(bit_array.array, axis=-1)[..., -bit_array.num_bits :].astype(bool)
 
         # Diagonalize
         result = diagonalize_fermionic_hamiltonian(