First pass at fixing chemistry costs to use QECGatesCost

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -66,6 +66,7 @@
 import qualtran.bloqs.chemistry.hubbard_model.qubitization
 import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_t
 import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_uv
+import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta
 import qualtran.bloqs.chemistry.pbc.first_quantization.projectile.select_and_prepare
 import qualtran.bloqs.chemistry.pbc.first_quantization.select_t
 import qualtran.bloqs.chemistry.pbc.first_quantization.select_uv
@@ -311,6 +312,7 @@
             qualtran.bloqs.chemistry.pbc.first_quantization.prepare_uv._PREPARE_UV,
             qualtran.bloqs.chemistry.pbc.first_quantization.select_t._SELECT_T,
             qualtran.bloqs.chemistry.pbc.first_quantization.select_uv._SELECT_UV,
+            qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta._PREPARE_ZETA,
         ],
         directory=f'{SOURCE_DIR}/bloqs/chemistry/pbc/first_quantization',
     ),

qualtran/bloqs/chemistry/pbc/first_quantization/first_quantization.ipynb

@@ -822,6 +822,119 @@
     "show_call_graph(select_uv_g)\n",
     "show_counts_sigma(select_uv_sigma)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "69c98b9f",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.bloq_doc.md"
+   },
+   "source": [
+    "## `PrepareZetaState`\n",
+    "PREPARE the superpostion over $l$ weighted by $\\zeta_l$.\n",
+    "\n",
+    "See https://github.com/quantumlib/Qualtran/issues/473.\n",
+    "#### Parameters\n",
+    " - `num_bits_p`: The number of bits to represent each dimension of the momentum register.\n",
+    " - `eta`: The number of electrons.\n",
+    " - `m_param`: $\\mathcal{M}$ in the reference.\n",
+    " - `lambda_zeta`: sum of nuclear charges. \n",
+    "\n",
+    "#### Registers\n",
+    " - `l`: the register indexing the atomic number. \n",
+    "\n",
+    "#### References\n",
+    " - [Fault-Tolerant Quantum Simulations of Chemistry in First Quantization](https://arxiv.org/abs/2105.12767). page 23-24, last 3 paragraphs.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "074bc1b8",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd3071f7",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31ee49a4",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.prepare_zeta"
+   },
+   "outputs": [],
+   "source": [
+    "num_atoms = 10\n",
+    "lambda_zeta = 10\n",
+    "num_bits_nuc_pos = 8\n",
+    "\n",
+    "prepare_zeta = PrepareZetaState(\n",
+    "    num_atoms=num_atoms, lambda_zeta=lambda_zeta, num_bits_nuc_pos=num_bits_nuc_pos\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4fb19c71",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6498eee0",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([prepare_zeta],\n",
+    "           ['`prepare_zeta`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1e644753",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbef8e6e",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "prepare_zeta_g, prepare_zeta_sigma = prepare_zeta.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(prepare_zeta_g)\n",
+    "show_counts_sigma(prepare_zeta_sigma)"
+   ]
   }
  ],
  "metadata": {

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

@@ -12,14 +12,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare import (
     UniformSuperpostionIJFirstQuantization,
 )
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_t import (
     _prepare_t,
     PrepareTFirstQuantization,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prepare_t(bloq_autotester):
@@ -33,15 +33,12 @@ def test_prepare_kinetic_t_counts():
     n_eta = (eta - 1).bit_length()
     expected_cost = (14 * n_eta + 8 * b_r - 36) + 2 * (2 * num_bits_p + 9)
     uni = UniformSuperpostionIJFirstQuantization(eta, num_bits_rot_aa=b_r)
-    _, counts = uni.call_graph()
-    qual_cost = counts[Toffoli()]
+
+    qual_cost = get_cost_value(uni, QECGatesCost()).total_toffoli_only()
     uni = UniformSuperpostionIJFirstQuantization(eta, num_bits_rot_aa=b_r).adjoint()
-    _, counts = uni.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(uni, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantization(num_bits_p, eta, num_bits_rot_aa=b_r)
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantization(num_bits_p, eta, num_bits_rot_aa=b_r).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     assert qual_cost == expected_cost

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

@@ -19,7 +19,7 @@
 import numpy as np
 from attrs import evolve, frozen
 
-from qualtran import Bloq, QAny, Register, Signature
+from qualtran import Bloq, bloq_example, BloqDocSpec, QAny, Register, Signature
 from qualtran.bloqs.basic_gates import Toffoli
 
 if TYPE_CHECKING:
@@ -65,3 +65,22 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         else:
 
             return {Toffoli(): self.lambda_zeta}
+
+
+@bloq_example
+def _prepare_zeta() -> PrepareZetaState:
+    num_atoms = 10
+    lambda_zeta = 10
+    num_bits_nuc_pos = 8
+
+    prepare_zeta = PrepareZetaState(
+        num_atoms=num_atoms, lambda_zeta=lambda_zeta, num_bits_nuc_pos=num_bits_nuc_pos
+    )
+    return prepare_zeta
+
+
+_PREPARE_ZETA = BloqDocSpec(
+    bloq_cls=PrepareZetaState,
+    import_line='from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState',
+    examples=(_prepare_zeta,),
+)

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

@@ -11,8 +11,8 @@
 #  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.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState
+from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import _prepare_zeta
 
 
-def test_uniform_superposition_ij():
-    prep = PrepareZetaState(num_atoms=10, lambda_zeta=20, num_bits_nuc_pos=8)
+def test_prepare_zeta(bloq_autotester):
+    bloq_autotester(_prepare_zeta)

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

@@ -14,13 +14,13 @@
 
 import pytest
 
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.prepare_t import (
     _prep_power_two_proj,
     _prep_t_proj,
     PreparePowerTwoStateWithProj,
     PrepareTFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prep_t_proj(bloq_autotester):
@@ -39,13 +39,11 @@ def test_prepare_kinetic_t_proj_counts():
     expected_cost = 2 * (2 * num_bits_n + 9) + 2 * (num_bits_n - num_bits_p) + 20
     qual_cost = 0
     prep = PrepareTFirstQuantizationWithProj(num_bits_p, num_bits_n, eta, num_bits_rot_aa=b_r)
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantizationWithProj(
         num_bits_p, num_bits_n, eta, num_bits_rot_aa=b_r
     ).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     assert qual_cost == expected_cost
 
 

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

@@ -11,11 +11,11 @@
 #  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.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.select_t import (
     _sel_t_proj,
     SelectTFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_sel_t_proj(bloq_autotester):
@@ -25,5 +25,5 @@ def test_sel_t_proj(bloq_autotester):
 def test_select_kinetic_t_counts():
     num_bits_n = 6
     sel = SelectTFirstQuantizationWithProj(num_bits_n, 10)
-    _, counts = sel.call_graph()
-    assert counts[Toffoli()] == 5 * (num_bits_n - 1) + 2 + 1
+    toffolis = get_cost_value(sel, QECGatesCost()).total_toffoli_only()
+    assert toffolis == 5 * (num_bits_n - 1) + 2 + 1

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

@@ -11,11 +11,11 @@
 #  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.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.select_t import (
     _select_t,
     SelectTFirstQuantization,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_select_t(bloq_autotester):
@@ -25,5 +25,5 @@ def test_select_t(bloq_autotester):
 def test_select_kinetic_t_counts():
     num_bits_p = 6
     sel = SelectTFirstQuantization(num_bits_p, 10)
-    _, counts = sel.call_graph()
-    assert counts[Toffoli()] == 5 * (num_bits_p - 1) + 2
+    toffolis = get_cost_value(sel, QECGatesCost()).total_toffoli_only()
+    assert toffolis == 5 * (num_bits_p - 1) + 2

qualtran/bloqs/chemistry/resource_estimation.ipynb

@@ -84,7 +84,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "2fb17f7f",
    "metadata": {},
    "outputs": [],
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "b9991178",
    "metadata": {},
    "outputs": [],
@@ -202,7 +202,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "5efd66ae",
    "metadata": {},
    "outputs": [],
@@ -305,7 +305,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "614bfb21",
    "metadata": {},
    "outputs": [],
@@ -346,9 +346,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "_, sigma = df_bloq.call_graph()\n",
     "refl_cost_and_qpe = (df_bloq.num_aux -1).bit_length() + (num_spin_orb // 2 - 1).bit_length() + num_bits_state_prep + 1 + 2\n",
-    "print(f'qualtran cost = {sigma[TGate()] // 4} vs paper = {21753 - refl_cost_and_qpe}')"
+    "print(f'qualtran cost = {get_toffoli_counts(df_bloq)} vs paper = {21753 - refl_cost_and_qpe}')"
    ]
   },
   {
@@ -363,7 +362,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "id": "349b70ab",
    "metadata": {},
    "outputs": [],
@@ -384,7 +383,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "id": "2ef506ab",
    "metadata": {},
    "outputs": [],
@@ -411,7 +410,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "id": "e34dcb1c",
    "metadata": {},
    "outputs": [],

qualtran/bloqs/chemistry/trotter/grid_ham/inverse_sqrt_test.py

@@ -41,7 +41,6 @@ def test_newton_raphson_inverse_sqrt_bloq_counts():
     poly_bitsize = 15
     target_bitsize = 22
     bloq = NewtonRaphsonApproxInverseSquareRoot(int_bitsize, poly_bitsize, target_bitsize)
-    _, counts = bloq.call_graph()
     cost_square = poly_bitsize**2 // 2 - 4
     cost_scale = poly_bitsize * (2 * int_bitsize - 1) - int_bitsize**2
     cost_mult = 2 * (target_bitsize**2 - target_bitsize - 1)