Use full call graph to build bloq JSON (quantumlib#1482)

Use full call graph to build bloq JSON, accounting for Adjoint and Controlled gates
KuantumBS · Oct 30, 2024 · bc155b6 · bc155b6
1 parent b3e0748
commit bc155b6
Showing 1 changed file with 69 additions and 26 deletions.
diff --git a/dev_tools/ui-export.ipynb b/dev_tools/ui-export.ipynb
@@ -53,21 +53,37 @@
     "import hashlib\n",
     "import json\n",
     "\n",
-    "from qualtran import CompositeBloq\n",
+    "from qualtran import Adjoint, CompositeBloq, Controlled\n",
     "from qualtran.bloqs.rotations.programmable_rotation_gate_array import ProgrammableRotationGateArray\n",
     "\n",
     "def bloq_attrs(bloq):\n",
     "  if isinstance(bloq, CompositeBloq):\n",
     "    return {}\n",
     "  if isinstance(bloq, ProgrammableRotationGateArray):\n",
     "    return {}\n",
+    "  if isinstance(bloq, (Adjoint, Controlled)):\n",
+    "    return bloq_attrs(bloq.subbloq)\n",
     "\n",
     "  return attrs.asdict(bloq)\n",
     "\n",
     "def bloq_filename(bloq):\n",
-    "  unhashed = json.dumps(bloq_attrs(bloq), cls=BloqEncoder)\n",
-    "\n",
-    "  return hashlib.md5(unhashed.encode(), usedforsecurity=False).hexdigest() + '.json'"
+    "  attrs_dict = bloq_attrs(bloq)\n",
+    "  attrs_keys = list(attrs_dict.keys())\n",
+    "  attrs_keys.sort()\n",
+    "\n",
+    "  prefix = ''\n",
+    "  if isinstance(bloq, Adjoint):\n",
+    "    prefix = 'Adjoint_'\n",
+    "  if isinstance(bloq, Controlled):\n",
+    "    prefix = 'Controlled_'\n",
+    "\n",
+    "  attrs_list = [\n",
+    "    [key, attrs_dict[key]]\n",
+    "    for key in attrs_keys\n",
+    "  ] if attrs_keys else []\n",
+    "  unhashed = json.dumps(attrs_list, cls=BloqEncoder)\n",
+    "\n",
+    "  return prefix + hashlib.md5(unhashed.encode(), usedforsecurity=False).hexdigest() + '.json'"
    ]
   },
   {
@@ -113,56 +129,83 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
     "import os\n",
     "from pathlib import Path\n",
     "\n",
     "from qualtran_dev_tools.notebook_specs import NB_BY_SECTION\n",
     "from qualtran_dev_tools.parse_docstrings import get_markdown_docstring\n",
+    "\n",
+    "for section in NB_BY_SECTION:\n",
+    "  for notebook_spec in section[1]:\n",
+    "    for bloq_spec in notebook_spec.bloq_specs:\n",
+    "      Path(f'ui_export/{bloq_spec.bloq_cls.__name__}').mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "      doc_name = f'ui_export/{bloq_spec.bloq_cls.__name__}/docs.txt'\n",
+    "      if not os.path.isfile(doc_name):\n",
+    "        with open(doc_name, 'w') as doc_file:\n",
+    "          doc_file.write('\\n'.join(get_markdown_docstring(bloq_spec.bloq_cls)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "from qualtran_dev_tools.all_call_graph import get_all_call_graph\n",
+    "from qualtran_dev_tools.notebook_specs import NB_BY_SECTION\n",
+    "from qualtran import Adjoint, Controlled\n",
     "from qualtran.drawing.musical_score import get_musical_score_data\n",
     "\n",
+    "examples = [\n",
+    "  example\n",
+    "  for section in NB_BY_SECTION\n",
+    "  for notebook_spec in section[1]\n",
+    "  for bloq_spec in notebook_spec.bloq_specs\n",
+    "  for example in bloq_spec.examples\n",
+    "]\n",
+    "\n",
+    "call_graph = get_all_call_graph(examples)\n",
+    "\n",
     "def bloq_score(bloq):\n",
     "  try:\n",
     "    return get_musical_score_data(bloq.decompose_bloq())\n",
     "  except:\n",
     "    return None\n",
     "\n",
-    "def write_example(bloq):\n",
-    "  call_graph, _ = bloq.call_graph(max_depth=1)\n",
+    "def bloq_name(bloq):\n",
+    "  if (isinstance(bloq, (Adjoint, Controlled))):\n",
+    "    return bloq_name(bloq.subbloq)\n",
     "\n",
-    "  for child_bloq, _ in call_graph.succ[bloq].items():\n",
-    "    write_example(child_bloq)\n",
+    "  return bloq.__class__.__name__\n",
     "\n",
-    "  file_name = f'ui_export/{bloq.__class__.__name__}/{bloq_filename(bloq)}'\n",
+    "def write_example(bloq):\n",
+    "  file_name = f'ui_export/{bloq_name(bloq)}/{bloq_filename(bloq)}'\n",
     "  if not os.path.isfile(file_name):\n",
     "    bloq_dict = {\n",
-    "      'name': bloq.__class__.__name__,\n",
+    "      'name': str(bloq),\n",
     "      'attrs': bloq_attrs(bloq),\n",
-    "      'score': bloq_score(bloq),\n",
+    "      'msd': bloq_score(bloq),\n",
     "      'callees': list(\n",
     "        {\n",
-    "          'name': child_bloq.__class__.__name__,\n",
+    "          'name': bloq_name(child_bloq),\n",
     "          'filename': bloq_filename(child_bloq)\n",
     "        }\n",
-    "        for child_bloq, _ in call_graph.succ[bloq].items()\n",
+    "        for child_bloq in call_graph.neighbors(bloq)\n",
     "      )\n",
     "    }\n",
     "\n",
+    "    Path(f'ui_export/{bloq_name(bloq)}').mkdir(parents=True, exist_ok=True)\n",
+    "\n",
     "    with open(file_name, 'w') as f:\n",
     "      json.dump(bloq_dict, f, indent=2, cls=BloqEncoder)\n",
     "\n",
-    "for section in NB_BY_SECTION:\n",
-    "  for notebook_spec in section[1]:\n",
-    "    for bloq_spec in notebook_spec.bloq_specs:\n",
-    "      Path(f'ui_export/{bloq_spec.bloq_cls.__name__}').mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "      doc_name = f'ui_export/{bloq_spec.bloq_cls.__name__}/docs.txt'\n",
-    "      if not os.path.isfile(doc_name):\n",
-    "        with open(doc_name, 'w') as doc_file:\n",
-    "          doc_file.write('\\n'.join(get_markdown_docstring(bloq_spec.bloq_cls)))\n",
-    "\n",
-    "      for example in bloq_spec.examples:\n",
-    "        write_example(example.make())"
+    "for bloq, _ in call_graph.nodes.items():\n",
+    "  write_example(bloq)"
    ]
   }
  ],