Bloqs for one hot encoding

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -131,6 +131,7 @@
 import qualtran.bloqs.state_preparation.state_preparation_via_rotation
 import qualtran.bloqs.swap_network.cswap_approx
 import qualtran.bloqs.swap_network.multiplexed_cswap
+import qualtran.bloqs.swap_network.one_hot_encoding
 import qualtran.bloqs.swap_network.swap_with_zero
 
 from .jupyter_autogen import NotebookSpecV2
@@ -227,6 +228,14 @@
             qualtran.bloqs.swap_network.multiplexed_cswap._MULTIPLEXED_CSWAP_DOC,
         ],
     ),
+    NotebookSpecV2(
+        title='One Hot Encodings',
+        module=qualtran.bloqs.swap_network.one_hot_encoding,
+        bloq_specs=[
+            qualtran.bloqs.swap_network.one_hot_encoding._ONE_HOT_LOG_DEPTH_DOC,
+            qualtran.bloqs.swap_network.one_hot_encoding._ONE_HOT_LINEAR_DEPTH_DOC,
+        ],
+    ),
     NotebookSpecV2(
         title='Global Phase',
         module=qualtran.bloqs.basic_gates.global_phase,

docs/bloqs/index.rst

@@ -42,6 +42,7 @@ Bloqs Library
     basic_gates/states_and_effects.ipynb
     basic_gates/swap.ipynb
     swap_network/swap_network.ipynb
+    swap_network/one_hot_encoding.ipynb
     basic_gates/global_phase.ipynb
     basic_gates/identity.ipynb
     bookkeeping/bookkeeping.ipynb

qualtran/bloqs/swap_network/one_hot_encoding.ipynb

@@ -0,0 +1,236 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "fc2b66f2",
+   "metadata": {
+    "cq.autogen": "title_cell"
+   },
+   "source": [
+    "# One Hot Encodings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4fe93874",
+   "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": "cf33d078",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.bloq_doc.md"
+   },
+   "source": [
+    "## `OneHotLogDepth`\n",
+    "Log depth one hot encoding using N - 1 CSWAPs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d873a11f",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.swap_network import OneHotLogDepth"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ddd4c000",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bd0a15fa",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.one_hot_log_depth"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import BQUInt\n",
+    "\n",
+    "one_hot_log_depth = OneHotLogDepth(BQUInt(4, 14))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6d5915cb",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4b05b47",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([one_hot_log_depth],\n",
+    "           ['`one_hot_log_depth`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "05b1373f",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7973a854",
+   "metadata": {
+    "cq.autogen": "OneHotLogDepth.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "one_hot_log_depth_g, one_hot_log_depth_sigma = one_hot_log_depth.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(one_hot_log_depth_g)\n",
+    "show_counts_sigma(one_hot_log_depth_sigma)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f440d529",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.bloq_doc.md"
+   },
+   "source": [
+    "## `OneHotLinearDepth`\n",
+    "Linear depth one hot encoding using N - 1 CSWAPs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "916169e5",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.swap_network import OneHotLinearDepth"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d0647430",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "68bd4068",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.one_hot_linear_depth"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran import BQUInt\n",
+    "\n",
+    "one_hot_linear_depth = OneHotLinearDepth(BQUInt(4, 14))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dde49844",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9440730a",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([one_hot_linear_depth],\n",
+    "           ['`one_hot_linear_depth`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "47ab8e36",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96546b81",
+   "metadata": {
+    "cq.autogen": "OneHotLinearDepth.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "one_hot_linear_depth_g, one_hot_linear_depth_sigma = one_hot_linear_depth.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(one_hot_linear_depth_g)\n",
+    "show_counts_sigma(one_hot_linear_depth_sigma)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

qualtran/bloqs/swap_network/one_hot_encoding.py

@@ -0,0 +1,175 @@
+#  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, Optional, Tuple
+
+import attrs
+import numpy as np
+
+from qualtran import (
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    BQUInt,
+    QBit,
+    Register,
+    Side,
+    Signature,
+    SoquetT,
+)
+from qualtran.bloqs.basic_gates import CNOT, XGate
+from qualtran.bloqs.mcmt import And, MultiTargetCNOT
+from qualtran.drawing import Circle, Text, TextBox, WireSymbol
+from qualtran.resource_counting.generalizers import ignore_split_join
+from qualtran.symbolics import SymbolicInt
+
+
+@attrs.frozen
+class CSwapViaAnd(Bloq):
+    """CSWAP(a, b, c) when c is guaranteed to be 0."""
+
+    cvs: tuple[int, int] = (1, 1)
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [
+                Register('ctrl', QBit()),
+                Register('x', QBit()),
+                Register('y', QBit(), side=Side.RIGHT),
+            ]
+        )
+
+    def build_composite_bloq(
+        self, bb: 'BloqBuilder', *, ctrl: 'SoquetT', x: 'SoquetT'
+    ) -> Dict[str, 'SoquetT']:
+        (ctrl, x), y = bb.add(And(*self.cvs), ctrl=[ctrl, x])
+        y, x = bb.add(CNOT(), ctrl=y, target=x)
+        return {'ctrl': ctrl, 'x': x, 'y': y}
+
+    def wire_symbol(self, reg: Optional['Register'], idx: Tuple[int, ...] = ()) -> 'WireSymbol':
+        if reg is None:
+            return Text('')
+        if reg.name == 'ctrl':
+            return Circle(filled=True)
+        else:
+            return TextBox('×')
+
+
+@attrs.frozen
+class OneHotLinearDepth(Bloq):
+    r"""Linear depth one hot encoding using N - 1 CSWAPs."""
+
+    selection_dtype: BQUInt
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [
+                Register('x', self.selection_dtype),
+                Register(
+                    'out', QBit(), shape=(self.selection_dtype.iteration_length,), side=Side.RIGHT
+                ),
+            ]
+        )
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> Dict[str, 'SoquetT']:
+        x = bb.split(x)[::-1]
+        out = [bb.allocate(dtype=QBit())]
+        out[0] = bb.add(XGate(), q=out[0])
+        for i in range(len(x)):
+            new_out = []
+            for j in range(2**i):
+                if j + 2**i < self.selection_dtype.iteration_length:
+                    x[i], out[j], out_k = bb.add(CSwapViaAnd(), ctrl=x[i], x=out[j])
+                    new_out.append(out_k)
+            out.extend(new_out)
+        return {'x': bb.join(x[::-1], dtype=self.selection_dtype), 'out': np.array(out)}
+
+
+@bloq_example(generalizer=[ignore_split_join])
+def _one_hot_linear_depth() -> OneHotLinearDepth:
+    from qualtran import BQUInt
+
+    one_hot_linear_depth = OneHotLinearDepth(BQUInt(4, 14))
+    return one_hot_linear_depth
+
+
+@attrs.frozen
+class Fanout(Bloq):
+    """Fanout via Multi-Target CNOT"""
+
+    n_copies: SymbolicInt
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [Register('x', QBit()), Register('y', QBit(), shape=(self.n_copies,), side=Side.RIGHT)]
+        )
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> Dict[str, 'SoquetT']:
+        y = bb.allocate(self.n_copies)
+        x, y = bb.add(MultiTargetCNOT(self.n_copies), control=x, targets=y)
+        return {'x': x, 'y': bb.split(y)}
+
+
+@attrs.frozen
+class OneHotLogDepth(Bloq):
+    r"""Log depth one hot encoding using N - 1 CSWAPs."""
+
+    selection_dtype: BQUInt
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [
+                Register('x', self.selection_dtype),
+                Register(
+                    'out', QBit(), shape=(self.selection_dtype.iteration_length,), side=Side.RIGHT
+                ),
+            ]
+        )
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> Dict[str, 'SoquetT']:
+        x = bb.split(x)[::-1]
+        out = [bb.allocate(dtype=QBit())]
+        out[0] = bb.add(XGate(), q=out[0])
+        for i in range(len(x)):
+            new_out = []
+            n_alloc = max(0, min(self.selection_dtype.iteration_length, 2 ** (i + 1)) - 2**i)
+            if n_alloc:
+                x[i], xx = bb.add(Fanout(n_alloc), x=x[i])
+            for j in range(n_alloc):
+                assert j + 2**i < self.selection_dtype.iteration_length
+                xx[j], out[j], out_k = bb.add(CSwapViaAnd(), ctrl=xx[j], x=out[j])
+                new_out.append(out_k)
+            out.extend(new_out)
+            if n_alloc:
+                x[i] = bb.add(Fanout(n_alloc).adjoint(), x=x[i], y=xx)
+        return {'x': bb.join(x[::-1], dtype=self.selection_dtype), 'out': np.array(out)}
+
+
+@bloq_example(generalizer=[ignore_split_join])
+def _one_hot_log_depth() -> OneHotLogDepth:
+    from qualtran import BQUInt
+
+    one_hot_log_depth = OneHotLogDepth(BQUInt(4, 14))
+    return one_hot_log_depth
+
+
+_ONE_HOT_LINEAR_DEPTH_DOC = BloqDocSpec(
+    bloq_cls=OneHotLinearDepth, examples=(_one_hot_linear_depth,)
+)
+_ONE_HOT_LOG_DEPTH_DOC = BloqDocSpec(bloq_cls=OneHotLogDepth, examples=(_one_hot_log_depth,))

qualtran/bloqs/swap_network/one_hot_encoding_test.py

@@ -0,0 +1,63 @@
+#  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 numpy as np
+import pytest
+
+from qualtran import BQUInt
+from qualtran.bloqs.swap_network.one_hot_encoding import OneHotLinearDepth, OneHotLogDepth
+from qualtran.resource_counting import get_cost_value, QECGatesCost, QubitCount
+
+
+@pytest.mark.parametrize('n, ilen', [(3, 8), (4, 14)])
+def test_one_hot_linear_depth_classical_action(n, ilen):
+    bloq = OneHotLinearDepth(BQUInt(n, ilen))
+    for x in range(ilen):
+        x_out, out = bloq.call_classically(x=x)
+        assert x == x_out
+        assert out[x] == 1 and np.all(out[:x] == 0) and np.all(out[x + 1 :] == 0)
+
+
+@pytest.mark.parametrize('n, ilen', [(3, 8), (4, 14)])
+def test_one_hot_log_depth_classical_action(n, ilen):
+    bloq = OneHotLogDepth(BQUInt(n, ilen))
+    for x in range(ilen):
+        x_out, out = bloq.call_classically(x=x)
+        assert x == x_out
+        assert out[x] == 1 and np.all(out[:x] == 0) and np.all(out[x + 1 :] == 0)
+
+
+@pytest.mark.parametrize('n, ilen', [(3, 8), (4, 14), (5, 30), (6, 60), (7, 120)])
+def test_one_hot_linear_depth_gate_counts(n, ilen):
+    bloq = OneHotLinearDepth(BQUInt(n, ilen))
+    # N - 1 AND gates.
+    assert get_cost_value(bloq, QECGatesCost()).and_bloq == ilen - 1
+    assert get_cost_value(bloq, QECGatesCost()).total_t_count() == 4 * ilen - 4
+    # Linear depth.
+    assert ilen // 2 < len(bloq.decompose_bloq().to_cirq_circuit()) < ilen
+    # Qubit Counts
+    assert get_cost_value(bloq, QubitCount()) == n + ilen
+
+
+@pytest.mark.parametrize('n, ilen', [(3, 8), (4, 14), (5, 30), (6, 60), (7, 120)])
+def test_one_hot_log_depth_gate_counts(n, ilen):
+    bloq = OneHotLogDepth(BQUInt(n, ilen))
+    # N - 1 AND gates.
+    assert get_cost_value(bloq, QECGatesCost()).and_bloq == ilen - 1
+    assert get_cost_value(bloq, QECGatesCost()).total_t_count() == 4 * ilen - 4
+    # Log depth.
+    assert len(bloq.decompose_bloq().to_cirq_circuit()) == n + 2
+    # O(N) additional qubits help achieve log depth
+    assert n + ilen < get_cost_value(bloq, QubitCount()) < n + 2 * ilen