From 188b663545cf63e88071e391bf7c3ddbba92ce07 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <7265746+anurudhp@users.noreply.github.com>
Date: Thu, 7 Nov 2024 12:33:11 -0800
Subject: [PATCH] fix QPE test failure, add non-slow bloq example autotest
 (#1487)

* fix QPE test failure, add non-slow bloq example autotest

* fix bug: `cirq.ControlledGate` fails when it can't detect that bloq is unitary

* regen notebook

* add comment on controlled unitary failure
---
 qualtran/_infra/controlled.py                 | 16 +++++++++++---
 .../phase_estimation/qubitization_qpe.ipynb   | 21 ++++++++++++++++---
 .../phase_estimation/qubitization_qpe.py      | 14 +++++++++++--
 .../phase_estimation/qubitization_qpe_test.py | 10 +++++++++
 qualtran/conftest.py                          |  1 +
 5 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/qualtran/_infra/controlled.py b/qualtran/_infra/controlled.py
index f2f946030..dca518e93 100644
--- a/qualtran/_infra/controlled.py
+++ b/qualtran/_infra/controlled.py
@@ -458,9 +458,19 @@ def _unitary_(self):
             # subbloq is a cirq gate, use the cirq-style API to derive a unitary.
             import cirq
 
-            return cirq.unitary(
-                cirq.ControlledGate(self.subbloq, control_values=self.ctrl_spec.to_cirq_cv())
-            )
+            # TODO It would be ideal to use `tensor_contract` always,
+            #      but at the moment it's about 5-10x slower than `cirq.unitary`.
+            #      So we default to `cirq.unitary`, and only use `tensor_contract` if it fails.
+            #      https://github.com/quantumlib/Qualtran/issues/1336
+            # TODO `cirq.ControlledGate` fails to correctly verify `subbloq` using
+            #      a compute-uncompute `And` pair is unitary.
+            #      https://github.com/quantumlib/Qualtran/issues/1488
+            try:
+                return cirq.unitary(
+                    cirq.ControlledGate(self.subbloq, control_values=self.ctrl_spec.to_cirq_cv())
+                )
+            except ValueError:
+                pass  # use the tensor contraction instead
         if all(reg.side == Side.THRU for reg in self.subbloq.signature):
             # subbloq has only THRU registers, so the tensor contraction corresponds
             # to a unitary matrix.
diff --git a/qualtran/bloqs/phase_estimation/qubitization_qpe.ipynb b/qualtran/bloqs/phase_estimation/qubitization_qpe.ipynb
index 92ab548b9..b99e4b7c6 100644
--- a/qualtran/bloqs/phase_estimation/qubitization_qpe.ipynb
+++ b/qualtran/bloqs/phase_estimation/qubitization_qpe.ipynb
@@ -61,7 +61,6 @@
     "\n",
     "#### Parameters\n",
     " - `walk`: Bloq representing the Qubitization walk operator to run the phase estimation protocol on.\n",
-    " - `m_bits`: Bitsize of the phase register to be used during phase estimation.\n",
     " - `ctrl_state_prep`: Bloq to prepare the control state on the phase register. Defaults to `OnEach(self.m_bits, Hadamard())`.\n",
     " - `qft_inv`: Bloq to apply inverse QFT on the phase register. Defaults to `QFTTextBook(self.m_bits).adjoint()`  \n",
     "\n",
@@ -193,6 +192,22 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "49743657",
+   "metadata": {
+    "cq.autogen": "QubitizationQPE.qubitization_qpe_ising"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.chemistry.ising.walk_operator import get_walk_operator_for_1d_ising_model\n",
+    "from qualtran.bloqs.phase_estimation import RectangularWindowState\n",
+    "\n",
+    "walk, _ = get_walk_operator_for_1d_ising_model(4, 0.1)\n",
+    "qubitization_qpe_ising = QubitizationQPE(walk, RectangularWindowState(4))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "b19f9365",
@@ -213,8 +228,8 @@
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([qubitization_qpe_hubbard_model_small, qubitization_qpe_sparse_chem, qubitization_qpe_chem_thc],\n",
-    "           ['`qubitization_qpe_hubbard_model_small`', '`qubitization_qpe_sparse_chem`', '`qubitization_qpe_chem_thc`'])"
+    "show_bloqs([qubitization_qpe_hubbard_model_small, qubitization_qpe_sparse_chem, qubitization_qpe_chem_thc, qubitization_qpe_ising],\n",
+    "           ['`qubitization_qpe_hubbard_model_small`', '`qubitization_qpe_sparse_chem`', '`qubitization_qpe_chem_thc`', '`qubitization_qpe_ising`'])"
    ]
   },
   {
diff --git a/qualtran/bloqs/phase_estimation/qubitization_qpe.py b/qualtran/bloqs/phase_estimation/qubitization_qpe.py
index 9cc1b3a8b..39b608ea5 100644
--- a/qualtran/bloqs/phase_estimation/qubitization_qpe.py
+++ b/qualtran/bloqs/phase_estimation/qubitization_qpe.py
@@ -64,7 +64,6 @@ class QubitizationQPE(GateWithRegisters):
     Args:
         walk: Bloq representing the Qubitization walk operator to run the phase estimation protocol
             on.
-        m_bits: Bitsize of the phase register to be used during phase estimation.
         ctrl_state_prep: Bloq to prepare the control state on the phase register. Defaults to
             `OnEach(self.m_bits, Hadamard())`.
         qft_inv: Bloq to apply inverse QFT on the phase register. Defaults to
@@ -119,7 +118,7 @@ def decompose_from_registers(
         qpre_reg = quregs['qpe_reg']
 
         yield self.ctrl_state_prep.on(*qpre_reg)
-        yield walk_controlled.on_registers(**walk_regs, control=qpre_reg[-1])
+        yield walk_controlled.on_registers(**walk_regs, ctrl=qpre_reg[-1])
         walk = self.walk**2
         for i in range(self.m_bits - 2, -1, -1):
             yield reflect_controlled.on_registers(control=qpre_reg[i], **reflect_regs)
@@ -143,6 +142,16 @@ def __str__(self) -> str:
         return f'QubitizationQPE[{self.m_bits}]'
 
 
+@bloq_example
+def _qubitization_qpe_ising() -> QubitizationQPE:
+    from qualtran.bloqs.chemistry.ising.walk_operator import get_walk_operator_for_1d_ising_model
+    from qualtran.bloqs.phase_estimation import RectangularWindowState
+
+    walk, _ = get_walk_operator_for_1d_ising_model(4, 0.1)
+    qubitization_qpe_ising = QubitizationQPE(walk, RectangularWindowState(4))
+    return qubitization_qpe_ising
+
+
 @bloq_example
 def _qubitization_qpe_hubbard_model_small() -> QubitizationQPE:
     import numpy as np
@@ -252,5 +261,6 @@ def _qubitization_qpe_sparse_chem() -> QubitizationQPE:
         _qubitization_qpe_hubbard_model_small,
         _qubitization_qpe_sparse_chem,
         _qubitization_qpe_chem_thc,
+        _qubitization_qpe_ising,
     ),
 )
diff --git a/qualtran/bloqs/phase_estimation/qubitization_qpe_test.py b/qualtran/bloqs/phase_estimation/qubitization_qpe_test.py
index 2888ef7da..2d6ae4fef 100644
--- a/qualtran/bloqs/phase_estimation/qubitization_qpe_test.py
+++ b/qualtran/bloqs/phase_estimation/qubitization_qpe_test.py
@@ -21,6 +21,7 @@
 from qualtran.bloqs.phase_estimation.qubitization_qpe import (
     _qubitization_qpe_chem_thc,
     _qubitization_qpe_hubbard_model_small,
+    _qubitization_qpe_ising,
     _qubitization_qpe_sparse_chem,
     QubitizationQPE,
 )
@@ -29,6 +30,10 @@
 from qualtran.testing import execute_notebook
 
 
+def test_ising_example(bloq_autotester):
+    bloq_autotester(_qubitization_qpe_ising)
+
+
 @pytest.mark.slow
 def test_qubitization_qpe_bloq_autotester(bloq_autotester):
     bloq_autotester(_qubitization_qpe_hubbard_model_small)
@@ -103,3 +108,8 @@ def test_qubitization_phase_estimation_of_walk(num_terms: int, use_resource_stat
 @pytest.mark.notebook
 def test_phase_estimation_of_qubitized_hubbard_model():
     execute_notebook('phase_estimation_of_quantum_walk')
+
+
+@pytest.mark.notebook
+def test_notebook():
+    execute_notebook('qubitization_qpe')
diff --git a/qualtran/conftest.py b/qualtran/conftest.py
index ede967863..ac6137d71 100644
--- a/qualtran/conftest.py
+++ b/qualtran/conftest.py
@@ -88,6 +88,7 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
         'qubitization_qpe_chem_thc',  # too slow
         'walk_op_chem_sparse',
         'qubitization_qpe_sparse_chem',  # too slow
+        'qubitization_qpe_ising',
         'trott_unitary',
         'symbolic_hamsim_by_gqsp',
         'gf16_addition',  # cannot serialize QGF