Skip to content

Implementation guide

Jump to bottom
lajerm edited this page Aug 11, 2026 · 3 revisions

PETSc backend implementation guide

This page describes how the backend-neutral EDRIXS interface fits together and what needs to be implemented in the PETSc backend. It is intended as a practical companion to the restructuring plan.

The main principle is that physical problem definition is backend independent, while construction and numerical solution of many-body operators are delegated to a backend.

For a detailed description of the reference SciPy implementation, see the SciPy backend implementation note.

Repository structure

The relevant parts of the repository are

edrixs/
    models.py
    solvers.py
    _solvers_helpers.py
    fock_basis.py
    krylov.py

    scipy_backend/
        __init__.py
        scipy_backend.py

    petsc_backend/
        __init__.py
        petsc_backend.py

tests/
    test_solvers_dispatch.py
    test_scipy_operator_construction.py
    solver_consistency_tests/

benchmark_checks/
    example_reference_data.zip
    run_and_compare_scipy_examples.py
    run_lanio3_2v1c_scipy.py
    run_lanio3_thin_scipy.py
    run_nio_aim_xas_scipy.py
    run_pu_o45_scipy.py
    run_u_l3_scipy.py
    run_uru2si2_scipy.py

The PETSc work should primarily be confined to edrixs/petsc_backend/petsc_backend.py, together with PETSc-specific tests and benchmark scripts as they are added.

The backend-independent API should not need to know the internal PETSc/SLEPc algorithm being used.

How a calculation reaches a backend

A typical calculation proceeds as follows.

1. Define the physical model

For example,

from edrixs.models import model_1v1c

problem = model_1v1c(...)

The model functions return backend-independent orbital-space data:

emat_i, umat_i, basis_i,
emat_n, umat_n, basis_n,
trans_mat

emat_i and emat_n contain two-fermion terms, umat_i and umat_n contain four-fermion terms, and the FockBasis objects describe the many-body basis.

trans_mat contains the orbital-space components of the photon transition operator. There are normally three components for a dipole transition or five for a quadrupole transition.

The model layer should remain independent of PETSc.

2. Build backend-owned many-body operators

The public call is

from edrixs.solvers import get_ops

hmat_i, hmat_n, trans_ops = get_ops(
    *problem,
    backend="petsc",
    backend_kws={...},
)

get_ops calls the public build_op function for the initial Hamiltonian, intermediate Hamiltonian, and each component of the transition operator.

The backend entry point is therefore

build_op_petsc(emat, umat, lb, rb=None, *, backend_kws=None)

This is a general Fock-space operator constructor, not specifically a Hamiltonian constructor.

It must construct an operator of the form

two-fermion contribution from emat
+
four-fermion contribution from umat

where either contribution may be absent (None).

lb is the output/left Fock basis and rb is the input/right Fock basis. If rb is None, it defaults to lb.

For Hamiltonians, normally lb == rb.

For photon transition operators,

lb = basis_n
rb = basis_i
umat = None

so the resulting operator can be rectangular and maps the initial-state Fock space into the intermediate-state Fock space.

The PETSc implementation must therefore not assume that every operator is square.

3. Ground-state diagonalization

The public call is

eval_i, evec_i = ed(
    hmat_i,
    num_evals=num_evals,
    backend="petsc",
    backend_kws={...},
)

The PETSc entry point is

ed_petsc(hmat_i, num_evals=1, *, backend_kws=None)

It should return

eval_i : shape (num_evals,)
evec_i : shape (dimension, num_evals)

with eigenvectors stored by column.

The public API deliberately passes backend-specific solver controls in the dictionary backend_kws. PETSc/SLEPc-specific options therefore belong in this dictionary rather than in the public ed() signature.

For example, PETSc-specific choices such as eigensolver type, tolerances, iteration limits, subspace dimensions, or GPU-related options can be interpreted inside ed_petsc.

The SciPy implementation is useful as a behavioral reference, but the PETSc implementation does not need to copy the SciPy numerical algorithm.

4. XAS

The public call is

xas(
    eval_i,
    evec_i,
    hmat_n,
    trans_ops,
    ominc,
    gamma_c=...,
    pol_type=...,
    backend="petsc",
    backend_kws={...},
)

The PETSc entry point is

xas_petsc(
    eval_i,
    evec_i,
    hmat_n,
    trans_op,
    ominc,
    *,
    gamma_c=0.1,
    thin=1.0,
    phi=0.0,
    pol_type=None,
    temperature=1.0,
    scatter_axis=None,
    backend_kws=None,
)

The returned spectrum should have shape

(len(ominc), len(pol_type))

after default polarization handling.

The public physical arguments should retain the same meaning for every backend. PETSc-specific numerical choices belong in backend_kws.

The SciPy implementation currently uses Lanczos continued fractions. PETSc may use an equivalent PETSc/SLEPc implementation as appropriate; numerical method details are backend-specific.

5. RIXS

The public call is

rixs(
    eval_i,
    evec_i,
    hmat_i,
    hmat_n,
    trans_ops,
    ominc,
    eloss,
    gamma_c=...,
    gamma_f=...,
    pol_type=...,
    skip_gs=False,
    backend="petsc",
    backend_kws={...},
)

The PETSc entry point is

rixs_petsc(
    eval_i,
    evec_i,
    hmat_i,
    hmat_n,
    trans_op,
    ominc,
    eloss,
    *,
    gamma_c=0.1,
    gamma_f=0.01,
    thin=1.0,
    thout=1.0,
    phi=0.0,
    pol_type=None,
    temperature=1.0,
    scatter_axis=None,
    skip_gs=False,
    return_poles=False,
    backend_kws=None,
)

The spectrum should have shape

(len(ominc), len(eloss), len(pol_type))

after default polarization handling.

skip_gs is part of the released public behavior and must be retained.

When skip_gs=True, transitions back into the retained initial-state subspace should be removed from the final-state spectrum. In the present SciPy implementation this is done by projecting the final RIXS vector out of the subspace spanned by evec_i before the final-state Lanczos calculation.

The PETSc implementation should reproduce this behavior even if the underlying numerical implementation is different.

return_poles is also part of the current public/backend contract and should be implemented consistently with the public behavior of the SciPy backend.

Functions to implement in petsc_backend.py

The current PETSc file is deliberately a stub. The functions that need real implementations are

owns_operator_petsc(operator)

build_op_petsc(
    emat, umat, lb, rb=None, *, backend_kws=None
)

ed_petsc(
    hmat_i, num_evals=1, *, backend_kws=None
)

xas_petsc(
    eval_i, evec_i, hmat_n, trans_op, ominc, *,
    gamma_c=0.1, thin=1.0, phi=0.0, pol_type=None,
    temperature=1.0, scatter_axis=None, backend_kws=None
)

rixs_petsc(
    eval_i, evec_i, hmat_i, hmat_n, trans_op, ominc, eloss, *,
    gamma_c=0.1, gamma_f=0.01, thin=1.0, thout=1.0, phi=0.0,
    pol_type=None, temperature=1.0, scatter_axis=None,
    skip_gs=False, return_poles=False, backend_kws=None
)

These signatures should be kept synchronized with the backend-neutral wrappers in solvers.py.

Backend ownership and inference

solvers.ed, solvers.xas, and solvers.rixs allow the backend to be supplied explicitly:

backend="petsc"

but they can also infer the backend from the supplied operators.

This is why owns_operator_petsc() is part of the backend contract.

It should return True for PETSc objects owned by this backend and False otherwise.

Once PETSc operators are correctly recognized, calls such as

ed(hmat_i)

can resolve to ed_petsc without the user explicitly repeating backend="petsc".

The implementation should avoid importing petsc4py when the PETSc backend is not being used. The current stub deliberately imports PETSc lazily.

backend_kws

Backend-specific controls are deliberately passed as one dictionary.

For example,

eval_i, evec_i = ed(
    hmat_i,
    num_evals=10,
    backend="petsc",
    backend_kws={
        # PETSc/SLEPc-specific options
    },
)

The backend implementation is responsible for interpreting and validating the contents of this dictionary.

The public solvers.py interface should not accumulate PETSc-specific keyword arguments.

Likewise, build_op, xas, and rixs each receive their own backend_kws dictionary.

This allows SciPy and PETSc to expose different numerical controls without changing the common physical API.

What should remain backend independent

The following should normally not be reimplemented in the PETSc backend:

  • physical model construction in models.py;
  • FockBasis and Fock-basis generation in fock_basis.py;
  • the public build_op, get_ops, ed, xas, and rixs wrappers in solvers.py;
  • physical meanings of polarization, broadening, temperature, incident-energy and energy-loss arguments;
  • public output shapes;
  • skip_gs behavior.

Backend-independent helper code can be shared where useful. Backend-specific numerical data structures, operator assembly, eigensolvers, Krylov methods, linear solves, and PETSc/SLEPc configuration belong in the PETSc backend.

SciPy as the reference implementation

edrixs/scipy_backend/scipy_backend.py is the existing working implementation of the staged API.

It is useful for understanding

  • operator shapes and left/right basis conventions;
  • the meaning of num_evals;
  • the XAS and RIXS output layouts;
  • dipole and quadrupole transition handling;
  • skip_gs;
  • the expected relationship between public functions and backend entry points.

It should not be treated as a requirement to reproduce SciPy implementation details. PETSc/SLEPc should use the numerical methods and data structures appropriate to that backend.

Tests

The repository-level tests live under tests/.

tests/test_solvers_dispatch.py checks the backend-neutral dispatch layer. Equivalent PETSc coverage should establish that

  • PETSc operators are recognized by backend inference;
  • backend_kws reaches the PETSc implementation;
  • XAS and RIXS return arrays of the expected shape;
  • skip_gs has the expected effect.

The consistency tests live under

tests/solver_consistency_tests/

These are the natural place to compare SciPy and PETSc on the same small physical problems.

The consistency tests should be small enough for CI. Their purpose is correctness across implementations rather than performance benchmarking.

Benchmark checks

The larger legacy-reference checks live in

benchmark_checks/

They serve a different purpose from the unit and consistency tests.

The six intended benchmark calculations are

LaNiO3_2v1c
LaNiO3_thin
NiO_AIM_XAS
Pu_O45
U_L3
URu2Si2

The reference spectra and eigenvalues are stored in

benchmark_checks/example_reference_data.zip

Prepare the reference data

From the repository root:

cd benchmark_checks
unzip example_reference_data.zip

After extraction there should be a directory

benchmark_checks/example_reference_data/

containing one directory for each benchmark case.

Do not modify the reference data when implementing a new backend.

Run the existing SciPy benchmark check

With the EDRIXS environment activated, run

python run_and_compare_scipy_examples.py

The runner

  1. runs each benchmark script;
  2. writes fresh calculation output below example_validation_output/;
  3. compares calculated eigenvalues with the legacy eigvals.dat;
  4. checks XAS/RIXS energy grids;
  5. compares XAS/RIXS intensities using the normalized integrated spectral difference;
  6. prints PASS or FAIL for each comparison.

The default comparison tolerances are

spectral normalized difference: 1e-6
eigenvalue absolute tolerance:  1e-7 eV
energy-grid absolute tolerance:  1e-10

The output should be read case by case. The runner exits with status 0 when all required comparisons pass and status 1 when one or more comparisons fail.

If the calculation outputs have already been generated, they can be checked again without rerunning the calculations:

python run_and_compare_scipy_examples.py --compare-only

PETSc benchmark validation

The current benchmark scripts and runner are named for and explicitly select the SciPy backend. The reference data and comparison logic are nevertheless backend independent.

For PETSc validation, retain exactly the same

  • physical model parameters;
  • incident-energy and energy-loss grids;
  • broadenings;
  • polarization choices;
  • number of retained states;
  • output-file conventions;
  • comparison tolerances.

Then either

  1. add corresponding run_*_petsc.py benchmark scripts, or
  2. later parameterize the benchmark scripts/runner so that the backend can be selected without duplicating the physical problem definition.

The important requirement is that a PETSc benchmark solves the same physical problem and writes the same comparison files. The reference data should remain unchanged.

A PETSc benchmark run should ultimately provide the same type of report:

LaNiO3_2v1c
  PASS/FAIL eigenvalues
  PASS/FAIL xas.dat
  PASS/FAIL rixs_pi.dat

...

Pu_O45
  PASS/FAIL eigenvalues
  PASS/FAIL xas.dat
  PASS/FAIL rixs_pi.dat

and likewise for the other four cases.

Parallelism

The current backend interface does not expose threading or multiprocessing controls.

Definition of a usable PETSc backend

At a minimum, the PETSc backend is integrated when the following command chain works without PETSc-specific changes to the model layer or public solver API:

problem = model_1v1c(...)

hmat_i, hmat_n, trans_ops = get_ops(
    *problem,
    backend="petsc",
    backend_kws={...},
)

eval_i, evec_i = ed(
    hmat_i,
    num_evals=...,
    backend="petsc",
    backend_kws={...},
)

Ixas = xas(
    eval_i,
    evec_i,
    hmat_n,
    trans_ops,
    ominc,
    backend="petsc",
    backend_kws={...},
)

Irixs = rixs(
    eval_i,
    evec_i,
    hmat_i,
    hmat_n,
    trans_ops,
    ominc,
    eloss,
    backend="petsc",
    backend_kws={...},
)

The corresponding small consistency tests should agree with SciPy, and the larger benchmark checks should then be used to identify any remaining numerical differences on realistic examples.

Clone this wiki locally