diff --git a/README.md b/README.md index e091a44..b2a14fe 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,7 @@ matter for it. SBD can serve as the eigensolver backend for qiskit-addon-sqd's SQD workflow. -**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. +**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. SBD-selected carryover and subspace policies need `0.14.0` or higher. ### Plain SQD @@ -98,9 +98,18 @@ parameter feeds the loop, and complete example. qiskit-addon-sqd is the orchestrator in that recipe: it owns the loop -(sampling, configuration recovery, subsampling), and SBD is plugged in -purely as the per-batch eigensolver (`sci_solver=sbd_solver` above) with no -say in how the subspace grows between iterations. +(sampling, configuration recovery, subsampling), and SBD is plugged in as the +per-batch eigensolver (`sci_solver=sbd_solver` above). By default SBD has no say +in how the subspace grows between iterations: SQD chooses the carryover from the +amplitudes SBD returns. + +With `qiskit-addon-sqd >= 0.14.0`, SBD can choose the carryover itself instead: +set `"carryover_type"` in `sbd_config` and SQD takes SBD's selection as given. With +type 1, the amplitudes never leave SBD, which saves a full wavefunction write and +read every iteration. See +[SBD-selected carryover](examples/tpb/README.md#sbd-selected-carryover) for what +each type does and costs, and the section after it for running the loop with a +different subspace policy such as Trim SQD. ### SQD with subspace enlargement diff --git a/examples/tpb/README.md b/examples/tpb/README.md index 3efe7d3..bb893b5 100644 --- a/examples/tpb/README.md +++ b/examples/tpb/README.md @@ -114,10 +114,12 @@ determinant list, giving a 275 × 275 = 75,625-determinant subspace at MPI decomposition flags, and the SQD tolerances `--energy_tol` / `--occupancies_tol` / `--sqd_carryover_threshold`. Inner-solver flags are prefixed (`--sbd_eps`, `--sbd_max_it`, `--sbd_method`, ...) and have sensible defaults; the -old unprefixed spellings still work. Run `python run_sqd_sbd.py --help` for the -full list, which is grouped by layer. +old unprefixed spellings still work. `--sbd_carryover_type` lets SBD choose what +carries into the next iteration instead of SQD; see +[SBD-selected carryover](#sbd-selected-carryover). Run +`python run_sqd_sbd.py --help` for the full list, which is grouped by layer. -**Requirements:** see [Integration with qiskit-addon-sqd](../../README.md#integration-with-qiskit-addon-sqd) in the Python Bindings README (`pyscf`, `qiskit`, `qiskit-addon-sqd`). +**Requirements:** see [Integration with qiskit-addon-sqd](../../README.md#integration-with-qiskit-addon-sqd) in the Python Bindings README (`pyscf`, `qiskit`, `qiskit-addon-sqd`). `--sbd_carryover_type` needs `qiskit-addon-sqd >= 0.14.0`; everything else works from `0.13.1`. See [SQD Parameters](#sqd-parameters) below for the full reference, grouped by SQD loop / SBD solver / MPI grid / checkpointing. @@ -202,10 +204,13 @@ strs_a = include_a ++ carryover_strings_a ++ samples_a then dedupe, trunc 1. **`include_a`** — configurations passed as `include_configurations`. Static: fixed before the loop, present every iteration, never updated. -2. **`carryover_strings_a`** — from the *previous* iteration's wavefunction. Every - determinant whose `|coefficient|` is at least `--sqd_carryover_threshold` - survives, ranked by `|c|^2`. -3. **`samples_a`** — drawn fresh this iteration, sorted by marginal probability. +2. **`carryover_strings_a`** — from the *previous* iteration's result, in descending + order of marginal weight (a string's `|c|^2` summed over the other spin sector). By + default SQD selects them: every determinant whose `|coefficient|` is at least + `--sqd_carryover_threshold` survives. With `--sbd_carryover_type` set, SBD selects + them instead and SQD takes that selection as given; see + [SBD-selected carryover](#sbd-selected-carryover). +3. **`samples_a`** — drawn fresh this iteration, sorted by how often each was sampled. The order matters when `max_dim` is set: `include` and `carryover` are kept ahead of fresh samples, so if those two already fill the cap, this iteration's new samples are @@ -221,7 +226,9 @@ raw samples are only filtered by electron count (Hamming-weight postselection). So exactly two things flow from iteration N to N+1, and neither is a tolerance: the **average orbital occupancies** (into recovery, source 3) and the -**wavefunction amplitudes** (into carryover, source 2). +**carryover** (source 2). By default the carryover is derived from the wavefunction +amplitudes; with `--sbd_carryover_type 1` SBD hands it over directly and the +amplitudes never leave SBD. ### SQD loop parameters @@ -241,7 +248,7 @@ prints an OOM warning when it detects this. | `--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` | +| `--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.** Has no effect when `--sbd_carryover_type` is non-zero (the driver warns) | `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` | | `--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 | @@ -280,6 +287,91 @@ For reference, upstream's own `TPB_SBD` struct defaults are looser still (`max_i `eps=1e-4`), and `run_sbd_diag.py` uses `eps=1e-3`. On the h2o counts case, `eps=1e-5` and `eps=1e-8` give the same energy to ten decimal places. +### SBD-selected carryover + +By default SQD picks the next iteration's carryover by thresholding the amplitudes +SBD returns. These flags let SBD pick it instead, using its own carryover rules, and +report the result through `SCIResult.carryover` (`qiskit-addon-sqd >= 0.14.0`). SQD +then uses SBD's selection as given: `--sqd_carryover_threshold` stops applying, and the +order SBD's carryover comes in is the order a `--max_dim` truncation keeps. + +| Parameter | What it controls | Default | +|-----------|-----------------|---------| +| `--sbd_carryover_type` | `0`: off, SQD selects. `1`: rank half-determinants by marginal weight. `2`: type 1 plus all single excitations of what it selected. `3`: rank whole determinants by amplitude, then single-excite | `0` | +| `--sbd_carryover_ratio` | Types 1 and 2: fraction of half-determinants to keep, **used only when non-zero**. Ignored by type 3 | `0.1` (upstream SBD's own default is `0.0`) | +| `--sbd_carryover_threshold` | Types 1 and 2: cumulative-weight cutoff, **used only when the ratio is 0**. Type 3: a direct cutoff on determinant weight | `1e-4` | + +SBD reads exactly one of ratio and threshold for each type and silently ignores the +other, so with the defaults, `--sbd_carryover_threshold` does nothing for types 1 and 2. +The driver warns when you pass a setting that cannot take effect. + +**What each type costs.** Type 1 is the one that saves work. Its ranking is exactly the +marginal weight SQD orders carryover by, so nothing else needs the eigenvector, and the +driver skips the `|adet| × |bdet|` wavefunction dump entirely. At large subspace size +that dump is a multi-GB write-then-read every iteration. Types 2 and 3 still write the +dump. They add determinants by single excitation, and SBD re-sorts the result into its +canonical order (which it needs for binary search), so the driver re-ranks their +carryover from the amplitudes: strings that were in the subspace come first, by weight, +and the generated ones, which have no weight yet, come after. + +A checkpoint written with `--sbd_carryover_type` set holds SBD's carryover strings, not +the whole solved subspace, so `--resume_from` seeds from that smaller set. + +### Using a different subspace policy (Trim SQD) + +`qiskit-addon-sqd >= 0.14.0` lets the loop run a different per-iteration schedule, +passed as a `policy`. The driver always uses the default, `StandardPolicy`, and does +not expose others as flags. To try **Trim SQD** (`qiskit_addon_sqd.trim.TrimPolicy`), +pass it from Python, with SBD as the solver as usual: + +```python +from functools import partial + +from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian +from qiskit_addon_sqd.trim import TrimPolicy +from sbd.sbd_solver import solve_sci_batch + +num_batches = 4 +policy = TrimPolicy( + trim_ratio=1 / num_batches, # see the note on trim_ratio below + carryover_threshold=1e-4, # the threshold goes INSIDE the policy +) + +result = diagonalize_fermionic_hamiltonian( + hcore, eri, bit_array, + samples_per_batch=3000, num_batches=num_batches, + norb=norb, nelec=nelec, symmetrize_spin=True, + sci_solver=partial(solve_sci_batch, sbd_config={"eps": 1e-5, "max_it": 10}), + policy=policy, +) +# Pass the threshold to the policy only. Also passing carryover_threshold= here +# raises a ValueError. +``` + +What changes compared to the default: + +- **The batches are disjoint.** One pool of `samples_per_batch × num_batches` + bitstrings is divided among them, instead of each batch being subsampled + independently. +- **Each iteration costs `num_batches + 1` SBD calls**, not `num_batches`. Every + batch is diagonalized, its highest-weight strings are kept, the survivors of all + batches are merged, and the merged subspace is diagonalized once more. That last + diagonalization's energy is what the iteration reports. +- **`trim_ratio` sets how big the merged subspace is.** It holds roughly + `trim_ratio × num_batches` of one batch's strings per spin sector (less where + batches share strings). TrimPolicy's default of `0.1` suits about 10 batches; with + fewer, the merged subspace comes out much smaller than a batch. `1 / num_batches` + keeps it about one batch in size, which is the regime Trim SQD is designed for. + That is our rule of thumb, not a tuned value. The addon's defaults are not tuned + either. +- **With `carryover_type != 0` in `sbd_config`, the policy's carryover settings do + nothing.** The trim schedule still runs, but each batch is screened by SBD's own + selection rather than by `trim_ratio`, so tune `ratio` / `threshold` in `sbd_config` + instead. + +Trim SQD pays off when the samples far outnumber the subspace you can afford to +diagonalize. + ### MPI grid parameters | Parameter | What it controls | Default | @@ -409,3 +501,5 @@ phases, not a fault. ## See Also - [Repository README](../../README.md) — Installation, API reference +- [qiskit-addon-sqd](https://github.com/Qiskit/qiskit-addon-sqd) — the SQD loop, + `SubspacePolicy`, `StandardPolicy` and `TrimPolicy` diff --git a/examples/tpb/run_sqd_sbd.py b/examples/tpb/run_sqd_sbd.py index bf43889..4b37132 100644 --- a/examples/tpb/run_sqd_sbd.py +++ b/examples/tpb/run_sqd_sbd.py @@ -52,6 +52,16 @@ from qiskit.primitives import BitArray from qiskit_addon_sqd.fermion import SCIResult, diagonalize_fermionic_hamiltonian +# Subspace policies arrived in qiskit-addon-sqd 0.14.0 (#369). On an older addon the +# loop's own default schedule is the only one, selected by passing +# carryover_threshold= directly, so keep that path working rather than raising. +# Other policies (e.g. Trim SQD) are not exposed as flags; examples/tpb/README.md +# shows how to pass one from Python. +try: + from qiskit_addon_sqd.fermion import StandardPolicy +except ImportError: # qiskit-addon-sqd < 0.14.0 + StandardPolicy = None + def parse_args(): p = argparse.ArgumentParser( @@ -107,11 +117,14 @@ def parse_args(): "host-side and superlinear in determinants per spin, and " "grows between iterations as carryover accumulates. " "Unset means no cap.") - sqd.add_argument("--sqd_carryover_threshold", type=float, default=1e-4, + sqd.add_argument("--sqd_carryover_threshold", type=float, default=None, + metavar="FLOAT", help="SQD keeps every determinant whose |coefficient| is at " "least this, and carries it into the next iteration's " - "subspace. Lower it to carry more. Distinct from " - "--sbd_carryover_threshold, which SQD does not use.") + "subspace. Lower it to carry more. Default 1e-4. Distinct from " + "--sbd_carryover_threshold, which SQD does not use. Has " + "NO effect when --sbd_carryover_type is non-zero: the " + "loop then takes SBD's selection as given.") sqd.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 " @@ -201,6 +214,39 @@ def parse_args(): "representation (RIKEN default 20). Must be <= 63: " "bitadvance() shifts a 64-bit size_t by this amount, " "so 64 is undefined behavior.") + sbd.add_argument("--sbd_carryover_type", type=int, default=0, + choices=[0, 1, 2, 3], + help="Let SBD pick the next iteration's determinants instead " + "of SQD thresholding the amplitudes. 0 (default): off. " + "1: rank half-determinants by marginal weight. 2: type 1 " + "plus all single excitations of what it selected. 3: rank " + "whole determinants by amplitude, then singles-extend. " + "Type 1 ALSO skips the wavefunction dump, because its " + "ranking already is the marginal weight the loop wants and " + "nothing else reads the eigenvector -- that is the win at " + "large subspace size. Types 2 and 3 keep the dump: they " + "singles-extend and SBD re-sorts the result into canonical " + "order, so the amplitudes are needed to restore the " + "descending-weight order that SCIResult.carryover requires. " + "Needs qiskit-addon-sqd >= 0.14.0 for SCIResult.carryover " + "(added in #369); on an older addon the carryover is " + "ignored and the dump always happens.") + sbd.add_argument("--sbd_carryover_ratio", type=float, default=None, + metavar="FLOAT", + help="Fraction of half-determinants to keep, ranked by " + "weight. Applies to --sbd_carryover_type 1 and 2 only, " + "and ONLY when non-zero: set it to 0 to select by " + "--sbd_carryover_threshold instead. Ignored outright by " + "type 3. Default 0.1 (note upstream SBD's own default is " + "0.0).") + sbd.add_argument("--sbd_carryover_threshold", type=float, default=None, + metavar="FLOAT", + help="Meaning depends on the type. Types 1 and 2: truncate " + "once cumulative weight reaches 1 - this, used ONLY when " + "--sbd_carryover_ratio is 0. Type 3: a direct cutoff -- " + "drop determinants with weight below this. Default 1e-4. " + "Distinct from --sqd_carryover_threshold, which is the " + "SQD layer's own knob.") # ---- MPI decomposition: hardware shape, not physics ---------------------- mpi = p.add_argument_group( @@ -255,12 +301,68 @@ def load_counts_as_bitarray(counts_path, num_bits): return BitArray.from_bool_array(bool_matrix) +def build_policy(args, rank): + """Build the SQD subspace policy, and warn when its threshold cannot take effect. + + Returns ``(policy, kwargs)``, where ``kwargs`` is what to pass to + ``diagonalize_fermionic_hamiltonian``. Since qiskit-addon-sqd 0.14.0 the loop + raises if given both ``policy=`` and ``carryover_threshold=``, so the threshold + travels inside the policy; that keeps this driver usable as a template for a + different policy (see examples/tpb/README.md). On an older addon there are no + policies and the threshold is passed directly, which selects the same schedule. + """ + threshold_given = args.sqd_carryover_threshold is not None + threshold = 1e-4 if args.sqd_carryover_threshold is None else args.sqd_carryover_threshold + args.sqd_carryover_threshold = threshold + + # A solver that reports its own carryover bypasses the policy's carryover + # selection: _select_carryover returns it before thresholding anything. + if rank == 0 and args.sbd_carryover_type != 0 and threshold_given: + print(f"WARNING: --sqd_carryover_threshold {threshold:g}: no effect with " + f"--sbd_carryover_type {args.sbd_carryover_type}: SBD then selects " + "the carryover itself and the loop takes that selection as given. " + "Tune the selection with --sbd_carryover_ratio / " + "--sbd_carryover_threshold instead.\n") + + if StandardPolicy is None: + return None, {"carryover_threshold": threshold} + policy = StandardPolicy(carryover_threshold=threshold) + return policy, {"policy": policy} + + def main(): args = parse_args() comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() + # SBD reads exactly one of ratio/threshold per carryover type and silently + # ignores the other (sbdiag.h, the carryover_type branch), so a setting that + # cannot take effect is worth saying out loud rather than leaving the user to + # wonder why tuning it changed nothing. Defaults are filled in after this so + # "was it passed?" is answerable at all. + ratio_given = args.sbd_carryover_ratio is not None + threshold_given = args.sbd_carryover_threshold is not None + if not ratio_given: + args.sbd_carryover_ratio = 0.1 + if not threshold_given: + args.sbd_carryover_threshold = 1e-4 + if rank == 0 and args.sbd_carryover_type != 0: + if (threshold_given and args.sbd_carryover_type in (1, 2) + and args.sbd_carryover_ratio != 0): + print(f"WARNING: --sbd_carryover_threshold " + f"{args.sbd_carryover_threshold:g} has no effect with " + f"--sbd_carryover_type {args.sbd_carryover_type} while " + f"--sbd_carryover_ratio is {args.sbd_carryover_ratio:g}. Types 1 " + "and 2 read the threshold only when the ratio is 0. Pass " + "--sbd_carryover_ratio 0 to select by threshold.\n") + if ratio_given and args.sbd_carryover_type == 3: + print(f"WARNING: --sbd_carryover_ratio {args.sbd_carryover_ratio:g} has " + "no effect with --sbd_carryover_type 3, which selects by " + "--sbd_carryover_threshold alone.\n") + + _, policy_kwargs = build_policy(args, rank) + norb, nelec_total, ms2 = parse_fcidump_header(args.fcidump) num_elec_a = (nelec_total + ms2) // 2 num_elec_b = (nelec_total - ms2) // 2 @@ -344,7 +446,9 @@ def main(): if rank == 0: print(f"Resuming from {args.resume_from}: iteration {last['iteration']}, " f"{len(last['ci_strs_a'])} alpha / {len(last['ci_strs_b'])} beta " - "strings carried in as include_configurations") + "strings carried in as include_configurations " + "(a checkpoint written under --sbd_carryover_type != 0 holds " + "SBD's carryover, not the whole solved subspace)") include_configurations = (include_a, include_b) if (include_a or include_b) else None if args.counts: @@ -384,6 +488,9 @@ def main(): "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, + "carryover_type": args.sbd_carryover_type, + "ratio": args.sbd_carryover_ratio, + "threshold": args.sbd_carryover_threshold, } sbd_solver = partial( @@ -409,8 +516,14 @@ def callback(results: list[SCIResult]): print(f"Iteration {iteration}") for i, r in enumerate(results): total_e = r.energy + nuclear_repulsion_energy - dim = np.prod(r.sci_state.amplitudes.shape) - print(f" Batch {i}: E={total_e:.10f}, dim={dim:_}") + if r.sci_state is None: + # SBD selected the carryover itself, so no eigenvector came + # back. The solved dimension is printed by the wrapper. + co_a, co_b = r.carryover + extent = f"carryover {len(co_a):_}a/{len(co_b):_}b" + else: + extent = f"dim={np.prod(r.sci_state.amplitudes.shape):_}" + print(f" Batch {i}: E={total_e:.10f}, {extent}") due = (iteration % args.checkpoint_frequency == 0 or iteration == args.max_iterations) if args.checkpoint_path and due: @@ -423,13 +536,21 @@ def callback(results: list[SCIResult]): # final state rather than whatever iteration happened to land on # a multiple of --checkpoint_frequency. r = results[0] + # With no eigenvector, checkpoint SBD's carryover instead of the + # solved subspace: same JSON keys, so --resume_from is unchanged, + # but a resume then seeds from what the loop itself carries + # forward rather than from everything that was diagonalized. + strs_a, strs_b = ( + r.carryover if r.sci_state is None + else (r.sci_state.ci_strs_a, r.sci_state.ci_strs_b) + ) entry = { "iteration": iteration, "energy": r.energy + nuclear_repulsion_energy, "occupancies_a": r.orbital_occupancies[0].tolist(), "occupancies_b": r.orbital_occupancies[1].tolist(), - "ci_strs_a": [int(x) for x in r.sci_state.ci_strs_a], - "ci_strs_b": [int(x) for x in r.sci_state.ci_strs_b], + "ci_strs_a": [int(x) for x in strs_a], + "ci_strs_b": [int(x) for x in strs_b], } checkpoint_history.append(entry) tmp = Path(args.checkpoint_path).with_suffix(".tmp") @@ -456,6 +577,10 @@ def callback(results: list[SCIResult]): f"--sbd_method {args.method} --sbd_eps {args.eps:g} " f"--sbd_max_it {args.max_it} --sbd_max_nb {args.max_nb} " f"--sbd_bit_length {args.bit_length}") + print(" " + f"--sbd_carryover_type {args.sbd_carryover_type} " + f"--sbd_carryover_ratio {args.sbd_carryover_ratio:g} " + f"--sbd_carryover_threshold {args.sbd_carryover_threshold:g}") print("MPI grid : " f"--task_comm_size {args.task_comm_size} " f"--adet_comm_size {args.adet_comm_size} " @@ -475,7 +600,6 @@ def callback(results: list[SCIResult]): max_iterations=args.max_iterations, energy_tol=args.energy_tol, occupancies_tol=args.occupancies_tol, - carryover_threshold=args.sqd_carryover_threshold, max_dim=args.max_dim, include_configurations=include_configurations, initial_occupancies=initial_occupancies, @@ -483,6 +607,9 @@ def callback(results: list[SCIResult]): symmetrize_spin=bool(args.symmetrize_spin), callback=callback, seed=rand_seed, + # Either policy= or, on an addon older than 0.14.0, carryover_threshold= + # -- never both, which the loop rejects. See build_policy. + **policy_kwargs, ) except RuntimeError as e: if "Failed to open FCIDUMP" in str(e): diff --git a/python/sbd_solver.py b/python/sbd_solver.py index d22b13e..e8b4508 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -171,13 +171,31 @@ def _solve_sci_core( # same value the C++ engine will use to interpret them. sbd_data = _create_sbd_config(sbd_config, backend, device_config) + # ``SCIResult.carryover`` must be ordered by DESCENDING MARGINAL WEIGHT: the loop + # takes a solver's carryover as given (fermion.py, ``_select_carryover`` returns it + # before any ranking) and a later truncation to ``max_dim`` keeps the leading + # entries, so a wrong order silently drops the strings that mattered. + # + # Only carryover_type 1 satisfies that by itself. It ranks half-determinants by + # p(D^alpha) = sum_beta |c|^2 -- exactly the addon's marginal weight -- and writes + # them out in that order, so its eigenvector has no remaining consumer and the + # |adet| x |bdet| dump can be skipped entirely. + # + # Types 2 and 3 singles-extend their selection, and SinglesExtendHalfdets ends in + # sort_bitarray (upstream extend.h), which re-sorts into CANONICAL order because + # SBD indexes subspaces by binary search. That destroys the weight order, and the + # generated determinants never had a weight at all. So those types keep the dump + # and we re-rank their carryover from the amplitudes below. + skip_wf = sbd_data.carryover_type == 1 and _addon_accepts_carryover() + adet = _ci_strings_to_sbd_dets(strings_a, norb, backend, sbd_data.bit_length) bdet = _ci_strings_to_sbd_dets(strings_b, norb, backend, sbd_data.bit_length) # Use .bin extension to trigger SBD's fast binary write path # (SaveMatrixFormWF in restart.h checks extension: .bin -> raw doubles) wf_dump_file = sbd_dir / "wavefunction.bin" - sbd_data.dump_matrix_form_wf = str(wf_dump_file) + if not skip_wf: + sbd_data.dump_matrix_form_wf = str(wf_dump_file) results = backend.tpb_diag( mpi_comm, sbd_data, fcidump, adet, bdet, loadname="", savename="" @@ -187,9 +205,20 @@ def _solve_sci_core( mpi_comm.Barrier() if mpi_rank != 0: + # Placeholder: only rank 0 is given the energy and the wavefunction. With + # skip_wf there is no SCIState to stand in, so carry an empty carryover + # instead -- a result with neither field satisfies no caller, and + # qiskit-addon-sqd rejects it outright. + placeholder = ( + {"carryover": (np.array([], dtype=np.int64), np.array([], dtype=np.int64))} + if skip_wf + else {} + ) return SCIResult( 0.0, - SCIState( + None + if skip_wf + else SCIState( amplitudes=np.empty((0, 0), dtype=np.float64), ci_strs_a=np.array([], dtype=np.int64), ci_strs_b=np.array([], dtype=np.int64), @@ -200,6 +229,7 @@ def _solve_sci_core( np.zeros(norb, dtype=np.float64), np.zeros(norb, dtype=np.float64), ), + **placeholder, ) # --- rank 0 only --- @@ -227,45 +257,209 @@ def _solve_sci_core( # amplitude magnitude (fermion.py, _carryover_* / weights_a / weights_b), so # every iteration after the first was seeded from uniform weights. See #19. # - # The dump is requested unconditionally above, so neither a missing file nor a - # size mismatch is a situation to paper over: both mean the run did not do what - # was asked, and a loud failure is the only honest response. + # The dump is requested above whenever these amplitudes are needed at all, so + # neither a missing file nor a size mismatch is a situation to paper over: both + # mean the run did not do what was asked, and a loud failure is the only honest + # response. # Label the amplitudes with the lists that were actually diagonalized, not with # the inputs. _ci_strings_to_sbd_dets ends in sort_bitarray, which sorts into # SBD's canonical order AND removes duplicates, so adet/bdet can differ from # strings_a/strings_b -- and the dump is written in adet/bdet order. Deriving the # labels from adet/bdet makes the amplitudes and their labels come from one list # by construction, instead of relying on the two orders agreeing. - solved_strings_a = _sbd_dets_to_ci_strings(adet, norb, backend, sbd_data.bit_length) - solved_strings_b = _sbd_dets_to_ci_strings(bdet, norb, backend, sbd_data.bit_length) - n_a = len(solved_strings_a) - n_b = len(solved_strings_b) - if not wf_dump_file.exists(): + if skip_wf: + # Report the dimension sci_state would otherwise have carried, so the + # per-iteration series survives without the eigenvector. + n_a, n_b = len(adet), len(bdet) + print(f" [sbd] solved subspace {n_a:_} x {n_b:_} = {n_a * n_b:_}") + sci_state = None + else: + solved_strings_a = _sbd_dets_to_ci_strings( + adet, norb, backend, sbd_data.bit_length + ) + solved_strings_b = _sbd_dets_to_ci_strings( + bdet, norb, backend, sbd_data.bit_length + ) + n_a = len(solved_strings_a) + n_b = len(solved_strings_b) + if not wf_dump_file.exists(): + raise RuntimeError( + f"SBD wrote no wavefunction dump at {wf_dump_file}. It is requested " + "whenever the amplitudes are needed, so a missing file means the " + "diagonalization did not complete as expected." + ) + flat = np.fromfile(str(wf_dump_file), dtype=np.float64) + if flat.size != n_a * n_b: + raise RuntimeError( + f"wavefunction dump at {wf_dump_file} holds {flat.size} amplitudes, " + f"expected {n_a * n_b} ({n_a} alpha x {n_b} beta over the " + "diagonalized subspace). SaveMatrixFormWF writes the full subspace; " + "a different size means the dump and the subspace have diverged." + ) + + sci_state = SCIState( + amplitudes=flat.reshape(n_a, n_b), + ci_strs_a=solved_strings_a, + ci_strs_b=solved_strings_b, + norb=norb, + nelec=nelec, + ) + + rdm1, rdm2 = assemble_rdms(results, norb) + + # When SBD selected its own carryover, hand it to the loop rather than letting it + # re-derive the same choice by thresholding our amplitudes. qiskit-addon-sqd takes + # a solver's carryover as given -- ``_select_carryover`` returns it before applying + # any threshold, ratio or ranking -- so ``carryover_threshold`` does not apply AND + # the order we hand over is the order a ``max_dim`` truncation will keep. + carryover_a, carryover_b = _extract_carryover( + results, norb, backend, sbd_data.bit_length + ) + if carryover_a is not None and sci_state is not None: + # Types 2/3 come back in canonical order (see the skip_wf comment); rank them + # the way the addon would have. Type 1 never reaches here with a sci_state. + carryover_a, carryover_b = _rank_carryover_from_amplitudes( + carryover_a, carryover_b, sci_state + ) + extra = {} + if carryover_a is not None and _addon_accepts_carryover(): + extra["carryover"] = (carryover_a, carryover_b) + + if skip_wf and "carryover" not in extra: raise RuntimeError( - f"SBD wrote no wavefunction dump at {wf_dump_file}. It is requested on " - "every call, so a missing file means the diagonalization did not " - "complete as expected." + f"carryover_type={sbd_data.carryover_type} was requested, so the " + "wavefunction was not dumped, but SBD selected no carryover " + "determinants -- leaving nothing to build the next subspace from. " + f"Either the selection was too aggressive (ratio={sbd_data.ratio}, " + f"threshold={sbd_data.threshold}) or this SBD build does not implement " + f"carryover_type={sbd_data.carryover_type}. Set carryover_type=0 to go " + "back to selecting from the amplitudes." ) - flat = np.fromfile(str(wf_dump_file), dtype=np.float64) - if flat.size != n_a * n_b: + + return SCIResult( + energy, sci_state, orbital_occupancies=occupancies, rdm1=rdm1, rdm2=rdm2, **extra + ) + + +def _addon_accepts_carryover() -> bool: + """True when the installed qiskit-addon-sqd has ``SCIResult.carryover``. + + The field arrived with qiskit-addon-sqd#369. On an older addon, passing it + would be a TypeError, so the carryover is simply not reported and the loop + falls back to thresholding the amplitudes as it always has. + """ + if SCIResult is None: + return False + import dataclasses + + return any(f.name == "carryover" for f in dataclasses.fields(SCIResult)) + + +def _rank_by_marginal_weight( + carryover: np.ndarray, solved: np.ndarray, weights: np.ndarray +) -> np.ndarray: + """Order one spin sector's carryover by descending marginal weight. + + ``SCIResult.carryover`` is required to be in that order, and the loop applies no + ranking of its own. Carryover strings that were in the diagonalized subspace get + their weight from the amplitudes; the rest -- the ones carryover_type 2/3 generate + by single excitation -- were never in the subspace and so have no weight at all. + Those go last, which is the right priority: a truncation to ``max_dim`` then sheds + weightless additions before it sheds a ranked parent. + + Mirrors the addon's own ordering (``_rank_carryover``, a stable argsort of negated + marginal weights) so the two agree on ties. + + Raises if none of the carryover is in the subspace, since the returned order + would then be SBD's canonical order rather than a ranking. That should not + happen: types 2/3 extend from the determinants SBD selected out of the + subspace, and those parents are part of the carryover. + """ + if carryover.size == 0: + return carryover + if solved.size: + # Look each carryover string up in the solved list. searchsorted needs a + # sorted haystack, so sort once and carry the permutation to recover + # weight positions. + order = np.argsort(solved, kind="stable") + sorted_solved = solved[order] + pos = np.minimum(np.searchsorted(sorted_solved, carryover), sorted_solved.size - 1) + found = sorted_solved[pos] == carryover + else: + found = np.zeros(carryover.size, dtype=bool) + if not found.any(): raise RuntimeError( - f"wavefunction dump at {wf_dump_file} holds {flat.size} amplitudes, " - f"expected {n_a * n_b} ({n_a} alpha x {n_b} beta over the " - "diagonalized subspace). SaveMatrixFormWF writes the full subspace; a " - "different size means the dump and the subspace have diverged." + f"none of the {carryover.size} carryover determinants were in the " + "diagonalized subspace, so none of them has a marginal weight and the " + "order returned here would be SBD's canonical order, not a weight " + "ranking. qiskit-addon-sqd truncates the carryover to max_dim by keeping " + "the leading entries, so it would keep an arbitrary subset." ) - sci_state = SCIState( - amplitudes=flat.reshape(n_a, n_b), - ci_strs_a=solved_strings_a, - ci_strs_b=solved_strings_b, - norb=norb, - nelec=nelec, + ranked = carryover[found] + ranked_weights = weights[order[pos[found]]] + # Stable argsort of the negated weights, matching the addon exactly. + ranked = ranked[np.argsort(-ranked_weights, kind="stable")] + return np.concatenate([ranked, carryover[~found]]) + + +def _rank_carryover_from_amplitudes( + carryover_a: np.ndarray, + carryover_b: np.ndarray, + sci_state, +) -> tuple[np.ndarray, np.ndarray]: + """Re-rank a carryover SBD returned in canonical order, using the amplitudes. + + Needed for carryover_type 2/3 only; type 1 already comes out weight-ordered. + + The ranking is PER SECTOR, where the loop would rank the two sectors jointly under + ``symmetrize_spin``. That is safe here, and not by luck. ``symmetrize_spin`` hands + the solver the SAME array for both sectors, so the subspace is symmetric and the + amplitude matrix is symmetric (or antisymmetric, for a triplet at Ms=0); marginal + weights square the amplitudes, so the two sectors' weights are equal either way + and the per-sector rankings coincide. Singles extension is spin-agnostic, so types + 2/3 inherit it. Measured on h2o: for a symmetric subspace the two carryover lists + come out bit-identical for types 1, 2 and 3. + + That matters because ``batch_to_ci_strings`` IGNORES ``carryover_strings_b`` under + ``symmetrize_spin`` -- it assumes the two are equal -- so a solver whose lists + differed there would silently lose its beta selection. Ours do not differ. When + ``symmetrize_spin`` is off, both lists are used and may legitimately differ. + """ + probabilities = np.abs(sci_state.amplitudes) ** 2 + weights_a = probabilities.sum(axis=1) + weights_b = probabilities.sum(axis=0) + return ( + _rank_by_marginal_weight(carryover_a, sci_state.ci_strs_a, weights_a), + _rank_by_marginal_weight(carryover_b, sci_state.ci_strs_b, weights_b), ) - rdm1, rdm2 = assemble_rdms(results, norb) - return SCIResult(energy, sci_state, orbital_occupancies=occupancies, rdm1=rdm1, rdm2=rdm2) +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, but they are empty unless the caller asked for a carryover type -- + ``_create_sbd_config`` leaves ``carryover_type`` at 0, so the common path pays + nothing here and gets ``(None, None)``. + + 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 -- the shape + ``SCIResult.carryover`` wants. + """ + co_a = results.get("carryover_adet") + co_b = results.get("carryover_bdet") + if not co_a and not co_b: + return None, None + empty = np.array([], dtype=np.int64) + carryover_a = ( + _sbd_dets_to_ci_strings(co_a, norb, backend, bit_length) if co_a else empty + ) + carryover_b = ( + _sbd_dets_to_ci_strings(co_b, norb, backend, bit_length) if co_b else empty + ) + return carryover_a, carryover_b def assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarray | None]: @@ -528,13 +722,13 @@ def _create_sbd_config(config_dict: dict | None = None, backend=None, device_con sbd_data.init = 0 sbd_data.do_shuffle = 0 sbd_data.do_rdm = 0 - # SBD's carryover is NOT consumed on this path: it is SBD's own iterative - # mechanism (its CLI writes it out with --carryover_adetfile and you re-run), - # whereas the SQD loop selects its own determinants from the amplitudes we - # return. Since _solve_sci_core discards results["carryover_*"], asking SBD to - # compute it is pure work -- for carryover_type=2 that includes building - # singles-extended determinant lists. Default it off; a caller who wants it - # can still set carryover_type through sbd_config. + # Off by default, but no longer discarded: since qiskit-addon-sqd#369 the loop + # accepts a solver's own carryover (SCIResult.carryover) instead of thresholding + # the amplitudes, and _solve_sci_core reports it when SBD computed one. Default + # stays 0 because computing it is real work -- carryover_type=2 builds + # singles-extended determinant lists -- and because the amplitude-threshold path + # remains the behaviour callers expect. Set carryover_type through sbd_config to + # let SBD choose the next iteration's determinants instead. sbd_data.carryover_type = 0 sbd_data.ratio = 0.1 sbd_data.threshold = 1e-4 diff --git a/releasenotes/notes/sbd-selected-carryover-37126cc4b5f6f04c.yaml b/releasenotes/notes/sbd-selected-carryover-37126cc4b5f6f04c.yaml new file mode 100644 index 0000000..74cb372 --- /dev/null +++ b/releasenotes/notes/sbd-selected-carryover-37126cc4b5f6f04c.yaml @@ -0,0 +1,22 @@ +--- +features: + - | + SBD can now choose the carryover for the next SQD iteration. When + ``carryover_type`` is non-zero in ``sbd_config``, + :func:`sbd.sbd_solver.solve_sci` and :func:`sbd.sbd_solver.solve_sci_batch` + return SBD's selection in ``SCIResult.carryover``, which ``qiskit-addon-sqd`` + 0.14.0 and later uses as given in place of its own amplitude threshold. With + an older ``qiskit-addon-sqd`` the selection is not reported and the default + behavior is unchanged. + + With ``carryover_type = 1`` the wavefunction is no longer written to disk and + read back. ``SCIResult.sci_state`` is ``None``, since SBD's ranking is already + the marginal weight the loop orders carryover by. Types 2 and 3 still return + the wavefunction, which is used to put their carryover in descending order of + weight, as ``SCIResult.carryover`` requires. The default, + ``carryover_type = 0``, is unchanged. + - | + The ``examples/tpb/run_sqd_sbd.py`` driver gains ``--sbd_carryover_type``, + ``--sbd_carryover_ratio`` and ``--sbd_carryover_threshold``, and warns when a + setting cannot take effect. ``examples/tpb/README.md`` also shows how to run + the loop with a different subspace policy, such as Trim SQD. diff --git a/test/test_sqd_integration.py b/test/test_sqd_integration.py index 36ef21c..19bc00a 100644 --- a/test/test_sqd_integration.py +++ b/test/test_sqd_integration.py @@ -175,9 +175,12 @@ def _assert_result_is_consistent(result, nelec=NELEC): assert alpha_occupancies.sum() == pytest.approx(nelec[0], abs=1e-6) assert beta_occupancies.sum() == pytest.approx(nelec[1], abs=1e-6) + # With carryover_type=1 the solver reports SBD's carryover instead of the + # eigenvector, so there is no state to check. state = result.sci_state - assert state.amplitudes.shape == (len(state.ci_strs_a), len(state.ci_strs_b)) - assert np.isfinite(state.amplitudes).all() + if state is not None: + assert state.amplitudes.shape == (len(state.ci_strs_a), len(state.ci_strs_b)) + assert np.isfinite(state.amplitudes).all() # --- solve_sci_batch over a fixed subspace ------------------------------------------- @@ -244,7 +247,17 @@ def test_published_energy_mpi(data_dir, device_config): # tests above are what guard the number. -def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_path): +def _check_diagonalize_fermionic_hamiltonian( + data_dir, device_config, counts_path, carryover_type=0 +): + """Run the qiskit-addon-sqd loop with SBD as the solver. + + ``carryover_type=1`` drives the SBD-selected carryover path end to end: the loop + must accept ``sci_state=None``, take SBD's carryover in place of its own + threshold, and build the next iteration from it. ``symmetrize_spin`` is on, and + under it the loop ignores ``carryover_strings_b``, so the two lists must agree -- + otherwise the beta selection would be silently dropped. + """ from functools import partial from mpi4py import MPI @@ -256,7 +269,10 @@ def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_pat from qiskit.primitives import BitArray from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian - from sbd.sbd_solver import solve_sci_batch + from sbd.sbd_solver import _addon_accepts_carryover, solve_sci_batch + + if carryover_type and not _addon_accepts_carryover(): + pytest.skip("installed qiskit-addon-sqd has no SCIResult.carryover") comm = MPI.COMM_WORLD hcore, eri, nuclear_repulsion = _load_hamiltonian(data_dir) @@ -270,10 +286,16 @@ def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_pat # eight digits; the loop only needs each iteration to converge. max_nb matches the # example notebook and upstream's own default. bit_length is inherited from # SOLVER_CONFIG and must not be overridden -- see the note there. - sbd_config=_sbd_config(comm, eps=1e-8, max_it=10, max_nb=10), + sbd_config=_sbd_config( + comm, eps=1e-8, max_it=10, max_nb=10, carryover_type=carryover_type + ), device_config=device_config, ) + # Every iteration's results, as the loop hands them over. The callback runs on + # the control process only, which is also the only rank that asserts below. + iterations = [] + result = diagonalize_fermionic_hamiltonian( hcore, eri, @@ -288,6 +310,7 @@ def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_pat # Seeded so a failure can be reproduced, though the assertions below do not # depend on which determinants the sampling happens to pick. seed=np.random.default_rng(42), + callback=iterations.append, ) if comm.Get_rank() != 0: @@ -295,6 +318,20 @@ def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_pat _assert_result_is_consistent(result) + # max_iterations=2, so the second iteration was built from the first's carryover. + assert len(iterations) == 2 + for results in iterations: + for r in results: + if carryover_type == 1: + assert r.sci_state is None + carryover_a, carryover_b = r.carryover + assert len(carryover_a) > 0 + np.testing.assert_array_equal(carryover_a, carryover_b) + else: + assert r.sci_state is not None + # SCIResult has no carryover field before qiskit-addon-sqd 0.14.0. + assert getattr(r, "carryover", None) is None + # A bracket rather than an equality: the subspace comes out of upstream's sampling. # The lower bound is the FCI energy, which no subspace of it can beat; the upper is # loose enough to tolerate a different selection but far below the -75.0 a badly @@ -303,14 +340,185 @@ def _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_pat assert -76.25 < total_energy < -75.9 +@pytest.mark.parametrize("carryover_type", [0, 1]) def test_diagonalize_fermionic_hamiltonian_standalone( - data_dir, device_config, counts_path + data_dir, device_config, counts_path, carryover_type ): """The self-consistent loop runs to completion in a single process.""" - _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_path) + _check_diagonalize_fermionic_hamiltonian( + data_dir, device_config, counts_path, carryover_type + ) @pytest.mark.mpi -def test_diagonalize_fermionic_hamiltonian_mpi(data_dir, device_config, counts_path): +@pytest.mark.parametrize("carryover_type", [0, 1]) +def test_diagonalize_fermionic_hamiltonian_mpi( + data_dir, device_config, counts_path, carryover_type +): """The self-consistent loop runs to completion across the launched ranks.""" - _check_diagonalize_fermionic_hamiltonian(data_dir, device_config, counts_path) + _check_diagonalize_fermionic_hamiltonian( + data_dir, device_config, counts_path, carryover_type + ) + + +# --- SBD-selected carryover skips the wavefunction dump ------------------------------ +# +# With carryover_type != 0 SBD ranks and selects the next subspace in C++, so the +# eigenvector has no consumer and is not dumped. The point of the test is the absence +# of the file: a run that quietly kept writing it would still give the right energy. + + +def _check_sbd_carryover_skips_wavefunction_dump(data_dir, device_config, tmp_path): + """``carryover_type=1`` returns no SCI state and writes no wavefunction. + + Only type 1 qualifies: it ranks by the marginal weight the loop wants, so the + eigenvector has no remaining consumer. Types 2/3 singles-extend and are re-sorted + canonically by SBD, so they need the amplitudes to restore the required order -- + see ``test_sbd_carryover_is_weight_ordered``. + """ + from mpi4py import MPI + + from sbd.sbd_solver import _addon_accepts_carryover, solve_sci_batch + + if not _addon_accepts_carryover(): + pytest.skip( + "installed qiskit-addon-sqd has no SCIResult.carryover (needs #369); " + "the solver keeps returning sci_state so the loop has something to use" + ) + + comm = MPI.COMM_WORLD + hcore, eri, _ = _load_hamiltonian(data_dir) + strings = _read_alpha_determinants( + data_dir / "h2o" / "h2o-1em3-alpha.txt", limit=SMALL_SUBSPACE_DETS + ) + + results = solve_sci_batch( + [(strings, strings)], + hcore, + eri, + norb=NORB, + nelec=NELEC, + sbd_config=_sbd_config(comm, carryover_type=1), + device_config=device_config, + temp_dir=tmp_path, + clean_temp_dir=False, + ) + result = results[0] + + # Holds on every rank: no rank materializes the eigenvector, and each still + # returns something the loop can use. + assert result.sci_state is None + assert result.carryover is not None + + # Only rank 0 is given the energy and the selected determinants. + if comm.Get_rank() != 0: + return + + carryover_a, carryover_b = result.carryover + assert len(carryover_a) > 0 + assert len(carryover_b) > 0 + # The energy is unaffected by how the next subspace gets chosen. + assert result.energy == pytest.approx(SMALL_SUBSPACE_ENERGY, abs=1e-8) + # clean_temp_dir=False above keeps the directory, so this is a real check that + # nothing was written rather than a check that it was tidied away. + assert not list(tmp_path.rglob("wavefunction.bin")) + + +def test_sbd_carryover_skips_wavefunction_dump_standalone( + data_dir, device_config, tmp_path +): + """``carryover_type=1`` skips the dump in a single process.""" + _check_sbd_carryover_skips_wavefunction_dump(data_dir, device_config, tmp_path) + + +@pytest.mark.mpi +def test_sbd_carryover_skips_wavefunction_dump_mpi(data_dir, device_config, tmp_path): + """``carryover_type=1`` skips the dump across the launched ranks.""" + _check_sbd_carryover_skips_wavefunction_dump(data_dir, device_config, tmp_path) + + +def _check_sbd_carryover_is_weight_ordered( + data_dir, device_config, tmp_path, carryover_type +): + """Types 2/3 keep the dump and hand over a weight-ordered carryover. + + ``SCIResult.carryover`` must be in descending marginal-weight order, because + ``_select_carryover`` returns a solver's carryover untouched and a later + truncation to ``max_dim`` keeps the leading entries. SBD returns types 2/3 in + canonical order (``SinglesExtendHalfdets`` ends in ``sort_bitarray``), so the + solver re-ranks them from the amplitudes. Without that, ``max_dim`` would keep + the numerically smallest strings and could discard the high-weight parents. + """ + from mpi4py import MPI + + from sbd.sbd_solver import _addon_accepts_carryover, solve_sci_batch + + if not _addon_accepts_carryover(): + pytest.skip("installed qiskit-addon-sqd has no SCIResult.carryover") + + comm = MPI.COMM_WORLD + hcore, eri, _ = _load_hamiltonian(data_dir) + strings = _read_alpha_determinants( + data_dir / "h2o" / "h2o-1em3-alpha.txt", limit=SMALL_SUBSPACE_DETS + ) + results = solve_sci_batch( + [(strings, strings)], + hcore, + eri, + norb=NORB, + nelec=NELEC, + sbd_config=_sbd_config(comm, carryover_type=carryover_type), + device_config=device_config, + temp_dir=tmp_path, + clean_temp_dir=False, + ) + result = results[0] + + if comm.Get_rank() != 0: + return + + # These types need the amplitudes, so the dump is taken rather than skipped. + assert result.sci_state is not None + assert result.carryover is not None + assert result.energy == pytest.approx(SMALL_SUBSPACE_ENERGY, abs=1e-8) + + probabilities = np.abs(result.sci_state.amplitudes) ** 2 + for strings_co, solved, weights in ( + (result.carryover[0], result.sci_state.ci_strs_a, probabilities.sum(axis=1)), + (result.carryover[1], result.sci_state.ci_strs_b, probabilities.sum(axis=0)), + ): + assert len(strings_co) > 0 + order = np.argsort(solved, kind="stable") + sorted_solved = solved[order] + pos = np.minimum( + np.searchsorted(sorted_solved, strings_co), sorted_solved.size - 1 + ) + in_subspace = sorted_solved[pos] == strings_co + # Without this, both assertions below would pass on an empty selection. + assert in_subspace.any() + # Strings that were in the subspace come first, ranked by descending weight; + # the singles-generated ones have no weight and must follow. + assert np.all(np.diff(in_subspace.astype(int)) <= 0) + ranked = weights[order[pos[in_subspace]]] + assert np.all(np.diff(ranked) <= 1e-12) + + +@pytest.mark.parametrize("carryover_type", [2, 3]) +def test_sbd_carryover_is_weight_ordered_standalone( + data_dir, device_config, tmp_path, carryover_type +): + """Types 2/3 hand over a weight-ordered carryover in a single process.""" + _check_sbd_carryover_is_weight_ordered( + data_dir, device_config, tmp_path, carryover_type + ) + + +@pytest.mark.mpi +@pytest.mark.parametrize("carryover_type", [2, 3]) +def test_sbd_carryover_is_weight_ordered_mpi( + data_dir, device_config, tmp_path, carryover_type +): + """Types 2/3 hand over a weight-ordered carryover across the launched ranks.""" + _check_sbd_carryover_is_weight_ordered( + data_dir, device_config, tmp_path, carryover_type + )