From a9d0644c741e75a40e0e90c661325399dd9fe465 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Mon, 2 Nov 2026 15:18:22 -0500 Subject: [PATCH 01/16] Let sbd_config set the Thrust options; bind thrust_collapse_loops On the Thrust backend, _create_sbd_config applied DeviceConfig after sbd_config, and DeviceConfig always wrote use_precalculated_dets and max_memory_gb_for_determinants from its own defaults. A value set in sbd_config was reset, so run_sqd_sbd.py and run_sqd_enlarge_subspace_sbd.py ignored --sbd_use_precalculated_dets and --sbd_max_memory_gb_for_determinants. DeviceConfig is now applied first and writes only the arguments it was given; those two arguments are deprecated in favour of sbd_config, which takes precedence over them. Also bind upstream's Thrust-only TPB field thrust_collapse_loops, present at the vendored pin but not exposed, with --thrust_collapse_loops in run_sbd_diag.py and --sbd_thrust_collapse_loops in the two SQD drivers. Co-Authored-By: Claude Opus 5.5 --- docs/apidocs/sbd.rst | 5 + examples/tpb/README.md | 1 + examples/tpb/run_sbd_diag.py | 5 + examples/tpb/run_sqd_enlarge_subspace_sbd.py | 2 + examples/tpb/run_sqd_sbd.py | 5 + python/bindings.cpp | 6 + python/device_config.py | 68 +++++-- python/sbd_solver.py | 11 +- ...st-options-reach-sbd-8c1f4e2a96d07b35.yaml | 21 +++ test/test_thrust_options.py | 174 ++++++++++++++++++ 10 files changed, 280 insertions(+), 18 deletions(-) create mode 100644 releasenotes/notes/thrust-options-reach-sbd-8c1f4e2a96d07b35.yaml create mode 100644 test/test_thrust_options.py diff --git a/docs/apidocs/sbd.rst b/docs/apidocs/sbd.rst index 36dd52d..7a3d493 100644 --- a/docs/apidocs/sbd.rst +++ b/docs/apidocs/sbd.rst @@ -112,6 +112,11 @@ Fields of ``TPB_SBD`` only: * - ``max_memory_gb_for_determinants`` - Thrust backend only, and used only when ``use_precalculated_dets`` is off: a per-thread buffer cap in GB, or ``-1`` for none. + * - ``thrust_collapse_loops`` + - Thrust backend only, and used only when ``use_precalculated_dets`` is on. On + (the default), each task's loops run as one flattened GPU loop; off, as nested + loops, which need less GPU memory. Upstream turns it off by itself when the + flattened size would overflow. Energies do not change. Fields of ``GDB_SBD`` only: diff --git a/examples/tpb/README.md b/examples/tpb/README.md index 0671068..0020366 100644 --- a/examples/tpb/README.md +++ b/examples/tpb/README.md @@ -284,6 +284,7 @@ Per diagonalization, not per loop: | `--sbd_method` | 0=Davidson, 1=Davidson+Ham, 2=Lanczos, 3=Lanczos+Ham | `0` | | `--sbd_use_precalculated_dets` | Thrust only. `1` precomputes a determinant index for **every** (α,β) pair — the whole subspace, on the GPU. `0` uses per-thread storage: slower per matvec, far less memory | `1` | | `--sbd_max_memory_gb_for_determinants` | Thrust only, and **only consulted when `--sbd_use_precalculated_dets 0`** (`mult_thrust.h:257-273`). Caps the per-thread buffer in GB | `-1` (uncapped) | +| `--sbd_thrust_collapse_loops` | Thrust only, and **only consulted when `--sbd_use_precalculated_dets 1`**. `1` runs each task's loops as one flattened GPU loop; `0` as nested loops, which need less GPU memory. Upstream switches to `0` by itself if the flattened size would overflow. Energies are unchanged | `1` | For reference, upstream's own `TPB_SBD` struct defaults are looser still (`max_it=1`, `eps=1e-4`), and `run_sbd_diag.py` uses `eps=1e-3`. On the h2o counts case, diff --git a/examples/tpb/run_sbd_diag.py b/examples/tpb/run_sbd_diag.py index 15609e6..4b94a13 100644 --- a/examples/tpb/run_sbd_diag.py +++ b/examples/tpb/run_sbd_diag.py @@ -172,6 +172,9 @@ def parse_args(): parser.add_argument('--max_memory_gb_for_determinants', '--gpu-memory', type=int, default=-1, dest='max_memory_gb_for_determinants', help='Maximum GPU memory in GB for determinants (-1=auto, GPU only)') + parser.add_argument('--thrust_collapse_loops', type=int, default=1, choices=[0, 1], + help='Thrust only, and only with --use_precalculated_dets 1: 1 runs ' + 'the loops flattened (faster), 0 nested (less GPU memory)') args = parser.parse_args() if args.initial_bdet and not args.initial_adet: @@ -230,6 +233,8 @@ def main(): 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 hasattr(config, 'thrust_collapse_loops'): + config.thrust_collapse_loops = bool(args.thrust_collapse_loops) if rank == 0: print("Configuration:") diff --git a/examples/tpb/run_sqd_enlarge_subspace_sbd.py b/examples/tpb/run_sqd_enlarge_subspace_sbd.py index b446612..cc7674e 100644 --- a/examples/tpb/run_sqd_enlarge_subspace_sbd.py +++ b/examples/tpb/run_sqd_enlarge_subspace_sbd.py @@ -154,6 +154,7 @@ def parse_args(): 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_thrust_collapse_loops", type=int, default=1, choices=[0, 1]) sbd.add_argument("--sbd_bit_length", type=int, default=20, dest="bit_length") # ---- MPI decomposition ------------------------------------------------------ @@ -441,6 +442,7 @@ def main(): "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, + "thrust_collapse_loops": bool(args.sbd_thrust_collapse_loops), "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, } diff --git a/examples/tpb/run_sqd_sbd.py b/examples/tpb/run_sqd_sbd.py index 4b37132..468a5cc 100644 --- a/examples/tpb/run_sqd_sbd.py +++ b/examples/tpb/run_sqd_sbd.py @@ -208,6 +208,10 @@ def parse_args(): help="Thrust only, and only with --sbd_use_precalculated_dets 0: " "cap the per-thread determinant buffer in GB. -1 means " "uncapped.") + sbd.add_argument("--sbd_thrust_collapse_loops", type=int, default=1, choices=[0, 1], + help="Thrust only, and only with --sbd_use_precalculated_dets 1: " + "1 runs each task's loops as one flattened GPU loop (faster), " + "0 as nested loops (less GPU memory). Energies are unchanged.") sbd.add_argument("--sbd_bit_length", "--bit_length", type=int, default=20, dest="bit_length", help="Bits packed into each size_t of the bitstring " @@ -485,6 +489,7 @@ def main(): "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, + "thrust_collapse_loops": bool(args.sbd_thrust_collapse_loops), "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, diff --git a/python/bindings.cpp b/python/bindings.cpp index d88f275..3a20663 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -341,6 +341,12 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { "Use precalculated determinants (THRUST)") .def_readwrite("max_memory_gb_for_determinants", &sbd::tpb::SBD::max_memory_gb_for_determinants, "Maximum memory in GB for determinants (THRUST)") + .def_readwrite("thrust_collapse_loops", &sbd::tpb::SBD::thrust_collapse_loops, + "Run each task's bra x excitation loops as one flattened GPU " + "loop (THRUST; default true). Only used when " + "use_precalculated_dets is on; false uses nested loops, which " + "need less GPU memory. Upstream turns it off by itself when " + "the flattened size would overflow. Energies are unchanged.") #endif ; diff --git a/python/device_config.py b/python/device_config.py index 477c53a..ff19895 100644 --- a/python/device_config.py +++ b/python/device_config.py @@ -32,14 +32,20 @@ class DeviceConfig: device = DeviceConfig.auto() # best backend that is built device = DeviceConfig.cpu() # or force one - device = DeviceConfig.gpu(max_memory_gb=16) + device = DeviceConfig.gpu() results = solve_sci_batch(..., device_config=device) + + It chooses the backend only. The solver's options, including the Thrust-only + ``use_precalculated_dets`` and ``max_memory_gb_for_determinants``, go in + ``sbd_config``. The two ``DeviceConfig`` arguments of the same purpose are + deprecated: they still work, but a value given in ``sbd_config`` takes + precedence over them. """ def __init__(self, device: str = 'cpu', - use_precalculated_dets: bool = True, - max_memory_gb: int = -1, + use_precalculated_dets: bool | None = None, + max_memory_gb: int | None = None, use_gpu: bool | None = None): """ Initialize device configuration. @@ -49,8 +55,12 @@ def __init__(self, device: str = 'cpu', NVIDIA-only), 'gpu-omp' (OpenMP target offload, NVIDIA or AMD), or any alias known to ``sbd._device_aliases``. Default 'cpu'. - use_precalculated_dets: Use precalculated determinants (GPU only) - max_memory_gb: Maximum GPU memory in GB (-1 = auto) + use_precalculated_dets: Deprecated; set ``use_precalculated_dets`` in + ``sbd_config`` instead. ``None`` (the default) leaves the field to + ``sbd_config`` and then to upstream's default. + max_memory_gb: Deprecated; set ``max_memory_gb_for_determinants`` in + ``sbd_config`` instead. ``None`` (the default) leaves the field to + ``sbd_config`` and then to upstream's default. use_gpu: Deprecated boolean. If supplied without ``device``, ``True`` maps to 'gpu' (Thrust) and ``False`` to 'cpu' for backward compatibility with pre-OMP-offload code. @@ -60,11 +70,27 @@ def __init__(self, device: str = 'cpu', device = 'gpu' if use_gpu else 'cpu' self.device = device self.use_gpu = device != 'cpu' + if use_precalculated_dets is not None or max_memory_gb is not None: + import sys + import warnings + # Point at the caller's line, whether it called DeviceConfig(...) + # directly or through the constructors below (gpu_nvidia_omp goes + # through gpu_omp): skip every frame in this module. + stacklevel, frame = 1, sys._getframe(0) + while frame is not None and frame.f_code.co_filename == __file__: + stacklevel, frame = stacklevel + 1, frame.f_back + warnings.warn( + "DeviceConfig's use_precalculated_dets and max_memory_gb are " + "deprecated: set use_precalculated_dets and " + "max_memory_gb_for_determinants in sbd_config instead, which " + "takes precedence over them.", + DeprecationWarning, stacklevel=stacklevel, + ) self.use_precalculated_dets = use_precalculated_dets self.max_memory_gb = max_memory_gb @classmethod - def auto(cls, max_memory_gb: int = -1) -> 'DeviceConfig': + def auto(cls, max_memory_gb: int | None = None) -> 'DeviceConfig': """ Auto-detect the best available backend and use it. @@ -75,7 +101,7 @@ def auto(cls, max_memory_gb: int = -1) -> 'DeviceConfig': ``sbd.get_backend('auto')`` uses. Args: - max_memory_gb: Maximum GPU memory in GB (-1 = auto) + max_memory_gb: Deprecated; see :class:`DeviceConfig`. Returns: DeviceConfig for the best backend that is both built and runnable @@ -113,8 +139,8 @@ def cpu(cls) -> 'DeviceConfig': return cls(device='cpu') @classmethod - def gpu(cls, use_precalculated_dets: bool = True, - max_memory_gb: int = -1) -> 'DeviceConfig': + def gpu(cls, use_precalculated_dets: bool | None = None, + max_memory_gb: int | None = None) -> 'DeviceConfig': """Force NVHPC Thrust GPU execution. **NVIDIA only.** Requires SBD compiled with THRUST (the ``_core_gpu_thrust`` extension, @@ -129,7 +155,7 @@ def gpu(cls, use_precalculated_dets: bool = True, max_memory_gb=max_memory_gb) @classmethod - def gpu_omp(cls, max_memory_gb: int = -1) -> 'DeviceConfig': + def gpu_omp(cls, max_memory_gb: int | None = None) -> 'DeviceConfig': """Force OpenMP target-offload GPU execution. Works on NVIDIA **and AMD**. Requires SBD compiled with the OMP-offload backend (the @@ -154,7 +180,7 @@ def gpu_omp(cls, max_memory_gb: int = -1) -> 'DeviceConfig': return cls(device='gpu-omp', max_memory_gb=max_memory_gb) @classmethod - def gpu_nvidia_omp(cls, max_memory_gb: int = -1) -> 'DeviceConfig': + def gpu_nvidia_omp(cls, max_memory_gb: int | None = None) -> 'DeviceConfig': """Deprecated alias for :meth:`gpu_omp`. The old name dates from when SBD shipped a LLVM-with-NVPTX @@ -219,7 +245,11 @@ def _check_hip(cls) -> bool: def apply(self, sbd_config) -> None: """ - Apply device configuration to an SBD TPB_SBD configuration object. + Apply the deprecated Thrust arguments, if given, to a TPB_SBD object. + + Only the arguments that were passed are written; one left at ``None`` + does not touch the field. The solver wrappers call this before applying + ``sbd_config``, so a value there takes precedence. Args: sbd_config: sbd.TPB_SBD configuration object @@ -229,10 +259,16 @@ def apply(self, sbd_config) -> None: # sbdiag.h). The OMP-offload backend doesn't expose them — that # path doesn't have the precalculated-dets cache yet, so there's # nothing to wire up. - if self.device == 'gpu': + given = { + name: value for name, value in ( + ("use_precalculated_dets", self.use_precalculated_dets), + ("max_memory_gb_for_determinants", self.max_memory_gb), + ) if value is not None + } + if self.device == 'gpu' and given: try: - sbd_config.use_precalculated_dets = self.use_precalculated_dets - sbd_config.max_memory_gb_for_determinants = self.max_memory_gb + for name, value in given.items(): + setattr(sbd_config, name, value) except AttributeError: print("WARNING: Thrust GPU knobs not exposed by this " "backend module — likely a CPU/OMP-offload build " @@ -242,6 +278,8 @@ def apply(self, sbd_config) -> None: def __repr__(self) -> str: if self.device == 'cpu': return "DeviceConfig(CPU)" + if self.use_precalculated_dets is None and self.max_memory_gb is None: + return f"DeviceConfig(device={self.device!r})" return (f"DeviceConfig(device={self.device!r}, " f"precalc_dets={self.use_precalculated_dets}, " f"max_mem={self.max_memory_gb}GB)") diff --git a/python/sbd_solver.py b/python/sbd_solver.py index c507617..ef51a9c 100644 --- a/python/sbd_solver.py +++ b/python/sbd_solver.py @@ -790,14 +790,19 @@ def _create_sbd_config(config_dict: dict | None = None, backend=None, device_con sbd_data.threshold = 1e-4 sbd_data.bit_length = SBD_DEFAULT_BIT_LENGTH + # The device config first, then sbd_config, so a value given in sbd_config is + # the one SBD gets. Applied last, DeviceConfig used to reset sbd_config's + # use_precalculated_dets and max_memory_gb_for_determinants to its own + # defaults on the Thrust backend, which silently disabled both SQD drivers' + # flags for them. It now writes only the (deprecated) arguments it was given. + if device_config is not None: + device_config.apply(sbd_data) + if config_dict: for key, value in config_dict.items(): if hasattr(sbd_data, key): setattr(sbd_data, key, value) - if device_config is not None: - device_config.apply(sbd_data) - return sbd_data diff --git a/releasenotes/notes/thrust-options-reach-sbd-8c1f4e2a96d07b35.yaml b/releasenotes/notes/thrust-options-reach-sbd-8c1f4e2a96d07b35.yaml new file mode 100644 index 0000000..b14df48 --- /dev/null +++ b/releasenotes/notes/thrust-options-reach-sbd-8c1f4e2a96d07b35.yaml @@ -0,0 +1,21 @@ +--- +features: + - | + ``TPB_SBD`` gains ``thrust_collapse_loops`` on the Thrust backend, upstream SBD's + switch between flattened (the default) and nested GPU loops when + ``use_precalculated_dets`` is on. The TPB drivers expose it as + ``--thrust_collapse_loops`` and ``--sbd_thrust_collapse_loops``. +fixes: + - | + On the Thrust backend, ``use_precalculated_dets`` and + ``max_memory_gb_for_determinants`` given in ``sbd_config`` were reset to their + defaults whenever a ``DeviceConfig`` was passed, so ``run_sqd_sbd.py`` and + ``run_sqd_enlarge_subspace_sbd.py`` ignored ``--sbd_use_precalculated_dets`` and + ``--sbd_max_memory_gb_for_determinants``. A value in ``sbd_config`` now takes + precedence. +deprecations: + - | + The ``use_precalculated_dets`` and ``max_memory_gb`` arguments of ``DeviceConfig`` + and its constructors are deprecated: set ``use_precalculated_dets`` and + ``max_memory_gb_for_determinants`` in ``sbd_config`` instead. They still apply + when given, below anything set in ``sbd_config``. diff --git a/test/test_thrust_options.py b/test/test_thrust_options.py new file mode 100644 index 0000000..51538f7 --- /dev/null +++ b/test/test_thrust_options.py @@ -0,0 +1,174 @@ +# 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. + +"""The Thrust-only TPB options reach SBD as given, whichever way they are set. + +``use_precalculated_dets``, ``max_memory_gb_for_determinants`` and +``thrust_collapse_loops`` exist on the Thrust backend only. They are set through +``sbd_config``; ``DeviceConfig`` used to carry its own copies of the first two and +apply them after ``sbd_config``, which reset whatever the caller had set there, and +with it both SQD drivers' flags. These check the precedence, the deprecated +``DeviceConfig`` route, what the two SQD drivers hand SBD, and that +``thrust_collapse_loops`` does not change the energy. +""" + +from __future__ import annotations + +import os +import pathlib +import subprocess +import sys +import warnings + +import pytest + +import sbd +from sbd import sbd_solver +from sbd.device_config import DeviceConfig + +REPO = pathlib.Path(__file__).resolve().parents[1] +FIELDS = ("use_precalculated_dets", "max_memory_gb_for_determinants", + "thrust_collapse_loops") +NON_DEFAULT = {"use_precalculated_dets": False, "max_memory_gb_for_determinants": 7, + "thrust_collapse_loops": False} +# One word per determinant, as the rest of the suite's in-process h2o cases. +BIT_LENGTH = 63 + + +def _thrust(backend): + if not backend.__name__.endswith("_core_gpu_thrust"): + pytest.skip("these fields exist on the gpu (Thrust) backend only") + + +def _fields(config): + return {name: getattr(config, name) for name in FIELDS} + + +def test_the_fields_are_on_the_thrust_backend_only(backend): + config = backend.TPB_SBD() + on_thrust = backend.__name__.endswith("_core_gpu_thrust") + assert all(hasattr(config, name) is on_thrust for name in FIELDS) + + +def test_the_defaults_are_upstreams(backend): + _thrust(backend) + assert _fields(sbd_solver._create_sbd_config(None, backend, DeviceConfig.gpu())) == { + "use_precalculated_dets": True, "max_memory_gb_for_determinants": -1, + "thrust_collapse_loops": True} + + +@pytest.mark.parametrize("device_config", [None, "gpu()", "gpu-key"]) +def test_sbd_config_reaches_sbd(backend, device_config): + """With or without a DeviceConfig, sbd_config's values are what SBD gets.""" + _thrust(backend) + device_config = {None: None, "gpu()": DeviceConfig.gpu(), + "gpu-key": DeviceConfig(device="gpu")}[device_config] + config = sbd_solver._create_sbd_config(dict(NON_DEFAULT), backend, device_config) + assert _fields(config) == NON_DEFAULT + + +def test_the_deprecated_device_config_arguments_still_apply(backend): + _thrust(backend) + with pytest.warns(DeprecationWarning, match="sbd_config"): + device_config = DeviceConfig.gpu(use_precalculated_dets=False, max_memory_gb=5) + config = sbd_solver._create_sbd_config(None, backend, device_config) + assert (config.use_precalculated_dets, config.max_memory_gb_for_determinants) == (False, 5) + # ... but sbd_config takes precedence over them. + config = sbd_solver._create_sbd_config( + {"use_precalculated_dets": True, "max_memory_gb_for_determinants": 9}, + backend, device_config) + assert (config.use_precalculated_dets, config.max_memory_gb_for_determinants) == (True, 9) + + +@pytest.mark.parametrize("make", [ + lambda: DeviceConfig(device="gpu", max_memory_gb=5), + lambda: DeviceConfig.gpu(max_memory_gb=5), + lambda: DeviceConfig.gpu_omp(max_memory_gb=5), + lambda: DeviceConfig.gpu_nvidia_omp(max_memory_gb=5), +]) +def test_the_deprecation_warning_names_the_callers_line(make): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + make() + deprecations = [w for w in caught if "sbd_config" in str(w.message)] + assert len(deprecations) == 1 and deprecations[0].filename == __file__ + + +def test_a_device_config_without_them_does_not_warn(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + DeviceConfig.gpu(), DeviceConfig.gpu_omp(), DeviceConfig.cpu() + DeviceConfig(device="gpu") + + +def test_collapse_loops_does_not_change_the_energy(backend, data_dir): + _thrust(backend) + fcidump = sbd.LoadFCIDump(str(data_dir / "h2o" / "fcidump.txt")) + norb = int(fcidump.header["NORB"]) + alpha = (data_dir / "h2o" / "h2o-1em3-alpha.txt").read_text().split()[:16] + half = backend.sort_bitarray([backend.from_string(s, BIT_LENGTH, norb) for s in alpha]) + + def energy(collapse): + config = backend.TPB_SBD() + config.eps, config.max_it, config.max_nb, config.bit_length = 1e-10, 200, 8, BIT_LENGTH + config.use_precalculated_dets, config.thrust_collapse_loops = True, collapse + return sbd.tpb_diag(fcidump, half, half, config)["energy"] + + assert energy(False) == pytest.approx(energy(True), abs=1e-9) + + +# --- what the two SQD drivers hand SBD ----------------------------------------------- + +# Runs a driver with _create_sbd_config wrapped, prints the first config it builds, and +# stops there: what is under test is the config, not the solve. +SPY = r""" +import runpy, sys +from sbd import sbd_solver +fields = ("use_precalculated_dets", "max_memory_gb_for_determinants", + "thrust_collapse_loops") +real = sbd_solver._create_sbd_config +def spy(*args, **kwargs): + config = real(*args, **kwargs) + print("CONFIG " + repr({name: getattr(config, name) for name in fields}), flush=True) + raise SystemExit(0) +sbd_solver._create_sbd_config = spy +script = sys.argv[1] +sys.argv = [script] + sys.argv[2:] +runpy.run_path(script, run_name="__main__") +""" + +DRIVERS = ("run_sqd_sbd.py", "run_sqd_enlarge_subspace_sbd.py") + + +def _driver_config(script, flags): + tpb = REPO / "examples" / "tpb" + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run( + [sys.executable, "-c", SPY, script, "--device", "gpu", + "--fcidump", str(REPO / "vendor" / "sbd-upstream" / "data" / "h2o" / "fcidump.txt"), + "--counts", "count_dict_h2o.json", "--max_iterations", "1", *flags], + cwd=tpb, capture_output=True, text=True, timeout=600, env=env) + lines = [l for l in done.stdout.splitlines() if l.startswith("CONFIG ")] + assert lines, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + return eval(lines[0][len("CONFIG "):]) + + +@pytest.mark.parametrize("script", DRIVERS) +def test_the_sqd_drivers_flags_reach_sbd(backend, script): + _thrust(backend) + assert _driver_config(script, []) == { + "use_precalculated_dets": True, "max_memory_gb_for_determinants": -1, + "thrust_collapse_loops": True} + assert _driver_config(script, [ + "--sbd_use_precalculated_dets", "0", "--sbd_max_memory_gb_for_determinants", "7", + "--sbd_thrust_collapse_loops", "0"]) == NON_DEFAULT From 96459f1f9dcc24e565e5324a4e1b436f29631cf5 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Mon, 2 Nov 2026 15:34:09 -0500 Subject: [PATCH 02/16] Refuse GDB on the OpenMP-offload backend Upstream SBD has no OpenMP-offload GDB code, so gdb_diag on the gpu-omp module ran entirely on the host while a GPU had been asked for; only run_gdb_diag.py warned. gdb_diag now raises ValueError on that module, and run_gdb_diag.py and run_gdb_heatbath.py exit with an error for --device gpu-omp, or auto resolving to it, pointing at cpu, or gpu (Thrust, NVIDIA GPUs only). Co-Authored-By: Claude Opus 5.5 --- docs/index.rst | 5 ++-- examples/gdb/README.md | 17 +++++------ examples/gdb/run_gdb_diag.py | 23 +++++++++------ examples/gdb/run_gdb_heatbath.py | 19 +++++++++++-- python/__init__.py | 10 +++---- python/bindings.cpp | 13 +++++++++ ...uses-offload-backend-4e7b0c19a2d6f853.yaml | 9 ++++++ test/test_gdb_drivers.py | 14 ++++++++++ test/test_gdb_equivalence.py | 28 ++++++++++++++++++- 9 files changed, 111 insertions(+), 27 deletions(-) create mode 100644 releasenotes/notes/gdb-refuses-offload-backend-4e7b0c19a2d6f853.yaml diff --git a/docs/index.rst b/docs/index.rst index fdc0aed..3d29027 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -84,8 +84,9 @@ call, with ``device=`` or a :class:`~sbd.device_config.DeviceConfig`: - NVIDIA GPUs, Thrust - NVIDIA HPC SDK (``nvc++``) * - ``gpu-omp`` - - NVIDIA or AMD GPUs, OpenMP offload. TPB only: a GDB run under ``gpu-omp`` - executes on the host, so use ``gpu`` for GDB on GPUs + - NVIDIA or AMD GPUs, OpenMP offload. TPB only: ``gdb_diag`` refuses this + backend, as upstream has no OpenMP-offload GDB code; use ``cpu``, or ``gpu`` + for GDB on NVIDIA GPUs - ``nvc++`` on NVIDIA, ``amdclang++`` on AMD The GPU backends pass device memory to MPI, so they expect a GPU-aware MPI. The diff --git a/examples/gdb/README.md b/examples/gdb/README.md index 6130254..7b9562b 100644 --- a/examples/gdb/README.md +++ b/examples/gdb/README.md @@ -138,7 +138,7 @@ Solver, MPI and output — the same meanings as in `run_gdb_diag.py`: | Parameter | What it controls | Default | |---|---|---| -| `--device` | `cpu` or `gpu` (Thrust) are the two real choices; `gpu-omp` and `auto` run GDB on the host, see [Choosing a backend](#choosing-a-backend) | `cpu` | +| `--device` | `cpu`, or `gpu` (Thrust, NVIDIA only). `gpu-omp` is refused, and so is `auto` when it resolves to `gpu-omp`; see [Choosing a backend](#choosing-a-backend) | `cpu` | | `--method` | 0=Davidson, 1=Davidson storing the Hamiltonian. GDB has no Lanczos | `0` | | `--tolerance` / `--iteration` / `--block` | Davidson residual tolerance, iteration cap, and basis-vector count. Also accepted as `--eps` / `--max_it` / `--max_nb`, matching SBD's own names | `1e-6` / `30` / `10` | | `--bit_length` | Bits per packed word: even, at most 62. 64 is undefined behavior in SBD's redistribution. An odd value breaks SBD's alpha/beta conversion once a determinant spans two words, so the heatbath expansion crashes or returns a wrong energy | `62` | @@ -439,10 +439,12 @@ already count-balanced. --device gpu # NVHPC Thrust, NVIDIA only -- the only GPU backend with GDB kernels ``` -`gpu-omp` and `auto` are accepted too, but for GDB neither is a GPU path: `gpu-omp` has -no GDB kernels and silently diagonalizes on the host (the drivers warn), and `auto` -prefers Thrust but falls back to `gpu-omp` where Thrust is not built — landing on the -host as well. **For GDB, pick `cpu` or `gpu` explicitly.** +`gpu-omp` is **refused**: upstream SBD has no OpenMP-offload GDB kernels, so it would +diagonalize on the host while a GPU was asked for. `sbd.gdb_diag` raises `ValueError` +on that backend, and both drivers exit with an error before reading any input. `auto` +prefers Thrust but falls back to `gpu-omp` where Thrust is not built, and is refused +when it does. **For GDB, pick `cpu` or `gpu` explicitly.** On AMD GPUs that means +`cpu`: there is no Thrust build for AMD. `sbd.available_backends()` reports what this install actually built — a static scan, so it is safe to call outside `mpirun` — and `sbd.loaded_backends()` reports what the @@ -458,9 +460,8 @@ not route device work through OpenMP. ### On GPUs GDB has **Thrust kernels only** (`--device gpu`, NVIDIA). There is no `omp target` -code under `include/sbd/chemistry/gdb/` at all, so under `--device gpu-omp` a GDB run -pins a device and then diagonalizes on the host; the driver warns when it resolves -there. AMD has no Thrust build, so AMD GDB is CPU-only. +code under `include/sbd/chemistry/gdb/` at all, so `--device gpu-omp` is refused +rather than run on the host. AMD has no Thrust build, so AMD GDB is CPU-only. The Thrust path additionally requires **`helper == 1`** (`gdb/mult_thrust.h:310-314`, checked on every kernel launch), so every rank must go to `t_comm_size × b_comm_size`. diff --git a/examples/gdb/run_gdb_diag.py b/examples/gdb/run_gdb_diag.py index da857a1..6bc7a3d 100755 --- a/examples/gdb/run_gdb_diag.py +++ b/examples/gdb/run_gdb_diag.py @@ -168,9 +168,10 @@ def parse_args(): parser.add_argument('--device', default='cpu', choices=['cpu', 'gpu', 'gpu-omp', 'auto'], - help="Compute backend. 'gpu' is NVHPC Thrust, which is the " - "only GPU backend with GDB kernels; 'gpu-omp' has none, " - "so GDB runs on the host there") + help="Compute backend: 'cpu', or 'gpu' (NVHPC Thrust, NVIDIA " + "only), the one GPU backend with GDB kernels. 'gpu-omp' is " + "refused, as upstream has no OpenMP-offload GDB code; 'auto' " + "is refused too if it resolves to 'gpu-omp'") # --- input ------------------------------------------------------------ parser.add_argument('--skip-weight-check', action='store_true', @@ -357,6 +358,17 @@ def main(): rank = sbd.get_rank() size = sbd.get_world_size() device = sbd.get_device() + # gdb_diag refuses the OpenMP-offload backend itself; refusing here first says + # so before any input is read, and covers --device auto resolving to it. Every + # rank resolves the same device, so every rank exits. + if device == 'gpu-omp': + if rank == 0: + print(f"ERROR: GDB does not run on the OpenMP-offload backend " + f"(--device {args.device} -> 'gpu-omp'): upstream SBD has no " + f"OpenMP-offload GDB kernels. Use --device cpu, or --device gpu " + f"(NVHPC Thrust, NVIDIA GPUs only; there is no GDB GPU path for AMD " + f"GPUs).", file=sys.stderr) + return 2 if rank == 0: print("=" * 70) @@ -490,11 +502,6 @@ def main(): print(" ERROR: GDB on the Thrust backend requires helper == 1, i.e. " "t_comm_size * b_comm_size == ranks. Give every rank to " "--b_comm_size (with --t_comm_size <= it).") - if device == 'gpu-omp': - print(" WARNING: the OMP-offload backend has no GDB kernels -- there is " - "no `omp target` code under include/sbd/chemistry/gdb/ at all. " - "This run pins a device and then diagonalizes on the host. Use " - "--device gpu (Thrust, NVIDIA only) for GDB on a GPU.") print() print("Subspace:") print(f" Source: {source}") diff --git a/examples/gdb/run_gdb_heatbath.py b/examples/gdb/run_gdb_heatbath.py index 2f9b744..c7bbacc 100755 --- a/examples/gdb/run_gdb_heatbath.py +++ b/examples/gdb/run_gdb_heatbath.py @@ -81,9 +81,12 @@ def parse_args(): parser.add_argument('--device', default='cpu', choices=['cpu', 'gpu', 'gpu-omp', 'auto'], - help="Compute backend. 'gpu' (Thrust) is the only GPU backend " - "with GDB kernels, and it requires the helper dimension " - "to be 1, i.e. t_comm_size * b_comm_size == ranks") + help="Compute backend: 'cpu', or 'gpu' (NVHPC Thrust, NVIDIA " + "only), the one GPU backend with GDB kernels, which requires " + "the helper dimension to be 1, i.e. t_comm_size * " + "b_comm_size == ranks. 'gpu-omp' is refused, as upstream has " + "no OpenMP-offload GDB code; 'auto' is refused too if it " + "resolves to 'gpu-omp'") parser.add_argument('--fcidump', default=_DEFAULT_FCIDUMP, help='FCIDUMP file defining the Hamiltonian') @@ -411,6 +414,16 @@ def main(): comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() + # As in run_gdb_diag.py: gdb_diag refuses this backend itself, and this says so + # before the seed is built, on every rank, --device auto included. + if sbd.get_device() == 'gpu-omp': + if rank == 0: + print(f"ERROR: GDB does not run on the OpenMP-offload backend " + f"(--device {args.device} -> 'gpu-omp'): upstream SBD has no " + f"OpenMP-offload GDB kernels. Use --device cpu, or --device gpu " + f"(NVHPC Thrust, NVIDIA GPUs only; there is no GDB GPU path for AMD " + f"GPUs).", file=sys.stderr) + return 2 fcidump = backend.LoadFCIDump(args.fcidump) header = fcidump.header diff --git a/python/__init__.py b/python/__init__.py index e0f1c7e..ee77754 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -797,11 +797,11 @@ def gdb_diag(fcidump, det, sbd_data, writes one file per b_comm position, ``f"{savename}{rank_b:06d}.bin"``, each holding only that shard; rank 0's file is not the whole wavefunction. GDB has no combined matrix-form dump. - device: Backend to run on: 'cpu', 'gpu' (Thrust), 'gpu-omp' (OpenMP - offload), or an alias accepted by :func:`init`. None uses the - default; see :func:`get_backend`. 'gpu-omp' is accepted, but SBD - has no OpenMP-offload GDB code, so the run executes on the host; - use 'gpu' for GDB on GPUs. + device: Backend to run on: 'cpu', 'gpu' (Thrust, NVIDIA GPUs only), or + an alias accepted by :func:`init`. None uses the default; see + :func:`get_backend`. 'gpu-omp' (OpenMP offload) is refused with + ``ValueError``: upstream SBD has no OpenMP-offload GDB code, so it + would run on the host. There is no GDB GPU path for AMD GPUs. determinant_distribution: How to place determinants across b_comm — one of ``input`` (keep the shards as given), ``equal-bra-a`` (default; equal distinct-alpha count per rank, which is what balances the diff --git a/python/bindings.cpp b/python/bindings.cpp index d88f275..6844c13 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -719,6 +719,19 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { int determinant_grid_a, int determinant_grid_b) { +#ifdef USE_OMP_OFFLOAD + // Upstream has no OpenMP-offload GDB code -- there is no `omp target` + // anywhere under include/sbd/chemistry/gdb/ -- so this build would run + // GDB entirely on the host while the caller asked for a GPU. Refuse + // rather than do that. Every rank loads the same module, so every rank + // raises here and none is left waiting in a collective. + throw std::invalid_argument( + "gdb_diag is not supported on the OpenMP-offload backend ('gpu-omp'): " + "upstream SBD has no OpenMP-offload GDB kernels, so it would run on " + "the CPU. Use device='cpu', or device='gpu' (NVHPC Thrust, NVIDIA GPUs " + "only; there is no GDB GPU path for AMD GPUs)."); +#endif + MPI_Comm comm = get_mpi_comm(py_comm); int mpi_rank; MPI_Comm_rank(comm, &mpi_rank); int mpi_size; MPI_Comm_size(comm, &mpi_size); diff --git a/releasenotes/notes/gdb-refuses-offload-backend-4e7b0c19a2d6f853.yaml b/releasenotes/notes/gdb-refuses-offload-backend-4e7b0c19a2d6f853.yaml new file mode 100644 index 0000000..9a58db9 --- /dev/null +++ b/releasenotes/notes/gdb-refuses-offload-backend-4e7b0c19a2d6f853.yaml @@ -0,0 +1,9 @@ +--- +upgrade: + - | + ``gdb_diag`` now refuses the OpenMP-offload backend (``gpu-omp``) with + ``ValueError``, and ``run_gdb_diag.py`` and ``run_gdb_heatbath.py`` exit with an + error for ``--device gpu-omp``, or ``--device auto`` when it resolves to it. + Upstream SBD has no OpenMP-offload GDB code, so these used to run on the CPU + while a GPU had been asked for. Use ``cpu``, or ``gpu`` (Thrust) on NVIDIA GPUs; + on AMD GPUs, GDB is CPU-only. diff --git a/test/test_gdb_drivers.py b/test/test_gdb_drivers.py index 4f49a2c..358c92b 100644 --- a/test/test_gdb_drivers.py +++ b/test/test_gdb_drivers.py @@ -356,6 +356,20 @@ def test_drivers_refuse_an_unsafe_word_size(script, bit_length): assert "--bit_length must be even and between 2 and 62" in completed.stderr +@pytest.mark.parametrize("script", ["run_gdb_diag.py", "run_gdb_heatbath.py"]) +def test_drivers_refuse_the_offload_backend(script): + """GDB has no OpenMP-offload kernels upstream: refused, pointing at cpu and gpu.""" + import sbd + + if "gpu-omp" not in sbd.available_backends(): + pytest.skip("the OpenMP-offload backend was not built") + completed = _run(script, ["--device", "gpu-omp"]) + assert completed.returncode != 0 + assert "OpenMP-offload" in completed.stderr + assert "--device cpu" in completed.stderr and "--device gpu" in completed.stderr + assert "Ground state energy" not in completed.stdout + + def test_skip_weight_check_bypasses_the_validation(tmp_path): """The escape hatch must skip the check, not merely survive it. diff --git a/test/test_gdb_equivalence.py b/test/test_gdb_equivalence.py index 9b475ac..6d8d288 100644 --- a/test/test_gdb_equivalence.py +++ b/test/test_gdb_equivalence.py @@ -145,6 +145,11 @@ def _is_thrust(backend): return "thrust" in getattr(backend, "__name__", "") +def _is_offload(backend): + """Whether ``backend`` is the OpenMP-offload build, on which gdb_diag is refused.""" + return getattr(backend, "__name__", "").endswith("_core_gpu_omp_offload") + + _THRUST_HELPER = "requires h_comm_size == 1" @@ -166,7 +171,13 @@ def _uneven_shard(full, b_comm_size, rank): @pytest.fixture(scope="module") def h2o(data_dir, backend): - """The h2o FCIDUMP, its orbital count, the alpha-list path, and a truncation.""" + """The h2o FCIDUMP, its orbital count, the alpha-list path, and a truncation. + + Every solving test here takes this fixture, so this is where they skip on the + OpenMP-offload backend, which refuses gdb_diag (see the refusal test below). + """ + if _is_offload(backend): + pytest.skip("gdb_diag is refused on the OpenMP-offload backend") molecule_dir = data_dir / "h2o" fcidump = backend.LoadFCIDump(str(molecule_dir / "fcidump.txt")) norb = int(fcidump.header["NORB"]) @@ -729,3 +740,18 @@ def test_gdb_fe4s4_at_realistic_size(): _, energy, _, dim = line[-1].split() assert int(dim) == 59536 assert float(energy) == pytest.approx(FE4S4_ENERGY, abs=1e-8) + + +def test_gdb_is_refused_on_the_offload_backend(backend, data_dir): + """Upstream has no OpenMP-offload GDB code, so gpu-omp refuses rather than run on the host.""" + if not _is_offload(backend): + pytest.skip("the OpenMP-offload backend only (SBD_TEST_DEVICE=gpu-omp)") + fcidump = backend.LoadFCIDump(str(data_dir / "h2o" / "fcidump.txt")) + norb = int(fcidump.header["NORB"]) + alpha = _alpha_strings(data_dir / "h2o" / "h2o-1em3-alpha.txt", 2) + det = sbd.from_strings([_interleave(a, b) for a in alpha for b in alpha], + BIT_LENGTH, 2 * norb) + config = backend.GDB_SBD() + config.bit_length = BIT_LENGTH + with pytest.raises(ValueError, match=r"OpenMP-offload.*device='cpu'.*device='gpu'"): + sbd.gdb_diag(fcidump, det, config, device="gpu-omp") From 633bf7dda6617071582e9f1595aea249ca9e6187 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 30 Oct 2026 16:06:06 -0400 Subject: [PATCH 03/16] Update vendored SBD to 95bf45a; refuse one-thread Thrust solves; bind cpu_subspace - vendor/sbd-upstream moves from 93ebabe to r-ccs-cms/sbd 95bf45a (25 commits): the opposite-spin 2-RDM sign fix (#98), Davidson/Lanczos subspace refactor and TPB performance work (#95, #99), the CAOP stat evaluator (#94) and the CAOP spectrum app (#96, #100). - Upstream's Thrust builds now MPI_Abort the job when omp_get_max_threads() is 1, which ends Python with nothing to catch. tpb_diag, tpb_diag_from_files and gdb_diag on Thrust raise ValueError on every rank first instead. - TPB_SBD.cpu_subspace and GDB_SBD.cpu_subspace (Thrust only, as upstream declares them): keep the Davidson (and TPB's Lanczos) subspace vectors in host memory instead of GPU memory, as upstream's apps' --cpu_subspace. - Tests: test_rdm_open_shell.py, the opposite-spin 2-RDM block on upstream's H3 doublet against PySCF, which fails at the old pin on every backend and passes now; test_openmp_threads.py, the one-thread refusal on Thrust; test_cpu_subspace.py, the same energy with cpu_subspace on and off for TPB and GDB, and the field present on the Thrust backend only. Also expose it in the drivers as --cpu_subspace / --sbd_cpu_subspace; it is ignored on the backends that do not have it. Co-Authored-By: Claude Opus 5.5 --- docs/apidocs/sbd.rst | 8 ++ examples/gdb/run_gdb_diag.py | 7 + examples/gdb/run_gdb_heatbath.py | 6 + examples/tpb/README.md | 1 + examples/tpb/run_sbd_diag.py | 6 + examples/tpb/run_sqd_enlarge_subspace_sbd.py | 4 + examples/tpb/run_sqd_sbd.py | 5 + python/bindings.cpp | 41 ++++++ ...sbd-upstream-95bf45a-4d7e2a19c0b3f856.yaml | 23 ++++ test/test_cpu_subspace.py | 122 ++++++++++++++++++ test/test_openmp_threads.py | 65 ++++++++++ test/test_rdm_open_shell.py | 90 +++++++++++++ vendor/sbd-upstream | 2 +- 13 files changed, 379 insertions(+), 1 deletion(-) create mode 100644 releasenotes/notes/sbd-upstream-95bf45a-4d7e2a19c0b3f856.yaml create mode 100644 test/test_cpu_subspace.py create mode 100644 test/test_openmp_threads.py create mode 100644 test/test_rdm_open_shell.py diff --git a/docs/apidocs/sbd.rst b/docs/apidocs/sbd.rst index 7a3d493..91cc22c 100644 --- a/docs/apidocs/sbd.rst +++ b/docs/apidocs/sbd.rst @@ -117,6 +117,11 @@ Fields of ``TPB_SBD`` only: (the default), each task's loops run as one flattened GPU loop; off, as nested loops, which need less GPU memory. Upstream turns it off by itself when the flattened size would overflow. Energies do not change. + * - ``cpu_subspace`` + - Thrust backend only. Keep the Davidson or Lanczos subspace vectors in host + memory instead of GPU memory, so the GPU memory they use no longer grows with + ``max_nb``; the Hamiltonian is still applied on the GPU. Energies agree with + the default to within rounding. Fields of ``GDB_SBD`` only: @@ -149,6 +154,9 @@ Fields of ``GDB_SBD`` only: ``2*i + 1`` beta orbital ``i``, as in :func:`from_strings`), instead of ``init``. Empty (the default) uses ``init``. It must occur exactly once across all shards, or the solve raises ``ValueError``. + * - ``cpu_subspace`` + - Thrust backend only, as for ``TPB_SBD``: keep the Davidson subspace vectors in + host memory instead of GPU memory. Inputs ====== diff --git a/examples/gdb/run_gdb_diag.py b/examples/gdb/run_gdb_diag.py index 6bc7a3d..6ea2033 100755 --- a/examples/gdb/run_gdb_diag.py +++ b/examples/gdb/run_gdb_diag.py @@ -236,6 +236,10 @@ def parse_args(): dest='eps', help='Convergence tolerance') parser.add_argument('--max_time', type=float, default=1e10, help='Maximum wall time in seconds') + parser.add_argument('--cpu_subspace', type=int, default=0, choices=[0, 1], + help='Keep the Davidson subspace vectors in CPU memory instead of ' + 'GPU memory (Thrust backend only; ignored elsewhere). For runs ' + 'too large for GPU memory. Needs OMP_NUM_THREADS of 2 or more') parser.add_argument('--init', type=int, default=0, help='Initial vector policy') parser.add_argument('--seed', type=int, default=1729, @@ -470,6 +474,9 @@ def main(): config.max_nb = args.max_nb config.eps = args.eps config.max_time = args.max_time + # Thrust only; the other backends' GDB_SBD has no such field and needs none. + if hasattr(config, 'cpu_subspace'): + config.cpu_subspace = bool(args.cpu_subspace) config.init = args.init config.seed = args.seed config.initial_determinant_bitstring = args.initial_det diff --git a/examples/gdb/run_gdb_heatbath.py b/examples/gdb/run_gdb_heatbath.py index c7bbacc..d731a84 100755 --- a/examples/gdb/run_gdb_heatbath.py +++ b/examples/gdb/run_gdb_heatbath.py @@ -146,6 +146,10 @@ def parse_args(): dest='max_it', help='Maximum Davidson iterations per round') parser.add_argument('--block', '--max_nb', type=int, default=10, dest='max_nb', help='Maximum number of basis vectors') + parser.add_argument('--cpu_subspace', type=int, default=0, choices=[0, 1], + help='Keep the Davidson subspace vectors in CPU memory instead of ' + 'GPU memory (Thrust backend only; ignored elsewhere). For runs ' + 'too large for GPU memory. Needs OMP_NUM_THREADS of 2 or more') parser.add_argument('--bit_length', type=int, default=62, help='Bits per packed word: even, at most 62') parser.add_argument('--carryover_type', type=int, default=2, choices=[2, 3], @@ -524,6 +528,8 @@ def global_dim(local): config.eps = args.eps config.max_it = args.max_it config.max_nb = args.max_nb + if hasattr(config, 'cpu_subspace'): # Thrust only + config.cpu_subspace = bool(args.cpu_subspace) config.bit_length = args.bit_length config.b_comm_size = args.b_comm_size config.t_comm_size = args.t_comm_size diff --git a/examples/tpb/README.md b/examples/tpb/README.md index 0020366..390f4fb 100644 --- a/examples/tpb/README.md +++ b/examples/tpb/README.md @@ -285,6 +285,7 @@ Per diagonalization, not per loop: | `--sbd_use_precalculated_dets` | Thrust only. `1` precomputes a determinant index for **every** (α,β) pair — the whole subspace, on the GPU. `0` uses per-thread storage: slower per matvec, far less memory | `1` | | `--sbd_max_memory_gb_for_determinants` | Thrust only, and **only consulted when `--sbd_use_precalculated_dets 0`** (`mult_thrust.h:257-273`). Caps the per-thread buffer in GB | `-1` (uncapped) | | `--sbd_thrust_collapse_loops` | Thrust only, and **only consulted when `--sbd_use_precalculated_dets 1`**. `1` runs each task's loops as one flattened GPU loop; `0` as nested loops, which need less GPU memory. Upstream switches to `0` by itself if the flattened size would overflow. Energies are unchanged | `1` | +| `--sbd_cpu_subspace` | Thrust only. `1` keeps the Davidson subspace vectors in CPU memory instead of GPU memory; H is still applied on the GPU. For runs whose subspace vectors do not fit in GPU memory. Needs `OMP_NUM_THREADS` of 2 or more; ignored on the other backends | `0` | For reference, upstream's own `TPB_SBD` struct defaults are looser still (`max_it=1`, `eps=1e-4`), and `run_sbd_diag.py` uses `eps=1e-3`. On the h2o counts case, diff --git a/examples/tpb/run_sbd_diag.py b/examples/tpb/run_sbd_diag.py index 4b94a13..b3be090 100644 --- a/examples/tpb/run_sbd_diag.py +++ b/examples/tpb/run_sbd_diag.py @@ -175,6 +175,10 @@ def parse_args(): parser.add_argument('--thrust_collapse_loops', type=int, default=1, choices=[0, 1], help='Thrust only, and only with --use_precalculated_dets 1: 1 runs ' 'the loops flattened (faster), 0 nested (less GPU memory)') + parser.add_argument('--cpu_subspace', type=int, default=0, choices=[0, 1], + help='Keep the Davidson subspace vectors in CPU memory instead of GPU memory ' + '(Thrust backend only; ignored elsewhere). For runs too large ' + 'for GPU memory. Needs OMP_NUM_THREADS of 2 or more') args = parser.parse_args() if args.initial_bdet and not args.initial_adet: @@ -235,6 +239,8 @@ def main(): config.max_memory_gb_for_determinants = args.max_memory_gb_for_determinants if hasattr(config, 'thrust_collapse_loops'): config.thrust_collapse_loops = bool(args.thrust_collapse_loops) + if hasattr(config, 'cpu_subspace'): + config.cpu_subspace = bool(args.cpu_subspace) if rank == 0: print("Configuration:") diff --git a/examples/tpb/run_sqd_enlarge_subspace_sbd.py b/examples/tpb/run_sqd_enlarge_subspace_sbd.py index cc7674e..25e8f0a 100644 --- a/examples/tpb/run_sqd_enlarge_subspace_sbd.py +++ b/examples/tpb/run_sqd_enlarge_subspace_sbd.py @@ -155,6 +155,9 @@ def parse_args(): choices=[0, 1]) sbd.add_argument("--sbd_max_memory_gb_for_determinants", type=int, default=-1) sbd.add_argument("--sbd_thrust_collapse_loops", type=int, default=1, choices=[0, 1]) + sbd.add_argument("--sbd_cpu_subspace", type=int, default=0, choices=[0, 1], + help="Thrust only: keep the Davidson subspace vectors in CPU memory " + "instead of GPU memory. Needs OMP_NUM_THREADS of 2 or more.") sbd.add_argument("--sbd_bit_length", type=int, default=20, dest="bit_length") # ---- MPI decomposition ------------------------------------------------------ @@ -443,6 +446,7 @@ def main(): "use_precalculated_dets": bool(args.sbd_use_precalculated_dets), "max_memory_gb_for_determinants": args.sbd_max_memory_gb_for_determinants, "thrust_collapse_loops": bool(args.sbd_thrust_collapse_loops), + "cpu_subspace": bool(args.sbd_cpu_subspace), "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, } diff --git a/examples/tpb/run_sqd_sbd.py b/examples/tpb/run_sqd_sbd.py index 468a5cc..f986ff1 100644 --- a/examples/tpb/run_sqd_sbd.py +++ b/examples/tpb/run_sqd_sbd.py @@ -212,6 +212,10 @@ def parse_args(): help="Thrust only, and only with --sbd_use_precalculated_dets 1: " "1 runs each task's loops as one flattened GPU loop (faster), " "0 as nested loops (less GPU memory). Energies are unchanged.") + sbd.add_argument("--sbd_cpu_subspace", type=int, default=0, choices=[0, 1], + help="Thrust only: keep the Davidson subspace vectors in CPU memory " + "instead of GPU memory, for runs too large for GPU memory. " + "Needs OMP_NUM_THREADS of 2 or more.") sbd.add_argument("--sbd_bit_length", "--bit_length", type=int, default=20, dest="bit_length", help="Bits packed into each size_t of the bitstring " @@ -490,6 +494,7 @@ def main(): "use_precalculated_dets": bool(args.sbd_use_precalculated_dets), "max_memory_gb_for_determinants": args.sbd_max_memory_gb_for_determinants, "thrust_collapse_loops": bool(args.sbd_thrust_collapse_loops), + "cpu_subspace": bool(args.sbd_cpu_subspace), "adet_comm_size": args.adet_comm_size, "bdet_comm_size": args.bdet_comm_size, "task_comm_size": args.task_comm_size, diff --git a/python/bindings.cpp b/python/bindings.cpp index ded1ddc..778fcff 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -213,6 +213,30 @@ static bool sbd_all_ranks_ok(bool local_ok, MPI_Comm comm) { return all != 0; } +/** + * Refuse a Thrust solve with one OpenMP thread, before upstream aborts the job. + * + * Upstream's Thrust builds keep the Davidson subspace and part of the setup on the + * host and call MPI_Abort when omp_get_max_threads() == 1 (tpb/sbdiag.h, + * basic/subspace_thrust.h, caop/basic/sbdiag.h), which ends the Python process with + * no exception to catch. Raise ValueError on every rank first instead. CPU and + * OpenMP-offload builds are not affected. + */ +static void sbd_require_openmp_threads(const char* where, MPI_Comm comm) { +#ifdef SBD_THRUST + if (!sbd_all_ranks_ok(omp_get_max_threads() > 1, comm)) { + throw std::invalid_argument( + std::string(where) + " on the gpu (Thrust) backend needs more than one " + "OpenMP thread per rank: upstream SBD does part of the work on the host " + "and aborts the job otherwise. Set OMP_NUM_THREADS to 2 or more, and give " + "each rank that many cores (for example srun --cpus-per-task)"); + } +#else + (void)where; + (void)comm; +#endif +} + /** * Order-insensitive-after-sort 64-bit fingerprint of a shard (FNV-1a). * @@ -317,6 +341,13 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { &sbd::tpb::SBD::initial_bdeterminant_bitstring, "Beta half of the explicit start. Empty uses the alpha string " "for beta; it cannot be set without the alpha one") +#ifdef SBD_THRUST + .def_readwrite("cpu_subspace", &sbd::tpb::SBD::cpu_subspace, + "gpu (Thrust) backend only: keep the Davidson/Lanczos subspace " + "vectors in host memory instead of GPU memory, so the GPU memory " + "they use no longer grows with max_nb; H is still applied on the " + "GPU. Energies agree with the default to within rounding") +#endif .def_readwrite("do_shuffle", &sbd::tpb::SBD::do_shuffle, "Shuffle determinants flag") .def_readwrite("do_rdm", &sbd::tpb::SBD::do_rdm, @@ -421,6 +452,13 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { "Weight truncation threshold applied before heatbath expansion") .def_readwrite("heatbath_batch_size", &sbd::gdb::SBD::heatbath_batch_size, "Heatbath expansion batch size") +#ifdef SBD_THRUST + .def_readwrite("cpu_subspace", &sbd::gdb::SBD::cpu_subspace, + "gpu (Thrust) backend only: keep the Davidson subspace vectors in " + "host memory instead of GPU memory, so the GPU memory they use no " + "longer grows with max_nb; H is still applied on the GPU. Energies " + "agree with the default to within rounding") +#endif .def_readwrite("bit_length", &sbd::gdb::SBD::bit_length, "Bit length for determinant representation") ; @@ -565,6 +603,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { // Convert MPI communicator MPI_Comm comm = get_mpi_comm(py_comm); + sbd_require_openmp_threads("tpb_diag", comm); // Get MPI rank for GPU assignment int mpi_rank; @@ -739,6 +778,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { #endif MPI_Comm comm = get_mpi_comm(py_comm); + sbd_require_openmp_threads("gdb_diag", comm); int mpi_rank; MPI_Comm_rank(comm, &mpi_rank); int mpi_size; MPI_Comm_size(comm, &mpi_size); @@ -1189,6 +1229,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { // Convert MPI communicator MPI_Comm comm = get_mpi_comm(py_comm); + sbd_require_openmp_threads("tpb_diag_from_files", comm); // Get MPI rank for GPU assignment int mpi_rank; diff --git a/releasenotes/notes/sbd-upstream-95bf45a-4d7e2a19c0b3f856.yaml b/releasenotes/notes/sbd-upstream-95bf45a-4d7e2a19c0b3f856.yaml new file mode 100644 index 0000000..d8edb81 --- /dev/null +++ b/releasenotes/notes/sbd-upstream-95bf45a-4d7e2a19c0b3f856.yaml @@ -0,0 +1,23 @@ +--- +features: + - | + The vendored SBD is updated to r-ccs-cms/sbd ``95bf45a``. ``TPB_SBD`` and + ``GDB_SBD`` gain ``cpu_subspace`` on the Thrust backend, which keeps the Davidson + (and TPB's Lanczos) subspace vectors in host memory, so large runs need less GPU + memory. The drivers expose it as ``--cpu_subspace`` (``run_sbd_diag.py``, + ``run_gdb_diag.py``, ``run_gdb_heatbath.py``) and ``--sbd_cpu_subspace`` (the two + SQD drivers); the other backends ignore it. +fixes: + - | + The updated SBD fixes the sign of opposite-spin double excitations in the + 2-particle RDM (r-ccs-cms/sbd#98). On open-shell states the opposite-spin + (``ab``) block of ``two_p_rdm`` from ``tpb_diag`` and ``gdb_diag`` was wrong + before; the error cancelled in the spin-summed 2-RDM, so ``assemble_rdms`` and + ``SCIResult.rdm2`` were unaffected. +upgrade: + - | + On the Thrust (``gpu``) backend, upstream SBD now needs more than one OpenMP + thread per rank and otherwise aborts the whole MPI job. ``tpb_diag``, + ``tpb_diag_from_files`` and ``gdb_diag`` on that backend now raise + ``ValueError`` on every rank first when ``OMP_NUM_THREADS`` is 1. Set it to 2 or + more. The CPU and OpenMP-offload backends are not affected. diff --git a/test/test_cpu_subspace.py b/test/test_cpu_subspace.py new file mode 100644 index 0000000..eac46ed --- /dev/null +++ b/test/test_cpu_subspace.py @@ -0,0 +1,122 @@ +# 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. + +"""``cpu_subspace`` on the Thrust backend: the subspace vectors in host memory. + +Upstream SBD's Thrust builds keep the Davidson (and TPB's Lanczos) subspace vectors in +GPU memory by default; ``cpu_subspace`` keeps them in host memory instead, applying H on +the GPU through two GPU buffers. Upstream documents the energies as agreeing with the +default to within rounding, which is what is checked here, for both solvers. The field +exists on the Thrust backend only, as upstream declares it only under ``SBD_THRUST``. +""" + +from __future__ import annotations + +import itertools + +import pytest + +import sbd + +# One word per determinant, as the rest of the suite's in-process h2o cases (SBD fixes the +# word count for the life of the process, so a different one here breaks the tests after it). +BIT_LENGTH = 62 +ALPHA_LIMIT = 16 + + +def _thrust(backend): + if not backend.__name__.endswith("_core_gpu_thrust"): + pytest.skip("cpu_subspace exists on the gpu (Thrust) backend only") + + +def _alpha(data_dir): + return (data_dir / "h2o" / "h2o-1em3-alpha.txt").read_text().split()[:ALPHA_LIMIT] + + +def _interleave(alpha, beta): + a, b = alpha[::-1], beta[::-1] + return "".join(a[i] + b[i] for i in range(len(a)))[::-1] + + +@pytest.mark.parametrize("solver", ["tpb", "gdb"]) +def test_cpu_subspace_gives_the_default_energy(backend, data_dir, solver): + _thrust(backend) + fcidump = sbd.LoadFCIDump(str(data_dir / "h2o" / "fcidump.txt")) + norb = int(fcidump.header["NORB"]) + alpha = _alpha(data_dir) + + def solve(cpu_subspace): + if solver == "tpb": + config = backend.TPB_SBD() + config.eps, config.max_it, config.max_nb, config.bit_length = 1e-10, 200, 8, BIT_LENGTH + config.cpu_subspace = cpu_subspace + half = backend.sort_bitarray( + [backend.from_string(s, BIT_LENGTH, norb) for s in alpha]) + return sbd.tpb_diag(fcidump, half, half, config)["energy"] + config = backend.GDB_SBD() + config.eps, config.max_it, config.max_nb, config.bit_length = 1e-10, 200, 8, BIT_LENGTH + config.cpu_subspace = cpu_subspace + dets = sbd.sort_bitarray_array(sbd.from_strings( + [_interleave(a, b) for a, b in itertools.product(alpha, alpha)], + BIT_LENGTH, 2 * norb)) + return sbd.gdb_diag(fcidump, dets, config)["energy"] + + assert solve(True) == pytest.approx(solve(False), abs=1e-9) + + +@pytest.mark.parametrize("config", ["TPB_SBD", "GDB_SBD"]) +def test_cpu_subspace_is_on_the_thrust_backend_only(backend, config): + has = hasattr(getattr(backend, config)(), "cpu_subspace") + assert has is backend.__name__.endswith("_core_gpu_thrust") + + +# --- the drivers' --cpu_subspace ------------------------------------------------------ + +import os # noqa: E402 +import re # noqa: E402 +import subprocess # noqa: E402 +import sys # noqa: E402 + +REPO = __import__("pathlib").Path(__file__).resolve().parents[1] +H2O = REPO / "vendor" / "sbd-upstream" / "data" / "h2o" +DRIVERS = { + "tpb/run_sbd_diag.py": ["--adetfile", str(H2O / "h2o-1em3-alpha.txt"), "--bit_length", str(BIT_LENGTH)], + "gdb/run_gdb_diag.py": ["--from-alpha", str(H2O / "h2o-1em3-alpha.txt"), + "--alpha-limit", str(ALPHA_LIMIT), "--bit_length", str(BIT_LENGTH)], +} + + +def _driver(script, flags, threads): + path = REPO / "examples" / script + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env["OMP_NUM_THREADS"] = threads + return subprocess.run([sys.executable, path.name, "--device", "gpu", "--fcidump", + str(H2O / "fcidump.txt"), "--tolerance", "1e-8", *DRIVERS[script], + *flags], cwd=path.parent, capture_output=True, text=True, + timeout=600, env=env) + + +def _energy(stdout): + return float(re.search(r"Ground state energy:\s*(-?\d+\.\d+)", stdout).group(1)) + + +@pytest.mark.parametrize("script", sorted(DRIVERS)) +def test_the_drivers_cpu_subspace_flag(backend, script): + """--cpu_subspace 1 gives the default energy, and needs more than one thread.""" + _thrust(backend) + default = _driver(script, [], "2") + host = _driver(script, ["--cpu_subspace", "1"], "2") + assert default.returncode == 0 and host.returncode == 0, host.stderr[-2000:] + assert _energy(host.stdout) == pytest.approx(_energy(default.stdout), abs=1e-8) + one_thread = _driver(script, ["--cpu_subspace", "1"], "1") + assert one_thread.returncode != 0 + assert "OMP_NUM_THREADS" in one_thread.stdout + one_thread.stderr diff --git a/test/test_openmp_threads.py b/test/test_openmp_threads.py new file mode 100644 index 0000000..6164ea2 --- /dev/null +++ b/test/test_openmp_threads.py @@ -0,0 +1,65 @@ +# 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. + +"""On the Thrust backend, a solve with one OpenMP thread is refused, not aborted. + +Upstream SBD's Thrust builds do part of the work on the host and call MPI_Abort when +``omp_get_max_threads() == 1``, which ends the Python process with nothing to catch. +The bindings raise ``ValueError`` first. The thread count is fixed when the process +starts, so each case runs in a **subprocess** with ``OMP_NUM_THREADS=1``; only the +Thrust backend is checked, since the CPU and OpenMP-offload builds are not affected. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest + +import sbd + +SCRIPT = r""" +import sys +import sbd + +sbd.init(device="gpu") +data = sys.argv[1] +fcidump = sbd.LoadFCIDump(f"{data}/fcidump.txt") +alpha = [l.strip() for l in open(f"{data}/h2o-1em3-alpha.txt") if l.strip()][:8] +config = sbd.TPB_SBD() +config.bit_length = 20 +half = sbd.sort_bitarray([sbd.from_string(s, 20, 24) for s in alpha]) +try: + sbd.tpb_diag(fcidump, half, half, config) +except ValueError as e: + print("REFUSED", e) +else: + print("SOLVED") +""" + + +@pytest.mark.parametrize("threads, expected", [("1", "REFUSED"), ("2", "SOLVED")]) +def test_thrust_needs_more_than_one_openmp_thread(data_dir, threads, expected): + if "gpu" not in sbd.available_backends(): + pytest.skip("the Thrust backend was not built") + if os.environ.get("SBD_TEST_DEVICE") not in (None, "gpu"): + pytest.skip("the Thrust backend only (SBD_TEST_DEVICE=gpu or unset)") + env = {**os.environ, "OMP_NUM_THREADS": threads} + env = {k: v for k, v in env.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + done = subprocess.run([sys.executable, "-c", SCRIPT, str(data_dir / "h2o")], + capture_output=True, text=True, timeout=300, env=env) + assert done.returncode == 0, f"{done.stdout[-2000:]}\n{done.stderr[-2000:]}" + assert expected in done.stdout, done.stdout[-2000:] + if expected == "REFUSED": + assert "OMP_NUM_THREADS" in done.stdout diff --git a/test/test_rdm_open_shell.py b/test/test_rdm_open_shell.py new file mode 100644 index 0000000..4db7553 --- /dev/null +++ b/test/test_rdm_open_shell.py @@ -0,0 +1,90 @@ +# 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. + +"""The opposite-spin block of the 2-RDM on an open-shell state, against PySCF. + +Upstream SBD fixed a sign in opposite-spin double excitations of the 2-RDM +(r-ccs-cms/sbd#98), shown there on an asymmetric H3 doublet. The error cancels in +the spin-summed 2-RDM but not in the ``ab`` block of ``two_p_rdm``, so this checks +that block, on the same molecule (here in ROHF orbitals), over the full space, +against PySCF's ``make_rdm12s``. Three orbitals is one word per determinant, which SBD +fixes for the life of a process, so the solve runs in a **subprocess**. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys + +import numpy as np +import pytest + +pytest.importorskip("pyscf", reason="pyscf builds the Hamiltonian and the reference") + +from pyscf import ao2mo, fci, gto, scf # noqa: E402 +from pyscf.fci import cistring # noqa: E402 +from pyscf.tools import fcidump # noqa: E402 + +# One word per string here, as in the rest of the suite's TPB tests. +BIT_LENGTH = 20 + +SOLVE = r""" +import json, sys +import sbd +a = json.loads(sys.argv[1]) +sbd.init(device=a["device"]) +def strings(occupations): + return sbd.sort_bitarray([sbd.from_string(s, a["bits"], a["norb"]) for s in occupations]) +config = sbd.TPB_SBD() +config.eps, config.max_it, config.bit_length, config.do_rdm = 1e-12, 300, a["bits"], 1 +result = sbd.tpb_diag(sbd.LoadFCIDump(a["fcidump"]), strings(a["alpha"]), strings(a["beta"]), + config) +print("RESULT " + json.dumps({"energy": result["energy"], + "two_p_rdm": [list(b) for b in result["two_p_rdm"]]})) +""" + + +def test_opposite_spin_2rdm_block_matches_pyscf(device_config, tmp_path): + mol = gto.M(atom="H 0 0 0; H 0 0 0.8; H 0 0 2.2", basis="sto-3g", spin=1, verbose=0) + mf = scf.ROHF(mol).run(conv_tol=1e-12) + c, norb, (na, nb) = mf.mo_coeff, mol.nao, mol.nelec + h1 = c.T @ scf.hf.get_hcore(mol) @ c + eri = ao2mo.restore(1, ao2mo.kernel(mol, c), norb) + path = tmp_path / "fcidump.txt" + fcidump.from_integrals(str(path), h1, eri, norb, na + nb, nuc=0.0, ms=na - nb) + + def strings(n): + return [format(int(s), f"0{norb}b") for s in cistring.make_strings(range(norb), n)] + + args = {"device": device_config.device if device_config is not None else "cpu", + "fcidump": str(path), "bits": BIT_LENGTH, "norb": norb, + "alpha": strings(na), "beta": strings(nb)} + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run([sys.executable, "-c", SOLVE, json.dumps(args)], + capture_output=True, text=True, timeout=300, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + result = json.loads([l for l in done.stdout.splitlines() + if l.startswith("RESULT ")][-1][len("RESULT "):]) + + energy, vector = fci.direct_spin1.kernel(h1, eri, norb, (na, nb), tol=1e-14) + assert result["energy"] == pytest.approx(energy, abs=1e-9) + (_, _), (rdm_aa, rdm_ab, rdm_bb) = fci.direct_spin1.make_rdm12s(vector, norb, (na, nb)) + # two_p_rdm[s + 2t] holds at i + L j + L^2 k + L^3 l; + # reshaped in Fortran order and transposed (0, 2, 1, 3), it is PySCF's p q r s order. + blocks = np.asarray(result["two_p_rdm"]) + block = lambda s: np.reshape(blocks[s], (norb,) * 4, order="F").transpose(0, 2, 1, 3) + np.testing.assert_allclose(block(0), rdm_aa, atol=1e-9) + np.testing.assert_allclose(block(2), rdm_ab, atol=1e-9) # the block #98 fixed + np.testing.assert_allclose(block(3), rdm_bb, atol=1e-9) diff --git a/vendor/sbd-upstream b/vendor/sbd-upstream index 93ebabe..95bf45a 160000 --- a/vendor/sbd-upstream +++ b/vendor/sbd-upstream @@ -1 +1 @@ -Subproject commit 93ebabec6c841d0ade706018b053795d8c1db9bd +Subproject commit 95bf45a5be9ffd977fac7b0aaa359833caf3401a From b14cfd7c1c3ed4350ae549d7ef0277eadd30dd46 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 30 Oct 2026 13:23:58 -0400 Subject: [PATCH 04/16] Bind upstream SBD's CAOP solver: caop_diag_from_files, CAOP_SBD, run_caop_diag.py CAOP diagonalizes a Hamiltonian written as creation/annihilation operators (hard-core bosons, spins, fermions: lattice models such as Hubbard) over a selected basis, as upstream's caop_selected_basis_diagonalization app does, from a --hamfile and basis files. The library is templated on the scalar type, so real and complex amplitudes are both bound in the existing modules and chosen per call (complex=); upstream's app picks one with -D_COMPLEX. CAOP runs on cpu and gpu (Thrust). Upstream has no OpenMP-offload path for it, so that module sets __sbd_has_caop__ = False and the Python wrapper refuses it. examples/caop/run_caop_diag.py takes upstream's option names and by default runs upstream's 8-site hard-core boson example. It exits with sys.exit(main()); switch to sbd.mpi_util.run(main) once #74 lands. Tests (each solve in a subprocess, since CAOP fixes the process-wide word count): upstream's boson chain over its 16-state basis, and upstream's 4-site Hubbard ring (t=-1, U=4) as a --hamfile in upstream's block layout, both against an exact diagonalization built in NumPy from the same terms; the ring also against GDB on upstream's FCIDUMP of the model; complex amplitudes give the real energy; the offload backend and an empty basis are refused. Co-Authored-By: Claude Opus 5.5 --- docs/apidocs/sbd.rst | 10 + examples/caop/README.md | 101 ++++ examples/caop/run_caop_diag.py | 172 +++++++ python/__init__.py | 71 +++ python/bindings.cpp | 298 +++++++++++ .../notes/caop-solver-1f6b2c9e4a7d3058.yaml | 11 + test/test_caop.py | 485 ++++++++++++++++++ 7 files changed, 1148 insertions(+) create mode 100644 examples/caop/README.md create mode 100644 examples/caop/run_caop_diag.py create mode 100644 releasenotes/notes/caop-solver-1f6b2c9e4a7d3058.yaml create mode 100644 test/test_caop.py diff --git a/docs/apidocs/sbd.rst b/docs/apidocs/sbd.rst index 91cc22c..2c876a6 100644 --- a/docs/apidocs/sbd.rst +++ b/docs/apidocs/sbd.rst @@ -167,6 +167,16 @@ Inputs .. autofunction:: from_strings .. autofunction:: sort_bitarray_array +CAOP: creation/annihilation-operator models +============================================ + +For lattice and spin models written as creation and annihilation operators, over a +selected basis, as upstream SBD's ``caop_selected_basis_diagonalization`` app. CAOP runs +on the ``cpu`` and ``gpu`` (Thrust) backends. + +.. autofunction:: CAOP_SBD +.. autofunction:: caop_diag_from_files + Backends and devices ==================== diff --git a/examples/caop/README.md b/examples/caop/README.md new file mode 100644 index 0000000..55ac100 --- /dev/null +++ b/examples/caop/README.md @@ -0,0 +1,101 @@ +# CAOP — models written with creation and annihilation operators + +`run_caop_diag.py` runs upstream SBD's +[`caop_selected_basis_diagonalization`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop_selected_basis_diagonalization) +app from Python: the ground state of a Hamiltonian written as operator terms, over a +selected basis of occupation bitstrings, for lattice and spin models (Hubbard, +hard-core bosons, spins). For molecules, use TPB or GDB. + +**The options are upstream's, with the same names and meanings;** see +[upstream's README](https://github.com/r-ccs-cms/sbd/blob/main/apps/caop_selected_basis_diagonalization/README.md) +for each one and for its `gen_ham.py` and `gen_bits.py`. This page covers what is +specific to the Python driver and what we found running it. + +```bash +cd examples/caop +mpirun -np 4 python run_caop_diag.py --b_comm_size 4 # upstream's 8-site boson chain +``` + +## Differences from upstream's app + +- `--device cpu` or `gpu` (Thrust, NVIDIA only); see *GPU*. +- `--complex` chooses complex amplitudes per run; upstream's app chooses at compile time + (`-D_COMPLEX`). Needed for `(re,im)` coefficients, e.g. `(0,1) cdag 2 c 0`, and for `sy`. +- Defaults are upstream's, except `--method 1`, `--iteration 10` and + `--do_redist_basis 1`, taken from upstream's `run.sh`, and `--bit_length 20`. With no + `--hamfile`/`--basisfiles`, it runs upstream's example files. +- Not exposed: `--init` (only `0`, random, exists) and `--h_comm_size` (listed upstream + but never read; see *MPI*). + +## Preparing the input files + +**A fermion Hamiltonian by hand.** `gen_ham.py` writes hard-core boson chains only. For +fermions, list every term, each with its Hermitian conjugate. Upstream's Hubbard +convention puts alpha site `p` at mode `p` and beta site `p` at mode `L + p`. A 2-site +Hubbard model, `t = 1`, `U = 4`: + +```text +-1 # fermions +-1.0 cdag 0 c 1 # alpha hopping, -t, both directions +-1.0 cdag 1 c 0 +-1.0 cdag 2 c 3 # beta hopping +-1.0 cdag 3 c 2 +4.0 cdag 0 cdag 2 c 2 c 0 # U n_{0,up} n_{0,down} +4.0 cdag 1 cdag 3 c 3 c 1 +``` + +**The basis files must be sorted, and in order across files.** Upstream checks this as +it loads them, before any option is read, and aborts the whole MPI job otherwise: + +- each file ascending as binary numbers, without duplicates; +- with several files, every state in a file greater than every state in the one before. + +`--do_sort_basis` does not repair files that break this. `gen_bits.py` meets it with +several `-o` files by default; with one file, add `--unique`, since its default +single-file output is unsorted. The sorted half-filling basis for the Hubbard file +above is `0101`, `0110`, `1001`, `1010`. For a fixed number of each spin, build the +basis yourself from every alpha and beta occupation: `gen_bits.py` fixes only the total. + +## Tips + +- **Keep `--do_redist_basis 1`.** Without it the states stay where the files put them: + one file means one rank holds them all, which unbalances memory. +- **Converge before you compare.** The `1e-4` tolerance stops early; use + `--tolerance 1e-8` and raise `--iteration` for energies to compare. +- **Fix the start.** The default start is random. Pass `--seed`, or + `--initial_determinant_bitstring` (e.g. a Néel state). +- **Carryover type 1** needs `--carryover_ratio` (the default `0.0` keeps nothing) and + `--do_redist_basis 1`: upstream counts the kept states from each rank's share, so an + uneven basis keeps too many. +- **Carryover type 2** returns a grown basis; the driver prints this rank's share. Feed + it to the next run's `--basisfiles` to iterate. The expansion runs on the CPU. + +## MPI + +The ranks form a grid `ranks = h × b × t`. **b** (`--b_comm_size`) splits the basis and +every vector, and is the only one that divides memory. **t** (`--t_comm_size`) and +**h** split the work of applying the Hamiltonian (**h** also splits its terms), each +holding a full copy of its basis piece. Upstream sets `h = ranks / (b × t)`. The rank +count must be a multiple of `b × t`, and `t ≤ b` is the useful range. The simplest +layout is `--b_comm_size` equal to the rank count. + +## GPU + +Upstream's GPU code applies the Hamiltonian on the fly, so only `--method 0` and `2` +use the GPU. The Davidson/Lanczos vector work, the diagonal and the heat-bath expansion +stay on the CPU, and the vector crosses to the GPU at each application. Methods `1` and +`3` build and apply a stored Hamiltonian on the CPU even on `--device gpu`, and the +driver says so. There is no OpenMP-offload (`gpu-omp`) CAOP code, so no AMD GPU path. +Run one rank per GPU. For a small basis, compare with `--device cpu` and a stored +Hamiltonian, which can come close to the GPU. + +## Checks + +Upstream accepts these and returns a wrong energy, a smaller grid than asked, or hangs. +The bindings refuse them on every rank before any work: + +- a rank count that is not a multiple of `b × t`, or smaller than it; +- Davidson with `--numivec` ≥ `--block`, or `0`; +- `--method` outside `0`–`3`; +- on `gpu`, methods `0`/`2` with `t > b`; +- a Hamiltonian file upstream cannot read, including `(re,im)` without `--complex`. diff --git a/examples/caop/run_caop_diag.py b/examples/caop/run_caop_diag.py new file mode 100644 index 0000000..a646dec --- /dev/null +++ b/examples/caop/run_caop_diag.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 + +# This code is a Qiskit project. +# +# (C) Copyright IBM 2026. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Ground state of a creation/annihilation-operator Hamiltonian over a selected basis. + +The Python counterpart of upstream SBD's ``caop_selected_basis_diagonalization`` app, +with its option names. With no arguments it runs upstream's own example: an 8-site +extended hard-core boson chain over a 16-state basis:: + + mpirun -np 4 python run_caop_diag.py --b_comm_size 4 --do_redist_basis 1 + +The Hamiltonian file's first non-comment line selects the statistics (positive: +hard-core bosons or spins; zero or negative: fermions); each further line is a term, +a coefficient then operator-mode pairs, e.g. ``-1.0 cdag 0 c 1``. The basis files hold +one bitstring of ``--system_size`` characters per line, rightmost = mode 0. +""" + +import argparse +import os +import sys + +_APP = "../../vendor/sbd-upstream/apps/caop_selected_basis_diagonalization" + + +def parse_args(): + parser = argparse.ArgumentParser( + description="CAOP selected-basis diagonalization", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument('--device', default='cpu', choices=['cpu', 'gpu'], + help="'cpu', or 'gpu' (NVHPC Thrust, NVIDIA GPUs only). " + "Upstream's CAOP code has no OpenMP-offload path, so there " + "is no 'gpu-omp' and no CAOP GPU path for AMD GPUs") + parser.add_argument('--hamfile', default=f"{_APP}/hamiltonian.txt", + help='Hamiltonian file') + parser.add_argument('--basisfiles', default=f"{_APP}/basis.txt", + help='Comma-separated basis files') + parser.add_argument('--system_size', type=int, default=8, + help='Number of modes, the length of each basis bitstring') + parser.add_argument('--bit_length', type=int, default=20, + help='Bits stored per 64-bit word of a bitstring') + parser.add_argument('--complex', action='store_true', + help='Complex amplitudes, for a Hamiltonian with complex ' + "coefficients (upstream's -D_COMPLEX build)") + parser.add_argument('--t_comm_size', type=int, default=1, + help='Task communicator size') + parser.add_argument('--b_comm_size', type=int, default=1, + help='Basis communicator size: shards of the basis') + parser.add_argument('--method', type=int, default=1, choices=[0, 1, 2, 3], + help='0/1 Davidson, Hamiltonian on the fly/stored; ' + '2/3 Lanczos, on the fly/stored') + parser.add_argument('--iteration', type=int, default=10, help='Restart cycles') + parser.add_argument('--block', type=int, default=10, + help='Largest reduced-space dimension before a restart') + parser.add_argument('--numivec', type=int, default=1, + help='Initial vectors for Davidson') + parser.add_argument('--tolerance', type=float, default=1e-4, + help='Convergence threshold on the residual (upstream\'s default)') + parser.add_argument('--initial_determinant_bitstring', default='', + help='Start with unit weight on this basis bitstring') + parser.add_argument('--seed', type=int, default=0, + help='Seed for the random initial vector; 0 (upstream\'s default) ' + 'picks a random one, which SBD prints') + parser.add_argument('--do_sort_basis', type=int, default=0, choices=[0, 1], + help='Sort the basis across ranks (for overlapping files)') + parser.add_argument('--do_redist_basis', type=int, default=1, choices=[0, 1], + help='Redistribute the basis evenly over --b_comm_size') + parser.add_argument('--carryover_type', type=int, default=0, choices=[0, 1, 2], + help='0 none, 1 keep the largest-weight fraction, ' + '2 deterministic heat-bath expansion') + parser.add_argument('--carryover_ratio', type=float, default=0.0, + help='carryover_type 1: fraction kept') + parser.add_argument('--heatbath_cutoff', type=float, default=1e-4, + help='carryover_type 2: term acceptance cutoff') + parser.add_argument('--heatbath_truncation', type=float, default=0.0, + help='carryover_type 2: drop a parent unless |c_j|^2 exceeds this') + parser.add_argument('--heatbath_batch_size', type=int, default=1000000, + help='carryover_type 2: candidate batch size per rank before ' + 'merging') + parser.add_argument('--heatbath_parent_distribution', type=int, default=1, + choices=[0, 1], + help='carryover_type 2: 0 keeps the basis layout for the ' + 'expansion, 1 deals parents out by weight') + parser.add_argument('--loadname', default='', + help='Load an initial wavefunction from this prefix') + parser.add_argument('--savename', default='', + help='Save the final wavefunction under this prefix') + args = parser.parse_args() + if args.loadname and args.initial_determinant_bitstring: + parser.error("--loadname and --initial_determinant_bitstring both set the " + "starting vector, and --loadname would win silently: give one") + return args + + +def _print_gpu_guidance(method): + """What the GPU does for this method, so a user can judge whether to use it. + + Upstream's CAOP code has a GPU (Thrust) kernel only for applying H on the fly, + methods 0 and 2. The stored-Hamiltonian methods 1 and 3 build and apply H on the + CPU in every build, so on 'gpu' they run exactly as on 'cpu'. + """ + if method in (1, 3): + print(f"NOTE: --method {method} stores the Hamiltonian, and upstream SBD builds and " + f"applies it on the CPU only: on --device gpu this run uses no GPU and " + f"runs as it would on --device cpu. For the GPU kernel, use --method " + f"{method - 1} (the Hamiltonian applied on the fly).") + else: + print(f"NOTE: --method {method} applies the Hamiltonian on the GPU. The rest of " + f"each Davidson/Lanczos step (vector updates, orthogonalization, the " + f"reduced-space solve) stays on the CPU, and the vector is copied to and " + f"from the GPU on every application. For small bases that overhead can " + f"outweigh the GPU kernel: compare against --device cpu, including " + f"--method {method + 1}, which stores the Hamiltonian.") + + +def main(): + args = parse_args() + import sbd + + sbd.init(device=args.device) + rank = sbd.get_rank() + if args.savename: + directory = os.path.dirname(os.path.abspath(args.savename)) + if not os.path.isdir(directory): + print(f"[rank {rank}] ERROR: --savename {args.savename!r}: directory " + f"{directory!r} does not exist; create it first", file=sys.stderr) + return 1 + + config = sbd.CAOP_SBD() + config.system_size, config.bit_length = args.system_size, args.bit_length + config.t_comm_size, config.b_comm_size = args.t_comm_size, args.b_comm_size + config.method, config.max_it, config.max_nb = args.method, args.iteration, args.block + config.max_iv, config.eps = args.numivec, args.tolerance + config.initial_determinant_bitstring = args.initial_determinant_bitstring + config.seed = args.seed + config.do_sort_basis = bool(args.do_sort_basis) + config.do_redist_basis = bool(args.do_redist_basis) + config.carryover_type, config.ratio = args.carryover_type, args.carryover_ratio + config.heatbath_cutoff = args.heatbath_cutoff + config.heatbath_truncation = args.heatbath_truncation + config.heatbath_batch_size = args.heatbath_batch_size + config.heatbath_parent_distribution = args.heatbath_parent_distribution + + if rank == 0 and args.device == 'gpu': + _print_gpu_guidance(args.method) + + results = sbd.caop_diag_from_files( + args.hamfile, [f for f in args.basisfiles.split(',') if f], config, + loadname=args.loadname, savename=args.savename, complex=args.complex) + if rank == 0: + print("=" * 70) + print(f"Device: {args.device.upper()}" + + (" (complex amplitudes)" if args.complex else "")) + print(f"Ground state energy: {results['energy']:.10f}") + if args.carryover_type: + print(f"Carryover basis on this rank: {len(results['carryover_basis'])}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/__init__.py b/python/__init__.py index ee77754..ddd85f7 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -679,6 +679,75 @@ def tpb_diag_from_files(fcidumpfile, adetfile, sbd_data, ) +def _caop_backend(device): + """The backend module for a CAOP call, refusing one that does not bind CAOP.""" + _ensure_initialized() + backend = get_backend(device) + if not getattr(backend, "__sbd_has_caop__", False): + raise ValueError( + f"the {backend.__name__.split('.')[-1]} backend has no CAOP solver: upstream " + "SBD's CAOP code has CPU and Thrust paths only. Use device='cpu', or " + "device='gpu' (NVHPC Thrust, NVIDIA GPUs only; there is no CAOP GPU path for " + "AMD GPUs)") + return backend + + +def CAOP_SBD(device=None): + """Create a CAOP_SBD configuration object, for :func:`caop_diag_from_files`. + + The fields are those of upstream SBD's ``caop_selected_basis_diagonalization`` + app; ``system_size`` (the number of modes) and ``bit_length`` must be set. + """ + return _caop_backend(device).CAOP_SBD() + + +def caop_diag_from_files(hamfile, basisfiles, sbd_data, loadname="", savename="", + complex=False, device=None): + """Diagonalize a creation/annihilation-operator Hamiltonian over a selected basis. + + The CAOP solver of upstream SBD, as its ``caop_selected_basis_diagonalization`` + app runs it. The Hamiltonian is read from a text file whose first non-comment line + selects the statistics (positive: hard-core bosons or spins; zero or negative: + fermions) and whose other lines are terms, a coefficient then operator-site pairs, + for example ``-1.0 cdag 0 c 1``. The basis is one or more files of bitstrings, one + per line, ``system_size`` characters each. + + Args: + hamfile: Path to the Hamiltonian file. + basisfiles: Path, or list of paths, to the basis files. + sbd_data: CAOP_SBD configuration object. + loadname: Path prefix to load an initial wavefunction from (optional). + savename: Path prefix to save the final wavefunction to (optional). + complex: Solve with complex amplitudes, for a Hamiltonian with complex + coefficients, which the file writes as ``(re,im)``, or one using ``sy``. + Upstream's app chooses this at compile time (``-D_COMPLEX``); here both + are built, and this chooses per call. + device: ``'cpu'`` or ``'gpu'`` (Thrust). The OpenMP-offload backend has no + CAOP solver and is refused. + + Returns: + dict with ``energy`` and ``carryover_basis``, this rank's share of the basis + ``carryover_type`` selected, as an ``(n, words)`` array. + + Raises: + ValueError: before any work, on every rank together, for inputs upstream + would get wrong silently or hang on: a rank count that is not a multiple + of ``t_comm_size * b_comm_size``; Davidson (``method`` 0 or 1) without + ``1 <= max_iv < max_nb``; a ``method`` outside 0-3; on the ``gpu`` + backend with ``method`` 0 or 2, ``t_comm_size > b_comm_size``; and a + Hamiltonian file that cannot be read, including a ``(re,im)`` + coefficient without ``complex=True``. Also for an empty ``basisfiles`` + and a backend with no CAOP solver (``gpu-omp``). See "Checks" in + ``examples/caop/README.md`` for what upstream does with each. + """ + if isinstance(basisfiles, (str, os.PathLike)): + basisfiles = [basisfiles] + backend = _caop_backend(device) + return backend.caop_diag_from_files( + _global_comm, sbd_data, str(hamfile), [str(f) for f in basisfiles], + loadname, savename, complex) + + def tpb_diag(fcidump, adet, bdet, sbd_data, loadname="", savename="", device=None): """ @@ -948,6 +1017,8 @@ def print_info(): 'sort_bitarray', 'sort_bitarray_array', 'tpb_diag_from_files', + 'CAOP_SBD', + 'caop_diag_from_files', 'tpb_diag', 'gdb_diag', diff --git a/python/bindings.cpp b/python/bindings.cpp index 778fcff..3c647c2 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -37,7 +37,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -213,6 +216,94 @@ static bool sbd_all_ranks_ok(bool local_ok, MPI_Comm comm) { return all != 0; } +/** + * The problem the lowest-numbered rank found, on every rank; empty if none did. + * + * For checks only some ranks can make (the one that reads a file), so every rank + * raises the same, informative message rather than "another rank failed". + */ +static std::string sbd_first_problem(const std::string& local, MPI_Comm comm) { + int rank = 0, size = 1; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + int mine = local.empty() ? size : rank, first = size; + MPI_Allreduce(&mine, &first, 1, MPI_INT, MPI_MIN, comm); + if (first == size) return std::string(); + int length = rank == first ? int(local.size()) : 0; + MPI_Bcast(&length, 1, MPI_INT, first, comm); + std::string message(rank == first ? local : std::string(size_t(length), ' ')); + MPI_Bcast(message.data(), length, MPI_CHAR, first, comm); + return message; +} + +#ifndef USE_OMP_OFFLOAD +/** + * Read a CAOP Hamiltonian file's coefficients on every rank, as upstream's loader + * would, and refuse together if any cannot be read. + * + * Upstream reads the file on h-rank 0 only and sends the terms on + * (caop/basic/loadmodel.h:92-96), so a line it cannot read raises on that rank while + * the others wait for terms: the job hangs on more than one rank. This uses + * upstream's line rules (blank and '#' lines skipped, the first other line the + * statistics, then one term per line, coefficient first). A real call reads a + * coefficient as a double, so "(re,im)" is refused there with a pointer to + * complex=True. + */ +static void caop_check_hamfile(const std::string& hamfile, bool complex, MPI_Comm comm, + const char* where) { + std::string problem; + std::ifstream ifs(hamfile); + if (!ifs) { + problem = "cannot open the Hamiltonian file " + hamfile; + } else { + std::string line; + bool statistics_read = false; + size_t number = 0; + while (problem.empty() && std::getline(ifs, line)) { + ++number; + const auto first = line.find_first_not_of(" \t"); + if (first == std::string::npos || line[first] == '#') continue; + std::istringstream iss(line); + if (!statistics_read) { + int statistics; + if (!(iss >> statistics)) { + problem = "line " + std::to_string(number) + " of " + hamfile + + ": the first line must be the statistics (an integer: positive " + "for hard-core bosons or spins, zero or negative for fermions)"; + } + statistics_read = true; + continue; + } + bool read; + if (complex) { + std::complex c; + read = static_cast(iss >> c); + } else { + double c; + read = static_cast(iss >> c); + } + if (!read) { + const bool looks_complex = line[first] == '('; + problem = "line " + std::to_string(number) + " of " + hamfile + + ": cannot read the coefficient '" + + line.substr(first, line.find_first_of(" \t", first) - first) + "'" + + (looks_complex && !complex + ? ". It is written as a complex number (re,im): use complex=True " + "(--complex in the drivers)" + : ". A coefficient is a real number, or (re,im) with complex=True"); + } + } + if (problem.empty() && !statistics_read) { + problem = hamfile + " has no statistics line"; + } + } + const std::string first = sbd_first_problem(problem, comm); + if (!first.empty()) { + throw std::invalid_argument(std::string(where) + ": " + first); + } +} +#endif + /** * Refuse a Thrust solve with one OpenMP thread, before upstream aborts the job. * @@ -1359,6 +1450,213 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { // Version information // ======================================================================== + // ======================================================================== + // CAOP: selected-basis diagonalization of a creation/annihilation-operator + // Hamiltonian (upstream apps/caop_selected_basis_diagonalization) + // ======================================================================== + // + // Upstream's CAOP code has CPU and Thrust paths and no OpenMP-offload one, so the + // offload module does not bind it; __sbd_has_caop__ says which modules do. The + // library is templated on the scalar type, so both real and complex + // Hamiltonians are bound here, chosen per call; upstream's app picks one at + // compile time with -D_COMPLEX. +#ifndef USE_OMP_OFFLOAD + m.attr("__sbd_has_caop__") = true; + + py::class_(m, "CAOP_SBD", py::module_local(), + "Configuration for CAOP diagonalization") + .def(py::init<>()) + .def_readwrite("t_comm_size", &sbd::caop::SBD::t_comm_size, + "Task communicator size (column-wise Hamiltonian application)") + .def_readwrite("b_comm_size", &sbd::caop::SBD::b_comm_size, + "Basis communicator size: how many shards the basis is split into") + // h_comm_size is deliberately NOT exposed, as on TPB_SBD and GDB_SBD. + // Upstream declares the field but never reads it: caop::diag derives + // h_comm_size = mpi_size / (t_comm_size * b_comm_size) + // (caop/basic/sbdiag.h:263, :564). An attribute that silently does nothing + // is worse than an absent one; change the rank count or the other two + // sizes instead. + .def_readwrite("method", &sbd::caop::SBD::method, + "0/1: Davidson, Hamiltonian on the fly / stored; " + "2/3: Lanczos, on the fly / stored") + .def_readwrite("max_it", &sbd::caop::SBD::max_it, "Restart cycles") + .def_readwrite("max_nb", &sbd::caop::SBD::max_nb, + "Largest reduced-space dimension before a restart") + .def_readwrite("max_iv", &sbd::caop::SBD::max_iv, + "Initial vectors (initial subspace dimension) for Davidson") + .def_readwrite("eps", &sbd::caop::SBD::eps, "Convergence threshold on the residual") + .def_readwrite("init", &sbd::caop::SBD::init, + "Initial vector; only 0 (random) is supported upstream") + .def_readwrite("initial_determinant_bitstring", + &sbd::caop::SBD::initial_determinant_bitstring, + "Start with unit weight on this basis bitstring (system_size " + "characters) instead of init. It must occur exactly once in the " + "basis. A loadname takes precedence") + .def_readwrite("seed", &sbd::caop::SBD::seed, "Seed for the random initial vector") + .def_readwrite("system_size", &sbd::caop::SBD::system_size, + "Number of modes, i.e. the length of each basis bitstring") + .def_readwrite("bit_length", &sbd::caop::SBD::bit_length, + "Bits stored per 64-bit word of a packed bitstring") + .def_readwrite("do_sort_basis", &sbd::caop::SBD::do_sort_basis, + "Sort the basis across ranks after loading (for overlapping files)") + .def_readwrite("do_redist_basis", &sbd::caop::SBD::do_redist_basis, + "Redistribute the basis evenly over b_comm after loading") + .def_readwrite("carryover_type", &sbd::caop::SBD::carryover_type, + "0: none; 1: keep the largest-weight fraction; 2: deterministic " + "heat-bath expansion") + .def_readwrite("ratio", &sbd::caop::SBD::ratio, + "carryover_type 1: fraction of basis states kept") + .def_readwrite("heatbath_cutoff", &sbd::caop::SBD::heatbath_cutoff, + "carryover_type 2: accept a term when |a_k c_j| exceeds this") + .def_readwrite("heatbath_truncation", &sbd::caop::SBD::heatbath_truncation, + "carryover_type 2: drop a parent unless |c_j|^2 exceeds this") + .def_readwrite("heatbath_batch_size", &sbd::caop::SBD::heatbath_batch_size, + "carryover_type 2: candidate batch size per rank") + .def_readwrite("heatbath_parent_distribution", + &sbd::caop::SBD::heatbath_parent_distribution, + "carryover_type 2: 0 keeps the basis layout for expansion, 1 " + "deals parents out by weight") + .def("__repr__", [](const sbd::caop::SBD& s) { + return "CAOP_SBD(method=" + std::to_string(s.method) + + ", system_size=" + std::to_string(s.system_size) + + ", bit_length=" + std::to_string(s.bit_length) + + ", b_comm_size=" + std::to_string(s.b_comm_size) + ")"; + }); + + m.def("caop_diag_from_files", + [](py::object py_comm, + const sbd::caop::SBD& sbd_data, + const std::string& hamfile, + const std::vector& basisfiles, + const std::string& loadname, + const std::string& savename, + bool complex) { + MPI_Comm comm = get_mpi_comm(py_comm); + sbd_require_openmp_threads("caop_diag_from_files", comm); + int mpi_rank; + MPI_Comm_rank(comm, &mpi_rank); + int mpi_size; + MPI_Comm_size(comm, &mpi_size); + + // ---- what upstream does not check --------------------------------- + // These depend only on the config and the rank count, which every rank + // sees identically, so every rank throws together without a vote. Each + // was observed upstream to return a wrong energy with exit 0, or hang. + const int t_comm_size = sbd_data.t_comm_size; + const int b_comm_size = sbd_data.b_comm_size; + if (t_comm_size < 1 || b_comm_size < 1) { + throw std::invalid_argument( + "caop_diag_from_files requires t_comm_size >= 1 and b_comm_size >= 1"); + } + // caop::diag derives h_comm_size = ranks / (t * b) by integer division + // (caop/basic/sbdiag.h:263) and setup_communicator then splits ranks into + // groups of t * b (helper.h:22-27). With a remainder, the last group is + // short, so ranks disagree on the grid (of 3 ranks with b = 2, one sees + // b = 1): a wrong energy, with exit 0. With fewer ranks than t * b the + // grid is consistent but smaller than asked (2 ranks, b = 4 run b = 2). + if (mpi_size % (t_comm_size * b_comm_size) != 0) { + throw std::invalid_argument( + "caop_diag_from_files: the rank count (" + std::to_string(mpi_size) + + ") must be a multiple of t_comm_size * b_comm_size (" + + std::to_string(t_comm_size) + " * " + std::to_string(b_comm_size) + + " = " + std::to_string(t_comm_size * b_comm_size) + + "). Upstream derives the helper size as ranks / (t * b) and " + "returns a wrong energy otherwise"); + } + // caop::diag's dispatch is `if (method == 0 || method == 2) {...} else + // {...}`, so any other value silently runs a stored-Hamiltonian method. + if (sbd_data.method < 0 || sbd_data.method > 3) { + throw std::invalid_argument( + "caop_diag_from_files: method must be 0-3 (0/1 Davidson, 2/3 " + "Lanczos; odd stores the Hamiltonian), got " + + std::to_string(sbd_data.method)); + } + // Davidson (methods 0 and 1) starts from max_iv vectors and grows the + // reduced space up to max_nb before a restart, so it needs room to grow: + // max_iv == max_nb returns a wrong energy, and max_iv > max_nb or 0 + // hangs. Lanczos (methods 2 and 3) does not use max_iv. + if ((sbd_data.method == 0 || sbd_data.method == 1) + && (sbd_data.max_iv < 1 || sbd_data.max_iv >= sbd_data.max_nb)) { + throw std::invalid_argument( + "caop_diag_from_files: Davidson (method 0 or 1) needs " + "1 <= max_iv < max_nb, got max_iv=" + std::to_string(sbd_data.max_iv) + + ", max_nb=" + std::to_string(sbd_data.max_nb) + + ". Upstream returns a wrong energy at max_iv == max_nb and hangs " + "above it or at 0"); + } +#ifdef SBD_THRUST + // make_slide deals the b ring stations out over the t task ranks + // (caop/basic/helper.h:41-55), so with t > b some ranks get none. The CPU + // mult guards an empty slide (mult.h:19), but the Thrust one reads + // slide_[0] unguarded (mult_thrust.h:628) and the job hangs; the + // stored-Hamiltonian methods run the CPU mult and are unaffected. The + // same t <= b holds for GDB on every backend. + if ((sbd_data.method == 0 || sbd_data.method == 2) && t_comm_size > b_comm_size) { + throw std::invalid_argument( + "caop_diag_from_files on the gpu (Thrust) backend with method " + + std::to_string(sbd_data.method) + " needs t_comm_size <= b_comm_size, " + "got t_comm_size=" + std::to_string(t_comm_size) + ", b_comm_size=" + + std::to_string(b_comm_size) + ": a task rank with no basis-ring " + "station hangs upstream's GPU kernel. Raise b_comm_size, lower " + "t_comm_size, or use device='cpu'"); + } +#endif +#ifdef SBD_THRUST + int numDevices, myDevice; +#ifdef __CUDACC__ + cudaGetDeviceCount(&numDevices); + myDevice = mpi_rank % numDevices; + cudaSetDevice(myDevice); +#else + hipGetDeviceCount(&numDevices); + myDevice = mpi_rank % numDevices; + hipSetDevice(myDevice); +#endif +#endif + if (basisfiles.empty()) { + throw std::invalid_argument("caop_diag_from_files needs at least one basis file"); + } + // Upstream reads the hamfile on h-rank 0 only and sends the terms on + // (caop/basic/loadmodel.h:92-96), so a file it cannot read hangs the other + // ranks; see caop_check_hamfile. + caop_check_hamfile(hamfile, complex, comm, "caop_diag_from_files"); + double energy = 0.0; + sbd::det_vector co_basis; + { + py::gil_scoped_release release; + if (complex) { + sbd::caop::diag>(comm, sbd_data, hamfile, basisfiles, + loadname, savename, energy, co_basis); + } else { + sbd::caop::diag(comm, sbd_data, hamfile, basisfiles, + loadname, savename, energy, co_basis); + } + } + const size_t words = co_basis.empty() ? 0 : co_basis.elem_size(); + py::array_t co_out(std::vector{ + static_cast(co_basis.size()), static_cast(words)}); + if (!co_basis.empty()) { + std::memcpy(co_out.mutable_data(), co_basis.cflat().data(), + co_basis.size() * words * sizeof(size_t)); + } + py::dict results; + results["energy"] = energy; + results["carryover_basis"] = co_out; + return results; + }, + "CAOP diagonalization of a Hamiltonian file over a basis given as files, " + "as upstream's caop_selected_basis_diagonalization app", + py::arg("comm"), + py::arg("sbd_data"), + py::arg("hamfile"), + py::arg("basisfiles"), + py::arg("loadname") = "", + py::arg("savename") = "", + py::arg("complex") = false); +#else + m.attr("__sbd_has_caop__") = false; +#endif + m.attr("__version__") = "1.2.0"; } diff --git a/releasenotes/notes/caop-solver-1f6b2c9e4a7d3058.yaml b/releasenotes/notes/caop-solver-1f6b2c9e4a7d3058.yaml new file mode 100644 index 0000000..db79ac7 --- /dev/null +++ b/releasenotes/notes/caop-solver-1f6b2c9e4a7d3058.yaml @@ -0,0 +1,11 @@ +--- +features: + - | + New :func:`sbd.caop_diag_from_files` and :class:`sbd.CAOP_SBD` bind upstream SBD's + CAOP solver: the ground state of a Hamiltonian written as creation and annihilation + operators (hard-core bosons, spins or fermions, as for lattice models such as the + Hubbard model), over a selected basis, as upstream's + ``caop_selected_basis_diagonalization`` app runs it. Real and complex amplitudes are + both available, chosen per call with ``complex=``. ``examples/caop/run_caop_diag.py`` + is the command-line driver, with upstream's option names. CAOP runs on the ``cpu`` + and ``gpu`` (Thrust) backends; upstream has no OpenMP-offload path for it. diff --git a/test/test_caop.py b/test/test_caop.py new file mode 100644 index 0000000..64e5f7a --- /dev/null +++ b/test/test_caop.py @@ -0,0 +1,485 @@ +# 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. + +"""CAOP selected-basis diagonalization, against exact references. + +The cases are upstream SBD's own: + +- ``hamiltonian.txt`` and ``basis.txt`` from ``apps/caop_selected_basis_diagonalization``: + an 8-site extended hard-core boson chain over a 16-state selected basis. +- The 4-site Hubbard ring of upstream's ``tests/functionality/test_caop_gdb_hubbard.cc`` + (t = -1, U = 4, two electrons of each spin), written here as a CAOP Hamiltonian file + in upstream's block layout: alpha site p is mode p, beta site p is mode L + p. + +Each is checked against an exact diagonalization built here in NumPy from the same +operator terms over the same basis, so the reference does not come from SBD. The +Hubbard ring is also checked against GDB on upstream's FCIDUMP of the same model, the +comparison upstream's own test makes. Every solve runs in a **subprocess**: CAOP fixes +the process-wide determinant word count, as GDB does. +""" + +from __future__ import annotations + +import itertools +import json +import os +import pathlib +import subprocess +import sys + +import numpy as np +import pytest + +import sbd + +REPO = pathlib.Path(__file__).resolve().parents[1] +UPSTREAM = REPO / "vendor" / "sbd-upstream" +CAOP_APP = UPSTREAM / "apps" / "caop_selected_basis_diagonalization" +HUBBARD_FCIDUMP = UPSTREAM / "tests" / "functionality" / "caop_gdb_hubbard4_fcidump.txt" + +SOLVE = r""" +import json, sys +import sbd +device, ham, basis, size, complex_ = sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4]), sys.argv[5] == "1" +sbd.init(device=device) +config = sbd.CAOP_SBD() +config.system_size, config.bit_length = size, 20 +config.method, config.max_it, config.max_nb, config.eps = 1, 200, 20, 1e-12 +config.do_redist_basis = True +result = sbd.caop_diag_from_files(ham, basis.split(","), config, complex=complex_) +print("RESULT " + json.dumps({"energy": result["energy"]})) +""" + + +def _device(device_config): + return device_config.device if device_config is not None else "cpu" + + +def _solve(device, ham, basis, size, complex_=False): + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run([sys.executable, "-c", SOLVE, device, str(ham), str(basis), + str(size), "1" if complex_ else "0"], + capture_output=True, text=True, timeout=600, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + line = [l for l in done.stdout.splitlines() if l.startswith("RESULT ")][-1] + return json.loads(line[len("RESULT "):])["energy"] + + +# --- an independent reference: the operator terms applied in NumPy ------------------ + +def _read_terms(path): + """(fermion, [(coefficient, [(op, mode), ...]), ...]) from a CAOP Hamiltonian file.""" + lines = [l.split("#")[0].strip() for l in pathlib.Path(path).read_text().splitlines()] + lines = [l for l in lines if l] + fermion = int(float(lines[0])) <= 0 + terms = [] + for line in lines[1:]: + fields = line.split() + ops = [(fields[k], int(fields[k + 1])) for k in range(1, len(fields), 2)] + terms.append((_coefficient(fields[0]), ops)) + return fermion, terms + + +def _coefficient(text): + """A coefficient as upstream writes it: a real number, or (re,im).""" + if text.startswith("("): + re, im = text.strip("()").split(",") + return complex(float(re), float(im)) + return complex(float(text)) + + +def _apply(state, ops, fermion): + """Apply a product of operators (rightmost first) to a basis state; (sign, state).""" + sign = 1.0 + for op, mode in reversed(ops): + occupied = (state >> mode) & 1 + if op in ("c", "b") and not occupied or op in ("cdag", "bdag") and occupied: + return 0.0, None + if fermion: + sign *= (-1) ** bin(state & ((1 << mode) - 1)).count("1") + state ^= 1 << mode + return sign, state + + +def _exact(path, states): + """Lowest eigenvalue of the Hamiltonian file restricted to ``states``.""" + fermion, terms = _read_terms(path) + index = {s: i for i, s in enumerate(states)} + h = np.zeros((len(states), len(states)), dtype=complex) + for j, s in enumerate(states): + for coefficient, ops in terms: + sign, t = _apply(s, ops, fermion) + if t is not None and t in index: + h[index[t], j] += coefficient * sign + assert np.allclose(h, h.conj().T), "the Hamiltonian is not Hermitian on this basis" + return float(np.linalg.eigvalsh(h)[0]) + + +def _bitstrings(path): + return [int(l.strip(), 2) for l in pathlib.Path(path).read_text().split()] + + +def _hubbard(tmp_path, sites=4, t=-1.0, u=4.0, ring=True): + """Upstream's Hubbard ring as a CAOP Hamiltonian file, and its full basis.""" + bonds = [(i, i + 1) for i in range(sites - 1)] + ([(0, sites - 1)] if ring else []) + lines = ["-1"] + for spin in (0, sites): + for i, j in bonds: + lines += [f"{t} cdag {i + spin} c {j + spin}", f"{t} cdag {j + spin} c {i + spin}"] + for i in range(sites): + lines.append(f"{u} cdag {i} cdag {i + sites} c {i + sites} c {i}") + ham = tmp_path / "hubbard.txt" + ham.write_text("\n".join(lines) + "\n") + half = [sum(1 << p for p in occ) for occ in itertools.combinations(range(sites), sites // 2)] + states = sorted(a | (b << sites) for a in half for b in half) + basis = tmp_path / "hubbard_basis.txt" + basis.write_text("\n".join(format(s, f"0{2 * sites}b") for s in states) + "\n") + return ham, basis, states + + +# --- tests ----------------------------------------------------------------------------- + +@pytest.fixture +def caop_device(device_config): + device = _device(device_config) + if device == "gpu-omp": + pytest.skip("upstream SBD's CAOP code has no OpenMP-offload path") + return device + + +def test_upstream_boson_chain_matches_exact_diagonalization(caop_device): + """Upstream's example: an 8-site hard-core boson chain over its 16-state basis.""" + ham, basis = CAOP_APP / "hamiltonian.txt", CAOP_APP / "basis.txt" + expected = _exact(ham, _bitstrings(basis)) + assert _solve(caop_device, ham, basis, 8) == pytest.approx(expected, abs=1e-9) + + +def test_complex_amplitudes_give_the_real_energy(caop_device): + """A real Hamiltonian solved with complex amplitudes gives the same energy.""" + ham, basis = CAOP_APP / "hamiltonian.txt", CAOP_APP / "basis.txt" + assert _solve(caop_device, ham, basis, 8, complex_=True) == pytest.approx( + _solve(caop_device, ham, basis, 8), abs=1e-9) + + +def test_hubbard_ring_matches_exact_diagonalization_and_gdb(caop_device, tmp_path): + """Upstream's 4-site Hubbard ring: CAOP, an exact diagonalization, and GDB agree. + + GDB solves upstream's FCIDUMP of the same model over the same determinants. Its + determinant bits interleave the spins, so each state is converted to that layout. + """ + ham, basis, states = _hubbard(tmp_path) + expected = _exact(ham, states) + assert _solve(caop_device, ham, basis, 8) == pytest.approx(expected, abs=1e-9) + + sites = 4 + interleaved = [] + for s in states: + bits = ["0"] * (2 * sites) + for p in range(sites): + if s >> p & 1: + bits[2 * sites - 1 - 2 * p] = "1" + if s >> (sites + p) & 1: + bits[2 * sites - 1 - (2 * p + 1)] = "1" + interleaved.append("".join(bits)) + gdb = r""" +import json, sys +import sbd +sbd.init(device=sys.argv[1]) +dets = sbd.sort_bitarray_array(sbd.from_strings(sys.argv[3].split(","), 20, 8)) +config = sbd.GDB_SBD(); config.eps, config.max_it, config.bit_length = 1e-12, 200, 20 +print("RESULT " + json.dumps(sbd.gdb_diag(sbd.LoadFCIDump(sys.argv[2]), dets, config)["energy"])) +""" + device = caop_device + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run([sys.executable, "-c", gdb, device, str(HUBBARD_FCIDUMP), + ",".join(interleaved)], capture_output=True, text=True, + timeout=600, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + gdb_energy = json.loads([l for l in done.stdout.splitlines() + if l.startswith("RESULT ")][-1][len("RESULT "):]) + assert gdb_energy == pytest.approx(expected, abs=1e-9) + + +def test_the_offload_backend_is_refused(): + if "gpu-omp" not in sbd.available_backends(): + pytest.skip("the OpenMP-offload backend was not built") + with pytest.raises(ValueError, match="no CAOP solver"): + sbd.CAOP_SBD(device="gpu-omp") + + +def test_a_basis_is_required(backend): + """The basis check comes before any solve, so it is safe in the suite's process.""" + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + with pytest.raises(ValueError, match="at least one basis file"): + sbd.caop_diag_from_files(CAOP_APP / "hamiltonian.txt", [], backend.CAOP_SBD()) + + +# --- the driver -------------------------------------------------------------------- + +def _run_driver(device, *flags): + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + return subprocess.run([sys.executable, "run_caop_diag.py", "--device", device, *flags], + cwd=REPO / "examples" / "caop", capture_output=True, text=True, + timeout=600, env=env) + + +def _driver_energy(stdout): + import re + + return float(re.search(r"Ground state energy:\s*(-?\d+\.\d+)", stdout).group(1)) + + +def test_the_driver_runs_upstreams_example(caop_device): + """With no input flags the driver runs upstream's example, to the exact energy. + + The tolerance is upstream's 1e-4 by default, so this tightens it to compare. + """ + done = _run_driver(caop_device, "--tolerance", "1e-12", "--iteration", "200") + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + expected = _exact(CAOP_APP / "hamiltonian.txt", _bitstrings(CAOP_APP / "basis.txt")) + assert _driver_energy(done.stdout) == pytest.approx(expected, abs=1e-9) + + +def test_the_driver_passes_the_seed_and_heatbath_flags(caop_device): + """The seed and heat-bath flags are accepted and reach SBD. + + A fixed seed is reported back by SBD, and a heat-bath carryover with a parent + truncation that keeps nothing returns an empty basis where the default does not. + """ + common = ("--tolerance", "1e-12", "--iteration", "200", "--carryover_type", "2", + "--heatbath_batch_size", "1000", "--heatbath_parent_distribution", "0") + seeded = _run_driver(caop_device, *common, "--seed", "1729") + assert seeded.returncode == 0, f"{seeded.stdout[-3000:]}\n{seeded.stderr[-3000:]}" + assert "seed: 1729 (user-specified)" in seeded.stdout + kept = int(seeded.stdout.split("Carryover basis on this rank:")[1].split()[0]) + nothing = _run_driver(caop_device, *common, "--heatbath_truncation", "2.0") + assert nothing.returncode == 0, f"{nothing.stdout[-3000:]}\n{nothing.stderr[-3000:]}" + dropped = int(nothing.stdout.split("Carryover basis on this rank:")[1].split()[0]) + assert kept > 0 and dropped < kept + + +# --- what the driver says about the GPU -------------------------------------------- + +@pytest.mark.parametrize("method", [0, 1]) +def test_the_driver_says_what_runs_on_the_gpu(method): + """On gpu the driver says whether this method uses the GPU; on cpu it says nothing. + + Upstream's GPU kernel is for applying the Hamiltonian on the fly (methods 0, 2); + the stored-Hamiltonian methods (1, 3) run on the CPU even on the Thrust backend. + One method of each kind covers both notes. + """ + if "gpu" not in sbd.available_backends(): + pytest.skip("the Thrust backend was not built") + flags = ("--method", str(method), "--tolerance", "1e-8", "--iteration", "50") + gpu = _run_driver("gpu", *flags) + assert gpu.returncode == 0, f"{gpu.stdout[-3000:]}\n{gpu.stderr[-3000:]}" + assert gpu.stdout.count("NOTE: --method") == 1 + if method == 1: + assert "uses no GPU" in gpu.stdout and "--method 0" in gpu.stdout + else: + assert "applies the Hamiltonian on the GPU" in gpu.stdout + assert "compare against --device cpu" in gpu.stdout + cpu = _run_driver("cpu", *flags) + assert cpu.returncode == 0 and "NOTE: --method" not in cpu.stdout + + +# --- what upstream does not check -------------------------------------------------- +# +# The bindings refuse these before any work, so they run in this process: a refused +# call never reaches caop::diag, which is what fixes the determinant word count. + +def _caop_config(backend, **fields): + config = backend.CAOP_SBD() + config.system_size, config.bit_length = 8, 20 + for name, value in fields.items(): + setattr(config, name, value) + return config + + +@pytest.mark.parametrize("fields, match", [ + ({"method": 1, "max_iv": 10, "max_nb": 10}, "1 <= max_iv < max_nb"), + ({"method": 0, "max_iv": 11, "max_nb": 10}, "1 <= max_iv < max_nb"), + ({"method": 1, "max_iv": 0}, "1 <= max_iv < max_nb"), + ({"method": 5}, "method must be 0-3"), +], ids=["iv-eq-nb", "iv-gt-nb", "iv-0", "method-5"]) +def test_settings_upstream_gets_wrong_are_refused(backend, fields, match): + """Davidson with numivec >= block returns a wrong energy upstream, or hangs above it + or at 0; a method outside 0-3 silently runs a stored-Hamiltonian one.""" + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + with pytest.raises(ValueError, match=match): + sbd.caop_diag_from_files(CAOP_APP / "hamiltonian.txt", [CAOP_APP / "basis.txt"], + _caop_config(backend, **fields)) + + +# Upstream reads the Hamiltonian on one rank and sends the terms on, so a file it cannot +# read raises there and leaves the other ranks waiting. The bindings read it first. +BAD_HAMFILES = { + "complex-on-real": ("-1\n(0,1) cdag 1 c 0\n(0,-1) cdag 0 c 1\n", "complex=True"), + "bad-coefficient": ("-1\nabc cdag 1 c 0\n", "cannot read the coefficient 'abc'"), + "no-statistics": ("# only a comment\n", "no statistics line"), +} + + +@pytest.mark.parametrize("name", sorted(BAD_HAMFILES)) +def test_a_hamiltonian_file_upstream_cannot_read_is_refused(backend, tmp_path, name): + """Refused naming the line, and a (re,im) coefficient on a real solve says to pass + complex=True rather than being read as its real part.""" + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + text, match = BAD_HAMFILES[name] + ham = tmp_path / "ham.txt" + ham.write_text(text) + with pytest.raises(ValueError, match=match.replace("(", r"\(").replace(")", r"\)")): + sbd.caop_diag_from_files(ham, [CAOP_APP / "basis.txt"], _caop_config(backend)) + + +def test_a_complex_hamiltonian_file_is_read_with_complex_amplitudes(backend, tmp_path): + """The same (re,im) file solves with complex=True, to the exact energy.""" + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + ham = tmp_path / "ham.txt" + ham.write_text("-1\n(0,1) cdag 2 c 0\n(0,-1) cdag 0 c 2\n(-0.5,0) cdag 2 c 1\n" + "(-0.5,0) cdag 1 c 2\n(0.5,0) cdag 2 c 2\n") + basis = tmp_path / "basis.txt" + basis.write_text("001\n010\n100\n") + config = _caop_config(backend, method=1, max_it=200, max_nb=3, eps=1e-12) + config.system_size = 3 + energy = sbd.caop_diag_from_files(ham, [basis], config, complex=True)["energy"] + assert energy == pytest.approx(_exact(ham, [1, 2, 4]), abs=1e-9) + + +# --- several ranks ----------------------------------------------------------------- +# +# These run in the ranks of `tox -e mpi` (SBD_TEST_NPROCS, 2 by default). CAOP's basis +# shares GDB's process-wide word count, and every case here is one word, so they share +# the job's process with the TPB and GDB tests. + +def _refusal_on_every_rank(backend, comm, match, **fields): + """Each rank calls; every rank must refuse, or one would wait in a collective.""" + try: + sbd.caop_diag_from_files(CAOP_APP / "hamiltonian.txt", [CAOP_APP / "basis.txt"], + _caop_config(backend, **fields)) + except ValueError as e: + outcome = str(e) + else: + outcome = "SOLVED" + outcomes = comm.allgather(outcome) + assert all(match in o for o in outcomes), outcomes + + +@pytest.fixture(scope="module") +def hubbard6(): + """A 6-site Hubbard ring at half filling, 400 states, in three uneven basis files. + + Written by rank 0 into one directory every rank reads, so the ranks share one basis. + Uneven files give the redistribution something to even out. + """ + import tempfile + + from mpi4py import MPI + + comm = MPI.COMM_WORLD + path = pathlib.Path(comm.bcast(tempfile.mkdtemp(prefix="caop-mpi-") + if comm.Get_rank() == 0 else None, root=0)) + if comm.Get_rank() == 0: + ham, basis, _ = _hubbard(path, sites=6) + rows = basis.read_text().split() + for k, chunk in enumerate((rows[:50], rows[50:250], rows[250:])): + (path / f"basis{k}.txt").write_text("\n".join(chunk) + "\n") + comm.Barrier() + rows = (path / "hubbard_basis.txt").read_text().split() + files = [path / f"basis{k}.txt" for k in range(3)] + # A fixed start, since the default start is random and every rank must agree. + return path / "hubbard.txt", files, [int(r, 2) for r in rows], rows[0] + + +@pytest.mark.mpi +@pytest.mark.parametrize("method", [0, 1, 2, 3]) +def test_the_basis_split_over_every_rank_gives_the_exact_energy(backend, caop_device, + hubbard6, method): + """The basis over all ranks (b_comm_size = ranks), each method, against NumPy.""" + from mpi4py import MPI + + comm = MPI.COMM_WORLD + ham, files, states, start = hubbard6 + config = _caop_config(backend, method=method, max_it=300, max_nb=20, eps=1e-12, + do_redist_basis=True, b_comm_size=comm.Get_size(), + initial_determinant_bitstring=start) + config.system_size = 12 + energy = sbd.caop_diag_from_files(ham, files, config)["energy"] + energies = comm.allgather(energy) + assert len(set(energies)) == 1, f"ranks disagree: {energies}" + assert energy == pytest.approx(_exact(ham, states), abs=1e-9) + + +@pytest.mark.mpi +def test_a_grid_the_ranks_cannot_make_is_refused(backend): + """A rank count that is not a multiple of t * b: upstream floors the helper size and + the ranks disagree on the grid (a wrong energy), or, with fewer ranks than t * b, + run a smaller grid than asked.""" + from mpi4py import MPI + + comm = MPI.COMM_WORLD + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + _refusal_on_every_rank(backend, comm, "multiple of t_comm_size", + b_comm_size=comm.Get_size() + 1) + + +@pytest.mark.mpi +def test_t_above_b_is_refused_for_the_gpu_kernel(backend): + """t > b gives task ranks no basis-ring station, which hangs upstream's GPU kernel. + + The CPU on-the-fly code guards an empty station list (mult.h:19); the Thrust one + reads its first entry unguarded (mult_thrust.h:628). + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + if not backend.__name__.endswith("_core_gpu_thrust"): + pytest.skip("the gpu (Thrust) backend only") + if comm.Get_size() < 2: + pytest.skip("needs at least two ranks for t > b") + _refusal_on_every_rank(backend, comm, "t_comm_size <= b_comm_size", method=0, + t_comm_size=comm.Get_size(), b_comm_size=1) + + +@pytest.mark.mpi +def test_a_hamiltonian_file_upstream_cannot_read_is_refused_on_every_rank(backend): + """Upstream reads the file on one rank, so a bad one used to leave the others + waiting for its terms: the job hung. Every rank now refuses together.""" + import tempfile + + from mpi4py import MPI + + comm = MPI.COMM_WORLD + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + path = pathlib.Path(comm.bcast(tempfile.mkdtemp(prefix="caop-bad-") + if comm.Get_rank() == 0 else None, root=0)) + if comm.Get_rank() == 0: + (path / "ham.txt").write_text(BAD_HAMFILES["complex-on-real"][0]) + comm.Barrier() + try: + sbd.caop_diag_from_files(path / "ham.txt", [CAOP_APP / "basis.txt"], + _caop_config(backend, b_comm_size=comm.Get_size())) + except ValueError as e: + outcome = str(e) + else: + outcome = "SOLVED" + outcomes = comm.allgather(outcome) + assert all("complex=True" in o for o in outcomes), outcomes From dc3051ec7b60077534c37dc2e6b79cdf0128bbeb Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 13:57:37 -0500 Subject: [PATCH 05/16] Bind upstream SBD's caop-spectrum: sbd.spectrum Green's functions of a CAOP state sbd.spectrum.green_function runs upstream's caop-spectrum on a ground state saved by caop_diag_from_files: block Lanczos in the excitation basis, for single-particle c/cdag on chosen orbitals or for general operators given as CAOP operator files, with upstream's deflation tolerances. green_function_from_coefficients reevaluates saved coefficients on a new grid or broadening without rerunning (--read-coefficients), and spectral_function gives the spectral matrix. Real and complex amplitudes, cpu and gpu (Thrust) backends. examples/caop/run_caop_spectrum.py runs it with upstream's option names and CSV output, and the CAOP README continues from the ground-state example to the Green's function. Co-Authored-By: Claude Opus 5.5 --- docs/apidocs/index.rst | 1 + docs/apidocs/sbd.rst | 3 +- docs/apidocs/sbd.spectrum.rst | 7 + examples/caop/README.md | 37 ++ examples/caop/run_caop_spectrum.py | 186 +++++++++ python/bindings.cpp | 237 ++++++++++++ python/spectrum.py | 144 +++++++ .../notes/caop-spectrum-8c3e5a1f7b2d4096.yaml | 19 + test/test_spectrum.py | 358 ++++++++++++++++++ 9 files changed, 991 insertions(+), 1 deletion(-) create mode 100644 docs/apidocs/sbd.spectrum.rst create mode 100644 examples/caop/run_caop_spectrum.py create mode 100644 python/spectrum.py create mode 100644 releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml create mode 100644 test/test_spectrum.py diff --git a/docs/apidocs/index.rst b/docs/apidocs/index.rst index 58288ee..e4fb4db 100644 --- a/docs/apidocs/index.rst +++ b/docs/apidocs/index.rst @@ -8,3 +8,4 @@ sbd sbd.sbd_solver sbd.device_config + sbd.spectrum diff --git a/docs/apidocs/sbd.rst b/docs/apidocs/sbd.rst index 2c876a6..bb89a36 100644 --- a/docs/apidocs/sbd.rst +++ b/docs/apidocs/sbd.rst @@ -172,7 +172,8 @@ CAOP: creation/annihilation-operator models For lattice and spin models written as creation and annihilation operators, over a selected basis, as upstream SBD's ``caop_selected_basis_diagonalization`` app. CAOP runs -on the ``cpu`` and ``gpu`` (Thrust) backends. +on the ``cpu`` and ``gpu`` (Thrust) backends. See :mod:`sbd.spectrum` for Green's +functions of a CAOP ground state. .. autofunction:: CAOP_SBD .. autofunction:: caop_diag_from_files diff --git a/docs/apidocs/sbd.spectrum.rst b/docs/apidocs/sbd.spectrum.rst new file mode 100644 index 0000000..25eba19 --- /dev/null +++ b/docs/apidocs/sbd.spectrum.rst @@ -0,0 +1,7 @@ +===================================================== +Green's functions (:mod:`sbd.spectrum`) +===================================================== + +.. automodule:: sbd.spectrum + :members: green_function, green_function_from_coefficients, spectral_function + :member-order: bysource diff --git a/examples/caop/README.md b/examples/caop/README.md index 55ac100..4933d0b 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -70,6 +70,43 @@ basis yourself from every alpha and beta occupation: `gen_bits.py` fixes only th - **Carryover type 2** returns a grown basis; the driver prints this rank's share. Feed it to the next run's `--basisfiles` to iterate. The expansion runs on the CPU. +## Green's function + +`run_caop_spectrum.py` continues from a saved ground state: upstream's +[`caop-spectrum`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop-spectrum) app +in Python, with its option names and CSV output; see +[its README](https://github.com/r-ccs-cms/sbd/blob/main/apps/caop-spectrum/README.md) +for the options, conventions and limits. For the 2-site Hubbard file above: + +```bash +# 1. ground state, saved; note the "Ground state energy:" line (10 decimals), not +# upstream's rounded "Energy =" log line +python run_caop_diag.py --hamfile hub2.txt --basisfiles hub2_basis.txt --system_size 4 \ + --tolerance 1e-10 --savename gs- +# 2. the N+1 sector, written sorted (here 3 electrons: 2 alpha + 1 beta, or 1 + 2) +printf '0111\n1011\n1101\n1110\n' > add_basis.txt +# 3. G+ for alpha site 0 and beta site 0, from the saved state +python run_caop_spectrum.py --hamfile hub2.txt --sites 4 --loadname gs- \ + --reference-energy --orbitals 0,2 --channel addition \ + --extra-detfile add_basis.txt --save-coefficients add.coeff --output add.csv \ + --omega-min -6 --omega-max 8 --points 141 --eta 0.1 +# 4. a narrower broadening from the saved coefficients, without rerunning +python run_caop_spectrum.py --read-coefficients add.coeff --output narrow.csv \ + --omega-min -6 --omega-max 8 --points 281 --eta 0.05 +``` + +- **Give the whole N±1 sector as `--extra-detfile`.** Upstream builds the excitation + basis only from what the operators reach from the state, which is usually not + enough. For orbitals of both spins, include both sectors (removing an alpha or a + beta electron gives different sectors). +- **Match the ground-state run:** `--sites` is its `--system_size`, `--bit-length` + its `--bit_length`, `--wavefunction-shards` its `--b_comm_size`, and + `--scalar-type complex` if it used `--complex`. The files record none of these. +- **Add both channels** for the full retarded Green's function: run `addition` and + `removal` on the same orbitals and grid, and sum. +- **Only `--hamiltonian-mode on-the-fly` uses the GPU**, as for the ground state; the + default `stored` runs on the CPU, and the driver says so. + ## MPI The ranks form a grid `ranks = h × b × t`. **b** (`--b_comm_size`) splits the basis and diff --git a/examples/caop/run_caop_spectrum.py b/examples/caop/run_caop_spectrum.py new file mode 100644 index 0000000..408c592 --- /dev/null +++ b/examples/caop/run_caop_spectrum.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 + +# This code is a Qiskit project. +# +# (C) Copyright IBM 2026. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Green's function of a CAOP ground state saved by run_caop_diag.py, by block Lanczos. + +The Python counterpart of upstream SBD's ``caop-spectrum`` app, with its option names +and its CSV output. Two modes, as upstream's: + +* generate: from a ground state saved with ``run_caop_diag.py --savename PREFIX``:: + + python run_caop_spectrum.py --hamfile ham.txt --sites 8 --loadname PREFIX \\ + --reference-energy E0 --orbitals 0,4 --channel addition \\ + --extra-detfile addition_basis.txt --save-coefficients add.coeff \\ + --output add.csv --omega-min -6 --omega-max 8 --points 141 --eta 0.1 + +* reevaluate: the saved coefficients on a new grid or broadening, with no + Hamiltonian or state:: + + python run_caop_spectrum.py --read-coefficients add.coeff --output narrow.csv \\ + --omega-min -6 --omega-max 8 --points 281 --eta 0.05 +""" + +import argparse +import sys + + +def _csv(text): + return [s for s in text.split(',') if s] + + +def parse_args(): + parser = argparse.ArgumentParser( + description="CAOP Green's function (upstream caop-spectrum)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument('--device', default='cpu', choices=['cpu', 'gpu'], + help="'cpu', or 'gpu' (NVHPC Thrust, NVIDIA GPUs only)") + parser.add_argument('--read-coefficients', default='', + help='Reevaluate these saved coefficients; needs only --output and ' + 'the grid') + parser.add_argument('--hamfile', help='Hamiltonian file the state was solved with') + parser.add_argument('--loadname', help='Prefix the ground state was saved under') + parser.add_argument('--wavefunction-shards', type=int, default=1, + help='Number of files the state was saved in (one per --b_comm_size ' + 'rank of the ground-state run)') + parser.add_argument('--sites', type=int, help="Modes, the state's --system_size") + parser.add_argument('--bit-length', type=int, default=20, + help="The state's --bit_length; the files do not record it") + parser.add_argument('--reference-energy', type=float, + help="The ground state's energy E0, as run_caop_diag.py printed it") + parser.add_argument('--orbitals', type=_csv, default=[], + help='Comma-separated modes for c / cdag') + parser.add_argument('--channel', default='addition', choices=['addition', 'removal']) + parser.add_argument('--applied-operator-type', default='single-particle', + choices=['single-particle', 'general']) + parser.add_argument('--operator-file', action='append', default=[], + help='A CAOP operator file, repeated once per observable, for ' + '--applied-operator-type general') + parser.add_argument('--extra-detfile', action='append', default=[], + help='Determinants added to the excitation basis, repeatable; ' + 'usually the whole N+1 or N-1 sector') + parser.add_argument('--steps', type=int, default=100, help='Most block-Lanczos steps') + parser.add_argument('--rank-atol', type=float, default=1e-13) + parser.add_argument('--rank-rtol', type=float, default=1e-12) + parser.add_argument('--hamiltonian-mode', default='stored', choices=['stored', 'on-the-fly'], + help="'on-the-fly' is the only mode that uses the GPU") + parser.add_argument('--h-comm-size', type=int, default=1) + parser.add_argument('--t-comm-size', type=int, default=1) + parser.add_argument('--scalar-type', default='real', choices=['real', 'complex']) + parser.add_argument('--wavefunction-type', default=None, choices=['real', 'complex'], + help='Type the state was saved with; default the scalar type') + parser.add_argument('--save-coefficients', default='', + help='Write the coefficients here, for --read-coefficients') + parser.add_argument('--output', default='', help='CSV of G and A on the grid') + parser.add_argument('--omega-min', type=float, default=-10.0) + parser.add_argument('--omega-max', type=float, default=10.0) + parser.add_argument('--points', type=int, default=401) + parser.add_argument('--eta', type=float, default=0.1) + args = parser.parse_args() + if args.read_coefficients: + if not args.output: + parser.error("--read-coefficients needs --output") + return args + missing = [f"--{n.replace('_', '-')}" for n in ('hamfile', 'loadname', 'sites', + 'reference_energy') + if getattr(args, n) is None] + if missing: + parser.error(f"generating needs {', '.join(missing)}") + general = args.applied_operator_type == 'general' + if general and (args.orbitals or not args.operator_file): + parser.error("--applied-operator-type general takes --operator-file, not --orbitals") + if not general and (args.operator_file or not args.orbitals): + parser.error("single-particle needs --orbitals (and no --operator-file)") + if not (args.save_coefficients or args.output): + parser.error("give --save-coefficients, --output, or both") + return args + + +def _write_csv(path, omegas, green, labels, general): + """Upstream's CSV: omega, a, b, then G and A, real and imaginary parts.""" + from sbd.spectrum import spectral_function + + spectral = spectral_function(green) + names = "operator_a,operator_b" if general else "orbital_a,orbital_b" + with open(path, "w") as f: + f.write(f"omega,{names},G_real,G_imag,A_real,A_imag\n") + for k, w in enumerate(omegas): + for i, a in enumerate(labels): + for j, b in enumerate(labels): + g, s = complex(green[k, i, j]), complex(spectral[k, i, j]) + f.write(",".join([f"{float(w):.17g}", str(a), str(b)] + + [f"{x:.17g}" for x in (g.real, g.imag, s.real, s.imag)]) + + "\n") + + +def main(): + args = parse_args() + import numpy as np + + import sbd + import sbd.spectrum as sp + + sbd.init(device=args.device) + rank = sbd.get_rank() + omegas = np.linspace(args.omega_min, args.omega_max, args.points) + general = args.applied_operator_type == 'general' + + if args.read_coefficients: + result = sp.green_function_from_coefficients(args.read_coefficients, omegas, + args.eta) + n = result["green_function"].shape[1] + if rank == 0: + _write_csv(args.output, omegas, result["green_function"], list(range(n)), + general=False) + print(f"Reevaluated {args.read_coefficients}: {n} x {n}, " + f"{len(omegas)} points, eta {args.eta} -> {args.output}") + return 0 + + if rank == 0 and args.device == 'gpu': + if args.hamiltonian_mode == 'stored': + print("NOTE: --hamiltonian-mode stored builds and applies the Hamiltonian on " + "the CPU only: on --device gpu this run uses no GPU. For the GPU kernel, " + "use --hamiltonian-mode on-the-fly.") + else: + print("NOTE: --hamiltonian-mode on-the-fly applies the Hamiltonian on the GPU; " + "the seed generation, QR and Lanczos vector work stay on the CPU. For a " + "small basis, compare against --device cpu with --hamiltonian-mode " + "stored.") + complex_ = args.scalar_type == 'complex' + wavefunction_complex = (args.wavefunction_type or args.scalar_type) == 'complex' + orbitals = [int(o) for o in args.orbitals] + result = sp.green_function( + args.loadname, args.hamfile, orbitals, args.channel, sites=args.sites, + bit_length=args.bit_length, reference_energy=args.reference_energy, + omegas=omegas, eta=args.eta, shards=args.wavefunction_shards, + extra_basis=args.extra_detfile, max_steps=args.steps, + h_comm_size=args.h_comm_size, t_comm_size=args.t_comm_size, + store_hamiltonian=args.hamiltonian_mode == 'stored', complex=complex_, + wavefunction_complex=wavefunction_complex, + save_coefficients=args.save_coefficients, operator_files=args.operator_file, + rank_atol=args.rank_atol, rank_rtol=args.rank_rtol, device=args.device) + if rank == 0: + labels = list(range(len(args.operator_file))) if general else orbitals + if args.output: + _write_csv(args.output, omegas, result["green_function"], labels, general) + print(f"Excitation basis: {result['basis_size']}, steps: {result['steps']}, " + f"stop: {result['stop_reason']}") + for path in (args.save_coefficients, args.output): + if path: + print(f"Wrote {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/bindings.cpp b/python/bindings.cpp index 3c647c2..a29906a 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -38,8 +38,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -1653,6 +1655,241 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { py::arg("loadname") = "", py::arg("savename") = "", py::arg("complex") = false); + + // CAOP spectrum: a Green's function of a saved CAOP ground state by block Lanczos + // (upstream apps/caop-spectrum). Upstream's pipeline, run in one call: read the + // checkpoint and the extra basis, load H, build the response coefficients with + // sbd::caop::spectrum::spectrum, then evaluate G(omega) on a grid. + m.def("caop_spectrum", + [](py::object py_comm, + const std::string& loadname, + size_t shards, + const std::string& hamfile, + const std::vector& orbitals, + bool addition, + size_t sites, + size_t bit_length, + double reference_energy, + const std::vector& extra_basis, + const std::vector& omegas, + double eta, + size_t max_steps, + size_t h_comm_size, + size_t t_comm_size, + bool store_hamiltonian, + bool complex, + bool wavefunction_complex, + const std::string& save_coefficients, + const std::vector& operator_files, + double rank_atol, + double rank_rtol) { + namespace cs = sbd::caop::spectrum; + const bool general = !operator_files.empty(); + if (wavefunction_complex && !complex) { + throw std::invalid_argument( + "caop_spectrum: a complex saved state needs complex=True"); + } + MPI_Comm comm = get_mpi_comm(py_comm); + sbd_require_openmp_threads("caop_spectrum", comm); + int world = 1, rank = 0; + MPI_Comm_size(comm, &world); + MPI_Comm_rank(comm, &rank); + if (h_comm_size == 0 || t_comm_size == 0 || world % (h_comm_size * t_comm_size)) { + throw std::invalid_argument( + "caop_spectrum: h_comm_size * t_comm_size must divide the rank count"); + } + if (general && !orbitals.empty()) { + throw std::invalid_argument( + "caop_spectrum: give orbitals (single-particle c/cdag) or operator " + "files (general operators), not both"); + } + if (!general && orbitals.empty()) { + throw std::invalid_argument( + "caop_spectrum needs at least one orbital or operator file"); + } + if (!(rank_atol >= 0) || !(rank_rtol >= 0)) { + throw std::invalid_argument( + "caop_spectrum: rank_atol and rank_rtol must be non-negative"); + } + if (sites == 0 || bit_length == 0 || bit_length > 8 * sizeof(size_t)) { + throw std::invalid_argument("caop_spectrum: invalid sites or bit_length"); + } + // Upstream reads the hamfile on h-rank 0 only, as for the solve; see + // caop_check_hamfile. The operator files are read the same way. + caop_check_hamfile(hamfile, complex, comm, "caop_spectrum"); + for (const auto& file : operator_files) { + caop_check_hamfile(file, complex, comm, "caop_spectrum"); + } +#ifdef SBD_THRUST + { + int numDevices = 1; +#ifdef __CUDACC__ + cudaGetDeviceCount(&numDevices); + cudaSetDevice(rank % numDevices); +#else + hipGetDeviceCount(&numDevices); + hipSetDevice(rank % numDevices); +#endif + } +#endif + std::vector> g; + size_t basis_size = 0, steps = 0; + std::string stop_reason; + auto run = [&](auto tag) { + using ElemT = decltype(tag); + sbd::det_vector::init_elem_size((sites + bit_length - 1) / bit_length); + const int b_size = int(world / h_comm_size / t_comm_size); + MPI_Comm h_comm, b_comm, t_comm; + sbd::setup_communicator(comm, int(h_comm_size), b_size, int(t_comm_size), + h_comm, b_comm, t_comm); + int h_rank = 0, t_rank = 0; + MPI_Comm_rank(h_comm, &h_rank); + MPI_Comm_rank(t_comm, &t_rank); + cs::Wavefunction state; + cs::Basis extra; + if (h_rank == 0 && t_rank == 0) { + // Read the checkpoint as the type it was saved with (real and complex + // files are not interchangeable), then promote a real state into a + // complex calculation, as upstream's app does. + if (wavefunction_complex) { + if constexpr (std::is_same_v>) { + state = cs::ReadWavefunctionCheckpoint( + loadname, shards, sites, bit_length, b_comm); + } + } else { + auto real_state = cs::ReadWavefunctionCheckpoint( + loadname, shards, sites, bit_length, b_comm); + state.basis = std::move(real_state.basis); + state.coefficients.assign(real_state.coefficients.begin(), + real_state.coefficients.end()); + } + if (!extra_basis.empty()) { + extra = cs::LoadBasis(extra_basis, sites, bit_length, b_comm); + } + } + sbd::GeneralOp h; + bool fermion = false; + sbd::load_GeneralOp_from_file(hamfile, h, fermion, h_comm, b_comm, t_comm); + sbd::sparse_solver::LanczosOptions options; + options.max_steps = max_steps; + options.rank_atol = rank_atol; + options.rank_rtol = rank_rtol; + // General operators, as upstream's app: each file is a CAOP operator, + // read on the h=t=0 ranks with the statistics of H, observables 0..n-1, + // and the addition-sign resolvent . + std::vector> applied; + std::vector observables = orbitals; + if (general) { + observables.resize(operator_files.size()); + std::iota(observables.begin(), observables.end(), size_t{0}); + if (h_rank == 0 && t_rank == 0) { + applied.resize(operator_files.size()); + for (size_t v = 0; v < applied.size(); ++v) { + bool op_fermion = false; + sbd::load_GeneralOp_from_file(operator_files[v], applied[v], + op_fermion, MPI_COMM_SELF, b_comm, + MPI_COMM_SELF); + if (op_fermion != fermion) { + throw std::invalid_argument( + "caop_spectrum: operator file " + operator_files[v] + + " has different statistics from the Hamiltonian"); + } + } + } + } + auto result = cs::spectrum(state, h, fermion, observables, + general || addition, sites, bit_length, + extra, reference_energy, h_comm, b_comm, + t_comm, options, + general ? &applied : nullptr, + store_hamiltonian); + result.operator_files = operator_files; + MPI_Comm_free(&t_comm); + MPI_Comm_free(&b_comm); + MPI_Comm_free(&h_comm); + basis_size = result.basis_size; + steps = result.coefficients.A.size(); + stop_reason = result.coefficients.stop_reason; + if (rank == 0 && !save_coefficients.empty()) { + cs::save_coefficients(save_coefficients, result, options); + } + const size_t n = observables.size(); + g.assign(omegas.size() * n * n, 0.0); + for (size_t k = 0; k < omegas.size(); ++k) { + const auto gk = cs::green_function(result, omegas[k], eta); + std::copy(gk.begin(), gk.end(), g.begin() + k * n * n); + } + }; + { + py::gil_scoped_release release; + if (complex) run(std::complex{}); + else run(double{}); + } + const size_t n = general ? operator_files.size() : orbitals.size(); + py::array_t> out(std::vector{ + static_cast(omegas.size()), static_cast(n), + static_cast(n)}); + std::copy(g.begin(), g.end(), out.mutable_data()); + py::dict results; + results["green_function"] = out; + results["basis_size"] = basis_size; + results["steps"] = steps; + results["stop_reason"] = stop_reason; + return results; + }, + "Green's function of a saved CAOP ground state by block Lanczos, as upstream's " + "caop-spectrum app", + py::arg("comm"), py::arg("loadname"), py::arg("shards"), py::arg("hamfile"), + py::arg("orbitals"), py::arg("addition"), py::arg("sites"), py::arg("bit_length"), + py::arg("reference_energy"), py::arg("extra_basis") = std::vector{}, + py::arg("omegas") = std::vector{}, py::arg("eta") = 0.05, + py::arg("max_steps") = 100, py::arg("h_comm_size") = 1, py::arg("t_comm_size") = 1, + py::arg("store_hamiltonian") = true, py::arg("complex") = false, + py::arg("wavefunction_complex") = false, py::arg("save_coefficients") = "", + py::arg("operator_files") = std::vector{}, + py::arg("rank_atol") = 1e-13, py::arg("rank_rtol") = 1e-12); + + // Upstream's --read-coefficients: evaluate saved block-Lanczos coefficients on a new + // frequency grid or broadening, without the Hamiltonian or the state. File-only and + // rank-local; the scalar type comes from the file's header. + m.def("caop_spectrum_from_coefficients", + [](const std::string& path, const std::vector& omegas, double eta) { + namespace cs = sbd::caop::spectrum; + std::vector> g; + size_t n = 0, basis_size = 0, steps = 0; + bool addition = true; + std::string stop_reason; + auto evaluate = [&](auto tag) { + using ElemT = decltype(tag); + auto result = cs::load_coefficients(path); + n = result.orbitals.size(); + basis_size = result.basis_size; + steps = result.coefficients.A.size(); + stop_reason = result.coefficients.stop_reason; + addition = result.addition; + g.assign(omegas.size() * n * n, 0.0); + for (size_t k = 0; k < omegas.size(); ++k) { + const auto gk = cs::green_function(result, omegas[k], eta); + std::copy(gk.begin(), gk.end(), g.begin() + k * n * n); + } + }; + if (cs::complex_coefficient_file(path)) evaluate(std::complex{}); + else evaluate(double{}); + py::array_t> out(std::vector{ + static_cast(omegas.size()), static_cast(n), + static_cast(n)}); + std::copy(g.begin(), g.end(), out.mutable_data()); + py::dict results; + results["green_function"] = out; + results["basis_size"] = basis_size; + results["steps"] = steps; + results["stop_reason"] = stop_reason; + results["addition"] = addition; + return results; + }, + "Green's function from coefficients saved by caop_spectrum, on a new grid, as " + "upstream's caop-spectrum --read-coefficients", + py::arg("path"), py::arg("omegas"), py::arg("eta")); #else m.attr("__sbd_has_caop__") = false; #endif diff --git a/python/spectrum.py b/python/spectrum.py new file mode 100644 index 0000000..bb9c0f5 --- /dev/null +++ b/python/spectrum.py @@ -0,0 +1,144 @@ +# 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. + +"""Green's functions of a CAOP ground state, by block Lanczos. + +The Python counterpart of upstream SBD's +`caop-spectrum `_ app. +The workflow, for a lattice model such as the Hubbard model: + +1. Write the model as a CAOP Hamiltonian file, and solve for its ground state with + :func:`sbd.caop_diag_from_files`, saving it with ``savename``. +2. Call :func:`green_function` with that state, the operator modes to add an electron + to (``channel="addition"``) or remove one from (``"removal"``), and a frequency + grid. It applies the creation or annihilation operators to the state and runs block + Lanczos in the excitation basis. +3. :func:`spectral_function` turns the result into the spectral function, whose peaks + are the addition or removal energies and whose weights are the spectral weights. + +The conventions are upstream's, for addition and removal respectively:: + + G+_ab(w) = + G-_ab(w) = + +Add the two channels on a common grid for the fermionic retarded Green's function. + +**The excitation basis.** The basis the Lanczos runs in is the union of the +determinants the operators generate from the state and any ``extra_basis`` files. +Upstream warns that the generated determinants alone are usually not enough: neither +mode completes the particle sector by itself, so give ``extra_basis`` the determinants +of the N+1 (addition) or N-1 (removal) sector you want included. +""" + +from __future__ import annotations + +import os + +import numpy as np + +import sbd + +__all__ = ["green_function", "green_function_from_coefficients", "spectral_function"] + + +def green_function(loadname, hamfile, orbitals, channel, *, sites, bit_length, + reference_energy, omegas, eta=0.05, shards=1, extra_basis=(), + max_steps=100, h_comm_size=1, t_comm_size=1, store_hamiltonian=True, + complex=False, wavefunction_complex=False, save_coefficients="", + operator_files=(), rank_atol=1e-13, rank_rtol=1e-12, device=None): + """The Green's function of a saved CAOP ground state on a frequency grid. + + Args: + loadname: The ``savename`` prefix the ground state was saved under, by + :func:`sbd.caop_diag_from_files` (files ``loadname000000.bin`` ...). + hamfile: The CAOP Hamiltonian file the state was solved with. + orbitals: Modes to apply the operators to (0-based CAOP mode indices). Empty + when ``operator_files`` is given. + channel: ``"addition"`` (creation operators, the N+1 sector) or + ``"removal"`` (annihilation operators, the N-1 sector). Ignored for + general operators, whose response is + ````. + sites: Number of modes, as the state's ``system_size``. + bit_length: Bits per packed word, as the state was saved with. + reference_energy: The ground state's energy E0, as the solve returned it. + omegas: Frequencies to evaluate the Green's function at. + eta: Broadening. + shards: Number of files the state was saved in. + extra_basis: Files of determinants to add to the excitation basis, one + bitstring of ``sites`` characters per line; see the module notes. + max_steps: Most block-Lanczos steps. + h_comm_size, t_comm_size: Helper and task communicator sizes; their product + must divide the rank count, and the rest of the ranks hold the basis. + store_hamiltonian: Store H for faster application (upstream's ``stored`` + mode), rather than applying it on the fly. + complex: Run with complex amplitudes, for a complex Hamiltonian. + wavefunction_complex: The saved state has complex amplitudes, as one solved + with ``caop_diag_from_files(..., complex=True)`` does; needs ``complex``. + Real and complex files are not interchangeable, so give the type the state + was saved with. A real state is promoted when ``complex`` is set. + save_coefficients: If given, write upstream's coefficient file here, from + which :func:`green_function_from_coefficients` re-evaluates the spectrum on + another grid or broadening without rerunning. + operator_files: CAOP operator files to apply instead of ``c``/``cdag`` on + ``orbitals`` (upstream's ``--applied-operator-type general``), one per + observable, with the Hamiltonian's statistics. They may be non-Hermitian + or change the particle number. + rank_atol, rank_rtol: Block-Lanczos deflation tolerances: a column whose norm + falls below them is dropped. Upstream's defaults. + device: ``"cpu"`` or ``"gpu"`` (Thrust); there is no OpenMP-offload path. + + Returns: + dict with ``green_function``, a complex array of shape + ``(len(omegas), n, n)`` for ``n`` orbitals or operators; ``basis_size``, the + excitation basis's size; ``steps``; and ``stop_reason``. + """ + if channel not in ("addition", "removal"): + raise ValueError(f"channel must be 'addition' or 'removal', got {channel!r}") + if isinstance(extra_basis, (str, os.PathLike)): + extra_basis = [extra_basis] + if isinstance(operator_files, (str, os.PathLike)): + operator_files = [operator_files] + backend = sbd._caop_backend(device) + return backend.caop_spectrum( + sbd._global_comm, str(loadname), int(shards), str(hamfile), + [int(o) for o in orbitals], channel == "addition", int(sites), int(bit_length), + float(reference_energy), [str(f) for f in extra_basis], + [float(w) for w in np.atleast_1d(omegas)], float(eta), int(max_steps), + int(h_comm_size), int(t_comm_size), bool(store_hamiltonian), bool(complex), + bool(wavefunction_complex), str(save_coefficients), + [str(f) for f in operator_files], float(rank_atol), float(rank_rtol)) + + +def green_function_from_coefficients(path, omegas, eta, *, device=None): + """The Green's function from coefficients saved by :func:`green_function`. + + Upstream's ``--read-coefficients``: the block-Lanczos coefficients fix the + response, so a new frequency grid or broadening needs neither the Hamiltonian nor + the state. Rank-local. It cannot extend the Lanczos run itself. + + Returns: + dict with ``green_function`` (``(len(omegas), n, n)``), ``basis_size``, + ``steps``, ``stop_reason`` and ``addition`` (the channel it was saved for). + """ + backend = sbd._caop_backend(device) + return backend.caop_spectrum_from_coefficients( + str(path), [float(w) for w in np.atleast_1d(omegas)], float(eta)) + + +def spectral_function(green): + """The spectral matrix ``-(G - G^dag) / (2 pi i)`` at each frequency. + + Its diagonal is ``-Im(G_aa) / pi``: peaks at the addition or removal energies, + with the spectral weights as their areas. + """ + g = np.asarray(green) + return -(g - np.conj(np.swapaxes(g, -1, -2))) / (2j * np.pi) diff --git a/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml b/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml new file mode 100644 index 0000000..70990d7 --- /dev/null +++ b/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml @@ -0,0 +1,19 @@ +--- +features: + - | + New :mod:`sbd.spectrum` binds upstream SBD's ``caop-spectrum``: the Green's + function of a CAOP ground state by block Lanczos. :func:`sbd.spectrum.green_function` + takes a state saved by :func:`sbd.caop_diag_from_files`, the modes to add an electron + to or remove one from, and a frequency grid, and returns the Green's function on it; + :func:`sbd.spectrum.spectral_function` gives the spectral function. The conventions + are upstream's. Real and complex amplitudes are both available. Like CAOP, it runs on + the ``cpu`` and ``gpu`` (Thrust) backends. + + As upstream's app, it also takes general operators instead of ``c``/``cdag`` + (``operator_files``) and the block-Lanczos deflation tolerances, and + :func:`sbd.spectrum.green_function_from_coefficients` reevaluates saved + coefficients on a new frequency grid or broadening without rerunning. + - | + New ``examples/caop/run_caop_spectrum.py`` runs it from the command line with + upstream's option names and CSV output, on a state saved by + ``run_caop_diag.py --savename``. diff --git a/test/test_spectrum.py b/test/test_spectrum.py new file mode 100644 index 0000000..91f86e6 --- /dev/null +++ b/test/test_spectrum.py @@ -0,0 +1,358 @@ +# 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. + +"""``sbd.spectrum``: the Green's function of a CAOP ground state, against an exact one. + +The model is upstream's 4-site Hubbard ring (t = -1, U = 4, two electrons of each spin), +the one ``test_caop.py`` solves. Its ground state is found with CAOP and saved, then the +Green's function is computed with block Lanczos, with the full N+1 or N-1 sector given as +``extra_basis``, as upstream's notes say it should be. The reference is the exact +Lehmann sum of upstream's formulas, from a dense diagonalization in NumPy of the same +operator terms. Each run is a **subprocess**, since CAOP fixes the process-wide word count. +""" + +from __future__ import annotations + +import itertools +import json +import os +import pathlib +import subprocess +import sys + +import numpy as np +import pytest + +import sbd + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from test_caop import _apply, _hubbard, _read_terms # noqa: E402 + +SITES, BITS, ETA = 4, 20, 0.1 +MODES = 2 * SITES +ORBITALS = [0, 1, SITES] # alpha sites 0 and 1, beta site 0 +OMEGAS = np.linspace(-6.0, 8.0, 15) + +RUN = r""" +import json, sys +import numpy as np +import sbd, sbd.spectrum as sp +a = json.loads(sys.argv[1]) +sbd.init(device=a["device"]) +config = sbd.CAOP_SBD() +config.system_size, config.bit_length = a["modes"], a["bits"] +config.method, config.max_it, config.max_nb, config.eps = 1, 200, 20, 1e-13 +ground = sbd.caop_diag_from_files(a["ham"], [a["basis"]], config, savename=a["state"], + complex=a["complex"]) +out = {"energy": ground["energy"]} +for channel in ("addition", "removal"): + r = sp.green_function(a["state"], a["ham"], a["orbitals"], channel, sites=a["modes"], + bit_length=a["bits"], reference_energy=ground["energy"], + omegas=a["omegas"], eta=a["eta"], extra_basis=[a[channel]], + max_steps=200, complex=a["complex"], + wavefunction_complex=a["complex"], + save_coefficients=a["save"] + channel) + g = r["green_function"] + out[channel] = {"real": g.real.tolist(), "imag": g.imag.tolist(), + "basis_size": r["basis_size"], "stop_reason": r["stop_reason"]} +print("RESULT " + json.dumps(out)) +""" + + +def _states(n_alpha, n_beta): + half = lambda n: [sum(1 << p for p in occ) for occ in itertools.combinations(range(SITES), n)] + return sorted(a | (b << SITES) for a in half(n_alpha) for b in half(n_beta)) + + +def _matrix(terms, fermion, states): + index = {s: i for i, s in enumerate(states)} + h = np.zeros((len(states), len(states))) + for j, s in enumerate(states): + for coefficient, ops in terms: + sign, t = _apply(s, ops, fermion) + if t is not None and t in index: + h[index[t], j] += coefficient.real * sign + return h + + +def _exact_green(ham, eta=ETA): + """Upstream's G+ and G- by the Lehmann sum over the N+1 and N-1 sectors.""" + fermion, terms = _read_terms(ham) + ground_states = _states(2, 2) + e0_all, v0_all = np.linalg.eigh(_matrix(terms, fermion, ground_states)) + e0, psi = e0_all[0], v0_all[:, 0] + + def apply_op(op, mode, vector, source, target): + out = np.zeros(len(target)) + index = {s: i for i, s in enumerate(target)} + for amplitude, s in zip(vector, source): + sign, t = _apply(s, [(op, mode)], fermion) + if t is not None: + out[index[t]] += sign * amplitude + return out + + result = {} + for channel, sector, op, sgn in (("addition", (3, 2), "cdag", +1), + ("removal", (2, 1), "c", -1)): + # Removal from beta modes lands in (2, 1); from alpha modes in (1, 2). Use the + # union of both N-1 sectors, and of both N+1 sectors, so every orbital fits. + states = sorted(set(_states(*sector)) | set(_states(*sector[::-1]))) + energies, vectors = np.linalg.eigh(_matrix(terms, fermion, states)) + projected = np.array([vectors.T @ apply_op(op, m, psi, ground_states, states) + for m in ORBITALS]) # (orbital, eigenstate) + g = np.zeros((len(OMEGAS), len(ORBITALS), len(ORBITALS)), dtype=complex) + for k, w in enumerate(OMEGAS): + if channel == "addition": # + denominator = w + e0 - energies + 1j * eta + else: # + denominator = w - e0 + energies + 1j * eta + g[k] = (projected.conj() / denominator) @ projected.T + if channel == "removal": + g[k] = g[k].T + result[channel] = (g, states) + return e0, result + + +@pytest.fixture +def caop_device(device_config): + device = device_config.device if device_config is not None else "cpu" + if device == "gpu-omp": + pytest.skip("upstream SBD's CAOP code has no OpenMP-offload path") + return device + + +def _run(caop_device, tmp_path, complex_=False): + ham, basis, _ = _hubbard(tmp_path) + e0, exact = _exact_green(ham) + sectors = {} + for channel, (_, states) in exact.items(): + path = tmp_path / f"{channel}_sector.txt" + path.write_text("\n".join(format(s, f"0{MODES}b") for s in states) + "\n") + sectors[channel] = str(path) + args = {"device": caop_device, "ham": str(ham), "basis": str(basis), "modes": MODES, + "bits": BITS, "state": str(tmp_path / "gs-"), "orbitals": ORBITALS, + "omegas": OMEGAS.tolist(), "eta": ETA, "complex": complex_, + "save": str(tmp_path / "coefficients-"), **sectors} + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run([sys.executable, "-c", RUN, json.dumps(args)], + capture_output=True, text=True, timeout=600, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + out = json.loads([l for l in done.stdout.splitlines() + if l.startswith("RESULT ")][-1][len("RESULT "):]) + return e0, exact, out + + +def test_green_function_matches_the_exact_lehmann_sum(caop_device, tmp_path): + e0, exact, out = _run(caop_device, tmp_path) + assert out["energy"] == pytest.approx(e0, abs=1e-9) + for channel in ("addition", "removal"): + got = np.array(out[channel]["real"]) + 1j * np.array(out[channel]["imag"]) + np.testing.assert_allclose(got, exact[channel][0], atol=1e-7, err_msg=channel) + assert (tmp_path / f"coefficients-{channel}").is_file() + + +def test_saved_coefficients_reevaluate_without_rerunning(caop_device, tmp_path): + """Upstream's --read-coefficients: the saved coefficients give the same G on the + same grid, and the exact G on a new broadening, with no Hamiltonian or state.""" + from sbd.spectrum import green_function_from_coefficients + + _, exact, out = _run(caop_device, tmp_path) + for channel in ("addition", "removal"): + path = tmp_path / f"coefficients-{channel}" + same = green_function_from_coefficients(path, OMEGAS, ETA, device=caop_device) + got = np.array(out[channel]["real"]) + 1j * np.array(out[channel]["imag"]) + np.testing.assert_allclose(same["green_function"], got, atol=1e-12) + assert same["addition"] is (channel == "addition") + # A new broadening, against the Lehmann sum at that broadening. + wider = green_function_from_coefficients(tmp_path / "coefficients-addition", OMEGAS, + 0.3, device=caop_device) + np.testing.assert_allclose(wider["green_function"], + _exact_green(_hubbard(tmp_path)[0], eta=0.3)[1]["addition"][0], + atol=1e-7) + + +def test_complex_amplitudes_give_the_same_green_function(caop_device, tmp_path): + """Solved and saved with complex amplitudes, the state is read back as complex.""" + _, exact, out = _run(caop_device, tmp_path, complex_=True) + for channel in ("addition", "removal"): + got = np.array(out[channel]["real"]) + 1j * np.array(out[channel]["imag"]) + np.testing.assert_allclose(got, exact[channel][0], atol=1e-7, err_msg=channel) + + +def test_spectral_function_is_minus_im_g_over_pi_on_the_diagonal(): + from sbd.spectrum import spectral_function + + rng = np.random.default_rng(0) + g = rng.normal(size=(5, 3, 3)) + 1j * rng.normal(size=(5, 3, 3)) + a = spectral_function(g) + np.testing.assert_allclose(np.diagonal(a, axis1=1, axis2=2).real, + -np.diagonal(g, axis1=1, axis2=2).imag / np.pi) + np.testing.assert_allclose(a, np.conj(np.swapaxes(a, 1, 2))) # Hermitian + + +def test_a_bad_channel_is_refused(backend): + from sbd.spectrum import green_function + + with pytest.raises(ValueError, match="channel"): + green_function("x-", "h.txt", [0], "both", sites=4, bit_length=20, + reference_energy=0.0, omegas=[0.0]) + + +def test_the_offload_backend_has_no_spectrum(): + if "gpu-omp" not in sbd.available_backends(): + pytest.skip("the OpenMP-offload backend was not built") + from sbd.spectrum import green_function + + with pytest.raises(ValueError, match="no CAOP solver"): + green_function("x-", "h.txt", [0], "addition", sites=4, bit_length=20, + reference_energy=0.0, omegas=[0.0], device="gpu-omp") + + +# --- general operators: upstream's spin-dimer example --------------------------------- + +DIMER = (pathlib.Path(__file__).resolve().parents[1] / "vendor" / "sbd-upstream" / "apps" + / "caop-spectrum" / "examples" / "spin_dimer") + +GENERAL = r""" +import json, sys +import numpy as np +import sbd, sbd.spectrum as sp +a = json.loads(sys.argv[1]) +sbd.init(device=a["device"]) +# The singlet of H = S0.S1 over {01, 10}, the state upstream's write_state.cc saves. +config = sbd.CAOP_SBD() +config.system_size, config.bit_length = 2, 20 +config.method, config.max_it, config.max_nb, config.eps = 1, 100, 2, 1e-13 +config.max_iv = 1 +ground = sbd.caop_diag_from_files(a["ham"], [a["basis"]], config, savename=a["state"]) +r = sp.green_function(a["state"], a["ham"], [], "addition", sites=2, bit_length=20, + reference_energy=ground["energy"], omegas=a["omegas"], eta=a["eta"], + operator_files=a["ops"], max_steps=8, complex=True) +g = r["green_function"] +print("RESULT " + json.dumps({"energy": ground["energy"], "real": g.real.tolist(), + "imag": g.imag.tolist()})) +""" + + +def test_general_operators_give_upstreams_spin_dimer_answer(caop_device, tmp_path): + """Upstream's example: S0.S1 with Sz and Sy on each site, from the singlet. + + G_aa = 1 / (4 (w - 1 + i eta)) for Sz/Sz and Sy/Sy, the opposite-site element its + negative, and the Sz/Sy cross elements zero (upstream's verify.py). + """ + basis = tmp_path / "basis.txt" + basis.write_text("01\n10\n") + omegas, eta = np.linspace(0.0, 2.0, 11), 0.05 + args = {"device": caop_device, "ham": str(DIMER / "ham.txt"), "basis": str(basis), + "state": str(tmp_path / "state-"), "omegas": omegas.tolist(), "eta": eta, + "ops": [str(DIMER / f) for f in ("sz0.txt", "sz1.txt", "sy0.txt", "sy1.txt")]} + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + done = subprocess.run([sys.executable, "-c", GENERAL, json.dumps(args)], + capture_output=True, text=True, timeout=600, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + out = json.loads([l for l in done.stdout.splitlines() + if l.startswith("RESULT ")][-1][len("RESULT "):]) + assert out["energy"] == pytest.approx(-0.75, abs=1e-12) + g = np.array(out["real"]) + 1j * np.array(out["imag"]) + same = 1 / (4 * (omegas - 1 + 1j * eta)) + expected = np.zeros((len(omegas), 4, 4), dtype=complex) + for block in (0, 2): # (Sz0, Sz1) and (Sy0, Sy1) + expected[:, block, block] = expected[:, block + 1, block + 1] = same + expected[:, block, block + 1] = expected[:, block + 1, block] = -same + np.testing.assert_allclose(g, expected, atol=1e-10) + + +def test_orbitals_and_operator_files_together_are_refused(backend): + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + from sbd.spectrum import green_function + + with pytest.raises(ValueError, match="not both"): + green_function("x-", DIMER / "ham.txt", [0], "addition", sites=2, bit_length=20, + reference_energy=0.0, omegas=[0.0], operator_files=[DIMER / "sz0.txt"]) + with pytest.raises(ValueError, match="non-negative"): + green_function("x-", DIMER / "ham.txt", [0], "addition", sites=2, bit_length=20, + reference_energy=0.0, omegas=[0.0], rank_atol=-1.0) + + +# --- the driver: the README's flow --------------------------------------------------- + +def test_the_drivers_flow_ground_state_then_green_function(caop_device, tmp_path): + """run_caop_diag.py --savename, then run_caop_spectrum.py on the saved state, and + --read-coefficients on its coefficients: the exact G+ of the 2-site Hubbard model.""" + caop = pathlib.Path(__file__).resolve().parents[1] / "examples" / "caop" + ham = tmp_path / "hub2.txt" + ham.write_text("-1\n-1.0 cdag 0 c 1\n-1.0 cdag 1 c 0\n-1.0 cdag 2 c 3\n" + "-1.0 cdag 3 c 2\n4.0 cdag 0 cdag 2 c 2 c 0\n4.0 cdag 1 cdag 3 c 3 c 1\n") + (tmp_path / "hub2_basis.txt").write_text("0101\n0110\n1001\n1010\n") + (tmp_path / "add_basis.txt").write_text("0111\n1011\n1101\n1110\n") + env = {k: v for k, v in os.environ.items() if not k.startswith(("PMI_", "PMIX_", "OMPI_"))} + env.setdefault("OMP_NUM_THREADS", "2") + + def run(script, *flags): + done = subprocess.run([sys.executable, str(caop / script), "--device", caop_device, + *flags], cwd=tmp_path, capture_output=True, text=True, + timeout=600, env=env) + assert done.returncode == 0, f"{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + return done.stdout + + ground = run("run_caop_diag.py", "--hamfile", "hub2.txt", "--basisfiles", "hub2_basis.txt", + "--system_size", "4", "--tolerance", "1e-10", "--block", "3", + "--savename", "gs-") + e0 = float(ground.split("Ground state energy:")[1].split()[0]) + assert e0 == pytest.approx((4 - np.sqrt(16 + 16)) / 2, abs=1e-9) # (U - sqrt(U^2+16t^2))/2 + grid = ["--omega-min", "-6", "--omega-max", "8", "--points", "15"] + run("run_caop_spectrum.py", "--hamfile", "hub2.txt", "--sites", "4", "--loadname", "gs-", + "--reference-energy", repr(e0), "--orbitals", "0,2", "--channel", "addition", + "--extra-detfile", "add_basis.txt", "--save-coefficients", "add.coeff", + "--output", "add.csv", *grid, "--eta", "0.1") + run("run_caop_spectrum.py", "--read-coefficients", "add.coeff", "--output", "narrow.csv", + *grid, "--eta", "0.05") + + terms = _read_terms(ham) + fermion = terms[0] + ground_states, add_states = [5, 6, 9, 10], [7, 11, 13, 14] + + def matrix(states): + index = {x: i for i, x in enumerate(states)} + h = np.zeros((len(states), len(states))) + for j, x in enumerate(states): + for c, ops in terms[1]: + sign, y = _apply(x, ops, fermion) + if y is not None and y in index: + h[index[y], j] += c.real * sign + return h + + e, v = np.linalg.eigh(matrix(ground_states)) + psi = v[:, 0] + ea, va = np.linalg.eigh(matrix(add_states)) + index = {x: i for i, x in enumerate(add_states)} + + def seed(mode): + out = np.zeros(len(add_states)) + for amplitude, x in zip(psi, ground_states): + sign, y = _apply(x, [("cdag", mode)], fermion) + if y is not None: + out[index[y]] += sign * amplitude + return va.T @ out + + projected = np.array([seed(0), seed(2)]) + omegas = np.linspace(-6, 8, 15) + for path, eta in (("add.csv", 0.1), ("narrow.csv", 0.05)): + rows = np.loadtxt(tmp_path / path, delimiter=",", skiprows=1) + got = (rows[:, 3] + 1j * rows[:, 4]).reshape(len(omegas), 2, 2) + expected = np.array([(projected.conj() / (w + e[0] - ea + 1j * eta)) @ projected.T + for w in omegas]) + np.testing.assert_allclose(got, expected, atol=1e-8, err_msg=path) + From d9af4e0190e80281be8c64d9b98360bb1ffbb50e Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 17:12:41 -0500 Subject: [PATCH 06/16] gdb_diag: refuse an empty Thrust shard after placement Upstream's Thrust GDB solve waits forever in its ring exchange when a b_comm shard is empty. Placement can empty a shard the caller filled: equal-bra-a, the default, deals out whole alpha strings, so fewer distinct alpha strings than b_comm_size leaves ranks with nothing. Check every rank after placement and raise on all of them together. Co-Authored-By: Claude Opus 5.5 --- python/bindings.cpp | 20 ++++++++++++++++++++ test/test_gdb_equivalence.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/python/bindings.cpp b/python/bindings.cpp index d88f275..500d22f 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -1054,6 +1054,26 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { } } +#ifdef SBD_THRUST + // ---- Thrust: every b shard must hold something, after placement ------- + // Upstream's Thrust mult never returns when a b shard is empty: the ring + // exchange (MpiSlider::Sync, gdb/helper_thrust.h) waits forever, while the + // CPU backend handles the same input. Placement can empty a shard the + // input filled: equal-bra-a deals out whole alpha strings, so fewer + // distinct alpha strings than b_comm_size leaves ranks with nothing. + if (b_comm_size > 1 && !sbd_all_ranks_ok(det.size() > 0, comm)) { + throw std::invalid_argument( + std::string("gdb_diag on the gpu (Thrust) backend needs determinants " + "on every b_comm shard after placement (determinant_" + "distribution='") + + sbd_det_distribution_name(distribution) + + "'), but at least one is empty: upstream's Thrust solve then waits " + "forever. Use fewer shards (a smaller b_comm_size), " + "determinant_distribution='input' with a non-empty slice per rank, " + "or device='cpu'"); + } +#endif + // ---- global dimension ------------------------------------------- // Summed over b_comm, since h_comm and t_comm hold replicas. unsigned long long local_dim = static_cast(det.size()); diff --git a/test/test_gdb_equivalence.py b/test/test_gdb_equivalence.py index 9b475ac..d6e5ac5 100644 --- a/test/test_gdb_equivalence.py +++ b/test/test_gdb_equivalence.py @@ -548,6 +548,36 @@ def test_gdb_placement_does_not_change_the_energy_mpi(backend, h2o, h2o_product, _check_placement_does_not_change_the_energy(backend, h2o, h2o_product, scheme) +@pytest.mark.mpi +def test_thrust_refuses_a_shard_that_placement_empties(backend, h2o): + """equal-bra-a deals out whole alpha strings, so one alpha string paired with many + betas lands on a single rank, though every rank passed determinants in. Upstream's + Thrust solve would then wait forever in its ring exchange; every rank refuses + instead. The CPU backend solves the same input. + """ + from mpi4py import MPI + + comm = MPI.COMM_WORLD + size = comm.Get_size() + if size < 2: + pytest.skip("needs 2 ranks or more") + fcidump, norb, _, alpha = h2o + product = sbd.sort_bitarray_array(sbd.from_strings( + [_interleave(alpha[0], b) for b in alpha], BIT_LENGTH, 2 * norb)) + shard = _shard(product, size, comm.Get_rank()) + if not _is_thrust(backend): + _gdb_energy(backend, fcidump, shard, norb, b_comm_size=size, + determinant_distribution="equal-bra-a") + return + refused = 0 + try: + _gdb_energy(backend, fcidump, shard, norb, b_comm_size=size, + determinant_distribution="equal-bra-a") + except ValueError as error: + refused = int("empty" in str(error)) + assert comm.allreduce(refused) == size + + @pytest.mark.mpi @pytest.mark.parametrize("scheme", ["count", "count-sorted"]) def test_gdb_count_schemes_rebalance_uneven_shards(backend, h2o, h2o_product, scheme): From 94a5dc9eade70f8dd89834dae5df3b14026aba03 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 17:45:33 -0500 Subject: [PATCH 07/16] CAOP solve: check the modes against system_size; fix the OpenMP team - caop_check_hamfile also reads each term's modes and the statistics, and caop_check_modes refuses a mode at or beyond system_size, which upstream's solve never checks. The spectrum and the variance reuse both. - omp_set_dynamic(0) before the solve, as upstream's caop-spectrum app: the stored Hamiltonian (methods 1/3) keeps per-thread slots built by a fixed team. - The Thrust t > b comment names the real cause: the mult returns early on an empty slide and skips the Allreduce the other ranks wait in. - README: the block runs from this folder, without cd. Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 1 - python/bindings.cpp | 50 ++++++++++++++++++++++++++++++++++++----- test/test_caop.py | 4 +++- 3 files changed, 47 insertions(+), 8 deletions(-) diff --git a/examples/caop/README.md b/examples/caop/README.md index 55ac100..03ed117 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -12,7 +12,6 @@ for each one and for its `gen_ham.py` and `gen_bits.py`. This page covers what i specific to the Python driver and what we found running it. ```bash -cd examples/caop mpirun -np 4 python run_caop_diag.py --b_comm_size 4 # upstream's 8-site boson chain ``` diff --git a/python/bindings.cpp b/python/bindings.cpp index f4b69e0..d5b6526 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -249,9 +249,15 @@ static std::string sbd_first_problem(const std::string& local, MPI_Comm comm) { * coefficient as a double, so "(re,im)" is refused there with a pointer to * complex=True. */ -static void caop_check_hamfile(const std::string& hamfile, bool complex, MPI_Comm comm, - const char* where) { +struct CaopFileFacts { + bool fermion = false; // statistics zero or negative, as upstream's loader reads it + size_t max_mode = 0; // highest mode index any term uses +}; + +static CaopFileFacts caop_check_hamfile(const std::string& hamfile, bool complex, + MPI_Comm comm, const char* where) { std::string problem; + CaopFileFacts facts; std::ifstream ifs(hamfile); if (!ifs) { problem = "cannot open the Hamiltonian file " + hamfile; @@ -271,6 +277,7 @@ static void caop_check_hamfile(const std::string& hamfile, bool complex, MPI_Com + ": the first line must be the statistics (an integer: positive " "for hard-core bosons or spins, zero or negative for fermions)"; } + facts.fermion = statistics <= 0; statistics_read = true; continue; } @@ -282,6 +289,14 @@ static void caop_check_hamfile(const std::string& hamfile, bool complex, MPI_Com double c; read = static_cast(iss >> c); } + // The operator tokens, each followed by its mode, up to a '#' comment. + std::string op; + size_t mode = 0; + while (read && iss >> op && op[0] != '#') { + if (iss >> mode) facts.max_mode = std::max(facts.max_mode, mode); + else problem = "line " + std::to_string(number) + " of " + hamfile + + ": operator " + op + " has no mode index after it"; + } if (!read) { const bool looks_complex = line[first] == '('; problem = "line " + std::to_string(number) + " of " + hamfile @@ -301,6 +316,22 @@ static void caop_check_hamfile(const std::string& hamfile, bool complex, MPI_Com if (!first.empty()) { throw std::invalid_argument(std::string(where) + ": " + first); } + return facts; +} + +/** + * Refuse, on every rank, a CAOP file whose terms use a mode at or beyond sites. + * + * Upstream checks max_index() on the h-sharded operator, so with h_comm_size > 1 only + * the shard holding the high mode throws and the others hang. + */ +static void caop_check_modes(const std::string& file, const CaopFileFacts& facts, + size_t sites, const char* where) { + if (facts.max_mode >= sites) { + throw std::invalid_argument( + std::string(where) + ": " + file + " uses mode " + std::to_string(facts.max_mode) + + ", beyond sites=" + std::to_string(sites) + " (modes are 0 .. sites-1)"); + } } #endif @@ -1607,9 +1638,9 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { #ifdef SBD_THRUST // make_slide deals the b ring stations out over the t task ranks // (caop/basic/helper.h:41-55), so with t > b some ranks get none. The CPU - // mult guards an empty slide (mult.h:19), but the Thrust one reads - // slide_[0] unguarded (mult_thrust.h:628) and the job hangs; the - // stored-Hamiltonian methods run the CPU mult and are unaffected. The + // mult guards an empty slide (mult.h:19); the Thrust one returns early + // (mult_thrust.h:583) and skips the Allreduce the other ranks wait in, so + // the job hangs; the stored-Hamiltonian methods run the CPU mult. The // same t <= b holds for GDB on every backend. if ((sbd_data.method == 0 || sbd_data.method == 2) && t_comm_size > b_comm_size) { throw std::invalid_argument( @@ -1639,7 +1670,14 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { // Upstream reads the hamfile on h-rank 0 only and sends the terms on // (caop/basic/loadmodel.h:92-96), so a file it cannot read hangs the other // ranks; see caop_check_hamfile. - caop_check_hamfile(hamfile, complex, comm, "caop_diag_from_files"); + // Upstream's solve does not check the terms' modes against system_size. + caop_check_modes(hamfile, + caop_check_hamfile(hamfile, complex, comm, + "caop_diag_from_files"), + sbd_data.system_size, "caop_diag_from_files"); + // As upstream's caop-spectrum app: the stored Hamiltonian (methods 1/3) keeps + // per-thread slots that need a fixed OpenMP team (caop/basic/mkham.h). + omp_set_dynamic(0); double energy = 0.0; sbd::det_vector co_basis; { diff --git a/test/test_caop.py b/test/test_caop.py index 64e5f7a..a2f16b1 100644 --- a/test/test_caop.py +++ b/test/test_caop.py @@ -326,11 +326,13 @@ def test_settings_upstream_gets_wrong_are_refused(backend, fields, match): # Upstream reads the Hamiltonian on one rank and sends the terms on, so a file it cannot -# read raises there and leaves the other ranks waiting. The bindings read it first. +# read raises there and leaves the other ranks waiting. The bindings read it first. A +# mode beyond system_size is not checked upstream at all. BAD_HAMFILES = { "complex-on-real": ("-1\n(0,1) cdag 1 c 0\n(0,-1) cdag 0 c 1\n", "complex=True"), "bad-coefficient": ("-1\nabc cdag 1 c 0\n", "cannot read the coefficient 'abc'"), "no-statistics": ("# only a comment\n", "no statistics line"), + "mode-beyond-sites": ("1\n1.0 bdag 8 b 0\n1.0 bdag 0 b 8\n", "mode 8, beyond sites=8"), } From b31312a1adef82ea93937dee25da9b5dffb4ca0c Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 17:46:50 -0500 Subject: [PATCH 08/16] CAOP spectrum: refuse on every rank what upstream checks on some ranks only; label reevaluated rows - On the Thrust backend with the on-the-fly Hamiltonian, t_comm_size above the b ranks is refused: the mult returns early on an empty slide and skips the Allreduce the other ranks wait in. - The Hamiltonian's and each operator file's modes are checked against sites, and each operator's statistics against H's, on every rank; upstream checks them on some ranks only, so a bad file hung the others. - omp_set_dynamic(0), as upstream's app: the stored Hamiltonian's per-thread slots. - caop_spectrum_from_coefficients returns orbitals and operator_files; the driver's --read-coefficients CSV labels its rows as the generating run did. - eta defaults to 0.1 and --points to 201, as upstream's caop-spectrum. Co-Authored-By: Claude Opus 5.5 --- examples/caop/run_caop_spectrum.py | 9 ++++--- python/bindings.cpp | 38 +++++++++++++++++++++++++++--- python/spectrum.py | 5 ++-- test/test_spectrum.py | 22 ++++++++++++++++- 4 files changed, 65 insertions(+), 9 deletions(-) diff --git a/examples/caop/run_caop_spectrum.py b/examples/caop/run_caop_spectrum.py index 408c592..ba0c5db 100644 --- a/examples/caop/run_caop_spectrum.py +++ b/examples/caop/run_caop_spectrum.py @@ -85,7 +85,7 @@ def parse_args(): parser.add_argument('--output', default='', help='CSV of G and A on the grid') parser.add_argument('--omega-min', type=float, default=-10.0) parser.add_argument('--omega-max', type=float, default=10.0) - parser.add_argument('--points', type=int, default=401) + parser.add_argument('--points', type=int, default=201) parser.add_argument('--eta', type=float, default=0.1) args = parser.parse_args() if args.read_coefficients: @@ -141,8 +141,11 @@ def main(): args.eta) n = result["green_function"].shape[1] if rank == 0: - _write_csv(args.output, omegas, result["green_function"], list(range(n)), - general=False) + # Label the rows as the generating run did: its orbitals, or for general + # operators the operator indices 0..n-1. + general = bool(result["operator_files"]) + labels = list(range(n)) if general else result["orbitals"] + _write_csv(args.output, omegas, result["green_function"], labels, general) print(f"Reevaluated {args.read_coefficients}: {n} x {n}, " f"{len(omegas)} points, eta {args.eta} -> {args.output}") return 0 diff --git a/python/bindings.cpp b/python/bindings.cpp index 3e08fad..cea45e2 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -1774,10 +1774,36 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { } // Upstream reads the hamfile on h-rank 0 only, as for the solve; see // caop_check_hamfile. The operator files are read the same way. - caop_check_hamfile(hamfile, complex, comm, "caop_spectrum"); + const auto ham_facts = caop_check_hamfile(hamfile, complex, comm, "caop_spectrum"); + caop_check_modes(hamfile, ham_facts, sites, "caop_spectrum"); + // Upstream compares each operator's statistics with H on the h=t=0 ranks + // only, so a mismatch there would hang the rest: refuse it here, everywhere. for (const auto& file : operator_files) { - caop_check_hamfile(file, complex, comm, "caop_spectrum"); + const auto facts = caop_check_hamfile(file, complex, comm, "caop_spectrum"); + caop_check_modes(file, facts, sites, "caop_spectrum"); + if (facts.fermion != ham_facts.fermion) { + throw std::invalid_argument( + "caop_spectrum: operator file " + file + + " has different statistics from the Hamiltonian"); + } } +#ifdef SBD_THRUST + // make_slide deals the b ring stations over the t task ranks + // (caop/basic/helper.h:41-55); with t > b some get none, and upstream's + // Thrust mult returns early on an empty slide (mult_thrust.h:583), skipping + // the Allreduce the other ranks wait in. The stored mode runs the CPU mult. + if (!store_hamiltonian && t_comm_size > world / (h_comm_size * t_comm_size)) { + throw std::invalid_argument( + "caop_spectrum on the gpu (Thrust) backend with the on-the-fly " + "Hamiltonian needs t_comm_size <= the b ranks (ranks / (h_comm_size * " + "t_comm_size)): a task rank with no basis-ring station hangs " + "upstream's GPU kernel. Lower t_comm_size, store the Hamiltonian, or " + "use device='cpu'"); + } +#endif + // As upstream's app: the stored Hamiltonian's per-thread slots need the + // same OpenMP team in every parallel region (caop/basic/mkham.h). + omp_set_dynamic(0); #ifdef SBD_THRUST { int numDevices = 1; @@ -1900,7 +1926,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { py::arg("comm"), py::arg("loadname"), py::arg("shards"), py::arg("hamfile"), py::arg("orbitals"), py::arg("addition"), py::arg("sites"), py::arg("bit_length"), py::arg("reference_energy"), py::arg("extra_basis") = std::vector{}, - py::arg("omegas") = std::vector{}, py::arg("eta") = 0.05, + py::arg("omegas") = std::vector{}, py::arg("eta") = 0.1, py::arg("max_steps") = 100, py::arg("h_comm_size") = 1, py::arg("t_comm_size") = 1, py::arg("store_hamiltonian") = true, py::arg("complex") = false, py::arg("wavefunction_complex") = false, py::arg("save_coefficients") = "", @@ -1917,10 +1943,14 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { size_t n = 0, basis_size = 0, steps = 0; bool addition = true; std::string stop_reason; + std::vector labels; + std::vector operator_files; auto evaluate = [&](auto tag) { using ElemT = decltype(tag); auto result = cs::load_coefficients(path); n = result.orbitals.size(); + labels = result.orbitals; + operator_files = result.operator_files; basis_size = result.basis_size; steps = result.coefficients.A.size(); stop_reason = result.coefficients.stop_reason; @@ -1943,6 +1973,8 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { results["steps"] = steps; results["stop_reason"] = stop_reason; results["addition"] = addition; + results["orbitals"] = labels; + results["operator_files"] = operator_files; return results; }, "Green's function from coefficients saved by caop_spectrum, on a new grid, as " diff --git a/python/spectrum.py b/python/spectrum.py index bb9c0f5..b5fbcf3 100644 --- a/python/spectrum.py +++ b/python/spectrum.py @@ -51,7 +51,7 @@ def green_function(loadname, hamfile, orbitals, channel, *, sites, bit_length, - reference_energy, omegas, eta=0.05, shards=1, extra_basis=(), + reference_energy, omegas, eta=0.1, shards=1, extra_basis=(), max_steps=100, h_comm_size=1, t_comm_size=1, store_hamiltonian=True, complex=False, wavefunction_complex=False, save_coefficients="", operator_files=(), rank_atol=1e-13, rank_rtol=1e-12, device=None): @@ -127,7 +127,8 @@ def green_function_from_coefficients(path, omegas, eta, *, device=None): Returns: dict with ``green_function`` (``(len(omegas), n, n)``), ``basis_size``, - ``steps``, ``stop_reason`` and ``addition`` (the channel it was saved for). + ``steps``, ``stop_reason``, ``addition`` (the channel it was saved for), and + the row labels: ``orbitals``, and ``operator_files`` for general operators. """ backend = sbd._caop_backend(device) return backend.caop_spectrum_from_coefficients( diff --git a/test/test_spectrum.py b/test/test_spectrum.py index 91f86e6..8f8f8b2 100644 --- a/test/test_spectrum.py +++ b/test/test_spectrum.py @@ -286,6 +286,25 @@ def test_orbitals_and_operator_files_together_are_refused(backend): reference_energy=0.0, omegas=[0.0], rank_atol=-1.0) +def test_operator_files_upstream_checks_on_some_ranks_only_are_refused(backend, tmp_path): + """A mode beyond sites, or an operator with other statistics than H: upstream + raises on the h=t=0 ranks or on one h shard only, which hangs the other ranks.""" + if not getattr(backend, "__sbd_has_caop__", False): + pytest.skip("this backend has no CAOP solver") + from sbd.spectrum import green_function + + wide = tmp_path / "wide.txt" + wide.write_text("1\n1.0 sz 0 sz 2\n") + with pytest.raises(ValueError, match="mode 2, beyond sites=2"): + green_function("x-", wide, [0], "addition", sites=2, bit_length=20, + reference_energy=0.0, omegas=[0.0]) + fermion = tmp_path / "fermion.txt" + fermion.write_text("-1\n1.0 cdag 0 c 1\n") + with pytest.raises(ValueError, match="different statistics"): + green_function("x-", DIMER / "ham.txt", [], "addition", sites=2, bit_length=20, + reference_energy=0.0, omegas=[0.0], operator_files=[fermion]) + + # --- the driver: the README's flow --------------------------------------------------- def test_the_drivers_flow_ground_state_then_green_function(caop_device, tmp_path): @@ -355,4 +374,5 @@ def seed(mode): expected = np.array([(projected.conj() / (w + e[0] - ea + 1j * eta)) @ projected.T for w in omegas]) np.testing.assert_allclose(got, expected, atol=1e-8, err_msg=path) - + # The reevaluated CSV keeps the generating run's orbital labels. + assert sorted(set(rows[:, 1].astype(int))) == [0, 2], path From b2568aa726f472503a9fde705b13530ea25ecdc5 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 17:47:10 -0500 Subject: [PATCH 09/16] caop README: the Green's function block writes its own input files Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/examples/caop/README.md b/examples/caop/README.md index 7aeff7d..4c5ab99 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -78,6 +78,10 @@ in Python, with its option names and CSV output; see for the options, conventions and limits. For the 2-site Hubbard file above: ```bash +# 0. the 2-site Hubbard file above, and its half-filling basis +printf '%s\n' -1 '-1.0 cdag 0 c 1' '-1.0 cdag 1 c 0' '-1.0 cdag 2 c 3' '-1.0 cdag 3 c 2' \ + '4.0 cdag 0 cdag 2 c 2 c 0' '4.0 cdag 1 cdag 3 c 3 c 1' > hub2.txt +printf '%s\n' 0101 0110 1001 1010 > hub2_basis.txt # 1. ground state, saved; note the "Ground state energy:" line (10 decimals), not # upstream's rounded "Energy =" log line python run_caop_diag.py --hamfile hub2.txt --basisfiles hub2_basis.txt --system_size 4 \ From 5f51025b903178cfcd0d7b6b5553a979a77fab4c Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Tue, 3 Nov 2026 17:51:24 -0500 Subject: [PATCH 10/16] README: GDB on gpu-omp is refused, not run on the host Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index fbd9d28..bbde4dd 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ constraints. Reach for GDB when the subspace is not a product. `'gpu-omp'` (OpenMP target offload, **NVIDIA or AMD**). All the backends your toolchain supports can be built into one install; each is imported only when first used. **TPB runs on all three; GDB has Thrust kernels only**, so GDB on a - GPU is NVIDIA-only and under `'gpu-omp'` it falls back to the host — see + GPU is NVIDIA-only and `'gpu-omp'` refuses it — see [`examples/gdb/README.md`](examples/gdb/README.md) - MPI parallelization - Integration with [qiskit-addon-sqd](https://github.com/Qiskit/qiskit-addon-sqd) for SQD workflows From 937c9f4c200e6bb8976224caf965b7329c3c76b2 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Thu, 5 Nov 2026 15:00:43 -0500 Subject: [PATCH 11/16] test: the Green's function of a state saved over every rank, under MPI The ground state is solved and saved with b_comm_size = ranks, and the Green's function is computed from those shards on the same ranks; each rank must give the exact Lehmann sum. The other spectrum tests run in one process. Co-Authored-By: Claude Opus 5.5 --- test/test_spectrum.py | 53 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/test/test_spectrum.py b/test/test_spectrum.py index 8f8f8b2..ce6d499 100644 --- a/test/test_spectrum.py +++ b/test/test_spectrum.py @@ -189,6 +189,59 @@ def test_complex_amplitudes_give_the_same_green_function(caop_device, tmp_path): np.testing.assert_allclose(got, exact[channel][0], atol=1e-7, err_msg=channel) +@pytest.mark.mpi +def test_green_function_with_the_state_over_every_rank(caop_device): + """The ground state solved and saved over every rank (b_comm_size = ranks), and its + Green's function from those shards on the same ranks: the exact Lehmann sum on each. + + In the job's own ranks rather than a subprocess: one word per determinant, as the + CAOP solves under MPI in test_caop.py. + """ + import shutil + import tempfile + + from mpi4py import MPI + + from sbd.spectrum import green_function + + comm = MPI.COMM_WORLD + size = comm.Get_size() + path = pathlib.Path(comm.bcast(tempfile.mkdtemp(prefix="sbd-spectrum-") + if comm.Get_rank() == 0 else None, root=0)) + try: + if comm.Get_rank() == 0: + _hubbard(path) + comm.Barrier() + ham, basis = path / "hubbard.txt", path / "hubbard_basis.txt" + e0, exact = _exact_green(ham) + if comm.Get_rank() == 0: + for channel, (_, states) in exact.items(): + (path / f"{channel}_sector.txt").write_text( + "\n".join(format(s, f"0{MODES}b") for s in states) + "\n") + comm.Barrier() + config = sbd.CAOP_SBD() + config.system_size, config.bit_length = MODES, BITS + config.method, config.max_it, config.max_nb, config.eps = 1, 200, 20, 1e-13 + config.b_comm_size, config.do_redist_basis = size, True + # A fixed start, since the default start is random and every rank must agree. + config.initial_determinant_bitstring = basis.read_text().split()[0] + ground = sbd.caop_diag_from_files(ham, [basis], config, savename=str(path / "gs-"), + device=caop_device) + assert ground["energy"] == pytest.approx(e0, abs=1e-9) + for channel in ("addition", "removal"): + r = green_function(path / "gs-", ham, ORBITALS, channel, sites=MODES, + bit_length=BITS, reference_energy=ground["energy"], + omegas=OMEGAS, eta=ETA, shards=size, + extra_basis=[path / f"{channel}_sector.txt"], max_steps=200, + device=caop_device) + np.testing.assert_allclose(r["green_function"], exact[channel][0], atol=1e-7, + err_msg=f"{channel} on rank {comm.Get_rank()}") + finally: + comm.Barrier() + if comm.Get_rank() == 0: + shutil.rmtree(path, ignore_errors=True) + + def test_spectral_function_is_minus_im_g_over_pi_on_the_diagonal(): from sbd.spectrum import spectral_function From 7faa78e101fa1cc0af9b5673cf684578fb74f93c Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 6 Nov 2026 09:25:05 -0500 Subject: [PATCH 12/16] caop README: differences from upstream first, then the command and its defaults Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/examples/caop/README.md b/examples/caop/README.md index 03ed117..f68877c 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -11,10 +11,6 @@ hard-core bosons, spins). For molecules, use TPB or GDB. for each one and for its `gen_ham.py` and `gen_bits.py`. This page covers what is specific to the Python driver and what we found running it. -```bash -mpirun -np 4 python run_caop_diag.py --b_comm_size 4 # upstream's 8-site boson chain -``` - ## Differences from upstream's app - `--device cpu` or `gpu` (Thrust, NVIDIA only); see *GPU*. @@ -26,6 +22,27 @@ mpirun -np 4 python run_caop_diag.py --b_comm_size 4 # upstream's 8-site boso - Not exposed: `--init` (only `0`, random, exists) and `--h_comm_size` (listed upstream but never read; see *MPI*). +## Running it + +```bash +mpirun -np 4 python run_caop_diag.py --b_comm_size 4 +``` + +With only `--b_comm_size` given, this runs with the defaults: + +- **Input:** upstream's example, an 8-site hard-core boson chain: `hamiltonian.txt` and + `basis.txt` from `vendor/sbd-upstream/apps/caop_selected_basis_diagonalization/`, + with `--system_size 8`. For your own model give `--hamfile`, `--basisfiles` and + `--system_size` together. +- **Device:** `--device cpu`. `--device gpu` needs the Thrust build and two or more + OpenMP threads per rank; upstream's `run.sh` uses one, which the Thrust build refuses. +- **Amplitudes:** real. Add `--complex` for complex coefficients or `sy`. +- **Solver:** `--method 1` (Davidson, stored Hamiltonian) to `--tolerance 1e-4`, from a + random start (`--seed 0`). Converge further and fix the seed before comparing + energies; see *Tips*. +- **MPI:** the basis split over all 4 ranks (`--b_comm_size` equal to the rank count), + the simplest layout; see *MPI*. + ## Preparing the input files **A fermion Hamiltonian by hand.** `gen_ham.py` writes hard-core boson chains only. For From 73ce93c5c721ebf6b0d69865eee39ca9225f1302 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 6 Nov 2026 09:53:13 -0500 Subject: [PATCH 13/16] spectrum: say what runs where -- MPI + OpenMP on the CPU, only the on-the-fly H on the GPU The README, the release note and green_function's docstring said the spectrum runs on the cpu and gpu backends "like CAOP", which overstated the GPU part. As upstream's, only applying the Hamiltonian on the fly (--hamiltonian-mode on-the-fly, store_hamiltonian=False) runs on the GPU; the stored Hamiltonian, the default, and the rest of the work run on the CPU. The README now says so up front, with both the driver option and the Python argument. Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 20 ++++++++++++++++--- python/spectrum.py | 7 +++++-- .../notes/caop-spectrum-8c3e5a1f7b2d4096.yaml | 6 ++++-- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/examples/caop/README.md b/examples/caop/README.md index ba6723b..3efbce9 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -92,7 +92,23 @@ basis yourself from every alpha and beta occupation: `gen_bits.py` fixes only th [`caop-spectrum`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop-spectrum) app in Python, with its option names and CSV output; see [its README](https://github.com/r-ccs-cms/sbd/blob/main/apps/caop-spectrum/README.md) -for the options, conventions and limits. For the 2-site Hubbard file above: +for the options, conventions and limits. + +**It runs with MPI + OpenMP on the CPU.** The one exception, as upstream's, is applying +the Hamiltonian on the fly, which runs on the GPU on `--device gpu` (Thrust, NVIDIA +only): + +| | Driver option | Python (`sbd.spectrum.green_function`) | Runs on | +|---|---|---|---| +| Stored Hamiltonian (default) | `--hamiltonian-mode stored` | `store_hamiltonian=True` | CPU, also on `--device gpu` | +| Hamiltonian on the fly | `--hamiltonian-mode on-the-fly` | `store_hamiltonian=False` | GPU on `--device gpu`, else CPU | + +Everything else stays on the CPU whichever you choose: building the excitation basis, +applying the operators to the state, the block-Lanczos vector work, and evaluating the +Green's function from the coefficients. On `--device gpu` the driver prints which of +the two it ran. There is no `gpu-omp` path. + +For the 2-site Hubbard file above: ```bash # 0. the 2-site Hubbard file above, and its half-filling basis @@ -124,8 +140,6 @@ python run_caop_spectrum.py --read-coefficients add.coeff --output narrow.csv \ `--scalar-type complex` if it used `--complex`. The files record none of these. - **Add both channels** for the full retarded Green's function: run `addition` and `removal` on the same orbitals and grid, and sum. -- **Only `--hamiltonian-mode on-the-fly` uses the GPU**, as for the ground state; the - default `stored` runs on the CPU, and the driver says so. ## MPI diff --git a/python/spectrum.py b/python/spectrum.py index b5fbcf3..8b14869 100644 --- a/python/spectrum.py +++ b/python/spectrum.py @@ -79,7 +79,8 @@ def green_function(loadname, hamfile, orbitals, channel, *, sites, bit_length, h_comm_size, t_comm_size: Helper and task communicator sizes; their product must divide the rank count, and the rest of the ranks hold the basis. store_hamiltonian: Store H for faster application (upstream's ``stored`` - mode), rather than applying it on the fly. + mode), rather than applying it on the fly. Only ``False`` uses the GPU + on ``device="gpu"``; the stored Hamiltonian is built and applied on the CPU. complex: Run with complex amplitudes, for a complex Hamiltonian. wavefunction_complex: The saved state has complex amplitudes, as one solved with ``caop_diag_from_files(..., complex=True)`` does; needs ``complex``. @@ -94,7 +95,9 @@ def green_function(loadname, hamfile, orbitals, channel, *, sites, bit_length, or change the particle number. rank_atol, rank_rtol: Block-Lanczos deflation tolerances: a column whose norm falls below them is dropped. Upstream's defaults. - device: ``"cpu"`` or ``"gpu"`` (Thrust); there is no OpenMP-offload path. + device: ``"cpu"`` or ``"gpu"`` (Thrust); there is no OpenMP-offload path. On + ``"gpu"`` only the on-the-fly Hamiltonian application runs on the GPU; the + rest, as everything on ``"cpu"``, runs on the CPU with MPI and OpenMP. Returns: dict with ``green_function``, a complex array of shape diff --git a/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml b/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml index 70990d7..730d1fe 100644 --- a/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml +++ b/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml @@ -6,8 +6,10 @@ features: takes a state saved by :func:`sbd.caop_diag_from_files`, the modes to add an electron to or remove one from, and a frequency grid, and returns the Green's function on it; :func:`sbd.spectrum.spectral_function` gives the spectral function. The conventions - are upstream's. Real and complex amplitudes are both available. Like CAOP, it runs on - the ``cpu`` and ``gpu`` (Thrust) backends. + are upstream's. Real and complex amplitudes are both available. It runs with MPI and + OpenMP on the CPU; as upstream's, only applying the Hamiltonian on the fly + (``store_hamiltonian=False``) runs on the GPU, on the ``gpu`` (Thrust) backend. The + stored Hamiltonian, the default, stays on the CPU there too. As upstream's app, it also takes general operators instead of ``c``/``cdag`` (``operator_files``) and the block-Lanczos deflation tolerances, and From 0a2d83261ea629f522231d42bf022a202c868c42 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 6 Nov 2026 09:58:36 -0500 Subject: [PATCH 14/16] caop README: the spectrum's backend table names the driver option only Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/caop/README.md b/examples/caop/README.md index 3efbce9..493873f 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -98,10 +98,10 @@ for the options, conventions and limits. the Hamiltonian on the fly, which runs on the GPU on `--device gpu` (Thrust, NVIDIA only): -| | Driver option | Python (`sbd.spectrum.green_function`) | Runs on | -|---|---|---|---| -| Stored Hamiltonian (default) | `--hamiltonian-mode stored` | `store_hamiltonian=True` | CPU, also on `--device gpu` | -| Hamiltonian on the fly | `--hamiltonian-mode on-the-fly` | `store_hamiltonian=False` | GPU on `--device gpu`, else CPU | +| `--hamiltonian-mode` | Runs on | +|---|---| +| `stored` (default) | CPU, also on `--device gpu` | +| `on-the-fly` | GPU on `--device gpu`, else CPU | Everything else stays on the CPU whichever you choose: building the excitation basis, applying the operators to the state, the block-Lanczos vector work, and evaluating the From 3b5574d4a084c4d4834db3335252ded9386dac02 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 6 Nov 2026 10:14:45 -0500 Subject: [PATCH 15/16] caop README: say up front how the solve and the spectrum drivers chain Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/caop/README.md b/examples/caop/README.md index 493873f..c909a53 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -11,6 +11,13 @@ hard-core bosons, spins). For molecules, use TPB or GDB. for each one and for its `gen_ham.py` and `gen_bits.py`. This page covers what is specific to the Python driver and what we found running it. +**The two drivers chain:** `run_caop_diag.py --savename gs-` solves for the ground state +and saves it; `run_caop_spectrum.py --loadname gs-` then reads it, with the same +Hamiltonian file and the ground-state energy, and computes its Green's function. The +second run must match the first's `--system_size`, `--bit_length`, `--b_comm_size` and +`--complex`, which the saved files do not record. See +[Green's function](#greens-function) for a worked example. + ## Differences from upstream's app - `--device cpu` or `gpu` (Thrust, NVIDIA only); see *GPU*. From 9501fa4e35a0358edf0feb7b61cba3907afe39c7 Mon Sep 17 00:00:00 2001 From: Sophia Wen Date: Fri, 6 Nov 2026 10:18:27 -0500 Subject: [PATCH 16/16] caop README: introduce both drivers at the top, and how they chain Co-Authored-By: Claude Opus 5.5 --- examples/caop/README.md | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/examples/caop/README.md b/examples/caop/README.md index c909a53..fe4644b 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -1,23 +1,24 @@ # CAOP — models written with creation and annihilation operators -`run_caop_diag.py` runs upstream SBD's -[`caop_selected_basis_diagonalization`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop_selected_basis_diagonalization) -app from Python: the ground state of a Hamiltonian written as operator terms, over a -selected basis of occupation bitstrings, for lattice and spin models (Hubbard, -hard-core bosons, spins). For molecules, use TPB or GDB. - -**The options are upstream's, with the same names and meanings;** see -[upstream's README](https://github.com/r-ccs-cms/sbd/blob/main/apps/caop_selected_basis_diagonalization/README.md) -for each one and for its `gen_ham.py` and `gen_bits.py`. This page covers what is -specific to the Python driver and what we found running it. - -**The two drivers chain:** `run_caop_diag.py --savename gs-` solves for the ground state -and saves it; `run_caop_spectrum.py --loadname gs-` then reads it, with the same -Hamiltonian file and the ground-state energy, and computes its Green's function. The -second run must match the first's `--system_size`, `--bit_length`, `--b_comm_size` and -`--complex`, which the saved files do not record. See +Two drivers run upstream SBD's CAOP apps from Python, for lattice and spin models +(Hubbard, hard-core bosons, spins) written as operator terms over a selected basis of +occupation bitstrings. For molecules, use TPB or GDB. + +- **`run_caop_diag.py`**: the ground state, as upstream's + [`caop_selected_basis_diagonalization`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop_selected_basis_diagonalization). +- **`run_caop_spectrum.py`**: the Green's function of a saved ground state, as upstream's + [`caop-spectrum`](https://github.com/r-ccs-cms/sbd/tree/main/apps/caop-spectrum). + +They chain: `run_caop_diag.py --savename gs-` saves the ground state, and +`run_caop_spectrum.py --loadname gs-` reads it, with the same Hamiltonian file and the +ground-state energy. The second run must match the first's `--system_size`, +`--bit_length`, `--b_comm_size` and `--complex`, which the saved files do not record. See [Green's function](#greens-function) for a worked example. +**The options are upstream's, with the same names and meanings:** see each app's README, +linked above, and upstream's `gen_ham.py` and `gen_bits.py` for writing input files. This +page covers what differs in the Python drivers, and the checks they add. + ## Differences from upstream's app - `--device cpu` or `gpu` (Thrust, NVIDIA only); see *GPU*.