Cleanup redundant code for computing t complexity (#1317)

* move _HasEps

* reduce dependence on `_get_all_rotation_types`

* remove `_get_all_rotation_types`

* mypy

* reorganize tests for `t_counts_from_sigma`

* mypy

* fix error

* replace `t_counts_from_sigma(bloq)` with `get_cost_value(bloq, QECGatesCost())`

* revert test

* link issue

---------

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

Author: anurudhp

qualtran/bloqs/basic_gates/rotation.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 Optional, Protocol, runtime_checkable, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import attrs
 import cirq
@@ -27,13 +27,6 @@
 from qualtran.symbolics import SymbolicFloat
 
 
-@runtime_checkable
-class _HasEps(Protocol):
-    """Protocol for typing `RotationBloq` base class mixin that has accuracy specified as eps."""
-
-    eps: float
-
-
 @frozen
 class ZPowGate(CirqGateAsBloqBase):
     r"""A gate that rotates around the Z axis of the Bloch sphere.
@@ -115,7 +108,7 @@ def _z_pow() -> ZPowGate:
 class CZPowGate(CirqGateAsBloqBase):
     exponent: float = 1.0
     global_shift: float = 0.0
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")
@@ -183,7 +176,7 @@ class XPowGate(CirqGateAsBloqBase):
     """
     exponent: Union[sympy.Expr, float] = 1.0
     global_shift: float = 0.0
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")
@@ -253,7 +246,7 @@ class YPowGate(CirqGateAsBloqBase):
     """
     exponent: Union[sympy.Expr, float] = 1.0
     global_shift: float = 0.0
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")
@@ -321,7 +314,7 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
 @frozen
 class Rx(CirqGateAsBloqBase):
     angle: Union[sympy.Expr, float]
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")
@@ -342,7 +335,7 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
 @frozen
 class Ry(CirqGateAsBloqBase):
     angle: Union[sympy.Expr, float]
-    eps: float = 1e-11
+    eps: SymbolicFloat = 1e-11
 
     def decompose_bloq(self) -> 'CompositeBloq':
         raise DecomposeTypeError(f"{self} is atomic")

qualtran/bloqs/chemistry/trotter/hubbard/qpe_cost_optimization.ipynb

@@ -111,7 +111,7 @@
     "import numpy as np\n",
     "import sympy\n",
     "\n",
-    "from qualtran.resource_counting.t_counts_from_sigma import _get_all_rotation_types\n",
+    "from qualtran.resource_counting.classify_bloqs import bloq_is_rotation\n",
     "from qualtran.resource_counting.generalizers import PHI\n",
     "from qualtran.cirq_interop.t_complexity_protocol import TComplexity\n",
     "from qualtran import Bloq\n",
@@ -130,11 +130,10 @@
     "\n",
     "\n",
     "def t_and_rot_counts_from_sigma(sigma: Dict['Bloq', Union[int, 'sympy.Expr']]) -> Tuple[int, int]:\n",
-    "    rotation_types = _get_all_rotation_types()\n",
     "    ret = sigma.get(TGate(), 0)\n",
     "    n_rot = 0\n",
     "    for bloq, counts in sigma.items():\n",
-    "        if isinstance(bloq, rotation_types):\n",
+    "        if bloq_is_rotation(bloq):\n",
     "            n_rot += counts\n",
     "    return ret, n_rot\n",
     "\n",

qualtran/bloqs/data_loading/select_swap_qrom_test.py

@@ -27,7 +27,7 @@
 )
 from qualtran.cirq_interop.t_complexity_protocol import t_complexity, TComplexity
 from qualtran.cirq_interop.testing import assert_circuit_inp_out_cirqsim
-from qualtran.resource_counting.t_counts_from_sigma import t_counts_from_sigma
+from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
 from qualtran.testing import assert_valid_bloq_decomposition
 
 
@@ -187,8 +187,9 @@ def test_qroam_t_complexity():
     qroam = SelectSwapQROM.build_from_data(
         [1, 2, 3, 4, 5, 6, 7, 8], target_bitsizes=(4,), log_block_sizes=(2,)
     )
-    _, sigma = qroam.call_graph()
-    assert t_counts_from_sigma(sigma) == qroam.t_complexity().t == 192
+    gate_counts = get_cost_value(qroam, QECGatesCost())
+    assert gate_counts == GateCounts(t=192, clifford=1082)
+    assert qroam.t_complexity() == TComplexity(t=192, clifford=1082)
 
 
 def test_qroam_many_registers():

qualtran/resource_counting/t_counts_from_sigma.py

@@ -11,45 +11,27 @@
 #  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 inspect
-import sys
-from typing import cast, Mapping, Optional, Tuple, Type, TYPE_CHECKING
+from typing import Mapping
 
 import cirq
 
+from qualtran import Bloq, Controlled
 from qualtran.symbolics import ceil, SymbolicInt
 
-if TYPE_CHECKING:
-    from qualtran import Bloq
-    from qualtran.bloqs.basic_gates.rotation import _HasEps
 
-
-def _get_all_rotation_types() -> Tuple[Type['_HasEps'], ...]:
-    """Returns all classes defined in bloqs.basic_gates which have an attribute `eps`."""
-    from qualtran.bloqs.basic_gates import GlobalPhase
-    from qualtran.bloqs.basic_gates.rotation import _HasEps
-
-    bloqs_to_exclude = [GlobalPhase]
-
-    return tuple(
-        cast(Type['_HasEps'], v)  # Can't use `issubclass` with protocols with attributes.
-        for (_, v) in inspect.getmembers(sys.modules['qualtran.bloqs.basic_gates'], inspect.isclass)
-        if isinstance(v, _HasEps) and v not in bloqs_to_exclude
-    )
-
-
-def t_counts_from_sigma(
-    sigma: Mapping['Bloq', SymbolicInt],
-    rotation_types: Optional[Tuple[Type['_HasEps'], ...]] = None,
-) -> SymbolicInt:
+def t_counts_from_sigma(sigma: Mapping['Bloq', SymbolicInt]) -> SymbolicInt:
     """Aggregates T-counts from a sigma dictionary by summing T-costs for all rotation bloqs."""
     from qualtran.bloqs.basic_gates import TGate
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+    from qualtran.resource_counting.classify_bloqs import bloq_is_rotation
 
-    if rotation_types is None:
-        rotation_types = _get_all_rotation_types()
     ret = sigma.get(TGate(), 0) + sigma.get(TGate().adjoint(), 0)
     for bloq, counts in sigma.items():
-        if isinstance(bloq, rotation_types) and not cirq.has_stabilizer_effect(bloq):
+        if bloq_is_rotation(bloq) and not cirq.has_stabilizer_effect(bloq):
+            if isinstance(bloq, Controlled):
+                # TODO native controlled rotation bloqs missing (CRz, CRy etc.)
+                #      https://github.com/quantumlib/Qualtran/issues/878
+                bloq = bloq.subbloq
+            assert hasattr(bloq, 'eps')
             ret += ceil(TComplexity.rotation_cost(bloq.eps)) * counts
     return ret

qualtran/resource_counting/t_counts_from_sigma_test.py

@@ -19,31 +19,14 @@
     Rx,
     Ry,
     Rz,
-    SU2RotationGate,
     TGate,
     Toffoli,
     XPowGate,
     YPowGate,
     ZPowGate,
 )
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
-from qualtran.resource_counting.t_counts_from_sigma import (
-    _get_all_rotation_types,
-    t_counts_from_sigma,
-)
-
-
-def test_all_rotation_types():
-    assert set(_get_all_rotation_types()) == {
-        CZPowGate,
-        Rx,
-        Ry,
-        Rz,
-        XPowGate,
-        YPowGate,
-        ZPowGate,
-        SU2RotationGate,
-    }
+from qualtran.resource_counting.t_counts_from_sigma import t_counts_from_sigma
 
 
 def test_t_counts_from_sigma():