Introduce types to t_complexity (#835)

* Typed T complexity

* extensive testing and diff reduction

* fix notebooks

* fix notebook

* regenerate notebook

qualtran/bloqs/arithmetic/t_complexity_of_comparison_gates.ipynb

@@ -78,7 +78,7 @@
     "from typing import Sequence\n",
     "\n",
     "import cirq\n",
-    "from qualtran.cirq_interop.t_complexity_protocol import t_complexity\n",
+    "from qualtran.cirq_interop.t_complexity_protocol import t_complexity_compat\n",
     "from qualtran.bloqs.mcmt import And, MultiAnd"
    ]
   },
@@ -135,7 +135,7 @@
    "outputs": [],
    "source": [
     "# This should have T count of 4*(3 qubits) - 4 = 8\n",
-    "t_complexity(equality_circuit)"
+    "t_complexity_compat(equality_circuit)"
    ]
   },
   {
@@ -302,7 +302,7 @@
    "outputs": [],
    "source": [
     "# This should have T count of 4*(2 qubits) - 4 = 4\n",
-    "t_complexity(quantum_quantum_equality)"
+    "t_complexity_compat(quantum_quantum_equality)"
    ]
   },
   {
@@ -679,7 +679,7 @@
    "outputs": [],
    "source": [
     "# This should have T count of 8*(3 qubits) - 4 = 20\n",
-    "t_complexity(quantum_compare)"
+    "t_complexity_compat(quantum_compare)"
    ]
   },
   {
@@ -820,7 +820,7 @@
    "outputs": [],
    "source": [
     "# This should have T complexity of 6*(2 qubits) + O(1)\n",
-    "t_complexity(circuit)"
+    "t_complexity_compat(circuit)"
    ]
   },
   {
@@ -959,7 +959,7 @@
    "outputs": [],
    "source": [
     "# T count should be 4*(3 qubits) = 12\n",
-    "t_complexity(less_than_circuit)"
+    "t_complexity_compat(less_than_circuit)"
    ]
   },
   {
@@ -1037,9 +1037,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }

qualtran/bloqs/phase_estimation/phase_estimation_of_quantum_walk.ipynb

@@ -142,7 +142,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from qualtran.cirq_interop.t_complexity_protocol import t_complexity\n",
+    "from qualtran.cirq_interop.t_complexity_protocol import t_complexity_compat\n",
     "\n",
     "num_sites: int = 200\n",
     "eps: float = 1e-5\n",
@@ -151,7 +151,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(circuit[1:-1])\n",
+    "%time result = t_complexity_compat(circuit[1:-1])\n",
     "print(result)"
    ]
   },
@@ -178,8 +178,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",
-    "t_complexity.cache_clear()\n",
-    "%time result = t_complexity(circuit[1:-1])\n",
+    "%time result = t_complexity_compat(circuit[1:-1])\n",
     "print(result)"
    ]
   },
@@ -191,7 +190,6 @@
    "source": [
     "from qualtran.bloqs.phase_estimation.qubitization_qpe import _qubitization_qpe_hubbard_model_large\n",
     "qpe = _qubitization_qpe_hubbard_model_large.make()\n",
-    "t_complexity.cache_clear()\n",
     "%time result = qpe.t_complexity()\n",
     "print(result)"
    ]

qualtran/bloqs/rotations/hamming_weight_phasing.ipynb

@@ -108,7 +108,7 @@
    "outputs": [],
    "source": [
     "hamming_weight_phasing = HammingWeightPhasing(4, np.pi / 2.0)\n",
-    "print(\"Applying this unitary to |1111> should be the identity, and |0101> will flip the sign.\")"
+    "# Applying this unitary to |1111> should be the identity, and |0101> will flip the sign."
    ]
   },
   {

qualtran/bloqs/rotations/hamming_weight_phasing.py

@@ -110,7 +110,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
 @bloq_example
 def _hamming_weight_phasing() -> HammingWeightPhasing:
     hamming_weight_phasing = HammingWeightPhasing(4, np.pi / 2.0)
-    print("Applying this unitary to |1111> should be the identity, and |0101> will flip the sign.")
+    # Applying this unitary to |1111> should be the identity, and |0101> will flip the sign.
     return hamming_weight_phasing
 
 

qualtran/cirq_interop/_bloq_to_cirq_test.py

@@ -25,7 +25,7 @@
 from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd
 from qualtran.bloqs.state_preparation import PrepareUniformSuperposition
 from qualtran.cirq_interop._bloq_to_cirq import BloqAsCirqGate, CirqQuregT
-from qualtran.cirq_interop.t_complexity_protocol import t_complexity
+from qualtran.cirq_interop.t_complexity_protocol import t_complexity_compat
 from qualtran.testing import execute_notebook
 
 if TYPE_CHECKING:
@@ -225,7 +225,7 @@ def test_bloq_as_cirq_gate_for_mod_exp():
     # newly allocated RIGHT registers in the decomposition to the one's specified by the user
     # when constructing the original operation (in this case, register `x`).
     circuit = cirq.Circuit(op, cirq.decompose_once(op))
-    assert t_complexity(circuit) == 2 * mod_exp.t_complexity()
+    assert t_complexity_compat(circuit) == 2 * mod_exp.t_complexity()
     cirq.testing.assert_has_diagram(
         circuit,
         '''
@@ -250,7 +250,7 @@ def test_bloq_as_cirq_gate_for_mod_exp():
     # Whereas when directly applying a cirq gate on qubits to get an operations, we need to
     # specify both input and output registers.
     circuit = cirq.Circuit(gate.on_registers(**out_regs), decomposed_circuit)
-    assert t_complexity(circuit) == 2 * mod_exp.t_complexity()
+    assert t_complexity_compat(circuit) == 2 * mod_exp.t_complexity()
     # Notice the newly allocated qubits _C(0) and _C(1) for output register x.
     cirq.testing.assert_has_diagram(
         circuit,

qualtran/cirq_interop/_cirq_to_bloq.py

@@ -48,7 +48,7 @@
     split_qubits,
 )
 from qualtran.cirq_interop._interop_qubit_manager import InteropQubitManager
-from qualtran.cirq_interop.t_complexity_protocol import t_complexity, TComplexity
+from qualtran.cirq_interop.t_complexity_protocol import _from_directly_countable_cirq, TComplexity
 
 if TYPE_CHECKING:
     import quimb.tensor as qtn
@@ -110,7 +110,10 @@ def my_tensors(
         )
 
     def _t_complexity_(self) -> 'TComplexity':
-        return t_complexity(self.cirq_gate)
+        t_count = _from_directly_countable_cirq(self.cirq_gate)
+        if t_count is None:
+            raise ValueError(f"Cirq gate must be directly countable, not {self.cirq_gate}")
+        return t_count
 
     def as_cirq_op(
         self, qubit_manager: 'cirq.QubitManager', **in_quregs: 'CirqQuregT'
@@ -335,6 +338,7 @@ def cirq_gate_to_bloq(gate: cirq.Gate) -> Bloq:
         CZPowGate,
         GlobalPhase,
         Hadamard,
+        Identity,
         Rx,
         Ry,
         Rz,
@@ -380,6 +384,7 @@ def cirq_gate_to_bloq(gate: cirq.Gate) -> Bloq:
         cirq.CZ: CZ(),
         cirq.SWAP: TwoBitSwap(),
         cirq.CSWAP: CSwap(1),
+        cirq.I: Identity(),
     }
     if gate in CIRQ_GATE_TO_BLOQ_MAP:
         return CIRQ_GATE_TO_BLOQ_MAP[gate]

qualtran/cirq_interop/jupyter_tools.py

@@ -62,7 +62,7 @@ def circuit_with_costs(circuit: 'cirq.AbstractCircuit') -> 'cirq.AbstractCircuit
     """Annotates each operation in the circuit with its T-complexity cost."""
 
     def _map_func(op: cirq.Operation, _):
-        t_cost = t_complexity_protocol.t_complexity(op)
+        t_cost = t_complexity_protocol.t_complexity_compat(op)
         return op.with_tags(f't:{t_cost.t:g},r:{t_cost.rotations:g}')
 
     return cirq.map_operations(circuit, map_func=_map_func)

qualtran/cirq_interop/t_complexity.ipynb

@@ -62,10 +62,8 @@
    "source": [
     "# And of two qubits\n",
     "gate = And() # create an And gate\n",
-    "# create an operation\n",
-    "operation = gate.on_registers(**get_named_qubits(gate.signature))\n",
-    "# this operation doesn't directly support TComplexity but it's decomposable and its components are simple.\n",
-    "print(t_complexity(operation))"
+    "# this gate doesn't directly support TComplexity but it's decomposable and its components are simple.\n",
+    "print(t_complexity(gate))"
    ]
   },
   {
@@ -84,9 +82,8 @@
    "outputs": [],
    "source": [
     "gate = And() ** -1 # adjoint of And\n",
-    "operation = gate.on_registers(**get_named_qubits(gate.signature))\n",
     "# the deomposition is H, measure, CZ, and Reset\n",
-    "print(t_complexity(operation))"
+    "print(t_complexity(gate))"
    ]
   },
   {
@@ -105,9 +102,8 @@
    "outputs": [],
    "source": [
     "n = 5\n",
-    "gate = MultiAnd((1, )*n)\n",
-    "operation = gate.on_registers(**get_named_qubits(gate.signature))\n",
-    "print(t_complexity(operation))"
+    "gate = MultiAnd((1,)*n)\n",
+    "print(t_complexity(gate))"
    ]
   },
   {
@@ -117,15 +113,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def Generate(n_max: int = 10):\n",
+    "def get_t_count_plot_data_for_and(n_max: int = 10):\n",
     "    \"\"\"Returns the #T when the number of qubits is between 3 and n_max inclusive\"\"\"\n",
     "    n_controls = []\n",
     "    t_count = []\n",
     "    for n in range(3, n_max + 2):\n",
     "        n_controls.append(n)\n",
     "        gate = MultiAnd(cvs=(1, )*n)\n",
-    "        op = gate.on_registers(**get_named_qubits(gate.signature))\n",
-    "        c = t_complexity(op)\n",
+    "        c = t_complexity(gate)\n",
     "        t_count.append(c.t)\n",
     "    return n_controls, t_count"
    ]
@@ -137,8 +132,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "n_controls, t_count = Generate()\n",
-    "plt.plot(n_controls, t_count, label='T')\n",
+    "n_controls, t_count = get_t_count_plot_data_for_and()\n",
+    "plt.plot(n_controls, t_count, 'o-', label='T')\n",
     "plt.ylabel('count')\n",
     "plt.xlabel('number of qubits')\n",
     "plt.title('And gate')\n",
@@ -163,7 +158,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.8"
   },
   "vscode": {
    "interpreter": {