Use QECGatesCost for phase_estimation_of_quantum_walk (#1383)

* Use QECGatesCost for phase_estimation_of_quantum_walk

* put bloq example back

Author: mpharrigan

qualtran/bloqs/phase_estimation/phase_estimation_of_quantum_walk.ipynb

@@ -7,27 +7,6 @@
     "# Phase Estimation of Quantum Walks"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#  Copyright 2023 Google LLC\n",
-    "#\n",
-    "#  Licensed under the Apache License, Version 2.0 (the \"License\");\n",
-    "#  you may not use this file except in compliance with the License.\n",
-    "#  You may obtain a copy of the License at\n",
-    "#\n",
-    "#      https://www.apache.org/licenses/LICENSE-2.0\n",
-    "#\n",
-    "#  Unless required by applicable law or agreed to in writing, software\n",
-    "#  distributed under the License is distributed on an \"AS IS\" BASIS,\n",
-    "#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
-    "#  See the License for the specific language governing permissions and\n",
-    "#  limitations under the License."
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -51,6 +30,15 @@
     "from qualtran.bloqs.chemistry.hubbard_model.qubitization import get_walk_operator_for_hubbard_model"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Phase estimation circuit\n",
+    "\n",
+    "We start by quickly sketching the phase estimation circuit in terms of the walk operator."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -129,6 +117,44 @@
     "print(circuit)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Computing costs\n",
+    "\n",
+    "Usually, you'd define a Bloq that captures the entire phase estimation circuit, but we can use a little helper function to compute the gate counts directly from the bloqs encountered within the Cirq circuit."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran import Bloq\n",
+    "from qualtran.resource_counting import get_cost_value, QECGatesCost, GateCounts\n",
+    "\n",
+    "def get_qec_gates_cost_for_circuit(circuit):\n",
+    "    # Usually, you'd combine this into a bloq of its own, but we\n",
+    "    # use this helper function to add up the costs of the bloqs\n",
+    "    # found within the circuit.\n",
+    "    cost_key = QECGatesCost()\n",
+    "    costs_cache = {}\n",
+    "    total_cost = cost_key.zero()\n",
+    "    for op in circuit.all_operations():\n",
+    "        if not isinstance(op.gate, Bloq):\n",
+    "            # Skip state prep and QFT for now\n",
+    "            continue\n",
+    "    \n",
+    "        bloq = op.gate\n",
+    "        total_cost += get_cost_value(bloq, cost_key, costs_cache=costs_cache)\n",
+    "    \n",
+    "    return total_cost\n",
+    "\n",
+    "get_qec_gates_cost_for_circuit(circuit)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -151,7 +177,7 @@
     "walk, _ = get_walk_operator_for_1d_ising_model(num_sites, eps)\n",
     "\n",
     "circuit = cirq.Circuit(phase_estimation(walk, m=m_bits))\n",
-    "%time result = t_complexity_compat(circuit[1:-1])\n",
+    "%time result = get_qec_gates_cost_for_circuit(circuit)\n",
     "print(result)"
    ]
   },
@@ -178,7 +204,7 @@
     "m_bits = int(np.ceil(np.log2(qlambda * np.pi * np.sqrt(2) / delta_E)))\n",
     "walk = get_walk_operator_for_hubbard_model(x_dim, y_dim, t, mu)\n",
     "circuit = cirq.Circuit(phase_estimation(walk, m=m_bits))\n",
-    "%time result = t_complexity_compat(circuit[1:-1])\n",
+    "%time result = get_qec_gates_cost_for_circuit(circuit)\n",
     "print(result)"
    ]
   },
@@ -188,9 +214,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Or, we can just use the included bloq example directly\n",
+    "\n",
     "from qualtran.bloqs.phase_estimation.qubitization_qpe import _qubitization_qpe_hubbard_model_large\n",
     "qpe = _qubitization_qpe_hubbard_model_large.make()\n",
-    "%time result = qpe.t_complexity()\n",
+    "%time result = get_cost_value(qpe, QECGatesCost())\n",
     "print(result)"
    ]
   },
@@ -211,9 +239,7 @@
     "from qualtran.drawing import show_flame_graph\n",
     "\n",
     "qpe_small = _qubitization_qpe_hubbard_model_small.make()\n",
-    "\n",
     "print(qpe_small.t_complexity())\n",
-    "\n",
     "show_flame_graph(qpe_small)"
    ]
   }