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

---
 qualtran/resource_counting/__init__.py        |  1 +
 qualtran/resource_counting/_success_prob.py   | 45 +++++++++++++++++++
 .../resource_counting/_success_prob_test.py   | 26 +++++++++++
 3 files changed, 72 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 228b3d1932..8f13b7cbd5 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -33,5 +33,6 @@
 
 from ._bloq_counts import BloqCount, QECGatesCost
 from ._qubit_counts import QubitCount
+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}