diff --git a/docs/apidocs/index.rst b/docs/apidocs/index.rst index 63080ca..96b8eb9 100644 --- a/docs/apidocs/index.rst +++ b/docs/apidocs/index.rst @@ -9,4 +9,5 @@ sbd.sbd_solver sbd.analysis sbd.device_config + sbd.spectrum sbd.mpi_util 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 f68877c..fe4644b 100644 --- a/examples/caop/README.md +++ b/examples/caop/README.md @@ -1,15 +1,23 @@ # 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. +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. -**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. +- **`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 @@ -86,6 +94,61 @@ 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. + +**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): + +| `--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 +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 +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 \ + --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. + ## 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..ba0c5db --- /dev/null +++ b/examples/caop/run_caop_spectrum.py @@ -0,0 +1,189 @@ +#!/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=201) + 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: + # 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 + + 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 0b0bf21..9e65627 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -38,8 +38,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -1881,6 +1883,273 @@ 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. + 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) { + 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; +#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.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") = "", + 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; + 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; + 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; + results["orbitals"] = labels; + results["operator_files"] = operator_files; + 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..8b14869 --- /dev/null +++ b/python/spectrum.py @@ -0,0 +1,148 @@ +# 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.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): + """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. 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``. + 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. 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 + ``(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``, ``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( + 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..730d1fe --- /dev/null +++ b/releasenotes/notes/caop-spectrum-8c3e5a1f7b2d4096.yaml @@ -0,0 +1,21 @@ +--- +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. 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 + :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..ce6d499 --- /dev/null +++ b/test/test_spectrum.py @@ -0,0 +1,431 @@ +# 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) + + +@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 + + 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) + + +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): + """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) + # The reevaluated CSV keeps the generating run's orbital labels. + assert sorted(set(rows[:, 1].astype(int))) == [0, 2], path