Use rotation eps in GateCounts

qualtran/cirq_interop/t_complexity_protocol.py

@@ -20,7 +20,7 @@
 
 from qualtran import Bloq, Controlled, DecomposeNotImplementedError, DecomposeTypeError
 from qualtran.resource_counting import SympySymbolAllocator
-from qualtran.symbolics import ceil, log2, SymbolicFloat, SymbolicInt
+from qualtran.symbolics import ceil, SymbolicFloat, SymbolicInt
 
 from .decompose_protocol import _decompose_once_considering_known_decomposition
 
@@ -38,7 +38,9 @@ class TComplexity:
 
     @staticmethod
     def rotation_cost(eps: SymbolicFloat) -> SymbolicFloat:
-        return ceil(1.149 * log2(1.0 / eps) + 9.2)
+        from qualtran.resource_counting import GateCounts
+
+        return GateCounts.rotation_t_cost(eps)
 
     def t_incl_rotations(self, eps: float = 1e-11) -> SymbolicInt:
         """Return the total number of T gates after compiling rotations"""

qualtran/resource_counting/_bloq_counts.py

@@ -12,15 +12,16 @@
 #  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, Iterator, Mapping, Sequence, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import networkx as nx
+import numpy as np
 import sympy
 from attrs import field, frozen
 
-from qualtran.symbolics import SymbolicInt
+from qualtran.symbolics import ceil, is_symbolic, log2, ssum, SymbolicFloat, SymbolicInt
 
 from ._call_graph import get_bloq_callee_counts
 from ._costing import CostKey
@@ -115,6 +116,16 @@ def __str__(self):
         return f'{self.gateset_name} counts'
 
 
+FloatRepr_T = Union[str, sympy.Expr]
+"""The type to represent floats as, to use as safe keys in mappings."""
+
+
+def _mapping_to_counter(mapping: Mapping[FloatRepr_T, int]) -> Counter[FloatRepr_T]:
+    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.
@@ -128,8 +139,38 @@ class GateCounts:
     cswap: SymbolicInt = 0
     and_bloq: SymbolicInt = 0
     clifford: SymbolicInt = 0
-    rotation: SymbolicInt = 0
     measurement: SymbolicInt = 0
+    binned_rotation_epsilons: Counter[FloatRepr_T] = field(
+        factory=Counter, converter=_mapping_to_counter, eq=lambda d: tuple(d.items())
+    )
+
+    @classmethod
+    def from_rotation_with_eps(
+        cls, eps: SymbolicFloat, *, n_rotations: int = 1, eps_repr_prec: int = 10
+    ):
+        """Construct a GateCount with a rotation of precision `eps`.
+
+        Formats the value of `eps` as a string using `np.format_float_scientific`,
+        to use as a safe dictionary key. If `eps` is symbolic, it is used as-is.
+
+        Args:
+            eps: precision to synthesize the rotation(s).
+            eps_repr_prec: number of digits to approximate `eps` to. Uses 10 by default.
+                           See `np.format_float_scientific` for more details.
+                           If `eps` is symbolic, this parameter is ignored.
+            n_rotations: number of rotations, defaults to 1.
+        """
+        if is_symbolic(eps):
+            eps_bin: FloatRepr_T = eps
+        else:
+            eps_bin = np.format_float_scientific(eps, precision=eps_repr_prec, unique=False)
+        return cls(binned_rotation_epsilons=Counter({eps_bin: n_rotations}))
+
+    def iter_rotations_with_epsilon(self) -> Iterator[tuple[SymbolicFloat, SymbolicInt]]:
+        """Iterate through the rotation precisions (epsilon) and their frequency."""
+        for eps_bin, n_rot in self.binned_rotation_epsilons.items():
+            eps: SymbolicFloat = eps_bin if is_symbolic(eps_bin) else float(eps_bin)
+            yield eps, n_rot
 
     def __add__(self, other):
         if not isinstance(other, GateCounts):
@@ -141,8 +182,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,
+            binned_rotation_epsilons=self.binned_rotation_epsilons + other.binned_rotation_epsilons,
         )
 
     def __mul__(self, other):
@@ -152,8 +193,10 @@ 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,
+            binned_rotation_epsilons=Counter(
+                {eps_bin: other * n_rot for eps_bin, n_rot in self.binned_rotation_epsilons.items()}
+            ),
         )
 
     def __rmul__(self, other):
@@ -174,36 +217,56 @@ def _is_nonzero(v):
                 return True
             return maybe_nonzero
 
-        return {k: v for k, v in d.items() if _is_nonzero(v)}
+        def _keep(key, value) -> bool:
+            if key == 'binned_rotation_epsilons':
+                return value
+            return _is_nonzero(value)
+
+        return {k: v for k, v in d.items() if _keep(k, v)}
+
+    @staticmethod
+    def rotation_t_cost(eps: SymbolicFloat) -> SymbolicInt:
+        """T-cost of a single Z rotation with precision `eps`.
+
+        References:
+            [Efficient synthesis of universal Repeat-Until-Success circuits](https://arxiv.org/abs/1404.5320)
+            Bocharov et. al. 2014. Page 4, Paragraph "Simulation Results."
+        """
+        return ceil(1.149 * log2(1.0 / eps) + 9.2)
+
+    def total_rotations_as_t(self) -> SymbolicInt:
+        """Total number of T Gates for the rotations."""
+        return ssum(
+            n_rotations * self.rotation_t_cost(eps)
+            for eps, n_rotations in self.iter_rotations_with_epsilon()
+        )
 
     def total_t_count(
-        self,
-        ts_per_toffoli: int = 4,
-        ts_per_cswap: int = 7,
-        ts_per_and_bloq: int = 4,
-        ts_per_rotation: int = 11,
+        self, ts_per_toffoli: int = 4, ts_per_cswap: int = 7, ts_per_and_bloq: int = 4
     ) -> int:
         """Get the total number of T Gates for the `GateCounts` object.
 
         This simply multiplies each gate type by its cost in terms of T gates, which is configurable
         via the arguments to this method.
-
-        The default value for `ts_per_rotation` assumes the rotation is approximated using
-        `Mixed fallback` protocol with error budget 1e-3.
         """
         return (
             self.t
             + ts_per_toffoli * self.toffoli
             + ts_per_cswap * self.cswap
             + ts_per_and_bloq * self.and_bloq
-            + ts_per_rotation * self.rotation
+            + self.total_rotations_as_t()
         )
 
-    def total_t_and_ccz_count(self, ts_per_rotation: int = 11) -> Dict[str, SymbolicInt]:
+    def total_t_and_ccz_count(self) -> Dict[str, SymbolicInt]:
         n_ccz = self.toffoli + self.cswap + self.and_bloq
-        n_t = self.t + ts_per_rotation * self.rotation
+        n_t = self.t + self.total_rotations_as_t()
         return {'n_t': n_t, 'n_ccz': n_ccz}
 
+    @property
+    def rotations_ignoring_eps(self) -> SymbolicInt:
+        """Total number of rotations, ignoring the individual precisions."""
+        return ssum(self.binned_rotation_epsilons.values())
+
     def total_beverland_count(self) -> Dict[str, SymbolicInt]:
         r"""Counts used by Beverland. et. al. using notation from the reference.
 
@@ -216,17 +279,20 @@ def total_beverland_count(self) -> Dict[str, SymbolicInt]:
            Toffoli gates. Since we don't compile the 'layers' explicitly, we set this to be the
            number of rotations.
 
+        Note: This costing method ignores the individual rotation precisions (`eps`).
+
         Reference:
             https://arxiv.org/abs/2211.07629.
             Equation D3.
         """
         toffoli = self.toffoli + self.and_bloq + self.cswap
+        rotation = self.rotations_ignoring_eps
         return {
             'meas': self.measurement,
-            'R': self.rotation,
+            'R': rotation,
             'T': self.t,
             'Tof': toffoli,
-            'D_R': self.rotation,
+            'D_R': rotation,
         }
 
 
@@ -239,6 +305,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
@@ -264,7 +331,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.from_rotation_with_eps(bloq.eps)
 
         # Recursive case
         totals = GateCounts()

qualtran/resource_counting/_bloq_counts_test.py

@@ -18,6 +18,7 @@
 from qualtran.bloqs.basic_gates import Hadamard, TGate, Toffoli
 from qualtran.bloqs.for_testing.costing import make_example_costing_bloqs
 from qualtran.resource_counting import BloqCount, GateCounts, get_cost_value, QECGatesCost
+from qualtran.symbolics import ceil, log2
 
 
 def test_bloq_count():
@@ -59,6 +60,24 @@ def test_gate_counts():
     assert str(gc2) == 't: n, cswap: 2'
 
 
+def test_gate_counts_rotations():
+    gc = GateCounts.from_rotation_with_eps(1e-10, n_rotations=4)
+    assert gc == GateCounts(binned_rotation_epsilons={'1.0000000000e-10': 4})
+
+    eps = sympy.Symbol(r"\epsilon")
+    gc_symb = GateCounts.from_rotation_with_eps(eps, n_rotations=6)
+    assert gc_symb == GateCounts(binned_rotation_epsilons={eps: 6})
+
+
+def test_gate_counts_rotations_to_t():
+    gc = GateCounts.from_rotation_with_eps(1e-10, n_rotations=4)
+    assert gc.total_rotations_as_t() == 192
+
+    eps = sympy.Symbol(r"\epsilon")
+    gc_symb = GateCounts.from_rotation_with_eps(eps, n_rotations=6)
+    assert gc_symb.total_rotations_as_t() == 6 * ceil(1.149 * log2(1.0 / eps) + 9.2)
+
+
 def test_qec_gates_cost():
     algo = make_example_costing_bloqs()
     gc = get_cost_value(algo, QECGatesCost())
@@ -77,12 +96,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.from_rotation_with_eps(1e-11),
+        ],
         [
             rotations.phase_gradient.PhaseGradientUnitary(
                 bitsize=10, exponent=1, is_controlled=False, eps=1e-10
             ),
-            GateCounts(rotation=10),
+            GateCounts.from_rotation_with_eps(1e-11, n_rotations=10),
         ],
         # Recursive
         [mcmt.MultiControlX(cvs=(1, 1, 1)), GateCounts(and_bloq=2, measurement=2, clifford=3)],

qualtran/surface_code/algorithm_summary_test.py

@@ -37,13 +37,18 @@
         [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.from_rotation_with_eps(1e-11)
+            ),
         ],
         [
             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.from_rotation_with_eps(1e-11, n_rotations=10),
+            ),
         ],
         [
             mcmt.MultiControlX(cvs=(1, 1, 1)),

qualtran/surface_code/beverland_et_al_model.ipynb

@@ -208,11 +208,13 @@
    "source": [
     "qd_alg = AlgorithmSummary(\n",
     "    n_algo_qubits = 100,\n",
-    "    n_logical_gates = GateCounts(\n",
-    "        rotation=30_100,\n",
-    "        # Note in the paper the number of measurements\n",
-    "        # has an extra zero which we assume to be a typo.\n",
-    "        measurement=1.4e5,\n",
+    "    n_logical_gates = (\n",
+    "        GateCounts.from_rotation_with_eps(0, n_rotations=30_100)\n",
+    "        + GateCounts(\n",
+    "            # Note in the paper the number of measurements\n",
+    "            # has an extra zero which we assume to be a typo.\n",
+    "            measurement=int(1.4e5)\n",
+    "        )\n",
     "    ),\n",
     "    n_rotation_layers = 501\n",
     ")"
@@ -394,11 +396,10 @@
     "chem_alg = AlgorithmSummary(\n",
     "    n_algo_qubits=1318,\n",
     "    n_logical_gates=GateCounts(\n",
-    "        rotation=2.06e8,\n",
     "        measurement=1.37e9,\n",
     "        toffoli=1.35e11,\n",
     "        t=5.53e7,\n",
-    "    ),\n",
+    "    ) + GateCounts.from_rotation_with_eps(0, n_rotations=2.06e8),\n",
     "    n_rotation_layers=2.05e8,\n",
     ")\n",
     "chem_alg"
@@ -569,13 +570,12 @@
     "shor_alg = AlgorithmSummary(\n",
     "    n_algo_qubits=12581,\n",
     "    n_logical_gates=GateCounts(\n",
-    "        rotation=12,\n",
     "        measurement=1.08e9,\n",
     "        # Note in the paper the number of Toffoli operations is 3.73e10.\n",
     "        # However we assume that the exponent has a typo and that the number is 3.73e9.\n",
     "        toffoli=3.73e9,\n",
     "        t=12,\n",
-    "    ),\n",
+    "    ) + GateCounts.from_rotation_with_eps(0, n_rotations=12),\n",
     "    n_rotation_layers=12,\n",
     ")"
    ]

qualtran/surface_code/beverland_et_al_model.py

@@ -117,8 +117,8 @@ def n_discrete_logical_gates(
         rotation_model: Cost model used to compute the number of T gates
             needed to approximate rotations.
     """
-    n_rotations: SymbolicInt = alg.n_logical_gates.rotation
-    ret = attrs.evolve(alg.n_logical_gates, rotation=0)
+    n_rotations: SymbolicInt = alg.n_logical_gates.rotations_ignoring_eps
+    ret = attrs.evolve(alg.n_logical_gates, binned_rotation_epsilons={})
     if n_rotations > 0:
         ret = (
             ret