From e45086563762b357ab9eeffdd9d5a542f589ea12 Mon Sep 17 00:00:00 2001 From: Anurudh Peduri Date: Tue, 3 Dec 2024 03:37:52 +0100 Subject: [PATCH] Implement Kikuchi guiding state preparation --- qualtran/bloqs/max_k_xor_sat/__init__.py | 15 + qualtran/bloqs/max_k_xor_sat/guiding_state.py | 445 ++++++++++++++++++ .../bloqs/max_k_xor_sat/guiding_state_test.py | 132 ++++++ .../guiding_state_tutorial.ipynb | 212 +++++++++ qualtran/bloqs/max_k_xor_sat/kxor_instance.py | 273 +++++++++++ .../bloqs/max_k_xor_sat/kxor_instance_test.py | 30 ++ 6 files changed, 1107 insertions(+) create mode 100644 qualtran/bloqs/max_k_xor_sat/__init__.py create mode 100644 qualtran/bloqs/max_k_xor_sat/guiding_state.py create mode 100644 qualtran/bloqs/max_k_xor_sat/guiding_state_test.py create mode 100644 qualtran/bloqs/max_k_xor_sat/guiding_state_tutorial.ipynb create mode 100644 qualtran/bloqs/max_k_xor_sat/kxor_instance.py create mode 100644 qualtran/bloqs/max_k_xor_sat/kxor_instance_test.py diff --git a/qualtran/bloqs/max_k_xor_sat/__init__.py b/qualtran/bloqs/max_k_xor_sat/__init__.py new file mode 100644 index 000000000..827b32662 --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 .guiding_state import GuidingState, SimpleGuidingState +from .kxor_instance import Constraint, KXorInstance diff --git a/qualtran/bloqs/max_k_xor_sat/guiding_state.py b/qualtran/bloqs/max_k_xor_sat/guiding_state.py new file mode 100644 index 000000000..012237dbd --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/guiding_state.py @@ -0,0 +1,445 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +r"""Prepare the guiding state for a kXOR instance $\mathcal{I}$ with +Kikuchi parameter $\ell$. + +References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Section 4.4.1, Theorem 4.15. +""" +from functools import cached_property + +from attrs import evolve, field, frozen + +from qualtran import ( + Bloq, + bloq_example, + BloqBuilder, + BloqDocSpec, + CtrlSpec, + DecomposeTypeError, + QAny, + QBit, + QUInt, + Register, + Signature, + Soquet, + SoquetT, +) +from qualtran.bloqs.arithmetic.lists import HasDuplicates, SortInPlace +from qualtran.bloqs.basic_gates import Hadamard, OnEach, XGate +from qualtran.bloqs.bookkeeping import Partition +from qualtran.bloqs.mcmt import MultiControlX +from qualtran.bloqs.state_preparation.prepare_base import PrepareOracle +from qualtran.bloqs.state_preparation.sparse_state_preparation_via_rotations import ( + SparseStatePreparationViaRotations, +) +from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator +from qualtran.symbolics import ceil, is_symbolic, log2, pi, SymbolicFloat, SymbolicInt + +from .kxor_instance import KXorInstance + + +@frozen +class SimpleGuidingState(PrepareOracle): + r"""Prepare the guiding state for $\ell = k$. + + Given an kXOR instance $\mathcal{I}$, prepare the guiding state for + parameter $\ell = k$ (i.e. $c = 1$), defined in Eq 134: + $$ + |\phi\rangle + \propto + |\Gamma^k(\mathcal{A})\rangle + = + \frac{1}{\sqrt{\tilde{m}}} + \sum_{S \in {[n] \choose k}} B_\mathcal{I}(S) |S\rangle + $$ + + Here, $\tilde{m}$ is the number of constraints in the input instance $\mathcal{I}$, + and $\mathcal{A} = \sqrt{\frac{{n\choose k}}{\tilde{m}}} \mathcal{I}$. + + This bloq has a gate cost of $O(\tilde{m} \log n)$ (see Eq 142 and paragraph below). + + Args: + inst: the kXOR instance $\mathcal{I}$. + eps: Precision of the prepared state (defaults to 1e-6). + + Registers: + S: a scope of $k$ variables, each in $[n]$. + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Equation 134. + """ + + inst: KXorInstance + eps: SymbolicFloat = field(default=1e-6, kw_only=True) + + @property + def signature(self) -> 'Signature': + return Signature.build_from_dtypes(S=QAny(self.target_bitsize)) + + @property + def target_bitsize(self): + """number of bits to represent a k-subset S""" + return self.inst.k * self.inst.index_bitsize + + @property + def selection_registers(self) -> tuple[Register, ...]: + return (Register('S', QAny(self.target_bitsize)),) + + @property + def phasegrad_bitsize(self) -> SymbolicInt: + return ceil(log2(2 * pi(self.eps) / self.eps)) + + @property + def _state_prep_bloq(self) -> SparseStatePreparationViaRotations: + N = 2**self.target_bitsize + + if self.inst.is_symbolic(): + bloq = SparseStatePreparationViaRotations.from_n_coeffs( + N, self.inst.num_unique_constraints, phase_bitsize=self.phasegrad_bitsize + ) + else: + assert not is_symbolic(self.inst.batched_scopes) + + bloq = SparseStatePreparationViaRotations.from_coefficient_map( + N, + {self.inst.scope_as_int(S): B_I for S, B_I in self.inst.batched_scopes}, + self.phasegrad_bitsize, + ) + + bloq = evolve(bloq, target_bitsize=self.target_bitsize) + return bloq + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': + return {self._state_prep_bloq: 1} + + +@bloq_example +def _simple_guiding_state() -> SimpleGuidingState: + from qualtran.bloqs.max_k_xor_sat import Constraint, KXorInstance + + inst = KXorInstance( + n=4, + k=2, + constraints=( + Constraint(S=(0, 1), b=1), + Constraint(S=(2, 3), b=-1), + Constraint(S=(1, 2), b=1), + ), + ) + simple_guiding_state = SimpleGuidingState(inst) + return simple_guiding_state + + +@bloq_example +def _simple_guiding_state_symb() -> SimpleGuidingState: + import sympy + + from qualtran.bloqs.max_k_xor_sat import KXorInstance + + n, m, k = sympy.symbols("n m k", positive=True, integer=True) + inst = KXorInstance.symbolic(n=n, m=m, k=k) + simple_guiding_state_symb = SimpleGuidingState(inst) + return simple_guiding_state_symb + + +_SIMPLE_GUIDING_STATE_DOC = BloqDocSpec( + bloq_cls=SimpleGuidingState, examples=[_simple_guiding_state_symb, _simple_guiding_state] +) + + +@frozen +class ProbabilisticUncompute(Bloq): + """Probabilistically uncompute a register using hadamards, and mark success in a flag qubit + + Apply hadamards to the register, and mark the flag conditioned on all input qubits being 0. + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Eq. 129 and Eq. 130. + """ + + bitsize: SymbolicInt + + @property + def signature(self) -> 'Signature': + return Signature.build_from_dtypes(q=QAny(self.bitsize), flag=QBit()) + + def build_composite_bloq( + self, bb: 'BloqBuilder', q: 'Soquet', flag: 'Soquet' + ) -> dict[str, 'SoquetT']: + q = bb.add(OnEach(self.bitsize, Hadamard()), q=q) + q, flag = bb.add( + XGate().controlled(CtrlSpec(qdtypes=QAny(self.bitsize), cvs=0)), ctrl=q, q=flag + ) + return {'q': q, 'flag': flag} + + +@frozen +class GuidingState(PrepareOracle): + r"""Prepare a guiding state for a kXOR instance with parameter $\ell$. + + Given an kXOR instance $\mathcal{I}$, and parameter $\ell$ (a multiple of $k$), + we want to prepare the unit-length guiding state $|\mathbb{\Psi}\rangle$ (Eq 135): + + $$ + |\mathbb{\Psi}\rangle + \propto + |\Gamma^\ell(\mathcal{A})\rangle + \propto + \sum_{T \in {[n] \choose \ell}} + \sum_{\{S_1, \ldots, S_c\} \in \text{Part}_k(T)} + \left( + \prod_{j = 1}^c B_{\mathcal{I}}(S) + \right) + |T\rangle + $$ + + This bloq prepares the state (Eq 136): + $$ \beta |\mathbb{\Psi}\rangle |0^{\ell \log \ell + 3}\rangle + + |\perp\rangle |1\rangle + $$ + where $\beta \ge \Omega(1 / \ell^{\ell/2})$, + and $\tilde{m}$ is the number of constraints in $\mathcal{I}$. + + This has a gate cost of $O(\ell \tilde{m} \log n)$. + + Args: + inst: the kXOR instance $\mathcal{I}$. + ell: the Kikuchi parameter $\ell$. + amplitude_good_part: (optional) the amplitude $\beta$ of the guiding state $|\Psi\rangle$ + Defaults to $\beta = 0.99 / \ell^{\ell/2}$. + eps: Precision of the prepared state (defaults to 1e-6). + + Registers: + T: $\ell$ indices each in $[n]$. + ancilla (RIGHT): (entangled) $\ell\log\ell+3$ ancilla qubits used for state preparation. + The all zeros state of the ancilla is the good subspace. + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Section 4.4.1 "Preparing the guiding state", Theorem 4.15. Eq 136. + """ + + inst: KXorInstance + ell: SymbolicInt + amplitude_good_part: SymbolicFloat = field(kw_only=True) + eps: SymbolicFloat = field(default=1e-6, kw_only=True) + + @amplitude_good_part.default + def _default_amplitude(self): + return self.coeff_good + + @property + def signature(self) -> 'Signature': + return Signature( + [ + Register('T', QAny(self.target_bitsize)), + Register('ancilla', QAny(self.ancilla_bitsize)), + ] + ) + + @property + def target_bitsize(self) -> SymbolicInt: + return self._index_dtype.num_qubits * self.ell + + @property + def ancilla_bitsize(self) -> SymbolicInt: + r"""total number of entangled ancilla. + + $\ell \log \ell$ for sorting, and 3 flag qubits. + """ + return self.sort_ancilla_bitsize + 3 + + @property + def selection_registers(self) -> tuple[Register, ...]: + return (Register('T', QAny(self.target_bitsize)),) + + @property + def junk_registers(self) -> tuple[Register, ...]: + return (Register('ancilla', QAny(self.ancilla_bitsize)),) + + @property + def sort_ancilla_bitsize(self): + r"""Number of entangled ancilla generated by the sorting algorithm. + + This is a sequence of $\ell$ numbers, each in $[\ell]$, therefore is $\ell \lceil \log \ell \rceil$. + """ + logl = ceil(log2(self.ell)) + return self.ell * logl + + @property + def c(self) -> SymbolicInt: + r"""Value of $c = \ell / k$.""" + c = self.ell // self.inst.k + try: + return int(c) + except TypeError: + pass + return c + + @property + def simple_guiding_state(self) -> SimpleGuidingState: + r"""The simple guiding state $|\phi\rangle$ + + This is the simple guiding state defined in Eq. 142, + which is proportional to $|\Gamma^k\rangle$ (Eq. 134). + We will use $c$ copies of this state to prepare the required guiding state. + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Section 4.4.1 "Preparing the guiding state", Eq. 134. + """ + return SimpleGuidingState(self.inst, eps=self.eps / self.c) + + @property + def _index_dtype(self) -> QUInt: + return QUInt(self.inst.index_bitsize) + + def build_composite_bloq( + self, bb: 'BloqBuilder', T: 'Soquet', ancilla: 'Soquet' + ) -> dict[str, 'SoquetT']: + if is_symbolic(self.c): + raise DecomposeTypeError(f"cannot decompose {self} with symbolic c=l/k={self.c}") + + partition_ancilla = Partition( + self.ancilla_bitsize, + ( + Register('ancilla', QAny(self.sort_ancilla_bitsize)), + Register('flags', QBit(), shape=(3,)), + ), + ) + + ancilla, [flag_duplicates, flag_uncompute, flag] = bb.add(partition_ancilla, x=ancilla) + + # Equation 144: |Phi> = |phi>^{\otimes c} + partition_T_to_S = Partition( + self.target_bitsize, + (Register('S', dtype=QAny(self.simple_guiding_state.target_bitsize), shape=(self.c,)),), + ) + S = bb.add(partition_T_to_S, x=T) + for i in range(self.c): + S[i] = bb.add(self.simple_guiding_state, S=S[i]) + T = bb.add(partition_T_to_S.adjoint(), S=S) + + # sort T using `l log l` entangled clean ancilla + T, ancilla = bb.add_and_partition( + SortInPlace(self.ell, self._index_dtype), + [ + (Register('T', QAny(self.target_bitsize)), ['xs']), + (Register('ancilla', QAny(self.sort_ancilla_bitsize)), ['pi']), + ], + T=T, + ancilla=ancilla, + ) + + # mark if T has duplicates (i.e. not disjoint) (Eq 145) + T, flag_duplicates = bb.add_and_partition( + HasDuplicates(self.ell, self._index_dtype), + [ + (Register('T', QAny(self.target_bitsize)), ['xs']), + (Register('flag', QBit()), ['flag']), + ], + T=T, + flag=flag_duplicates, + ) + + # probabilistically uncompute the sorting ancilla, and mark in a flag bit + # note: flag is 0 for success (like syscall/c exit codes) + ancilla, flag_uncompute = bb.add( + ProbabilisticUncompute(self.sort_ancilla_bitsize), q=ancilla, flag=flag_uncompute + ) + + # compute the overall flag using OR, to obtain Eq 130. + [flag_duplicates, flag_uncompute], flag = bb.add( + MultiControlX(cvs=(0, 0)), controls=[flag_duplicates, flag_uncompute], target=flag + ) + flag = bb.add(XGate(), q=flag) + + # join all the ancilla into a single bag of bits + ancilla = bb.add( + partition_ancilla.adjoint(), + ancilla=ancilla, + flags=[flag_duplicates, flag_uncompute, flag], + ) + + return {'T': T, 'ancilla': ancilla} + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> BloqCountDictT: + return { + self.simple_guiding_state: self.c, + SortInPlace(self.ell, self._index_dtype): 1, + HasDuplicates(self.ell, self._index_dtype): 1, + ProbabilisticUncompute(self.sort_ancilla_bitsize): 1, + MultiControlX(cvs=(0, 0)): 1, + XGate(): 1, + } + + @cached_property + def coeff_good(self): + """lower bound on beta, the coefficient of the good state. + + Sentence below Eq. 147. + """ + return 0.99 / 2 ** (self.sort_ancilla_bitsize / 2) + + +@bloq_example +def _guiding_state() -> GuidingState: + from qualtran.bloqs.max_k_xor_sat import Constraint, KXorInstance + + inst = KXorInstance( + n=4, + k=2, + constraints=( + Constraint(S=(0, 1), b=1), + Constraint(S=(2, 3), b=-1), + Constraint(S=(1, 2), b=1), + ), + ) + guiding_state = GuidingState(inst, ell=4) + return guiding_state + + +@bloq_example +def _guiding_state_symb() -> GuidingState: + import sympy + + from qualtran.bloqs.max_k_xor_sat import KXorInstance + + n, m, k = sympy.symbols("n m k", positive=True, integer=True) + inst = KXorInstance.symbolic(n=n, m=m, k=k) + c = 2 + guiding_state_symb = GuidingState(inst, ell=c * inst.k) + return guiding_state_symb + + +@bloq_example +def _guiding_state_symb_c() -> GuidingState: + import sympy + + from qualtran.bloqs.max_k_xor_sat import KXorInstance + + n, m, c = sympy.symbols("n m c", positive=True, integer=True) + k = sympy.symbols("k", positive=True, integer=True, even=True) + inst = KXorInstance.symbolic(n=n, m=m, k=k) + guiding_state_symb_c = GuidingState(inst, ell=c * k) + return guiding_state_symb_c + + +_GUIDING_STATE_DOC = BloqDocSpec( + bloq_cls=GuidingState, examples=[_guiding_state_symb_c, _guiding_state_symb, _guiding_state] +) diff --git a/qualtran/bloqs/max_k_xor_sat/guiding_state_test.py b/qualtran/bloqs/max_k_xor_sat/guiding_state_test.py new file mode 100644 index 000000000..f95c33ca9 --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/guiding_state_test.py @@ -0,0 +1,132 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 unittest.mock import ANY + +import pytest +import sympy + +import qualtran.testing as qlt_testing +from qualtran.bloqs.max_k_xor_sat.guiding_state import ( + _guiding_state, + _guiding_state_symb, + _guiding_state_symb_c, + _simple_guiding_state, + _simple_guiding_state_symb, +) +from qualtran.resource_counting import big_O, GateCounts, get_cost_value, QECGatesCost +from qualtran.symbolics import bit_length, ceil, log2 + + +@pytest.mark.parametrize( + "bloq_ex", + [ + _simple_guiding_state, + _simple_guiding_state_symb, + _guiding_state, + _guiding_state_symb, + _guiding_state_symb_c, + ], + ids=lambda b: b.name, +) +def test_examples(bloq_autotester, bloq_ex): + if bloq_autotester.check_name == 'serialize': + pytest.skip() + + bloq_autotester(bloq_ex) + + +def test_t_cost_simple(): + bloq = _simple_guiding_state() + gc = get_cost_value(bloq, QECGatesCost()) + B_GRAD = bloq.phasegrad_bitsize + + assert gc == GateCounts(and_bloq=24, toffoli=3 * (B_GRAD - 2), clifford=ANY, measurement=ANY) + + +def test_t_cost_simple_symb(): + bloq = _simple_guiding_state_symb() + gc = get_cost_value(bloq, QECGatesCost()) + B_GRAD = bloq.phasegrad_bitsize + + n, m, k = bloq.inst.n, bloq.inst.m, bloq.inst.k + klogn = k * ceil(log2(n)) + # https://github.com/quantumlib/Qualtran/issues/1341 + klogn_roundtrip = bit_length(2**klogn - 1) + + assert gc == GateCounts( + # O(k m log n) + and_bloq=4 * m + (2 * m + 1) * (klogn_roundtrip - 1) - 4, + toffoli=2 * (B_GRAD - 2), + clifford=ANY, + measurement=ANY, + ) + + +def test_t_cost(): + bloq = _guiding_state() + gc = get_cost_value(bloq, QECGatesCost()) + B_GRAD = bloq.simple_guiding_state.phasegrad_bitsize + + assert gc == GateCounts( + and_bloq=352, toffoli=6 * (B_GRAD - 2), cswap=192, clifford=ANY, measurement=ANY + ) + + +@pytest.mark.parametrize("bloq_ex", [_guiding_state_symb, _guiding_state_symb_c]) +def test_t_cost_symb_c(bloq_ex): + bloq = bloq_ex() + gc = get_cost_value(bloq, QECGatesCost()) + B_GRAD = bloq.simple_guiding_state.phasegrad_bitsize + + n, m, k = bloq.inst.n, bloq.inst.m, bloq.inst.k + l, c = bloq.ell, bloq.c + + logn = ceil(log2(n)) + logl = ceil(log2(l)) + + klogn = k * logn + # https://github.com/quantumlib/Qualtran/issues/1341 + klogn_roundtrip = bit_length(2**klogn - 1) + + assert gc == GateCounts( + and_bloq=( + 6 * l**2 * (2 * logn + 1) + + l * logl + + l + + c * (4 * m + (2 * m + 1) * (klogn_roundtrip - 1) - 4) + + (l - 1) * logn + - 2 + ), + toffoli=c * (2 * (B_GRAD - 2)), + cswap=6 * l**2 * logn, + clifford=ANY, + measurement=ANY, + ) + + # verify big_O + t_cost = gc.total_t_count() + t_cost = sympy.sympify(t_cost) + t_cost = t_cost.subs(klogn_roundtrip, klogn) + t_cost = t_cost.simplify() + assert t_cost in big_O(l * m * logn + l**2 * logn + B_GRAD * c) + + +@pytest.mark.notebook +def test_notebook(): + qlt_testing.execute_notebook('guiding_state') + + +@pytest.mark.notebook +def test_tutorial(): + qlt_testing.execute_notebook('guiding_state_tutorial') diff --git a/qualtran/bloqs/max_k_xor_sat/guiding_state_tutorial.ipynb b/qualtran/bloqs/max_k_xor_sat/guiding_state_tutorial.ipynb new file mode 100644 index 000000000..82681faf6 --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/guiding_state_tutorial.ipynb @@ -0,0 +1,212 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import sympy\n", + "from qualtran.drawing import show_bloq, show_call_graph" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "Let us start with a kXOR instance with $n$ variables and $m$ constraints." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.bloqs.max_k_xor_sat.kxor_instance import KXorInstance\n", + "\n", + "n, m, k = sympy.symbols(\"n m k\", positive=True, integer=True)\n", + "inst = KXorInstance.symbolic(n, m, k)\n", + "inst" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "We first prepare the guiding state to use in the guided sparse hamiltonian algorithm.\n", + "The guiding state is defined by the instance, and a parameter $\\ell$ (a multiple of $k$)\n", + "\n", + "From Theorem 4.15 of the paper, this should be a circuit of $O(\\ell m \\log n)$ gates,\n", + "and prepare the state $\\beta |\\Psi\\rangle|0^{\\ell \\log \\ell}\\rangle + |\\perp\\rangle|1\\rangle$,\n", + "where $\\beta \\ge 0.99 / \\ell^{\\ell/2}$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.bloqs.max_k_xor_sat.guiding_state import GuidingState\n", + "\n", + "c = sympy.symbols(\"c\", positive=True, integer=True)\n", + "l = c * k\n", + "guiding_state = GuidingState(inst, l)\n", + "show_call_graph(guiding_state.call_graph(max_depth=1)[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "guiding_state_3 = GuidingState(inst, 3 * k)\n", + "show_bloq(guiding_state_3.decompose_bloq())" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "We can also build the guiding state for a concrete (non symbolic) instance:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "inst = KXorInstance.random_instance(n=20, m=100, k=4, planted_advantage=0.8, rng=np.random.default_rng(100))\n", + "guiding_state_concrete = GuidingState(inst, ell=12)\n", + "show_bloq(guiding_state_concrete)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "show_bloq(guiding_state_concrete.decompose_bloq())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "show_bloq(guiding_state_concrete.decompose_bloq().flatten_once())" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "Let us evaluate the gate cost for the above bloqs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.resource_counting import get_cost_value, QECGatesCost\n", + "\n", + "get_cost_value(guiding_state_concrete, QECGatesCost())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "gc = get_cost_value(guiding_state, QECGatesCost())\n", + "t_cost = gc.total_t_count(ts_per_toffoli=4, ts_per_cswap=4, ts_per_and_bloq=4)\n", + "t_cost" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.symbolics import ceil, log2, floor\n", + "from qualtran.resource_counting import big_O\n", + "\n", + "# simplify some expressions that sympy could not\n", + "klogn = k * ceil(log2(n))\n", + "klogn_long = ceil(log2(floor(2**klogn)))\n", + "t_cost = t_cost.subs(klogn_long, klogn)\n", + "t_cost = t_cost.simplify()\n", + "\n", + "# replace l with a symbol\n", + "l_symb = sympy.symbols(r\"\\ell\", positive=True, integer=True)\n", + "t_cost = t_cost.subs(c * k, l_symb)\n", + "\n", + "big_O(t_cost) # matches paper Theorem 4.15 (as c, l are O(m))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "show_call_graph(guiding_state_concrete, max_depth=3)" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "As we know that $c = \\ell/k \\le \\ell$ and $\\ell \\le m$, the above expression matches the paper result of $O(\\ell m \\log_2(n))$ 1/2-qubit gates.\n", + "" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/qualtran/bloqs/max_k_xor_sat/kxor_instance.py b/qualtran/bloqs/max_k_xor_sat/kxor_instance.py new file mode 100644 index 000000000..913ae06ae --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/kxor_instance.py @@ -0,0 +1,273 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 itertools +from collections import defaultdict +from functools import cached_property +from typing import cast, Sequence, TypeAlias, Union + +import numpy as np +import sympy +from attrs import evolve, field, frozen +from numpy.typing import NDArray + +from qualtran.symbolics import bit_length, ceil, HasLength, is_symbolic, log2, slen, SymbolicInt + +Scope: TypeAlias = Union[tuple[int, ...], HasLength] +"""A subset of variables""" + + +def _sort_scope(S: Scope) -> Scope: + if is_symbolic(S): + return S + return tuple(sorted(S)) + + +@frozen +class Constraint: + """A single kXOR constraint. + + Definition 2.1. + + Note: n, k are not stored here, but only in the instance. + + Attributes: + S: the scope - subset of `[n]` of size k. + b: +1 or -1. + """ + + S: Scope = field(converter=_sort_scope) + b: SymbolicInt = field() + + @classmethod + def random(cls, n: int, k: int, *, rng: np.random.Generator): + """Single random constraint, Notation 2.3.""" + S = tuple(rng.choice(n, k, replace=False)) + b = rng.choice([-1, +1]) + return cls(S, b) + + @classmethod + def random_planted(cls, n: int, k: int, *, rho: float, z: NDArray, rng: np.random.Generator): + """Single planted constraint, Notation 2.4.""" + S = tuple(rng.choice(n, k, replace=False)) + eta = (-1) ** (rng.random() < (1 + rho) / 2) # i.e. expectation rho. + unplanted = cls(S, 1) # supporting constraint to evaluate z^S + b = eta * unplanted.evaluate(z) + return cls(S, b) + + @classmethod + def symbolic(cls, n: SymbolicInt, ix: int): + return cls(HasLength(n), sympy.Symbol(f"b_{ix}")) + + def is_symbolic(self): + return is_symbolic(self.S, self.b) + + def evaluate(self, x: NDArray[np.integer]): + return np.prod(x[np.array(self.S)]) + + +@frozen +class KXorInstance: + r"""A kXOR instance $\mathcal{I}$. + + Definition 2.1: A kXOR instance $\mathcal{I}$ over variables indexed by $[n]$ + consists of a multiset of constraints $\mathcal{C} = (S, b)$, where each scope + $S \subseteq [n]$ has cardinality $k$, and each right-hand side $b \in \{\pm 1\}$. + + Attributes: + n: number of variables. + k: number of variables per clause. + constraints: a tuple of `m` Constraints. + max_rhs: maximum value of the RHS polynomial $B_\mathcal{I}(S)$. + see default constructor for default value. In case the instance is symbolic, + the user can specify an expression for this, to avoid the default value. + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Definition 2.1. + """ + + n: SymbolicInt + k: SymbolicInt + constraints: Union[tuple[Constraint, ...], HasLength] + max_rhs: SymbolicInt = field() + + @max_rhs.default + def _default_max_rhs(self): + """With very high probability, the max entry will be quite small. + + This is a classical preprocesing step. Time $m$. + """ + if is_symbolic(self.constraints) or is_symbolic(*self.constraints): + # user did not provide a value, assume some small constant + return 2 + + # instance is not symbolic, so we can compute the exact value. + assert isinstance(self.batched_scopes, tuple) + return max(abs(b) for _, b in self.batched_scopes) + + @cached_property + def m(self): + return slen(self.constraints) + + @classmethod + def random_instance( + cls, n: int, m: int, k: int, *, planted_advantage: float = 0, rng: np.random.Generator + ): + r"""Generate a random kXOR instance with the given planted advantage. + + `planted_advantage=0` generates random instances, and `1` generates a + linear system with a solution. + + Args: + n: number of variables + m: number of clauses + k: number of terms per clause + planted_advantage: $\rho$ + rng: random generator + + References: + [Quartic quantum speedups for planted inference](https://arxiv.org/abs/2406.19378v1) + Notation 2.4. + """ + # planted vector + z = rng.choice([-1, +1], size=n) + + # constraints + constraints = tuple( + Constraint.random_planted(n=n, k=k, rho=planted_advantage, z=z, rng=rng) + for _ in range(m) + ) + + return cls(n=n, k=k, constraints=constraints) + + @classmethod + def symbolic(cls, n: SymbolicInt, m: SymbolicInt, k: SymbolicInt, *, max_rhs: SymbolicInt = 2): + """Create a symbolic instance with n variables, m constraints.""" + constraints = HasLength(m) + return cls(n=n, k=k, constraints=constraints, max_rhs=max_rhs) + + def is_symbolic(self): + if is_symbolic(self.n, self.m, self.k, self.constraints): + return True + assert isinstance(self.constraints, tuple) + return is_symbolic(*self.constraints) + + def subset(self, indices: Union[Sequence[int], HasLength]) -> 'KXorInstance': + """Pick a subset of clauses defined by the set of indices provided.""" + if self.is_symbolic() or is_symbolic(indices): + return evolve(self, constraints=HasLength(slen(indices))) + assert isinstance(self.constraints, tuple) + + constraints = tuple(self.constraints[i] for i in indices) + return evolve(self, constraints=constraints) + + @cached_property + def index_bitsize(self): + """number of bits required to represent the index of a variable, i.e. `[n]` + + We assume zero-indexing. + """ + return ceil(log2(self.n)) + + @cached_property + def num_unique_constraints(self) -> SymbolicInt: + return slen(self.batched_scopes) + + @cached_property + def batched_scopes(self) -> Union[tuple[tuple[Scope, int], ...], HasLength]: + r"""Group all the constraints by Scope, and add up the $b$ values. + + This is a classical preprocessing step. Time $k m \log m$. + """ + if self.is_symbolic(): + return HasLength(self.m) + + assert isinstance(self.constraints, tuple) + + batches: dict[Scope, int] = defaultdict(lambda: 0) + for con in self.constraints: + assert isinstance(con.S, tuple) + batches[con.S] += con.b + + batches_sorted = sorted(batches.items(), key=lambda c: c[1]) + return tuple(batches_sorted) + + @cached_property + def rhs_sum_bitsize(self): + r"""number of bits to represent the RHS polynomial $B_{\mathcal{I}}(S)$.""" + return bit_length(2 * self.max_rhs) + + def scope_as_int(self, S: Scope) -> int: + r"""Convert a scope into a single integer. + + Given a scope `S = (x_1, x_2, ..., x_k)`, and a bitsize `r` for each index, + the integer representation is given by concatenating `r`-bit unsigned repr + of each `x_i`. That is, $\sum_i r^{k - i} x_i$. + + This uses Big-endian representation, like all qualtran dtypes. + + The bitsize `r` is picked as `ceil(log(n))` for an n-variable instance. + """ + assert not is_symbolic(S) + + bitsize = self.index_bitsize + + result = 0 + for x in S: + result = (result << bitsize) + x + return result + + def brute_force_sparsity(self, ell: int) -> int: + r"""Compute the sparsity of the Kikuchi matrix with parameter $\ell$ by brute force. + + Takes time `O(C(n, l) * m * l)`. Extremely slow, use with caution. + """ + assert isinstance(self.n, int) + s = 0 + for S in itertools.combinations(range(self.n), ell): + nz = 0 + for U, _ in cast(tuple, self.batched_scopes): + T = set(S).symmetric_difference(U) + if len(T) == ell: + nz += 1 + s = max(s, nz) + return s + + +def example_kxor_instance() -> KXorInstance: + n, k = 10, 4 + cs = ( + Constraint((0, 1, 2, 3), -1), + Constraint((0, 2, 4, 5), 1), + Constraint((0, 3, 4, 5), 1), + Constraint((0, 3, 4, 5), 1), + Constraint((1, 2, 3, 4), -1), + Constraint((1, 3, 4, 5), -1), + Constraint((1, 3, 4, 5), -1), + Constraint((2, 3, 4, 5), 1), + ) + inst = KXorInstance(n, k, cs) + return inst diff --git a/qualtran/bloqs/max_k_xor_sat/kxor_instance_test.py b/qualtran/bloqs/max_k_xor_sat/kxor_instance_test.py new file mode 100644 index 000000000..700dfd6eb --- /dev/null +++ b/qualtran/bloqs/max_k_xor_sat/kxor_instance_test.py @@ -0,0 +1,30 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 numpy as np +import pytest + +from .kxor_instance import KXorInstance + + +@pytest.mark.slow +@pytest.mark.parametrize("rho", [0, 0.8, 0.9]) +def test_max_rhs(rho: float): + rng = np.random.default_rng(402) + + rhs = [] + for i in range(100): + inst = KXorInstance.random_instance(n=100, m=1000, k=4, planted_advantage=rho, rng=rng) + rhs.append(inst.max_rhs) + + assert max(rhs) == 2