First preparations for making And an atomic/leaf bloq (#1347)

* First preparations for making And an atomic/leaf bloq

* fix notebook

due to change away from defaultdict

* format

Author: mpharrigan

qualtran/bloqs/arithmetic/controlled_add_or_subtract_test.py

@@ -36,9 +36,7 @@
     _ctrl_add_or_sub_unsigned,
     ControlledAddOrSubtract,
 )
-from qualtran.bloqs.basic_gates import TGate
-from qualtran.resource_counting import get_cost_value
-from qualtran.resource_counting._bloq_counts import BloqCount
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_examples(bloq_autotester):
@@ -107,4 +105,6 @@ def test_t_complexity():
     dtype = QUInt(n)
     bloq = ControlledAddOrSubtract(dtype, dtype)
 
-    assert get_cost_value(bloq, BloqCount.for_gateset('t')) == {TGate(): 4 * n - 4}
+    counts = get_cost_value(bloq, QECGatesCost()).total_t_and_ccz_count()
+    assert counts['n_t'] == 0, 'toffoli only'
+    assert counts['n_ccz'] == n - 1

qualtran/bloqs/arithmetic/permutation_test.py

@@ -37,9 +37,10 @@
     Permutation,
     PermutationCycle,
 )
-from qualtran.bloqs.basic_gates import CNOT, TGate, XGate
-from qualtran.bloqs.bookkeeping import Allocate, ArbitraryClifford, Free
-from qualtran.resource_counting.generalizers import ignore_split_join
+from qualtran.bloqs.basic_gates import CNOT, XGate
+from qualtran.bloqs.bookkeeping import Allocate, Free
+from qualtran.bloqs.mcmt import And
+from qualtran.resource_counting.generalizers import generalize_cvs, ignore_split_join
 from qualtran.symbolics import ceil, log2, slen
 
 
@@ -62,11 +63,14 @@ def test_permutation_cycle_unitary_and_call_graph():
         bloq.tensor_contract(), np.array([[0, 0, 1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
     )
 
-    _, sigma = bloq.call_graph(generalizer=ignore_split_join)
+    cv = sympy.Symbol('cv')
+    _, sigma = bloq.call_graph(
+        generalizer=[ignore_split_join, generalize_cvs], keep=lambda b: isinstance(b, And)
+    )
     assert sigma == {
         CNOT(): 8,
-        TGate(): 16,
-        ArbitraryClifford(n=2): 76,
+        And(cv1=cv, cv2=cv): 4,
+        And(cv1=cv, cv2=cv).adjoint(): 4,
         Allocate(QBit()): 1,
         Free(QBit()): 1,
     }
@@ -76,10 +80,10 @@ def test_permutation_cycle_symbolic_call_graph():
     bloq = _permutation_cycle_symb()
     logN, L = ceil(log2(bloq.N)), slen(bloq.cycle)
 
-    _, sigma = bloq.call_graph()
+    _, sigma = bloq.call_graph(keep=lambda b: isinstance(b, And))
     assert sigma == {
-        ArbitraryClifford(n=2): (L + 1) * (13 * logN - 13),
-        TGate(): (L + 1) * (4 * logN - 4),
+        And(): (L + 1) * (logN - 1),
+        And().adjoint(): (L + 1) * (logN - 1),
         CNOT(): L * logN + L + 1,
     }
 
@@ -103,12 +107,12 @@ def test_permutation_unitary_and_call_graph():
         ),
     )
 
-    _, sigma = bloq.call_graph(generalizer=ignore_split_join)
+    _, sigma = bloq.call_graph(generalizer=ignore_split_join, keep=lambda b: isinstance(b, And))
     assert sigma == {
         CNOT(): 17,
-        TGate(): 56,
+        And(): 56 // 4,
+        And().adjoint(): 56 // 4,
         XGate(): 56,
-        ArbitraryClifford(n=2): 182,
         Allocate(QBit()): 2,
         Free(QBit()): 2,
     }
@@ -130,9 +134,9 @@ def test_permutation_symbolic_call_graph():
     logN = ceil(log2(N))
     bloq = _permutation_symb()
 
-    _, sigma = bloq.call_graph()
+    _, sigma = bloq.call_graph(keep=lambda b: isinstance(b, And))
     assert sigma == {
-        ArbitraryClifford(n=2): (N + 1) * (13 * logN - 13),
-        TGate(): (N + 1) * (4 * logN - 4),
+        And().adjoint(): (N + 1) * (logN - 1),
+        And(): (N + 1) * (logN - 1),
         CNOT(): N * logN + N + 1,
     }

qualtran/bloqs/chemistry/df/double_factorization_test.py

@@ -25,6 +25,7 @@
 from qualtran.bloqs.state_preparation.prepare_uniform_superposition import (
     PrepareUniformSuperposition,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.testing import execute_notebook
 
 
@@ -71,7 +72,7 @@ def test_compare_cost_to_openfermion():
         num_bits_rot_aa_inner=7,
         num_bits_rot=num_bits_rot,
     )
-    _, counts = bloq.call_graph()
+    t_counts = get_cost_value(bloq, QECGatesCost()).total_t_count()
     # https://github.com/quantumlib/OpenFermion/issues/839
     of_cost = compute_cost(
         num_spin_orb, lambd, 1e-3, num_aux, num_eig, num_bits_state_prep, num_bits_rot, 10_000
@@ -83,16 +84,14 @@ def test_compare_cost_to_openfermion():
     prog_rot_qrom_diff = 60
     missing_toffoli = 4  # need one more toffoli for second application of CZ
     swap_cost = 4 * (7 - 4) * num_spin_orb // 2
-    qual_cost = (
-        counts[TGate()] - inner_prep_qrom_diff - prog_rot_qrom_diff + missing_toffoli - swap_cost
-    )
+    qual_cost = t_counts - inner_prep_qrom_diff - prog_rot_qrom_diff + missing_toffoli - swap_cost
     # correct the expected cost by using a different uniform superposition algorithm
     # see: https://github.com/quantumlib/Qualtran/issues/611
     eta = power_two(num_aux + 1)
     cost1a = 4 * 2 * (3 * nl - 3 * eta + 2 * 7 - 9)
     prep = PrepareUniformSuperposition(num_aux + 1)
-    cost1a_mod = prep.call_graph()[1][TGate()]
-    cost1a_mod += prep.adjoint().call_graph()[1][TGate()]
+    cost1a_mod = get_cost_value(prep, QECGatesCost()).total_t_count()
+    cost1a_mod += get_cost_value(prep.adjoint(), QECGatesCost()).total_t_count()
     delta_uni_prep = cost1a_mod - cost1a
     qual_cost -= delta_uni_prep
     inner_refl = num_bits_state_prep + 1

qualtran/bloqs/chemistry/df/prepare_test.py

@@ -27,6 +27,7 @@
 from qualtran.bloqs.state_preparation.prepare_uniform_superposition import (
     PrepareUniformSuperposition,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prep_inner(bloq_autotester):
@@ -48,13 +49,11 @@ def test_outerprep_t_counts():
     outer_prep = OuterPrepareDoubleFactorization(
         num_aux, num_bits_state_prep=num_bits_state_prep, num_bits_rot_aa=num_bits_rot_aa
     )
-    _, counts = outer_prep.call_graph()
-    toff = counts[TGate()] // 4
+    toff = get_cost_value(outer_prep, QECGatesCost()).total_t_and_ccz_count()['n_ccz']
     outer_prep = OuterPrepareDoubleFactorization(
         num_aux, num_bits_state_prep=num_bits_state_prep, num_bits_rot_aa=num_bits_rot_aa
     ).adjoint()
-    _, counts = outer_prep.call_graph()
-    toff += counts[TGate()] // 4
+    toff += get_cost_value(outer_prep, QECGatesCost()).total_t_and_ccz_count()['n_ccz']
     # The output size for the QROM for the first state preparation in Eq. (C27)
     eta = power_two(num_aux + 1)
     nl = num_aux.bit_length()
@@ -65,8 +64,8 @@ def test_outerprep_t_counts():
     # correct the expected cost by using a different uniform superposition algorithm
     # https://github.com/quantumlib/Qualtran/issues/611
     prep = PrepareUniformSuperposition(num_aux + 1)
-    cost1a_mod = prep.call_graph()[1][TGate()] // 4
-    cost1a_mod += prep.adjoint().call_graph()[1][TGate()] // 4
+    cost1a_mod = get_cost_value(prep, QECGatesCost()).total_t_and_ccz_count()['n_ccz']
+    cost1a_mod += get_cost_value(prep.adjoint(), QECGatesCost()).total_t_and_ccz_count()['n_ccz']
     assert cost1a != cost1a_mod
     assert toff == cost1a_mod + cost1b + cost1cd
 
@@ -107,17 +106,15 @@ def test_inner_prepare_t_counts():
         num_bits_rot_aa=num_bits_rot_aa,
         num_bits_state_prep=num_bits_state_prep,
     )
-    _, counts = in_prep.call_graph()
-    toff = counts[TGate()] // 4
+    toff = get_cost_value(in_prep, QECGatesCost()).total_t_and_ccz_count()['n_ccz']
     in_prep = InnerPrepareDoubleFactorization(
         num_aux=num_aux,
         num_spin_orb=num_spin_orb,
         num_eig=num_eig,
         num_bits_rot_aa=num_bits_rot_aa,
         num_bits_state_prep=num_bits_state_prep,
     ).adjoint()
-    _, counts = in_prep.call_graph()
-    toff += counts[TGate()] // 4
+    toff += get_cost_value(in_prep, QECGatesCost()).total_t_and_ccz_count()['n_ccz']
     toff *= 2  # cost is for the two applications of the in-prep, in-prep^
     # application of ciruit.
     # captured from cost3 in openfermion df.compute_cost

qualtran/bloqs/chemistry/pbc/first_quantization/prepare_nu_test.py

@@ -12,8 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_nu import PrepareNuState
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.testing import assert_valid_bloq_decomposition
 
 
@@ -37,11 +37,9 @@ def test_prepare_nu_t_counts():
     eq_90 = 3 * num_bits_p**2 + 15 * num_bits_p - 7 + 4 * num_bits_m * (num_bits_p + 1)
     assert expected_cost == eq_90 + 5
     prep = PrepareNuState(num_bits_p, m_param)
-    _, counts = prep.call_graph()
-    qual_cost = counts[TGate()]
+    qual_cost = get_cost_value(prep, QECGatesCost()).total_t_count()
     prep = PrepareNuState(num_bits_p, m_param).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[TGate()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_t_count()
     qual_cost //= 4
     comp_diff = 1
     assert qual_cost == expected_cost - comp_diff

qualtran/bloqs/chemistry/pbc/first_quantization/prepare_uv_test.py

@@ -14,11 +14,11 @@
 
 import numpy as np
 
-from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_uv import (
     _prepare_uv,
     PrepareUVFirstQuantization,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prepare_uv(bloq_autotester):
@@ -41,13 +41,11 @@ def test_prepare_uv_t_counts():
     prep = PrepareUVFirstQuantization(
         num_bits_p, eta, num_atoms, m_param, lambda_zeta, num_bits_nuc_pos
     )
-    _, counts = prep.call_graph()
-    qual_cost = counts[TGate()]
+    qual_cost = get_cost_value(prep, QECGatesCost()).total_t_count()
     prep = PrepareUVFirstQuantization(
         num_bits_p, eta, num_atoms, m_param, lambda_zeta, num_bits_nuc_pos
     ).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[TGate()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_t_count()
     qual_cost //= 4
     comp_diff = 1
     assert qual_cost == expected_cost - comp_diff

qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_nu_test.py

@@ -12,12 +12,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.prepare_nu import (
     _prep_mu_proj,
     _prep_nu_proj,
     PrepareNuStateWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prepare_num(bloq_autotester):
@@ -44,11 +44,9 @@ def test_prepare_nu_with_proj_t_counts():
     )
     assert expected_cost == eq_c6 + 5
     prep = PrepareNuStateWithProj(num_bits_p, num_bits_n, m_param)
-    _, counts = prep.call_graph()
-    qual_cost = counts[TGate()]
+    qual_cost = get_cost_value(prep, QECGatesCost()).total_t_count()
     prep = PrepareNuStateWithProj(num_bits_p, num_bits_n, m_param).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[TGate()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_t_count()
     qual_cost //= 4
     comp_diff = 1
     assert qual_cost == expected_cost - comp_diff

qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_uv_test.py

@@ -14,11 +14,11 @@
 
 import numpy as np
 
-from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.prepare_uv import (
     _prep_uv_proj,
     PrepareUVFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prep_uv_proj(bloq_autotester):
@@ -42,13 +42,11 @@ def test_prepare_uv_t_counts():
     prep = PrepareUVFirstQuantizationWithProj(
         num_bits_p, num_bits_n, eta, num_atoms, m_param, lambda_zeta, num_bits_nuc_pos
     )
-    _, counts = prep.call_graph()
-    qual_cost = counts[TGate()]
+    qual_cost = get_cost_value(prep, QECGatesCost()).total_t_count()
     prep = PrepareUVFirstQuantizationWithProj(
         num_bits_p, num_bits_n, eta, num_atoms, m_param, lambda_zeta, num_bits_nuc_pos
     ).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[TGate()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_t_count()
     qual_cost //= 4
     comp_diff = 1
     assert qual_cost == expected_cost - comp_diff

qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_and_prepare_test.py

@@ -22,6 +22,7 @@
     PrepareFirstQuantizationWithProj,
     SelectFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.testing import assert_valid_bloq_decomposition, execute_notebook
 
 
@@ -45,13 +46,12 @@ def test_select_t_costs():
     num_atoms = 10
     lambda_zeta = 10
     num_bits_nuc_pos = 41
-    cost = 0
 
     sel_first_quant = SelectFirstQuantizationWithProj(
         num_bits_p, num_bits_n, eta, num_atoms, lambda_zeta, num_bits_nuc_pos=num_bits_nuc_pos
     )
     assert_valid_bloq_decomposition(sel_first_quant)
-    cost += sel_first_quant.call_graph()[1][TGate()]
+    cost = get_cost_value(sel_first_quant, QECGatesCost()).total_t_count()
 
     # Swaps
     expected_cost = 7 * (12 * eta * num_bits_p + 6 * num_bits_n) + 4 * (4 * eta - 6)  #

qualtran/bloqs/chemistry/sparse/select_bloq_test.py

@@ -12,7 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.chemistry.sparse.select_bloq import _sel_sparse
 
 
@@ -22,6 +21,6 @@ def test_prep_inner(bloq_autotester):
 
 def test_decompose_bloq_counts():
     sel = _sel_sparse()
-    cost_decomp = sel.decompose_bloq().call_graph()[1][TGate()]
-    cost_call = sel.call_graph()[1][TGate()]
+    cost_decomp = sel.decompose_bloq().call_graph()[1]
+    cost_call = sel.call_graph()[1]
     assert cost_call == cost_decomp