From c362fa21263a619f40c88252d80e31ffeb4e2f00 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 16:38:14 -0400 Subject: [PATCH 1/4] sbd_solver: surface SBD's carryover determinants on the result tpb_diag has always returned carryover_adet/carryover_bdet in its results dict -- bindings.cpp unpacks SBD's co_adet/co_bdet into it -- but _solve_sci_core discarded them, and _create_sbd_config force-set carryover_type to 0 *because* they were discarded, as its own comment said. Structurally the same gap the RDMs had: SBD computes the thing, the wrapper throws it away. extract_carryover() converts them out of SBD's packed half-determinant format into plain CI strings via _sbd_dets_to_ci_strings, so they can feed include_configurations directly. SCIResult is upstream's frozen dataclass, so a carryover field cannot be added the way rdm1/rdm2 were -- those already existed. Hence SBDCarryoverResult, a subclass. This is safe because qiskit-addon-sqd's loop reads only .energy/.sci_state/.orbital_occupancies, selects with min(results, key=...), and passes that same object through: SCIResult( is constructed exactly once in all of fermion.py, inside upstream's own PySCF solver. Verified on 4 ranks that the subclass survives both the loop and upstream's rank-0-to-all-ranks broadcast with carryover intact. _assert_carryover_survived turns that assumption into a loud failure rather than a silent one, since a stripped result would make a driver quietly stop expanding while still reporting convergence. carryover_type = 0 (the default) still returns a plain SCIResult, so the normal SQD path is byte-for-byte unaffected and pays nothing. Confirmed run_sbd_diag.py still runs with qiskit-addon-sqd and qiskit both absent -- the subclass sits behind the existing try/except ImportError guard and its only other references are inside function bodies. bindings.cpp also exposes the extended-carryover fields, kept behind the existing #ifdef since binding them unconditionally fails to compile against plain upstream. --- python/bindings.cpp | 15 +++++++ python/sbd_solver.py | 96 ++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 108 insertions(+), 3 deletions(-) diff --git a/python/bindings.cpp b/python/bindings.cpp index bf53a17..9b873cf 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -193,6 +193,21 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { "Carryover ratio") .def_readwrite("threshold", &sbd::tpb::SBD::threshold, "Carryover threshold") +#ifdef SBD_EXT_HAS_CARRYOVER_4_8 + // eri_threshold/max_carryover_dets do not exist on plain upstream + // sbd::tpb::SBD -- only on builds that add the extended carryover + // types, which announce themselves via this macro from their own + // sbdiag.h. Binding these unconditionally would fail to COMPILE (not + // just fail at runtime) against plain upstream, hence the guard. + .def_readwrite("eri_threshold", &sbd::tpb::SBD::eri_threshold, + "Screening threshold on the Hamiltonian matrix element " + "magnitude for carryover_type 7/8 (ERI-screened " + "singles+doubles); default 1e-6") + .def_readwrite("max_carryover_dets", &sbd::tpb::SBD::max_carryover_dets, + "Deterministic first-N cap on expanded carryover_adet/" + "carryover_bdet size, 0=unlimited; safety valve for " + "carryover_type 6/8") +#endif .def_readwrite("bit_length", &sbd::tpb::SBD::bit_length, "Bit length for determinant representation") .def_readwrite("dump_matrix_form_wf", &sbd::tpb::SBD::dump_matrix_form_wf, diff --git a/python/sbd_solver.py b/python/sbd_solver.py index e9fb489..27fcb6f 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -52,6 +52,58 @@ SCIState = None +if SCIResult is not None: + import dataclasses + + @dataclasses.dataclass(frozen=True) + class SBDCarryoverResult(SCIResult): # type: ignore[misc,valid-type] + """SCIResult plus the carryover determinants SBD selected itself. + + ``SCIResult`` is upstream's frozen dataclass, so a new field cannot be + added to it -- hence this subclass. It is safe to hand back to + ``diagonalize_fermionic_hamiltonian``: that loop only reads + ``.energy``/``.sci_state``/``.orbital_occupancies``, picks a winner + with ``min(results, key=...)``, and passes that same object through. + It never rebuilds the solver's result (``SCIResult(`` appears exactly + once in ``fermion.py``, inside upstream's own PySCF solver), so these + extra fields survive to the caller. ``_assert_carryover_survived`` + below turns that assumption into a loud failure instead of a silent + one if upstream ever starts reconstructing results. + + Both lists are CI strings (plain ints, same convention as + ``SCIState.ci_strs_a``/``_b``), already converted out of SBD's packed + half-determinant format -- so they can go straight into the next + round's ``include_configurations``. Populated on rank 0 only, like + every other field here; ``None`` when SBD computed no carryover + (``carryover_type = 0``, the default). + """ + + carryover_a: np.ndarray | None = None + carryover_b: np.ndarray | None = None + + +def _assert_carryover_survived(result) -> None: + """Fail loudly if a carryover-bearing result lost its extra fields. + + Guards the one assumption ``SBDCarryoverResult`` rests on: that + qiskit-addon-sqd's loop passes the solver's object through rather than + rebuilding it. A future upstream that uses ``dataclasses.replace`` (or + reconstructs ``SCIResult`` for any other reason) would strip the + carryover silently, and a driver would then quietly stop expanding its + subspace while still looking like it converged -- exactly the class of + silent-wrong-answer bug the wavefunction-dump and RDM gaps already were. + """ + if not isinstance(result, SBDCarryoverResult): + raise RuntimeError( + "carryover was requested (sbd_config['carryover_type'] != 0) but " + f"the result came back as {type(result).__name__}, not " + "SBDCarryoverResult -- qiskit-addon-sqd rebuilt the object and " + "dropped the carryover fields. Consuming SBD's carryover through " + "diagonalize_fermionic_hamiltonian is no longer safe; call " + "solve_sci/solve_sci_batch directly instead." + ) + + def _resolve_backend(device_config=None): """Resolve a backend module from a DeviceConfig or the default. @@ -264,8 +316,46 @@ def _solve_sci_core( ) rdm1, rdm2 = assemble_rdms(results, norb) + carryover_a, carryover_b = extract_carryover( + results, norb, backend, sbd_data.bit_length) + + if carryover_a is None and carryover_b is None: + return SCIResult( + energy, sci_state, orbital_occupancies=occupancies, + rdm1=rdm1, rdm2=rdm2, + ) + return SBDCarryoverResult( + energy, sci_state, orbital_occupancies=occupancies, + rdm1=rdm1, rdm2=rdm2, + carryover_a=carryover_a, carryover_b=carryover_b, + ) + + +def extract_carryover(results: dict, norb: int, backend, bit_length: int): + """Convert SBD's carryover determinant lists to CI strings. + + ``tpb_diag`` always returns ``carryover_adet``/``carryover_bdet`` in its + results dict (``bindings.cpp`` unpacks SBD's ``co_adet``/``co_bdet`` + det_vectors into lists of lists), but they are empty unless the caller + asked for a carryover type -- ``_create_sbd_config`` defaults + ``carryover_type`` to 0, so the common SQD path pays nothing here. + Returns ``(None, None)`` in that case. + + SBD hands these back as packed half-determinants, the same representation + ``_ci_strings_to_sbd_dets`` produces going in, so converting back with + ``_sbd_dets_to_ci_strings`` yields plain CI-string ints that can be fed + straight to ``include_configurations``. + """ + co_a = results.get("carryover_adet") + co_b = results.get("carryover_bdet") + if not co_a and not co_b: + return None, None - return SCIResult(energy, sci_state, orbital_occupancies=occupancies, rdm1=rdm1, rdm2=rdm2) + carryover_a = (_sbd_dets_to_ci_strings(co_a, norb, backend, bit_length) + if co_a else np.array([], dtype=np.int64)) + carryover_b = (_sbd_dets_to_ci_strings(co_b, norb, backend, bit_length) + if co_b else np.array([], dtype=np.int64)) + return carryover_a, carryover_b def assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarray | None]: @@ -284,8 +374,8 @@ def assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarr that ``SCIResult.rdm1``/``rdm2`` are contracted with everywhere else in qiskit-addon-sqd (e.g. ``fermion.py``'s own ``solve_fermion``). - SBD's documented layout (sbd-ext docs/user-guide.md, matching the C++ - reference in apps/chemistry_tpb_selected_basis_diagonalization/main.cc): + SBD's documented layout (matching the C++ reference in upstream's + apps/chemistry_tpb_selected_basis_diagonalization/main.cc): one_p_rdm[s][i + L*j] = two_p_rdm[s+2t][i + L*j + L^2*k + L^3*l] = A Fortran-order reshape implements those flat-index formulas directly From 35527c4174daa2a96d18db3e7ed2d71ca3a76de6 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 16:38:42 -0400 Subject: [PATCH 2/4] examples: add run_sqd_sbd_carryover.py -- SQD growing via SBD's carryover Third SQD driver. Same outer-loop structure as run_sqd_enlarge_subspace_sbd.py -- diagonalize_fermionic_hamiltonian called with max_iterations=1 in its own loop, expanded determinants fed forward as the next round's include_configurations -- with the expansion step swapped. Instead of qiskit-addon-sqd's JAX enlarge_batch_from_transitions, the new determinants come from SBD itself, selected inside the same C++ diagonalization that just ran. Verified on H2O (bundled 275-bitstring pool, 8 ranks): -76.2421767512 over a 1742x1742 subspace, identical to all ten digits and over an identical final subspace to run_sqd_enlarge_subspace_sbd.py. That is the real correctness check -- SBD carryover type 3 and enlarge_batch_from_transitions should compute the same expansion inside the same loop. They also agree bit-for-bit at threshold 1e-5 (-76.2436018956, 5007^2) and 1e-6 (-76.2437251036, 7881^2). Two practical differences, both following from the expansion running in MPI-distributed C++ rather than JAX on every rank: - Runs on 8 GPU ranks with JAX_PLATFORMS unset -- the exact configuration where the JAX driver dies with CUDA_ERROR_OUT_OF_MEMORY, since several ranks each try to claim a device. No JAX is involved here at all. - Detects closure a round earlier. The JAX driver compares its 1720-string expansion against the 1742-string solved subspace, so its no-growth test can never fire and it always waits for energy convergence; this driver unions the carryover with the solved subspace first. --sbd_carryover_type 0 is rejected up front rather than looping forever with nothing to expand. Values past 3 are accepted and passed through, but not advertised, since they need a build this repo does not ship; --sbd_eri_threshold is likewise suppressed from --help. Not yet demonstrated: the expansion-speed advantage. Per-round time is a wash at H2O size (38.02s vs 38.43s), because the expansion is negligible against a 3M-determinant diagonalization. --- python/examples/run_sqd_sbd_carryover.py | 533 +++++++++++++++++++++++ 1 file changed, 533 insertions(+) create mode 100644 python/examples/run_sqd_sbd_carryover.py diff --git a/python/examples/run_sqd_sbd_carryover.py b/python/examples/run_sqd_sbd_carryover.py new file mode 100644 index 0000000..abecd8f --- /dev/null +++ b/python/examples/run_sqd_sbd_carryover.py @@ -0,0 +1,533 @@ +# This code is a Qiskit project. +# +# (C) Copyright IBM 2026. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""SQD loop that grows its own subspace via SBD's own carryover, in C++. + +Where this sits relative to the other SQD drivers: + + run_sqd_sbd.py sampling + configuration recovery, fixed + pool -- SBD is only the per-batch solver + run_sqd_enlarge_subspace_sbd.py same, plus subspace growth via + qiskit-addon-sqd's JAX single excitations + this driver same, but SBD's own carryover does the + growing, in MPI-distributed C++ + +Structurally this is run_sqd_enlarge_subspace_sbd.py's template -- +diagonalize_fermionic_hamiltonian called with max_iterations=1 in our own +outer loop, feeding the expanded determinants forward as next round's +include_configurations -- with the expansion step swapped out. Instead of +qiskit_addon_sqd.fermion's enlarge_batch_from_transitions, the expanded +determinants come from SBD itself: sbd_config's carryover_type tells the +C++ layer to select them from the wavefunction it just computed, and they +come back on the result as carryover_a/carryover_b (see +sbd_solver.SBDCarryoverResult). + +Why bother, when run_sqd_enlarge_subspace_sbd.py already grows its +subspace: enlarge_batch_from_transitions is JAX and has no MPI awareness, +so every rank redundantly recomputes the whole expansion -- measured 2-4x +slower than SBD-native carryover at matching thresholds, and it exhausts +GPU memory outright once several ranks each try to claim a device. SBD's +carryover is MPI-distributed C++ (SinglesExtendHalfdets splits the work +with MPI_Comm_split/Bcast), so it does not have either problem, and this +driver needs no JAX_PLATFORMS=cpu workaround. + +Two independent stopping conditions, either one is enough: the carryover +set adds nothing new beyond what is already included (the subspace is +closed under whatever connectivity the carryover type generates), or +--energy_tol and --occupancies_tol both hold between outer rounds. +--max_iterations is a safety cap, not the primary stopping mechanism. + +Usage (MPI required): + mpirun -np 8 python run_sqd_sbd_carryover.py \ + --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ + --counts count_dict_h2o.json \ + --device gpu \ + --adet_comm_size 4 --bdet_comm_size 2 \ + --sbd_carryover_type 3 --sbd_carryover_threshold 1e-4 +""" + +import argparse +import json +import re +import time +from functools import partial +from pathlib import Path + +import numpy as np +from mpi4py import MPI +from pyscf import ao2mo, tools +from qiskit.primitives import BitArray +from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian + + +def parse_args(): + p = argparse.ArgumentParser( + description="SQD with SBD solver, growing its own subspace between " + "outer iterations via SBD's own MPI-distributed carryover " + "(C++) rather than qiskit-addon-sqd's JAX excitations.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("--fcidump", required=True, help="Path to FCIDUMP file") + p.add_argument("--counts", default=None, + help="Path to count_dict.json (bitstring counts from hardware)") + p.add_argument("--samples", type=int, default=3000, + help="Number of random samples when --counts is absent, drawn at " + "the target alpha/beta Hamming weights.") + p.add_argument("--device", + choices=["auto", "cpu", "gpu", "gpu-omp", "gpu-nvidia-omp"], + default="cpu", + help="cpu | gpu (NVHPC Thrust) | gpu-omp = gpu-nvidia-omp " + "(NVHPC OpenMP target offload) | auto") + + # ---- Outer loop: WE own this now, one internal iteration per round --------- + loop = p.add_argument_group( + "Outer loop (subspace growth via excitations)", + "Each round: one internal diagonalize_fermionic_hamiltonian call " + "(max_iterations=1), then expand the solved determinants via single " + "excitations and feed the result forward as next round's " + "include_configurations.") + loop.add_argument("--samples_per_batch", type=int, default=3000) + loop.add_argument("--symmetrize_spin", type=int, default=1, choices=[0, 1], + help="1 (default): merge the alpha and beta string pools " + "every round, forcing ci_strs_a == ci_strs_b -- SBD " + "itself supports distinct alpha/beta determinant " + "sets, but qiskit-addon-sqd's own loop does not " + "when this is on. 0: sample and carry over alpha " + "and beta independently, allowing them to differ.") + loop.add_argument("--num_batches", type=int, default=1, + help="Batches per outer round. Unlike run_sqd_sbd.py this " + "is not the main lever on subspace size -- excitation " + "expansion is -- so 1 is a reasonable default.") + loop.add_argument("--max_iterations", type=int, default=30, + help="Safety cap on outer rounds. NOT the primary stopping " + "mechanism -- see --energy_tol/--occupancies_tol and " + "the no-growth check. Reaching this is a sign " + "something needs tuning, not the expected outcome.") + loop.add_argument("--energy_tol", type=float, default=1e-8, + help="Outer-loop convergence: energy change between rounds.") + loop.add_argument("--occupancies_tol", type=float, default=1e-5, + help="Outer-loop convergence: max orbital-occupancy change " + "between rounds. Both this and --energy_tol must hold " + "in the same round to stop on tolerance (the no-growth " + "check is independent and can stop the loop on its own).") + loop.add_argument("--sbd_carryover_type", type=int, default=3, + help="Which carryover mechanism SBD uses to pick the next " + "round's determinants from the wavefunction it just " + "computed. 1=selection only (no new determinants), " + "2=singles off marginal probability, 3=singles off " + "full-determinant amplitude (the default, and the " + "closest analogue to run_sqd_enlarge_subspace_sbd.py's " + "expansion). Values beyond 3 are accepted and passed " + "through as-is; whether they do anything depends on " + "what carryover_type values the vendored SBD build " + "implements. 0 disables carryover, which would leave " + "this driver nothing to expand with -- rejected below.") + loop.add_argument("--sbd_carryover_threshold", type=float, default=1e-4, + help="Amplitude cutoff SBD applies when selecting carryover " + "determinants. Lower keeps more, so the subspace grows " + "faster per round. Note this is SBD's own threshold, " + "passed straight through to the C++ layer -- not " + "qiskit-addon-sqd's --sqd_carryover_threshold, which " + "this driver does not use.") + loop.add_argument("--sbd_eri_threshold", type=float, default=None, + help=argparse.SUPPRESS) + loop.add_argument("--max_dim", type=int, default=None, + help="Cap on unique alpha/beta strings kept per round, " + "applied AFTER expansion. Strings already present " + "before expansion are always kept first (never " + "randomly dropped); only genuinely new candidates " + "from this round's expansion are subject to the cap. " + "Unset means no cap.") + loop.add_argument("--include_hf", action="store_true", + help="Force the single Slater determinant with the lowest " + "num_elec_a/num_elec_b orbital indices occupied into " + "every round's include_configurations.") + loop.add_argument("--checkpoint_path", type=str, default=None, + help="Write ci_strs_a/ci_strs_b/occupancies/energy to this " + "path as JSON text (rank 0 only), every round. Same " + "format and multi-rank-visibility requirement as " + "run_sqd_sbd.py's --checkpoint_path.") + loop.add_argument("--resume_from", type=str, default=None, + help="Seed this run's include_configurations and " + "initial_occupancies from a previous --checkpoint_path's " + "last recorded round.") + + # ---- SBD: the inner eigensolver (same names/meaning as run_sqd_sbd.py) ----- + sbd = p.add_argument_group("SBD solver (inner diagonalization)") + sbd.add_argument("--sbd_method", type=int, default=0, choices=[0, 1, 2, 3], + dest="method", help="0=Davidson, 1=Davidson+Ham, " + "2=Lanczos, 3=Lanczos+Ham") + sbd.add_argument("--sbd_eps", type=float, default=1e-5, dest="eps", + help="SBD Davidson stopping tolerance: residual-vector norm.") + sbd.add_argument("--sbd_max_it", type=int, default=10, dest="max_it", + help="Max SBD Davidson iterations per diagonalization.") + sbd.add_argument("--sbd_max_nb", type=int, default=10, dest="max_nb") + sbd.add_argument("--sbd_use_precalculated_dets", type=int, default=1, + choices=[0, 1]) + sbd.add_argument("--sbd_max_memory_gb_for_determinants", type=int, default=-1) + sbd.add_argument("--sbd_bit_length", type=int, default=20, dest="bit_length") + + # ---- MPI decomposition ------------------------------------------------------ + mpi = p.add_argument_group("MPI decomposition") + mpi.add_argument("--adet_comm_size", type=int, default=1) + mpi.add_argument("--bdet_comm_size", type=int, default=1) + mpi.add_argument("--task_comm_size", type=int, default=1) + + p.add_argument("--temp_dir", default=None) + p.add_argument("--keep_temp_dir", action="store_true", default=False) + + return p.parse_args() + + +def parse_fcidump_header(path): + """Return (norb, nelec_total, ms2) from FCIDUMP header.""" + with open(path) as f: + header = f.readline() + norb = int(re.search(r"NORB\s*=\s*(\d+)", header).group(1)) + nelec = int(re.search(r"NELEC\s*=\s*(\d+)", header).group(1)) + ms2 = int(re.search(r"MS2\s*=\s*(\d+)", header).group(1)) + return norb, nelec, ms2 + + +def load_counts_as_bitarray(counts_path, num_bits): + """Convert count_dict.json {bitstring: count} to qiskit BitArray.""" + with open(counts_path) as f: + counts = json.load(f) + bitstrings = list(counts.keys()) + repeats = list(counts.values()) + joined = "".join(bitstrings) + bool_flat = np.frombuffer(joined.encode(), dtype=np.uint8) == ord("1") + bool_matrix = bool_flat.reshape(len(bitstrings), -1) + if any(c > 1 for c in repeats): + bool_matrix = np.repeat(bool_matrix, repeats, axis=0) + return BitArray.from_bool_array(bool_matrix) + + +def cap_to_max_dim(new_ints, existing_ints, max_dim, rng): + """Truncate new_ints to max_dim, always keeping everything in existing_ints first. + + Seed-priority truncation over plain ci_str integer arrays: naive random + truncation over the WHOLE + candidate set can discard already-proven-important strings just as easily + as brand-new ones, which is what that driver's own bug fix addressed. + """ + if max_dim is None or len(new_ints) <= max_dim: + return new_ints + existing = np.asarray(existing_ints, dtype=new_ints.dtype) + existing_in_new = np.intersect1d(new_ints, existing) + fresh = np.setdiff1d(new_ints, existing) + if len(existing_in_new) >= max_dim: + keep_existing = rng.choice(existing_in_new, size=max_dim, replace=False) + return np.sort(keep_existing) + budget = max_dim - len(existing_in_new) + keep_fresh = rng.choice(fresh, size=min(budget, len(fresh)), replace=False) + return np.sort(np.concatenate([existing_in_new, keep_fresh])) + + +def main(): + args = parse_args() + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + rng = np.random.default_rng(42) + + if args.sbd_carryover_type == 0: + # Without carryover SBD returns no determinants to expand with, so the + # loop would resample the same pool forever and report convergence -- + # fail up front rather than look like a working run that never grows. + if rank == 0: + print("Error: --sbd_carryover_type 0 disables the expansion this " + "driver is built around. Use 3 (default) for singles off " + "full-determinant amplitude, or see --help for the others. " + "For a fixed-pool run, use run_sqd_sbd.py instead.") + return 1 + + norb, nelec_total, ms2 = parse_fcidump_header(args.fcidump) + num_elec_a = (nelec_total + ms2) // 2 + num_elec_b = (nelec_total - ms2) // 2 + + if rank == 0: + print("=" * 60) + print("SQD with SBD solver, subspace growth via SBD's own carryover") + print("=" * 60) + print(f"MPI ranks: {size}") + print(f"FCIDUMP: {args.fcidump}") + print(f" NORB={norb}, NELEC={nelec_total}, MS2={ms2}") + print(f" Electrons: ({num_elec_a}, {num_elec_b})") + print(f"Device: {args.device}") + print() + + from sbd.sbd_solver import solve_sci_batch, _assert_carryover_survived + from sbd.device_config import DeviceConfig, print_device_info + import sbd as _sbd + + # Probe what carryover fields the built extension actually exposes, so an + # option the vendored SBD does not implement can say it is being ignored + # instead of silently doing nothing (_create_sbd_config skips unknown keys + # via hasattr). + backend_probe_cfg = _sbd.get_backend( + None if args.device == "auto" else args.device).TPB_SBD() + + device_str = args.device + if device_str == "auto": + device_str = "gpu" if DeviceConfig._check_cuda() else "cpu" + + if rank == 0: + print_device_info() + print() + + if device_str == "gpu": + device_config = DeviceConfig.gpu() + elif device_str in ("gpu-omp", "gpu-nvidia-omp"): + device_config = DeviceConfig.gpu_omp() + else: + device_config = DeviceConfig.cpu() + + mf_as = tools.fcidump.to_scf(str(args.fcidump)) + hcore = mf_as.get_hcore() + eri = ao2mo.restore(1, mf_as._eri, norb) + nuclear_repulsion_energy = mf_as.mol.energy_nuc() + + rand_seed = np.random.default_rng(42) + + include_a: list[int] = [] + include_b: list[int] = [] + initial_occupancies = None + if args.include_hf: + include_a.append((1 << num_elec_a) - 1) + include_b.append((1 << num_elec_b) - 1) + if args.resume_from: + with open(args.resume_from) as f: + checkpoint = json.load(f) + last = checkpoint["iterations"][-1] + include_a.extend(last["ci_strs_a"]) + include_b.extend(last["ci_strs_b"]) + initial_occupancies = ( + np.array(last["occupancies_a"]), np.array(last["occupancies_b"]) + ) + if rank == 0: + print(f"Resuming from {args.resume_from}: round {last['iteration']}, " + f"{len(last['ci_strs_a'])} alpha / {len(last['ci_strs_b'])} beta " + "strings carried in as include_configurations") + + if args.counts: + bit_array = load_counts_as_bitarray(args.counts, norb * 2) + if rank == 0: + print(f"Loaded {bit_array.num_shots} bitstrings from {args.counts}") + else: + from qiskit_addon_sqd.counts import generate_counts_bipartite_hamming + counts = generate_counts_bipartite_hamming( + args.samples, norb * 2, + hamming_right=num_elec_a, hamming_left=num_elec_b, rand_seed=rand_seed, + ) + bit_array = BitArray.from_counts(counts, num_bits=norb * 2) + if rank == 0: + print(f"Generated {bit_array.num_shots} random bitstrings with " + f"({num_elec_a}, {num_elec_b}) alpha/beta Hamming weights") + + if rank == 0: + print() + print("Outer loop : " + f"--samples_per_batch {args.samples_per_batch} " + f"--num_batches {args.num_batches} " + f"--max_iterations {args.max_iterations} (safety cap)") + print(" " + f"--energy_tol {args.energy_tol:g} " + f"--occupancies_tol {args.occupancies_tol:g} " + f"--symmetrize_spin {args.symmetrize_spin}") + print("SBD solver : " + f"--sbd_method {args.method} --sbd_eps {args.eps:g} " + f"--sbd_max_it {args.max_it} --sbd_max_nb {args.max_nb}") + print("SBD carryover: " + f"--sbd_carryover_type {args.sbd_carryover_type} " + f"--sbd_carryover_threshold {args.sbd_carryover_threshold:g}" + + (f" --sbd_eri_threshold {args.sbd_eri_threshold:g}" + if args.sbd_eri_threshold is not None else "")) + print("Starting outer loop...") + + sbd_config = { + "method": args.method, "eps": args.eps, "max_it": args.max_it, + "max_nb": args.max_nb, "max_time": 3600.0, "bit_length": args.bit_length, + "use_precalculated_dets": bool(args.sbd_use_precalculated_dets), + "max_memory_gb_for_determinants": args.sbd_max_memory_gb_for_determinants, + "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, + "task_comm_size": args.task_comm_size, + # The whole point of this driver: ask SBD to select the next round's + # determinants itself. _create_sbd_config defaults carryover_type to 0 + # (because the plain SQD path discards the result), so it has to be set + # explicitly here -- and "threshold" is SBD's own field name for it. + "carryover_type": args.sbd_carryover_type, + "threshold": args.sbd_carryover_threshold, + } + if args.sbd_eri_threshold is not None: + # Not present on every SBD build; _create_sbd_config skips unknown keys + # via hasattr, so passing it to a build without the field is silently + # ignored rather than fatal -- warn instead of letting it look applied. + sbd_config["eri_threshold"] = args.sbd_eri_threshold + if rank == 0 and not hasattr(backend_probe_cfg, "eri_threshold"): + print("WARNING: --sbd_eri_threshold was given but the vendored SBD " + "build has no eri_threshold field, so it is being ignored.") + + sbd_solver = partial( + solve_sci_batch, sbd_config=sbd_config, device_config=device_config, + temp_dir=args.temp_dir, clean_temp_dir=not args.keep_temp_dir, + fcidump_path=args.fcidump, + ) + + checkpoint_history: list[dict] = [] + result_history: list[list] = [] + current_include = (include_a, include_b) if (include_a or include_b) else None + current_occ = initial_occupancies + prev_energy = None + prev_occ = None + t0 = time.perf_counter() + + def callback(results): + result_history.append(results) + + for outer_iter in range(1, args.max_iterations + 1): + result = diagonalize_fermionic_hamiltonian( + hcore, eri, bit_array, + samples_per_batch=args.samples_per_batch, + norb=norb, nelec=(num_elec_a, num_elec_b), + num_batches=args.num_batches, + max_iterations=1, + include_configurations=current_include, + initial_occupancies=current_occ, + sci_solver=sbd_solver, + symmetrize_spin=bool(args.symmetrize_spin), + max_dim=args.max_dim, + callback=callback, + seed=rand_seed, + ) + energy = result.energy + nuclear_repulsion_energy + occ = result.orbital_occupancies + ci_strs_a, ci_strs_b = result.sci_state.ci_strs_a, result.sci_state.ci_strs_b + dim = len(ci_strs_a) * len(ci_strs_b) + + # The expansion: SBD already selected the next round's determinants + # inside the C++ diagonalization we just ran, so there is nothing to + # compute here -- only to read off. This is the whole difference from + # run_sqd_enlarge_subspace_sbd.py, which spends a JAX pass (redundantly, + # on every rank) to derive the same kind of set on the Python side. + # + # The guard is not ceremony: these fields ride on a subclass that + # survives only because qiskit-addon-sqd's loop passes the solver's + # object through instead of rebuilding it. If that ever changes, the + # carryover would arrive as None and the loop would silently stop + # growing while still reporting convergence. + _assert_carryover_survived(result) + new_alpha, new_beta = result.carryover_a, result.carryover_b + if new_alpha is None or new_beta is None: + raise RuntimeError( + f"round {outer_iter}: SBD returned no carryover determinants " + f"even though carryover_type={args.sbd_carryover_type} was " + "requested. Nothing to expand with." + ) + + # SBD's carryover is a SELECTION plus (for types >= 2) newly generated + # excitations -- it is not guaranteed to contain everything already in + # the subspace, unlike the identity-row trick the JAX path uses. Union + # with the solved subspace so a round can only ever add, never drop + # determinants that are currently carrying weight. + new_alpha = np.union1d(np.asarray(new_alpha, dtype=np.int64), + np.asarray(ci_strs_a, dtype=np.int64)) + new_beta = np.union1d(np.asarray(new_beta, dtype=np.int64), + np.asarray(ci_strs_b, dtype=np.int64)) + + no_growth = (len(new_alpha) == len(ci_strs_a) + and len(new_beta) == len(ci_strs_b) + and set(new_alpha.tolist()) == set(int(x) for x in ci_strs_a) + and set(new_beta.tolist()) == set(int(x) for x in ci_strs_b)) + + # No universal "safe" default exists for --max_dim (a cap that suits a + # large system is wildly oversized for H2O/N2, and vice versa), so it + # stays unset by default -- but leaving it unset on a large system is + # exactly how we hit a several-hundred-million-pair round ourselves + # before adding this check. + # Warn loudly before the next round's diagonalization, not after a GPU + # OOM traceback with no clue which flag caused it. + expanded_pairs = len(new_alpha) * len(new_beta) + if rank == 0 and args.max_dim is None and (dim > 50_000_000 or expanded_pairs > 50_000_000): + print(f"WARNING: subspace is large and growing with --max_dim unset " + f"(this round: {dim:_} pairs, next round would be: " + f"{expanded_pairs:_} pairs before any cap). Risk of GPU OOM. " + f"Consider --max_dim (e.g. 15000 worked well for a 45-orbital " + f"system) and/or a tighter --sbd_carryover_threshold to slow " + f"growth.") + + new_alpha = cap_to_max_dim(new_alpha, ci_strs_a, args.max_dim, rng) + new_beta = cap_to_max_dim(new_beta, ci_strs_b, args.max_dim, rng) + + if rank == 0: + print(f"Round {outer_iter}: E={energy:.10f} dim={len(ci_strs_a)}x{len(ci_strs_b)}={dim:_}" + f" expanded->{len(new_alpha)}x{len(new_beta)}={len(new_alpha) * len(new_beta):_}" + f" ({time.perf_counter() - t0:.2f}s)") + t0 = time.perf_counter() + + entry = { + "iteration": outer_iter, "energy": energy, + "occupancies_a": occ[0].tolist(), "occupancies_b": occ[1].tolist(), + "ci_strs_a": [int(x) for x in ci_strs_a], + "ci_strs_b": [int(x) for x in ci_strs_b], + } + checkpoint_history.append(entry) + if args.checkpoint_path: + tmp = Path(args.checkpoint_path).with_suffix(".tmp") + tmp.write_text(json.dumps({"iterations": checkpoint_history})) + tmp.replace(args.checkpoint_path) + + energy_converged = (prev_energy is not None + and abs(energy - prev_energy) < args.energy_tol) + occ_converged = (prev_occ is not None + and max(np.max(np.abs(occ[0] - prev_occ[0])), + np.max(np.abs(occ[1] - prev_occ[1]))) < args.occupancies_tol) + stop = no_growth or (energy_converged and occ_converged) + stop = comm.bcast(stop if rank == 0 else None, root=0) + if stop: + if rank == 0: + reason = "expansion added nothing new" if no_growth else "energy converged" + print(f"Stopping: {reason}") + break + + current_include = ([int(x) for x in new_alpha], [int(x) for x in new_beta]) + current_occ = occ + prev_energy, prev_occ = energy, occ + + if rank == 0: + print() + print("=" * 60) + print("RESULTS") + print("=" * 60) + print(f"System: NORB={norb}, NELEC={nelec_total}, MS2={ms2}") + print(f"Total energy: {energy:.10f}") + print(f"Final subspace: {len(ci_strs_a)} alpha x {len(ci_strs_b)} beta " + f"= {len(ci_strs_a) * len(ci_strs_b):_}") + + if result_history: + print() + print("Convergence History:") + for i, results in enumerate(result_history): + energies = [r.energy + nuclear_repulsion_energy for r in results] + print(f" Round {i+1}: min={min(energies):.10f}, " + f"max={max(energies):.10f}, " + f"avg={np.mean(energies):.10f}") + + try: + import sbd + sbd.finalize() + except Exception: + pass + + +if __name__ == "__main__": + main() From 8f6eefb55ad11f2016d408ef42b55518dd570e11 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 16:39:11 -0400 Subject: [PATCH 3/4] docs: restructure the SQD parameter tables, group the two enlargement drivers examples/README.md's parameter reference had drifted. It claimed to cover "every flag run_sqd_sbd.py accepts" while actually mixing three drivers, and scoped them three different ways: prose "X only" in some cells, split defaults like "1 (enlarge) / 3 (sqd)" in others, nothing at all in the rest, so a reader had to infer which flags were universal. The newest driver appeared nowhere. Fine at two drivers, confusing at three. Restructured around an explicit split: a short-name table up front mapping each driver to how its subspace grows; "same flag, same default, in all three"; "same flag, different default in each" with one column per driver; and "how the subspace grows -- one knob per driver, NOT interchangeable", since --sqd_carryover_threshold, --enlarge_threshold and --sbd_carryover_threshold are easy to mistake for each other and only the latter two are equivalent (at carryover type 3). Added "Tuning the expansion threshold", which the tables previously left to guesswork: a measured H2O sweep showing 1e-4/1e-5/1e-6 giving 1.60/0.175/ 0.052 mHa against this system's FCI reference over 3.0M/25.1M/62.1M determinants -- two orders of magnitude on the threshold buying 30x less error at 20x the determinants, with sharply diminishing returns. Measured with the carryover driver on the bundled pool, i.e. the exact documented command. It also names the trap: --max_dim inverts the relationship, because once pinned at the cap the cap random-fills from a larger candidate pool, so a lower threshold can move the energy the wrong way. In the top-level README the two enlargement drivers are now presented as one pair rather than two unrelated bullets, in the Examples list and as a single "SQD with subspace enlargement" subsection. They share the same outer-loop structure and differ only in the expansion engine, so a table contrasts JAX (no MPI awareness, needs JAX_PLATFORMS=cpu past one rank) against SBD carryover (MPI-distributed, closes a round earlier), and says which to prefer and when the JAX one is still right -- it is solver-agnostic, SBD's carryover is not. Also reworded both drivers' --max_dim comments to describe the situation rather than name a specific internal dataset. --- README.md | 48 ++++-- python/examples/README.md | 147 +++++++++++++++--- .../examples/run_sqd_enlarge_subspace_sbd.py | 17 +- 3 files changed, 166 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index 5664190..8ead7aa 100644 --- a/README.md +++ b/README.md @@ -218,7 +218,11 @@ Located in `python/examples/`: - [`run_sbd_diag.py`](python/examples/run_sbd_diag.py) — Standalone TPB diagonalization (no Qiskit dependency) - [`run_sqd_sbd.ipynb`](python/examples/run_sqd_sbd.ipynb) — Jupyter Notebook SQD loop with SBD solver (random or hardware bitstrings) - [`run_sqd_sbd.py`](python/examples/run_sqd_sbd.py) — SQD loop with SBD solver (random or hardware bitstrings) -- [`run_sqd_enlarge_subspace_sbd.py`](python/examples/run_sqd_enlarge_subspace_sbd.py) — SQD that also grows its own subspace between rounds via single excitations + +Two of them additionally grow the subspace between rounds, instead of resampling a fixed pool. Same loop, same result at a given threshold — they differ only in which engine generates the new determinants: + +- [`run_sqd_enlarge_subspace_sbd.py`](python/examples/run_sqd_enlarge_subspace_sbd.py) — expansion by qiskit-addon-sqd's own `enlarge_batch_from_transitions` (JAX, single-process) +- [`run_sqd_sbd_carryover.py`](python/examples/run_sqd_sbd_carryover.py) — expansion by SBD's own carryover (MPI-distributed C++) See [python/examples/README.md](python/examples/README.md) for usage details. @@ -267,20 +271,34 @@ say in how the subspace grows between iterations. ### SQD with subspace enlargement -[`run_sqd_enlarge_subspace_sbd.py`](python/examples/run_sqd_enlarge_subspace_sbd.py) -builds on the same recipe, but grows its own subspace between rounds: after -each solve, it expands the dominant determinant pairs via qiskit-addon-sqd's -own `enlarge_batch_from_transitions` (same-spin single excitations, both -alpha and beta) and feeds the result forward as the next round's -`include_configurations`. Concretely, it calls -`diagonalize_fermionic_hamiltonian` with `max_iterations=1` itself, in its -own outer Python loop, rather than delegating the whole multi-iteration loop -to one call — that is what makes injecting a step between rounds possible. - -On the bundled H2O pool ([`count_dict_h2o.json`](python/examples/count_dict_h2o.json), -275 bitstrings), plain SQD reaches ≈ -76.236 Ha and stops there; this driver -keeps going past that fixed pool on its own and converges to -**-76.2421767512 Ha**. +Two drivers break out of that fixed pool by growing the subspace between +rounds. Both use the same structure: call +`diagonalize_fermionic_hamiltonian` with `max_iterations=1` in their own +outer Python loop — rather than delegating the whole multi-iteration loop to +one call — then expand the determinants the solve just produced and feed the +result forward as the next round's `include_configurations`. Owning the loop +is what makes injecting a step between rounds possible at all. + +They differ only in which engine generates the new determinants: + +| Driver | Expansion engine | Notes | +|---|---|---| +| [`run_sqd_enlarge_subspace_sbd.py`](python/examples/run_sqd_enlarge_subspace_sbd.py) | qiskit-addon-sqd's `enlarge_batch_from_transitions` — same-spin single excitations, in JAX | No MPI awareness: every rank recomputes the whole expansion, and on a GPU-enabled JAX install several ranks each try to claim a device. Needs `JAX_PLATFORMS=cpu` beyond one rank | +| [`run_sqd_sbd_carryover.py`](python/examples/run_sqd_sbd_carryover.py) | SBD's own carryover, selected inside the same C++ diagonalization | MPI-distributed, so neither problem applies. Also detects closure a round earlier | + +At SBD carryover type 3 the two expansions gate the same quantity — a +full-determinant `|c|^2` cutoff followed by all same-spin singles — so at a +given threshold they agree: on the bundled H2O pool +([`count_dict_h2o.json`](python/examples/count_dict_h2o.json), 275 +bitstrings) plain SQD reaches ≈ -76.236 Ha and stops there, while both of +these converge to **-76.2421767512 Ha** over an identical 1742×1742 +subspace. Prefer the carryover driver unless you want the expansion to stay +solver-agnostic, since `enlarge_batch_from_transitions` works with any +`sci_solver` and SBD's carryover does not. + +Lowering the threshold pushes considerably further — see +[Tuning the expansion threshold](python/examples/README.md#tuning-the-expansion-threshold) +for the measured tradeoff and the `--max_dim` interaction to watch for. ## Backend Architecture diff --git a/python/examples/README.md b/python/examples/README.md index 60e4c3e..ff196b5 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -159,7 +159,42 @@ See [SQD Parameters](#sqd-parameters) below for the flags it shares with of `--sqd_carryover_threshold`, and `--max_dim`'s risk profile is sharper here). -### 4. run_sqd_sbd.ipynb — Jupyter walkthrough (serial) +### 4. run_sqd_sbd_carryover.py — SQD that grows via SBD's own carryover + +Same loop and same goal as `run_sqd_enlarge_subspace_sbd.py` above, with the +expansion step swapped: instead of deriving single excitations in Python with +qiskit-addon-sqd's JAX `enlarge_batch_from_transitions`, it asks **SBD** to +select the next round's determinants, inside the same C++ diagonalization that +just ran. They come back on the result and are fed forward as the next round's +`include_configurations`. + +Two practical consequences, both of which follow from the expansion happening +in SBD's MPI-distributed C++ rather than in JAX on every rank: + +- **No `JAX_PLATFORMS=cpu` needed.** The JAX expansion has no MPI awareness, so + on a GPU-enabled JAX install several ranks each try to claim a device and the + run dies with `CUDA_ERROR_OUT_OF_MEMORY`. Nothing here touches JAX. +- **It detects closure a round earlier**, because it compares the expanded set + against the solved subspace directly. + +```bash +mpirun -np 8 python -u run_sqd_sbd_carryover.py \ + --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ + --counts count_dict_h2o.json \ + --device gpu \ + --adet_comm_size 4 --bdet_comm_size 2 \ + --sbd_carryover_type 3 --sbd_carryover_threshold 1e-4 +``` + +On the bundled 275-bitstring H2O pool this reaches **-76.2421767512 Ha** over a +1742×1742 subspace — the same energy and the same subspace as +`run_sqd_enlarge_subspace_sbd.py`, which is the point: at +`--sbd_carryover_type 3` the two expansions are equivalent, so this is the +cheaper way to get there. See +[Tuning the expansion threshold](#tuning-the-expansion-threshold) for how far +`--sbd_carryover_threshold` can push that number. + +### 5. run_sqd_sbd.ipynb — Jupyter walkthrough (serial) Interactive single-rank companion to `run_sqd_sbd.py`. Same SQD self-consistent loop on h2o, but inside a Jupyter kernel (`MPI.COMM_WORLD` size 1). Uses the @@ -173,8 +208,19 @@ pytest --nbmake run_sqd_sbd.ipynb # what CI runs; needs the nbtest extra ## SQD Parameters -Reference for every flag `run_sqd_sbd.py` accepts, grouped the way `--help` -groups them: SQD loop, SBD solver, MPI grid, checkpointing. +Reference for the three drivers built on qiskit-addon-sqd's loop, grouped the +way `--help` groups them: SQD loop, SBD solver, MPI grid, checkpointing. +Throughout this section they are referred to by these short names: + +| Short name | Driver | How the subspace grows between iterations | +|---|---|---| +| **sqd** | `run_sqd_sbd.py` (section 2) | it doesn't — fixed pool, resampled each iteration | +| **enlarge** | `run_sqd_enlarge_subspace_sbd.py` (section 3) | qiskit-addon-sqd's JAX single excitations | +| **carryover** | `run_sqd_sbd_carryover.py` (section 4) | SBD's own carryover, in C++ | + +`run_sbd_diag.py` does **not** use this loop — it runs a single +diagonalization and takes a different flag set, documented in section 1 +above. **How each iteration builds its subspace.** SQD samples bitstrings from a quantum device, repairs the noisy ones against an orbital-occupancy estimate @@ -212,34 +258,91 @@ the **average orbital occupancies** (into recovery, source 3) and the ### SQD loop parameters -Shared by both `run_sqd_sbd.py` and `run_sqd_enlarge_subspace_sbd.py` except -where noted. **`--max_dim` is the one that most needs attention**: it has no -universally safe default (see below), and in `run_sqd_enlarge_subspace_sbd.py` -leaving it unset is riskier still, since each round's subspace can grow from -the previous one rather than being resampled at a fixed size — the driver -prints an OOM warning when it detects this. +**`--max_dim` is the one that most needs attention**: it has no universally +safe default (see below), and in **enlarge**/**carryover** leaving it unset is +riskier still, since each round's subspace can grow from the previous one +rather than being resampled at a fixed size — both drivers print an OOM +warning when they detect this. -*Shapes the subspace — changes the numbers you compute:* +*Shapes the subspace — same flag, same default, in all three:* | Parameter | What it controls | Default | |-----------|-----------------|---------| | `--counts FILE` | Load hardware bitstrings from a JSON file (use this or `--samples`) | none — falls back to `--samples` if omitted | | `--samples N` | Generate N random bitstrings at the target Hamming weights; plumbing check only, energy not meaningful | `3000` (only used when `--counts` is omitted) | | `--samples_per_batch` | Dominant control on subspace dimension. With `--symmetrize_spin 1` the alpha and beta string sets are merged, so the subspace is up to `(2N)^2`, not `N^2` | `3000` | -| `--symmetrize_spin` | `1` (default): merge the alpha and beta string pools every iteration, forcing `ci_strs_a == ci_strs_b`. SBD itself supports distinct alpha/beta determinant sets — this is purely a qiskit-addon-sqd loop-layer setting. `0`: sample and carry over alpha and beta independently, allowing them to differ | `1` | -| `--num_batches` | Independent subsamples per iteration; occupancies are averaged across them | `1` (`run_sqd_enlarge_subspace_sbd.py`) / `3` (`run_sqd_sbd.py`) | -| `--sqd_carryover_threshold` | `run_sqd_sbd.py` only. `\|coefficient\|` cutoff for carrying a determinant into the next iteration's sample pool. **Lower it to carry more** | `1e-4` | -| `--enlarge_threshold` | `run_sqd_enlarge_subspace_sbd.py` only — the analogous "carry more" knob for that driver, but structurally different: it gates which *pairs* get expanded into single excitations via `enlarge_batch_from_transitions`, not which determinants survive into resampling. **Lower it to expand more pairs per round** | `1e-4` | +| `--symmetrize_spin` | `1`: merge the alpha and beta string pools every iteration, forcing `ci_strs_a == ci_strs_b`. SBD itself supports distinct alpha/beta determinant sets — this is purely a qiskit-addon-sqd loop-layer setting. `0`: sample and carry over alpha and beta independently, allowing them to differ | `1` | | `--max_dim` | **Critical.** Cap on strings per spin sector, so the subspace cannot exceed `max_dim^2`. The main brake on runaway cost — no fixed value is safe for every system, since the right cap depends on available memory and orbital count. Start from a value known to work at a similar orbital count (e.g. `15000` was used for a 45-orbital system) and adjust down if you see an OOM | unset (no cap) | | `--include_hf` | Force the single Slater determinant with the lowest `num_elec_a`/`num_elec_b` orbital indices occupied into `include_configurations`, every iteration. Cheap correctness check: that determinant's own diagonal energy is an exact lower bound on what a subspace containing it can do — if forcing it in moves the result, the sampled pool was missing it (and probably its low-excitation neighbors too) | off | - -*Decides when to stop — changes nothing about the subspace:* - -| Parameter | What it controls | Default | -|-----------|-----------------|---------| -| `--max_iterations` | Hard cap on loop iterations (not the inner `--sbd_max_it`). In `run_sqd_enlarge_subspace_sbd.py` this is a safety cap only — the loop normally stops earlier, once a round adds no new determinants or both tolerances below are met | `30` (`run_sqd_enlarge_subspace_sbd.py`) / `5` (`run_sqd_sbd.py`) | -| `--energy_tol` | Iteration-to-iteration change in energy | `1e-8` | -| `--occupancies_tol` | Largest change in any single orbital occupancy — an infinity norm, not an average | `1e-5` | +| `--energy_tol` | Stop when the iteration-to-iteration energy change falls below this | `1e-8` | +| `--occupancies_tol` | Stop when the largest change in any single orbital occupancy falls below this — an infinity norm, not an average | `1e-5` | + +*Same flag in all three, but a different default in each:* + +| Parameter | What it controls | sqd | enlarge | carryover | +|-----------|-----------------|-----|---------|-----------| +| `--num_batches` | Independent subsamples per iteration; occupancies are averaged across them. Not the main lever on subspace size in enlarge/carryover — the expansion is | `3` | `1` | `1` | +| `--max_iterations` | Hard cap on loop iterations (**not** the inner `--sbd_max_it`). In enlarge/carryover it is a safety cap only: the loop normally stops earlier, once a round adds no new determinants or both tolerances above are met | `5` | `30` | `30` | + +*How the subspace grows — one knob per driver, and they are **not** +interchangeable:* + +| Driver | Parameter | What it gates | Default | +|--------|-----------|---------------|---------| +| **sqd** | `--sqd_carryover_threshold` | `\|coefficient\|` cutoff for carrying a determinant into the next iteration's **sample pool**. Generates no new determinants — pure selection. **Lower it to carry more** | `1e-4` | +| **enlarge** | `--enlarge_threshold` | `\|amplitude\|^2` cutoff on determinant **pairs**, which are then expanded into all same-spin single excitations via `enlarge_batch_from_transitions`. **Lower it to expand from more pairs** | `1e-4` | +| **carryover** | `--sbd_carryover_type` | Which mechanism SBD uses to pick the next round's determinants: `1` selection only, `2` singles off marginal probability, `3` singles off full-determinant amplitude | `3` | +| **carryover** | `--sbd_carryover_threshold` | The cutoff SBD applies while doing the above. **Meaning depends on the type**: for `3` it is a full-determinant `\|c\|^2` cutoff, for `1`/`2` a marginal (half-determinant) probability | `1e-4` | + +At `--sbd_carryover_type 3`, **carryover**'s threshold gates the same quantity +as **enlarge**'s — both are a full-determinant `|c|^2` cutoff followed by all +same-spin singles. Verified: at `1e-4` on the bundled H2O pool the two drivers +reach an identical energy over an identical final subspace. At other carryover +types the quantity differs, so the values are **not** transferable. + +#### Tuning the expansion threshold + +The threshold is the main accuracy lever in **enlarge**/**carryover**, and the +default `1e-4` is deliberately conservative. Lowering it prunes less, so more +determinants survive into the next round and the subspace grows — which lowers +(improves) the energy, since a variational subspace method can only get better +as the subspace grows. + +Sweeping it on the bundled 275-bitstring H2O pool, with **carryover** at type 3 +(reproducible with the command in section 4 above, changing only +`--sbd_carryover_threshold`): + +| threshold | energy (Ha) | gap to FCI | final subspace | +|---|---|---|---| +| `1e-4` (default) | -76.2421767512 | 1.60 mHa | 1742² = 3.0M | +| `1e-5` | -76.2436018956 | 0.175 mHa | 5007² = 25.1M | +| `1e-6` | -76.2437251036 | **0.052 mHa** | 7881² = 62.1M | + +against this system's FCI reference of `-76.24377680`. Two things to read off +it. First, the payoff is real: two orders of magnitude on the threshold buys +**30x** less error. Second, it is bought with determinants — 20x more of them — +and the returns diminish sharply, with `1e-5`→`1e-6` costing 2.5x the subspace +for a further 0.12 mHa. Sweep downward until the gain stops being worth the +cost for your purpose, rather than reaching for the smallest value. + +**The one trap: `--max_dim` inverts this.** Everything above assumes the +subspace is allowed to grow freely. Once a run is pinned at the cap, lowering +the threshold can make the result *worse*, not better. The cap keeps the +determinants already in the subspace and then fills the remaining budget +**randomly** from this round's new candidates, so a lower threshold means a +much larger candidate pool sampled just as thinly — you pay for generating +them and then discard most, with no ranking to decide which survive. + +The symptom is easy to spot: the final subspace printed at the end is exactly +`max_dim x max_dim`, and the energy moved the wrong way when you tightened the +threshold. We have seen this reverse the ordering outright on a 45-orbital +system at `--max_dim 15000`, where `1e-4` beat `1e-5` comfortably. If that +happens, raise `--max_dim` (memory permitting) or back the threshold off — +tightening it further will not help. + +So, in practice: leave `--max_dim` unset while you sweep the threshold, and +only introduce it once you know the subspace size you are aiming at. Set both +at once and it is hard to tell which one is limiting the answer. **Both stopping criteria must hold in the same iteration.** `fermion.py`'s convergence check combines the energy-change test and the occupancy-change test diff --git a/python/examples/run_sqd_enlarge_subspace_sbd.py b/python/examples/run_sqd_enlarge_subspace_sbd.py index 4f88587..d1e7536 100644 --- a/python/examples/run_sqd_enlarge_subspace_sbd.py +++ b/python/examples/run_sqd_enlarge_subspace_sbd.py @@ -20,15 +20,15 @@ enlarge_batch_from_transitions (single excitations, both spin channels), and feeds the result forward as next round's include_configurations. That is the same general idea as SBD's own --sbd_carryover_type 2/3 (see -run_sbd_selected_ci.py in python/experimental/), implemented instead with +run_sqd_sbd_carryover.py, which uses exactly that), implemented instead with qiskit-addon-sqd's own excitation-generation utility, so it works with any sci_solver, not just SBD, and needs only plain upstream SBD when SBD is used as the solver here. -Two independent stopping conditions, either one is enough (matching -run_sbd_selected_ci.py's own two): the enlarged set adds nothing new beyond -what's already included (closed under single-excitation connectivity), or ---energy_tol and --occupancies_tol both hold between outer rounds. +Two independent stopping conditions, either one is enough: the enlarged set +adds nothing new beyond what's already included (closed under +single-excitation connectivity), or --energy_tol and --occupancies_tol both +hold between outer rounds. --max_iterations is a safety cap, not the primary stopping mechanism -- a run reaching it before either real criterion is a sign something needs tuning, not the expected happy path. @@ -278,10 +278,9 @@ def enlarge_via_singles(ci_strs_a, ci_strs_b, amplitudes, norb, threshold, def cap_to_max_dim(new_ints, existing_ints, max_dim, rng): """Truncate new_ints to max_dim, always keeping everything in existing_ints first. - Same seed-priority logic as run_sbd_selected_ci.py's _cap_to_max_dim, ported - to plain ci_str integer arrays: naive random truncation over the WHOLE - candidate set can discard already-proven-important strings just as easily - as brand-new ones, which is what that driver's own bug fix addressed. + Seed-priority truncation over plain ci_str integer arrays: naive random + truncation over the WHOLE candidate set can discard + already-proven-important strings just as easily as brand-new ones. """ if max_dim is None or len(new_ints) <= max_dim: return new_ints From 18bb575c0dcfd342e6586d6aa1b43e0241dfc9ba Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 18 Sep 2026 18:36:29 -0400 Subject: [PATCH 4/4] Fix carryover driver freezing at --max_dim, correct the speed claim Two problems, both found by benchmarking on a 45-orbital system. H2O could not have caught either. 1. The driver unioned SBD's carryover with the solved subspace before capping, so that the no-growth check had something to compare. But cap_to_max_dim keeps everything already in the subspace first, so once the solved subspace is itself max_dim strings the union makes existing consume the entire budget, every new candidate is discarded, and the subspace freezes. --energy_tol then reads the frozen energy as convergence. On a 45-orbital system at --max_dim 15000 it "converged" after four rounds at -296.0744895995, 3.41 Ha above where the same expansion actually reaches. Invisible with --max_dim unset, which is all H2O exercised. Fixed by forwarding the raw carryover and deriving no-growth from a subset test (carryover subset of solved) instead of a union. Not preserving the solved subspace is correct, not a regression: pruning is what carryover is for, run_sqd_enlarge_subspace_sbd.py's expansion prunes the same way (its identity row preserves the thresholded pairs, not the whole subspace), and qiskit-addon-sqd's own carryover plus fresh sampling re-supply anything that still matters. That driver escapes the deadlock precisely because it does not union. After the fix the two drivers agree round for round on that system -- identical energy, subspace size and expansion size at all 8 rounds, converging to -299.4859985587 over 15000^2. H2O unchanged at -76.2421767512 over 1742^2. 2. The docstring claimed the MPI-distributed expansion was "2-4x faster" than the JAX path. Measured head to head, it is not: 2248 s versus 2250 s. Each round is dominated by configuration recovery and by diagonalizing a 225M-determinant subspace, so where the expansion runs is not measurable. The original figure came from comparing against a driver with no configuration recovery at all, so it measured recovery overhead and was attributed to the wrong thing. Reworded here and in both READMEs: the reason to prefer this driver is operational -- no JAX in the expansion path, so no GPU contention between ranks and no JAX_PLATFORMS=cpu workaround -- and explicitly not speed. --- README.md | 2 +- python/examples/README.md | 11 ++-- python/examples/run_sqd_sbd_carryover.py | 68 ++++++++++++++++-------- 3 files changed, 55 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 8ead7aa..6ec972f 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ They differ only in which engine generates the new determinants: | Driver | Expansion engine | Notes | |---|---|---| | [`run_sqd_enlarge_subspace_sbd.py`](python/examples/run_sqd_enlarge_subspace_sbd.py) | qiskit-addon-sqd's `enlarge_batch_from_transitions` — same-spin single excitations, in JAX | No MPI awareness: every rank recomputes the whole expansion, and on a GPU-enabled JAX install several ranks each try to claim a device. Needs `JAX_PLATFORMS=cpu` beyond one rank | -| [`run_sqd_sbd_carryover.py`](python/examples/run_sqd_sbd_carryover.py) | SBD's own carryover, selected inside the same C++ diagonalization | MPI-distributed, so neither problem applies. Also detects closure a round earlier | +| [`run_sqd_sbd_carryover.py`](python/examples/run_sqd_sbd_carryover.py) | SBD's own carryover, selected inside the same C++ diagonalization | No JAX involved, so neither problem applies and no env workaround is needed. Not faster, though: measured identical wall time on a 45-orbital system | At SBD carryover type 3 the two expansions gate the same quantity — a full-determinant `|c|^2` cutoff followed by all same-spin singles — so at a diff --git a/python/examples/README.md b/python/examples/README.md index ff196b5..09301c7 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -168,14 +168,17 @@ select the next round's determinants, inside the same C++ diagonalization that just ran. They come back on the result and are fed forward as the next round's `include_configurations`. -Two practical consequences, both of which follow from the expansion happening -in SBD's MPI-distributed C++ rather than in JAX on every rank: +What that buys, and what it does not: - **No `JAX_PLATFORMS=cpu` needed.** The JAX expansion has no MPI awareness, so on a GPU-enabled JAX install several ranks each try to claim a device and the run dies with `CUDA_ERROR_OUT_OF_MEMORY`. Nothing here touches JAX. -- **It detects closure a round earlier**, because it compares the expanded set - against the solved subspace directly. +- **It is not faster.** Measured head to head on a 45-orbital system (8 ranks, + `--max_dim 15000`, threshold `1e-4`): 2248 s against 2250 s for the JAX path, + reaching bit-identical energies and subspace sizes at every round. Each round + is dominated by configuration recovery and by diagonalizing the subspace, so + where the expansion runs makes no measurable difference. Pick this driver for + the operational reason above, not for speed. ```bash mpirun -np 8 python -u run_sqd_sbd_carryover.py \ diff --git a/python/examples/run_sqd_sbd_carryover.py b/python/examples/run_sqd_sbd_carryover.py index abecd8f..ffd32f6 100644 --- a/python/examples/run_sqd_sbd_carryover.py +++ b/python/examples/run_sqd_sbd_carryover.py @@ -32,13 +32,24 @@ sbd_solver.SBDCarryoverResult). Why bother, when run_sqd_enlarge_subspace_sbd.py already grows its -subspace: enlarge_batch_from_transitions is JAX and has no MPI awareness, -so every rank redundantly recomputes the whole expansion -- measured 2-4x -slower than SBD-native carryover at matching thresholds, and it exhausts -GPU memory outright once several ranks each try to claim a device. SBD's -carryover is MPI-distributed C++ (SinglesExtendHalfdets splits the work -with MPI_Comm_split/Bcast), so it does not have either problem, and this -driver needs no JAX_PLATFORMS=cpu workaround. +subspace. The honest answer is operational, not performance: +enlarge_batch_from_transitions is JAX and has no MPI awareness, so every +rank redundantly recomputes the whole expansion, and on a GPU-enabled JAX +install several ranks each try to claim a device and the run dies with +CUDA_ERROR_OUT_OF_MEMORY -- which is why that driver needs +JAX_PLATFORMS=cpu beyond one rank. Nothing here touches JAX, so neither +applies. + +It is NOT faster. Measured head to head on a 45-orbital system, 8 ranks, +--max_dim 15000, threshold 1e-4: 2248 s here versus 2250 s for the JAX +path, reaching bit-identical energies and subspace sizes at every one of +the 8 rounds. The expansion is simply not where the time goes -- each +round is dominated by configuration recovery over the sample pool and by +diagonalizing a 225M-determinant subspace, so making the expansion +MPI-distributed buys nothing measurable. (An earlier "2-4x faster" note +was a misattribution: that gap was against a driver with no configuration +recovery at all, so it measured recovery overhead, not the expansion +engine.) Two independent stopping conditions, either one is enough: the carryover set adds nothing new beyond what is already included (the subspace is @@ -434,20 +445,35 @@ def callback(results): "requested. Nothing to expand with." ) - # SBD's carryover is a SELECTION plus (for types >= 2) newly generated - # excitations -- it is not guaranteed to contain everything already in - # the subspace, unlike the identity-row trick the JAX path uses. Union - # with the solved subspace so a round can only ever add, never drop - # determinants that are currently carrying weight. - new_alpha = np.union1d(np.asarray(new_alpha, dtype=np.int64), - np.asarray(ci_strs_a, dtype=np.int64)) - new_beta = np.union1d(np.asarray(new_beta, dtype=np.int64), - np.asarray(ci_strs_b, dtype=np.int64)) - - no_growth = (len(new_alpha) == len(ci_strs_a) - and len(new_beta) == len(ci_strs_b) - and set(new_alpha.tolist()) == set(int(x) for x in ci_strs_a) - and set(new_beta.tolist()) == set(int(x) for x in ci_strs_b)) + new_alpha = np.asarray(new_alpha, dtype=np.int64) + new_beta = np.asarray(new_beta, dtype=np.int64) + + # Stop when the carryover proposes nothing the solved subspace does not + # already hold. Deliberately a SUBSET test against the solved strings, + # not a union folded into what gets forwarded below: + # + # unioning the carryover with the solved subspace looks safer -- "a + # round can only add, never drop a determinant carrying weight" -- but + # it deadlocks the loop the moment --max_dim binds. cap_to_max_dim + # keeps everything already in the subspace first, so once the solved + # subspace is itself max_dim strings the union makes existing == the + # whole budget, every new candidate is discarded, the subspace freezes, + # and --energy_tol reads the frozen energy as convergence. Measured: on + # a 45-orbital system at --max_dim 15000 that bottomed out 3.4 Ha above + # where the same expansion reaches without the union, "converging" after + # four rounds. It cannot show up when --max_dim is unset, which is why + # small-system testing missed it. + # + # Not unioning means a low-amplitude determinant can leave the subspace. + # That is fine and intended -- pruning is what carryover is for, it is + # exactly what run_sqd_enlarge_subspace_sbd.py's expansion does too + # (its identity row preserves the thresholded pairs, not the whole + # subspace), and qiskit-addon-sqd's own carryover plus fresh sampling + # re-supply anything that still matters. + solved_a = set(int(x) for x in ci_strs_a) + solved_b = set(int(x) for x in ci_strs_b) + no_growth = (set(new_alpha.tolist()) <= solved_a + and set(new_beta.tolist()) <= solved_b) # No universal "safe" default exists for --max_dim (a cap that suits a # large system is wildly oversized for H2O/N2, and vice versa), so it