From f147315a7663f5a876af46797ca9b4ebfea1104f Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Mon, 7 Nov 2022 14:43:30 -0800
Subject: [PATCH] Improve generic select documentation (#105)

* Improve generic select documentation

* try new docstring

* re-render

* re-render after merging main

* one line

* Re-render
---
 cirq_qubitization/generic_select.ipynb        | 39 ++++++++++---------
 cirq_qubitization/generic_select.py           | 31 ++++++++-------
 .../jupyter_autogen_factories.py              | 18 +++------
 3 files changed, 42 insertions(+), 46 deletions(-)

diff --git a/cirq_qubitization/generic_select.ipynb b/cirq_qubitization/generic_select.ipynb
index 9a93d9aba..917db3b13 100644
--- a/cirq_qubitization/generic_select.ipynb
+++ b/cirq_qubitization/generic_select.ipynb
@@ -7,7 +7,9 @@
     "cq.autogen": "title_cell"
    },
    "source": [
-    "# Generic Select"
+    "# Generic Select\n",
+    "\n",
+    "Gates for applying generic selected unitaries."
    ]
   },
   {
@@ -35,12 +37,21 @@
    },
    "source": [
     "## `GenericSelect`\n",
-    "Gate that implements SELECT for a Hamiltonian expressed as an LCU.\n",
+    "A SELECT gate for selecting and applying operators from an array of `PauliString`s.\n",
+    "\n",
+    "$$\n",
+    "\\mathrm{SELECT} = \\sum_{l}|l \\rangle \\langle l| \\otimes U_l\n",
+    "$$\n",
+    "\n",
+    "Where $U_l$ is a member of the Pauli group.\n",
     "\n",
-    "Recall: SELECT = \\sum_{l}|l><l| \\otimes U_{l}\n",
+    "This gate uses the unary iteration scheme to apply `select_unitaries[selection]` to `target`\n",
+    "controlled on the single-bit `control` register.\n",
     "\n",
-    "The first log(L) qubits is the index register and the last M qubits are the system\n",
-    "register U_{l} is applied to"
+    "#### Parameters\n",
+    " - `selection_bitsize`: The size of the indexing `select` register. This should be at least `log2(len(select_unitaries))`\n",
+    " - `target_bitsize`: The size of the `target` register.\n",
+    " - `select_unitaries`: List of `DensePauliString`s to apply to the `target` register. Each dense pauli string must contain `target_bitsize` terms.\n"
    ]
   },
   {
@@ -53,21 +64,13 @@
    "outputs": [],
    "source": [
     "from cirq_qubitization.generic_select import GenericSelect\n",
-    "from cirq_qubitization.generic_select_test import get_1d_ising_hamiltonian\n",
-    "\n",
-    "num_sites = 4\n",
-    "target_bitsize = num_sites\n",
-    "num_select_unitaries = 2 * num_sites\n",
-    "\n",
-    "# PBC Ising in 1-D has num_sites ZZ operations and num_sites X operations.\n",
-    "# Thus, 2 * num_sites Pauli ops\n",
-    "selection_bitsize = int(np.ceil(np.log(num_select_unitaries)))\n",
     "\n",
-    "target = cirq.LineQubit.range(target_bitsize)  # placeholder\n",
-    "ham = get_1d_ising_hamiltonian(target, 1, 1)\n",
-    "dense_ham = [tt.dense(target) for tt in ham]\n",
+    "target_bitsize = 4\n",
+    "us = ['XIXI', 'YIYI', 'ZZZZ', 'ZXYZ']\n",
+    "us = [cirq.DensePauliString(u) for u in us]\n",
+    "selection_bitsize = int(np.ceil(np.log2(len(us))))\n",
     "g = cq_testing.GateHelper(\n",
-    "    GenericSelect(selection_bitsize, target_bitsize, select_unitaries=dense_ham)\n",
+    "    GenericSelect(selection_bitsize, target_bitsize, select_unitaries=us)\n",
     ")\n",
     "\n",
     "display_gate_and_compilation(g, vertical=True)"
diff --git a/cirq_qubitization/generic_select.py b/cirq_qubitization/generic_select.py
index 06b34757b..5fa55d96e 100644
--- a/cirq_qubitization/generic_select.py
+++ b/cirq_qubitization/generic_select.py
@@ -1,3 +1,4 @@
+"""Gates for applying generic selected unitaries."""
 from typing import Sequence, List, Tuple
 from functools import cached_property
 import cirq
@@ -6,12 +7,23 @@
 
 
 class GenericSelect(unary_iteration.UnaryIterationGate):
-    r"""Gate that implements SELECT for a Hamiltonian expressed as an LCU.
+    r"""A SELECT gate for selecting and applying operators from an array of `PauliString`s.
 
-    Recall: SELECT = \sum_{l}|l><l| \otimes U_{l}
+    $$
+    \mathrm{SELECT} = \sum_{l}|l \rangle \langle l| \otimes U_l
+    $$
 
-    The first log(L) qubits is the index register and the last M qubits are the system
-    register U_{l} is applied to
+    Where $U_l$ is a member of the Pauli group.
+
+    This gate uses the unary iteration scheme to apply `select_unitaries[selection]` to `target`
+    controlled on the single-bit `control` register.
+
+    Args:
+        selection_bitsize: The size of the indexing `select` register. This should be at least
+            `log2(len(select_unitaries))`
+        target_bitsize: The size of the `target` register.
+        select_unitaries: List of `DensePauliString`s to apply to the `target` register. Each
+            dense pauli string must contain `target_bitsize` terms.
     """
 
     def __init__(
@@ -20,17 +32,6 @@ def __init__(
         target_bitsize: int,
         select_unitaries: List[cirq.DensePauliString],
     ):
-        """An implementation of the SELECT unitary using the `UnaryIterationGate`
-
-        Args:
-            selection_bitsize: Number of qubits needed for select register. This is ceil(log2(len(select_unitaries)))
-            target_bitsize: number of qubits in the target register.
-            select_unitaries: List of DensePauliString's to apply to target register. Each dense
-                pauli string must contain `target_bitsize` terms.
-
-        Raises:
-            ValueError if any(len(dps) != target_bitsize for dps in select_unitaries).
-        """
         if any(len(dps) != target_bitsize for dps in select_unitaries):
             raise ValueError(
                 f"Each dense pauli string in `select_unitaries` should contain "
diff --git a/cirq_qubitization/jupyter_autogen_factories.py b/cirq_qubitization/jupyter_autogen_factories.py
index 9091bb14c..f27fbf698 100644
--- a/cirq_qubitization/jupyter_autogen_factories.py
+++ b/cirq_qubitization/jupyter_autogen_factories.py
@@ -56,20 +56,12 @@ def _make_MultiTargetCSwapApprox():
 
 def _make_GenericSelect():
     from cirq_qubitization.generic_select import GenericSelect
-    from cirq_qubitization.generic_select_test import get_1d_ising_hamiltonian
 
-    num_sites = 4
-    target_bitsize = num_sites
-    num_select_unitaries = 2 * num_sites
-
-    # PBC Ising in 1-D has num_sites ZZ operations and num_sites X operations.
-    # Thus, 2 * num_sites Pauli ops
-    selection_bitsize = int(np.ceil(np.log(num_select_unitaries)))
-
-    target = cirq.LineQubit.range(target_bitsize)  # placeholder
-    ham = get_1d_ising_hamiltonian(target, 1, 1)
-    dense_ham = [tt.dense(target) for tt in ham]
-    return GenericSelect(selection_bitsize, target_bitsize, select_unitaries=dense_ham)
+    target_bitsize = 4
+    us = ['XIXI', 'YIYI', 'ZZZZ', 'ZXYZ']
+    us = [cirq.DensePauliString(u) for u in us]
+    selection_bitsize = int(np.ceil(np.log2(len(us))))
+    return GenericSelect(selection_bitsize, target_bitsize, select_unitaries=us)
 
 
 def _make_GenericSubPrepare():