From 8fc12b22812dd7f1e7546448b8d4b68790c7d1a4 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <7265746+anurudhp@users.noreply.github.com>
Date: Thu, 29 Aug 2024 17:25:55 -0700
Subject: [PATCH] `is_zero`: catch errors thrown by `Expr.equals(0)`, and
 assume it is non-zero. (#1366)

* is_zero: catch errors thrown by `Expr.equals(0)`, and assume it is non-zero.

* check for complexity match more carefully

* lint
---
 .../hamiltonian_simulation_by_gqsp_test.py      | 17 +++++++++++++----
 .../polynomial/jacobi_anger_approximations.py   |  7 +++++--
 qualtran/symbolics/math_funcs.py                |  7 ++++++-
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py b/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py
index 2a706d49d..5d0e68be1 100644
--- a/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py
+++ b/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py
@@ -34,8 +34,7 @@
 )
 from qualtran.bloqs.qubitization.qubitization_walk_operator import QubitizationWalkOperator
 from qualtran.cirq_interop import BloqAsCirqGate
-from qualtran.cirq_interop.t_complexity_protocol import TComplexity
-from qualtran.resource_counting import big_O, BloqCount, get_cost_value
+from qualtran.resource_counting import big_O, BloqCount, get_cost_value, QECGatesCost, QubitCount
 from qualtran.symbolics import Shaped
 
 
@@ -47,6 +46,11 @@ def test_symbolic_examples(bloq_autotester):
     bloq_autotester(_symbolic_hamsim_by_gqsp)
 
 
+def test_qubit_count_symbolic():
+    bloq = _symbolic_hamsim_by_gqsp()
+    _ = get_cost_value(bloq, QubitCount())
+
+
 @pytest.mark.parametrize("bitsize", [1, 2])
 @pytest.mark.parametrize("t", [2, 3, 5, pytest.param(10, marks=pytest.mark.slow)])
 @pytest.mark.parametrize("precision", [1e-5, 1e-7])
@@ -108,7 +112,12 @@ def test_hamiltonian_simulation_by_gqsp_t_complexity():
     counts = get_cost_value(hubbard_time_evolution_by_gqsp, BloqCount.for_gateset('t+tof+cswap'))
     assert t_comp.t == counts[TwoBitCSwap()] * 7 + counts[TGate()]
 
+
+def test_symbolic_t_cost():
     symbolic_hamsim_by_gqsp = _symbolic_hamsim_by_gqsp()
+    gc = get_cost_value(symbolic_hamsim_by_gqsp, QECGatesCost())
+    t_cost = gc.total_t_count()
+
     tau, t, inv_eps = sympy.symbols(r"\tau t \epsilon^{-1}", positive=True)
-    T = big_O(tau * t + sympy.log(2 * inv_eps) / sympy.log(sympy.log(2 * inv_eps)))
-    assert symbolic_hamsim_by_gqsp.t_complexity() == TComplexity(t=T, clifford=T, rotations=T)  # type: ignore[arg-type]
+    O_t_cost = sympy.O(t_cost, (tau, sympy.oo), (t, sympy.oo), (inv_eps, sympy.oo))
+    assert O_t_cost == big_O(tau * t + sympy.log(2 * inv_eps) / sympy.log(sympy.log(2 * inv_eps)))
diff --git a/qualtran/linalg/polynomial/jacobi_anger_approximations.py b/qualtran/linalg/polynomial/jacobi_anger_approximations.py
index 5dd339d1d..07a84e70c 100644
--- a/qualtran/linalg/polynomial/jacobi_anger_approximations.py
+++ b/qualtran/linalg/polynomial/jacobi_anger_approximations.py
@@ -18,7 +18,6 @@
 import sympy
 from numpy.typing import NDArray
 
-from qualtran.resource_counting import big_O
 from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt
 
 
@@ -40,6 +39,8 @@ def degree_jacobi_anger_approximation(t: SymbolicFloat, *, precision: SymbolicFl
         d = \mathcal{O}(t + \frac{\log(1/\epsilon)}{\log\log(1/\epsilon)})
     $$
 
+    It returns `d` up to a symbolic constant. To ignore the constant, please use `big_O` on the final expression.
+
     Args:
         t: scale of the exponent in the function to approximate.
         precision: $\epsilon$ in the above polynomial approximation
@@ -48,7 +49,9 @@ def degree_jacobi_anger_approximation(t: SymbolicFloat, *, precision: SymbolicFl
         Truncation degree $d$ as defined above.
     """
     if is_symbolic(t, precision):
-        return big_O(t + sympy.log(1 / precision) / sympy.log(sympy.log(1 / precision)))
+        # use a symbol for the constant.
+        c_ja = sympy.Symbol("C_{JA}", positive=True)
+        return c_ja * (t + sympy.log(1 / precision) / sympy.log(sympy.log(1 / precision)))
 
     def term_too_small(n: int) -> bool:
         return bool(np.isclose(scipy.special.jv(n, t), 0, atol=float(precision)))
diff --git a/qualtran/symbolics/math_funcs.py b/qualtran/symbolics/math_funcs.py
index d588386d2..7ac2e496d 100644
--- a/qualtran/symbolics/math_funcs.py
+++ b/qualtran/symbolics/math_funcs.py
@@ -323,5 +323,10 @@ def is_zero(x: SymbolicInt) -> bool:
     or could not be symbolically symplified to a zero.
     """
     if is_symbolic(x):
-        return x.equals(0)
+        try:
+            zero = x.equals(0)
+        except ValueError:
+            zero = False
+        return zero
+
     return x == 0