From 4156b56713e3d9353e597df59170e2586e2acf3f Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 24 Apr 2024 11:56:20 -0700
Subject: [PATCH 1/2] success prob

---
 qualtran/resource_counting/__init__.py        |  2 +
 qualtran/resource_counting/_success_prob.py   | 45 +++++++++++++++++++
 .../resource_counting/_success_prob_test.py   | 26 +++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 qualtran/resource_counting/_success_prob.py
 create mode 100644 qualtran/resource_counting/_success_prob_test.py

diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py
index 0f9d1885f2..8288931a6c 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -31,4 +31,6 @@
 
 from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
 
+from ._success_prob import SuccessProb
+
 from . import generalizers
diff --git a/qualtran/resource_counting/_success_prob.py b/qualtran/resource_counting/_success_prob.py
new file mode 100644
index 0000000000..24c6826772
--- /dev/null
+++ b/qualtran/resource_counting/_success_prob.py
@@ -0,0 +1,45 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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 abc
+import logging
+from typing import Callable, Dict, Generic, Sequence, Tuple, TYPE_CHECKING
+
+import sympy
+from attrs import frozen
+
+from qualtran import Bloq, DecomposeNotImplementedError, DecomposeTypeError
+
+from . import CostValT
+from ._call_graph import get_bloq_callee_counts
+from ._costing import CostKey
+
+logger = logging.getLogger(__name__)
+
+
+@frozen
+class SuccessProb(CostKey[float]):
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], float]) -> float:
+        tot: float = 1.0
+        callees = get_bloq_callee_counts(bloq)
+        logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees))
+        for callee, n in callees:
+            v = get_callee_cost(callee)
+            tot *= v**n
+        return tot
+
+    def zero(self) -> CostValT:
+        return 1.0  # under multiplication, 1 is the identity.
+
+    def __str__(self):
+        return 'success prob'
diff --git a/qualtran/resource_counting/_success_prob_test.py b/qualtran/resource_counting/_success_prob_test.py
new file mode 100644
index 0000000000..8b1fdef54a
--- /dev/null
+++ b/qualtran/resource_counting/_success_prob_test.py
@@ -0,0 +1,26 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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.
+from qualtran.bloqs.for_testing.costing import CostingBloq
+from qualtran.resource_counting import get_cost_cache, get_cost_value, SuccessProb
+
+
+def test_coin_flip():
+    flip = CostingBloq('CoinFlip', num_qubits=1, static_costs=[(SuccessProb(), 0.5)])
+    algo = CostingBloq('Algo', num_qubits=0, callees=[(flip, 4)])
+
+    p = get_cost_value(algo, SuccessProb())
+    assert p == 0.5**4
+
+    costs = get_cost_cache(algo, SuccessProb())
+    assert costs == {algo: p, flip: 0.5}

From 57f81ab403fdd065c548cae52a5104fefd9c01f5 Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Mon, 20 May 2024 15:12:31 -0700
Subject: [PATCH 2/2] docs

---
 qualtran/bloqs/for_testing/costing.py          | 11 ++++++++++-
 qualtran/resource_counting/_success_prob.py    | 18 +++++++++++-------
 .../resource_counting/_success_prob_test.py    |  2 +-
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/qualtran/bloqs/for_testing/costing.py b/qualtran/bloqs/for_testing/costing.py
index cc0e24114e..465bf3d64d 100644
--- a/qualtran/bloqs/for_testing/costing.py
+++ b/qualtran/bloqs/for_testing/costing.py
@@ -24,6 +24,13 @@ def _convert_callees(callees: Sequence[BloqCountT]) -> Tuple[BloqCountT, ...]:
     return tuple(callees)
 
 
+def _convert_static_costs(
+    static_costs: Sequence[Tuple[CostKey, Any]]
+) -> Tuple[Tuple[CostKey, Any], ...]:
+    # Convert to tuples in a type-checked way.
+    return tuple(static_costs)
+
+
 @frozen
 class CostingBloq(Bloq):
     """A bloq that lets you set the costs via attributes."""
@@ -31,7 +38,9 @@ class CostingBloq(Bloq):
     name: str
     num_qubits: int
     callees: Sequence[BloqCountT] = field(converter=_convert_callees, factory=tuple)
-    static_costs: Sequence[Tuple[CostKey, Any]] = field(converter=tuple, factory=tuple)
+    static_costs: Sequence[Tuple[CostKey, Any]] = field(
+        converter=_convert_static_costs, factory=tuple
+    )
 
     @property
     def signature(self) -> 'Signature':
diff --git a/qualtran/resource_counting/_success_prob.py b/qualtran/resource_counting/_success_prob.py
index 24c6826772..4aab783bdd 100644
--- a/qualtran/resource_counting/_success_prob.py
+++ b/qualtran/resource_counting/_success_prob.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -11,16 +11,13 @@
 #  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 abc
 import logging
-from typing import Callable, Dict, Generic, Sequence, Tuple, TYPE_CHECKING
+from typing import Callable
 
-import sympy
 from attrs import frozen
 
-from qualtran import Bloq, DecomposeNotImplementedError, DecomposeTypeError
+from qualtran import Bloq
 
-from . import CostValT
 from ._call_graph import get_bloq_callee_counts
 from ._costing import CostKey
 
@@ -29,6 +26,13 @@
 
 @frozen
 class SuccessProb(CostKey[float]):
+    """The success probability of a bloq.
+
+    A bloq's success probability is the multiplicative product of its callees'
+    success probabilities. Bloqs that have a specific success probability should override
+    `my_static_costs` to provide their actual success probability.
+    """
+
     def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], float]) -> float:
         tot: float = 1.0
         callees = get_bloq_callee_counts(bloq)
@@ -38,7 +42,7 @@ def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], float]) -> f
             tot *= v**n
         return tot
 
-    def zero(self) -> CostValT:
+    def zero(self) -> float:
         return 1.0  # under multiplication, 1 is the identity.
 
     def __str__(self):
diff --git a/qualtran/resource_counting/_success_prob_test.py b/qualtran/resource_counting/_success_prob_test.py
index 8b1fdef54a..3c4eeb552d 100644
--- a/qualtran/resource_counting/_success_prob_test.py
+++ b/qualtran/resource_counting/_success_prob_test.py
@@ -1,4 +1,4 @@
-#  Copyright 2023 Google LLC
+#  Copyright 2024 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.