Add GF2Square bloq for squaring over GF($2^m$)

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -107,6 +107,7 @@
 import qualtran.bloqs.factoring.mod_exp
 import qualtran.bloqs.gf_arithmetic.gf2_addition
 import qualtran.bloqs.gf_arithmetic.gf2_multiplication
+import qualtran.bloqs.gf_arithmetic.gf2_square
 import qualtran.bloqs.hamiltonian_simulation.hamiltonian_simulation_by_gqsp
 import qualtran.bloqs.mcmt.and_bloq
 import qualtran.bloqs.mcmt.controlled_via_and
@@ -567,6 +568,11 @@
         module=qualtran.bloqs.gf_arithmetic.gf2_addition,
         bloq_specs=[qualtran.bloqs.gf_arithmetic.gf2_addition._GF2_ADDITION_DOC],
     ),
+    NotebookSpecV2(
+        title='GF($2^m$) Square',
+        module=qualtran.bloqs.gf_arithmetic.gf2_square,
+        bloq_specs=[qualtran.bloqs.gf_arithmetic.gf2_square._GF2_SQUARE_DOC],
+    ),
 ]
 
 

docs/bloqs/index.rst

@@ -93,6 +93,7 @@ Bloqs Library
 
     gf_arithmetic/gf2_multiplication.ipynb
     gf_arithmetic/gf2_addition.ipynb
+    gf_arithmetic/gf2_square.ipynb
 
 .. toctree::
     :maxdepth: 2

qualtran/bloqs/gf_arithmetic/__init__.py

@@ -14,3 +14,4 @@
 
 from qualtran.bloqs.gf_arithmetic.gf2_addition import GF2Addition
 from qualtran.bloqs.gf_arithmetic.gf2_multiplication import GF2Multiplication
+from qualtran.bloqs.gf_arithmetic.gf2_square import GF2Square

qualtran/bloqs/gf_arithmetic/gf2_square.ipynb

@@ -0,0 +1,173 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cb46f029",
+   "metadata": {
+    "cq.autogen": "title_cell"
+   },
+   "source": [
+    "# GF($2^m$) Square"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6054ebf1",
+   "metadata": {
+    "cq.autogen": "top_imports"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import Bloq, CompositeBloq, BloqBuilder, Signature, Register\n",
+    "from qualtran import QBit, QInt, QUInt, QAny\n",
+    "from qualtran.drawing import show_bloq, show_call_graph, show_counts_sigma\n",
+    "from typing import *\n",
+    "import numpy as np\n",
+    "import sympy\n",
+    "import cirq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abc9f3d5",
+   "metadata": {
+    "cq.autogen": "GF2Square.bloq_doc.md"
+   },
+   "source": [
+    "## `GF2Square`\n",
+    "In place squaring for elements in GF($2^m$)\n",
+    "\n",
+    "The bloq implements in-place squaring of a quantum registers storing elements\n",
+    "from GF($2^m$). Specifically, it implements the transformation\n",
+    "\n",
+    "$$\n",
+    "    |a\\rangle \\rightarrow |a^2\\rangle\n",
+    "$$\n",
+    "\n",
+    "The key insight is that for elements in GF($2^m$),\n",
+    "$$\n",
+    "    a^2 =a_0 + a_1 x^2 + a_2 x^4 + ... + a_{n-1} x^{2(n - 1)}\n",
+    "$$\n",
+    "\n",
+    "Thus, squaring can be implemented via a linear reversible circuit using only CNOT gates.\n",
+    "\n",
+    "#### Parameters\n",
+    " - `bitsize`: The degree $m$ of the galois field $GF(2^m)$. Also corresponds to the number of qubits in the input register to be squared. \n",
+    "\n",
+    "#### Registers\n",
+    " - `x`: Input THRU register of size $m$ that stores elements from $GF(2^m)$.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c78c541d",
+   "metadata": {
+    "cq.autogen": "GF2Square.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.gf_arithmetic import GF2Square"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3867aabe",
+   "metadata": {
+    "cq.autogen": "GF2Square.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10c374a4",
+   "metadata": {
+    "cq.autogen": "GF2Square.gf16_square"
+   },
+   "outputs": [],
+   "source": [
+    "gf16_square = GF2Square(4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "34d1aa8e",
+   "metadata": {
+    "cq.autogen": "GF2Square.gf2_square_symbolic"
+   },
+   "outputs": [],
+   "source": [
+    "import sympy\n",
+    "\n",
+    "m = sympy.Symbol('m')\n",
+    "gf2_square_symbolic = GF2Square(m)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "40f24bac",
+   "metadata": {
+    "cq.autogen": "GF2Square.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "218453ac",
+   "metadata": {
+    "cq.autogen": "GF2Square.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([gf16_square, gf2_square_symbolic],\n",
+    "           ['`gf16_square`', '`gf2_square_symbolic`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ffc34750",
+   "metadata": {
+    "cq.autogen": "GF2Square.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "05bdb33f",
+   "metadata": {
+    "cq.autogen": "GF2Square.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "gf16_square_g, gf16_square_sigma = gf16_square.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(gf16_square_g)\n",
+    "show_counts_sigma(gf16_square_sigma)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

qualtran/bloqs/gf_arithmetic/gf2_square.py

@@ -0,0 +1,125 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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 functools import cached_property
+from typing import Dict, Set, TYPE_CHECKING, Union
+
+import attrs
+import numpy as np
+from galois import GF, Poly
+
+from qualtran import Bloq, bloq_example, BloqDocSpec, DecomposeTypeError, QGF, Register, Signature
+from qualtran.bloqs.gf_arithmetic.gf2_multiplication import SynthesizeLRCircuit
+from qualtran.symbolics import is_symbolic, Shaped, SymbolicInt
+
+if TYPE_CHECKING:
+    from qualtran import BloqBuilder, Soquet
+    from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
+    from qualtran.simulation.classical_sim import ClassicalValT
+
+
+@attrs.frozen
+class GF2Square(Bloq):
+    r"""In place squaring for elements in GF($2^m$)
+
+    The bloq implements in-place squaring of a quantum registers storing elements
+    from GF($2^m$). Specifically, it implements the transformation
+
+    $$
+        |a\rangle \rightarrow |a^2\rangle
+    $$
+
+    The key insight is that for elements in GF($2^m$),
+    $$
+        a^2 =a_0 + a_1 x^2 + a_2 x^4 + ... + a_{n-1} x^{2(n - 1)}
+    $$
+
+    Thus, squaring can be implemented via a linear reversible circuit using only CNOT gates.
+
+    Args:
+        bitsize: The degree $m$ of the galois field $GF(2^m)$. Also corresponds to the number of
+            qubits in the input register to be squared.
+
+    Registers:
+        x: Input THRU register of size $m$ that stores elements from $GF(2^m)$.
+    """
+
+    bitsize: SymbolicInt
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature([Register('x', dtype=self.qgf)])
+
+    @cached_property
+    def qgf(self) -> QGF:
+        return QGF(characteristic=2, degree=self.bitsize)
+
+    @cached_property
+    def squaring_matrix(self) -> np.ndarray:
+        r"""$m \times m$ matrix that maps the input $x^{i}$ to $x^{2 * i} % P(x)$"""
+        m = int(self.bitsize)
+        f = self.qgf.gf_type.irreducible_poly
+        M = np.zeros((m, m))
+        alpha = [0] * m
+        for i in range(m):
+            # x ** (2 * i) % f
+            alpha[-i - 1] = 1
+            coeffs = ((Poly(alpha, GF(2)) * Poly(alpha, GF(2))) % f).coeffs.tolist()[::-1]
+            coeffs = coeffs + [0] * (m - len(coeffs))
+            M[i] = coeffs
+            alpha[-i - 1] = 0
+        return np.transpose(M)
+
+    @cached_property
+    def synthesize_squaring_matrix(self) -> SynthesizeLRCircuit:
+        m = self.bitsize
+        return (
+            SynthesizeLRCircuit(Shaped((m, m)))
+            if is_symbolic(m)
+            else SynthesizeLRCircuit(self.squaring_matrix)
+        )
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> Dict[str, 'Soquet']:
+        if is_symbolic(self.bitsize):
+            raise DecomposeTypeError(f"Cannot decompose symbolic {self}")
+        x = bb.split(x)[::-1]
+        x = bb.add(self.synthesize_squaring_matrix, q=x)
+        x = bb.join(x[::-1], dtype=self.qgf)
+        return {'x': x}
+
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
+        return {self.synthesize_squaring_matrix: 1}
+
+    def on_classical_vals(self, *, x) -> Dict[str, 'ClassicalValT']:
+        assert isinstance(x, self.qgf.gf_type)
+        return {'x': x**2}
+
+
+@bloq_example
+def _gf16_square() -> GF2Square:
+    gf16_square = GF2Square(4)
+    return gf16_square
+
+
+@bloq_example
+def _gf2_square_symbolic() -> GF2Square:
+    import sympy
+
+    m = sympy.Symbol('m')
+    gf2_square_symbolic = GF2Square(m)
+    return gf2_square_symbolic
+
+
+_GF2_SQUARE_DOC = BloqDocSpec(bloq_cls=GF2Square, examples=(_gf16_square, _gf2_square_symbolic))

qualtran/bloqs/gf_arithmetic/gf2_square_test.py

@@ -0,0 +1,52 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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.
+
+import pytest
+import sympy
+from galois import GF
+
+from qualtran.bloqs.gf_arithmetic.gf2_square import _gf2_square_symbolic, _gf16_square, GF2Square
+from qualtran.resource_counting import get_cost_value, QECGatesCost
+from qualtran.symbolics import ceil, log2
+from qualtran.testing import assert_consistent_classical_action
+
+
+def test_gf16_square(bloq_autotester):
+    bloq_autotester(_gf16_square)
+
+
+def test_gf2_square_symbolic(bloq_autotester):
+    bloq_autotester(_gf2_square_symbolic)
+
+
+def test_gf2_square_classical_sim_quick():
+    m = 2
+    bloq = GF2Square(m)
+    GFM = GF(2**m)
+    assert_consistent_classical_action(bloq, x=GFM.elements)
+
+
+def test_gf2_square_resource():
+    bloq = _gf2_square_symbolic.make()
+    m = bloq.bitsize
+    assert get_cost_value(bloq, QECGatesCost()).total_t_count() == 0
+    assert sympy.simplify(get_cost_value(bloq, QECGatesCost()).clifford - ceil(m**2 / log2(m))) == 0
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize('m', [3, 4, 5])
+def test_gf2_square_classical_sim(m):
+    bloq = GF2Square(m)
+    GFM = GF(2**m)
+    assert_consistent_classical_action(bloq, x=GFM.elements)

qualtran/conftest.py

@@ -94,6 +94,8 @@ def assert_bloq_example_serializes_for_pytest(bloq_ex: BloqExample):
         'gf2_addition_symbolic',  # cannot serialize QGF
         'gf16_multiplication',  # cannot serialize QGF
         'gf2_multiplication_symbolic',  # cannot serialize QGF
+        'gf16_square',  # cannot serialize QGF
+        'gf2_square_symbolic',  # cannot serialize QGF
         'gqsp_1d_ising',
         'auto_partition',
         'unitary_block_encoding',