Implement Kikuchi graph oracles and block-encoding (#1555)

part of #1348 

Implements the Kikuchi graph adjacency matrix and adjacency list
oracles, along with ones to load and index the constraints. And uses the
oracles to build a sparse matrix block-encoding of the Kikuchi matrix.

---------

Co-authored-by: Matthew Harrigan <[email protected]>

Author: anurudhp

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -107,6 +107,7 @@
 import qualtran.bloqs.multiplexers.black_box_select
 import qualtran.bloqs.multiplexers.select_base
 import qualtran.bloqs.multiplexers.select_pauli_lcu
+import qualtran.bloqs.optimization.k_xor_sat
 import qualtran.bloqs.optimization.k_xor_sat.kikuchi_guiding_state
 import qualtran.bloqs.phase_estimation.lp_resource_state
 import qualtran.bloqs.phase_estimation.qubitization_qpe
@@ -839,14 +840,41 @@
 # -----   Optimization   ---------------------------------------------------
 # --------------------------------------------------------------------------
 OPTIMIZATION: List[NotebookSpecV2] = [
+    # -----   Algorithm  ------------------------------------------
+    NotebookSpecV2(
+        title='kXOR: Instance load Oracles',
+        module=qualtran.bloqs.optimization.k_xor_sat.load_kxor_instance,
+        bloq_specs=[qualtran.bloqs.optimization.k_xor_sat.load_kxor_instance._LOAD_INSTANCE_DOC],
+    ),
     NotebookSpecV2(
         title='Planted Noisy kXOR - Kikuchi Guiding State',
         module=qualtran.bloqs.optimization.k_xor_sat.kikuchi_guiding_state,
         bloq_specs=[
             qualtran.bloqs.optimization.k_xor_sat.kikuchi_guiding_state._SIMPLE_GUIDING_STATE_DOC,
             qualtran.bloqs.optimization.k_xor_sat.kikuchi_guiding_state._GUIDING_STATE_DOC,
         ],
-    )
+    ),
+    NotebookSpecV2(
+        title='Planted Noisy kXOR: Kikuchi Adjacency List',
+        module=qualtran.bloqs.optimization.k_xor_sat.kikuchi_adjacency_list,
+        bloq_specs=[
+            qualtran.bloqs.optimization.k_xor_sat.kikuchi_adjacency_list._KIKUCHI_NONZERO_INDEX_DOC
+        ],
+    ),
+    NotebookSpecV2(
+        title='Planted Noisy kXOR: Kikuchi Adjacency Matrix',
+        module=qualtran.bloqs.optimization.k_xor_sat.kikuchi_adjacency_matrix,
+        bloq_specs=[
+            qualtran.bloqs.optimization.k_xor_sat.kikuchi_adjacency_matrix._KIKUCHI_MATRIX_ENTRY_DOC
+        ],
+    ),
+    NotebookSpecV2(
+        title='Planted Noisy kXOR: Block-encoding the Kikuchi Matrix',
+        module=qualtran.bloqs.optimization.k_xor_sat.kikuchi_block_encoding,
+        bloq_specs=[
+            qualtran.bloqs.optimization.k_xor_sat.kikuchi_block_encoding._KIKUCHI_HAMILTONIAN_DOC
+        ],
+    ),
 ]
 
 # --------------------------------------------------------------------------

docs/bloqs/index.rst

@@ -156,7 +156,11 @@ Bloqs Library
     :maxdepth: 2
     :caption: Optimization:
 
+    optimization/k_xor_sat/load_kxor_instance.ipynb
     optimization/k_xor_sat/kikuchi_guiding_state.ipynb
+    optimization/k_xor_sat/kikuchi_adjacency_list.ipynb
+    optimization/k_xor_sat/kikuchi_adjacency_matrix.ipynb
+    optimization/k_xor_sat/kikuchi_block_encoding.ipynb
 
 .. toctree::
     :maxdepth: 2

qualtran/bloqs/block_encoding/sparse_matrix_hermitian.py

@@ -15,18 +15,15 @@
 from collections import Counter
 from functools import cached_property
 
-import attrs
 import numpy as np
 import sympy
 from attrs import frozen
 
 from qualtran import (
-    AddControlledT,
     Bloq,
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
-    CtrlSpec,
     DecomposeTypeError,
     QAny,
     QBit,
@@ -35,10 +32,10 @@
     Soquet,
     SoquetT,
 )
-from qualtran.bloqs.basic_gates import CSwap, Ry, Swap
+from qualtran.bloqs.basic_gates import Ry, Swap
 from qualtran.bloqs.block_encoding import BlockEncoding
 from qualtran.bloqs.block_encoding.sparse_matrix import RowColumnOracle
-from qualtran.bloqs.bookkeeping import Partition
+from qualtran.bloqs.bookkeeping import Always, Partition
 from qualtran.bloqs.bookkeeping.auto_partition import AutoPartition, Unused
 from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity
 from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare
@@ -133,18 +130,14 @@ class SparseMatrixHermitian(BlockEncoding):
     col_oracle: RowColumnOracle
     entry_oracle: SqrtEntryOracle
     eps: SymbolicFloat
-    is_controlled: bool = False
 
     def __attrs_post_init__(self):
         if self.col_oracle.system_bitsize != self.entry_oracle.system_bitsize:
             raise ValueError("column and entry oracles must have same bitsize")
 
     @cached_property
     def signature(self) -> Signature:
-        n_ctrls = 1 if self.is_controlled else 0
-
         return Signature.build_from_dtypes(
-            ctrl=QAny(n_ctrls),
             system=QAny(self.system_bitsize),
             ancilla=QAny(self.ancilla_bitsize),
             resource=QAny(self.resource_bitsize),  # if resource_bitsize is 0, not present
@@ -190,29 +183,23 @@ def diffusion(self):
     def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT:
         counts = Counter[Bloq]()
 
-        counts[self.diffusion] += 1
-        counts[self.col_oracle] += 1
-        counts[self.entry_oracle] += 1
-        if self.is_controlled:
-            counts[CSwap(self.system_bitsize)] += 1
-            counts[CSwap(1)] += 1
-        else:
-            counts[Swap(self.system_bitsize)] += 1
-            counts[Swap(1)] += 1
-        counts[self.entry_oracle.adjoint()] += 1
-        counts[self.col_oracle.adjoint()] += 1
-        counts[self.diffusion.adjoint()] += 1
+        counts[Always(self.diffusion)] += 1
+        counts[Always(self.col_oracle)] += 1
+        counts[Always(self.entry_oracle)] += 1
+        counts[Swap(self.system_bitsize)] += 1
+        counts[Swap(1)] += 1
+        counts[Always(self.entry_oracle.adjoint())] += 1
+        counts[Always(self.col_oracle.adjoint())] += 1
+        counts[Always(self.diffusion.adjoint())] += 1
 
         return counts
 
     def build_composite_bloq(
-        self, bb: BloqBuilder, system: SoquetT, ancilla: SoquetT, **soqs
+        self, bb: BloqBuilder, system: SoquetT, ancilla: SoquetT
     ) -> dict[str, SoquetT]:
         if is_symbolic(self.system_bitsize) or is_symbolic(self.col_oracle.num_nonzero):
             raise DecomposeTypeError(f"Cannot decompose symbolic {self=}")
 
-        ctrl = soqs.pop('ctrl', None)
-
         assert not isinstance(ancilla, np.ndarray)
         partition_ancilla = Partition(
             n=self.ancilla_bitsize,
@@ -225,42 +212,24 @@ def build_composite_bloq(
 
         a, l, b = bb.add(partition_ancilla, x=ancilla)
 
-        l = bb.add(self.diffusion, target=l)
-        l, system = bb.add(self.col_oracle, l=l, i=system)
-        b, l, system = bb.add(self.entry_oracle, q=b, i=l, j=system)
+        l = bb.add(Always(self.diffusion), target=l)
+        l, system = bb.add(Always(self.col_oracle), l=l, i=system)
+        b, l, system = bb.add(Always(self.entry_oracle), q=b, i=l, j=system)
 
-        if self.is_controlled:
-            ctrl, l, system = bb.add(CSwap(self.system_bitsize), ctrl=ctrl, x=l, y=system)
-            ctrl, a, b = bb.add(CSwap(1), ctrl=ctrl, x=a, y=b)
-        else:
-            l, system = bb.add(Swap(self.system_bitsize), x=l, y=system)
-            a, b = bb.add(Swap(1), x=a, y=b)
+        l, system = bb.add(Swap(self.system_bitsize), x=l, y=system)
+        a, b = bb.add(Swap(1), x=a, y=b)
 
-        b, l, system = bb.add(self.entry_oracle.adjoint(), q=b, i=l, j=system)
-        l, system = bb.add(self.col_oracle.adjoint(), l=l, i=system)
-        l = bb.add(self.diffusion.adjoint(), target=l)
+        b, l, system = bb.add(Always(self.entry_oracle.adjoint()), q=b, i=l, j=system)
+        l, system = bb.add(Always(self.col_oracle.adjoint()), l=l, i=system)
+        l = bb.add(Always(self.diffusion.adjoint()), target=l)
 
         ancilla = bb.add(partition_ancilla.adjoint(), a=a, l=l, b=b)
 
-        out_soqs = {"system": system, "ancilla": ancilla}
-        if self.is_controlled:
-            out_soqs |= {"ctrl": ctrl}
-        return out_soqs
+        return {"system": system, "ancilla": ancilla}
 
     def adjoint(self) -> 'SparseMatrixHermitian':
         return self
 
-    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
-
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self,
-            ctrl_spec,
-            current_ctrl_bit=1 if self.is_controlled else None,
-            bloq_with_ctrl=self if self.is_controlled else attrs.evolve(self, is_controlled=True),
-            ctrl_reg_name='ctrl',
-        )
-
 
 @frozen
 class UniformSqrtEntryOracle(SqrtEntryOracle):

qualtran/bloqs/bookkeeping/always.py

@@ -11,7 +11,7 @@
 #  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.
-from typing import Iterable, Optional, Sequence
+from typing import Iterable, Optional, Sequence, Union
 
 import attrs
 
@@ -25,6 +25,7 @@
     Signature,
     SoquetT,
 )
+from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
 
 
 @attrs.frozen
@@ -67,6 +68,11 @@ def signature(self) -> 'Signature':
     def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str, 'SoquetT']:
         return bb.add_d(self.subbloq, **soqs)
 
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', set['BloqCountT']]:
+        return self.subbloq.build_call_graph(ssa)
+
     def get_ctrl_system(
         self, ctrl_spec: Optional['CtrlSpec'] = None
     ) -> tuple['Bloq', 'AddControlledT']:

qualtran/bloqs/optimization/k_xor_sat/__init__.py

@@ -11,5 +11,9 @@
 #  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.
+from .kikuchi_adjacency_list import KikuchiNonZeroIndex
+from .kikuchi_adjacency_matrix import KikuchiMatrixEntry
+from .kikuchi_block_encoding import KikuchiHamiltonian, KikuchiMatrixEntry, KikuchiNonZeroIndex
 from .kikuchi_guiding_state import GuidingState, SimpleGuidingState
 from .kxor_instance import Constraint, KXorInstance
+from .load_kxor_instance import LoadConstraintScopes, LoadUniqueScopeIndex, PRGAUniqueConstraintRHS