diff --git a/README.md b/README.md index e091a44..361320d 100644 --- a/README.md +++ b/README.md @@ -4,19 +4,37 @@ Python bindings for the Selected Basis Diagonalization (SBD) library, with CPU a ## Overview -SBD (Selected Basis Diagonalization) is a high-performance library for quantum chemistry calculations. The Python bindings provide access to SBD's **Tensor-Product Basis (TPB)** diagonalization method on CPU and GPU. +SBD (Selected Basis Diagonalization) is a high-performance library for quantum +chemistry calculations. These bindings expose two of its diagonalization methods on CPU +and GPU. They differ in the shape of the subspace they span, which is what decides +which one you want: + +- **Tensor-Product Basis (TPB)** — the subspace is the Cartesian product of an alpha + and a beta determinant list, so its dimension is `|adet| × |bdet|`. The mature path, + and the one the qiskit-addon-sqd integration uses. +- **General-Determinant Basis (GDB)** — the subspace is the explicit determinant list + you pass, so it can be an arbitrary *sparse* set rather than a product. Experimental: + newer, a smaller tested surface, and still evolving as upstream SBD does. + +If your subspace is a product, prefer TPB — it represents that case with two +half-determinant lists instead of every product element, and needs no extra +constraints. Reach for GDB when the subspace is not a product. **Key Features:** -- **TPB diagonalization** for quantum chemistry Hamiltonians +- **TPB and GDB diagonalization** for quantum chemistry Hamiltonians - Three backends, selected per call at runtime via `device=`: `'cpu'` (host OpenMP), `'gpu'` (NVHPC Thrust/CUDA, NVIDIA only) and `'gpu-omp'` (OpenMP target offload, **NVIDIA or AMD**). All the backends your toolchain supports can be built into one install; each is imported only when - first used + first used. **TPB runs on all three; GDB has Thrust kernels only**, so GDB on a + GPU is NVIDIA-only and under `'gpu-omp'` it falls back to the host — see + [`examples/gdb/README.md`](examples/gdb/README.md) - MPI parallelization - Integration with [qiskit-addon-sqd](https://github.com/Qiskit/qiskit-addon-sqd) for SQD workflows -In addition to TPB, this package also contains experimental support for SBD's **General-Determinant Basis (GDB)** method. However, SBD's **Creation/Annihilation operator (CAOP)** method is currently not supported by this wrapper; users who need it should reference and use the C++ CLI apps in the upstream submodule (`vendor/sbd-upstream/apps/`). +SBD's **Creation/Annihilation operator (CAOP)** method is not supported by this +wrapper; users who need it should reference and use the C++ CLI apps in the upstream +submodule (`vendor/sbd-upstream/apps/`). > [!NOTE] > This package is newly open-sourced. The Python API follows semantic versioning, but the build configuration and GPU backends have been exercised on a limited set of platforms — please report issues. @@ -58,11 +76,23 @@ matter for it. - [`examples/tpb/`](examples/tpb/README.md) — tensor-product basis: standalone TPB diagonalization, the SQD loops, and the subspace-enlargement driver. +- [`examples/gdb/`](examples/gdb/README.md) — general determinant basis: standalone + GDB diagonalization over an explicit determinant list, and an iterative + heatbath-expansion driver. ## Integration with qiskit-addon-sqd SBD can serve as the eigensolver backend for qiskit-addon-sqd's SQD workflow. +> [!NOTE] +> This integration is **TPB-only**, and not merely for want of plumbing. +> `solve_sci`/`solve_sci_batch` call `tpb_diag`, and the addon's interface describes a +> product subspace by construction: `ci_strings` is a `(strings_a, strings_b)` pair and +> `SCIState.amplitudes` is an `|a| × |b|` matrix. A sparse determinant list cannot be +> expressed that way without padding back up to the full product, which discards the +> reason to use GDB. For GDB, call `gdb_diag` directly or use the drivers in +> [`examples/gdb/`](examples/gdb/README.md). + **Note:** Requires [qiskit-addon-sqd](https://github.com/Qiskit/qiskit-addon-sqd) with distributed (SPMD) support — `diagonalize_fermionic_hamiltonian` calling `sci_solver` on every MPI rank. This is available in `qiskit-addon-sqd` version `0.13.1` or higher. ### Plain SQD @@ -192,37 +222,27 @@ Total MPI ranks = `task_comm_size × adet_comm_size × bdet_comm_size`. config = sbd.GDB_SBD() ``` -Shares `method`, `max_it`, `max_nb`, `eps`, `max_time`, `init`, `do_shuffle`, -`do_rdm`, `carryover_type`, `ratio`, `threshold` and `bit_length` with `TPB_SBD`, -and replaces the determinant communicators with a single basis communicator: +Shares `max_it`, `max_nb`, `eps`, `max_time`, `init`, `do_rdm`, `carryover_type`, +`ratio`, `threshold` and `bit_length` with `TPB_SBD`, and replaces the determinant +communicators with a single basis communicator: | Attribute | Default | Description | |-----------|---------|-------------| -| `b_comm_size` | 1 | Basis communicator size — must be 1 for `gdb_diag`, see below | -| `t_comm_size` | 1 | Task communicator size — must be 1 while `b_comm_size` is, see below | +| `method` | 0 | 0=Davidson, 1=Davidson+Ham. GDB has no Lanczos, so TPB's 2 and 3 are rejected | +| `b_comm_size` | 1 | Basis communicator size — how many shards the determinant list is split into | +| `t_comm_size` | 1 | Task communicator size — must not exceed `b_comm_size` | | `seed` | 1729 | Seed for the initial vector | | `heatbath_cutoff` | 1e-4 | Heatbath expansion cutoff | | `heatbath_truncation` | 0.0 | Weight truncation applied before heatbath expansion | | `heatbath_batch_size` | 200000000 | Heatbath expansion batch size | -**Why both must be 1**, since the two constraints have different owners: - -`b_comm_size == 1` is a limitation of *this wrapper*, not of SBD. Upstream's in-memory -`gdb::diag` expects each rank to pass **its own shard** of the determinant list — that -is what upstream's file-based entry point hands it, after distributing determinant -files across `b_comm`. This wrapper passes the whole list from every rank, which is -only consistent with a single basis block, so it rejects anything else rather than -have each rank diagonalize the full subspace while believing it held a shard. Upstream -itself runs with a split basis: its own `run.sh` for the GDB app passes -`--b_comm_size 2`. +Total MPI ranks = `t_comm_size × b_comm_size × helper`, where the helper dimension is +the derived quotient and is not settable. With `b_comm_size > 1` each rank passes its +own shard of the determinant list rather than the whole list. -`t_comm_size == 1` then follows from *upstream's* algorithm rather than from us. GDB's -matrix-vector product rotates the ket around `b_comm` as a ring, so there are exactly -`b_comm_size` ring stations and one "task" is one station — meaning `t_comm_size` -cannot exceed `b_comm_size`. With the basis in a single block there is a single task. - -Ranks are not wasted in the meantime: the derived helper dimension, -`ranks / (t_comm_size × b_comm_size)`, absorbs them and does not change the energy. +See [`examples/gdb/README.md`](examples/gdb/README.md) for the decomposition, the shard +contract, the six determinant-placement schemes, and why the helper dimension must be 1 +on the Thrust backend. ### Diagonalization @@ -250,11 +270,19 @@ results = sbd.gdb_diag(fcidump, det, sbd_data, spin-beta orbital `i`. The determinants must be distinct; they are sorted into SBD's canonical order internally, which `sort_bitarray` reproduces. -`gdb_diag` does not return the wavefunction amplitudes, because SBD's `gdb::diag` -has no in-memory output for them. Passing `savename` makes SBD write them to -`f"{savename}000000.bin"` instead: two `size_t` headers -(`n_dets`, `words_per_det`), then `n_dets × words_per_det` `size_t` determinant -words in canonical order, then `n_dets` `float64` amplitudes. +`det` may be a single `(ndets, words)` array or **this rank's shard of one**: +`sbd_data.b_comm_size` decides which. At `1` every rank passes the whole basis; above +`1` every rank passes its own shard, the union over b_comm positions being the basis. +Sharded input must be globally sorted and disjoint; that and the rank-layout constraints +are checked, and raise rather than silently diagonalizing the wrong subspace. +[`examples/gdb/README.md`](examples/gdb/README.md) has the shard contract, the +placement schemes, and which returned values are replicated versus sharded. + +`gdb_diag` returns no wavefunction amplitudes — `gdb::diag` has no in-memory output for +them — and for most uses none are needed: the energy, density and RDMs come back +directly, and an iterative heatbath run gets its next subspace from `carryover_det`. If +you do want the amplitudes, `savename` writes them to disk; the file layout is in +[`examples/gdb/README.md`](examples/gdb/README.md). The optional `device` parameter overrides the default set by `init()`. @@ -294,13 +322,13 @@ prefer it. Tracked as **Ranks die with `Bus error` or `SIGSEGV` inside the MPI's own copy path** (`MPIR_Localcopy`, `ucp_worker_progress`, ...) **on a GPU backend:** the MPI is not GPU-aware and was handed a device pointer. Rebuild UCX `--with-cuda` / `--with-rocm`, and confirm with `ucx_info -d | grep -i 'Transport: cuda'` (or `rocm`). Two things mislead here. A partly GPU-aware stack fails in only one place: an MPICH with GPU support *disabled* over a CUDA-aware UCX ran OMP-offload fine and crashed only in Thrust, because the inter-rank path went through UCX while the local-copy path did not. And on AMD a non-ROCm-aware MPI does not crash at all — ROCm maps device memory into the process address space, so the host copy succeeds and merely stages everything through the host aperture (measured on MI250X, XNACK off, 8 ranks) — so a working AMD run is not evidence that the MPI is ROCm-aware. -**`GDB Thrust mult does not support h_comm_size > 1` from `gdb_diag` on more than one -rank:** GDB's Thrust kernels never implemented the helper dimension, and the helper -dimension is `ranks / (t_comm_size × b_comm_size)`. Since `gdb_diag` requires -`b_comm_size == 1` (see [Configuration](#configuration)), which forces `t_comm_size` to -1, every rank you add lands in the helper dimension — so GPU GDB is limited to a single -rank in this release. Run GDB on one GPU, or on the CPU backend, where the helper -dimension is unconstrained. TPB is unaffected and shards over `adet_comm_size` / -`bdet_comm_size` as usual. +**GDB on more than one GPU refuses to start, naming the helper dimension:** GDB's Thrust +kernels never implemented that dimension, and it is the quotient +`ranks / (t_comm_size × b_comm_size)` — so any rank you do not assign to `t` or `b` lands +there. Give every rank to the basis: `-np 4` with `b_comm_size = 4` leaves a helper +dimension of 1. Leaving `b_comm_size` at 1 puts *all* ranks in the helper dimension, +which is why multi-GPU GDB requires a split basis. `gdb_diag` checks this before doing +any work rather than letting the kernel throw mid-launch. The CPU backend has no such +restriction, and TPB is unaffected. **Repository:** https://github.com/Qiskit/sbd-eigensolver-python diff --git a/examples/gdb/README.md b/examples/gdb/README.md new file mode 100644 index 0000000..9bd3458 --- /dev/null +++ b/examples/gdb/README.md @@ -0,0 +1,511 @@ +# GDB examples — general determinant basis + +GDB spans the subspace with the determinants it is given, rather than with the +Cartesian product of an alpha and a beta list that TPB uses. TPB's dimension is +`|adet| x |bdet|`; GDB's is exactly the number of determinants passed. That is what +makes it the right solver for an arbitrary sparse subspace — a set of sampled +bitstrings used *as sampled*, with no product completion. + +For TPB and the SQD loops, see [`../tpb/README.md`](../tpb/README.md). + +## The default data + +Both drivers default `--fcidump` and `--detfiles` to upstream's own GDB app data under +`vendor/sbd-upstream/apps/chemistry_gdb_selected_basis_diagonalization/`: the Fe4S4 +FCIDUMP at 36 orbitals, and its four `det0.txt`-`det3.txt` holding 14,884 determinants +each, 59,536 in total. **Any command below that passes neither flag runs exactly that +case** — the determinants come from those four files, not from nowhere. Upstream +publishes no reference energy for it, so treat it as a self-consistency benchmark +rather than a validation target. + +## run_gdb_diag.py — standalone GDB diagonalization + +Stays on the in-memory entry point (`sbd.gdb_diag`) throughout: determinant text is +read in Python and handed to the binding as a list. SBD's own file-based entry +point is deliberately not used. + +```bash +# The default case: Fe4S4, upstream's four det files, 59,536 determinants. +python run_gdb_diag.py + +# The same run with the defaults written out -- this is the --detfiles syntax to +# copy for your own data. Files are comma-separated and concatenated in Python +# into one in-memory list, so their combined order must be sorted and disjoint. +GDB=../../vendor/sbd-upstream/apps/chemistry_gdb_selected_basis_diagonalization +python run_gdb_diag.py --fcidump $GDB/fcidump_Fe4S4.txt \ + --detfiles $GDB/det0.txt,$GDB/det1.txt,$GDB/det2.txt,$GDB/det3.txt + +# Shard the basis: Fe4S4's four files, one per rank. For this data that is +# already the balanced globally-sorted split, so no redistribution is needed. +OMP_NUM_THREADS=8 mpirun -np 4 python run_gdb_diag.py --b_comm_size 4 + +# Spend ranks on both named dimensions (t <= b, and t*b must divide the ranks) +mpirun -np 8 python run_gdb_diag.py --b_comm_size 4 --t_comm_size 2 + +# Choose how determinants are placed across b_comm +mpirun -np 4 python run_gdb_diag.py --b_comm_size 4 \ + --determinant_distribution grid-cyclic + +# A subspace TPB can also express: interleave an alpha list with itself into the +# full |A|^2 product basis. This is the cross-check against tpb_diag -- same +# subspace, two independent solvers, energies must agree. +python run_gdb_diag.py \ + --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ + --from-alpha ../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt \ + --alpha-limit 30 + +# Grow the subspace with SBD's own heatbath expansion (one round; the expanded +# list comes back as carryover_det) +python run_gdb_diag.py --carryover_type 2 --heatbath_cutoff 1e-4 + +# On GPUs, use --device gpu (Thrust). It requires helper == 1, so every rank has +# to go to t*b -- which means --b_comm_size is not optional for multi-GPU GDB. +mpirun -np 4 python run_gdb_diag.py --device gpu --b_comm_size 4 + +# One rank, one GPU is the exception: b=1 already leaves helper == 1. +python run_gdb_diag.py --device gpu +``` + +The determinant bit order is the one the upstream app documents: reading from the +right, alpha orbital 1, beta orbital 1, alpha orbital 2, and so on — i.e. bit +`2*i` is alpha orbital `i` and bit `2*i + 1` is beta orbital `i`. + +## run_gdb_heatbath.py — grow the subspace, round after round + +Selected CI (HCI) driven from Python: diagonalize, let SBD expand the subspace from +the resulting wavefunction, diagonalize the larger subspace, repeat. The loop needs +no extra machinery — `carryover_type` 2 and 3 return the parents *together with* the +new candidates, so one round's result **is** the next round's subspace. + +```bash +# Fe4S4 from upstream's shipped subspace (the default data above), one cutoff +python run_gdb_heatbath.py --cutoffs 1e-3 + +# A ladder, stopping before the subspace passes 2M determinants. The last cutoff +# dominates the cost, so --max_dim is the brake. +python run_gdb_heatbath.py --cutoffs 1e-3,1e-4,1e-5 --max_dim 2000000 + +# The no-input null: the Hartree-Fock determinant alone. Run it serially -- one +# determinant cannot be sharded. +OMP_NUM_THREADS=48 python run_gdb_heatbath.py --subspace-from hf --cutoffs 1e-3,1e-4 + +# Sharded. At --b_comm_size == ranks with t=1 the expansion comes back already +# sharded for the next round, so nothing has to be gathered. +mpirun -np 4 python run_gdb_heatbath.py --b_comm_size 4 --cutoffs 1e-3 +mpirun -np 4 python run_gdb_heatbath.py --b_comm_size 4 --device gpu --cutoffs 1e-4 + +# Record the (dimension, energy) series for a comparison table +python run_gdb_heatbath.py --cutoffs 1e-3,1e-4 --log ladder.json +``` + +### Parameters + +Input flags match `run_gdb_diag.py`: `--fcidump`, `--detfiles` and `--alpha-limit` are +spelled and defaulted identically, and the alpha list is accepted as either +`--alpha-file` or `--from-alpha` by both drivers, so a command that feeds one its data +feeds the other. The one flag that is *not* shared is `--seed`: in +`run_gdb_diag.py` it is the integer RNG seed for a random initial vector, and this +driver has none — `--subspace-from` selects the starting subspace and is unrelated. + +Seed — where the starting subspace comes from: + +| Parameter | What it controls | Default | +|---|---|---| +| `--subspace-from` | `files` reads `--detfiles`; `hf` starts from the single Hartree-Fock determinant; `from-alpha` builds the `\|A\|^2` product of an alpha list; `strings` reads full determinants, e.g. sampled configurations | `files` | +| `--fcidump` | FCIDUMP defining the Hamiltonian | Fe4S4, see [The default data](#the-default-data) | +| `--detfiles` | `--subspace-from files`: comma-separated determinant files, concatenated in Python. Their combined order must be sorted and disjoint | upstream's four Fe4S4 files | +| `--alpha-file` (or `--from-alpha`) / `--alpha-limit` | `--subspace-from from-alpha`: the alpha list, and a cap on how many of its strings to keep. The product costs `N^2` determinants, so this is the size dial | none / `0` (all) | +| `--strings-file` | `--subspace-from strings`: a file of `2*norb`-bit determinants | none | + +Ladder — how far the expansion is pushed. A rung runs rounds at one cutoff until the +energy stops moving, then the next cutoff begins: + +| Parameter | What it controls | Default | +|---|---|---| +| `--cutoffs` | The ladder itself: comma-separated `heatbath_cutoff` values, smallest step last. Each rung admits candidates whose estimated contribution exceeds it | `1e-3` | +| `--max_rounds` | Cap on rounds **per rung**, so a rung that never converges cannot run forever | `8` | +| `--energy_tol` | Advance to the next rung once `\|dE\|` between rounds falls below this (Hartree) | `1e-5` | +| `--max_dim` | Stop before diagonalizing a subspace larger than this. The brake that keeps rows of a comparison cost-matched | `0` (no cap) | +| `--carryover_type` | Heatbath variant, 2 or 3. Types 0 and 1 do not expand, so they cannot drive the loop | `2` | +| `--heatbath_truncation` | Weight threshold applied to **parents**, before expanding — not the cutoff. See the warning below; leave it at 0 | `0.0` | +| `--heatbath_batch_size` | Expansion batch size per rank | `1000000` | + +Solver, MPI and output — the same meanings as in `run_gdb_diag.py`: + +| Parameter | What it controls | Default | +|---|---|---| +| `--device` | `cpu` or `gpu` (Thrust) are the two real choices; `gpu-omp` and `auto` run GDB on the host, see [Choosing a backend](#choosing-a-backend) | `cpu` | +| `--method` | 0=Davidson, 1=Davidson storing the Hamiltonian. GDB has no Lanczos | `0` | +| `--tolerance` / `--iteration` / `--block` | Davidson residual tolerance, iteration cap, and basis-vector count. Also accepted as `--eps` / `--max_it` / `--max_nb`, matching SBD's own names | `1e-6` / `30` / `10` | +| `--bit_length` | Bits per packed word | `64` | +| `--b_comm_size` | Basis shards — the only dimension that divides memory | `1` | +| `--t_comm_size` | Tasks per ring station; must not exceed `--b_comm_size` | `1` | +| `--determinant_distribution` | Placement across `b_comm`; see [Placement across `b_comm`](#placement-across-b_comm) | `equal-bra-a` | +| `--log FILE` | Write the per-round `(rung, cutoff, round, dimension, energy, delta_energy, seconds)` series as JSON — the machine-readable form of the table the driver prints | none | + +### Why a cutoff *ladder* + +`heatbath_cutoff` admits a candidate when its estimated contribution exceeds the +threshold, so **for a fixed cutoff the subspace reaches a self-consistent size and +stops growing** — more rounds then buy nothing. Going deeper means lowering the +cutoff. So `--cutoffs` takes a ladder: rounds run at one cutoff until the energy +stops moving (`--energy_tol`, the standard HCI criterion) or the subspace stops +growing, then the next rung begins. `--max_dim` caps the whole run, which is what +keeps rows of a comparison cost-matched rather than cutoff-matched. + +On Fe4S4 (36 orbitals, 54 electrons) starting from upstream's four shipped files, +a `1e-3,1e-4` ladder walks the subspace out like this: + +| rung | dimension | energy | +|---|---|---| +| seed (upstream's four files) | 59,536 | −326.6982518821 | +| `1e-3` round 1 | 63,569 | −326.7298282871 | +| `1e-3` settled | 63,887 | −326.7302438694 | +| `1e-4` round 1 | 568,538 | −326.7773323983 | +| `1e-4` round 3 | 755,107 | −326.7839718082 | + +`1e-3` gains ~32 mHa for 7% more determinants and then saturates — that is the fixed +point for that cutoff, not convergence of the method. One step to `1e-4` multiplies the +subspace by roughly nine. Use `--max_dim` to stop before a rung outgrows your memory. + +The energy is variational, so it must fall monotonically; the driver flags a rise, +which would mean the subspace shrank or a round failed to converge. For h2o and n2 the +FCI limits in that basis (−76.24377680 and −109.04874199) are hard ceilings a correct +run can never cross, which makes a sparse run self-checking even without a reference. + +### What the rungs cost + +Each rung's expansion is much larger than the one before, so the last cutoff you name +dominates the run — a rung is reachable or not rather than merely slow, since past some +point the expansion outgrows what can be diagonalized. `--max_dim` is what makes that a +clean stop with a `stop_reason` instead of running out of memory mid-round. + +### Seeds + +`--subspace-from` selects where the starting subspace comes from, so one driver produces +every row of a seed comparison: + +| `--subspace-from` | starting subspace | +|---|---| +| `files` *(default)* | determinant files, defaulting to upstream's four Fe4S4 files | +| `hf` | the Hartree-Fock determinant alone, built from the FCIDUMP header — the no-input null | +| `from-alpha` | an alpha list interleaved with itself, i.e. a TPB-shaped product space | +| `strings` | an arbitrary bitstring file, e.g. sampled configurations | + +Which seed is best is system-dependent, and worth measuring rather than assuming. +The `hf` seed is the useful baseline: it starts from a single determinant and needs no +input subspace at all, so it shows what the classical expansion achieves on its own. +On Fe4S4 it climbs from −326.5243550095 (the HF energy) into the same range the +file-seeded run reaches, at a comparable dimension — so on that system the expansion +is doing most of the work and a pre-existing subspace adds little. Whether that holds +for your system is exactly the kind of thing to check with `--log` and a dimension +sweep. + +Compare seeds **at matched dimension**, not at matched cutoff. Different seeds reach +different sizes from the same cutoff, so a cutoff-matched comparison mostly reports +subspace size rather than seed quality. `--max_dim` and the `--log` series are there +for that. + +### Feeding QPU samples in: `--subspace-from strings` + +`--strings-file` takes a plain text file, one `2*norb`-bit determinant per line — not a +counts JSON. Converting a qiskit-addon-sqd counts file means two steps, and the second +one is not optional: + +1. **Postselect** on the target Hamming weight per spin sector. `recover_configurations` + and friends do this in the SQD loop; doing it here keeps determinants with the wrong + electron count out of the subspace. +2. **Interleave the bits.** qiskit-addon-sqd emits `[beta | alpha]` — the two halves + concatenated — while GDB wants them interleaved, bit `2*i` alpha orbital `i` and bit + `2*i + 1` beta orbital `i`. The drivers expose the conversion as `interleave()`. + +```python +import json, pathlib, sys +sys.path.insert(0, ".") # run from examples/gdb/ +from run_gdb_heatbath import interleave + +counts = json.loads(pathlib.Path("../tpb/count_dict_h2o.json").read_text()) +norb = len(next(iter(counts))) // 2 +n_alpha = n_beta = 5 # the sector you are solving + +dets = { + interleave(key[norb:], key[:norb]) # key is [beta | alpha] + for key in counts + if key[norb:].count("1") == n_alpha and key[:norb].count("1") == n_beta +} +pathlib.Path("dets.txt").write_text("\n".join(sorted(dets)) + "\n") +``` + +A `set` drops duplicate samples, which GDB rejects rather than silently dedupes, and +`sorted()` gives the ordering the shard contract wants. Then: + +```bash +python run_gdb_heatbath.py \ + --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ + --subspace-from strings --strings-file dets.txt \ + --cutoffs 1e-3,1e-4 +``` + +On the bundled 275-bitstring h2o file that is a 275-determinant sparse subspace at +−76.0723973374, which the ladder then grows — as opposed to the 75,625-determinant +product space TPB would build from the same samples. + +**Both drivers check this for you.** Every determinant's alpha and beta electron +counts are compared against the FCIDUMP's `NELEC`/`MS2` before anything is solved, and a +mismatch is refused with the count of offending determinants and the first one's index. +Skipping the interleave is worth guarding because it does not fail cleanly on its own: +the concatenated strings diagonalize to a plausible-looking number (−66.8043 for the +file above, against −76.0724 done right) and only abort later, inside the heatbath +expansion, with `std::out_of_range`. The occupation density cannot catch it either — +permuting bits preserves how many are set, so it still sums to 10. + +The check costs about 0.04 s per million determinants, which is not measurable against +the read and pack that precede it; `--skip-weight-check` turns it off if you are +deliberately mixing spin sectors. + +### The loop needs no amplitudes + +A natural question, since `gdb_diag` does not return the wavefunction: the ladder never +needs it. The amplitudes are what drive the selection — weight truncation keeps +determinants by `|c|`, and heatbath scoring is essentially `|c_i · H_ij|` — but SBD +consumes them internally (`WeightTruncation` then `HeatbathExpansion`, which takes the +coefficients as an input) and hands back only the expanded determinant list. That list +is the next subspace, so the loop closes with nothing but determinants crossing the +Python boundary. + +Where amplitudes *would* be needed is Python-side selection — deciding yourself which +determinants to expand, as `../tpb/run_sqd_enlarge_subspace_sbd.py` does for TPB. For +GDB that means reading them back from the per-shard `savename` files, since +`gdb::diag` has no in-memory amplitude output. + +### `--heatbath_truncation` is not the cutoff + +It discards **parents** by weight *before* expansion starts, and its default of 0 +(keep every parent) is almost always what you want. Setting it to `1e-4` on the +59,536-determinant Fe4S4 wavefunction cut the subspace to 506 rather than growing it. + +## Expected results + +| case | energy | note | +|---|---|---| +| h2o, first 24 alpha interleaved into a 576-determinant product basis | **-76.0588897208** | matches `tpb_diag` on the same subspace to 1.4e-14 | +| h2o, full 275² = 75,625-determinant interleave | **-76.2359466308** | against the **-76.23594663** published for that alpha list | +| Fe4S4, 59,536 determinants, 36 orbitals | **-326.6982518821** | measured here; upstream publishes no reference for this case | + +The h2o product-basis energy is unchanged across `b_comm_size` 1, 2 and 4, across +`t_comm_size` 1 and 2, and across all six placement schemes. Fe4S4 is likewise +identical at `b=1`, `b=4` (one file per rank) and `b=2, t=2`. + +A note on what sharding buys: `b_comm_size` divides **memory**, not necessarily time. +At modest sizes the ring communication can offset the extra parallelism, so the reason +to shard is to hold a subspace that would not fit on one rank — and, on GPUs, because +it is the only way to use more than one card at all (see [On GPUs](#on-gpus)). Measure +on your own hardware rather than extrapolating from anyone else's. + +## Preparing pre-split determinant files (optional) + +Neither driver needs pre-split files: both read a list and slice it deterministically, +so one big file is always correct. Splitting is an **I/O and memory optimization** — +with one file every rank reads the whole thing and keeps only its slice, whereas with +one file per rank the read is parallel and peak memory is a single shard. Both drivers +take that path automatically when the number of `--detfiles` equals `--b_comm_size`. + +Upstream already ships the tool for producing them: **`apps/gen_dets`** (built as +`gdet`). It takes an alpha determinant list, forms the full alpha × beta product, +sorts it, and writes it as N shards: + +```bash +cd vendor/sbd-upstream/apps/gen_dets +# edit Configuration for your compiler, then: +make +mpirun -np 1 ./gdet --adetfile AlphaDets.txt \ + --detfiles det0.txt,det1.txt,det2.txt,det3.txt \ + --bit_length 20 --norb 36 +``` + +The number of output files comes from `--detfiles`, **not** from the rank count — rank +0 writes all of them itself, splitting with `get_mpi_range`, so `-np 1` is enough +(upstream's own `run.sh` passes `-np 4`, which works but is not required). Set +`--norb` to the orbital count and give one file per rank you intend to run on, e.g. +eight files for eight GPUs. + +The result satisfies what `gdb_diag` requires of input shards — each shard sorted, +shards disjoint, and the concatenation globally sorted so shard *i* sits strictly +below shard *i+1* — because `get_mpi_range` is the same `q = N/p` +remainder-to-the-low-ranks split the drivers use. + +Two gaps worth knowing. `gdet` only takes an **alpha list** and can only produce the +alpha × beta *product*, so there is no upstream way to shard an arbitrary sparse or +heatbath-expanded determinant list into files; feed those to the drivers as a single +file (or in memory, which is what `run_gdb_heatbath.py` does between rounds). And +upstream's error messages point at a `scripts/sort-basis-shards.py` that is not +shipped — there is no `scripts/` directory in the tree. + +### Note on Fe4S4's shipped basis + +Worth knowing when reading benchmark numbers: `AlphaDets.txt` next to `gen_dets` +holds exactly **244** alpha determinants, 244² = **59,536**, and the four det files +total 59,536 and are byte-identical to the GDB app's. So upstream's Fe4S4 "GDB" +benchmark subspace is the **full Cartesian product** of 244 alpha determinants with +itself — a TPB-shaped space written out as an explicit determinant list, not a sparse +one. TPB fed `AlphaDets.txt` returns the same −326.6982518821 that GDB returns on the +four files. + +## MPI decomposition + +GDB decomposes as `t_comm_size × b_comm_size × helper`, with its own field names +rather than TPB's `adet`/`bdet`/`task`. Only two are settable; `helper` is derived: + +``` +helper = ranks / (t_comm_size × b_comm_size) +``` + +**`t_comm_size × b_comm_size` must divide the rank count exactly.** Upstream takes +that quotient by integer division and never checks the remainder, which silently +produces communicators of unequal size and a rank alone in its own basis ring, so +`gdb_diag` refuses it. + +These examples set `OMP_NUM_THREADS` in the shell rather than through `mpirun`, because +the flag for that is implementation-specific — `-x VAR=VAL` on Open MPI, `-env VAR VAL` +on MPICH, and each rejects the other's spelling. Setting it in the shell works on both +for a single-node run. + +### What each dimension buys + +| | parallelizes | shards memory? | constraint | +|---|---|---|---| +| **`b_comm_size`** | the basis itself — each rank owns a block, and the blocks form the ring the ket rotates around | **yes — the only one that does** | ≥ 1 | +| **`t_comm_size`** | the ring's stations, each task rank starting at its own offset | no — costs `t` replicas of the ket | **`t ≤ b`** | +| **`helper`** | rows within a block (`idet % helper`), reduced by an allreduce | no — every helper rank builds the full excitation lookup | none on CPU; **must be 1 on Thrust** | + +`t ≤ b` is structural, not a limitation of this wrapper: the matvec rotates the ket +around `b_comm` as a ring (`gdb/mult.h:39-41`, `:196-202`), so there are exactly +`b_comm_size` stations and one task is one station. Upstream does not check it and a +starved task rank dereferences an empty lookup (`gdb/helper.h:761`), so `gdb_diag` +rejects it rather than letting it segfault. + +Because only `b` shards memory, `--b_comm_size 1` on many ranks means every rank +holds the whole basis *and* the whole excitation lookup no matter how many ranks +there are. That is the wall to watch at scale. + +### The shard contract + +With `--b_comm_size R > 1` each rank passes **its own shard**, not the whole list: + +- the shard index is `rank % R`, and ranks sharing one must pass **identical** + determinants (the in-memory path does not broadcast the list); +- shards must be **globally sorted and disjoint** — shard `i` strictly below shard + `i + 1`, which slicing a globally sorted list gives you for free; +- the union must be the basis you meant. That part cannot be checked here, so + compare the returned `global_dim` against what you expect. + +Everything else is checked and raises rather than quietly diagonalizing the wrong +subspace. + +### Placement across `b_comm` + +`--determinant_distribution` chooses who owns what. All six agree on the energy — +placement is a load-balancing decision — so pick by balance, not by result. + +| scheme | equalizes | leaves uneven | +|---|---|---| +| `input` | nothing; your slices are kept | whatever imbalance you passed | +| `equal-bra-a` *(default)* | distinct **alpha** strings per rank, which is what the matvec's outer loop runs over | determinant counts, when one alpha has many betas | +| `count` | determinant counts | one alpha's determinants can straddle ranks | +| `count-sorted` | as `count`, plus a reorder within each rank | same | +| `grid-cyclic` | spreads alpha and beta keys cyclically over a `grid_a × grid_b` grid | no count rebalance after | +| `grid-cyclic-balanced` | `grid-cyclic`, then a count rebalance | — | + +The grid schemes exist because a one-dimensional alpha split can be badly skewed for +an irregular subspace — exactly the sampled-bitstring case. `--determinant_grid_a` +and `--determinant_grid_b` must multiply to `b_comm_size`; give both or neither, and +the default is the factor pair nearest square. + +For Fe4S4 the four shipped files are already the balanced globally-sorted split +(14,884 each, disjoint, individually sorted, concatenation in order), so at +`--b_comm_size 4` the driver hands file *i* to shard *i* and `input` placement is +already count-balanced. + +### Choosing a backend + +`--device` selects the compute backend per run: + +```bash +--device cpu # host OpenMP (default) +--device gpu # NVHPC Thrust, NVIDIA only -- the only GPU backend with GDB kernels +``` + +`gpu-omp` and `auto` are accepted too, but for GDB neither is a GPU path: `gpu-omp` has +no GDB kernels and silently diagonalizes on the host (the drivers warn), and `auto` +prefers Thrust but falls back to `gpu-omp` where Thrust is not built — landing on the +host as well. **For GDB, pick `cpu` or `gpu` explicitly.** + +`sbd.available_backends()` reports what this install actually built — a static scan, so +it is safe to call outside `mpirun` — and `sbd.loaded_backends()` reports what the +current process has imported. + +One interaction to know about: **`'cpu'` and `'gpu-omp'` must not be used in the same +process.** Both link the same OpenMP runtime and the CPU module is built without offload +support, so whichever loads first initializes that runtime; if it is the CPU backend, the +offload backend can no longer acquire a device and silently runs on the host — right +answers, exit 0, idle GPU. `'cpu'` and `'gpu'` (Thrust) coexist fine, since Thrust does +not route device work through OpenMP. + +### On GPUs + +GDB has **Thrust kernels only** (`--device gpu`, NVIDIA). There is no `omp target` +code under `include/sbd/chemistry/gdb/` at all, so under `--device gpu-omp` a GDB run +pins a device and then diagonalizes on the host; the driver warns when it resolves +there. AMD has no Thrust build, so AMD GDB is CPU-only. + +The Thrust path additionally requires **`helper == 1`** (`gdb/mult_thrust.h:310-314`, +checked on every kernel launch), so every rank must go to `t_comm_size × b_comm_size`. +Since `helper = ranks / (t × b)`, leaving the basis in one block caps GPU GDB at a +single rank — the helper dimension absorbs every rank you add. So **multi-GPU GDB +requires `--b_comm_size`**, e.g. `-np 4 --b_comm_size 4`. Asking for more ranks than +`t × b` is refused up front, naming the helper dimension, rather than throwing from +inside a kernel launch. + +Only the Davidson runs on the device. `gdb/expansion.h` and `gdb/carryover.h` contain +no `thrust::` code and are included outside any `SBD_THRUST` guard (`inc_all.h:23`), so +heatbath expansion and carryover selection run on the host with OpenMP even in a GPU +build. Keep `OMP_NUM_THREADS` generous on GPU runs — one thread per rank throttles that +half of every iteration — and expect less than full GPU utilisation during a ladder for +the same reason. + +### `method` + +Only 0 and 1. TPB's 2 and 3 select Lanczos, which GDB does not implement; `gdb_diag` +rejects them rather than letting `gdb::diag` fall through both of its branches and +return an uninitialized energy. + +### Outputs, and what is replicated + +| key | distribution | +|---|---| +| `energy` | replicated, bit-identical on every rank | +| `density`, `one_p_rdm`, `two_p_rdm` | replicated on every rank | +| `carryover_det` | **this rank's shard**, as an `(n, words)` array. For `carryover_type` 1 it is split over `b_comm` *and duplicated* across the helper dimension, so gathering means one representative per `rank % b_comm_size`; for types 2 and 3 it is split over the world communicator with no duplication, so a plain allgather is correct | +| `local_dim` / `global_dim` | determinants on this rank, and summed over `b_comm` | +| `savename` | optional, and off by default: writes the amplitudes to one file per shard rather than returning them — see below | + +### Getting the amplitudes, if you want them + +Usually you do not. `gdb::diag` has no in-memory output for the wavefunction, and +nothing in the normal flow needs it: the energy, density and RDMs are returned directly, +and a heatbath ladder takes its next subspace from `carryover_det` (see [The loop needs +no amplitudes](#the-loop-needs-no-amplitudes)). GDB is also not wired into +qiskit-addon-sqd, whose `SCIState` would be the usual consumer — that path is TPB-only. + +When you do want them — your own analysis, or a selection step written in Python — pass +`savename` and read the files back. SBD writes one per b_comm position, +`f"{savename}{rank_b:06d}.bin"`: a single `…000000.bin` at `b_comm_size 1`, otherwise +`b_comm_size` files each holding only that shard, so no single file is the whole +wavefunction. Each is two `size_t` headers (`n_dets`, `words_per_det`), then +`n_dets × words_per_det` `size_t` determinant words in canonical order, then `n_dets` +`float64` amplitudes. + +## See Also + +- [`../tpb/README.md`](../tpb/README.md) — TPB, the SQD loops, and the bundled test data +- [Repository README](../../README.md) — installation, API reference diff --git a/examples/gdb/run_gdb_diag.py b/examples/gdb/run_gdb_diag.py new file mode 100755 index 0000000..8683e6d --- /dev/null +++ b/examples/gdb/run_gdb_diag.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python3 + +# 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. + +""" +Standalone GDB diagonalization over an explicit determinant list, in memory. + +GDB (general determinant basis) spans the subspace with the determinants it is +given. TPB spans it with the Cartesian product of an alpha and a beta list, so +TPB's dimension is |adet| x |bdet| while GDB's is exactly the number of +determinants passed. That is the whole point of GDB: an arbitrary sparse +subspace -- for example a set of sampled bitstrings used as sampled, with no +product completion. + +This driver stays on the in-memory entry point (``sbd.gdb_diag``) throughout. +Determinant text is read in Python and handed to the binding as a list; SBD's +own file-based entry point is deliberately not used. + +The determinant list can be sharded: with ``--b_comm_size R`` each rank passes only +its own slice, which is the only way GDB's memory scales, because b_comm is the one +dimension that divides the basis (and with it the excitation lookup). The derived +helper dimension divides work but not storage. ``--t_comm_size`` must not exceed +``--b_comm_size`` (one task per basis-ring station), and their product must divide +the rank count. + +Usage: + # Fe4S4, upstream's own GDB data: 4 files, 59,536 determinants, 36 orbitals. + # Note that upstream publishes no reference energy for this case. + python run_gdb_diag.py + + # Shard Fe4S4's four files one per rank: file i goes to b_comm position i, + # which for this data is already the balanced globally-sorted split + mpirun -np 4 -x OMP_NUM_THREADS=8 python run_gdb_diag.py --b_comm_size 4 + + # Spend ranks on both named dimensions (t <= b, and t*b must divide the ranks) + mpirun -np 8 python run_gdb_diag.py --b_comm_size 4 --t_comm_size 2 + + # Choose how determinants are placed across b_comm + mpirun -np 4 python run_gdb_diag.py --b_comm_size 4 \ + --determinant_distribution grid-cyclic + + # A subspace TPB can also express: interleave an alpha list with itself to + # form the full |A|^2 product basis. This is the cross-check against + # tpb_diag -- same subspace, two independent solvers, energies must agree. + python run_gdb_diag.py \ + --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ + --from-alpha ../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt \ + --alpha-limit 30 + + # Grow the subspace with SBD's own heatbath expansion (one round; the + # expanded list comes back as carryover_det) + python run_gdb_diag.py --carryover_type 2 --heatbath_cutoff 1e-4 +""" + +import argparse +import sys +import time + +import numpy as np + +# Fe4S4 is the only general-determinant data upstream ships, and it lives in the +# app directory rather than under data/ (which is alpha-only, i.e. TPB). +_UPSTREAM = "../../vendor/sbd-upstream" +_GDB_APP = f"{_UPSTREAM}/apps/chemistry_gdb_selected_basis_diagonalization" +_DEFAULT_FCIDUMP = f"{_GDB_APP}/fcidump_Fe4S4.txt" +_DEFAULT_DETFILES = ",".join(f"{_GDB_APP}/det{i}.txt" for i in range(4)) + + +# One 256-entry popcount table, applied to the packed determinant words by viewing +# them as bytes. numpy.bitwise_count would be ~9x faster but needs numpy 2.0, and +# this package's floor is 1.19; at 0.018 s per million determinants the table is +# already a ~6% addition to the read-and-pack this replaces nothing of. +_POPCOUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8) + +# Interleaved layout: alpha sits on the even bit positions, beta on the odd ones. +_ALPHA_MASK = np.uint64(0x5555555555555555) +_BETA_MASK = np.uint64(0xAAAAAAAAAAAAAAAA) + + +def _popcount_rows(words): + """Bits set per row of a (n, words) uint64 array.""" + return _POPCOUNT[words.view(np.uint8)].reshape(words.shape[0], -1).sum( + axis=1, dtype=np.int64) + + +def spin_weights(det): + """Per-determinant (n_alpha, n_beta) for a packed interleaved determinant array.""" + words = np.ascontiguousarray(det, dtype=np.uint64) + if words.ndim == 1: + words = words.reshape(1, -1) + return _popcount_rows(words & _ALPHA_MASK), _popcount_rows(words & _BETA_MASK) + + +def check_spin_weights(det, nelec, ms2, bit_length=64): + """Refuse a determinant list whose electron counts per spin are not uniform. + + A wrong bit order is the failure this catches, and it is worth catching here + because it does not fail cleanly downstream: concatenated ``[beta | alpha]`` + strings (what qiskit-addon-sqd emits) diagonalize to a plausible-looking + energy and only abort later inside the heatbath expansion. The occupation + density cannot catch it either, since permuting bits preserves how many are + set, so it still sums to the right electron count. + + Costs about 0.018 s per million determinants -- a few percent of the read and + pack that precede it. + """ + if bit_length != 64: + return # the masks above assume 64-bit words + want_a = (nelec + ms2) // 2 + want_b = nelec - want_a + got_a, got_b = spin_weights(det) + bad = (got_a != want_a) | (got_b != want_b) + n_bad = int(bad.sum()) + if not n_bad: + return + first = int(np.argmax(bad)) + raise ValueError( + f"{n_bad} of {len(bad)} determinants do not have {want_a} alpha and " + f"{want_b} beta electrons (first at index {first}: " + f"{int(got_a[first])} alpha, {int(got_b[first])} beta).\n" + " The usual cause is bit order: GDB expects alpha and beta INTERLEAVED " + "(bit 2*i alpha orbital i, bit 2*i+1 beta orbital i), while " + "qiskit-addon-sqd emits them CONCATENATED as [beta | alpha]. Interleave " + "before packing -- see examples/gdb/README.md.\n" + " The other cause is samples that were never postselected on the target " + "Hamming weight.\n" + " Pass --skip-weight-check to proceed anyway." + ) + + +def looks_like_a_single_determinant(density, tolerance=1e-6): + """True when every occupancy is 0 or 1: a one-determinant wavefunction. + + For a subspace of more than one determinant that means Davidson never iterated -- + the start vector had no coupling to the rest, so the residual was zero at once. + """ + return all(min(abs(x), abs(1.0 - x)) < tolerance for x in density) + + +def parse_args(): + """Parse command line arguments for all GDB_SBD parameters.""" + parser = argparse.ArgumentParser( + description="GDB diagonalization over an explicit determinant list", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + parser.add_argument('--device', default='cpu', + choices=['cpu', 'gpu', 'gpu-omp', 'auto'], + help="Compute backend. 'gpu' is NVHPC Thrust, which is the " + "only GPU backend with GDB kernels; 'gpu-omp' has none, " + "so GDB runs on the host there") + + # --- input ------------------------------------------------------------ + parser.add_argument('--skip-weight-check', action='store_true', + dest='skip_weight_check', + help='Do not verify that every determinant has the expected ' + 'alpha/beta electron count. The check costs about ' + '0.02 s per million determinants and catches a wrong ' + 'bit order, which otherwise returns a plausible wrong ' + 'energy with no error at all') + parser.add_argument('--fcidump', default=_DEFAULT_FCIDUMP, + help='FCIDUMP file defining the Hamiltonian') + parser.add_argument('--detfiles', default=_DEFAULT_DETFILES, + help='Comma-separated files of full determinants, one ' + '2*norb-bit string per line. Concatenated in Python ' + 'and passed as a single in-memory list') + parser.add_argument('--from-alpha', '--alpha-file', default='', metavar='FILE', + dest='from_alpha', + help='Instead of --detfiles, read a norb-bit alpha list and ' + 'form the full |A|^2 product basis by interleaving it ' + 'with itself. This is the subspace TPB would build from ' + 'the same file, so the two solvers are comparable') + parser.add_argument('--alpha-limit', type=int, default=0, dest='alpha_limit', + help='With --from-alpha, keep only the first N alpha strings. ' + '0 means all, which costs |A|^2 determinants') + + # --- MPI decomposition ------------------------------------------------ + parser.add_argument('--b_comm_size', type=int, default=1, + help='Basis communicator size, i.e. how many shards the ' + 'determinant list is split into. The only dimension that ' + 'divides memory') + parser.add_argument('--t_comm_size', type=int, default=1, + help='Task communicator size. Must not exceed b_comm_size: GDB ' + 'runs one task per basis-ring station and there are ' + 'exactly b_comm_size of them') + parser.add_argument('--determinant_distribution', default='', + choices=['', 'input', 'equal-bra-a', 'count', 'count-sorted', + 'grid-cyclic', 'grid-cyclic-balanced'], + help="How to place determinants across b_comm. Default " + "equal-bra-a, which equalizes distinct alpha strings per " + "rank -- the quantity that balances the matvec") + parser.add_argument('--determinant_grid_a', type=int, default=0, + help='Grid rows for the grid-cyclic schemes; with ' + '--determinant_grid_b must multiply to b_comm_size') + parser.add_argument('--determinant_grid_b', type=int, default=0, + help='Grid columns; see --determinant_grid_a') + + # --- diagonalization -------------------------------------------------- + # GDB implements Davidson only: there is no Lanczos anywhere under + # include/sbd/chemistry/gdb/. gdb::diag is `if(method==0){} else if(method==1){}` + # with no else, and `energy` is assigned only inside those branches + # (gdb/sbdiag.h:418, :501), so method 2 or 3 returns an uninitialized double on + # the CPU backend. The Thrust path is accidentally safe -- sbdiag.h:210-211 does + # `method &= 1`. Hence choices=[0, 1], unlike run_sbd_diag.py's [0, 1, 2, 3]. + parser.add_argument('--method', type=int, default=0, choices=[0, 1], + help='0=Davidson, 1=Davidson storing the Hamiltonian. GDB has ' + 'no Lanczos, so TPB methods 2 and 3 do not exist here') + parser.add_argument('--iteration', '--max_it', type=int, default=100, + dest='max_it', help='Maximum Davidson iterations') + parser.add_argument('--block', '--max_nb', type=int, default=10, + dest='max_nb', help='Maximum number of basis vectors') + parser.add_argument('--tolerance', '--eps', type=float, default=1e-6, + dest='eps', help='Convergence tolerance') + parser.add_argument('--max_time', type=float, default=1e10, + help='Maximum wall time in seconds') + parser.add_argument('--init', type=int, default=0, + help='Initial vector policy') + parser.add_argument('--seed', type=int, default=1729, + help='Seed for a random initial vector (init != 0)') + parser.add_argument('--bit_length', type=int, default=64, + help='Bits per packed word. Words per determinant is ' + 'ceil(2*norb / bit_length)') + + # --- RDMs ------------------------------------------------------------- + parser.add_argument('--rdm_output', default='', metavar='FILE', + help='Save spin-summed rdm1 and rdm2 to one .npz file. ' + 'Implies do_rdm=1') + + # --- carryover / expansion ------------------------------------------- + parser.add_argument('--carryover_type', type=int, default=0, + choices=[0, 1, 2, 3], + help='0=off, 1=keep the top --ratio fraction by weight, ' + '2/3=weight-truncate then heatbath-expand (2 and 3 are ' + 'the two heatbath variants). 2/3 GROW the subspace') + parser.add_argument('--carryover_ratio', '--ratio', type=float, default=0.0, + dest='ratio', + help='Fraction of determinants kept by carryover_type 1') + parser.add_argument('--carryover_threshold', '--threshold', type=float, + default=0.01, dest='threshold', + help='Weight threshold for carryover_type 1') + parser.add_argument('--heatbath_cutoff', type=float, default=1e-4, + help='Integral-magnitude cutoff admitting a candidate ' + 'determinant (carryover_type 2/3)') + parser.add_argument('--heatbath_truncation', type=float, default=0.0, + help='Weight threshold applied to parents before expanding') + parser.add_argument('--heatbath_batch_size', type=int, default=1000000, + help='Heatbath expansion batch size per rank') + + # --- wavefunction I/O ------------------------------------------------- + parser.add_argument('--loadname', default='', + help='Load an initial wavefunction from this file') + parser.add_argument('--savename', default='', + help='Save the final wavefunction under this prefix') + + return parser.parse_args() + + +def read_det_strings(paths): + """Read full-determinant bitstrings from one or more text files. + + Returns them in file order, unfiltered: gdb_diag sorts its own input and + rejects duplicates, so both properties are left for it to enforce rather + than silently repaired here. + """ + strings = [] + for path in paths: + with open(path, encoding='utf-8') as handle: + for line in handle: + line = line.strip() + if line: + strings.append(line) + return strings + + +def interleave_spin_strings(alpha, beta): + """Interleave an alpha and a beta bitstring into one GDB determinant string. + + Both inputs are norb-bit strings written most-significant-first, as in + SBD's alpha determinant files. The result is 2*norb bits in the order the + GDB app documents: reading from the right, alpha orbital 1, beta orbital 1, + alpha orbital 2, and so on -- i.e. bit 2*i is alpha orbital i and bit + 2*i + 1 is beta orbital i, which is what from_string() then packs. + """ + a_rev = alpha[::-1] + b_rev = beta[::-1] + return ''.join(a_rev[i] + b_rev[i] for i in range(len(a_rev)))[::-1] + + +def shard_bounds(total, b_comm_size, index): + """This rank's slice of a globally sorted list of ``total`` determinants. + + Mirrors SBD's own ``q = N/p``, remainder-to-the-low-ranks split + (``balanced_begin``, framework/bit_manipulation.h:1523-1531), so an ``input`` + placement is already the count-balanced one. + """ + quotient, remainder = divmod(total, b_comm_size) + begin = index * quotient + min(index, remainder) + return begin, begin + quotient + (1 if index < remainder else 0) + + +def product_basis(alpha_strings): + """Every (alpha, beta) pair from one list: the subspace TPB spans from it.""" + return [interleave_spin_strings(a, b) + for a in alpha_strings for b in alpha_strings] + + +def main(): + args = parse_args() + + import sbd + from sbd.sbd_solver import assemble_rdms + + sbd.init(device=args.device) + + rank = sbd.get_rank() + size = sbd.get_world_size() + device = sbd.get_device() + + if rank == 0: + print("=" * 70) + print("SBD GDB - general determinant basis, in memory") + print("=" * 70) + sbd.print_info() + print() + + # Every rank builds the same list: that is what the in-memory entry point + # expects, and it is why b_comm_size must be 1. + fcidump = sbd.LoadFCIDump(args.fcidump) + norb = int(fcidump.header["NORB"]) + nelec = int(fcidump.header["NELEC"]) + ms2 = int(fcidump.header.get("MS2", 0)) + total_bits = 2 * norb + + b_size = args.b_comm_size + shard_index = rank % b_size if b_size > 1 else 0 + + if args.from_alpha: + alpha = read_det_strings([args.from_alpha]) + if args.alpha_limit: + alpha = alpha[:args.alpha_limit] + bad = [s for s in alpha if len(s) != norb] + if bad: + print(f"ERROR: --from-alpha strings must be {norb} bits " + f"(NORB from the FCIDUMP); found one of length {len(bad[0])}", + file=sys.stderr) + return 1 + det_strings = product_basis(alpha) + source = (f"{args.from_alpha} ({len(alpha)} alpha strings -> " + f"{len(alpha)}^2 product determinants)") + whole = sbd.sort_bitarray_array( + sbd.from_strings(det_strings, args.bit_length, total_bits)) + begin, end = shard_bounds(whole.shape[0], b_size, shard_index) + det = whole[begin:end] if b_size > 1 else whole + placement = (f"sliced from the globally sorted list " + f"[{begin}:{end}]") if b_size > 1 else "whole basis" + else: + paths = [p for p in args.detfiles.split(',') if p] + if not paths: + print("ERROR: no determinant files given", file=sys.stderr) + return 1 + # When the files already partition the basis -- each sorted, disjoint, and + # in globally sorted order, as gdet emits and Fe4S4's four files are -- + # a rank can read only its own share and nothing is read twice. + # + # The share must be a CONTIGUOUS BLOCK of files, not a stride: shard i has + # to hold a strictly lower range than shard i+1, so with 24 files over 8 + # ranks it is files 0-2, 3-5, ... and NOT 0,8,16. A strided assignment + # would interleave the ranges and fail the sorted-disjoint check. + took_file_block = b_size > 1 and len(paths) % b_size == 0 + if took_file_block: + per_rank = len(paths) // b_size + first = shard_index * per_rank + mine = paths[first:first + per_rank] + placement = (f"file {first}" if per_rank == 1 else + f"files {first}-{first + per_rank - 1}") + placement += f" of {len(paths)}" + elif b_size > 1: + mine = paths + placement = "sliced from the globally sorted concatenation" + else: + mine = paths + placement = "whole basis" + det_strings = read_det_strings(mine) + bad = [s for s in det_strings if len(s) != total_bits] + if bad: + print(f"ERROR: determinant strings must be {total_bits} bits " + f"(2 * NORB); found one of length {len(bad[0])}", file=sys.stderr) + return 1 + if not det_strings: + print("ERROR: no determinants read", file=sys.stderr) + return 1 + det = sbd.sort_bitarray_array( + sbd.from_strings(det_strings, args.bit_length, total_bits)) + if b_size > 1 and not took_file_block: + begin, end = shard_bounds(det.shape[0], b_size, shard_index) + det = det[begin:end] + placement += f" [{begin}:{end}]" + source = f"{len(paths)} file(s): {', '.join(paths)}" + + if det.shape[0] == 0: + print("ERROR: no determinants read", file=sys.stderr) + return 1 + if not args.skip_weight_check: + try: + check_spin_weights(det, nelec, ms2, args.bit_length) + except ValueError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + words = det.shape[1] + + config = sbd.GDB_SBD() + config.t_comm_size = args.t_comm_size + config.b_comm_size = args.b_comm_size + config.method = args.method + config.max_it = args.max_it + config.max_nb = args.max_nb + config.eps = args.eps + config.max_time = args.max_time + config.init = args.init + config.seed = args.seed + config.do_rdm = 1 if args.rdm_output else 0 + config.bit_length = args.bit_length + config.carryover_type = args.carryover_type + config.ratio = args.ratio + config.threshold = args.threshold + config.heatbath_cutoff = args.heatbath_cutoff + config.heatbath_truncation = args.heatbath_truncation + config.heatbath_batch_size = args.heatbath_batch_size + + if rank == 0: + grid = args.t_comm_size * args.b_comm_size + helper = size // grid if grid and size % grid == 0 else 0 + print("Configuration:") + print(f" Device: {device}") + print(f" Method: {'Davidson' if args.method == 0 else 'Davidson + stored H'}") + print(f" Max iterations: {config.max_it}") + print(f" Tolerance: {config.eps}") + print(f" MPI ranks: {size}") + print(f" Decomposition: t_comm_size={args.t_comm_size} " + f"b_comm_size={args.b_comm_size} helper={helper}") + print(f" Placement: {args.determinant_distribution or 'equal-bra-a (default)'}") + if args.b_comm_size == 1 and size > 1: + print(" NOTE: b_comm_size=1 means every rank holds the whole basis and " + "the whole excitation lookup. Only b_comm_size shards memory -- the " + "helper dimension divides work but not storage.") + if device == 'gpu' and helper != 1: + print(" ERROR: GDB on the Thrust backend requires helper == 1, i.e. " + "t_comm_size * b_comm_size == ranks. Give every rank to " + "--b_comm_size (with --t_comm_size <= it).") + if device == 'gpu-omp': + print(" WARNING: the OMP-offload backend has no GDB kernels -- there is " + "no `omp target` code under include/sbd/chemistry/gdb/ at all. " + "This run pins a device and then diagonalizes on the host. Use " + "--device gpu (Thrust, NVIDIA only) for GDB on a GPU.") + print() + print("Subspace:") + print(f" Source: {source}") + print(f" This rank: {placement}") + print(f" Determinants on this rank: {det.shape[0]}") + print(f" Orbitals: {norb} ({total_bits} spin orbitals), electrons: {nelec}") + print(f" Packing: {words} word(s) of {args.bit_length} bits") + print() + print("Running GDB diagonalization...") + print() + + start = time.perf_counter() + results = sbd.gdb_diag( + fcidump, det, config, + loadname=args.loadname, savename=args.savename, + determinant_distribution=args.determinant_distribution, + determinant_grid_a=args.determinant_grid_a, + determinant_grid_b=args.determinant_grid_b, + ) + elapsed = time.perf_counter() - start + + if rank != 0: + return 0 + + print("=" * 70) + print("Results") + print("=" * 70) + print(f"Device: {device.upper()}") + print(f"Ground state energy: {results['energy']:.10f} Hartree") + print(f"Subspace dimension: {results['global_dim']} " + f"({results['local_dim']} on this rank)") + print(f"Placement applied: {results['determinant_distribution']}") + print(f"Wall time: {elapsed:.2f} s") + + density = results['density'] + combined = [density[2 * i] + density[2 * i + 1] for i in range(len(density) // 2)] + print(f"Density: {np.round(combined, 6).tolist()}") + print(f" (sums to {sum(combined):.6f}; should equal the electron count {nelec})") + if results["global_dim"] > 1 and looks_like_a_single_determinant(density): + print(" WARNING: every occupancy is 0 or 1, so the wavefunction is a single " + "determinant") + print(" and Davidson did not iterate (look for 'tol=0' above). The energy is " + "a diagonal") + print(" element -- a valid upper bound, not the subspace's ground state. " + "Expanding the") + print(" subspace once (--carryover_type 2) gives Davidson something to " + "couple.") + + carryover = results['carryover_det'] + if args.carryover_type: + print(f"Carryover determinants on this rank: {carryover.shape[0]}") + if size > 1: + print(" NOTE: the carryover list is this rank's shard, not the whole " + "list -- " + "HeatbathExpansion ends with sort_global_bitarray + " + "redistribution_bitarray over the WORLD communicator " + "(gdb/expansion.h), so each rank holds a different slice and a " + "driver that feeds it back must allgather first.") + if args.carryover_type in (2, 3): + print(" The expanded list includes its parents: " + "local_heatbath_expansion starts from `edet = det` " + "(gdb/expansion.h:543), so this is the next subspace, not just " + "the new candidates.") + + if args.rdm_output: + # assemble_rdms' layout was verified against PySCF on the TPB path; the + # keys and flat-index convention are shared with GDB, but that + # verification has not been repeated here. + rdm1, rdm2 = assemble_rdms(results, norb) + if rdm1 is None: + print("No RDMs returned (do_rdm was 0)") + else: + np.savez(args.rdm_output, rdm1=rdm1, rdm2=rdm2) + print(f"\nRDMs saved to {args.rdm_output}") + print(f"1-RDM trace: {np.trace(rdm1):.6f} " + f"(should equal the electron count {nelec})") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/gdb/run_gdb_heatbath.py b/examples/gdb/run_gdb_heatbath.py new file mode 100755 index 0000000..9ec03c2 --- /dev/null +++ b/examples/gdb/run_gdb_heatbath.py @@ -0,0 +1,587 @@ +#!/usr/bin/env python3 + +# 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. + +""" +Grow a GDB subspace with SBD's own heatbath expansion, round after round. + +This is selected CI (HCI) driven from Python: diagonalize, let SBD expand the +subspace from the resulting wavefunction, diagonalize the larger subspace, repeat. +No new machinery is needed for the loop itself -- ``carryover_type`` 2 and 3 return +the parents *together with* the new candidates (``local_heatbath_expansion`` opens +with ``edet = det``), so the result of one round IS the next round's subspace. + +**The cutoff is the dial, and the loop converges per cutoff.** Expansion admits a +candidate when its estimated contribution exceeds ``heatbath_cutoff``, so for a +fixed cutoff the subspace reaches a self-consistent size and stops growing -- +further rounds buy nothing. Going deeper means lowering the cutoff. Hence a +*ladder*: spend rounds at one cutoff until the energy stops moving, then step to +the next. ``--cutoffs`` takes the whole ladder. + +Do not confuse ``--heatbath_truncation`` with the cutoff: it discards *parents* +by weight before expansion even starts, and its default of 0 (keep everything) is +almost always what you want. Setting it to 1e-4 on a 59,536-determinant Fe4S4 +wavefunction cut the subspace to 506. + +Seeds, so the same driver produces every row of a seed comparison: + + --subspace-from files determinant files (default: upstream's four Fe4S4 files) + --subspace-from hf the Hartree-Fock determinant alone, the no-input null + --subspace-from from-alpha an alpha list interleaved with itself (a TPB-shaped space) + --subspace-from strings an arbitrary bitstring file, e.g. sampled configurations + +Usage: + # Fe4S4 from upstream's shipped subspace, one cutoff + python run_gdb_heatbath.py --cutoffs 1e-3 + + # A ladder, stopping if the subspace would pass 2M determinants + python run_gdb_heatbath.py --cutoffs 1e-3,1e-4,1e-5 --max_dim 2000000 + + # The null hypothesis: no input subspace at all, just Hartree-Fock + python run_gdb_heatbath.py --subspace-from hf --cutoffs 1e-3,1e-4 + + # Sharded across 4 ranks/GPUs. At --b_comm_size == ranks with t=1 the + # expansion comes back already sharded for the next round, no gather needed. + mpirun -np 4 python run_gdb_heatbath.py --b_comm_size 4 --cutoffs 1e-3,1e-4 + mpirun -np 4 python run_gdb_heatbath.py --b_comm_size 4 --device gpu --cutoffs 1e-4 + + # Record the (dimension, energy) series for a comparison table + python run_gdb_heatbath.py --cutoffs 1e-3,1e-4 --log fe4s4_ladder.json +""" + +import argparse +import json +import sys +import time + +import numpy as np + +_UPSTREAM = "../../vendor/sbd-upstream" +_GDB_APP = f"{_UPSTREAM}/apps/chemistry_gdb_selected_basis_diagonalization" +_DEFAULT_FCIDUMP = f"{_GDB_APP}/fcidump_Fe4S4.txt" +_DEFAULT_DETFILES = ",".join(f"{_GDB_APP}/det{i}.txt" for i in range(4)) + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Iteratively grow a GDB subspace by heatbath expansion", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + parser.add_argument('--device', default='cpu', + choices=['cpu', 'gpu', 'gpu-omp', 'auto'], + help="Compute backend. 'gpu' (Thrust) is the only GPU backend " + "with GDB kernels, and it requires the helper dimension " + "to be 1, i.e. t_comm_size * b_comm_size == ranks") + parser.add_argument('--fcidump', default=_DEFAULT_FCIDUMP, + help='FCIDUMP file defining the Hamiltonian') + + # --- the seed --------------------------------------------------------- + parser.add_argument('--skip-weight-check', action='store_true', + dest='skip_weight_check', + help='Do not verify that every seed determinant has the ' + 'expected alpha/beta electron count. The check costs ' + 'about 0.02 s per million determinants and catches a ' + 'wrong bit order, which otherwise diagonalizes to a ' + 'plausible wrong energy and aborts during expansion') + parser.add_argument('--subspace-from', default='files', + dest='subspace_from', + choices=['files', 'hf', 'from-alpha', 'strings'], + help='Where the starting subspace comes from') + parser.add_argument('--detfiles', default=_DEFAULT_DETFILES, + help='--subspace-from files: comma-separated files of 2*norb-bit ' + 'determinant strings') + parser.add_argument('--alpha-file', '--from-alpha', default='', dest='alpha_file', + help='--subspace-from from-alpha: a norb-bit alpha determinant list, ' + 'interleaved with itself into the full product basis') + parser.add_argument('--alpha-limit', type=int, default=0, dest='alpha_limit', + help='--subspace-from from-alpha: keep only the first N alpha strings ' + '(the product costs N^2 determinants)') + parser.add_argument('--strings-file', default='', dest='strings_file', + help='--subspace-from strings: a file of 2*norb-bit determinant ' + 'strings, e.g. sampled configurations') + + # --- the ladder ------------------------------------------------------- + parser.add_argument('--cutoffs', default='1e-3', + help='Comma-separated heatbath cutoffs, smallest step last. ' + 'Each is a rung: rounds run at that cutoff until the ' + 'energy stops moving, then the next rung begins') + parser.add_argument('--max_rounds', type=int, default=8, + help='Maximum rounds per rung') + parser.add_argument('--energy_tol', type=float, default=1e-5, + help='Advance to the next rung once |dE| between rounds falls ' + 'below this (Hartree)') + parser.add_argument('--max_dim', type=int, default=0, + help='Stop before diagonalizing a subspace larger than this ' + '(0 = no cap). The cap is what keeps rows of a ' + 'comparison cost-matched') + parser.add_argument('--heatbath_truncation', type=float, default=0.0, + help='Weight threshold applied to PARENTS before expanding. ' + 'Leave at 0 unless you mean to expand only from the ' + 'dominant determinants') + parser.add_argument('--heatbath_batch_size', type=int, default=1000000, + help='Heatbath expansion batch size per rank') + + # --- diagonalization -------------------------------------------------- + parser.add_argument('--method', type=int, default=0, choices=[0, 1], + help='0=Davidson, 1=Davidson storing the Hamiltonian. GDB has ' + 'no Lanczos, so TPB methods 2 and 3 do not exist here') + parser.add_argument('--tolerance', '--eps', type=float, default=1e-6, + dest='eps', help='Davidson convergence tolerance') + parser.add_argument('--iteration', '--max_it', type=int, default=30, + dest='max_it', help='Maximum Davidson iterations per round') + parser.add_argument('--block', '--max_nb', type=int, default=10, + dest='max_nb', help='Maximum number of basis vectors') + parser.add_argument('--bit_length', type=int, default=64, + help='Bits per packed word') + parser.add_argument('--carryover_type', type=int, default=2, choices=[2, 3], + help='Heatbath variant: 2 or 3. Types 0 and 1 do not expand, ' + 'so they cannot drive this loop') + + # --- decomposition ---------------------------------------------------- + parser.add_argument('--b_comm_size', type=int, default=1, + help='Basis shards. The only dimension that divides memory') + parser.add_argument('--t_comm_size', type=int, default=1, + help='Task communicator size; must not exceed b_comm_size') + parser.add_argument('--determinant_distribution', default='', + choices=['', 'input', 'equal-bra-a', 'count', 'count-sorted', + 'grid-cyclic', 'grid-cyclic-balanced'], + help='How determinants are placed across b_comm') + + parser.add_argument('--log', default='', metavar='FILE', + help='Write the per-round (dimension, energy) series as JSON') + + args = parser.parse_args() + + # One determinant cannot be sharded: rank 0 owns it, the rest get empty shards, + # and a single parent gives OpenMP nothing to divide either. + if args.subspace_from == 'hf' and args.b_comm_size > 1: + print("WARNING: --subspace-from hf is a single determinant, so " + f"--b_comm_size {args.b_comm_size} leaves the", file=sys.stderr) + print(" other ranks with empty shards and runs the first rounds on one " + "core. Use -np 1.", file=sys.stderr) + + return args + + +def read_strings(paths): + """Read bitstrings from one or more text files, in file order.""" + out = [] + for path in paths: + with open(path, encoding='utf-8') as handle: + out.extend(line.strip() for line in handle if line.strip()) + return out + + +def interleave(alpha, beta): + """Interleave a norb-bit alpha and beta string into one 2*norb-bit determinant. + + Bit ``2 * i`` is alpha orbital ``i`` and bit ``2 * i + 1`` is beta orbital ``i``, + counting from the right -- the order the GDB app documents and ``from_strings`` + packs. + """ + a_rev, b_rev = alpha[::-1], beta[::-1] + return ''.join(a_rev[i] + b_rev[i] for i in range(len(a_rev)))[::-1] + + +# One 256-entry popcount table, applied to the packed determinant words by viewing +# them as bytes. numpy.bitwise_count would be ~9x faster but needs numpy 2.0, and +# this package's floor is 1.19; at 0.018 s per million determinants the table is +# already a ~6% addition to the read-and-pack this replaces nothing of. +_POPCOUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8) + +# Interleaved layout: alpha sits on the even bit positions, beta on the odd ones. +_ALPHA_MASK = np.uint64(0x5555555555555555) +_BETA_MASK = np.uint64(0xAAAAAAAAAAAAAAAA) + + +def _popcount_rows(words): + """Bits set per row of a (n, words) uint64 array.""" + return _POPCOUNT[words.view(np.uint8)].reshape(words.shape[0], -1).sum( + axis=1, dtype=np.int64) + + +def spin_weights(det): + """Per-determinant (n_alpha, n_beta) for a packed interleaved determinant array.""" + words = np.ascontiguousarray(det, dtype=np.uint64) + if words.ndim == 1: + words = words.reshape(1, -1) + return _popcount_rows(words & _ALPHA_MASK), _popcount_rows(words & _BETA_MASK) + + +def check_spin_weights(det, nelec, ms2, bit_length=64): + """Refuse a determinant list whose electron counts per spin are not uniform. + + A wrong bit order is the failure this catches, and it is worth catching here + because it does not fail cleanly downstream: concatenated ``[beta | alpha]`` + strings (what qiskit-addon-sqd emits) diagonalize to a plausible-looking + energy and only abort later inside the heatbath expansion. The occupation + density cannot catch it either, since permuting bits preserves how many are + set, so it still sums to the right electron count. + + Costs about 0.018 s per million determinants -- a few percent of the read and + pack that precede it. + """ + if bit_length != 64: + return # the masks above assume 64-bit words + want_a = (nelec + ms2) // 2 + want_b = nelec - want_a + got_a, got_b = spin_weights(det) + bad = (got_a != want_a) | (got_b != want_b) + n_bad = int(bad.sum()) + if not n_bad: + return + first = int(np.argmax(bad)) + raise ValueError( + f"{n_bad} of {len(bad)} determinants do not have {want_a} alpha and " + f"{want_b} beta electrons (first at index {first}: " + f"{int(got_a[first])} alpha, {int(got_b[first])} beta).\n" + " The usual cause is bit order: GDB expects alpha and beta INTERLEAVED " + "(bit 2*i alpha orbital i, bit 2*i+1 beta orbital i), while " + "qiskit-addon-sqd emits them CONCATENATED as [beta | alpha]. Interleave " + "before packing -- see examples/gdb/README.md.\n" + " The other cause is samples that were never postselected on the target " + "Hamming weight.\n" + " Pass --skip-weight-check to proceed anyway." + ) + + +def looks_like_a_single_determinant(density, tolerance=1e-6): + """True when every occupancy is 0 or 1: a one-determinant wavefunction. + + For a subspace of more than one determinant that means Davidson never iterated -- + the start vector had no coupling to the rest, so the residual was zero at once. + """ + return all(min(abs(x), abs(1.0 - x)) < tolerance for x in density) + + +def hartree_fock_string(norb, nelec, ms2): + """The Hartree-Fock determinant: the lowest orbitals doubly occupied. + + Built from the FCIDUMP header alone, so ``--subspace-from hf`` needs no input subspace at + all. That is the point of it: if expansion from HF reaches the same place as + expansion from a sampled subspace, the sampling added nothing. + """ + n_alpha = (nelec + ms2) // 2 + n_beta = nelec - n_alpha + if n_alpha > norb or n_beta > norb: + raise ValueError( + f"cannot place {n_alpha} alpha and {n_beta} beta electrons in {norb} " + f"orbitals") + alpha = ''.join('1' if i < n_alpha else '0' for i in range(norb))[::-1] + beta = ''.join('1' if i < n_beta else '0' for i in range(norb))[::-1] + return interleave(alpha, beta) + + +def shard_bounds(total, b_comm_size, index): + """This rank's slice of a globally sorted list. + + Mirrors SBD's ``q = N/p``, remainder-to-the-low-ranks split + (``balanced_begin``, framework/bit_manipulation.h:1523-1531). + """ + quotient, remainder = divmod(total, b_comm_size) + begin = index * quotient + min(index, remainder) + return begin, begin + quotient + (1 if index < remainder else 0) + + +def build_seed(args, sbd, norb, nelec, ms2, rank): + """The starting determinant array for this rank, plus a description.""" + total_bits = 2 * norb + + if args.subspace_from == 'hf': + strings = [hartree_fock_string(norb, nelec, ms2)] + source = f"Hartree-Fock determinant ({nelec} electrons, MS2={ms2})" + elif args.subspace_from == 'from-alpha': + if not args.alpha_file: + raise ValueError("--subspace-from from-alpha requires --alpha-file") + alpha = read_strings([args.alpha_file]) + if args.alpha_limit: + alpha = alpha[:args.alpha_limit] + bad = [s for s in alpha if len(s) != norb] + if bad: + raise ValueError( + f"--alpha-file strings must be {norb} bits, found {len(bad[0])}") + strings = [interleave(a, b) for a in alpha for b in alpha] + source = f"{args.alpha_file}: {len(alpha)} alpha -> {len(alpha)}^2 product" + else: + paths = ([args.strings_file] if args.subspace_from == 'strings' + else [p for p in args.detfiles.split(',') if p]) + if not paths or not all(paths): + raise ValueError(f"--subspace-from {args.subspace_from} requires input file(s)") + strings = read_strings(paths) + source = f"{len(paths)} file(s): {', '.join(paths)}" + + bad = [s for s in strings if len(s) != total_bits] + if bad: + raise ValueError( + f"determinant strings must be {total_bits} bits (2 * NORB), " + f"found one of length {len(bad[0])}") + if not strings: + raise ValueError("the seed is empty") + + # Sorted globally here so that slicing gives the disjoint, ordered shards + # gdb_diag requires. + whole = sbd.sort_bitarray_array( + sbd.from_strings(strings, args.bit_length, total_bits, device=args.device), + device=args.device) + if args.b_comm_size == 1: + return whole, source + begin, end = shard_bounds(whole.shape[0], args.b_comm_size, + rank % args.b_comm_size) + return whole[begin:end], source + + +def reshard(det, args, comm, rank, size): + """Put the expanded list back into the shape the next round expects. + + ``HeatbathExpansion`` ends with a global sort and a count rebalance over the + WORLD communicator (gdb/expansion.h:821-822), so the result is already + globally sorted, balanced and disjoint across world ranks. When + ``b_comm_size == ranks`` and ``t_comm_size == 1`` the b_comm position equals the + world rank, so that is *exactly* a valid next-round shard and nothing needs to + move -- which is the configuration a GPU run uses anyway, since Thrust requires + the helper dimension to be 1. + + Any other grid means world position and b_comm position disagree, and the only + portable fix from Python is to gather and re-slice. That materializes the whole + list on every rank, so it is the slow path and says so. + """ + if args.b_comm_size == 1 or size == 1: + if size == 1: + return det + # Every rank must hold the whole basis; the shards are disjoint pieces of it. + gathered = comm.allgather(det) + return np.concatenate([g for g in gathered if g.size], axis=0) + + if args.b_comm_size == size and args.t_comm_size == 1: + return det # already aligned: world rank == b_comm position + + if rank == 0: + print(" NOTE: b_comm_size != ranks (or t_comm_size > 1), so the expanded " + "list has to be gathered and re-sliced, which puts the whole list on " + "every rank. Use --b_comm_size == ranks with --t_comm_size 1 to avoid " + "it.", flush=True) + gathered = comm.allgather(det) + whole = np.concatenate([g for g in gathered if g.size], axis=0) + begin, end = shard_bounds(whole.shape[0], args.b_comm_size, + rank % args.b_comm_size) + return whole[begin:end] + + +def main(): + args = parse_args() + + import sbd + from mpi4py import MPI + + sbd.init(device=args.device) + backend = sbd.get_backend(args.device) + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + + fcidump = backend.LoadFCIDump(args.fcidump) + header = fcidump.header + norb = int(header["NORB"]) + nelec = int(header["NELEC"]) + ms2 = int(header.get("MS2", 0)) + + cutoffs = [float(c) for c in args.cutoffs.split(',') if c.strip()] + if not cutoffs: + print("ERROR: --cutoffs is empty", file=sys.stderr) + return 1 + + try: + det, source = build_seed(args, sbd, norb, nelec, ms2, rank) + if not args.skip_weight_check: + check_spin_weights(det, nelec, ms2, args.bit_length) + except ValueError as exc: + if rank == 0: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + + def global_dim(local): + """Determinants across the whole basis, counted once each. + + Summing over the world communicator would over-count: ranks sharing a b_comm + position hold identical shards, so with t_comm_size or the helper dimension + above 1 each shard is counted once per replica. Contribute only from the ranks + that are one-per-b-position -- with the layout + ``rank = h*(b*t) + t_index*b + b_index``, those are exactly ranks + ``0 .. b_comm_size - 1`` -- which gives the b_comm sum without building a + sub-communicator. + """ + if args.b_comm_size == 1: + return int(local.shape[0]) + mine = int(local.shape[0]) if rank < args.b_comm_size else 0 + return comm.allreduce(mine, op=MPI.SUM) + + # Collective: every rank must call it, so it cannot live inside a rank guard. + # This deadlocks otherwise -- rank 0 waits in the allreduce while the others walk + # into gdb_diag's own collectives. It hides in serial testing because the + # b_comm_size == 1 path short-circuits without communicating. + starting_dim = global_dim(det) + if args.b_comm_size > 1 and det.shape[0] == 0: + # Printed by whichever rank is starved, which is the informative one. + print(f" NOTE: rank {rank}'s shard is empty -- the seed has fewer " + f"determinants than b_comm_size. 'count' and 'count-sorted' placement " + f"cannot take an empty shard; 'input' and the grid schemes can.", + flush=True) + + if rank == 0: + grid = args.t_comm_size * args.b_comm_size + helper = size // grid if grid and size % grid == 0 else 0 + print("=" * 78) + print("GDB + heatbath expansion - iterative subspace growth") + print("=" * 78) + print(f" Device: {sbd.get_device()} ranks: {size}") + print(f" Decomposition: t={args.t_comm_size} b={args.b_comm_size} " + f"helper={helper}") + print(f" System: {norb} orbitals ({2 * norb} spin orbitals), " + f"{nelec} electrons") + print(f" Seed: {source}") + print(f" Starting dimension: {starting_dim}") + print(f" Cutoff ladder: {cutoffs}") + print(f" Per rung: up to {args.max_rounds} rounds, advance when " + f"|dE| < {args.energy_tol:g}") + if args.max_dim: + print(f" Dimension cap: {args.max_dim}") + if args.heatbath_truncation: + print(f" WARNING: --heatbath_truncation {args.heatbath_truncation:g} " + f"discards parents BEFORE expanding, which can shrink the " + f"subspace rather than grow it.") + print() + print(f" {'rung':>4} {'round':>5} {'dimension':>12} {'energy':>18} " + f"{'dE':>12} {'secs':>7}") + + history = [] + previous_energy = None # for the displayed dE, continuous across rungs + stop_reason = "ladder complete" + + for rung, cutoff in enumerate(cutoffs): + rung_converged = False + # Deliberately NOT seeded from the previous rung's energy. Round 0 of a new + # rung re-diagonalizes the subspace the previous rung already expanded, so + # its energy is essentially unchanged -- testing convergence against that + # would declare the rung done before its smaller cutoff had expanded + # anything at all. + rung_baseline = None + for round_index in range(args.max_rounds): + dim = global_dim(det) + if args.max_dim and dim > args.max_dim: + stop_reason = f"dimension {dim} exceeds --max_dim {args.max_dim}" + if rank == 0: + print(f" stopping: {stop_reason}") + return _finish(args, history, stop_reason, rank) + + config = backend.GDB_SBD() + config.method = args.method + config.eps = args.eps + config.max_it = args.max_it + config.max_nb = args.max_nb + config.bit_length = args.bit_length + config.b_comm_size = args.b_comm_size + config.t_comm_size = args.t_comm_size + config.carryover_type = args.carryover_type + config.heatbath_cutoff = cutoff + config.heatbath_truncation = args.heatbath_truncation + config.heatbath_batch_size = args.heatbath_batch_size + + start = time.perf_counter() + result = sbd.gdb_diag( + fcidump, det, config, device=args.device, + determinant_distribution=args.determinant_distribution, + ) + elapsed = time.perf_counter() - start + + energy = result["energy"] + delta = None if previous_energy is None else energy - previous_energy + if rank == 0: + shown = "" if delta is None else f"{delta:+.6f}" + print(f" {rung:>4} {round_index:>5} {dim:>12} {energy:>18.10f} " + f"{shown:>12} {elapsed:>7.1f}", flush=True) + # Only the seed round can be a no-op; after that the subspace contains + # its own excitation neighbours. + if (rank == 0 and rung == 0 and round_index == 0 and dim > 1 + and looks_like_a_single_determinant(result["density"])): + print(" NOTE: the seed diagonalization returned a single " + "determinant, so Davidson did") + print(" not iterate there -- the ladder effectively starts from one " + "determinant.", flush=True) + history.append({ + "rung": rung, "cutoff": cutoff, "round": round_index, + "dimension": dim, "energy": energy, + "delta_energy": delta, "seconds": elapsed, + }) + + grown = result["carryover_det"] + det = reshard(grown, args, comm, rank, size) + new_dim = global_dim(det) + if rank == 0 and new_dim < dim: + print(f" NOTE: the expansion returned fewer determinants " + f"({new_dim} < {dim}). With --heatbath_truncation > 0 the " + f"parents are pruned before expanding.") + + # Energy-targeted advance, which is the standard HCI protocol: a rung is + # done when another round at the same cutoff no longer moves the answer. + rung_delta = None if rung_baseline is None else energy - rung_baseline + rung_baseline = energy + previous_energy = energy + if rung_delta is not None and abs(rung_delta) < args.energy_tol: + rung_converged = True + break + # A rung is equally done if the subspace stopped growing: the expansion + # has reached its fixed point for this cutoff. + if new_dim == dim: + rung_converged = True + if rank == 0: + print(f" rung {rung} at cutoff {cutoff:g}: subspace stopped " + f"growing at {dim} -- its fixed point for this cutoff") + break + + if rank == 0 and not rung_converged: + print(f" rung {rung} at cutoff {cutoff:g}: hit --max_rounds " + f"{args.max_rounds} without converging") + + return _finish(args, history, stop_reason, rank) + + +def _finish(args, history, stop_reason, rank): + """Print the summary and optionally write the series.""" + if rank != 0: + return 0 + if not history: + print("no rounds completed") + return 1 + first, last = history[0], history[-1] + print() + print("=" * 78) + print(f" {stop_reason}") + print(f" Start: dim {first['dimension']}, E = {first['energy']:.10f}") + print(f" End: dim {last['dimension']}, E = {last['energy']:.10f}") + print(f" Gain: {last['energy'] - first['energy']:+.6f} Hartree over " + f"{len(history)} round(s), " + f"{sum(h['seconds'] for h in history):.1f} s total") + print(" The energy is variational, so it must not rise: a positive gain means " + "the subspace shrank or a round failed to converge.") + if args.log: + with open(args.log, "w", encoding="utf-8") as handle: + json.dump({"stop_reason": stop_reason, "rounds": history}, handle, + indent=2) + print(f" Series written to {args.log}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/tpb/README.md b/examples/tpb/README.md index 3efe7d3..1794ea7 100644 --- a/examples/tpb/README.md +++ b/examples/tpb/README.md @@ -408,4 +408,6 @@ phases, not a fault. ## See Also +- [`../gdb/README.md`](../gdb/README.md) — the general determinant basis, which + decomposes over MPI differently - [Repository README](../../README.md) — Installation, API reference diff --git a/python/__init__.py b/python/__init__.py index d74a2fe..75dfd92 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -535,6 +535,28 @@ def from_string(s, bit_length, total_bit_length, device=None): return get_backend(device).from_string(s, bit_length, total_bit_length) +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 + 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. + """ + _ensure_initialized() + return get_backend(device).from_strings(list(strings), bit_length, total_bit_length) + + +def sort_bitarray_array(dets, device=None): + """Sort packed determinants into canonical order, array in and array out. + + The array form of :func:`sort_bitarray`, for a determinant list held as an + ``(n, words)`` array. Deduplicates, like the list form. + """ + _ensure_initialized() + return get_backend(device).sort_bitarray_array(dets) + + def sort_bitarray(dets, device=None): """Sort determinants into canonical order, removing duplicates. @@ -596,7 +618,9 @@ def tpb_diag(fcidump, adet, bdet, sbd_data, def gdb_diag(fcidump, det, sbd_data, - loadname="", savename="", device=None): + loadname="", savename="", device=None, + determinant_distribution="", determinant_grid_a=0, + determinant_grid_b=0): """ Perform GDB diagonalization over an explicit list of determinants. @@ -605,29 +629,79 @@ def gdb_diag(fcidump, det, sbd_data, themselves, so an arbitrary sparse subspace can be diagonalized. Each determinant is a ``2 * norb``-bit configuration packed into words of - ``sbd_data.bit_length`` bits, as returned by :func:`from_string`. Bit ``2 * i`` - is the occupation of spin-alpha orbital ``i`` and bit ``2 * i + 1`` that of - spin-beta orbital ``i``. + ``sbd_data.bit_length`` bits. Bit ``2 * i`` is the occupation of spin-alpha + orbital ``i`` and bit ``2 * i + 1`` that of spin-beta orbital ``i``. + :func:`from_strings` packs a list of bitstrings into the expected + ``(ndets, words)`` array; a nested list is accepted and converted. + + **The shard contract.** ``sbd_data.b_comm_size`` decides what ``det`` means: + + - ``1`` — every rank passes the whole basis. + - ``> 1`` — every rank passes **its own shard**, and the union over b_comm + positions is the basis. This is the only way GDB's memory scales, because + b_comm is the only dimension that divides the basis (and with it the + excitation lookup); the derived helper dimension divides work but not + storage. + + The shard index is ``rank % b_comm_size``, and ranks sharing one must pass + identical determinants — the in-memory path does not broadcast the list. + Shards must be globally sorted and disjoint: shard ``i`` strictly below shard + ``i + 1``. All of this is checked, and a violation raises rather than + silently diagonalizing the wrong subspace. What cannot be checked is + *completeness* — that the union is the basis you meant — so compare the + returned ``global_dim`` against what you expect. + + ``t_comm_size`` must not exceed ``b_comm_size``: GDB runs one task per + basis-ring station and there are exactly ``b_comm_size`` of them. + ``t_comm_size * b_comm_size`` must divide the rank count exactly. On the + Thrust backend the helper dimension must be 1, i.e. every rank goes to + ``t_comm_size * b_comm_size``. Args: fcidump: FCIDump object. - det: Determinants spanning the subspace. Must be distinct; they are - sorted into SBD's canonical order internally, which - :func:`sort_bitarray` reproduces. - sbd_data: GDB_SBD configuration object. ``b_comm_size`` must be 1. - loadname: Path to load initial wavefunction (optional). + det: Determinants spanning the subspace, or this rank's shard of them, as + an ``(ndets, words)`` array. Must be distinct. + sbd_data: GDB_SBD configuration object. ``method`` must be 0 or 1 — GDB + implements Davidson only, so TPB's Lanczos methods 2 and 3 do not + exist here and are rejected. + loadname: Path to load an initial wavefunction (optional). savename: Path prefix to save the final wavefunction to (optional). SBD - writes ``f"{savename}000000.bin"``, holding the determinants in - canonical order and their amplitudes. + 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). + 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 + matvec), ``count``, ``count-sorted``, ``grid-cyclic`` or + ``grid-cyclic-balanced``. Ignored when ``b_comm_size`` is 1. + determinant_grid_a: Grid rows for the grid-cyclic schemes. With + ``determinant_grid_b``, must multiply to ``b_comm_size``; give both or + neither, and the default is the factor pair nearest square. + determinant_grid_b: Grid columns; see ``determinant_grid_a``. Returns: - dict with keys: energy, density, carryover_det, one_p_rdm, two_p_rdm. + dict with keys ``energy``, ``density``, ``carryover_det``, ``one_p_rdm``, + ``two_p_rdm``, ``local_dim``, ``global_dim`` and + ``determinant_distribution``. What is replicated and what is sharded: + + - ``energy`` — replicated, bit-identical on every rank. + - ``density``, ``one_p_rdm``, ``two_p_rdm`` — replicated on every rank. + - ``carryover_det`` — **this rank's shard**, as an ``(n, words)`` array, + not the whole list. For ``carryover_type`` 1 it is split over b_comm and + *duplicated* across the helper dimension, so gathering means taking one + representative per ``rank % b_comm_size``. For types 2 and 3 it is split + over the world communicator with no duplication, so a plain allgather is + correct; those types return the parents together with the new + candidates, i.e. the next subspace rather than only the additions. + - ``local_dim`` / ``global_dim`` — determinants on this rank, and summed + over b_comm. """ _ensure_initialized() backend = get_backend(device) return backend.gdb_diag( - _global_comm, sbd_data, fcidump, det, loadname, savename + _global_comm, sbd_data, fcidump, det, loadname, savename, + determinant_distribution, determinant_grid_a, determinant_grid_b, ) @@ -738,7 +812,9 @@ def print_info(): 'LoadAlphaDets', 'makestring', 'from_string', + 'from_strings', 'sort_bitarray', + 'sort_bitarray_array', 'tpb_diag_from_files', 'tpb_diag', 'gdb_diag', diff --git a/python/bindings.cpp b/python/bindings.cpp index 5aa3c3b..b68f262 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -34,8 +34,14 @@ // repo rather than in the vendored upstream submodule. #include +#include +#include +#include +#include + #include #include +#include #include #include #include @@ -134,6 +140,102 @@ static void sbd_pin_offload_device(int mpi_rank) { #define SBD_MODULE_NAME _core #endif +// --------------------------------------------------------------------------- +// GDB distributed-basis helpers +// +// GDB's matvec is systolic over b_comm: each rank owns the bra rows of its basis +// block, and the ket vector plus its index map rotate around the ring +// (gdb/mult.h:39-41 for each rank's initial offset, :196-202 for the hops). Two +// consequences drive everything below. +// +// 1. b_comm is the ONLY dimension that shards memory. h_comm is a row stride +// within a block (mult.h:70) closed by an allreduce (:208), and MakeHelpers +// ignores it, so every h-rank builds the full excitation lookup. +// 2. The ring has exactly mpi_size_b stations and one "task" is one station, so +// t_comm_size <= b_comm_size is structural. +// +// Upstream checks neither, nor that t*b divides the rank count. The failures are +// a segfault and silently ragged communicators respectively, so they are checked +// here instead. +// --------------------------------------------------------------------------- + +/** GDB determinant placement over b_comm. The six schemes of the upstream app. */ +enum class SbdDetDistribution { + input, equal_bra_a, count, count_sorted, grid_cyclic, grid_cyclic_balanced +}; + +/** + * Resolve a scheme name, accepting underscores for dashes as the app does. + * + * Empty means the default. The app reaches its default through three legacy + * booleans with a priority rule (do_redist_alpha_eq, which defaults true, then + * do_sort_det, then do_redist_det -- main.cc:33-56); those are deliberately not + * exposed here, since they encode a four-way choice as three flags where two are + * unreachable unless the first is explicitly zeroed. The name says it directly. + */ +static SbdDetDistribution sbd_resolve_det_distribution(const std::string& raw) { + std::string name = raw; + std::replace(name.begin(), name.end(), '_', '-'); + if (name.empty() || name == "equal-bra-a") return SbdDetDistribution::equal_bra_a; + if (name == "input") return SbdDetDistribution::input; + if (name == "count") return SbdDetDistribution::count; + if (name == "count-sorted") return SbdDetDistribution::count_sorted; + if (name == "grid-cyclic") return SbdDetDistribution::grid_cyclic; + if (name == "grid-cyclic-balanced") return SbdDetDistribution::grid_cyclic_balanced; + throw std::invalid_argument( + "unknown determinant_distribution '" + raw + "'; expected one of: input, " + "equal-bra-a (default), count, count-sorted, grid-cyclic, grid-cyclic-balanced"); +} + +static const char* sbd_det_distribution_name(SbdDetDistribution d) { + switch (d) { + case SbdDetDistribution::input: return "input"; + case SbdDetDistribution::equal_bra_a: return "equal-bra-a"; + case SbdDetDistribution::count: return "count"; + case SbdDetDistribution::count_sorted: return "count-sorted"; + case SbdDetDistribution::grid_cyclic: return "grid-cyclic"; + case SbdDetDistribution::grid_cyclic_balanced: return "grid-cyclic-balanced"; + } + return "unknown"; +} + +/** + * True on every rank iff local_ok holds on every rank of comm. + * + * Every per-rank check below routes through this. Throwing on one rank only + * would leave the others inside the next collective, so the verdict has to be + * unanimous before anyone raises. + */ +static bool sbd_all_ranks_ok(bool local_ok, MPI_Comm comm) { + int local = local_ok ? 1 : 0; + int all = 0; + MPI_Allreduce(&local, &all, 1, MPI_INT, MPI_LAND, comm); + return all != 0; +} + +/** + * Order-insensitive-after-sort 64-bit fingerprint of a shard (FNV-1a). + * + * Used to compare the replicas that ranks sharing a b_comm position must pass: + * the in-memory diag never broadcasts the determinant list (contrast the + * file-based overload's MpiBcast(det,0,h_comm) at gdb/sbdiag.h:764), so a + * divergence there is silently wrong rather than detected. Taken after the local + * sort, so a caller that supplies the same set in a different order still agrees. + */ +static uint64_t sbd_shard_fingerprint(const sbd::det_vector& det) { + uint64_t h = 1469598103934665603ULL; + const auto& flat = det.cflat(); + const unsigned char* bytes = reinterpret_cast(flat.data()); + const size_t n = flat.size() * sizeof(size_t); + for (size_t i = 0; i < n; ++i) { + h ^= static_cast(bytes[i]); + h *= 1099511628211ULL; + } + // Fold in the row count so an empty shard and a zero-filled one differ. + h ^= static_cast(det.size()) + 0x9e3779b97f4a7c15ULL; + return h; +} + PYBIND11_MODULE(SBD_MODULE_NAME, m) { // Set module docstring based on backend #ifdef SBD_THRUST @@ -245,8 +347,13 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { "Task communicator size") .def_readwrite("b_comm_size", &sbd::gdb::SBD::b_comm_size, "Basis communicator size") + // GDB implements Davidson only -- there is no Lanczos anywhere under + // include/sbd/chemistry/gdb/, so TPB's methods 2 and 3 do not exist + // here. gdb_diag rejects them rather than passing them through: see the + // check in gdb_diag for what upstream does with an out-of-range value. .def_readwrite("method", &sbd::gdb::SBD::method, - "Diagonalization method (0=Davidson, 1=Davidson+Ham, 2=Lanczos, 3=Lanczos+Ham)") + "Diagonalization method (0=Davidson, 1=Davidson+Ham). GDB has " + "no Lanczos; TPB's 2 and 3 are not available") .def_readwrite("max_it", &sbd::gdb::SBD::max_it, "Maximum number of iterations") .def_readwrite("max_nb", &sbd::gdb::SBD::max_nb, @@ -259,8 +366,15 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { "Initialization method") .def_readwrite("seed", &sbd::gdb::SBD::seed, "Seed for the initial vector") - .def_readwrite("do_shuffle", &sbd::gdb::SBD::do_shuffle, - "Shuffle determinants flag") + // do_shuffle is deliberately NOT exposed, for the same reason as + // h_comm_size: GDB declares the field (gdb/sbdiag.h:25), parses it from + // argv (:94) and copies it into a local (:219), then never reads it + // again -- the only other "shuffle" under gdb/ is an unrelated + // warp-shuffle comment in mult_thrust.h. TPB does use it + // (tpb/sbdiag.h:888-900), which is why it stays on TPB_SBD. An + // attribute that silently does nothing is worse than an absent one, + // and determinant placement is controlled by + // gdb_diag's determinant_distribution argument instead. .def_readwrite("do_rdm", &sbd::gdb::SBD::do_rdm, "Calculate RDM flag (0=density only, 1=full RDM)") .def_readwrite("carryover_type", &sbd::gdb::SBD::carryover_type, @@ -318,6 +432,85 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { py::arg("bit_length"), py::arg("total_bit_length")); + // Packing one determinant per call from Python is the bottleneck at scale: + // issue #31 measures it at tens of thousands of strings, and a sampled + // subspace is far larger. Loop in C++ and hand back one array instead. + m.def("from_strings", + [](const std::vector& strings, size_t bit_length, + size_t total_bit_length) { + const size_t words = + (total_bit_length + bit_length - 1) / bit_length; + const std::vector shape{ + static_cast(strings.size()), + static_cast(words)}; + py::array_t out(shape); + size_t* data = out.mutable_data(); + std::memset(data, 0, strings.size() * words * sizeof(size_t)); + for (size_t r = 0; r < strings.size(); ++r) { + const std::string& s = strings[r]; + if (s.size() != total_bit_length) { + throw std::invalid_argument( + "from_strings: expected every bitstring to be " + + std::to_string(total_bit_length) + " characters, got " + + std::to_string(s.size()) + " at index " + + std::to_string(r)); + } + size_t* row = data + r * words; + for (size_t i = 0; i < total_bit_length; ++i) { + if (s[total_bit_length - 1 - i] == '1') { + row[i / bit_length] |= + (static_cast(1) << (i % bit_length)); + } + } + } + return out; + }, + "Pack many bitstrings into one (n, words) array, as from_string does " + "for a single determinant", + py::arg("strings"), + py::arg("bit_length"), + py::arg("total_bit_length")); + + // Array in, array out, so a sharded determinant list can be put in canonical + // order without a round trip through Python lists. Deduplicates locally, as + // sbd::sort_bitarray does. + m.def("sort_bitarray_array", + [](py::array_t arr) + -> py::array_t { + if (arr.ndim() != 2) { + // An empty shard is legal and carries no width. + if (arr.ndim() <= 1 && arr.size() == 0) { + const std::vector empty_shape{ + static_cast(0), + static_cast(0)}; + return py::array_t(empty_shape); + } + throw std::invalid_argument( + "sort_bitarray_array expects a 2-D (ndets, words) array"); + } + const size_t n = static_cast(arr.shape(0)); + const size_t words = static_cast(arr.shape(1)); + std::vector> rows(n, std::vector(words)); + const size_t* in = arr.data(); + for (size_t r = 0; r < n; ++r) { + std::memcpy(rows[r].data(), in + r * words, + words * sizeof(size_t)); + } + sbd::sort_bitarray(rows); + const std::vector shape{ + static_cast(rows.size()), + static_cast(words)}; + py::array_t out(shape); + size_t* data = out.mutable_data(); + for (size_t r = 0; r < rows.size(); ++r) { + std::memcpy(data + r * words, rows[r].data(), + words * sizeof(size_t)); + } + return out; + }, + "Sort packed determinants into canonical order, removing duplicates", + py::arg("dets")); + m.def("sort_bitarray", [](std::vector>& dets) { sbd::sort_bitarray(dets); @@ -417,37 +610,385 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { // ======================================================================== // Main GDB diagonalization function (data structure version) // - // The determinant list is the subspace itself: unlike TPB, the subspace is - // not the Cartesian product of two half-determinant lists, so an arbitrary - // sparse set of determinants can be diagonalized. + // The determinant list IS the subspace: unlike TPB, it is not the Cartesian + // product of two half-determinant lists, so an arbitrary sparse set can be + // diagonalized. + // + // THE SHARD CONTRACT. The list may be distributed: + // + // b_comm_size == 1 -> every rank passes the WHOLE basis. + // b_comm_size > 1 -> every rank passes ITS OWN SHARD, and the union over + // b_comm positions is the basis. + // + // Sharding is what makes GDB's memory scale: b_comm is the only dimension + // that divides the basis (and hence the excitation lookup), so at + // b_comm_size == 1 every rank stores everything no matter how many ranks + // there are. Ranks sharing a b_comm position -- i.e. one h_comm, which is + // world_rank % b_comm_size -- must pass IDENTICAL shards, because the + // in-memory gdb::diag never broadcasts the list the way the file-based + // overload does (gdb/sbdiag.h:764). That is checked, not trusted. + // + // Input shards must be globally sorted and disjoint. This is stricter than + // upstream (grid-cyclic tolerates arbitrary shards), but it is what slicing a + // globally sorted list gives you for free, and it lets a caller's partition + // be checked rather than silently producing a wrong subspace. Completeness -- + // that the union is the basis you MEANT -- cannot be checked here, so the + // global dimension is returned for the caller to assert on. // ======================================================================== m.def("gdb_diag", [](py::object py_comm, const sbd::gdb::SBD& sbd_data, const sbd::FCIDump& fcidump, - const std::vector>& det, + py::array_t det_in, const std::string& loadname, - const std::string& savename) { + const std::string& savename, + const std::string& determinant_distribution, + int determinant_grid_a, + int determinant_grid_b) { - // Convert MPI communicator MPI_Comm comm = get_mpi_comm(py_comm); + int mpi_rank; MPI_Comm_rank(comm, &mpi_rank); + int mpi_size; MPI_Comm_size(comm, &mpi_size); - int mpi_rank; - MPI_Comm_rank(comm, &mpi_rank); + const int b_comm_size = sbd_data.b_comm_size; + const int t_comm_size = sbd_data.t_comm_size; - // Every rank passes the whole determinant list, which is what - // sbd::gdb::diag expects only when the basis is not split over - // b_comm. Splitting it would require distributing the determinants - // over b_comm first, as the file-based entry point does. - if (sbd_data.b_comm_size != 1) { + // ---- grid shape ------------------------------------------------- + // These depend only on the config, which every rank passes + // identically, so they can throw directly without a vote. + if (b_comm_size < 1 || t_comm_size < 1) { + throw std::invalid_argument( + "gdb_diag requires b_comm_size >= 1 and t_comm_size >= 1"); + } + // GDB runs one task per basis-ring station and there are exactly + // b_comm_size of them, so a larger t_comm_size starves a rank and + // upstream's MakeHelpers then dereferences an empty lookup + // (gdb/helper.h:736-739 then :761) -- a segfault, not an error. + if (t_comm_size > b_comm_size) { + throw std::invalid_argument( + "gdb_diag requires t_comm_size <= b_comm_size: GDB runs one task " + "per basis-ring station and there are exactly b_comm_size of them, " + "so a larger t_comm_size starves a rank and upstream's MakeHelpers " + "then dereferences an empty lookup (segfault, not an error)"); + } + // gdb::diag dispatches with `if (method == 0) {...} else if (method == 1) + // {...}` and no else, assigning `energy` only inside those branches + // (gdb/sbdiag.h:418, :501). A method of 2 or 3 -- valid for TPB, where + // they select Lanczos -- runs no diagonalization at all and leaves + // `energy` uninitialized. The Thrust build masks the value first + // (`method &= 1`, sbdiag.h:210-211), so it is the CPU and OMP-offload + // backends that would return garbage. Reject it on every backend. + if (sbd_data.method != 0 && sbd_data.method != 1) { throw std::invalid_argument( - "gdb_diag requires b_comm_size == 1; distribute work over " - "t_comm_size, and over the derived helper dimension by " - "changing the rank count (h_comm_size is not settable)"); + "gdb_diag requires method 0 or 1 (Davidson, or Davidson storing " + "the Hamiltonian); GDB has no Lanczos, so TPB's methods 2 and 3 " + "do not exist here and would return an uninitialized energy"); } - if (det.empty()) { - throw std::invalid_argument("gdb_diag requires at least one determinant"); + const long long named_grid = + static_cast(b_comm_size) * static_cast(t_comm_size); + // h_comm_size is derived by integer division upstream with no check, + // so a rank count that is not a multiple of b*t silently yields + // ragged communicators: unequal h_comm sizes in one run, and a rank + // alone in its own b_comm believing the ring has one station. + if (named_grid > mpi_size || mpi_size % named_grid != 0) { + throw std::invalid_argument( + "gdb_diag requires t_comm_size * b_comm_size to divide the rank " + "count exactly (the helper dimension is the quotient); got " + + std::to_string(t_comm_size) + " * " + std::to_string(b_comm_size) + + " against " + std::to_string(mpi_size) + " ranks"); + } + const int h_comm_size = static_cast(mpi_size / named_grid); +#ifdef SBD_THRUST + // gdb/mult_thrust.h:310-314 throws on every kernel launch unless the + // helper dimension is trivial. Say so before any work happens. + if (h_comm_size != 1) { + throw std::invalid_argument( + "GDB on the Thrust backend requires h_comm_size == 1, i.e. " + "t_comm_size * b_comm_size == ranks; got a helper dimension of " + + std::to_string(h_comm_size) + ". Spend every rank on " + "b_comm_size (and t_comm_size <= b_comm_size) instead"); + } +#endif + + // ---- placement scheme ------------------------------------------- + const SbdDetDistribution distribution = + sbd_resolve_det_distribution(determinant_distribution); + const bool uses_grid = + distribution == SbdDetDistribution::grid_cyclic || + distribution == SbdDetDistribution::grid_cyclic_balanced; + int grid_a = determinant_grid_a; + int grid_b = determinant_grid_b; + if (uses_grid) { + if (grid_a < 0 || grid_b < 0) { + throw std::invalid_argument( + "determinant grid dimensions must be positive"); + } + if ((grid_a == 0) != (grid_b == 0)) { + throw std::invalid_argument( + "specify both determinant grid dimensions or neither"); + } + if (grid_a == 0) { + // The factor pair of b_comm_size nearest square, as main.cc does. + grid_a = static_cast(std::sqrt(static_cast(b_comm_size))); + while (grid_a > 1 && b_comm_size % grid_a != 0) --grid_a; + if (grid_a < 1) grid_a = 1; + grid_b = b_comm_size / grid_a; + } + if (static_cast(grid_a) * grid_b != b_comm_size) { + throw std::invalid_argument( + "determinant grid dimensions must multiply to b_comm_size (" + + std::to_string(b_comm_size) + ")"); + } + } else if (grid_a != 0 || grid_b != 0) { + throw std::invalid_argument( + "determinant grid dimensions require a grid-cyclic distribution"); + } + + // ---- orbital count and packing width ---------------------------- + size_t norb = 0; + for (const auto& kv : fcidump.header) { + if (kv.first == std::string("NORB")) { + norb = static_cast(std::atoi(kv.second.c_str())); + } + } + if (norb == 0) { + throw std::invalid_argument( + "gdb_diag could not read NORB from the FCIDUMP header"); + } + const size_t total_bits = 2 * norb; + const size_t bit_length = sbd_data.bit_length; + if (bit_length == 0) { + throw std::invalid_argument("gdb_diag requires bit_length >= 1"); + } + const size_t expected_words = (total_bits + bit_length - 1) / bit_length; + + // ---- the determinant buffer ------------------------------------- + // An empty shard is legal (a rank may own nothing), and numpy gives a + // 1-D zero-length array for an empty list, which carries no width -- + // hence the agreement step below rather than trusting shape(1). + const py::ssize_t ndim = det_in.ndim(); + size_t n_local = 0; + size_t local_words = 0; + bool shape_ok = true; + if (ndim == 2) { + n_local = static_cast(det_in.shape(0)); + local_words = static_cast(det_in.shape(1)); + if (n_local > 0 && local_words == 0) shape_ok = false; + } else if (ndim <= 1 && det_in.size() == 0) { + n_local = 0; + local_words = 0; + } else { + shape_ok = false; + } + if (!sbd_all_ranks_ok(shape_ok, comm)) { + throw std::invalid_argument( + "gdb_diag expects det as a 2-D (ndets, words) array of packed " + "determinants, as from_strings() returns; an empty shard may be " + "an empty array"); + } + + unsigned long long words_local = static_cast(local_words); + unsigned long long words_max = 0; + MPI_Allreduce(&words_local, &words_max, 1, MPI_UNSIGNED_LONG_LONG, + MPI_MAX, comm); + if (words_max == 0) { + throw std::invalid_argument( + "gdb_diag requires at least one determinant somewhere on the " + "communicator"); + } + const size_t words = static_cast(words_max); + if (!sbd_all_ranks_ok(local_words == 0 || local_words == words, comm)) { + throw std::invalid_argument( + "gdb_diag requires every rank's determinants to have the same " + "number of packed words"); + } + if (words != expected_words) { + throw std::invalid_argument( + "gdb_diag got determinants packed into " + std::to_string(words) + + " word(s), but NORB=" + std::to_string(norb) + " with bit_length=" + + std::to_string(bit_length) + " implies " + + std::to_string(expected_words) + + "; pack with the same bit_length the config carries"); + } + + // det_vector's row width is a process-global property fixed by the + // first container built, so set it while the GIL is held and report a + // mismatch before any work is done. The half-determinant width is set + // too, matching the upstream app: the grid-cyclic path builds + // half-determinant containers before any diagonalization runs. + try { + sbd::det_vector::init_elem_size(words); + sbd::det_vector::init_elem_size( + (norb + bit_length - 1) / bit_length); + } catch (const std::length_error&) { + throw std::invalid_argument( + "gdb_diag was already called in this process with a different " + "number of words per determinant, which SBD fixes for the " + "lifetime of the process. Keep norb and bit_length fixed, or " + "run the new problem in a fresh process."); + } + + sbd::det_vector det; + det.resize(n_local); + if (n_local > 0) { + std::memcpy(det.flat().data(), det_in.data(), + n_local * words * sizeof(size_t)); + } + + // SBD indexes the subspace with binary searches, so each shard must be + // in canonical order; an unsorted list silently yields a wrong energy. + // sort_bitarray also removes duplicates, which would leave part of the + // subspace unreachable, so reject a shrink rather than dropping them. + const size_t n_before = det.size(); + sbd::sort_bitarray(det); + if (!sbd_all_ranks_ok(det.size() == n_before, comm)) { + throw std::invalid_argument( + "gdb_diag requires distinct determinants"); + } + + // ---- communicators ---------------------------------------------- + // Built here, before diag builds its own set, because the placement + // schemes redistribute over b_comm. With both named dimensions + // trivial, h_comm is the whole communicator, so skip the split. + const bool split_comms = (b_comm_size > 1 || t_comm_size > 1); + MPI_Comm h_comm = comm, b_comm = MPI_COMM_NULL, t_comm = MPI_COMM_NULL; + if (split_comms) { + sbd::gdb::DetBasisCommunicator(comm, h_comm_size, b_comm_size, + t_comm_size, h_comm, b_comm, t_comm); + } + struct CommGuard { + bool owns; MPI_Comm *h, *b, *t; + ~CommGuard() { + if (!owns) return; + if (*h != MPI_COMM_NULL) MPI_Comm_free(h); + if (*b != MPI_COMM_NULL) MPI_Comm_free(b); + if (*t != MPI_COMM_NULL) MPI_Comm_free(t); + } + } comm_guard{split_comms, &h_comm, &b_comm, &t_comm}; + + // ---- shard agreement across h_comm ------------------------------ + { + int h_size = 1; + MPI_Comm_size(h_comm, &h_size); + if (h_size > 1) { + unsigned long long fp = sbd_shard_fingerprint(det); + unsigned long long lo = 0, hi = 0; + MPI_Allreduce(&fp, &lo, 1, MPI_UNSIGNED_LONG_LONG, MPI_MIN, h_comm); + MPI_Allreduce(&fp, &hi, 1, MPI_UNSIGNED_LONG_LONG, MPI_MAX, h_comm); + if (!sbd_all_ranks_ok(lo == hi, comm)) { + throw std::invalid_argument( + "gdb_diag requires ranks sharing a b_comm position to pass " + "identical determinants: the shard index is " + "world_rank % b_comm_size, and the in-memory path does not " + "broadcast the list, so a divergence would silently " + "diagonalize different subspaces"); + } + } + } + + // ---- shards are globally sorted and disjoint --------------------- + if (b_comm_size > 1) { + int rank_b = 0, size_b = 0; + MPI_Comm_rank(b_comm, &rank_b); + MPI_Comm_size(b_comm, &size_b); + std::vector my_last(words, 0), prev_last(words, 0); + const bool have_mine = det.size() > 0; + if (have_mine) { + std::memcpy(my_last.data(), &det.cflat()[(det.size() - 1) * words], + words * sizeof(size_t)); + } + int send_n = have_mine ? 1 : 0; + int recv_n = 0; + const int dst = (rank_b + 1 < size_b) ? rank_b + 1 : MPI_PROC_NULL; + const int srcr = (rank_b > 0) ? rank_b - 1 : MPI_PROC_NULL; + MPI_Sendrecv(&send_n, 1, MPI_INT, dst, 91, + &recv_n, 1, MPI_INT, srcr, 91, b_comm, MPI_STATUS_IGNORE); + MPI_Sendrecv(my_last.data(), static_cast(words), + MPI_UNSIGNED_LONG, dst, 92, + prev_last.data(), static_cast(words), + MPI_UNSIGNED_LONG, srcr, 92, b_comm, MPI_STATUS_IGNORE); + // Only comparable when both sides hold something; an empty shard + // in between is skipped rather than treated as a failure. + bool ordered = true; + if (recv_n == 1 && have_mine) { + std::vector my_first(words, 0); + std::memcpy(my_first.data(), det.cflat().data(), + words * sizeof(size_t)); + const bool duplicate = + std::memcmp(my_first.data(), prev_last.data(), + words * sizeof(size_t)) == 0; + ordered = !duplicate && !sbd::less_from_back(my_first, prev_last); + } + if (!sbd_all_ranks_ok(ordered, comm)) { + throw std::invalid_argument( + "gdb_diag requires the input shards to be globally sorted and " + "disjoint: shard i must hold a strictly lower range than shard " + "i+1, where the shard index is world_rank % b_comm_size. Slice " + "a globally sorted determinant list to get this"); + } + } + + // ---- schemes that cannot take an empty shard -------------------- + // redistribution/reordering index config[0] unconditionally + // (framework/bit_manipulation.h:566, :577) and SaveWavefunction does + // basis[0].size() (caop/basic/restart.h:35), all UB when a rank owns + // nothing. redistribution_bitarray and the grid-cyclic path are safe. + const bool needs_nonempty = + distribution == SbdDetDistribution::count || + distribution == SbdDetDistribution::count_sorted || + !savename.empty(); + if (needs_nonempty && !sbd_all_ranks_ok(det.size() > 0, comm)) { + throw std::invalid_argument( + std::string("gdb_diag with determinant_distribution='") + + sbd_det_distribution_name(distribution) + + "'" + (savename.empty() ? "" : " or a savename") + + " requires a non-empty shard on every rank: upstream indexes the " + "first determinant unconditionally on those paths"); + } + + // ---- redistribute over b_comm ----------------------------------- + // Nothing to place with a single basis block, and equal-bra-a there + // would still pay an allgather of every alpha string, so skip it. + if (b_comm_size > 1 && distribution != SbdDetDistribution::input) { + py::gil_scoped_release release; + switch (distribution) { + case SbdDetDistribution::equal_bra_a: + sbd::redistribution_equal_bra_a(det, bit_length, total_bits, + b_comm); + break; + case SbdDetDistribution::count: + sbd::redistribution(det, bit_length, total_bits, b_comm); + break; + case SbdDetDistribution::count_sorted: + sbd::redistribution(det, bit_length, total_bits, b_comm); + sbd::reordering(det, bit_length, total_bits, b_comm); + break; + case SbdDetDistribution::grid_cyclic: + case SbdDetDistribution::grid_cyclic_balanced: + sbd::gdb::redistribution_grid_bra_ab_cyclic( + det, bit_length, total_bits, + static_cast(grid_a), static_cast(grid_b), + b_comm); + if (distribution == SbdDetDistribution::grid_cyclic_balanced) { + sbd::redistribution_bitarray(det, b_comm); + } + // Mandatory: the grid-cyclic output is locally unsorted, + // and MakeHelpers requires each shard in canonical order. + sbd::sort_bitarray(det); + break; + case SbdDetDistribution::input: + break; + } + } + + // ---- global dimension ------------------------------------------- + // Summed over b_comm, since h_comm and t_comm hold replicas. + unsigned long long local_dim = static_cast(det.size()); + unsigned long long global_dim = local_dim; + if (b_comm_size > 1) { + MPI_Allreduce(&local_dim, &global_dim, 1, MPI_UNSIGNED_LONG_LONG, + MPI_SUM, b_comm); } #ifdef SBD_THRUST @@ -468,76 +1009,80 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { sbd_pin_offload_device(mpi_rank); #endif - // Output variables - double energy; + double energy = 0.0; std::vector density; sbd::det_vector co_det; std::vector> one_p_rdm; std::vector> two_p_rdm; - // det_vector packs each determinant into a fixed number of words, - // which is a process-global property of the type: it is fixed by the - // first determinant container built in the process and cannot be - // changed afterwards. Set it here, while the GIL is still held, so - // that a mismatch is reported before any work is done. - try { - sbd::det_vector::init_elem_size(det[0].size()); - } catch (const std::length_error&) { - throw std::invalid_argument( - "gdb_diag was already called in this process with a different " - "number of words per determinant, which SBD fixes for the " - "lifetime of the process. Keep norb and bit_length fixed, or " - "run the new problem in a fresh process."); - } - - // SBD indexes the subspace with binary searches, so the determinants - // must be in canonical order; an unsorted list silently yields a - // wrong energy. Sort here rather than asking the caller to; the - // exposed sort_bitarray reproduces this order for a caller that - // needs it. sort_bitarray also removes duplicates, which would leave - // part of the subspace unreachable, so reject them instead of - // dropping them silently. - std::vector> det_sorted(det); - sbd::sort_bitarray(det_sorted); - if (det_sorted.size() != det.size()) { - throw std::invalid_argument( - "gdb_diag requires distinct determinants"); - } - - std::vector> co_det_vvs; - { - // Release GIL for long computation py::gil_scoped_release release; - sbd::det_vector det_packed(det_sorted.begin(), det_sorted.end()); - - sbd::gdb::diag(comm, sbd_data, fcidump, det_packed, + sbd::gdb::diag(comm, sbd_data, fcidump, det, loadname, savename, energy, density, co_det, one_p_rdm, two_p_rdm); - for (const auto& row : co_det) { - co_det_vvs.emplace_back(row.begin(), row.end()); + // With do_rdm == 0 the occupation density is computed only where + // mpi_rank_t == 0 (gdb/sbdiag.h:525) and left untouched + // elsewhere, so ranks off that plane would see an empty list. + // Fill them in so the Python contract does not depend on t. + if (t_comm_size > 1 && t_comm != MPI_COMM_NULL) { + unsigned long long n_den = density.size(), n_den_max = 0; + MPI_Allreduce(&n_den, &n_den_max, 1, MPI_UNSIGNED_LONG_LONG, + MPI_MAX, t_comm); + if (n_den_max > 0) { + density.resize(static_cast(n_den_max), 0.0); + MPI_Bcast(density.data(), static_cast(n_den_max), + MPI_DOUBLE, 0, t_comm); + } } } - // Return results as dictionary + // Carryover comes back as a shard, not a gathered list: for + // carryover_type 1 it is split over b_comm AND duplicated across + // h_comm (gdb/carryover.h), and for 2/3 it is split over the world + // communicator with no duplication (gdb/expansion.h:821-822). Handing + // back the raw shard keeps the memory win; the caller gathers if it + // wants the whole list, taking one representative per b_comm position + // for type 1. + const size_t n_co = co_det.size(); + // Explicit shape vector: a braced list here is ambiguous under g++ + // (array_t's ShapeContainer constructor versus its copy/move), which + // nvc++ and clang both accepted -- so the Linux CPU build broke while + // the macOS and Thrust builds passed. + const std::vector co_shape{ + static_cast(n_co), + static_cast(words)}; + py::array_t co_out(co_shape); + if (n_co > 0) { + std::memcpy(co_out.mutable_data(), co_det.cflat().data(), + n_co * words * sizeof(size_t)); + } + py::dict results; results["energy"] = energy; results["density"] = density; - results["carryover_det"] = co_det_vvs; + results["carryover_det"] = co_out; results["one_p_rdm"] = one_p_rdm; results["two_p_rdm"] = two_p_rdm; + results["local_dim"] = local_dim; + results["global_dim"] = global_dim; + results["determinant_distribution"] = + std::string(sbd_det_distribution_name(distribution)); return results; }, - "Perform GDB diagonalization over an explicit list of determinants", + "Perform GDB diagonalization over an explicit list of determinants, which " + "may be sharded across b_comm", py::arg("comm"), py::arg("sbd_data"), py::arg("fcidump"), py::arg("det"), py::arg("loadname") = "", - py::arg("savename") = ""); + py::arg("savename") = "", + py::arg("determinant_distribution") = "", + py::arg("determinant_grid_a") = 0, + py::arg("determinant_grid_b") = 0); // ======================================================================== // Main TPB diagonalization function (file-based version) diff --git a/test/conftest.py b/test/conftest.py index 49040d2..ad666dc 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -100,7 +100,16 @@ def backend(): """The SBD backend module to test, for tests calling the extension directly.""" import sbd - return sbd.get_backend(_requested_device()) + device = _requested_device() + # Pin the package default to the same backend this fixture hands out, or the two + # disagree on an install with more than one backend built. get_backend(None) here + # auto-resolves and PREFERS Thrust, while the module-level entry points reach + # _ensure_initialized() whose default is 'cpu' -- so a test would build its + # FCIDump and config from one module and call tpb_diag/gdb_diag in another, + # failing with "incompatible function arguments". Invisible on a CPU-only build, + # which is why it went unnoticed. + sbd.init(device=device or sbd.get_device()) + return sbd.get_backend(device) @pytest.fixture(scope="session") diff --git a/test/test_gdb_drivers.py b/test/test_gdb_drivers.py new file mode 100644 index 0000000..9eec318 --- /dev/null +++ b/test/test_gdb_drivers.py @@ -0,0 +1,336 @@ +# 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. + +"""Check the GDB example drivers end to end, as a user invokes them. + +``test_gdb_equivalence.py`` covers the binding: it calls ``sbd.gdb_diag`` +directly and proves the solver right. Nothing covered the layer a user actually +touches -- the argument parsing, the determinant text reading, the sharding +arithmetic, the default file paths -- so a driver could break while every +library test stayed green. + +These run the scripts in a **subprocess**, the same device the Fe4S4 case in +``test_gdb_equivalence.py`` uses and for the same reason: ``det_vector``'s row +width is fixed process-wide on first use, so h2o (one word) and Fe4S4 (two) +cannot share a process. Invoking the driver as a subprocess also means the test +exercises ``__main__``, argument defaults included, rather than importing a +function and bypassing them. + +The energies asserted here are not new reference values. The small h2o case +reproduces the same -76.0588897208 that ``test_gdb_equivalence.py`` anchors +against TPB, which ties the driver to the solver: if they ever disagree, the +driver's reading or sharding is at fault, not the diagonalization. +""" + +from __future__ import annotations + +import json +import pathlib +import re +import subprocess +import sys + +import pytest + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +DRIVER_DIR = REPO_ROOT / "examples" / "gdb" +H2O_DIR = REPO_ROOT / "vendor" / "sbd-upstream" / "data" / "h2o" +FE4S4_DIR = ( + REPO_ROOT + / "vendor" + / "sbd-upstream" + / "apps" + / "chemistry_gdb_selected_basis_diagonalization" +) + +# The 24x24 product of the h2o alpha list, interleaved into 576 full +# determinants. Asserted through the binding in test_gdb_equivalence.py; the +# driver must land on the same value from the same input files. +H2O_576_ENERGY = -76.0588897208 + +# Upstream's shipped GDB subspace: four files of 14,884 determinants over 36 +# orbitals. This is what both drivers use when neither --fcidump nor --detfiles +# is given, so a test of the defaults is also a test that those paths resolve. +FE4S4_DIM = 59536 +FE4S4_ENERGY = -326.6982518821 + +# Small enough to keep the guardrail cases near-instant: 12 alphas is a +# 144-determinant subspace, and these never reach a diagonalization anyway. +TINY_ALPHA_LIMIT = 12 + + +def _h2o_args(alpha_limit: int) -> list[str]: + """Point a driver at the bundled h2o data as an |A|^2 product subspace.""" + return [ + "--fcidump", str(H2O_DIR / "fcidump.txt"), + "--from-alpha", str(H2O_DIR / "h2o-1em3-alpha.txt"), + "--alpha-limit", str(alpha_limit), + ] + + +def _run(script: str, args: list[str], timeout: int = 900): + """Run a driver from its own directory, the way the README shows it.""" + if not (DRIVER_DIR / script).is_file(): + pytest.skip(f"driver not found: {DRIVER_DIR / script}") + return subprocess.run( + [sys.executable, script, *args], + cwd=DRIVER_DIR, + capture_output=True, + text=True, + timeout=timeout, + ) + + +def _require(path: pathlib.Path) -> None: + if not path.exists(): + pytest.skip(f"vendored upstream data not found at {path} (submodule not checked out?)") + + +def _assert_ok(completed) -> str: + assert completed.returncode == 0, ( + "driver exited non-zero:\n" + f"--- stdout ---\n{completed.stdout[-3000:]}\n" + f"--- stderr ---\n{completed.stderr[-3000:]}" + ) + return completed.stdout + + +def _energy(stdout: str) -> float: + match = re.search(r"Ground state energy:\s*(-?\d+\.\d+)", stdout) + assert match, f"no energy line in driver output:\n{stdout[-3000:]}" + return float(match.group(1)) + + +def _dimension(stdout: str) -> int: + match = re.search(r"Subspace dimension:\s*(\d+)", stdout) + assert match, f"no dimension line in driver output:\n{stdout[-3000:]}" + return int(match.group(1)) + + +# --- run_gdb_diag.py ---------------------------------------------------------- + + +def test_diag_reproduces_the_library_anchor(): + """The driver on h2o must equal what the binding gives on the same subspace.""" + _require(H2O_DIR / "h2o-1em3-alpha.txt") + stdout = _assert_ok(_run("run_gdb_diag.py", _h2o_args(24))) + assert _dimension(stdout) == 576 + assert _energy(stdout) == pytest.approx(H2O_576_ENERGY, abs=1e-8) + + +def test_diag_reports_a_sane_electron_count(): + """The occupation density must sum to the electron count, not merely exist. + + A subspace built with the wrong bit order still diagonalizes and still + returns a plausible-looking energy; the electron count is what catches it, + which is why the driver prints the sum rather than the vector alone. + """ + _require(H2O_DIR / "h2o-1em3-alpha.txt") + stdout = _assert_ok(_run("run_gdb_diag.py", _h2o_args(24))) + match = re.search(r"sums to (\d+\.\d+); should equal the electron count (\d+)", stdout) + assert match, f"no electron-count check in output:\n{stdout[-3000:]}" + assert float(match.group(1)) == pytest.approx(float(match.group(2)), abs=1e-6) + + +def test_diag_carryover_returns_parents_with_the_candidates(): + """``carryover_type 2`` must come back larger than the subspace it expanded. + + The expansion starts from ``edet = det`` (gdb/expansion.h:543), so the + returned list is the next subspace rather than only the additions -- the + property the heatbath driver's loop depends on. + """ + _require(H2O_DIR / "h2o-1em3-alpha.txt") + stdout = _assert_ok(_run("run_gdb_diag.py", [ + *_h2o_args(TINY_ALPHA_LIMIT), "--carryover_type", "2", + "--heatbath_cutoff", "1e-3", + ])) + match = re.search(r"Carryover determinants on this rank:\s*(\d+)", stdout) + assert match, f"no carryover line in output:\n{stdout[-3000:]}" + assert int(match.group(1)) > _dimension(stdout) + + +@pytest.mark.parametrize( + "extra, expected", + [ + (["--t_comm_size", "2", "--b_comm_size", "1"], "t_comm_size <= b_comm_size"), + (["--b_comm_size", "2"], "divide the rank count"), + ], + ids=["t-exceeds-b", "decomposition-does-not-divide-ranks"], +) +def test_diag_rejects_an_impossible_decomposition(extra, expected): + """Both guardrails must fail loudly on one rank rather than segfault. + + ``t > b`` is the dangerous one: upstream does not check it, and MakeHelpers + then dereferences an empty lookup, so without this guard the failure mode is + a segfault rather than a message. + """ + _require(H2O_DIR / "h2o-1em3-alpha.txt") + completed = _run("run_gdb_diag.py", [*_h2o_args(TINY_ALPHA_LIMIT), *extra]) + assert completed.returncode != 0, f"expected a refusal, got:\n{completed.stdout[-2000:]}" + combined = completed.stdout + completed.stderr + assert expected in combined, f"guardrail message missing:\n{combined[-3000:]}" + + +@pytest.mark.slow +def test_diag_default_files_are_the_fe4s4_subspace(): + """With no input flags at all, the driver must run upstream's Fe4S4 data. + + This is the only test of the default paths, and the README documents them as + the case every flagless command runs -- so if the vendored layout moves, this + is what says so. + """ + _require(FE4S4_DIR / "det0.txt") + stdout = _assert_ok(_run("run_gdb_diag.py", [], timeout=1800)) + assert _dimension(stdout) == FE4S4_DIM + assert _energy(stdout) == pytest.approx(FE4S4_ENERGY, abs=1e-8) + + +# --- run_gdb_heatbath.py ------------------------------------------------------ + + +def test_heatbath_ladder_grows_and_lowers_the_energy(tmp_path): + """Two rounds on h2o: the subspace must grow and the energy must not rise. + + The energy is variational in the subspace, and each round's subspace + contains the previous one, so a rise means the expansion dropped parents or + a round failed to converge. Round 0 diagonalizes the seed untouched, so it + must equal the same anchor the diagonalization driver reports. + """ + _require(H2O_DIR / "h2o-1em3-alpha.txt") + log = tmp_path / "ladder.json" + _assert_ok(_run("run_gdb_heatbath.py", [ + "--fcidump", str(H2O_DIR / "fcidump.txt"), + "--subspace-from", "from-alpha", + "--alpha-file", str(H2O_DIR / "h2o-1em3-alpha.txt"), + "--alpha-limit", "24", + "--cutoffs", "1e-3", + "--max_rounds", "2", + "--log", str(log), + ])) + + rounds = json.loads(log.read_text())["rounds"] + assert len(rounds) >= 2, f"expected at least two rounds, got {rounds}" + + assert rounds[0]["dimension"] == 576 + assert rounds[0]["energy"] == pytest.approx(H2O_576_ENERGY, abs=1e-8) + assert rounds[0]["delta_energy"] is None, "the seed round has nothing to compare against" + + dims = [entry["dimension"] for entry in rounds] + energies = [entry["energy"] for entry in rounds] + assert dims == sorted(dims), f"subspace shrank across rounds: {dims}" + assert dims[-1] > dims[0], f"expansion added nothing: {dims}" + for before, after in zip(energies, energies[1:]): + assert after <= before + 1e-10, f"energy rose across a round: {energies}" + + +@pytest.mark.parametrize( + "alpha_flag", ["--from-alpha", "--alpha-file"], ids=["from-alpha", "alpha-file"] +) +def test_both_drivers_accept_either_alpha_flag_spelling(alpha_flag): + """The two drivers grew different names for the same input; both now take both. + + Asserted on the diagonalization driver, where ``--from-alpha`` is the primary + name and ``--alpha-file`` the alias added for parity with the heatbath driver. + """ + _require(H2O_DIR / "h2o-1em3-alpha.txt") + stdout = _assert_ok(_run("run_gdb_diag.py", [ + "--fcidump", str(H2O_DIR / "fcidump.txt"), + alpha_flag, str(H2O_DIR / "h2o-1em3-alpha.txt"), + "--alpha-limit", "24", + ])) + assert _dimension(stdout) == 576 + assert _energy(stdout) == pytest.approx(H2O_576_ENERGY, abs=1e-8) + + +# --- spin-weight validation --------------------------------------------------- + +COUNTS_FILE = REPO_ROOT / "examples" / "tpb" / "count_dict_h2o.json" + + +def _interleave(alpha: str, beta: str) -> str: + """Bit 2*i alpha orbital i, bit 2*i+1 beta orbital i, counting from the right. + + Spelled out here rather than imported from the driver so the test does not + agree with the code under test by construction. + """ + a, b = alpha[::-1], beta[::-1] + return "".join(a[i] + b[i] for i in range(len(a)))[::-1] + + +def _counts_as_files(tmp_path): + """The bundled counts file written both correctly and incorrectly. + + qiskit-addon-sqd emits ``[beta | alpha]`` concatenated; GDB wants the two + interleaved. The "wrong" file is that raw concatenation, which is the mistake + a first conversion actually makes. + """ + counts = json.loads(COUNTS_FILE.read_text()) + norb = len(next(iter(counts))) // 2 + right = sorted({_interleave(k[norb:], k[:norb]) for k in counts}) + wrong = sorted(set(counts)) + good, bad = tmp_path / "right.txt", tmp_path / "wrong.txt" + good.write_text("\n".join(right) + "\n") + bad.write_text("\n".join(wrong) + "\n") + return good, bad + + +def _heatbath_on(strings_file, tmp_path, extra=()): + return _run("run_gdb_heatbath.py", [ + "--fcidump", str(H2O_DIR / "fcidump.txt"), + "--subspace-from", "strings", "--strings-file", str(strings_file), + "--cutoffs", "1e-3", "--max_rounds", "1", + "--log", str(tmp_path / "ladder.json"), *extra, + ]) + + +def test_interleaved_counts_are_accepted(tmp_path): + """The documented conversion of a counts file must run.""" + _require(COUNTS_FILE) + _require(H2O_DIR / "fcidump.txt") + good, _ = _counts_as_files(tmp_path) + _assert_ok(_heatbath_on(good, tmp_path)) + rounds = json.loads((tmp_path / "ladder.json").read_text())["rounds"] + assert rounds[0]["dimension"] == 275 + + +def test_concatenated_counts_are_refused_with_a_diagnosis(tmp_path): + """Feeding [beta | alpha] straight through must be caught before it is solved. + + Without the check this diagonalizes to a plausible-looking energy and only + aborts later inside the heatbath expansion with std::out_of_range, so the + failure gave no hint of its cause. The occupation density cannot catch it + either: permuting bits preserves how many are set. + """ + _require(COUNTS_FILE) + _require(H2O_DIR / "fcidump.txt") + _, bad = _counts_as_files(tmp_path) + completed = _heatbath_on(bad, tmp_path) + assert completed.returncode != 0, "mis-ordered determinants were accepted" + combined = completed.stdout + completed.stderr + assert "do not have 5 alpha and 5 beta electrons" in combined, combined[-2000:] + assert "INTERLEAVED" in combined, "the message should name the likely cause" + + +def test_skip_weight_check_bypasses_the_validation(tmp_path): + """The escape hatch must skip the check, not merely survive it. + + Asserted by the absence of the refusal: the run still fails downstream, since + the determinants really are malformed, but it fails in SBD rather than here. + """ + _require(COUNTS_FILE) + _require(H2O_DIR / "fcidump.txt") + _, bad = _counts_as_files(tmp_path) + completed = _heatbath_on(bad, tmp_path, extra=("--skip-weight-check",)) + combined = completed.stdout + completed.stderr + assert "do not have 5 alpha and 5 beta electrons" not in combined, ( + "--skip-weight-check did not skip the check" + ) diff --git a/test/test_gdb_equivalence.py b/test/test_gdb_equivalence.py new file mode 100644 index 0000000..d270346 --- /dev/null +++ b/test/test_gdb_equivalence.py @@ -0,0 +1,560 @@ +# 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. + +"""Check GDB against TPB on a subspace both can express. + +GDB spans the subspace with the determinants it is handed; TPB spans it with the +Cartesian product of an alpha and a beta determinant list. The product is a +subspace GDB can also express -- interleave every (alpha, beta) pair into one +full determinant -- so the two solvers can be pointed at exactly the same +Hilbert space and their energies must agree. That equality is the anchor here, +and it needs no pinned reference value. + +It also gives GDB a ground-truth number for free. The energies published with +the upstream TPB data are for the full product of the alpha list with itself, so +the full interleave of ``h2o-1em3-alpha.txt`` must reproduce the published +-76.23594663 -- see ``test_reference_energies.py``, which asserts the same value +through TPB. + +Before this file, GDB had no test anywhere: not in the wrapper, and not in +``vendor/sbd-upstream/tests/functionality`` apart from upstream's +``test_gdb_grid_distribution.cc``, which exercises the distribution helpers +rather than a diagonalization. +""" + +from __future__ import annotations + +import itertools +import pathlib +import subprocess +import sys + +import pytest + +import sbd + +# One word per determinant for h2o either way: 24 orbitals is 24 bits of alpha +# (TPB's half determinants) and 48 bits of alpha+beta (GDB's full ones), both +# under 64. ``det_vector::init_elem_size`` fixes the word count process-wide on +# first use, so every case in this file -- and in test_reference_energies.py, +# which shares the process -- must agree on it. +BIT_LENGTH = 64 + +# Enough alpha strings to make a non-trivial subspace, few enough that the +# product stays small: the |A|^2 scaling is the reason this is 24 and not 275. +ALPHA_LIMIT = 24 + +# Published for the full 275 x 275 product of h2o-1em3-alpha.txt. +H2O_1EM3_ENERGY = -76.23594663 + + +def _interleave(alpha: str, beta: str) -> str: + """Interleave two norb-bit strings into one 2*norb-bit GDB determinant. + + Bit ``2 * i`` is alpha orbital ``i`` and bit ``2 * i + 1`` is beta orbital + ``i``, counting bits from the right -- the order the GDB app's README + specifies ("the rightmost bit corresponds to alpha-spin orbital 1, the next + to beta-spin orbital 1") and the order ``from_string`` then packs. + """ + a_rev, b_rev = alpha[::-1], beta[::-1] + return "".join(a_rev[i] + b_rev[i] for i in range(len(a_rev)))[::-1] + + +def _alpha_strings(path, limit=None): + with open(path, encoding="utf-8") as handle: + strings = [line.strip() for line in handle if line.strip()] + return strings[:limit] if limit else strings + + +def _tpb_energy(backend, fcidump, alpha_strings, norb, **overrides): + """Diagonalize the product of ``alpha_strings`` with itself, via TPB.""" + config = backend.TPB_SBD() + config.eps = 1e-10 + config.max_it = 200 + config.bit_length = BIT_LENGTH + for name, value in overrides.items(): + setattr(config, name, value) + half = backend.sort_bitarray( + [backend.from_string(s, BIT_LENGTH, norb) for s in alpha_strings] + ) + return sbd.tpb_diag(fcidump, half, half, config)["energy"] + + +# Placement is a property of the call, not of the solver config, so these are +# kwargs on gdb_diag rather than fields to setattr onto GDB_SBD. +_CALL_KWARGS = ("determinant_distribution", "determinant_grid_a", "determinant_grid_b") + + +def _gdb_result(backend, fcidump, det, norb, **overrides): + """Diagonalize an explicit determinant list (or this rank's shard) via GDB. + + ``det`` may be a list of bitstrings, which is packed here, or an already + packed ``(n, words)`` array. + """ + config = backend.GDB_SBD() + config.eps = 1e-10 + config.max_it = 200 + config.bit_length = BIT_LENGTH + call = {k: overrides.pop(k) for k in list(overrides) if k in _CALL_KWARGS} + for name, value in overrides.items(): + setattr(config, name, value) + # Bitstrings need packing; anything else is already packed, as an array or as + # the nested sequence the binding's forcecast accepts. + if len(det) and isinstance(det[0], str): + det = sbd.from_strings(list(det), BIT_LENGTH, 2 * norb) + return sbd.gdb_diag(fcidump, det, config, **call) + + +def _gdb_energy(backend, fcidump, det, norb, **overrides): + """The energy alone, for the many cases that assert only on it.""" + return _gdb_result(backend, fcidump, det, norb, **overrides)["energy"] + + +def _shard(full, b_comm_size, rank): + """This rank's slice of a globally sorted list. + + Mirrors SBD's own ``q = N/p``, remainder-to-the-low-ranks split + (``balanced_begin``, framework/bit_manipulation.h:1523-1531), and indexes by + ``rank % b_comm_size`` because that is the b_comm position. + """ + if b_comm_size == 1: + return full + index = rank % b_comm_size + quotient, remainder = divmod(len(full), b_comm_size) + begin = index * quotient + min(index, remainder) + end = begin + quotient + (1 if index < remainder else 0) + return full[begin:end] + + +@pytest.fixture(scope="module") +def h2o(data_dir, backend): + """The h2o FCIDUMP, its orbital count, the alpha-list path, and a truncation.""" + molecule_dir = data_dir / "h2o" + fcidump = backend.LoadFCIDump(str(molecule_dir / "fcidump.txt")) + norb = int(fcidump.header["NORB"]) + alpha_path = molecule_dir / "h2o-1em3-alpha.txt" + return fcidump, norb, alpha_path, _alpha_strings(alpha_path, ALPHA_LIMIT) + + +def test_gdb_matches_tpb_on_the_same_subspace(backend, h2o): + """Both solvers give the same energy for the same Hilbert space. + + The strong form of the check: no reference value, so it cannot pass by + coincidence with a wrong Hamiltonian or a misread FCIDUMP. Only the + determinant *representation* differs -- |A| x |A| half determinants against + |A|^2 interleaved full ones. + """ + fcidump, norb, _, alpha = h2o + product = [_interleave(a, b) for a, b in itertools.product(alpha, alpha)] + assert len(product) == len(alpha) ** 2 + + tpb = _tpb_energy(backend, fcidump, alpha, norb) + gdb = _gdb_energy(backend, fcidump, product, norb) + assert gdb == pytest.approx(tpb, abs=1e-9) + + +def test_gdb_rejects_duplicate_determinants(backend, h2o): + """A repeated determinant is an error, not a silently smaller subspace. + + ``sort_bitarray`` deduplicates, so the binding compares sizes afterwards and + raises. Recorded here because TPB takes the opposite branch and drops + duplicates silently -- the divergence tracked in issue #31. If that is ever + reconciled in GDB's favour, this test is what has to change. + """ + fcidump, norb, _, alpha = h2o + doubled = [_interleave(alpha[0], alpha[0])] * 2 + with pytest.raises(ValueError, match="distinct determinants"): + _gdb_energy(backend, fcidump, doubled, norb) + + +def test_gdb_rejects_a_basis_split_the_ranks_cannot_tile(backend, h2o): + """b_comm_size > 1 needs ranks to match it. + + Splitting the basis is supported now, but the grid still has to tile: upstream + derives the helper dimension as ``ranks / (t * b)`` by integer division and + never checks the remainder, so asking for more basis blocks than there are + ranks would silently produce communicators of unequal size. Serially that + means any b_comm_size above 1 is refused. + """ + fcidump, norb, _, alpha = h2o + product = [_interleave(alpha[0], b) for b in alpha] + with pytest.raises(ValueError, match="divide the rank count"): + _gdb_energy(backend, fcidump, product, norb, b_comm_size=2) + + +def test_gdb_rejects_lanczos_methods(backend, h2o): + """Methods 2 and 3 are refused rather than returning an uninitialized energy. + + They are valid for TPB, where they select Lanczos. GDB has no Lanczos: + ``gdb::diag`` handles only ``method == 0`` and ``method == 1`` and assigns + ``energy`` only inside those branches, so passing 2 through would return + whatever was on the stack -- with no diagonalization having run. + """ + fcidump, norb, _, alpha = h2o + product = [_interleave(alpha[0], b) for b in alpha] + with pytest.raises(ValueError, match="method 0 or 1"): + _gdb_energy(backend, fcidump, product, norb, method=2) + + +@pytest.mark.slow +def test_gdb_reproduces_the_published_h2o_energy(backend, h2o): + """The full interleave reproduces the energy published for this alpha list. + + The subspace is 275^2 = 75,625 determinants -- the same Hilbert space + ``test_reference_energies.py`` covers through TPB in its fast tier, but GDB + carries its dimension explicitly rather than as a product, so the determinant + list itself is 75,625 entries. Marked slow for the cost of that rather than + for being intractable: measured at 5.5 s of diagonalization plus 0.2 s to + build the list, on 8 host threads, returning -76.2359466308 against the + published -76.23594663. Promote it to the fast tier if that budget is fine. + """ + fcidump, norb, alpha_path, _ = h2o + # The full list, not the fixture's ALPHA_LIMIT truncation: the published + # energy is for the product of all of it with itself. + alpha = _alpha_strings(alpha_path) + product = [_interleave(a, b) for a, b in itertools.product(alpha, alpha)] + assert len(product) == len(alpha) ** 2 + + energy = _gdb_energy(backend, fcidump, product, norb) + # Quoted to eight decimals upstream, so compared to that rather than to the + # solver's own 1e-10 convergence tolerance. + assert energy == pytest.approx(H2O_1EM3_ENERGY, abs=1e-8) + + +@pytest.mark.mpi +def test_gdb_under_mpi_matches_tpb_on_the_same_subspace(backend, h2o): + """Distributing GDB over ranks does not change the answer. + + GDB decomposes as ``t_comm_size x b_comm_size x helper``. Both named + dimensions are pinned to 1 on the in-memory path -- b_comm_size because every + rank passes the whole determinant list, t_comm_size because GDB enumerates one + task per b_comm rank and so cannot split a single task further (see + ``test_gdb_rejects_a_task_split``). Ranks therefore all land on the derived + helper dimension, and this checks that spending them there is harmless. + + The comparison is against a TPB run of the same subspace in the same process, + which reaches the same Hilbert space through an independent decomposition + (``adet_comm_size``). No pinned reference value is needed, and a broken helper + distribution shows up as disagreement. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + fcidump, norb, _, alpha = h2o + product = [_interleave(a, b) for a, b in itertools.product(alpha, alpha)] + + gdb = _gdb_energy(backend, fcidump, product, norb) + tpb = _tpb_energy(backend, fcidump, alpha, norb, adet_comm_size=size) + + if comm.Get_rank() == 0: + assert gdb == pytest.approx(tpb, abs=1e-9), ( + f"{size} ranks (helper={size}) gave {gdb}, but TPB on the same " + f"subspace gave {tpb}" + ) + + +def test_gdb_rejects_a_task_split(backend, h2o): + """t_comm_size > 1 is refused rather than segfaulting. + + ``t_comm_size <= b_comm_size`` is structural, not incidental. GDB's matvec + rotates the ket around ``b_comm`` as a ring (gdb/mult.h:39-41, :196-202), so + the ring has exactly ``b_comm_size`` stations and one "task" is one station; + ``t_comm`` parallelizes stations, and cannot have more workers than there are + stations. ``MakeHelpers`` encodes this as ``task_end = mpi_size_b`` split + across ``t_comm`` (gdb/helper.h:736-739). + + Upstream does not check it: a starved rank resizes ``exidx`` to zero and then + reads ``exidx[0].slide`` anyway because its ``task_begin`` is nonzero + (helper.h:761), which is a null dereference during helper construction -- + observed as EXC_BAD_ACCESS at 0x0 on 2 and 4 ranks before this guard existed. + With ``b_comm_size`` pinned to 1 in memory, that makes any ``t_comm_size > 1`` + fatal. + + The check runs serially because the guard is a configuration check that + precedes any communication. + """ + fcidump, norb, _, alpha = h2o + product = [_interleave(alpha[0], b) for b in alpha] + with pytest.raises(ValueError, match="t_comm_size <= b_comm_size"): + _gdb_energy(backend, fcidump, product, norb, t_comm_size=2) + + +# --------------------------------------------------------------------------- +# Distributed basis: b_comm_size > 1, where each rank passes only its shard +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def h2o_product(backend, h2o): + """The truncated product basis, packed and globally sorted.""" + fcidump, norb, _, alpha = h2o + strings = [_interleave(a, b) for a, b in itertools.product(alpha, alpha)] + packed = sbd.sort_bitarray_array( + sbd.from_strings(strings, BIT_LENGTH, 2 * norb) + ) + assert packed.shape[0] == len(alpha) ** 2 + return packed + + +def test_from_strings_matches_the_per_determinant_form(backend, h2o): + """The bulk packer agrees elementwise with from_string called in a loop. + + from_strings exists because the per-call form costs one boundary crossing per + determinant, which issue #31 measures as a real cost at tens of thousands of + strings. It is only worth having if it is the same function, hence this. + """ + fcidump, norb, _, alpha = h2o + strings = [_interleave(a, b) for a, b in itertools.product(alpha[:6], alpha[:6])] + bulk = sbd.from_strings(strings, BIT_LENGTH, 2 * norb) + one_by_one = [backend.from_string(s, BIT_LENGTH, 2 * norb) for s in strings] + assert bulk.tolist() == one_by_one + + +def test_gdb_accepts_a_nested_list(backend, h2o, h2o_product): + """A list of lists still works, so existing callers are unaffected. + + The binding takes a forcecast array, which converts a nested sequence, so the + numpy contract is an addition rather than a break. + """ + fcidump, norb, _, _ = h2o + as_list = [list(row) for row in h2o_product] + assert _gdb_energy(backend, fcidump, as_list, norb) == pytest.approx( + _gdb_energy(backend, fcidump, h2o_product, norb), abs=1e-12 + ) + + +def test_gdb_reports_the_dimensions_it_diagonalized(backend, h2o, h2o_product): + """global_dim is what the caller needs to check completeness itself. + + Whether the union of shards is the basis the caller *meant* cannot be checked + inside the binding, so the dimension is reported back instead. + """ + fcidump, norb, _, _ = h2o + result = _gdb_result(backend, fcidump, h2o_product, norb) + assert result["global_dim"] == h2o_product.shape[0] + assert result["local_dim"] == h2o_product.shape[0] + assert result["determinant_distribution"] == "equal-bra-a" + + +@pytest.mark.mpi +def test_gdb_sharded_basis_matches_the_whole_basis(backend, h2o, h2o_product): + """Splitting the basis across b_comm does not change the energy. + + The core claim of the distributed path: each rank passes only its slice, and + the answer still matches TPB on the same subspace. Checked against a TPB run + in this same process, which reaches the same Hilbert space through an + independent decomposition, so no pinned value is involved. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + fcidump, norb, _, alpha = h2o + + gdb = _gdb_energy(backend, fcidump, _shard(h2o_product, size, comm.Get_rank()), + norb, b_comm_size=size) + tpb = _tpb_energy(backend, fcidump, alpha, norb, adet_comm_size=size) + if comm.Get_rank() == 0: + assert gdb == pytest.approx(tpb, abs=1e-9) + + +@pytest.mark.mpi +def test_gdb_task_dimension_works_once_the_ring_has_stations(backend, h2o, h2o_product): + """t_comm_size > 1 becomes usable exactly when b_comm_size allows it. + + ``t <= b`` is structural -- one task per basis-ring station -- so with the + basis in one block the task dimension is unreachable. Shard the basis and it + opens up. Needs at least 4 ranks for t=2, b=2. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + if size < 4 or size % 4: + pytest.skip(f"needs a rank count divisible by 4 for t=2 x b=2, got {size}") + fcidump, norb, _, alpha = h2o + + gdb = _gdb_energy(backend, fcidump, _shard(h2o_product, 2, comm.Get_rank()), + norb, b_comm_size=2, t_comm_size=2) + tpb = _tpb_energy(backend, fcidump, alpha, norb, adet_comm_size=size) + if comm.Get_rank() == 0: + assert gdb == pytest.approx(tpb, abs=1e-9) + + +@pytest.mark.mpi +@pytest.mark.parametrize( + "scheme", + ["input", "equal-bra-a", "count", "count-sorted", + "grid-cyclic", "grid-cyclic-balanced"], +) +def test_gdb_placement_does_not_change_the_energy(backend, h2o, h2o_product, scheme): + """All six placement schemes agree. + + Placement is a load-balancing decision -- which rank owns which determinants, + and how the ring is laid out -- so it must not move the answer. ``input`` keeps + the caller's slices, ``equal-bra-a`` equalizes distinct alpha strings (what + actually balances the matvec, whose outer loop is over alpha), ``count`` and + ``count-sorted`` equalize determinant counts, and the two grid-cyclic schemes + spread alpha and beta keys over a 2-D rank grid. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + fcidump, norb, _, alpha = h2o + + gdb = _gdb_energy(backend, fcidump, _shard(h2o_product, size, comm.Get_rank()), + norb, b_comm_size=size, determinant_distribution=scheme) + tpb = _tpb_energy(backend, fcidump, alpha, norb, adet_comm_size=size) + if comm.Get_rank() == 0: + assert gdb == pytest.approx(tpb, abs=1e-9) + + +@pytest.mark.mpi +def test_gdb_rejects_shards_that_are_not_disjoint(backend, h2o, h2o_product): + """Overlapping shards are refused, not silently diagonalized. + + Every rank passing the whole list while claiming b_comm_size > 1 is the + mistake this catches: each rank would believe its copy was a shard, and the + ring would rotate duplicated blocks. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + if size < 2: + pytest.skip("needs at least 2 ranks to have distinct shards") + fcidump, norb, _, _ = h2o + with pytest.raises(ValueError, match="globally sorted and disjoint"): + _gdb_energy(backend, fcidump, h2o_product, norb, b_comm_size=size) + + +@pytest.mark.mpi +def test_gdb_rejects_a_rank_count_the_grid_does_not_tile(backend, h2o, h2o_product): + """t * b must divide the rank count exactly. + + Upstream derives the helper dimension by integer division and never checks the + remainder, which silently yields communicators of unequal size and a rank alone + in its own basis ring. Refuse it instead. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + fcidump, norb, _, _ = h2o + bad = size + 1 + with pytest.raises(ValueError, match="divide the rank count"): + _gdb_energy(backend, fcidump, _shard(h2o_product, bad, comm.Get_rank()), + norb, b_comm_size=bad) + + +def test_gdb_rejects_an_unknown_placement_scheme(backend, h2o, h2o_product): + """A misspelled scheme names the alternatives rather than falling back.""" + fcidump, norb, _, _ = h2o + with pytest.raises(ValueError, match="unknown determinant_distribution"): + _gdb_energy(backend, fcidump, h2o_product, norb, + determinant_distribution="equal-bra-alpha") + + +def test_gdb_accepts_underscores_in_the_scheme_name(backend, h2o, h2o_product): + """``grid_cyclic`` and ``grid-cyclic`` are the same scheme, as upstream has it.""" + fcidump, norb, _, _ = h2o + result = _gdb_result(backend, fcidump, h2o_product, norb, + determinant_distribution="grid_cyclic") + assert result["determinant_distribution"] == "grid-cyclic" + + +@pytest.mark.parametrize( + "grid_a,grid_b,scheme,message", + [ + (2, 0, "grid-cyclic", "both determinant grid dimensions or neither"), + (3, 3, "grid-cyclic", "multiply to b_comm_size"), + (1, 1, "equal-bra-a", "require a grid-cyclic distribution"), + ], +) +def test_gdb_validates_the_determinant_grid(backend, h2o, h2o_product, + grid_a, grid_b, scheme, message): + """The grid dimensions are checked against b_comm_size before any work. + + ``redistribution_grid_bra_ab_cyclic`` validates the product itself, but against + the communicator it is handed, so checking here names b_comm_size in the error. + """ + fcidump, norb, _, _ = h2o + with pytest.raises(ValueError, match=message): + _gdb_energy(backend, fcidump, h2o_product, norb, + determinant_distribution=scheme, + determinant_grid_a=grid_a, determinant_grid_b=grid_b) + + +# Upstream's only general-determinant data, and the only case here at a realistic +# size. It lives in the app directory rather than under data/, which is alpha-only. +FE4S4_DIR = ( + "vendor/sbd-upstream/apps/chemistry_gdb_selected_basis_diagonalization" +) +# Measured by this wrapper, NOT published upstream: run.sh ships no expected value +# and data/ has reference tables only for the TPB molecules. So this is a +# regression guard against our own verified result, not an independent check. It +# was identical across b_comm_size 1, 2 and 4 and across t_comm_size 1 and 2. +FE4S4_ENERGY = -326.6982518821 + + +@pytest.mark.slow +def test_gdb_fe4s4_at_realistic_size(): + """59,536 determinants over 36 orbitals, upstream's own GDB input. + + Everything else here is h2o at a few hundred to tens of thousands of + determinants; this is the only case with a real determinant list and a real + orbital count, and the only one whose four input files are a ready-made + four-way partition (equal sizes, disjoint, individually sorted, concatenation + globally sorted). Roughly 21 s on 8 host threads. + + Runs in a **subprocess**, because ``det_vector``'s row width is fixed for the + lifetime of a process: h2o packs into one 64-bit word (48 bits) and Fe4S4 needs + two (72 bits), so the two cannot share a process. That is the constraint + ``gdb_diag`` reports rather than a problem with this case, and forking is the + workaround it tells callers to use. + """ + root = pathlib.Path(__file__).resolve().parents[1] / FE4S4_DIR + if not root.is_dir(): + pytest.skip(f"upstream GDB app data not found at {root}") + + script = f""" +import sbd +root = {str(root)!r} +backend = sbd.get_backend() +fcidump = backend.LoadFCIDump(root + "/fcidump_Fe4S4.txt") +norb = int(fcidump.header["NORB"]) +strings = [] +for index in range(4): + with open(root + "/det%d.txt" % index) as handle: + strings.extend(line.strip() for line in handle if line.strip()) +assert norb == 36, norb +assert len(strings) == 59536, len(strings) +config = backend.GDB_SBD() +config.eps, config.max_it, config.bit_length = 1e-6, 30, {BIT_LENGTH} +det = sbd.from_strings(strings, {BIT_LENGTH}, 2 * norb) +result = sbd.gdb_diag(fcidump, det, config) +print("ENERGY", repr(result["energy"]), "DIM", result["global_dim"]) +""" + completed = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, text=True, timeout=1800, + ) + assert completed.returncode == 0, ( + f"subprocess failed:\n{completed.stdout[-2000:]}\n{completed.stderr[-2000:]}" + ) + line = [l for l in completed.stdout.splitlines() if l.startswith("ENERGY")] + assert line, f"no energy reported:\n{completed.stdout[-2000:]}" + _, energy, _, dim = line[-1].split() + assert int(dim) == 59536 + assert float(energy) == pytest.approx(FE4S4_ENERGY, abs=1e-8)