Fix t_complexity for symbolic ham. sim. by gqsp example (#944)

* investigate `t_complexity` failure for symbolic HamSimbyGQSP

* `UtilBloq` -> `_BookkeepingBloq`, improve docstring

---------

Co-authored-by: Matthew Harrigan <mpharrigan@google.com>

Author: anurudhp

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import cast, Dict, Tuple, TYPE_CHECKING, Union
+from typing import cast, Dict, Set, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import field, frozen
@@ -29,6 +29,7 @@
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder, SoquetT
+    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
 
 
 @frozen
@@ -176,6 +177,21 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
 
         return soqs
 
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        if self.is_symbolic():
+            from qualtran.bloqs.basic_gates.su2_rotation import SU2RotationGate
+
+            d = self.degree
+            return {
+                (self.walk_operator.prepare, 1),
+                (self.walk_operator.prepare.adjoint(), 1),
+                (self.walk_operator.controlled(control_values=[0]), d),
+                (self.walk_operator.adjoint().controlled(), d),
+                (SU2RotationGate.arbitrary(ssa), 2 * d + 1),
+            }
+
+        return super().build_call_graph(ssa)
+
 
 @bloq_example
 def _hubbard_time_evolution_by_gqsp() -> HamiltonianSimulationByGQSP:
@@ -192,9 +208,8 @@ def _symbolic_hamsim_by_gqsp() -> HamiltonianSimulationByGQSP:
 
     from qualtran.bloqs.hubbard_model import get_walk_operator_for_hubbard_model
 
-    walk_op = get_walk_operator_for_hubbard_model(2, 2, 1, 1)
-
-    t, inv_eps = sympy.symbols("t N")
+    tau, t, inv_eps = sympy.symbols(r"\tau t \epsilon^{-1}", positive=True)
+    walk_op = get_walk_operator_for_hubbard_model(2, 2, tau, 4 * tau)
     symbolic_hamsim_by_gqsp = HamiltonianSimulationByGQSP(walk_op, t=t, precision=1 / inv_eps)
     return symbolic_hamsim_by_gqsp
 

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pytest
 import scipy
+import sympy
 from numpy.typing import NDArray
 
 from qualtran.bloqs.for_testing.matrix_gate import MatrixGate
@@ -26,6 +27,8 @@
     verify_generalized_qsp,
 )
 from qualtran.bloqs.qubitization_walk_operator import QubitizationWalkOperator
+from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+from qualtran.resource_counting import big_O
 from qualtran.symbolics import Shaped
 
 from .hamiltonian_simulation_by_gqsp import (
@@ -100,3 +103,8 @@ def test_hamiltonian_simulation_by_gqsp(
 def test_hamiltonian_simulation_by_gqsp_t_complexity():
     hubbard_time_evolution_by_gqsp = _hubbard_time_evolution_by_gqsp.make()
     _ = hubbard_time_evolution_by_gqsp.t_complexity()
+
+    symbolic_hamsim_by_gqsp = _symbolic_hamsim_by_gqsp()
+    tau, t, inv_eps = sympy.symbols(r"\tau t \epsilon^{-1}", positive=True)
+    T = big_O(tau * t + sympy.log(inv_eps) / sympy.log(sympy.log(inv_eps)))
+    assert symbolic_hamsim_by_gqsp.t_complexity() == TComplexity(t=T, clifford=T, rotations=T)  # type: ignore[arg-type]

qualtran/bloqs/hubbard_model.py

@@ -66,6 +66,7 @@
 from qualtran.bloqs.state_preparation.prepare_uniform_superposition import (
     PrepareUniformSuperposition,
 )
+from qualtran.symbolics.math_funcs import acos, ssqrt
 
 if TYPE_CHECKING:
     from qualtran.symbolics import SymbolicFloat
@@ -308,7 +309,7 @@ def decompose_from_registers(
         temp = quregs['temp']
 
         N = self.x_dim * self.y_dim * 2
-        yield cirq.Ry(rads=2 * np.arccos(np.sqrt(self.t * N / self.l1_norm_of_coeffs))).on(*V)
+        yield cirq.Ry(rads=2 * acos(ssqrt(self.t * N / self.l1_norm_of_coeffs))).on(*V)
         yield cirq.Ry(rads=2 * np.arccos(np.sqrt(1 / 5))).on(*U).controlled_by(*V)
         yield PrepareUniformSuperposition(self.x_dim).on_registers(controls=[], target=p_x)
         yield PrepareUniformSuperposition(self.y_dim).on_registers(controls=[], target=p_y)

qualtran/bloqs/util_bloqs.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 """Bloqs for virtual operations and register reshaping."""
-
+import abc
 from functools import cached_property
 from typing import Any, Dict, Iterable, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
@@ -50,8 +50,32 @@
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
+class _BookkeepingBloq(Bloq, metaclass=abc.ABCMeta):
+    """Base class for utility bloqs used for bookkeeping.
+
+    This bloq:
+    - has trivial controlled versions, which pass through the control register.
+    - does not affect T complexity.
+    """
+
+    def get_ctrl_system(
+        self, ctrl_spec: Optional['CtrlSpec'] = None
+    ) -> Tuple['Bloq', 'AddControlledT']:
+        def add_controlled(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            # ignore `ctrl_soq` and pass it through for bookkeeping operation.
+            out_soqs = bb.add_t(self, **in_soqs)
+            return ctrl_soqs, out_soqs
+
+        return self, add_controlled
+
+    def _t_complexity_(self) -> 'TComplexity':
+        return TComplexity()
+
+
 @frozen
-class Split(Bloq):
+class Split(_BookkeepingBloq):
     """Split a bitsize `n` register into a length-`n` array-register.
 
     Attributes:
@@ -75,9 +99,6 @@ def adjoint(self) -> 'Bloq':
     def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]:
         return None, {'reg': reg.reshape((self.dtype.num_qubits, 1))}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
     def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
         return {'reg': ints_to_bits(np.array([reg]), self.dtype.num_qubits)[0]}
 
@@ -103,16 +124,6 @@ def add_my_tensors(
             )
         )
 
-    def get_ctrl_system(self, ctrl_spec=None) -> Tuple['Bloq', 'AddControlledT']:
-        def add_controlled(
-            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
-        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
-            # ignore `ctrl_soq` and pass it through for bookkeeping operation.
-            out_soqs = bb.add_t(self, **in_soqs)
-            return ctrl_soqs, out_soqs
-
-        return self, add_controlled
-
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
             return Text(self.pretty_name())
@@ -123,7 +134,7 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
 
 
 @frozen
-class Join(Bloq):
+class Join(_BookkeepingBloq):
     """Join a length-`n` array-register into one register of bitsize `n`.
 
     Attributes:
@@ -147,9 +158,6 @@ def adjoint(self) -> 'Bloq':
     def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]:
         return None, {'reg': reg.reshape(self.dtype.num_qubits)}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
     def add_my_tensors(
         self,
         tn: 'qtn.TensorNetwork',
@@ -175,18 +183,6 @@ def add_my_tensors(
     def on_classical_vals(self, reg: 'NDArray[np.uint]') -> Dict[str, int]:
         return {'reg': bits_to_ints(reg)[0]}
 
-    def get_ctrl_system(
-        self, ctrl_spec: Optional['CtrlSpec'] = None
-    ) -> Tuple['Bloq', 'AddControlledT']:
-        def add_controlled(
-            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
-        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
-            # ignore `ctrl_soq` and pass it through for bookkeeping operation.
-            out_soqs = bb.add_t(self, **in_soqs)
-            return ctrl_soqs, out_soqs
-
-        return self, add_controlled
-
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
             return Text('')
@@ -197,7 +193,7 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
 
 
 @frozen
-class Partition(Bloq):
+class Partition(_BookkeepingBloq):
     """Partition a generic index into multiple registers.
 
     Args:
@@ -240,9 +236,6 @@ def as_cirq_op(self, qubit_manager, **cirq_quregs) -> Tuple[None, Dict[str, 'Cir
         else:
             return None, {'x': np.concatenate([v.ravel() for _, v in cirq_quregs.items()])}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
     def add_my_tensors(
         self,
         tn: 'qtn.TensorNetwork',
@@ -323,7 +316,7 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
 
 
 @frozen
-class Allocate(Bloq):
+class Allocate(_BookkeepingBloq):
     """Allocate an `n` bit register.
 
     Attributes:
@@ -342,9 +335,6 @@ def adjoint(self) -> 'Bloq':
     def on_classical_vals(self) -> Dict[str, int]:
         return {'reg': 0}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
     def add_my_tensors(
         self,
         tn: 'qtn.TensorNetwork',
@@ -367,7 +357,7 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
 
 
 @frozen
-class Free(Bloq):
+class Free(_BookkeepingBloq):
     """Free (i.e. de-allocate) an `n` bit register.
 
     The tensor decomposition assumes the `n` bit register is uncomputed and is in the $|0^{n}>$
@@ -392,9 +382,6 @@ def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
             raise ValueError(f"Tried to free a non-zero register: {reg}.")
         return {}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
     def add_my_tensors(
         self,
         tn: 'qtn.TensorNetwork',
@@ -440,13 +427,14 @@ def _t_complexity_(self) -> 'TComplexity':
 
 
 @frozen
-class Cast(Bloq):
+class Cast(_BookkeepingBloq):
     """Cast a register from one n-bit QDType to another QDType.
 
 
     Args:
-        in_qdtype: Input QDType to cast from.
-        out_qdtype: Output QDType to cast to.
+        inp_dtype: Input QDType to cast from.
+        out_dtype: Output QDType to cast to.
+        shape: shape of the register to cast.
 
     Registers:
         in: input register to cast from.
@@ -501,9 +489,6 @@ def on_classical_vals(self, reg: int) -> Dict[str, 'ClassicalValT']:
     def as_cirq_op(self, qubit_manager, reg: 'CirqQuregT') -> Tuple[None, Dict[str, 'CirqQuregT']]:
         return None, {'reg': reg}
 
-    def _t_complexity_(self) -> 'TComplexity':
-        return TComplexity()
-
 
 @frozen
 class Power(GateWithRegisters):

qualtran/symbolics/math_funcs.py

@@ -41,6 +41,12 @@ def sabs(x: SymbolicFloat) -> SymbolicFloat:
     return cast(SymbolicFloat, abs(x))
 
 
+def ssqrt(x: SymbolicFloat) -> SymbolicFloat:
+    if isinstance(x, sympy.Basic):
+        return sympy.sqrt(x)
+    return np.sqrt(x)
+
+
 def ceil(x: SymbolicFloat) -> SymbolicInt:
     if not isinstance(x, sympy.Basic):
         return int(np.ceil(x))