prototype: use rotation eps in GateCounts

qualtran/resource_counting/_bloq_counts.py

@@ -12,8 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 import logging
-from collections import defaultdict
-from typing import Callable, Dict, Sequence, Tuple, TYPE_CHECKING
+from collections import Counter, defaultdict
+from typing import Callable, Dict, Mapping, Sequence, Tuple, TYPE_CHECKING
 
 import attrs
 import networkx as nx
@@ -112,6 +112,12 @@ def __str__(self):
         return f'{self.gateset_name} counts'
 
 
+def _mapping_to_counter(mapping: Mapping[float, int]) -> Counter[float]:
+    if isinstance(mapping, Counter):
+        return mapping
+    return Counter(mapping)
+
+
 @frozen(kw_only=True)
 class GateCounts:
     """A data class of counts of the typical target gates in a compilation.
@@ -125,8 +131,17 @@ class GateCounts:
     cswap: int = 0
     and_bloq: int = 0
     clifford: int = 0
-    rotation: int = 0
     measurement: int = 0
+    rotation_epsilons: Counter[float] = field(factory=Counter, converter=_mapping_to_counter)
+
+    @property
+    def rotation(self):
+        from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+
+        return sum(
+            n_rotations * int(TComplexity.rotation_cost(eps))
+            for eps, n_rotations in self.rotation_epsilons.items()
+        )
 
     def __add__(self, other):
         if not isinstance(other, GateCounts):
@@ -138,8 +153,8 @@ def __add__(self, other):
             cswap=self.cswap + other.cswap,
             and_bloq=self.and_bloq + other.and_bloq,
             clifford=self.clifford + other.clifford,
-            rotation=self.rotation + other.rotation,
             measurement=self.measurement + other.measurement,
+            rotation_epsilons=self.rotation_epsilons + other.rotation_epsilons,
         )
 
     def __mul__(self, other):
@@ -149,8 +164,8 @@ def __mul__(self, other):
             cswap=other * self.cswap,
             and_bloq=other * self.and_bloq,
             clifford=other * self.clifford,
-            rotation=other * self.rotation,
             measurement=other * self.measurement,
+            rotation_epsilons=Counter({k: other * v for k, v in self.rotation_epsilons.items()}),
         )
 
     def __rmul__(self, other):
@@ -167,7 +182,13 @@ def __str__(self):
 
     def asdict(self):
         d = attrs.asdict(self)
-        return {k: v for k, v in d.items() if v > 0}
+
+        def _keep(key, value) -> bool:
+            if key == 'rotation_epsilons':
+                return value
+            return value > 0
+
+        return {k: v for k, v in d.items() if _keep(k, v)}
 
     def total_t_count(
         self,
@@ -232,6 +253,7 @@ class QECGatesCost(CostKey[GateCounts]):
 
     def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts]) -> GateCounts:
         from qualtran.bloqs.basic_gates import TGate, Toffoli, TwoBitCSwap
+        from qualtran.bloqs.basic_gates.rotation import _HasEps
         from qualtran.bloqs.mcmt.and_bloq import And
 
         # T gates
@@ -257,7 +279,8 @@ def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts])
             return GateCounts(clifford=1)
 
         if bloq_is_rotation(bloq):
-            return GateCounts(rotation=1)
+            assert isinstance(bloq, _HasEps)
+            return GateCounts(rotation_epsilons={bloq.eps: 1})
 
         # Recursive case
         totals = GateCounts()

qualtran/resource_counting/_bloq_counts_test.py

@@ -74,12 +74,15 @@ def test_qec_gates_cost():
         # And
         [mcmt.And(), GateCounts(and_bloq=1)],
         # Rotations
-        [basic_gates.ZPowGate(exponent=0.1, global_shift=0.0, eps=1e-11), GateCounts(rotation=1)],
+        [
+            basic_gates.ZPowGate(exponent=0.1, global_shift=0.0, eps=1e-11),
+            GateCounts(rotation_epsilons={1e-11: 1}),
+        ],
         [
             rotations.phase_gradient.PhaseGradientUnitary(
                 bitsize=10, exponent=1, is_controlled=False, eps=1e-10
             ),
-            GateCounts(rotation=10),
+            GateCounts(rotation_epsilons={1e-10: 10}),
         ],
         # Recursive
         [

qualtran/surface_code/algorithm_summary_test.py

@@ -38,13 +38,17 @@
         [mcmt.And(), AlgorithmSummary(n_algo_qubits=3, n_logical_gates=GateCounts(and_bloq=1))],
         [
             basic_gates.ZPowGate(exponent=0.1, global_shift=0.0, eps=1e-11),
-            AlgorithmSummary(n_algo_qubits=1, n_logical_gates=GateCounts(rotation=1)),
+            AlgorithmSummary(
+                n_algo_qubits=1, n_logical_gates=GateCounts(rotation_epsilons={1e-11: 1})
+            ),
         ],
         [
             rotations.phase_gradient.PhaseGradientUnitary(
                 bitsize=10, exponent=1, is_controlled=False, eps=1e-10
             ),
-            AlgorithmSummary(n_algo_qubits=10, n_logical_gates=GateCounts(rotation=10)),
+            AlgorithmSummary(
+                n_algo_qubits=10, n_logical_gates=GateCounts(rotation_epsilons={1e-10: 10})
+            ),
         ],
         [
             mcmt.MultiControlPauli(cvs=(1, 1, 1), target_gate=cirq.X),

qualtran/surface_code/beverland_et_al_model.py

@@ -116,14 +116,12 @@ def n_discrete_logical_gates(
         rotation_model: Cost model used to compute the number of T gates
             needed to approximate rotations.
     """
-    n_rotations: int = alg.n_logical_gates.rotation
-    ret = attrs.evolve(alg.n_logical_gates, rotation=0)
-    if n_rotations > 0:
-        ret = (
-            ret
-            + rotation_model.preparation_overhead(eps_syn)
-            + n_rotations * rotation_model.rotation_cost(eps_syn / n_rotations)
-        )
+    rotation_epsilons: dict[float, int] = alg.n_logical_gates.rotation_epsilons
+    ret = attrs.evolve(alg.n_logical_gates, rotation_epsilons={})
+    if rotation_epsilons:
+        rotation_model.preparation_overhead(min(eps for eps in rotation_epsilons.values()))
+        for eps, n_rotations in rotation_epsilons.items():
+            ret += n_rotations * rotation_model.rotation_cost(eps)
     return ret
 
 

qualtran/surface_code/beverland_et_al_model_test.py

@@ -38,11 +38,13 @@ class Test:
     Test(
         alg=AlgorithmSummary(
             n_algo_qubits=100,
-            n_logical_gates=GateCounts(rotation=30_100, measurement=int(1.4e6)),
+            n_logical_gates=GateCounts(
+                rotation_epsilons={1e-3 / 30_000: 30_000}, measurement=int(1.4e6)
+            ),
             n_rotation_layers=501,
         ),
         error_budget=1e-3,
-        c_min=1.5e6,
+        c_min=2.5e6,
         time_steps=1.5e5,
         code_distance=9,
         t_states=602000,
@@ -51,7 +53,10 @@ class Test:
         alg=AlgorithmSummary(
             n_algo_qubits=1318,
             n_logical_gates=GateCounts(
-                t=int(5.53e7), rotation=int(2.06e8), toffoli=int(1.35e11), measurement=int(1.37e9)
+                t=int(5.53e7),
+                rotation_epsilons={1e-2 / 2.06e8: int(2.06e8)},
+                toffoli=int(1.35e11),
+                measurement=int(1.37e9),
             ),
             n_rotation_layers=int(2.05e8),
         ),
@@ -65,7 +70,10 @@ class Test:
         alg=AlgorithmSummary(
             n_algo_qubits=12581,
             n_logical_gates=GateCounts(
-                t=12, rotation=12, toffoli=int(3.73e9), measurement=int(1.08e9)
+                t=12,
+                rotation_epsilons={1 / 3 / 12: 12},
+                toffoli=int(3.73e9),
+                measurement=int(1.08e9),
             ),
             n_rotation_layers=12,
         ),

qualtran/surface_code/ui.py

@@ -522,24 +522,30 @@ def update(
     qec_name: str,
     magic_name: str,
     magic_count: int,
-    rotaion_model_name: str,
+    rotation_model_name: str,
 ):
     """Updates the visualization."""
     if any(x is None for x in [physical_error_rate, error_budget, *algorithm_data, magic_count]):
         raise PreventUpdate
 
     # TODO: We implicitly rely on the order of the input components
     qubits, measurements, ts, toffolis, rotations, n_rotation_layers = algorithm_data
+
+    rotation_eps = 1e-3 / rotations  # TODO this should not be a magic number
+
     algorithm = AlgorithmSummary(
         n_algo_qubits=qubits,
         n_logical_gates=GateCounts(
-            measurement=measurements, t=ts, toffoli=toffolis, rotation=rotations
+            measurement=measurements,
+            t=ts,
+            toffoli=toffolis,
+            rotation_epsilons={rotation_eps: rotations},
         ),
         n_rotation_layers=n_rotation_layers,
     )
     qec = _QEC_SCHEMES[qec_name]
     magic_factory = _MAGIC_FACTORIES[magic_name]
-    rotation_model = _ROTATION_MODELS[rotaion_model_name]
+    rotation_model = _ROTATION_MODELS[rotation_model_name]
     n_logical_gates = beverland_et_al_model.n_discrete_logical_gates(
         eps_syn=error_budget / 3, alg=algorithm, rotation_model=rotation_model
     )

qualtran/surface_code/ui_test.py

@@ -29,7 +29,7 @@ def test_ensure_support_for_all_supported_models(estimation_model: str):
         qec_name='GidneyFowler',
         magic_name='FifteenToOne733',
         magic_count=1,
-        rotaion_model_name='BeverlandEtAlRotationCost',
+        rotation_model_name='BeverlandEtAlRotationCost',
     )
 
 
@@ -81,7 +81,7 @@ def test_update(estimation_model: str, desired):
         qec_name='GidneyFowler',
         magic_name='FifteenToOne733',
         magic_count=1,
-        rotaion_model_name='BeverlandEtAlRotationCost',
+        rotation_model_name='BeverlandEtAlRotationCost',
     )
     assert (
         display_runtime,
@@ -104,7 +104,7 @@ def test_update_bad_input():
             qec_name='GidneyFowler',
             magic_name='FifteenToOne733',
             magic_count=1,
-            rotaion_model_name='BeverlandEtAlRotationCost',
+            rotation_model_name='BeverlandEtAlRotationCost',
         )