Support more symbolic decompositions (#1369)

* allow more symbolic decomp in Permutation bloqs

* `StatePreparationViaRotation` - support symbolic phasegrad bitsize

* permutation - add more symbolic examples

* state-prep-rotations - symbolic examples + autotest

* assert decomp exists

* regen notebooks

* update bloq_ex

Author: anurudhp

qualtran/bloqs/arithmetic/permutation.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "a48878df",
+   "id": "a697fa0f",
    "metadata": {
     "cq.autogen": "title_cell"
    },
@@ -13,7 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d083d346",
+   "id": "f3a56020",
    "metadata": {
     "cq.autogen": "top_imports"
    },
@@ -30,7 +30,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "56539662",
+   "id": "39be898d",
    "metadata": {
     "cq.autogen": "Permutation.bloq_doc.md"
    },
@@ -62,7 +62,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "df73124f",
+   "id": "88b7712d",
    "metadata": {
     "cq.autogen": "Permutation.bloq_doc.py"
    },
@@ -73,7 +73,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "31147b0d",
+   "id": "5c4722c6",
    "metadata": {
     "cq.autogen": "Permutation.example_instances.md"
    },
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ec6460b6",
+   "id": "8e38bf3a",
    "metadata": {
     "cq.autogen": "Permutation.permutation"
    },
@@ -96,7 +96,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0c8e04f0",
+   "id": "52218544",
    "metadata": {
     "cq.autogen": "Permutation.permutation_symb"
    },
@@ -113,7 +113,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8578d036",
+   "id": "cb1a6988",
    "metadata": {
     "cq.autogen": "Permutation.permutation_symb_with_cycles"
    },
@@ -132,7 +132,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "84b5606a",
+   "id": "f58b58f7",
    "metadata": {
     "cq.autogen": "Permutation.sparse_permutation"
    },
@@ -143,9 +143,26 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e644f2ac",
+   "metadata": {
+    "cq.autogen": "Permutation.sparse_permutation_with_symbolic_N"
+   },
+   "outputs": [],
+   "source": [
+    "import sympy\n",
+    "\n",
+    "N = sympy.symbols(\"N\", positive=True, integer=True)\n",
+    "sparse_permutation_with_symbolic_N = Permutation.from_partial_permutation_map(\n",
+    "    N, {0: 1, 1: 3, 2: 4, 3: 7}\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
-   "id": "369df9cb",
+   "id": "ace174d5",
    "metadata": {
     "cq.autogen": "Permutation.graphical_signature.md"
    },
@@ -156,20 +173,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "25b2dc2b",
+   "id": "1e55dcd4",
    "metadata": {
     "cq.autogen": "Permutation.graphical_signature.py"
    },
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([permutation, permutation_symb, permutation_symb_with_cycles, sparse_permutation],\n",
-    "           ['`permutation`', '`permutation_symb`', '`permutation_symb_with_cycles`', '`sparse_permutation`'])"
+    "show_bloqs([permutation, permutation_symb, permutation_symb_with_cycles, sparse_permutation, sparse_permutation_with_symbolic_N],\n",
+    "           ['`permutation`', '`permutation_symb`', '`permutation_symb_with_cycles`', '`sparse_permutation`', '`sparse_permutation_with_symbolic_N`'])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "ad4321ad",
+   "id": "32e2b5dd",
    "metadata": {
     "cq.autogen": "Permutation.call_graph.md"
    },
@@ -180,7 +197,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3bfbf3f7",
+   "id": "fab3fd69",
    "metadata": {
     "cq.autogen": "Permutation.call_graph.py"
    },
@@ -194,7 +211,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b3e895db",
+   "id": "408e209c",
    "metadata": {
     "cq.autogen": "PermutationCycle.bloq_doc.md"
    },
@@ -230,7 +247,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d569cd2c",
+   "id": "bc68fccd",
    "metadata": {
     "cq.autogen": "PermutationCycle.bloq_doc.py"
    },
@@ -241,7 +258,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a93c6e89",
+   "id": "de4c7922",
    "metadata": {
     "cq.autogen": "PermutationCycle.example_instances.md"
    },
@@ -252,19 +269,23 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "264ba946",
+   "id": "071de1e1",
    "metadata": {
-    "cq.autogen": "PermutationCycle.permutation_cycle"
+    "cq.autogen": "PermutationCycle.permutation_cycle_symb_N"
    },
    "outputs": [],
    "source": [
-    "permutation_cycle = PermutationCycle(4, (0, 1, 2))"
+    "import sympy\n",
+    "\n",
+    "N = sympy.symbols(\"n\", positive=True, integer=True)\n",
+    "cycle = (3, 1, 2)\n",
+    "permutation_cycle_symb_N = PermutationCycle(N, cycle)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "78572528",
+   "id": "fd61e92a",
    "metadata": {
     "cq.autogen": "PermutationCycle.permutation_cycle_symb"
    },
@@ -279,9 +300,21 @@
     "permutation_cycle_symb = PermutationCycle(N, cycle)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b818f3b1",
+   "metadata": {
+    "cq.autogen": "PermutationCycle.permutation_cycle"
+   },
+   "outputs": [],
+   "source": [
+    "permutation_cycle = PermutationCycle(4, (0, 1, 2))"
+   ]
+  },
   {
    "cell_type": "markdown",
-   "id": "87615783",
+   "id": "3d3c3e1b",
    "metadata": {
     "cq.autogen": "PermutationCycle.graphical_signature.md"
    },
@@ -292,20 +325,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f6961010",
+   "id": "de967993",
    "metadata": {
     "cq.autogen": "PermutationCycle.graphical_signature.py"
    },
    "outputs": [],
    "source": [
     "from qualtran.drawing import show_bloqs\n",
-    "show_bloqs([permutation_cycle, permutation_cycle_symb],\n",
-    "           ['`permutation_cycle`', '`permutation_cycle_symb`'])"
+    "show_bloqs([permutation_cycle_symb_N, permutation_cycle_symb, permutation_cycle],\n",
+    "           ['`permutation_cycle_symb_N`', '`permutation_cycle_symb`', '`permutation_cycle`'])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "16e1c105",
+   "id": "58deda5c",
    "metadata": {
     "cq.autogen": "PermutationCycle.call_graph.md"
    },
@@ -316,16 +349,16 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "27628f10",
+   "id": "ec16903f",
    "metadata": {
     "cq.autogen": "PermutationCycle.call_graph.py"
    },
    "outputs": [],
    "source": [
     "from qualtran.resource_counting.generalizers import ignore_split_join\n",
-    "permutation_cycle_g, permutation_cycle_sigma = permutation_cycle.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
-    "show_call_graph(permutation_cycle_g)\n",
-    "show_counts_sigma(permutation_cycle_sigma)"
+    "permutation_cycle_symb_N_g, permutation_cycle_symb_N_sigma = permutation_cycle_symb_N.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(permutation_cycle_symb_N_g)\n",
+    "show_counts_sigma(permutation_cycle_symb_N_sigma)"
    ]
   }
  ],

qualtran/bloqs/arithmetic/permutation.py

@@ -104,9 +104,6 @@ def signature(self) -> Signature:
     def bitsize(self):
         return bit_length(self.N - 1)
 
-    def is_symbolic(self):
-        return is_symbolic(self.N, self.cycle)
-
     def build_composite_bloq(self, bb: 'BloqBuilder', x: 'SoquetT') -> dict[str, 'SoquetT']:
         if is_symbolic(self.cycle):
             raise DecomposeTypeError(f"cannot decompose symbolic {self}")
@@ -126,7 +123,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'SoquetT') -> dict[str, 'So
         return {'x': x}
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
-        if self.is_symbolic():
+        if is_symbolic(self.cycle):
             x = ssa.new_symbol('x')
             cycle_len = slen(self.cycle)
             return {
@@ -143,6 +140,16 @@ def _permutation_cycle() -> PermutationCycle:
     return permutation_cycle
 
 
+@bloq_example
+def _permutation_cycle_symb_N() -> PermutationCycle:
+    import sympy
+
+    N = sympy.symbols("n", positive=True, integer=True)
+    cycle = (3, 1, 2)
+    permutation_cycle_symb_N = PermutationCycle(N, cycle)
+    return permutation_cycle_symb_N
+
+
 @bloq_example
 def _permutation_cycle_symb() -> PermutationCycle:
     import sympy
@@ -158,7 +165,7 @@ def _permutation_cycle_symb() -> PermutationCycle:
 _PERMUTATION_CYCLE_DOC = BloqDocSpec(
     bloq_cls=PermutationCycle,
     import_line='from qualtran.bloqs.arithmetic.permutation import PermutationCycle',
-    examples=[_permutation_cycle, _permutation_cycle_symb],
+    examples=[_permutation_cycle_symb_N, _permutation_cycle_symb, _permutation_cycle],
 )
 
 
@@ -261,7 +268,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', x: 'Soquet') -> dict[str, 'Soq
         return {'x': x}
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
-        if self.is_symbolic():
+        if is_symbolic(self.cycles):
             # worst case cost: single cycle of length N
             cycle = Shaped((self.N,))
             return {(PermutationCycle(self.N, cycle), 1)}
@@ -307,8 +314,25 @@ def _sparse_permutation() -> Permutation:
     return sparse_permutation
 
 
+@bloq_example
+def _sparse_permutation_with_symbolic_N() -> Permutation:
+    import sympy
+
+    N = sympy.symbols("N", positive=True, integer=True)
+    sparse_permutation_with_symbolic_N = Permutation.from_partial_permutation_map(
+        N, {0: 1, 1: 3, 2: 4, 3: 7}
+    )
+    return sparse_permutation_with_symbolic_N
+
+
 _PERMUTATION_DOC = BloqDocSpec(
     bloq_cls=Permutation,
     import_line='from qualtran.bloqs.arithmetic.permutation import Permutation',
-    examples=[_permutation, _permutation_symb, _permutation_symb_with_cycles, _sparse_permutation],
+    examples=[
+        _permutation,
+        _permutation_symb,
+        _permutation_symb_with_cycles,
+        _sparse_permutation,
+        _sparse_permutation_with_symbolic_N,
+    ],
 )