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
 
 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, log2, ssum, SymbolicFloat, SymbolicInt
 
 from ._call_graph import get_bloq_callee_counts
 from ._costing import CostKey
@@ -115,6 +116,12 @@ def __str__(self):
         return f'{self.gateset_name} counts'
 
 
+def _mapping_to_counter(mapping: Mapping[int, int]) -> Counter[int]:
+    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,21 +135,69 @@ class GateCounts:
     cswap: SymbolicInt = 0
     and_bloq: SymbolicInt = 0
     clifford: SymbolicInt = 0
-    rotation: SymbolicInt = 0
     measurement: SymbolicInt = 0
+    binned_rotation_epsilons: Counter[int] = field(
+        factory=Counter, converter=_mapping_to_counter, eq=lambda d: tuple(d.items())
+    )
+    eps_bin_prec: int = 20
+
+    @classmethod
+    def from_rotation_with_eps(cls, eps: float, *, eps_bin_prec: int = 20, n_rotations: int = 1):
+        """Construct a GateCount with a rotation of precision `eps`.
+
+        Args:
+            eps: precision to synthesize the rotation(s).
+            eps_bin_prec: number of bits to approximate `eps` to, defaults to 10.
+            n_rotations: number of rotations, defaults to 1.
+        """
+        eps_bin = int(np.ceil(eps * 2**eps_bin_prec))
+        # treat any eps < 2**(-eps_bin_prec) as 2**(-eps_bin_prec)
+        eps_bin = max(1, eps_bin)
+        return cls(binned_rotation_epsilons=Counter({eps_bin: n_rotations}))
+
+    def with_rotation_eps_bin_prec(self, new_eps_bin_prec: int) -> 'GateCounts':
+        """Returns `GateCounts` with a new bin precision for rotation epsilons."""
+        if new_eps_bin_prec == self.eps_bin_prec:
+            return self
+
+        def _get_new_eps_bin(eps_bin):
+            return int(eps_bin * 2 ** (new_eps_bin_prec - self.eps_bin_prec))
+
+        new_binned_rotation_epsilons = Counter(
+            {
+                _get_new_eps_bin(eps_bin): n_rot
+                for eps_bin, n_rot in self.binned_rotation_epsilons.items()
+            }
+        )
+
+        return attrs.evolve(
+            self,
+            binned_rotation_epsilons=new_binned_rotation_epsilons,
+            eps_bin_prec=new_eps_bin_prec,
+        )
+
+    def iter_rotations_with_epsilon(self) -> Iterator[tuple[float, SymbolicInt]]:
+        """Iterate through the rotation precisions (epsilon) and their frequency."""
+        for eps_bin, n_rot in self.binned_rotation_epsilons.items():
+            yield eps_bin / 2**self.eps_bin_prec, n_rot
 
     def __add__(self, other):
         if not isinstance(other, GateCounts):
             raise TypeError(f"Can only add other `GateCounts` objects, not {self}")
 
+        eps_bin_prec = max(self.eps_bin_prec, other.eps_bin_prec)
+        this = self.with_rotation_eps_bin_prec(eps_bin_prec)
+        other = other.with_rotation_eps_bin_prec(other.eps_bin_prec)
+
         return GateCounts(
-            t=self.t + other.t,
-            toffoli=self.toffoli + other.toffoli,
-            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,
+            t=this.t + other.t,
+            toffoli=this.toffoli + other.toffoli,
+            cswap=this.cswap + other.cswap,
+            and_bloq=this.and_bloq + other.and_bloq,
+            clifford=this.clifford + other.clifford,
+            measurement=this.measurement + other.measurement,
+            binned_rotation_epsilons=this.binned_rotation_epsilons + other.binned_rotation_epsilons,
+            eps_bin_prec=eps_bin_prec,
         )
 
     def __mul__(self, other):
@@ -152,8 +207,11 @@ 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()}
+            ),
+            eps_bin_prec=self.eps_bin_prec,
         )
 
     def __rmul__(self, other):
@@ -174,36 +232,59 @@ 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
+            if key == 'eps_bin_prec':
+                # rotations non-empty
+                return len(self.binned_rotation_epsilons) > 0
+            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)
+
+    @property
+    def rotation_to_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.rotation_to_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.rotation_to_t
         return {'n_t': n_t, 'n_ccz': n_ccz}
 
+    def n_rotation_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 +297,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.n_rotation_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 +323,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 +349,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

@@ -78,12 +78,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-10, n_rotations=10),
         ],
         # Recursive
         [

qualtran/surface_code/algorithm_summary_test.py

@@ -38,13 +38,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-10, n_rotations=10),
+            ),
         ],
         [
             mcmt.MultiControlPauli(cvs=(1, 1, 1), target_gate=cirq.X),

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.n_rotation_ignoring_eps()
+    ret = attrs.evolve(alg.n_logical_gates, binned_rotation_epsilons={})
     if n_rotations > 0:
         ret = (
             ret

qualtran/surface_code/beverland_et_al_model_test.py

@@ -38,7 +38,10 @@ class Test:
     Test(
         alg=AlgorithmSummary(
             n_algo_qubits=100,
-            n_logical_gates=GateCounts(rotation=30_100, measurement=int(1.4e6)),
+            n_logical_gates=(
+                GateCounts.from_rotation_with_eps(0, n_rotations=30_000)
+                + GateCounts(measurement=int(1.4e6))
+            ),
             n_rotation_layers=501,
         ),
         error_budget=1e-3,
@@ -50,8 +53,9 @@ class Test:
     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)
+            n_logical_gates=(
+                GateCounts(t=int(5.53e7), toffoli=int(1.35e11), measurement=int(1.37e9))
+                + GateCounts.from_rotation_with_eps(0, n_rotations=int(2.06e8))
             ),
             n_rotation_layers=int(2.05e8),
         ),
@@ -64,8 +68,9 @@ class Test:
     Test(
         alg=AlgorithmSummary(
             n_algo_qubits=12581,
-            n_logical_gates=GateCounts(
-                t=12, rotation=12, toffoli=int(3.73e9), measurement=int(1.08e9)
+            n_logical_gates=(
+                GateCounts(t=12, toffoli=int(3.73e9), measurement=int(1.08e9))
+                + GateCounts.from_rotation_with_eps(0, n_rotations=12)
             ),
             n_rotation_layers=12,
         ),
@@ -106,5 +111,6 @@ def test_t_states(test: Test):
     assert got == pytest.approx(test.t_states, rel=0.1)
 
 
+@pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook('beverland_et_al_model')