Merge remote-tracking branch 'origin/main' into 2025-09/soquetprotocol

Author: mpharrigan

.github/workflows/nightly.yaml

@@ -43,4 +43,4 @@ jobs:
           pip install -r dev_tools/requirements/envs/pytest.env.txt
           pip install --no-deps -e .
       - run: |
-          check/pytest
+          check/pytest --durations=10

.github/workflows/pr.yaml

@@ -76,7 +76,7 @@ jobs:
           pip install -r dev_tools/requirements/envs/pytest.env.txt
           pip install --no-deps -e .
       - run: |
-          python dev_tools/execute-notebooks.py
+          python dev_tools/execute-notebooks.py --n-workers=8
         env:
           NUMBA_NUM_THREADS: 4
 

dev_tools/conf/mypy.ini

@@ -21,7 +21,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # Non-Google
-[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*,nbformat.*,nbconvert.*,openfermion.*,pennylane.*]
+[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*,nbformat.*,nbconvert.*,openfermion.*,pennylane.*,mpmath.*]
 follow_imports = silent
 ignore_missing_imports = true
 

dev_tools/execute-notebooks.py

@@ -22,9 +22,13 @@ def parse_args():
     p.add_argument('--output-nbs', action=argparse.BooleanOptionalAction, default=True)
     p.add_argument('--output-md', action=argparse.BooleanOptionalAction, default=False)
     p.add_argument('--only-out-of-date', action=argparse.BooleanOptionalAction, default=True)
+    p.add_argument('--n-workers', type=int, default=None)
     args = p.parse_args()
     execute_and_export_notebooks(
-        output_nbs=args.output_nbs, output_md=args.output_md, only_out_of_date=args.only_out_of_date
+        output_nbs=args.output_nbs,
+        output_md=args.output_md,
+        only_out_of_date=args.only_out_of_date,
+        n_workers=args.n_workers,
     )
 
 

dev_tools/qualtran_dev_tools/notebook_execution.py

@@ -13,11 +13,14 @@
 #  limitations under the License.
 import asyncio
 import multiprocessing
+import random
 import subprocess
 import sys
+import time
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import attrs
 import filelock
 import nbconvert
 import nbformat
@@ -180,6 +183,13 @@ def execute_and_export_notebook(paths: _NBInOutPaths) -> Optional[Exception]:
     return None
 
 
+@attrs.frozen
+class _NotebookRunResult:
+    nb_in: Path
+    err: Optional[Exception]
+    duration_s: float
+
+
 class _NotebookRunClosure:
     """Used to run notebook execution logic in subprocesses."""
 
@@ -189,7 +199,7 @@ def __init__(self, reporoot: Path, output_nbs: bool, output_md: bool, only_out_o
         self.output_md = output_md
         self.only_out_of_date = only_out_of_date
 
-    def __call__(self, nb_rel_path: Path, sourceroot: Path) -> Tuple[Path, Optional[Exception]]:
+    def __call__(self, nb_rel_path: Path, sourceroot: Path) -> _NotebookRunResult:
         paths = _NBInOutPaths.from_nb_rel_path(
             nb_rel_path,
             reporoot=self.reporoot,
@@ -200,15 +210,21 @@ def __call__(self, nb_rel_path: Path, sourceroot: Path) -> Tuple[Path, Optional[
 
         if self.only_out_of_date and not paths.needs_reexport():
             print(f'{nb_rel_path} up to date')
-            return paths.nb_in, None
+            return _NotebookRunResult(paths.nb_in, None, 0.0)
 
+        start = time.time()
         err = execute_and_export_notebook(paths)
-        print(f"Exported {nb_rel_path}")
-        return paths.nb_in, err
+        end = time.time()
+        print(f"Exported {nb_rel_path} in {end-start:.2f} seconds.")
+        return _NotebookRunResult(paths.nb_in, err, duration_s=end - start)
 
 
 def execute_and_export_notebooks(
-    *, output_nbs: bool, output_md: bool, only_out_of_date: bool = True
+    *,
+    output_nbs: bool,
+    output_md: bool,
+    only_out_of_date: bool = True,
+    n_workers: Optional[int] = None,
 ):
     """Find, execute, and export all checked-in ipynbs.
 
@@ -217,23 +233,41 @@ def execute_and_export_notebooks(
         output_md: Whether to save the executed notebooks as markdown
         only_out_of_date: Only re-execute and re-export notebooks whose output files
             are out of date.
+        n_workers: If set to 1, do not use parallelization. If set to `None` (the detault),
+            `multiprocessing.Pool()` will be used, which uses the number of processors as
+            a default. Otherwise, this argument is passed to
+            `multiprocessing.Pool(n_workers)` to execute notebooks in parallel on this many
+            worker processes..
     """
     reporoot = get_git_root()
     nb_rel_paths = get_nb_rel_paths(sourceroot=reporoot / 'qualtran')
     nb_rel_paths += get_nb_rel_paths(sourceroot=reporoot / 'tutorials')
+    random.shuffle(nb_rel_paths)
+    print(f"Found {len(nb_rel_paths)} notebooks.")
     func = _NotebookRunClosure(
         reporoot=reporoot,
         output_nbs=output_nbs,
         output_md=output_md,
         only_out_of_date=only_out_of_date,
     )
-    with multiprocessing.Pool() as pool:
-        results = pool.starmap(func, nb_rel_paths)
-    bad_nbs = [nbname for nbname, err in results if err is not None]
+    if n_workers == 1:
+        print("(Not using multiprocessing, n_workers=1)")
+        results = [func(nb_rel_path, sourceroot) for nb_rel_path, sourceroot in nb_rel_paths]
+    else:
+        print(f"Multiprocessing with {n_workers=}")
+        with multiprocessing.Pool(n_workers, maxtasksperchild=1) as pool:
+            results = pool.starmap(func, nb_rel_paths)
+        assert results
+    bad_nbs = [result.nb_in for result in results if result.err is not None]
 
     if len(bad_nbs) > 0:
         print()
         print("Errors in notebooks:")
         for nb in bad_nbs:
             print(' ', nb)
         sys.exit(1)
+
+    duration_nbs = sorted(results, key=lambda r: r.duration_s, reverse=True)
+    print("Slowest 10 notebooks:")
+    for result in duration_nbs[:10]:
+        print(f'{result.duration_s:5.2f}s {result.nb_in}')

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -368,6 +368,24 @@
         ],
         directory=f'{SOURCE_DIR}/bloqs/chemistry/trotter/hubbard',
     ),
+    NotebookSpecV2(
+        title='Qubitized Hubbard',
+        module=qualtran.bloqs.chemistry.hubbard_model.qubitization,
+        path_stem='hubbard_model',
+        bloq_specs=[
+            qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard._SELECT_HUBBARD_DOC,
+            qualtran.bloqs.chemistry.hubbard_model.qubitization.prepare_hubbard._PREPARE_HUBBARD,
+        ],
+    ),
+    NotebookSpecV2(
+        title='Qubitized Hubbard Select',
+        module=qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard,
+        bloq_specs=[
+            qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard._SELECT_HUBBARD_DOC,
+            qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard._HUBBARD_MAJORANNA_OPERATOR_DOC,
+            qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard._HUBBARD_SPIN_UP_Z_DOC,
+        ],
+    ),
     NotebookSpecV2(
         title='Givens Rotations',
         module=qualtran.bloqs.chemistry.quad_fermion.givens_bloq,
@@ -930,15 +948,6 @@
             qualtran.bloqs.state_preparation.prepare_uniform_superposition._PREP_UNIFORM_DOC
         ],
     ),
-    NotebookSpecV2(
-        title='Qubitized Hubbard Model',
-        module=qualtran.bloqs.chemistry.hubbard_model.qubitization,
-        path_stem='hubbard_model',
-        bloq_specs=[
-            qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard._SELECT_HUBBARD,
-            qualtran.bloqs.chemistry.hubbard_model.qubitization.prepare_hubbard._PREPARE_HUBBARD,
-        ],
-    ),
     NotebookSpecV2(
         title='Apply to Lth Target',
         module=qualtran.bloqs.multiplexers.apply_gate_to_lth_target,

dev_tools/requirements/deps/runtime.txt

@@ -2,9 +2,9 @@
 attrs>=23.2.0
 cachetools>=5.3
 networkx
-numpy~=1.26
+numpy>=1.26,<3.0
 sympy
-cirq-core>=1.4.0,<1.6
+cirq-core~=1.4
 fxpmath
 galois
 
@@ -46,4 +46,7 @@ protobuf
 # typing
 typing_extensions>=4.10.0
 
+# rotation_synthesis
+mpmath
+
 # Note: use `pipreqs` to generate a list of dependencies based on the imports in our files.