Add SymmetricBanded matrix block encoding

qualtran/bloqs/block_encoding/block_encoding.ipynb

@@ -1412,9 +1412,28 @@
     "explicit_matrix_block_encoding = SparseMatrix(row_oracle, col_oracle, entry_oracle, eps=0)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8f51f49b",
+   "metadata": {
+    "cq.autogen": "SparseMatrix.symmetric_banded_matrix_block_encoding"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.block_encoding.sparse_matrix import SymmetricBandedRowColumnOracle\n",
+    "\n",
+    "row_oracle = SymmetricBandedRowColumnOracle(3, bandsize=1)\n",
+    "col_oracle = SymmetricBandedRowColumnOracle(3, bandsize=1)\n",
+    "entry_oracle = UniformEntryOracle(3, entry=0.3)\n",
+    "symmetric_banded_matrix_block_encoding = SparseMatrix(\n",
+    "    row_oracle, col_oracle, entry_oracle, eps=0\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
-   "id": "cd3c5958",
+   "id": "d3b284ec",
    "metadata": {
     "cq.autogen": "SparseMatrix.graphical_signature.md"
    },
@@ -1425,20 +1444,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3e309a00",
+   "id": "aeb878e9",
    "metadata": {
     "cq.autogen": "SparseMatrix.graphical_signature.py"
    },
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([sparse_matrix_block_encoding, explicit_matrix_block_encoding],\n",
-    "           ['`sparse_matrix_block_encoding`', '`explicit_matrix_block_encoding`'])"
+    "show_bloqs([sparse_matrix_block_encoding, explicit_matrix_block_encoding, symmetric_banded_matrix_block_encoding],\n",
+    "           ['`sparse_matrix_block_encoding`', '`explicit_matrix_block_encoding`', '`symmetric_banded_matrix_block_encoding`'])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "39cc32a3",
+   "id": "84a1bbe1",
    "metadata": {
     "cq.autogen": "SparseMatrix.call_graph.md"
    },
@@ -1449,7 +1468,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bcad0ddf",
+   "id": "2ca44b35",
    "metadata": {
     "cq.autogen": "SparseMatrix.call_graph.py"
    },

qualtran/bloqs/block_encoding/sparse_matrix.py

@@ -14,7 +14,7 @@
 
 import abc
 from functools import cached_property
-from typing import cast, Dict, Iterable, Tuple
+from typing import Dict, Iterable, Set, Tuple
 
 import numpy as np
 from attrs import field, frozen
@@ -35,6 +35,7 @@
     Soquet,
     SoquetT,
 )
+from qualtran.bloqs.arithmetic import Add, AddK
 from qualtran.bloqs.basic_gates import Ry, Swap
 from qualtran.bloqs.block_encoding import BlockEncoding
 from qualtran.bloqs.bookkeeping.auto_partition import AutoPartition, Unused
@@ -43,6 +44,8 @@
 from qualtran.bloqs.state_preparation.prepare_uniform_superposition import (
     PrepareUniformSuperposition,
 )
+from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+from qualtran.simulation.classical_sim import ClassicalValT
 from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt
 from qualtran.symbolics.math_funcs import bit_length
 
@@ -229,7 +232,8 @@ def build_composite_bloq(
         if is_symbolic(self.system_bitsize) or is_symbolic(self.row_oracle.num_nonzero):
             raise DecomposeTypeError(f"Cannot decompose symbolic {self=}")
 
-        ancilla_bits = bb.split(cast(Soquet, ancilla))
+        assert not isinstance(ancilla, np.ndarray)
+        ancilla_bits = bb.split(ancilla)
         q, l = ancilla_bits[0], bb.join(ancilla_bits[1:])
 
         unused = self.system_bitsize - bit_length(self.row_oracle.num_nonzero - 1)
@@ -240,10 +244,10 @@ def build_composite_bloq(
             ],
         )
         l = bb.add(diffusion, target=l)
-        l, system = bb.add_t(self.col_oracle, l=cast(Soquet, l), i=system)
-        q, l, system = bb.add_t(self.entry_oracle, q=q, i=l, j=system)
-        l, system = bb.add_t(Swap(self.system_bitsize), x=l, y=system)
-        l, system = bb.add_t(self.row_oracle.adjoint(), l=l, i=system)
+        l, system = bb.add(self.col_oracle, l=l, i=system)
+        q, l, system = bb.add(self.entry_oracle, q=q, i=l, j=system)
+        l, system = bb.add(Swap(self.system_bitsize), x=l, y=system)
+        l, system = bb.add(self.row_oracle.adjoint(), l=l, i=system)
         l = bb.add(diffusion.adjoint(), target=l)
 
         return {"system": system, "ancilla": bb.join(np.concatenate([[q], bb.split(l)]))}
@@ -281,6 +285,68 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So
         return soqs
 
 
+@frozen
+class SymmetricBandedRowColumnOracle(RowColumnOracle):
+    """Oracle specifying the non-zero rows and columns of a symmetric
+    [banded matrix](https://en.wikipedia.org/wiki/Band_matrix).
+
+    The symmetry here refers to the pattern of non-zero entries, not necessarily the entries themselves, which are determined separately by the `EntryOracle`.
+
+    Args:
+        system_bitsize: The number of bits used to represent an index.
+        bandsize: The number of pairs of non-zero off-main diagonals. A diagonal matrix has
+            bandsize 0 and a tridiagonal matrix has bandsize 1.
+
+    Registers:
+        l: As input, index specifying the `l`-th non-zero entry to find in row / column `i`.
+           As output, position of the `l`-th non-zero entry in row / column `i`.
+        i: The row / column index.
+    """
+
+    system_bitsize: SymbolicInt
+    bandsize: SymbolicInt
+
+    @cached_property
+    def num_nonzero(self) -> SymbolicInt:
+        return 2 * self.bandsize + 1
+
+    def __attrs_post_init__(self):
+        if is_symbolic(self.system_bitsize) or is_symbolic(self.bandsize):
+            return
+        if 2**self.system_bitsize < 2 * self.bandsize:
+            raise ValueError(
+                f"bandsize={self.bandsize} too large for system_bitsize={self.system_bitsize}"
+            )
+
+    def call_classically(self, l: ClassicalValT, i: ClassicalValT) -> Tuple[ClassicalValT, ...]:
+        if (
+            is_symbolic(self.bandsize)
+            or is_symbolic(self.system_bitsize)
+            or is_symbolic(self.num_nonzero)
+        ):
+            raise DecomposeTypeError(f"Cannot call symbolic {self=} classically")
+
+        assert not isinstance(l, np.ndarray) and not isinstance(i, np.ndarray)
+        if l >= self.num_nonzero:
+            raise IndexError("l out of bounds")
+        return ((l + i - self.bandsize) % (2**self.system_bitsize), i)
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set[BloqCountT]:
+        return {
+            (Add(QUInt(self.system_bitsize), QUInt(self.system_bitsize)), 1),
+            (AddK(self.system_bitsize, -self.bandsize, signed=True), 1),
+        }
+
+    def build_composite_bloq(self, bb: BloqBuilder, l: SoquetT, i: SoquetT) -> Dict[str, SoquetT]:
+        if is_symbolic(self.system_bitsize) or is_symbolic(self.bandsize):
+            raise DecomposeTypeError(f"Cannot decompose symbolic {self=}")
+
+        i, l = bb.add(Add(QUInt(self.system_bitsize), QUInt(self.system_bitsize)), a=i, b=l)
+        l = bb.add(AddK(self.system_bitsize, -self.bandsize, signed=True), x=l)
+
+        return {"l": l, "i": i}
+
+
 @frozen
 class UniformEntryOracle(EntryOracle):
     """Oracle specifying the entries of a matrix with uniform entries."""
@@ -293,7 +359,7 @@ def build_composite_bloq(
     ) -> Dict[str, SoquetT]:
         # Either Rx or Ry work here; Rx would induce a phase on the subspace with non-zero ancilla
         # See https://arxiv.org/abs/2302.10949 for reference that uses Rx
-        soqs["q"] = cast(Soquet, bb.add(Ry(2 * np.arccos(self.entry)), q=q))
+        soqs["q"] = bb.add(Ry(2 * np.arccos(self.entry)), q=q)
         return soqs
 
 
@@ -332,8 +398,8 @@ def __attrs_post_init__(self):
             raise ValueError("data must have dimension 2 ** self.system_bitsize")
         if not is_symbolic(self.entry_bitsize) and self.entry_bitsize < 1:
             raise ValueError("entry_bitsize must be >= 1")
-        if not all(x >= 0 and x < 1 for x in self.data.flat):
-            raise ValueError("entries must be >= 0 and < 1")
+        if not all(x >= 0 and x <= 1 for x in self.data.flat):
+            raise ValueError("entries must be >= 0 and <= 1")
 
     def build_composite_bloq(
         self, bb: BloqBuilder, q: SoquetT, i: SoquetT, j: SoquetT
@@ -344,20 +410,20 @@ def build_composite_bloq(
         data = np.arccos(self.data) / np.pi * 2**self.entry_bitsize
         qrom = QROM.build_from_data(data, target_bitsizes=(self.entry_bitsize,))
         target = bb.allocate(self.entry_bitsize)
-        i, j, target = bb.add_t(qrom, selection0=i, selection1=j, target0_=target)
+        i, j, target = bb.add(qrom, selection0=i, selection1=j, target0_=target)
         # perform fractional Ry
         # can't use StatePreparationViaRotations here because the coefficients depend on i, j
         # can't use QvrZPow here because CRz is not symmetric and we condition on target, not q
         # TODO: could potentially use RzViaPhaseGradient when it is done
-        target_bits = bb.split(cast(Soquet, target))
+        target_bits = bb.split(target)
         for k in range(len(target_bits)):
-            target_bits[k], q = bb.add_t(
+            target_bits[k], q = bb.add(
                 Ry(2 * np.pi * (2 ** -(k + 1))).controlled(), ctrl=target_bits[k], q=q
             )
         target = bb.join(target_bits)
         # unload from QROM
-        i, j, target = bb.add_t(qrom.adjoint(), selection0=i, selection1=j, target0_=target)
-        bb.free(cast(Soquet, target))
+        i, j, target = bb.add(qrom.adjoint(), selection0=i, selection1=j, target0_=target)
+        bb.free(target)
         return {"q": q, "i": i, "j": j}
 
 
@@ -390,8 +456,25 @@ def _explicit_matrix_block_encoding() -> SparseMatrix:
     return explicit_matrix_block_encoding
 
 
+@bloq_example
+def _symmetric_banded_matrix_block_encoding() -> SparseMatrix:
+    from qualtran.bloqs.block_encoding.sparse_matrix import SymmetricBandedRowColumnOracle
+
+    row_oracle = SymmetricBandedRowColumnOracle(3, bandsize=1)
+    col_oracle = SymmetricBandedRowColumnOracle(3, bandsize=1)
+    entry_oracle = UniformEntryOracle(3, entry=0.3)
+    symmetric_banded_matrix_block_encoding = SparseMatrix(
+        row_oracle, col_oracle, entry_oracle, eps=0
+    )
+    return symmetric_banded_matrix_block_encoding
+
+
 _SPARSE_MATRIX_DOC = BloqDocSpec(
     bloq_cls=SparseMatrix,
     import_line="from qualtran.bloqs.block_encoding import SparseMatrix",
-    examples=[_sparse_matrix_block_encoding, _explicit_matrix_block_encoding],
+    examples=[
+        _sparse_matrix_block_encoding,
+        _explicit_matrix_block_encoding,
+        _symmetric_banded_matrix_block_encoding,
+    ],
 )