diff --git a/qualtran/bloqs/block_encoding/block_encoding.ipynb b/qualtran/bloqs/block_encoding/block_encoding.ipynb index f89357053..310f2dc5f 100644 --- a/qualtran/bloqs/block_encoding/block_encoding.ipynb +++ b/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" }, diff --git a/qualtran/bloqs/block_encoding/sparse_matrix.py b/qualtran/bloqs/block_encoding/sparse_matrix.py index 17e4c0a32..cf285ddf9 100644 --- a/qualtran/bloqs/block_encoding/sparse_matrix.py +++ b/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, + ], ) diff --git a/qualtran/bloqs/block_encoding/sparse_matrix_test.py b/qualtran/bloqs/block_encoding/sparse_matrix_test.py index a7205939d..8db8778f4 100644 --- a/qualtran/bloqs/block_encoding/sparse_matrix_test.py +++ b/qualtran/bloqs/block_encoding/sparse_matrix_test.py @@ -17,13 +17,16 @@ import numpy as np import pytest +import qualtran.testing as qlt_testing from qualtran import BloqBuilder, QAny, Register, Signature, Soquet from qualtran.bloqs.basic_gates import IntEffect, IntState from qualtran.bloqs.block_encoding.sparse_matrix import ( _explicit_matrix_block_encoding, _sparse_matrix_block_encoding, + _symmetric_banded_matrix_block_encoding, ExplicitEntryOracle, SparseMatrix, + SymmetricBandedRowColumnOracle, TopLeftRowColumnOracle, UniformEntryOracle, ) @@ -37,6 +40,10 @@ def test_explicit_matrix(bloq_autotester): bloq_autotester(_explicit_matrix_block_encoding) +def test_symmetric_banded_matrix(bloq_autotester): + bloq_autotester(_symmetric_banded_matrix_block_encoding) + + def test_sparse_matrix_signature(): bloq = _sparse_matrix_block_encoding() assert bloq.signature == Signature( @@ -99,7 +106,7 @@ def test_explicit_entry_oracle(n, data): np.testing.assert_allclose(data, from_tensors, atol=0.003) -test_matrix = [ +topleft_matrix = [ [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0], @@ -125,5 +132,147 @@ def test_top_left_matrix(): bb.add(IntEffect(0, 3 + 1), val=ancilla) bloq = bb.finalize(system=system) + from_tensors = bloq.tensor_contract() * alpha + np.testing.assert_allclose(topleft_matrix, from_tensors, atol=0.003) + + +test_matrix = [ + [0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.3], + [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.3, 0.3, 0.3, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.3, 0.3, 0.3, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.3, 0.3, 0.3, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.3, 0.3, 0.3], + [0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.3, 0.3], +] + +test_matrix_nonzeros = [ + [7, 0, 1], + [0, 1, 2], + [1, 2, 3], + [2, 3, 4], + [3, 4, 5], + [4, 5, 6], + [5, 6, 7], + [6, 7, 0], +] + + +def test_symmetric_banded_row_column_oracle_classical(): + n = 3 + bloq = SymmetricBandedRowColumnOracle(n, bandsize=1) + + def test_entry(l, i): + l_out, i_out = bloq.call_classically(l=l, i=i) + assert i_out == i + assert l_out == test_matrix_nonzeros[i][l] + + for i in range(2**n): + for l in range(3): + test_entry(l, i) + + +def test_symmetric_banded_row_column_oracle(): + n = 3 + bloq = SymmetricBandedRowColumnOracle(n, bandsize=1) + + def test_entry(l, i): + bb = BloqBuilder() + l_soq = cast(Soquet, bb.add(IntState(l, n))) + i_soq = cast(Soquet, bb.add(IntState(i, n))) + l_soq, i_soq = bb.add_t(bloq, l=l_soq, i=i_soq) + bb.add(IntEffect(i, n), val=i_soq) + out = bb.finalize(l=l_soq) + np.testing.assert_allclose( + IntState(test_matrix_nonzeros[i][l], n).tensor_contract(), out.tensor_contract() + ) + + for i in range(2**n): + for l in range(3): + test_entry(l, i) + + +def test_symmetric_banded_row_column_matrix(): + n = 3 + bloq = _symmetric_banded_matrix_block_encoding() + alpha = bloq.alpha + + bb = BloqBuilder() + system = bb.add_register("system", n) + ancilla = cast(Soquet, bb.add(IntState(0, n + 1))) + system, ancilla = bb.add_t(bloq, system=system, ancilla=ancilla) + bb.add(IntEffect(0, n + 1), val=ancilla) + bloq = bb.finalize(system=system) + from_tensors = bloq.tensor_contract() * alpha np.testing.assert_allclose(test_matrix, from_tensors, atol=0.003) + + +@pytest.mark.slow +def test_matrix_stress(): + rs = np.random.RandomState(1234) + f = lambda: rs.randint(0, 10) / 10 + data = [ + [f(), f(), f(), 0.0, 0.0, 0.0, 0.0, 0.0], + [f(), f(), f(), f(), 0.0, 0.0, 0.0, 0.0], + [f(), f(), f(), f(), f(), 0.0, 0.0, 0.0], + [0.0, f(), f(), f(), f(), f(), 0.0, 0.0], + [0.0, 0.0, f(), f(), f(), f(), f(), 0.0], + [0.0, 0.0, 0.0, f(), f(), f(), f(), f()], + [0.0, 0.0, 0.0, 0.0, f(), f(), f(), f()], + [0.0, 0.0, 0.0, 0.0, 0.0, f(), f(), f()], + ] + n = 3 + row_oracle = SymmetricBandedRowColumnOracle(n, bandsize=2) + col_oracle = SymmetricBandedRowColumnOracle(n, bandsize=2) + entry_oracle = ExplicitEntryOracle(system_bitsize=n, data=np.array(data), entry_bitsize=7) + bloq = SparseMatrix(row_oracle, col_oracle, entry_oracle, eps=0) + alpha = bloq.alpha + + bb = BloqBuilder() + system = bb.add_register("system", n) + ancilla = cast(Soquet, bb.add(IntState(0, n + 1))) + system, ancilla = bb.add_t(bloq, system=system, ancilla=ancilla) + bb.add(IntEffect(0, n + 1), val=ancilla) + bloq = bb.finalize(system=system) + + from_tensors = bloq.tensor_contract() * alpha + np.testing.assert_allclose(data, from_tensors, atol=0.03) + + +def gen_vlasov_hamiltonian(n, alpha, m): + data = np.zeros((2**n, 2**n)) + data[0][1] = data[1][0] = np.sqrt((1 + alpha) / 2) + for i in range(2, m + 1): + data[i - 1][i] = data[i][i - 1] = np.sqrt(i / 2) + data /= np.max(data) + return data + + +@pytest.mark.slow +def test_vlasov_explicit(): + n = 3 + k = 2 + alpha = 2 / k**2 + data = gen_vlasov_hamiltonian(n, alpha, m=(2**n - 1)) + row_oracle = SymmetricBandedRowColumnOracle(n, bandsize=1) + col_oracle = SymmetricBandedRowColumnOracle(n, bandsize=1) + entry_oracle = ExplicitEntryOracle(system_bitsize=n, data=data, entry_bitsize=7) + bloq = SparseMatrix(row_oracle, col_oracle, entry_oracle, eps=0) + alpha = bloq.alpha + + bb = BloqBuilder() + system = bb.add_register("system", n) + ancilla = cast(Soquet, bb.add(IntState(0, n + 1))) + system, ancilla = bb.add_t(bloq, system=system, ancilla=ancilla) + bb.add(IntEffect(0, n + 1), val=ancilla) + bloq = bb.finalize(system=system) + + from_tensors = bloq.tensor_contract() * alpha + np.testing.assert_allclose(data, from_tensors, atol=0.02) + + +def test_symmetric_banded_counts(): + bloq = SymmetricBandedRowColumnOracle(3, bandsize=1) + qlt_testing.assert_equivalent_bloq_counts(bloq) diff --git a/qualtran/conftest.py b/qualtran/conftest.py index 60f2f8877..a17353e68 100644 --- a/qualtran/conftest.py +++ b/qualtran/conftest.py @@ -108,6 +108,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample): 'sparse_permutation', 'permutation_cycle_symb', 'explicit_matrix_block_encoding', # cannot serialize AutoPartition + 'symmetric_banded_matrix_block_encoding', # cannot serialize AutoPartition ]: pytest.xfail("Skipping serialization test for bloq examples that cannot yet be serialized.")