Add vectorized from_bits and to_bits

qualtran/_infra/data_types.py

@@ -50,7 +50,7 @@
 
 import abc
 from enum import Enum
-from typing import Any, cast, Iterable, List, Sequence, Union
+from typing import Any, Iterable, List, Sequence, Union
 
 import attrs
 import numpy as np
@@ -77,10 +77,30 @@ def get_classical_domain(self) -> Iterable[Any]:
     def to_bits(self, x) -> List[int]:
         """Yields individual bits corresponding to binary representation of x"""
 
+    def to_bits_array(self, x_array: NDArray[Any]) -> NDArray[np.uint8]:
+        """Yields an NDArray of bits corresponding to binary representations of the input elements.
+
+        Often, converting an array can be performed faster than converting each element individually.
+        This operation accepts any NDArray of values, and the output array satisfies
+        `output_shape = input_shape + (self.bitsize,)`.
+        """
+        return np.vectorize(
+            lambda x: np.asarray(self.to_bits(x), dtype=np.uint8), signature='()->(n)'
+        )(x_array)
+
     @abc.abstractmethod
     def from_bits(self, bits: Sequence[int]):
         """Combine individual bits to form x"""
 
+    def from_bits_array(self, bits_array: NDArray[np.uint8]):
+        """Combine individual bits to form classical values.
+
+        Often, converting an array can be performed faster than converting each element individually.
+        This operation accepts any NDArray of bits such that the last dimension equals `self.bitsize`,
+        and the output array satisfies `output_shape = input_shape[:-1]`.
+        """
+        return np.vectorize(self.from_bits, signature='(n)->()')(bits_array)
+
     @abc.abstractmethod
     def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
         """Raises an exception if `val` is not a valid classical value for this type.
@@ -90,17 +110,6 @@ def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
             debug_str: Optional debugging information to use in exception messages.
         """
 
-    @abc.abstractmethod
-    def is_symbolic(self) -> bool:
-        """Returns True if this qdtype is parameterized with symbolic objects."""
-
-    def iteration_length_or_zero(self) -> SymbolicInt:
-        """Safe version of iteration length.
-
-        Returns the iteration_length if the type has it or else zero.
-        """
-        return getattr(self, 'iteration_length', 0)
-
     def assert_valid_classical_val_array(self, val_array: NDArray[Any], debug_str: str = 'val'):
         """Raises an exception if `val_array` is not a valid array of classical values
         for this type.
@@ -116,6 +125,17 @@ def assert_valid_classical_val_array(self, val_array: NDArray[Any], debug_str: s
         for val in val_array.reshape(-1):
             self.assert_valid_classical_val(val)
 
+    @abc.abstractmethod
+    def is_symbolic(self) -> bool:
+        """Returns True if this qdtype is parameterized with symbolic objects."""
+
+    def iteration_length_or_zero(self) -> SymbolicInt:
+        """Safe version of iteration length.
+
+        Returns the iteration_length if the type has it or else zero.
+        """
+        return getattr(self, 'iteration_length', 0)
+
     def __str__(self):
         return f'{self.__class__.__name__}({self.num_qubits})'
 
@@ -214,9 +234,11 @@ def get_classical_domain(self) -> Iterable[int]:
 
     def to_bits(self, x: int) -> List[int]:
         """Yields individual bits corresponding to binary representation of x"""
+        if is_symbolic(self.bitsize):
+            raise ValueError(f"cannot compute bits with symbolic {self.bitsize=}")
+
         self.assert_valid_classical_val(x)
-        mask = (1 << cast(int, self.bitsize)) - 1
-        return QUInt(self.bitsize).to_bits(int(x) & mask)
+        return [int(b) for b in np.binary_repr(x, width=self.bitsize)]
 
     def from_bits(self, bits: Sequence[int]) -> int:
         """Combine individual bits to form x"""
@@ -324,10 +346,52 @@ def to_bits(self, x: int) -> List[int]:
         self.assert_valid_classical_val(x)
         return [int(x) for x in f'{int(x):0{self.bitsize}b}']
 
+    def to_bits_array(self, x_array: NDArray[np.integer]) -> NDArray[np.uint8]:
+        """Returns the big-endian bitstrings specified by the given integers.
+
+        Args:
+            x_array: An integer or array of unsigned integers.
+        """
+        if is_symbolic(self.bitsize):
+            raise ValueError(f"Cannot compute bits for symbolic {self.bitsize=}")
+
+        if self.bitsize > 64:
+            # use the default vectorized `to_bits`
+            return super().to_bits_array(x_array)
+
+        w = int(self.bitsize)
+        x = np.atleast_1d(x_array)
+        if not np.issubdtype(x.dtype, np.uint):
+            assert np.all(x >= 0)
+            assert np.iinfo(x.dtype).bits <= 64
+            x = x.astype(np.uint64)
+        assert w <= np.iinfo(x.dtype).bits
+        mask = 2 ** np.arange(w - 1, 0 - 1, -1, dtype=x.dtype).reshape((w, 1))
+        return (x & mask).astype(bool).astype(np.uint8).T
+
     def from_bits(self, bits: Sequence[int]) -> int:
         """Combine individual bits to form x"""
         return int("".join(str(x) for x in bits), 2)
 
+    def from_bits_array(self, bits_array: NDArray[np.uint8]) -> NDArray[np.integer]:
+        """Returns the integer specified by the given big-endian bitstrings.
+
+        Args:
+            bits_array: A bitstring or array of bitstrings, each of which has the 1s bit (LSB) at the end.
+        Returns:
+            An array of integers; one for each bitstring.
+        """
+        bitstrings = np.atleast_2d(bits_array)
+        if bitstrings.shape[1] != self.bitsize:
+            raise ValueError(f"Input bitsize {bitstrings.shape[1]} does not match {self.bitsize=}")
+
+        if self.bitsize > 64:
+            # use the default vectorized `from_bits`
+            return super().from_bits_array(bits_array)
+
+        basis = 2 ** np.arange(self.bitsize - 1, 0 - 1, -1, dtype=np.uint64)
+        return np.sum(basis * bitstrings, axis=1, dtype=np.uint64)
+
     def assert_valid_classical_val(self, val: int, debug_str: str = 'val'):
         if not isinstance(val, (int, np.integer)):
             raise ValueError(f"{debug_str} should be an integer, not {val!r}")
@@ -528,6 +592,13 @@ def from_bits(self, bits: Sequence[int]) -> Fxp:
         fxp_bin = "0b" + bits_bin[: -self.num_frac] + "." + bits_bin[-self.num_frac :]
         return Fxp(fxp_bin, dtype=self.fxp_dtype_str)
 
+    def from_bits_array(self, bits_array: NDArray[np.uint8]):
+        assert isinstance(self.bitsize, int), "cannot convert to bits for symbolic bitsize"
+        # TODO figure out why `np.vectorize` is not working here
+        return Fxp(
+            [self.from_bits(bitstring) for bitstring in bits_array.reshape(-1, self.bitsize)]
+        )
+
     def to_fixed_width_int(self, x: Union[float, Fxp]) -> int:
         """Returns the interpretation of the binary representation of `x` as an integer. Requires `x` to be nonnegative."""
         if x < 0:

qualtran/_infra/data_types_test.py

@@ -11,13 +11,15 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
 import math
 import random
+from typing import Any, Sequence, Union
 
+import cirq
 import numpy as np
 import pytest
 import sympy
+from numpy.typing import NDArray
 
 from qualtran.symbolics import is_symbolic
 
@@ -233,8 +235,20 @@ def test_single_qubit_consistency():
     assert check_dtypes_consistent(QFxp(1, 1), QBit())
 
 
-def test_to_and_from_bits():
-    # QInt
+def assert_to_and_from_bits_array_consistent(qdtype: QDType, values: Union[Sequence[Any], NDArray]):
+    values = np.asarray(values)
+    bits_array = qdtype.to_bits_array(values)
+
+    # individual values
+    for val, bits in zip(values.reshape(-1), bits_array.reshape(-1, qdtype.num_qubits)):
+        assert np.all(bits == qdtype.to_bits(val))
+
+    # round trip
+    values_roundtrip = qdtype.from_bits_array(bits_array)
+    assert np.all(values_roundtrip == values)
+
+
+def test_qint_to_and_from_bits():
     qint4 = QInt(4)
     assert [*qint4.get_classical_domain()] == [*range(-8, 8)]
     for x in range(-8, 8):
@@ -246,7 +260,10 @@ def test_to_and_from_bits():
     with pytest.raises(ValueError):
         QInt(4).to_bits(10)
 
-    # QUInt
+    assert_to_and_from_bits_array_consistent(qint4, range(-8, 8))
+
+
+def test_quint_to_and_from_bits():
     quint4 = QUInt(4)
     assert [*quint4.get_classical_domain()] == [*range(0, 16)]
     assert list(quint4.to_bits(10)) == [1, 0, 1, 0]
@@ -259,23 +276,66 @@ def test_to_and_from_bits():
     with pytest.raises(ValueError):
         quint4.to_bits(-1)
 
-    # BoundedQUInt
+    assert_to_and_from_bits_array_consistent(quint4, range(0, 16))
+
+
+def test_bits_to_int():
+    rs = np.random.RandomState(52)
+    bitstrings = rs.choice([0, 1], size=(100, 23))
+
+    nums = QUInt(23).from_bits_array(bitstrings)
+    assert nums.shape == (100,)
+
+    for num, bs in zip(nums, bitstrings):
+        ref_num = cirq.big_endian_bits_to_int(bs.tolist())
+        assert num == ref_num
+
+    # check one input bitstring instead of array of input bitstrings.
+    (num,) = QUInt(2).from_bits_array(np.array([1, 0]))
+    assert num == 2
+
+
+def test_int_to_bits():
+    rs = np.random.RandomState(52)
+    nums = rs.randint(0, 2**23 - 1, size=(100,), dtype=np.uint64)
+    bitstrings = QUInt(23).to_bits_array(nums)
+    assert bitstrings.shape == (100, 23)
+
+    for num, bs in zip(nums, bitstrings):
+        ref_bs = cirq.big_endian_int_to_bits(int(num), bit_count=23)
+        np.testing.assert_array_equal(ref_bs, bs)
+
+    # check bounds
+    with pytest.raises(AssertionError):
+        QUInt(8).to_bits_array(np.array([4, -2]))
+
+
+def test_bounded_quint_to_and_from_bits():
     bquint4 = BoundedQUInt(4, 12)
     assert [*bquint4.get_classical_domain()] == [*range(0, 12)]
     assert list(bquint4.to_bits(10)) == [1, 0, 1, 0]
     with pytest.raises(ValueError):
         BoundedQUInt(4, 12).to_bits(13)
 
-    # QBit
+    assert_to_and_from_bits_array_consistent(bquint4, range(0, 12))
+
+
+def test_qbit_to_and_from_bits():
     assert list(QBit().to_bits(0)) == [0]
     assert list(QBit().to_bits(1)) == [1]
     with pytest.raises(ValueError):
         QBit().to_bits(2)
 
-    # QAny
+    assert_to_and_from_bits_array_consistent(QBit(), [0, 1])
+
+
+def test_qany_to_and_from_bits():
     assert list(QAny(4).to_bits(10)) == [1, 0, 1, 0]
 
-    # QIntOnesComp
+    assert_to_and_from_bits_array_consistent(QAny(4), range(16))
+
+
+def test_qintonescomp_to_and_from_bits():
     qintones4 = QIntOnesComp(4)
     assert list(qintones4.to_bits(-2)) == [1, 1, 0, 1]
     assert list(qintones4.to_bits(2)) == [0, 0, 1, 0]
@@ -287,6 +347,10 @@ def test_to_and_from_bits():
     with pytest.raises(ValueError):
         qintones4.to_bits(-8)
 
+    assert_to_and_from_bits_array_consistent(qintones4, range(-7, 8))
+
+
+def test_qfxp_to_and_from_bits():
     # QFxp: Negative numbers are stored as ones complement
     qfxp_4_3 = QFxp(4, 3, True)
     assert list(qfxp_4_3.to_bits(0.5)) == [0, 1, 0, 0]
@@ -321,6 +385,11 @@ def test_to_and_from_bits():
     assert list(QFxp(7, 3, True).to_bits(-4.375)) == [1] + [0, 1, 1] + [1, 0, 1]
     assert list(QFxp(7, 3, True).to_bits(+4.625)) == [0] + [1, 0, 0] + [1, 0, 1]
 
+    assert_to_and_from_bits_array_consistent(QFxp(4, 3, False), [1 / 2, 1 / 4, 3 / 8])
+    assert_to_and_from_bits_array_consistent(
+        QFxp(4, 3, True), [1 / 2, -1 / 2, 1 / 4, -1 / 4, -3 / 8, 3 / 8]
+    )
+
 
 def test_iter_bits():
     assert QUInt(2).to_bits(0) == [0, 0]

qualtran/simulation/classical_sim.py

@@ -37,19 +37,18 @@
 ClassicalValT = Union[int, np.integer, NDArray[np.integer]]
 
 
-def bits_to_ints(bitstrings: Union[Sequence[int], NDArray[np.uint]]) -> NDArray[np.uint]:
+def bits_to_ints(bitstrings: Union[Sequence[int], NDArray[np.uint]]) -> NDArray[np.integer]:
     """Returns the integer specified by the given big-endian bitstrings.
 
     Args:
         bitstrings: A bitstring or array of bitstrings, each of which has the 1s bit (LSB) at the end.
     Returns:
         An array of integers; one for each bitstring.
     """
+    from qualtran import QUInt
+
     bitstrings = np.atleast_2d(bitstrings)
-    if bitstrings.shape[1] > 64:
-        raise NotImplementedError()
-    basis = 2 ** np.arange(bitstrings.shape[1] - 1, 0 - 1, -1, dtype=np.uint64)
-    return np.sum(basis * bitstrings, axis=1, dtype=np.uint64)
+    return QUInt(bitstrings.shape[1]).from_bits_array(bitstrings)
 
 
 def ints_to_bits(
@@ -61,8 +60,13 @@ def ints_to_bits(
         x: An integer or array of unsigned integers.
         w: The bit width of the returned bitstrings.
     """
+    from qualtran import QInt, QUInt
+
     x = np.atleast_1d(x)
-    return np.array([list(map(int, np.binary_repr(v, width=w))) for v in x], dtype=np.uint8)
+    if np.all(x >= 0):
+        return QUInt(w).to_bits_array(x)
+    else:
+        return QInt(w).to_bits_array(x)
 
 
 def _get_in_vals(

qualtran/simulation/classical_sim_test.py

@@ -170,33 +170,21 @@ def test_apply_classical_cbloq():
     np.testing.assert_array_equal(z, xarr)
 
 
-@pytest.mark.parametrize(
-    ['x', 'y', 'n_bits'],
-    [
-        (x, y, n_bits)
-        for n_bits in range(1, 5)
-        for x, y in itertools.product(range(1 << n_bits), repeat=2)
-    ],
-)
-def test_add_ints_unsigned(x, y, n_bits):
-    assert add_ints(x, y, num_bits=n_bits, is_signed=False) == (x + y) % (1 << n_bits)
+@pytest.mark.parametrize('n_bits', range(1, 5))
+def test_add_ints_unsigned(n_bits):
+    for x, y in itertools.product(range(1 << n_bits), repeat=2):
+        assert add_ints(x, y, num_bits=n_bits, is_signed=False) == (x + y) % (1 << n_bits)
 
 
-@pytest.mark.parametrize(
-    ['x', 'y', 'n_bits'],
-    [
-        (x, y, n_bits)
-        for n_bits in range(2, 5)
-        for x, y in itertools.product(range(-(2 ** (n_bits - 1)), 2 ** (n_bits - 1)), repeat=2)
-    ],
-)
-def test_add_ints_signed(x, y, n_bits):
+@pytest.mark.parametrize('n_bits', range(2, 5))
+def test_add_ints_signed(n_bits: int):
     half_n = 1 << (n_bits - 1)
     # Addition of signed ints `x` and `y` is a cyclic rotation of the interval [-2^(n-1), 2^(n-1)) by `y`.
     interval = [*range(-(2 ** (n_bits - 1)), 2 ** (n_bits - 1))]
-    i = x + half_n  # position of `x` in the interval
-    z = interval[(i + y) % len(interval)]  # rotate by `y`
-    assert add_ints(x, y, num_bits=n_bits, is_signed=True) == z
+    for x, y in itertools.product(range(-(2 ** (n_bits - 1)), 2 ** (n_bits - 1)), repeat=2):
+        i = x + half_n  # position of `x` in the interval
+        z = interval[(i + y) % len(interval)]  # rotate by `y`
+        assert add_ints(x, y, num_bits=n_bits, is_signed=True) == z
 
 
 @pytest.mark.notebook