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