Add HWP Trotter bloqs and costs to notebook. (#946)

* Small fixes to hubbard bloqs.

* Add cost notebook.

* Fix eps type.

* Add hamming weight phasing bloqs.

* Add costs to notebook.

* Add costs to notebook.

* Add hwp trotter factory.

* Fix up notebook.

* Fix formatting.

* Fix mypy errors.

* Fix serialization.

* Use symbolicfloat.

* Fix bug.

* Add notebook test.

* Add missing semicolon.

* Fix test.

---------

Co-authored-by: Matthew Harrigan <[email protected]>

qualtran/bloqs/chemistry/trotter/hubbard/hopping.py

@@ -15,12 +15,13 @@
 from functools import cached_property
 from typing import Set, TYPE_CHECKING, Union
 
-import sympy
 from attrs import frozen
 
 from qualtran import Bloq, bloq_example, BloqDocSpec, QAny, QBit, Register, Signature
 from qualtran.bloqs.basic_gates import Rz
 from qualtran.bloqs.qft.two_bit_ffft import TwoBitFFFT
+from qualtran.bloqs.rotations.hamming_weight_phasing import HammingWeightPhasing
+from qualtran.symbolics import SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -66,8 +67,8 @@ class HoppingPlaquette(Bloq):
         page 13 Eq. E4 and E5 (Appendix E)
     """
 
-    kappa: Union[float, sympy.Expr]
-    eps: Union[float, sympy.Expr] = 1e-9
+    kappa: Union[SymbolicFloat]
+    eps: Union[SymbolicFloat] = 1e-9
 
     @cached_property
     def signature(self) -> Signature:
@@ -109,10 +110,10 @@ class HoppingTile(Bloq):
         see Eq. 21 and App E.
     """
 
-    length: Union[int, sympy.Expr]
-    angle: Union[float, sympy.Expr]
+    length: Union[SymbolicInt]
+    angle: Union[SymbolicFloat]
     tau: float = 1.0
-    eps: Union[float, sympy.Expr] = 1e-9
+    eps: Union[SymbolicFloat] = 1e-9
     pink: bool = True
 
     def __attrs_post_init__(self):
@@ -134,6 +135,68 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         }
 
 
+@frozen
+class HoppingTileHWP(HoppingTile):
+    r"""Bloq implementing a "tile" of the one-body hopping unitary using Hamming weight phasing.
+
+    Implements the unitary
+    $$
+    e^{i H_h^{x}} = \prod_{k\sigma} e^{i t H_h^{x(k,\sigma)}}
+    $$
+    for a particular choise of of plaquette hamiltonian $H_h^x$, where $x$ = pink or gold.
+
+    Each plaquette Hamiltonian can be split into $L^2/4$ commuting terms. Each
+    term can be implemented using the 4-qubit HoppingPlaquette above. The
+    HoppingPlaquette bloq contains 2 arbitrary rotations which are flanked by Clifford operations.
+    All of the rotations within a HoppingTile have the same angle so we can use
+    HammingWeightPhaseing to reduce the number of T gates that need to be
+    synthesized. Accounting for spin there are then $2 \times 2 \times L^2/4$
+    arbitrary rotations in each Tile, but only  $L^2/2$ of them can be applied
+    at the same time due to the $e^{iXX} e^{iYY}$ circuit not permitting parallel $R_z$ gates.
+
+    Unlike in the HoppingTile implementation where we can neatly factor
+    everything into sub-bloqs, here we would need to apply any clifford and F
+    gates first in parallel then bulk apply the rotations in parallel using
+    HammingWeightPhasing and then apply another layer of clifford and F gates.
+
+    Args:
+        length: Lattice side length L.
+        angle: The prefactor scaling the Hopping hamiltonian in the unitary (`t` above).
+            This should contain any relevant prefactors including the time step
+            and any splitting coefficients.
+        tau: The Hopping hamiltonian parameter. Typically the hubbard model is
+            defined relative to $\tau$ so it's defaulted to 1.
+        eps: The precision of the single qubit rotations.
+        pink: The colour of the plaquette.
+
+    Registers:
+        system: The system register of size 2 `length`.
+
+    References:
+        [Early fault-tolerant simulations of the Hubbard model](
+            https://arxiv.org/abs/2012.09238) see Eq. 21 and App E.
+    """
+
+    def short_name(self) -> str:
+        l = 'p' if self.pink else 'g'
+        return f'H_h^{l}(HWP)'
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        # Page 5, text after Eq. 22. There are L^2 / 4 plaquettes of a given colour and x2 for spin.
+        # Each plaquette contributes 4 TwoBitFFFT gates and two arbitrary rotations.
+        # We use Hamming weight phasing to apply all 2 * L^2/4 (two for spin
+        # here) for both of these rotations.
+        return {
+            (TwoBitFFFT(0, 1, self.eps), 4 * self.length**2 // 2),
+            (
+                HammingWeightPhasing(
+                    2 * self.length**2 // 4, self.tau * self.angle, eps=self.eps
+                ),
+                2,
+            ),
+        }
+
+
 @bloq_example
 def _hopping_tile() -> HoppingTile:
     length = 8
@@ -162,3 +225,18 @@ def _plaquette() -> HoppingPlaquette:
     import_line='from qualtran.bloqs.chemistry.trotter.hubbard.hopping import HoppingPlaquette',
     examples=(_plaquette,),
 )
+
+
+@bloq_example
+def _hopping_tile_hwp() -> HoppingTileHWP:
+    length = 8
+    angle = 0.15
+    hopping_tile_hwp = HoppingTileHWP(length, angle)
+    return hopping_tile_hwp
+
+
+_HOPPING_TILE_HWP_DOC = BloqDocSpec(
+    bloq_cls=HoppingTileHWP,
+    import_line='from qualtran.bloqs.chemistry.trotter.hubbard.hopping import HoppingTileHWP',
+    examples=(_hopping_tile_hwp,),
+)

qualtran/bloqs/chemistry/trotter/hubbard/hopping_test.py

@@ -12,8 +12,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from qualtran import Bloq
-from qualtran.bloqs.basic_gates import Rz, TGate
-from qualtran.bloqs.chemistry.trotter.hubbard.hopping import _hopping_tile, _plaquette
+from qualtran.bloqs.basic_gates import Rz, TGate, ZPowGate
+from qualtran.bloqs.chemistry.trotter.hubbard.hopping import (
+    _hopping_tile,
+    _hopping_tile_hwp,
+    _plaquette,
+)
 from qualtran.bloqs.util_bloqs import ArbitraryClifford
 from qualtran.resource_counting.generalizers import PHI
 
@@ -27,8 +31,10 @@ def test_hopping_plaquette(bloq_autotester):
 
 
 def catch_rotations(bloq) -> Bloq:
-    if isinstance(bloq, Rz):
-        if abs(float(bloq.angle)) < 1e-12:
+    if isinstance(bloq, (Rz, ZPowGate)):
+        if isinstance(bloq, ZPowGate):
+            return Rz(angle=PHI)
+        elif abs(float(bloq.angle)) < 1e-12:
             return ArbitraryClifford(1)
         else:
             return Rz(angle=PHI)
@@ -40,3 +46,13 @@ def test_hopping_tile_t_counts():
     _, counts = bloq.call_graph(generalizer=catch_rotations)
     assert counts[TGate()] == 8 * bloq.length**2 // 2
     assert counts[Rz(PHI)] == 2 * bloq.length**2 // 2
+
+
+def test_hopping_tile_hwp_t_counts():
+    bloq = _hopping_tile_hwp()
+    _, counts = bloq.call_graph(generalizer=catch_rotations)
+    n_rot_par = bloq.length**2 // 2
+    assert counts[Rz(PHI)] == 2 * n_rot_par.bit_length()
+    assert counts[TGate()] == 8 * bloq.length**2 // 2 + 2 * 4 * (
+        n_rot_par - n_rot_par.bit_count()
+    )

qualtran/bloqs/chemistry/trotter/hubbard/interaction.py

@@ -16,11 +16,12 @@
 from functools import cached_property
 from typing import Set, TYPE_CHECKING, Union
 
-import sympy
 from attrs import frozen
 
 from qualtran import Bloq, bloq_example, BloqDocSpec, QAny, Register, Signature
 from qualtran.bloqs.basic_gates import Rz
+from qualtran.bloqs.rotations.hamming_weight_phasing import HammingWeightPhasing
+from qualtran.symbolics import SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
@@ -52,10 +53,10 @@ class Interaction(Bloq):
         Eq. 6 page 2 and page 13 paragraph 1.
     """
 
-    length: Union[int, sympy.Expr]
-    angle: Union[float, sympy.Expr]
-    hubb_u: Union[float, sympy.Expr]
-    eps: Union[float, sympy.Expr] = 1e-9
+    length: Union[SymbolicInt]
+    angle: Union[SymbolicFloat]
+    hubb_u: Union[SymbolicFloat]
+    eps: Union[SymbolicFloat] = 1e-9
 
     @cached_property
     def signature(self) -> Signature:
@@ -66,6 +67,50 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return {(Rz(angle=self.angle * self.hubb_u, eps=self.eps), self.length**2)}
 
 
+@frozen
+class InteractionHWP(Bloq):
+    r"""Bloq implementing the hubbard U part of the hamiltonian using Hamming weight phasing.
+
+    Specifically:
+    $$
+        U_I = e^{i t H_I}
+    $$
+    which can be implemented using equal angle single-qubit Z rotations.
+
+    Each interaction term can be implemented using a e^{iZZ} gate, which
+    decomposes into a single Rz gate flanked by cliffords. There are L^2
+    equal angle rotations in total all of which may be applied in parallel using HWP.
+
+    Args:
+        length: Lattice length L.
+        angle: The rotation angle for unitary.
+        hubb_u: The hubbard U parameter.
+        eps: The precision for single qubit rotations.
+
+    Registers:
+        system: The system register of size 2 `length`.
+
+    References:
+        [Early fault-tolerant simulations of the Hubbard model](
+            https://arxiv.org/abs/2012.09238) Eq. page 13 paragraph 1, and page
+            14 paragraph 3 right column. The apply 2 batches of $L^2/2$ rotations.
+    """
+
+    length: Union[SymbolicInt]
+    angle: Union[SymbolicFloat]
+    hubb_u: Union[SymbolicFloat]
+    eps: Union[SymbolicFloat] = 1e-9
+
+    @cached_property
+    def signature(self) -> Signature:
+        return Signature([Register('system', QAny(self.length), shape=(2,))])
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        return {
+            (HammingWeightPhasing(self.length**2 // 2, self.angle * self.hubb_u, eps=self.eps), 2)
+        }
+
+
 @bloq_example
 def _interaction() -> Interaction:
     length = 8
@@ -80,3 +125,19 @@ def _interaction() -> Interaction:
     import_line='from qualtran.bloqs.chemistry.trotter.hubbard.interaction import Interaction',
     examples=(_interaction,),
 )
+
+
+@bloq_example
+def _interaction_hwp() -> InteractionHWP:
+    length = 8
+    angle = 0.5
+    hubb_u = 4.0
+    interaction = InteractionHWP(length, angle, hubb_u)
+    return interaction
+
+
+_INTERACTION_HWP_DOC = BloqDocSpec(
+    bloq_cls=InteractionHWP,
+    import_line='from qualtran.bloqs.chemistry.trotter.hubbard.interaction import InteractionHWP',
+    examples=(_interaction_hwp,),
+)

qualtran/bloqs/chemistry/trotter/hubbard/interaction_test.py

@@ -11,8 +11,30 @@
 #  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.trotter.hubbard.interaction import _interaction
+from qualtran.bloqs.basic_gates import Rz, TGate
+from qualtran.bloqs.chemistry.trotter.hubbard.hopping_test import catch_rotations
+from qualtran.bloqs.chemistry.trotter.hubbard.interaction import _interaction, _interaction_hwp
+from qualtran.resource_counting.generalizers import PHI
 
 
 def test_hopping_tile(bloq_autotester):
     bloq_autotester(_interaction)
+
+
+def test_interaction_hwp(bloq_autotester):
+    bloq_autotester(_interaction_hwp)
+
+
+def test_interaction_hwp_bloq_counts():
+    bloq = _interaction_hwp()
+    _, counts = bloq.call_graph(generalizer=catch_rotations)
+    n_rot_par = bloq.length**2 // 2
+    assert counts[Rz(PHI)] == 2 * n_rot_par.bit_length()
+    assert counts[TGate()] == 2 * 4 * (n_rot_par - n_rot_par.bit_count())
+
+
+def test_interaction_bloq_counts():
+    bloq = _interaction()
+    _, counts = bloq.call_graph(generalizer=catch_rotations)
+    n_rot = bloq.length**2
+    assert counts[Rz(PHI)] == n_rot

qualtran/bloqs/chemistry/trotter/hubbard/qpe_cost_optimization.ipynb

@@ -446,6 +446,59 @@
     "print(rf\"T_{{opt}} = {t_opt:4.3e}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using Hamming Weight Phasing "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can compare this cost to that found using Hamming weight phasing for the equal angle rotations. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.chemistry.trotter.hubbard.trotter_step import build_plaq_hwp_unitary_second_order_suzuki\n",
+    "trotter_step_hwp = build_plaq_hwp_unitary_second_order_suzuki(length, hubb_u, timestep, eps=1e-10)\n",
+    "n_t_hwp, n_rot_hwp = t_and_rot_counts_from_sigma(trotter_step_hwp.call_graph(generalizer=catch_rotations)[1])\n",
+    "print(f\"N_T(HWP) = {n_t_hwp} vs {(3*length**2 // 2)*8}\")\n",
+    "print(f\"N_rot(HWP) = {n_rot_hwp} vs {(3 * length**2 + 2*length**2)}\")\n",
+    "delta_ht_opt, delta_ts_opt, delta_pe_opt, t_opt = minimize_linesearch(n_rot_hwp, n_t_hwp, xi_bound, prod_ord)\n",
+    "print(rf\"T_{{OPT}}(HWP) = {t_opt:4.3e}\")\n",
+    "# The reference counts Toffolis as 2 T gates, we count them as 4.\n",
+    "print(rf\"Reference value = {1.7e6 + 4 * 1.8e5:4.3e}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Our value is slightly higher as we included all terms in the Trotter step. The reference only counts one layer of interaction gates. Let's correct for that."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trotter_step_hwp = build_plaq_hwp_unitary_second_order_suzuki(length, hubb_u, timestep, eps=1e-10, strip_layer=True)\n",
+    "n_t_hwp, n_rot_hwp = t_and_rot_counts_from_sigma(trotter_step_hwp.call_graph(generalizer=catch_rotations)[1])\n",
+    "print(f\"N_T(HWP) = {n_t_hwp}\")\n",
+    "print(f\"N_rot(HWP) = {n_rot_hwp}\")\n",
+    "delta_ht_opt, delta_ts_opt, delta_pe_opt, t_opt = minimize_linesearch(n_rot_hwp, n_t_hwp, xi_bound, prod_ord)\n",
+    "print(rf\"T_{{OPT}}(HWP) = {t_opt:4.3e}\")\n",
+    "print(rf\"Reference value = {1.7e6 + 4 * 1.8e5:4.3e}\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},