From fd2f3f18c8f1f701108e0b4de63c19ee983d6a36 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 13 Oct 2026 13:36:37 -0400 Subject: [PATCH 1/2] docs: correct outdated and misleading API docstrings - tpb_diag, tpb_diag_from_files, gdb_diag and get_backend list 'gpu-omp' and the aliases; gdb_diag says gpu-omp runs GDB on the host. - get_backend says what None resolves to ('auto' before any init, 'cpu' once a diag call has initialized) and that cpu + gpu-omp must not share a process. - init lists gdb_diag among the per-call device overrides. - tpb_diag_from_files says its one file supplies both alpha and beta. - TPB_SBD, GDB_SBD, LoadFCIDump and LoadAlphaDets get real docstrings. - assemble_rdms cites upstream SBD's docs/user-guide.md, not a private repo. - DeviceConfig's example passes the config to solve_sci_batch instead of the undocumented apply(); auto() drops development history. - References to functions not on the API page become literals, so they no longer render as dead links; from_strings examples use bit_length 62. - The bit_length row covers odd values in all GDB use, not only the heatbath expansion: upstream's own diag gives a wrong plain-diagonalization energy at 31 on Fe4S4 (r-ccs-cms/sbd#104). Docstrings and one table row only; no code changes. Co-Authored-By: Claude Opus 5.5 --- docs/apidocs/sbd.rst | 9 +++-- python/__init__.py | 86 +++++++++++++++++++++++++++++++++-------- python/device_config.py | 36 +++++++---------- python/sbd_solver.py | 21 ++++------ 4 files changed, 95 insertions(+), 57 deletions(-) diff --git a/docs/apidocs/sbd.rst b/docs/apidocs/sbd.rst index 1a5e968..7484075 100644 --- a/docs/apidocs/sbd.rst +++ b/docs/apidocs/sbd.rst @@ -67,10 +67,11 @@ Fields shared by ``TPB_SBD`` and ``GDB_SBD``: reads differs, and SBD silently ignores the other. * - ``bit_length`` - Bits per word when packing determinants. At most 63; 64 is undefined - behavior in SBD's multi-rank redistribution. For GDB's heatbath expansion - (``carryover_type`` 2 or 3) it must also be even whenever a determinant spans - more than one word, or the expansion crashes or returns a wrong energy. The - ``examples/gdb`` drivers require an even value of at most 62. + behavior in SBD's multi-rank redistribution. For GDB it must also be even + whenever a determinant spans more than one word: an odd value breaks SBD's + alpha/beta conversion, which can crash or return a wrong energy, in plain + diagonalization as well as in the heatbath expansion. The ``examples/gdb`` + drivers require an even value of at most 62. Fields of ``TPB_SBD`` only: diff --git a/python/__init__.py b/python/__init__.py index fce06f9..46dc3e6 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -318,8 +318,9 @@ def init(device='cpu', comm_backend='mpi'): first use with ``device='cpu'`` and ``comm_backend='mpi'``. Call it explicitly only when you need GPU or want to control startup timing. - The device can be overridden per-call via the ``device`` parameter on - ``tpb_diag()``, ``tpb_diag_from_files()``, and ``get_backend()``. + The device can be overridden per call via the ``device`` parameter of + ``tpb_diag()``, ``tpb_diag_from_files()``, ``gdb_diag()`` and + ``get_backend()``. Args: device: Default compute device — 'cpu', 'gpu', 'gpu-omp', or 'auto'. @@ -413,11 +414,20 @@ def get_backend(device=None): Get the backend module for the given device. Auto-initializes SBD if needed. Passing ``device`` overrides the - default — this is how you switch between CPU and GPU within the - same process. + default, which is how you switch backends within one process. The one + combination to avoid is the CPU backend together with the OpenMP-offload + one ('gpu-omp'): they share an OpenMP runtime, and loading both leaves + offload regions running on the host. ``sbd.has_backend_conflict()`` + reports it. Args: - device: 'cpu', 'gpu', 'auto', or None (use default). + device: 'cpu', 'gpu' (Thrust), 'gpu-omp' (OpenMP offload), 'auto', an + alias accepted by :func:`init`, or None. None means the device set + by :func:`init`. If nothing has initialized yet it means 'auto', + which prefers Thrust, then OpenMP offload, then CPU, whereas the + first call of :func:`tpb_diag` or :func:`gdb_diag` initializes with + 'cpu'. Call :func:`init` with the device you want to avoid + depending on call order. Returns: The pybind11 backend module (_core_cpu, _core_gpu_thrust, or @@ -496,13 +506,19 @@ def barrier(): # --------------------------------------------------------------------------- def TPB_SBD(device=None): - """Create TPB_SBD configuration object.""" + """Create a TPB_SBD configuration object, for :func:`tpb_diag`. + + Its fields are listed in the tables below. + """ _ensure_initialized() return get_backend(device).TPB_SBD() def GDB_SBD(device=None): - """Create GDB_SBD configuration object.""" + """Create a GDB_SBD configuration object, for :func:`gdb_diag`. + + Its fields are listed in the tables below. + """ _ensure_initialized() return get_backend(device).GDB_SBD() @@ -514,13 +530,37 @@ def FCIDump(device=None): def LoadFCIDump(filename, device=None): - """Load FCIDUMP file.""" + """Load an FCIDUMP file. + + Args: + filename: Path to the FCIDUMP file. + device: Backend whose ``FCIDump`` type to build; None uses the default. + + Returns: + The ``FCIDump`` object that :func:`tpb_diag` and :func:`gdb_diag` take. + """ _ensure_initialized() return get_backend(device).LoadFCIDump(filename) def LoadAlphaDets(filename, bit_length, total_bit_length, device=None): - """Load alpha determinants from file.""" + """Load determinants from a file of bitstrings, one per line. + + Each line is packed as :func:`from_strings` does, in file order, without + sorting or removing duplicates; pass the result through + :func:`sort_bitarray` before using it as a TPB determinant list. + + Args: + filename: Path to the file. + bit_length: Bits per packed word. Must match the config's + ``bit_length``. + total_bit_length: Length of each bitstring, ``norb`` for the alpha or + beta strings of TPB. + device: Backend to use; None uses the default. + + Returns: + A list of packed determinants, each a list of words. + """ _ensure_initialized() return get_backend(device).LoadAlphaDets(filename, bit_length, total_bit_length) @@ -550,7 +590,7 @@ def from_string(s, bit_length, total_bit_length, device=None): def from_strings(strings, bit_length, total_bit_length, device=None): """Pack many bitstrings into one ``(n, words)`` array. - The bulk form of :func:`from_string`, looping in C++ instead of once per + The bulk form of ``sbd.from_string``, looping in C++ instead of once per determinant across the Python boundary. Prefer it for anything larger than a handful: the per-call form costs tens of thousands of round trips on the vendored inputs alone. @@ -563,13 +603,13 @@ def from_strings(strings, bit_length, total_bit_length, device=None): land in the *last* word rather than the first:: from_strings(["0011"], 20, 4) # -> array([[3]], dtype=uint64) - from_strings(["1" + "0" * 69], 63, 70) # -> array([[ 0, 64]], ...) - from_strings(["0" * 69 + "1"], 63, 70) # -> array([[1, 0]], ...) + from_strings(["1" + "0" * 69], 62, 70) # -> array([[ 0, 128]], ...) + from_strings(["0" * 69 + "1"], 62, 70) # -> array([[1, 0]], ...) Within a word the rightmost character is the least significant bit, which is why ``"0011"`` packs to ``3`` rather than ``12``. Getting this backwards produces a valid-looking but wrong subspace rather than an error, so check - round trips with :func:`makestring` when building determinants by hand. + round trips with ``sbd.makestring`` when building determinants by hand. ``bit_length`` must match the ``bit_length`` of the config passed to the diagonalizers, since it sets how many words each determinant occupies. @@ -604,13 +644,19 @@ def tpb_diag_from_files(fcidumpfile, adetfile, sbd_data, """ Perform TPB diagonalization from files. + The one determinant file supplies both the alpha and the beta list, so the + subspace is the product of that list with itself. For different alpha and + beta lists, as an open-shell case needs, load them and call :func:`tpb_diag`. + Args: fcidumpfile: Path to FCIDUMP file. - adetfile: Path to alpha determinants file. + adetfile: Path to the determinant file, used for both alpha and beta. sbd_data: TPB_SBD configuration object. loadname: Path prefix to load an initial wavefunction from (optional). savename: Path prefix to save the final wavefunction to (optional). - device: Override device ('cpu', 'gpu', or None for default). + device: Backend to run on: 'cpu', 'gpu' (Thrust), 'gpu-omp' (OpenMP + offload), or an alias accepted by :func:`init`. None uses the + default; see :func:`get_backend`. Returns: dict with keys: energy, density, carryover_adet, carryover_bdet, @@ -641,7 +687,9 @@ def tpb_diag(fcidump, adet, bdet, sbd_data, savename: Path prefix to save the final wavefunction to (optional). See the note below: for reading amplitudes back, ``sbd_data.dump_matrix_form_wf`` is usually the one you want. - device: Override device ('cpu', 'gpu', or None for default). + device: Backend to run on: 'cpu', 'gpu' (Thrust), 'gpu-omp' (OpenMP + offload), or an alias accepted by :func:`init`. None uses the + default; see :func:`get_backend`. Returns: dict with keys: energy, density, carryover_adet, carryover_bdet, @@ -733,7 +781,11 @@ def gdb_diag(fcidump, det, sbd_data, writes one file per b_comm position, ``f"{savename}{rank_b:06d}.bin"``, each holding only that shard; rank 0's file is not the whole wavefunction. GDB has no combined matrix-form dump. - device: Override device ('cpu', 'gpu', or None for default). + device: Backend to run on: 'cpu', 'gpu' (Thrust), 'gpu-omp' (OpenMP + offload), or an alias accepted by :func:`init`. None uses the + default; see :func:`get_backend`. 'gpu-omp' is accepted, but SBD + has no OpenMP-offload GDB code, so the run executes on the host; + use 'gpu' for GDB on GPUs. determinant_distribution: How to place determinants across b_comm — one of ``input`` (keep the shards as given), ``equal-bra-a`` (default; equal distinct-alpha count per rank, which is what balances the diff --git a/python/device_config.py b/python/device_config.py index f68aa0b..477c53a 100644 --- a/python/device_config.py +++ b/python/device_config.py @@ -27,18 +27,14 @@ class DeviceConfig: Usage:: - # Auto-detect (uses GPU if available) - config = DeviceConfig.auto() + from sbd.device_config import DeviceConfig + from sbd.sbd_solver import solve_sci_batch - # Force CPU - config = DeviceConfig.cpu() + device = DeviceConfig.auto() # best backend that is built + device = DeviceConfig.cpu() # or force one + device = DeviceConfig.gpu(max_memory_gb=16) - # Force GPU with specific settings - config = DeviceConfig.gpu(max_memory_gb=16) - - # Apply to SBD configuration - sbd_config = sbd.TPB_SBD() - config.apply(sbd_config) + results = solve_sci_batch(..., device_config=device) """ def __init__(self, device: str = 'cpu', @@ -72,17 +68,11 @@ def auto(cls, max_memory_gb: int = -1) -> 'DeviceConfig': """ Auto-detect the best available backend and use it. - Resolves against the backends that were actually COMPILED, not just the - hardware that is present. Detecting a GPU and returning 'gpu' - unconditionally was wrong in two ways: on an AMD host it selected the - Thrust/CUDA backend, which cannot exist there (upstream wires Thrust to - nvc++ -cuda), so a machine that reported "GPU detected (HIP)" then failed - to load a CUDA module; and on a CPU-only build with a GPU present it - picked a backend that was never built. - - Preference order matches sbd._resolve_device(): Thrust ('gpu') first - where it exists, since it is the long-validated NVIDIA default and keeps - more phases on the device, then OpenMP offload ('gpu-omp'), then CPU. + Chooses among the backends that were actually built, not merely the + hardware present, so an AMD host or a CPU-only build is never handed a + backend it lacks. The preference order is Thrust ('gpu') where it exists, + then OpenMP offload ('gpu-omp'), then CPU, the same order + ``sbd.get_backend('auto')`` uses. Args: max_memory_gb: Maximum GPU memory in GB (-1 = auto) @@ -158,8 +148,8 @@ def gpu_omp(cls, max_memory_gb: int = -1) -> 'DeviceConfig': one combination to avoid in a single process is this backend together with the CPU one: they share an OpenMP runtime (``libnvomp`` on NVIDIA, ``libomp`` on AMD), and loading ``_core_cpu`` first leaves it initialised - host-only, after which offload regions run on the host. See - :func:`sbd.has_backend_conflict`. + host-only, after which offload regions run on the host. + ``sbd.has_backend_conflict()`` reports that combination. """ return cls(device='gpu-omp', max_memory_gb=max_memory_gb) diff --git a/python/sbd_solver.py b/python/sbd_solver.py index e8b4508..b346158 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -470,25 +470,20 @@ def assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarr which only computes the diagonal density, not the correlation functions), matching SCIResult's own ``rdm1``/``rdm2`` default. - The reshape/transpose below is not a guess: it was verified against - PySCF's ``make_rdm1``/``make_rdm2`` on a fixed subspace, on all three - SBD backends (cpu, gpu-thrust, gpu-omp-offload) -- both element-wise on - the full tensors and via the energy identity + The layout conversion was verified against PySCF's ``make_rdm1`` and + ``make_rdm2`` on all three backends (cpu, gpu, gpu-omp), both element-wise + and through the energy identity ``E = einsum("pr,pr->",rdm1,hcore) + 0.5*einsum("prqs,prqs->",rdm2,eri)`` - that ``SCIResult.rdm1``/``rdm2`` are contracted with everywhere else in - qiskit-addon-sqd (e.g. ``fermion.py``'s own ``solve_fermion``). + with which ``qiskit-addon-sqd`` contracts ``SCIResult.rdm1`` and ``rdm2``. - 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 layout, as documented in upstream SBD's ``docs/user-guide.md``:: 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 - (arr_F[i, j] / arr_F[i, j, k, l]); rdm1 needs no further transpose - (it is symmetric here regardless), and rdm2's spin-summed block sum - needs axes (0, 2, 1, 3) to land in the "prqs" slot order SCIResult's - contract expects. + A Fortran-order reshape implements those flat-index formulas directly. + rdm1 needs no further transpose, and the spin-summed rdm2 is transposed + with axes ``(0, 2, 1, 3)`` into the "prqs" order ``SCIResult`` expects. """ one_p_rdm = results.get("one_p_rdm") two_p_rdm = results.get("two_p_rdm") From 27ed54857d5942670e573351a6d29c100145cd63 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Wed, 14 Oct 2026 15:42:31 -0400 Subject: [PATCH 2/2] docs: LoadAlphaDets sorts and removes duplicates Upstream's LoadAlphaDets ends with an unconditional sort_bitarray (makedeterminants.h:190), so the result is sorted and distinct, not in file order, and needs no further sort. Also note that only .txt and .bin are read: any other extension returns an empty list without an error. Co-Authored-By: Claude Opus 5.5 --- python/__init__.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/python/__init__.py b/python/__init__.py index 46dc3e6..d0a9b78 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -546,9 +546,14 @@ def LoadFCIDump(filename, device=None): def LoadAlphaDets(filename, bit_length, total_bit_length, device=None): """Load determinants from a file of bitstrings, one per line. - Each line is packed as :func:`from_strings` does, in file order, without - sorting or removing duplicates; pass the result through - :func:`sort_bitarray` before using it as a TPB determinant list. + Each line is packed as :func:`from_strings` does. Upstream sorts the + determinants and removes duplicates, so the result is ready to use as a + TPB determinant list and may hold fewer determinants than the file has + lines. + + The format is chosen by extension: ``.txt`` for one bitstring per line, + ``.bin`` for upstream's binary format. Any other extension returns an + empty list without an error. Args: filename: Path to the file. @@ -559,7 +564,8 @@ def LoadAlphaDets(filename, bit_length, total_bit_length, device=None): device: Backend to use; None uses the default. Returns: - A list of packed determinants, each a list of words. + A list of packed determinants, each a list of words, sorted and + distinct. """ _ensure_initialized() return get_backend(device).LoadAlphaDets(filename, bit_length, total_bit_length)