From a6fc9f3d341c398e4c2d775f68aaaf6fb3365bec Mon Sep 17 00:00:00 2001
From: Tanuj Khattar <tanujkhattar@google.com>
Date: Mon, 9 Sep 2024 10:46:52 -0700
Subject: [PATCH] Make symbolic call graph and costs of
 `ZPowUsingProgrammedAncilla` more accurate (#1389)

* Make symbolic call graph and costs of ZPowUsingProgrammedAncilla more accurate

* Fix pylint
---
 .../programmable_ancilla_rotation.py          |  2 +-
 .../programmable_ancilla_rotation_test.py     | 20 +++++++++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/qualtran/bloqs/rotations/programmable_ancilla_rotation.py b/qualtran/bloqs/rotations/programmable_ancilla_rotation.py
index e65779cb1..07b6097e9 100644
--- a/qualtran/bloqs/rotations/programmable_ancilla_rotation.py
+++ b/qualtran/bloqs/rotations/programmable_ancilla_rotation.py
@@ -138,7 +138,7 @@ def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT:
 
         if is_symbolic(self.n_rounds):
             phi = ssa.new_symbol(r"\phi")
-            eps = ssa.new_symbol(r"\epsilon")
+            eps = self.eps / n_rz
             resources[ZPowProgrammedAncilla(phi, eps)] += self.n_rounds
         else:
             for i in range(int(self.n_rounds)):
diff --git a/qualtran/bloqs/rotations/programmable_ancilla_rotation_test.py b/qualtran/bloqs/rotations/programmable_ancilla_rotation_test.py
index 2bb4b9bf6..1315f0b1e 100644
--- a/qualtran/bloqs/rotations/programmable_ancilla_rotation_test.py
+++ b/qualtran/bloqs/rotations/programmable_ancilla_rotation_test.py
@@ -26,6 +26,8 @@
     ZPowProgrammedAncilla,
     ZPowUsingProgrammedAncilla,
 )
+from qualtran.resource_counting.t_counts_from_sigma import t_counts_from_sigma
+from qualtran.symbolics import ceil, log2
 
 
 def test_rz_resource_state_examples(bloq_autotester):
@@ -52,6 +54,20 @@ def test_rz_via_par_examples(bloq_autotester):
     bloq_autotester(_zpow_using_programmed_ancilla_symb_rounds)
 
 
+def test_t_cost_including_rotations():
+    theta, eps, p = sympy.symbols(r'theta, \epsilon, p')
+    bloq = ZPowUsingProgrammedAncilla.from_failure_probability(
+        exponent=theta, max_fail_probability=p, eps=eps
+    )
+    _, sigma = bloq.call_graph()
+    n_rot = ceil(log2(1 / p))
+    eps_per_rot = eps / n_rot
+    # TODO(#1250): Ideally, we should be able to get this symbolic cost via more standard
+    # ways like `bloq.t_complexity()` or `get_cost_value(bloq, QECGatesCost())` but both
+    # of them ignore epsilon right now.
+    assert t_counts_from_sigma(sigma) == ceil(1.149 * log2(1.0 / eps_per_rot) + 9.2) * n_rot
+
+
 def test_rz_via_par_call_graphs():
     _, sigma_rz = ZPowUsingProgrammedAncilla(np.pi / 4).call_graph(max_depth=1)
     assert sigma_rz == {
@@ -74,13 +90,13 @@ def test_rz_via_par_call_graphs():
         Measure(): 3,
     }
 
-    phi0, eps0, n = sympy.symbols(r"_\phi0 _\epsilon0 n")
+    phi0, n = sympy.symbols(r"_\phi0 n")
     _, sigma_rz_symb_rounds = ZPowUsingProgrammedAncilla(phi, eps=eps, n_rounds=n).call_graph(
         max_depth=1
     )
     assert sigma_rz_symb_rounds == {
         CNOT(): n,
         XGate(): n,
-        ZPowProgrammedAncilla(phi0, eps0): n,
+        ZPowProgrammedAncilla(phi0, eps / n): n,
         Measure(): n,
     }