diff --git a/python/examples/README.md b/python/examples/README.md index d6d0700..da59ac4 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -28,12 +28,31 @@ mpirun -np 8 python -u run_sbd_diag.py \ --fcidump ../../vendor/sbd-upstream/data/n2/fcidump.txt \ --adetfile ../../vendor/sbd-upstream/data/n2/1em3-alpha.txt \ --adet_comm_size 2 --bdet_comm_size 2 --task_comm_size 2 + +# Retrieve the 1-/2-particle RDMs and save them to a file +mpirun -np 2 python -u run_sbd_diag.py --rdm_output /tmp/h2o_rdms.npz ``` -**Key options:** `--device`, `--fcidump`, `--adetfile`, `--adet_comm_size`, -`--bdet_comm_size`, `--task_comm_size`, `--method`, `--tolerance`, `--iteration`. -(These keep their unprefixed names here: this driver *is* SBD. The SQD driver -prefixes them `--sbd_*`.) Run `python run_sbd_diag.py --help` for the full list. +`--rdm_output` takes the file to save to (`/tmp/h2o_rdms.npz` above) and +writes **one** `.npz` file there holding both `rdm1` and `rdm2` together +(`data = np.load("/tmp/h2o_rdms.npz"); data["rdm1"]`, `data["rdm2"]`) — +unlike upstream SBD's own CLI, which writes two separate files +(`1pRDM.txt`/`2pRDM.txt`). It also prints `trace(rdm1)` and the natural +orbital occupations. Leaving it empty (the default) skips computing RDMs +entirely. + +By default beta determinants are derived from `--adetfile` alone (identical +to it, or a shuffled copy if `--shuffle` is set). `--symmetrize_spin 0` +loads `--adetfile` and `--bdetfile` as independent, genuinely distinct +alpha/beta determinant sets instead; `--bdetfile` is otherwise ignored +(with a warning) since symmetric mode always derives beta from alpha. + +**Key options:** `--device`, `--fcidump`, `--adetfile`, `--bdetfile`, +`--symmetrize_spin`, `--adet_comm_size`, `--bdet_comm_size`, +`--task_comm_size`, `--method`, `--tolerance`, `--iteration`, +`--rdm_output`. (These keep their unprefixed names here: this driver *is* +SBD. The SQD drivers prefix them `--sbd_*`.) Run `python run_sbd_diag.py +--help` for the full list. **Requirements:** `sbd`, `mpi4py` @@ -201,7 +220,8 @@ prints an OOM warning when it detects this. |-----------|-----------------|---------| | `--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` the alpha and beta string sets are merged, so the subspace is up to `(2N)^2`, not `N^2` | `3000` | +| `--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` | diff --git a/python/examples/run_sbd_diag.py b/python/examples/run_sbd_diag.py index 39adf63..4ac49cf 100644 --- a/python/examples/run_sbd_diag.py +++ b/python/examples/run_sbd_diag.py @@ -34,11 +34,22 @@ mpirun -np 8 python run_sbd_diag.py \ --fcidump ../../vendor/sbd-upstream/data/h2o/fcidump.txt \ --adetfile ../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt + + # Retrieve the 1-/2-particle RDMs: saves ONE .npz file holding both + # rdm1 and rdm2 (np.load(path)["rdm1"] / ["rdm2"]) -- not two separate + # files the way upstream SBD's own CLI does (1pRDM.txt/2pRDM.txt). + mpirun -np 8 python run_sbd_diag.py --rdm_output /tmp/h2o_rdms.npz + + # Distinct alpha and beta determinant files (default is beta = alpha) + mpirun -np 8 python run_sbd_diag.py --symmetrize_spin 0 \ + --adetfile alpha-dets.txt --bdetfile beta-dets.txt """ import argparse import sys +import numpy as np + def parse_args(): """Parse command line arguments for all TPB_SBD parameters""" parser = argparse.ArgumentParser( @@ -64,7 +75,11 @@ def parse_args(): parser.add_argument('--adetfile', default='../../vendor/sbd-upstream/data/h2o/h2o-1em3-alpha.txt', help='Path to alpha determinants file') parser.add_argument('--bdetfile', default='', - help='Path to beta determinants file (optional, uses adetfile if not specified)') + help='Path to beta determinants file, used only when ' + '--symmetrize_spin 0 (otherwise ignored with a ' + 'warning: symmetric mode always derives beta from ' + '--adetfile). Defaults to --adetfile itself when ' + 'left unset.') parser.add_argument('--loadname', default='', help='Load initial wavefunction from file') parser.add_argument('--savename', default='', @@ -94,10 +109,29 @@ def parse_args(): parser.add_argument('--init', type=int, default=0, help='Initialization method') parser.add_argument('--shuffle', '--do_shuffle', type=int, default=0, dest='do_shuffle', - help='Shuffle determinants (0=no, 1=yes)') - parser.add_argument('--rdm', '--do_rdm', type=int, default=0, choices=[0, 1], dest='do_rdm', - help='Calculate RDM (0=density only, 1=full RDM)') - + help='Shuffle determinants loaded from --adetfile before ' + 'mirroring/deriving beta from them (0=no, 1-4=yes, ' + 'different shuffle seeds -- see sbdiag.h). Only ' + 'takes effect when --symmetrize_spin 1 (default): ' + 'that is the code path that derives beta from a ' + 'single loaded list at all.') + parser.add_argument('--symmetrize_spin', type=int, default=1, choices=[0, 1], + help='1 (default): beta determinants are derived from ' + '--adetfile alone (identical to it, or a shuffled ' + 'copy if --shuffle is set) -- --bdetfile is ignored ' + 'with a warning if given. 0: load --adetfile and ' + '--bdetfile as independent, genuinely distinct ' + 'alpha/beta determinant sets (--shuffle has no ' + 'effect in this mode).') + parser.add_argument('--rdm_output', default='', metavar='FILE', + help='Save rdm1 and rdm2 together in ONE numpy .npz ' + 'file at this path (np.load(FILE)["rdm1"] / ' + '["rdm2"]) -- unlike upstream SBD\'s own CLI, ' + 'which writes two separate files (1pRDM.txt / ' + '2pRDM.txt). Also prints trace(rdm1) and the ' + 'natural orbital occupations. Empty (default): ' + 'no RDMs computed at all.') + # Carryover determinant selection parser.add_argument('--carryover_type', type=int, default=0, help='Carryover determinant selection type') @@ -123,12 +157,14 @@ def parse_args(): return parser.parse_args() + def main(): args = parse_args() # Import sbd — auto-initializes on first use, but we call init() # explicitly here to set the default device from --device flag. import sbd + from sbd.sbd_solver import assemble_rdms sbd.init(device=args.device) @@ -148,7 +184,10 @@ def main(): config.eps = args.eps config.method = args.method config.max_nb = args.max_nb - config.do_rdm = args.do_rdm + config.max_time = args.max_time + config.init = args.init + config.do_shuffle = args.do_shuffle + config.do_rdm = 1 if args.rdm_output else 0 config.bit_length = args.bit_length config.carryover_type = args.carryover_type config.ratio = args.ratio @@ -156,6 +195,13 @@ def main(): config.adet_comm_size = args.adet_comm_size config.bdet_comm_size = args.bdet_comm_size config.task_comm_size = args.task_comm_size + # Thrust-only fields -- absent from the CPU/OMP-offload backends' TPB_SBD, + # so guard with hasattr rather than assume, matching sbd_solver.py's own + # _create_sbd_config pattern for exactly this reason. + if hasattr(config, 'use_precalculated_dets'): + config.use_precalculated_dets = bool(args.use_precalculated_dets) + if hasattr(config, 'max_memory_gb_for_determinants'): + config.max_memory_gb_for_determinants = args.max_memory_gb_for_determinants if rank == 0: print("Configuration:") @@ -190,13 +236,37 @@ def main(): if args.dump_matrix_form_wf: config.dump_matrix_form_wf = args.dump_matrix_form_wf - results = sbd.tpb_diag_from_files( - fcidumpfile=args.fcidump, - adetfile=args.adetfile, - sbd_data=config, - loadname=args.loadname, - savename=args.savename, - ) + if args.symmetrize_spin: + if args.bdetfile and rank == 0: + print(f"WARNING: --bdetfile {args.bdetfile!r} is ignored because " + "--symmetrize_spin is 1 (default) -- symmetric mode always " + "derives beta from --adetfile alone. Pass --symmetrize_spin 0 " + "to use a distinct beta-determinant file.\n") + results = sbd.tpb_diag_from_files( + fcidumpfile=args.fcidump, + adetfile=args.adetfile, + sbd_data=config, + loadname=args.loadname, + savename=args.savename, + ) + else: + # No bound function accepts two separate determinant files + # directly (tpb_diag_from_files always derives beta from the one + # adetfile it's given) -- load both ourselves and call the + # data-structure entry point instead, mirroring what + # tpb_diag_from_files itself does internally after loading + # (sbdiag.h: LoadAlphaDets(...); sort_bitarray(adet);). + bdetfile = args.bdetfile or args.adetfile + fcidump = sbd.LoadFCIDump(args.fcidump) + norb = int(fcidump.header["NORB"]) + adet = sbd.sort_bitarray( + sbd.LoadAlphaDets(args.adetfile, args.bit_length, norb)) + bdet = sbd.sort_bitarray( + sbd.LoadAlphaDets(bdetfile, args.bit_length, norb)) + results = sbd.tpb_diag( + fcidump, adet, bdet, config, + loadname=args.loadname, savename=args.savename, + ) if rank == 0: print("="*70) @@ -214,6 +284,26 @@ def main(): print(f"Density: {combined_density}") print(f"Carryover determinants: {len(results['carryover_adet'])}") + + norb = len(density) // 2 + rdm1, rdm2 = assemble_rdms(results, norb) + if rdm1 is not None: + print() + print(f"1-RDM trace: {np.trace(rdm1):.6f} " + "(should equal the total electron count)") + occupations = np.sort(np.linalg.eigvalsh(rdm1))[::-1] + print(f"Natural orbital occupations (sorted): " + f"{np.round(occupations, 6).tolist()}") + print(" -- occupations near 2 or 0 indicate a " + "single-reference-like orbital; occupations near 1 " + "(or several clustered together) flag multi-reference " + "character / a candidate active space.") + # rdm1 is only non-None when --rdm_output was given (that's + # what turns config.do_rdm on above), so this always fires + # here -- no separate on/off check needed. + np.savez(args.rdm_output, rdm1=rdm1, rdm2=rdm2) + print(f"Saved rdm1/rdm2 to {args.rdm_output}") + print("="*70) print("\n✓ Calculation completed successfully!") print() diff --git a/python/examples/run_sqd_enlarge_subspace_sbd.py b/python/examples/run_sqd_enlarge_subspace_sbd.py index 4d2b062..bc2b63c 100644 --- a/python/examples/run_sqd_enlarge_subspace_sbd.py +++ b/python/examples/run_sqd_enlarge_subspace_sbd.py @@ -81,6 +81,13 @@ def parse_args(): "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 " @@ -136,8 +143,6 @@ def parse_args(): 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_do_rdm", type=int, default=0, dest="do_rdm") - sbd.add_argument("--sbd_do_shuffle", type=int, default=0, dest="do_shuffle") 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) @@ -369,7 +374,8 @@ def main(): print(" " f"--energy_tol {args.energy_tol:g} " f"--occupancies_tol {args.occupancies_tol:g} " - f"--enlarge_threshold {args.enlarge_threshold:g}") + f"--enlarge_threshold {args.enlarge_threshold: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}") @@ -377,8 +383,7 @@ def main(): sbd_config = { "method": args.method, "eps": args.eps, "max_it": args.max_it, - "max_nb": args.max_nb, "max_time": 3600.0, "do_rdm": args.do_rdm, - "do_shuffle": args.do_shuffle, "bit_length": args.bit_length, + "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, @@ -413,7 +418,7 @@ def callback(results): include_configurations=current_include, initial_occupancies=current_occ, sci_solver=sbd_solver, - symmetrize_spin=True, + symmetrize_spin=bool(args.symmetrize_spin), max_dim=args.max_dim, callback=callback, seed=rand_seed, diff --git a/python/examples/run_sqd_sbd.py b/python/examples/run_sqd_sbd.py index 8c836d5..e6999a0 100644 --- a/python/examples/run_sqd_sbd.py +++ b/python/examples/run_sqd_sbd.py @@ -81,6 +81,13 @@ def parse_args(): help="Dominant control on subspace size. With " "symmetrize_spin the alpha and beta string sets " "merge, so the subspace is up to (2N)^2.") + sqd.add_argument("--symmetrize_spin", type=int, default=1, choices=[0, 1], + help="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, 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.") sqd.add_argument("--num_batches", type=int, default=3) sqd.add_argument("--max_iterations", type=int, default=5, help="SQD self-consistent loop iterations. NOT the SBD " @@ -174,11 +181,6 @@ def parse_args(): "and cross-batch agreement.") sbd.add_argument("--sbd_max_nb", "--block", "--max_nb", type=int, default=10, dest="max_nb") - sbd.add_argument("--sbd_do_rdm", "--rdm", "--do_rdm", type=int, default=0, - dest="do_rdm", - help="0=density only (default, sufficient for SQD), 1=full RDM") - sbd.add_argument("--sbd_do_shuffle", "--shuffle", "--do_shuffle", type=int, - default=0, dest="do_shuffle") sbd.add_argument("--sbd_use_precalculated_dets", type=int, default=1, choices=[0, 1], help="Thrust only. 1 precomputes a determinant index for every " @@ -366,8 +368,6 @@ def main(): "max_it": args.max_it, "max_nb": args.max_nb, "max_time": 3600.0, - "do_rdm": args.do_rdm, - "do_shuffle": args.do_shuffle, "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, @@ -440,7 +440,8 @@ def callback(results: list[SCIResult]): print(" " f"--energy_tol {args.energy_tol:g} " f"--occupancies_tol {args.occupancies_tol:g} " - f"--sqd_carryover_threshold {args.sqd_carryover_threshold:g}") + f"--sqd_carryover_threshold {args.sqd_carryover_threshold: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} " @@ -469,7 +470,7 @@ def callback(results: list[SCIResult]): include_configurations=include_configurations, initial_occupancies=initial_occupancies, sci_solver=sbd_solver, - symmetrize_spin=True, + symmetrize_spin=bool(args.symmetrize_spin), callback=callback, seed=rand_seed, ) diff --git a/python/sbd_solver.py b/python/sbd_solver.py index e95a2d6..e9fb489 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -263,12 +263,12 @@ def _solve_sci_core( nelec=nelec, ) - rdm1, rdm2 = _assemble_rdms(results, norb) + rdm1, rdm2 = assemble_rdms(results, norb) return SCIResult(energy, sci_state, orbital_occupancies=occupancies, rdm1=rdm1, rdm2=rdm2) -def _assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarray | None]: +def assemble_rdms(results: dict, norb: int) -> tuple[np.ndarray | None, np.ndarray | None]: """Build spin-summed (rdm1, rdm2) from SBD's raw one_p_rdm/two_p_rdm. Returns (None, None) when ``do_rdm`` was 0 (SBD leaves these keys as