Vectorized variants of to_bits and from_bits

qualtran/_infra/data_types.py

@@ -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})'
 
@@ -324,10 +344,43 @@ 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=}")
+
+        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] > 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)
+
     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}")
@@ -486,12 +539,31 @@ def num_int(self) -> SymbolicInt:
         return self.bitsize - self.num_frac - int(self.signed)
 
     @property
-    def fxp_dtype_str(self) -> str:
-        return f'fxp-{"us"[self.signed]}{self.bitsize}/{self.num_frac}'
+    def fxp_dtype_template(self) -> Fxp:
+        """A template of the `Fxp` data type for classical values.
+
+        -   op_sizing='same' and const_op_sizing='same' ensure that the returned object is not resized
+            to a bigger fixed point number when doing operations with other Fxp objects.
+        -   shifting='trunc' ensures that when shifting the Fxp integer to left / right; the digits are
+            truncated and no rounding occurs
+        -   overflow='wrap' ensures that when performing operations where result overflows, the overflowed
+            digits are simply discarded.
+        """
+        if is_symbolic(self.bitsize) or is_symbolic(self.num_frac):
+            raise ValueError(
+                "Cannot construct Fxp template for symbolic bitsizes {self.bitsize=}, {self.num_frac=}"
+            )
 
-    @property
-    def _fxp_dtype(self) -> Fxp:
-        return Fxp(None, dtype=self.fxp_dtype_str)
+        return Fxp(
+            None,
+            n_word=self.bitsize,
+            n_frac=self.num_frac,
+            signed=self.signed,
+            op_sizing='same',
+            const_op_sizing='same',
+            shifting='trunc',
+            overflow='wrap',
+        )
 
     def is_symbolic(self) -> bool:
         return is_symbolic(self.bitsize, self.num_frac)
@@ -517,7 +589,7 @@ def to_bits(
             sign = int(x < 0)
             x = abs(x)
         fxp = x if isinstance(x, Fxp) else Fxp(x)
-        bits = [int(x) for x in fxp.like(self._fxp_dtype).bin()]
+        bits = [int(x) for x in fxp.like(self.fxp_dtype_template).bin()]
         if self.signed and not complement:
             bits[0] = sign
         return bits
@@ -526,7 +598,14 @@ def from_bits(self, bits: Sequence[int]) -> Fxp:
         """Combine individual bits to form x"""
         bits_bin = "".join(str(x) for x in bits[:])
         fxp_bin = "0b" + bits_bin[: -self.num_frac] + "." + bits_bin[-self.num_frac :]
-        return Fxp(fxp_bin, dtype=self.fxp_dtype_str)
+        return Fxp(fxp_bin, like=self.fxp_dtype_template)
+
+    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."""
@@ -546,19 +625,37 @@ def __attrs_post_init__(self):
     def get_classical_domain(self) -> Iterable[Fxp]:
         qint = QIntOnesComp(self.bitsize) if self.signed else QUInt(self.bitsize)
         for x in qint.get_classical_domain():
-            yield Fxp(x / 2**self.num_frac, dtype=self.fxp_dtype_str)
+            yield Fxp(x / 2**self.num_frac).like(self.fxp_dtype_template)
 
     def _assert_valid_classical_val(self, val: Union[float, Fxp], debug_str: str = 'val'):
         fxp_val = val if isinstance(val, Fxp) else Fxp(val)
-        if fxp_val.get_val() != fxp_val.like(self._fxp_dtype).get_val():
+        if fxp_val.get_val() != fxp_val.like(self.fxp_dtype_template).get_val():
             raise ValueError(
                 f"{debug_str}={val} cannot be accurately represented using Fxp {fxp_val}"
             )
 
     def assert_valid_classical_val(self, val: Union[float, Fxp], debug_str: str = 'val'):
-        # TODO: Asserting a valid value here opens a can of worms because classical data, except integers,
-        # is currently not propagated correctly through Bloqs
-        pass
+        assert isinstance(val, Fxp)
+        assert val.overflow == 'wrap'
+        assert val.shifting == 'trunc'
+        self._assert_valid_classical_val(val, debug_str)
+
+    def float_to_fxp(
+        self, val: Union[float, int], *, raw: bool = False, require_exact: bool = True
+    ) -> Fxp:
+        """Convert a floating point value to an Fxp constant of this dtype.
+
+        If `raw` is True, then returns `val / 2**self.n_frac` instead.
+
+        Args:
+            val: Floating point value.
+            raw: Convert from a raw integer value instead
+            require_exact: If True, represent the input `val` exactly and raise
+                           a ValueError if it cannot be represented.
+        """
+        if require_exact:
+            self._assert_valid_classical_val(val if not raw else val / 2**self.num_frac)
+        return Fxp(val, raw=raw, like=self.fxp_dtype_template)
 
     def __str__(self):
         if self.signed:

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
 
@@ -95,12 +97,12 @@ def test_qfxp():
     assert str(qfp_16) == 'QFxp(16, 15)'
     assert qfp_16.num_qubits == 16
     assert qfp_16.num_int == 1
-    assert qfp_16.fxp_dtype_str == 'fxp-u16/15'
+    assert qfp_16.fxp_dtype_template.dtype == 'fxp-u16/15'
     qfp_16 = QFxp(16, 15, signed=True)
     assert str(qfp_16) == 'QFxp(16, 15, True)'
     assert qfp_16.num_qubits == 16
     assert qfp_16.num_int == 0
-    assert qfp_16.fxp_dtype_str == 'fxp-s16/15'
+    assert qfp_16.fxp_dtype_template.dtype == 'fxp-s16/15'
     with pytest.raises(ValueError, match="num_qubits must be > 1."):
         QFxp(1, 1, signed=True)
     QFxp(1, 1, signed=False)
@@ -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(23).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]