GradSCF

Differentiable electronic structure in JAX.

GradSCF is a Python framework for molecular and periodic electronic-structure calculations, response theory, and learning within self-consistent calculations. It combines HF/DFT with MP, CI/FCI, CC, EOM-CC, GW, and BSE. External energy functionals and neural Gaussian basis models connect to physical equations through JAX automatic differentiation and shared implicit-response solvers.

Installation · Quick start · Methods · Differentiation · Learning · Examples · Citation

Highlights

• Ground states and correlation. Molecular and periodic HF/DFT, orbital-free DFT, and molecular MP, CI/FCI, and CC/QCI calculations.
• Excitations and spectra. TDHF/TDDFT, EOM excitation/ionization/attachment energies, GW quasiparticles, BSE optical response, and vibrational coupling.
• Memory-aware molecular response. Matrix-free TDA/full TDDFT, streamed AO grids, reusable MO grid fields, and coordinated vector/subspace budgets.
• Differentiable physical solutions. Array kernels expose parameter derivatives; converged electronic states and spectral problems use shared response rules with explicit validity conditions.
• External models in self-consistency. User-defined features and networks supply XC or kinetic energies; GradSCF computes potentials and response actions. Neural basis models optimize Gaussian contractions.
• Shared libcint integral interface. Gaussian integrals, exact packed ERIs, shell-direct J/K and explicit auxiliary RI share one production engine with operator-specific JVP/VJP rules.
• Two Python interfaces. Short calculation objects for interactive use, and functional array APIs for jax.jit, jax.grad, and custom objectives.

%%{init: {"htmlLabels": false, "themeVariables": {"fontFamily": "Arial, sans-serif", "fontSize": "16px"}, "flowchart": {"padding": 24, "nodeSpacing": 36, "rankSpacing": 48}}}%%
flowchart TB
    G["Geometry and basis"] --> I["Integrals and grids"]
    N["NNAO basis model"] -.-> I
    M["External XC / KEDF"] -.->|XC| K["KS-DFT"]
    M -.->|XC / KEDF| O["OFDFT"]
    I --> H["HF"]
    I --> K
    I --> O

    H --> P["MP / CI / FCI"]
    H --> C["CC / QCI"]
    H --> T["TDA / TDHF / TDDFT"]
    K --> T
    H --> W["GW"]
    K --> W
    W --> B["BSE"]
    C -->|CCSD| E["EOM-CCSD"]

    H --> D["DFPT"]
    K --> D
    O --> D
    P -->|Reference data| A["iKS assessment"]

Loading

Solid arrows show reference/preparation paths; dashed arrows show external-model inputs. Gaussian AO integrals use libcint; periodic grids and pseudopotentials retain their FFT/GTH paths. iKS here denotes fixed-reference functional assessment. Each method's guide specifies its supported references, representations and derivative contract.

Installation

Install from source on Linux or macOS. The package declares Python 3.10+; repository tests use Python 3.11+. Native CPU integrals require CMake 3.20+, a C99/C++17 compiler, and a BLAS library (Accelerate on macOS).

git clone --branch release/v1.0.0 https://github.com/STOKES-DOT/GradSCF.git
cd GradSCF
python -m venv .venv
source .venv/bin/activate
python -m pip install -e .
python -m gradscf.integrals._native.build
JAX_PLATFORMS=cpu JAX_ENABLE_X64=1 python -c \
  "from gradscf import gto, scf; mf = scf.RHF(gto.M(atom='H 0 0 0; H 0 0 .74', basis='sto-3g')).run(); assert mf.converged; print(float(mf.e_tot))"

Native sources are bundled; the build downloads no upstream source. Rebuild after changing JAX versions. See the native build guide.

Optional installation	Purpose
python -m pip install -e ".[upstreams]"	jax-xc for supported conventional XC functionals
python -m pip install -e ".[dev,comparison-tests]"	pytest and independent PySCF comparisons
python -m pip install -e ".[nnao]"	Pinned MACE-JAX dependency for neural basis models
python -m pip install -e ".[reproducibility]"	Plotting, HDF5, SciPy, and evaluation tools

PySCF is an optional independent reference for comparisons. Production Gaussian integrals use the bundled libcint engine without a PySCF Python runtime. Integral preparation executes on CPU; supported downstream JAX calculations can use GPU. Device coverage is specified for each calculation path.

Quick start

Run these independent snippets after installation and the native build. Set JAX_PLATFORMS=cpu before launching Python for CPU execution. Enable float64 before constructing arrays. Coordinates below are Angstrom; energies are Hartree.

HF and MP2

import jax
jax.config.update("jax_enable_x64", True)
from gradscf import gto, scf, mp

mol = gto.M(atom="H 0 0 0; H 0 0 .74", basis="sto-3g")
mf = scf.RHF(mol, conv_tol=1e-12).run()
assert mf.converged
pt = mp.MP2(mf).run()
print("HF:", float(mf.e_tot), "MP2:", float(pt.e_tot))

Use scf.UHF or scf.ROHF for their respective open-shell references. MP2/MP3 and generic MPn accept real canonical RHF/UHF. CI and CC have additional ROHF-based paths.

DFT and excitation energies

This snippet needs the [upstreams] extra. tdscf.TDA and tdscf.TDDFT provide excitation energies and transition properties for supported references and XC kernels. An HF reference selects TDHF response.

import jax
jax.config.update("jax_enable_x64", True)
from gradscf import gto, dft, tdscf

mol = gto.M(atom="O 0 0 0; H 0 .75 .58; H 0 -.75 .58", basis="sto-3g")
mf = dft.RKS(mol, xc="pbe", grids_level=1, conv_tol=1e-10,
             max_memory=60000).run()
assert mf.converged
td = tdscf.TDA(mf, nstates=3)
td.kernel()
assert td.converged
print("Excitation energies / Ha:", td.e)

Use tdscf.TDDFT(mf, nstates=3) for full response. The same facades dispatch dft.UKS references to spin-conserving unrestricted TDA/TDDFT. TD response inherits the SCF memory budget; pass max_memory to the TD constructor to override it. The default is 60000 decimal MB, a planning budget rather than a process-memory limit. See molecular response execution.

Differentiate a correlated energy

A dimensionless parameter scales the two-electron MO integrals at a fixed orbital frame. The derivative includes converged CC amplitude response; HF orbital and nuclear-coordinate response are outside this example.

import jax
jax.config.update("jax_enable_x64", True)
from gradscf import gto, scf, cc
from gradscf.scf.reference import reference_from_source

mf = scf.RHF(gto.M(atom="H 0 0 0; H 0 0 .74", basis="sto-3g")).run()
ref = reference_from_source(mf)
config = cc.CCConfig(conv_tol=1e-12, residual_tol=1e-11)

def energy(coupling):
    result = cc.run_cc(ref.h1, coupling * ref.eri, nocc=ref.nocc,
                       nuclear_repulsion=ref.nuclear_repulsion, config=config)
    return result.total_energy

value, derivative = jax.jit(jax.value_and_grad(energy))(1.0)
print("Energy / Ha:", float(value), "dE/dcoupling / Ha:", float(derivative))

For a complete implicit-HF-to-MP parameter derivative, see the MP gradient example.

Train an external energy functional

This pure-JAX network adds a learned correction to Dirac exchange. Features and architecture are defined outside GradSCF; the trainer reconverges the electronic state and differentiates its fixed point.

import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
from gradscf import gto, scf, dft, fci
from gradscf.model.neural_xc import training

def inputs(state):
    return {"rho": state.rho, "weights": state.weights}

def energy(params, x):
    rho = jnp.maximum(x["rho"], 1e-18)
    exchange = -.75 * (3 / jnp.pi)**(1 / 3) * rho**(4 / 3)
    correction = .01 * rho * jnp.tanh(params["w"] * jnp.log1p(rho))
    return jnp.sum(x["weights"] * (exchange + correction))

functional = dft.custom_xc(input_fn=inputs, energy_fn=energy)
mf = scf.RHF(gto.M(atom="H 0 0 0; H 0 0 .74", basis="sto-3g"),
             grids_level=0, conv_tol=1e-12).run()
target = float(fci.FCI(mf, solver="dense").run().e_tot)
trainer = training.Trainer(functional, params={"w": jnp.array(0.)})
trainer.mode = "implicit"
trainer.loss = {"energy": {"mse": 1., "mae": 1.}}
trainer.learning_rate = .001
trainer.run([training.Sample(mf, energy=target)], steps=3)
print(trainer.history["loss"])

This demonstrates the interface, not a transferable XC model. See the external Flax example, four-layer density-matrix MLP, and 100-step H2/FCI comparison.

Supported methods

Forward availability, derivative coverage, and physical accuracy are separate.

Family	Available calculations	Reference / scope and guide
Molecular HF/DFT	RHF/UHF/ROHF/GHF; RKS/UKS/ROKS/GKS; UHF/UKS stability	SCF, DFT; closed-shell HF uses scf.RHF
Periodic HF/DFT	Gamma/k-point calculations, GTH, FFT density fitting, bands	Periodic modules; method-specific periodic contracts
MP	MP2/MP3 and order-driven MPn; frozen orbitals, E2 spin components	MP, Taylor series; real canonical RHF/UHF, capacity-limited high-order engine
CI	Singlet/triplet CIS, singlet CIS(D), spin-conserving UCIS, rank-truncated CI	CI; real RHF/UHF/ROHF determinant spaces
FCI	Complete occupation-string spaces, core/active selection, RDMs, transition RDMs, spin diagnostics	FCI; real common spatial orbitals, integer spin populations
CC/QCI	CCS/CCD/CCSD/CC2/LCCD/LCCSD, QCISD, CCSD(T)/QCISD(T), Lambda, 1/2-RDMs	CC; restricted models and UCCSD/UCCD; model-specific properties
EOM-CCSD	EE singlet, IP/EA doublet energies, left/right states, first-order response	EOM; real restricted CCSD reference
TDHF/TDDFT	TDA/full response, transition properties and spectra; periodic q=0 response	Molecular facade, periodic response
GW	Molecular G0W0-CD, evGW0, evGW, restricted qsGW/matrix Matsubara scGW; Gamma/k-point GW	GW, scGW, periodic GW
BSE	Static TDA/full optics, singlet/triplet and unrestricted spin-conserving sectors	BSE; molecular reference and screening conventions
OFDFT	Molecular/periodic Gaussian amplitudes and FFT grids; TF, vW, TF+vW, WT, WGC99, external KEDFs	OFDFT; unpolarized density, WT/WGC periodic only
DFPT	Static molecular/periodic KS and OFDFT response; Taylor response series and supported finite-q perturbations	DFPT; converged-reference, k-mesh and perturbation contracts
Inverse KS references	RKS/UKS/GKS density targets and local-potential maps; finite-temperature matrix references; fixed-reference functional assessment	iKS; explicit finite-basis reference method, zero-temperature Wu–Yang/coupled inverse-GKS unfinished
Vibrational coupling	Fan/Debye–Waller self-energies, spectra, exciton–vibration optics, finite-space vibronic Hamiltonians	GW EP, BSE EP; fixed bath, BSE projection restricted TDA

Canonical real UHF CCSD(T), UCC Lambda and spin-resolved densities are available. The semicanonical triples option supports its specified noncanonical references through a tensor resolvent. ROHF-based UCCSD uses unrestricted amplitudes on common orbitals, not a separate spin-adapted ROCCSD formulation. See open-shell scope.

FCI active spaces provide a foundation for future multireference drivers; CASSCF orbital optimization is not implemented. Fixed-phonon GW/BSE coupling has no self-consistent phonon feedback; Matsubara scGW has no analytic continuation.

Differentiation

Choose the variable and observable before selecting a derivative path:

Target	Numerical response	Contract
SCF energy/state	Explicit JAX computation or implicit stationarity/fixed-point response	Convergence, occupations, integral/XC derivatives
MP energy/amplitudes	Specialized JAX kernels or residual-driven Taylor lifting	Canonical inputs, physical denominators; upstream HF response
CI/FCI/TDA states	Shared isolated-root response; FCI also exposes complete-subspace response	Coefficient objectives need an eigenvector response mode
CC energy/amplitudes	Shared nonlinear-root response	Converged amplitudes and checked adjoint solves
EOM states and cluster observables	First-order non-Hermitian isolated-root/pair and complete-cluster projector/energy-sum response	Real semisimple clusters with a resolved external boundary; state/Dyson properties require the specified isolated-pair mode
GW roots/self-consistency	Implicit QP roots and opt-in molecular outer response	GW contracts; method/spectral-gap restrictions
OFDFT density/energy	Explicit iterates or constrained implicit stationarity	Particle-number constraint, functional derivatives, resolved solution
Geometry/basis parameters	Integral rules composed with electronic response	Operator, layout, parameter and derivative-order coverage

Internal degeneracy can permit differentiable complete-subspace projectors and energy sums when the selected space is separated from excluded states. A numbered root or arbitrary eigenvector need not have a unique derivative. See spectral response.

Calculation objects validate eager inputs; use array kernels inside jax.jit/jax.grad, with topology, occupations and configuration static. Check convergence, residuals, denominators and validity flags. Higher derivatives are supported on specifically validated paths.

Full nuclear derivatives require electronic response and integral/grid/XC derivatives. Native full integrals have first/second geometry rules for the documented operators; packed production, auxiliary integrals and direct J/K parameter derivatives have different coverage. The generic mf.nuc_grad_method().kernel() remains disabled. See native contracts and the force-supervision example.

Learning with GradSCF

External XC. dft.custom_xc(input_fn=..., energy_fn=..., init_fn=...) connects external formulas and JAX/Flax/Haiku networks through the existing dft.Functional. dft.custom_xc(spec) instead constructs a complete traditional recipe, including exact exchange and additional physical terms. See the XC library for names and method scope. DensityInputs exposes the current density matrix, grid density and supported derivatives, kinetic-energy density, and fixed geometry/AO data. Features are rebuilt from each trial SCF density. The callback returns a real scalar XC energy in Hartree, including its baseline. GradSCF differentiates the whole feature-to-energy graph for the AO potential and Hessian-vector response. See functional contracts.

Sample and Trainer live in gradscf.model.neural_xc.training. They accept both external dft.Functional callbacks and the built-in Neural XC model. Each model owns its training and experiment workflows; there are no top-level gradscf.training or gradscf.workflows packages.

Mode	Forward and derivative
fixed_density	Evaluate the functional on the supplied fixed state
explicit	Run and differentiate the existing JAX SCF computation
implicit	Run SCF and differentiate its fixed-point condition

explicit is the public name; SCF uses lax.scan, not manual expansion of iteration blocks. The training guide defines labels, MSE/MAE losses and update rejection. Samples capture prepared states; reconstruct them after changing the source. External callbacks must support the derivative order and physical inputs required by their objective.

Neural basis sets. NNAO maps MACE outputs to normalized Gaussian contractions. Fixed-primitive direct J/K and RI/DF paths support self-consistent contraction optimization. Contracted joint workflows also optimize supported log-exponent scales, preserving normalization, trainable masks and explicit CPU-integral/model-device transfers. Coefficient, exponent and geometry derivatives have separate contracts in the NNAO guide. See the molecular tool, pilot validator, and shared trainer. Single-geometry or training-set optimization does not establish transferability.

Orbital-free learning. External KEDF callbacks receive density and representation features; their energy supplies potential and response. See neural kinetic energy. Additional dispersion models live in model.neural_d.

Inverse-reference assessment. dft.iks provides AO density targets, block-local potential maps and analytic/neural functional evaluation for RKS/UKS/GKS layouts. The explicitly selected thermal_matrix method constructs an offline finite-temperature finite-basis reference. Evaluation and training against that fixed reference are differentiable; inverse-map AD/JIT is not implemented. See the carbon FCI example for energy/operator training; its objective does not reconverge a candidate SCF density. Zero-temperature Wu–Yang and coupled inverse-GKS optimization/backward remain unfinished.

Model development and ownership

Keep a new neural model in src/gradscf/model/<name>/. Its network, physical integration, training and workflows belong together:

model/<name>/
  config.py, initialization.py   # architecture, physical assets and initialization
  encoder.py, inputs.py          # network and model-specific features
  integrals/                    # model-specific integral execution and caches
  training/                     # objectives, gradients, updates and checkpoints
  workflows/                    # data selection, training and evaluation entry points
  data/                         # required physical assets and provenance
  vendor/                       # pinned upstream source, when needed

Use only the files/subpackages the model needs; a small model can use a single training.py. Model-owned integral modules call the shared gradscf.integrals kernels. SCF, response and implicit/explicit derivative solvers remain in their physical modules. Training loops and workflow policy are maintained by each model, rather than a separate global framework.

The standard development sequence is:

1. Specify the predicted physical quantity, units, shapes, symmetries and fixed versus trainable parameters.
2. Implement initialization and network evaluation with explicit parameter and non-training state; importing a model must not start calculations or choose a device.
3. Assemble physical inputs using the existing normalization and integral operators; declare cache dependencies and supported derivative orders.
4. Define the model's objective and compose its gradient through the existing electronic solver. Variational basis optimization and supervised XC learning require different objectives.
5. Implement updates, failure handling and reproducible checkpoint recovery within the model. Invalid SCF/gradient candidates must not replace accepted parameters or optimizer state.
6. Verify initialization, directional finite differences, a small training update and checkpoint continuation before measuring larger workloads.

Current owners are neural_xc/training, neural_xc/workflows, and NNAO's integrals, training and workflows. The unchanged MACE snapshot stays inside NNAO at nnao/vendor/mace-jax. See the model development guide and NNAO guide for the full contracts and commands.

Numerical infrastructure and memory

gradscf.solvers owns Hermitian/non-Hermitian diagonalization, linear solves, nonlinear iteration and their derivative rules. Methods supply Hamiltonian actions, residuals and physical convergence conditions.

The integral layer uses libcint as its sole production Gaussian integral interface, with exact s4/s8 storage, shell-direct J/K and explicit auxiliary RI/DF. Ordinary calculations select the physical representation or device through their method controls. Retired integral-engine selectors (cpu, native, jax, jax_reference) are rejected. JAX continues to provide AO/grid evaluation, normalization, contractions and differentiation. Necessary scalar ECP C extensions retain their own provenance and licenses within the same private native interface.

The sibling AO-to-MO module owns shared selected dense, packed, DF and native-direct transformations for correlated and response methods. Complete four-index outputs remain quartic when explicitly requested.

• Packing is exact; direct J/K avoids a stored four-index tensor.
• RI/DF introduces an auxiliary-basis approximation. DF-MP2 retains occupied–virtual factors and streams slices; with_t2=False omits doubles output.
• MP3 and some CI/CC paths can expand DF inputs into full MO tensors. Accepting DF input is not a DF contraction algorithm.
• FCI transforms only active orbitals after core folding. BSE offers matrix-free Hamiltonian/screening actions with documented solver options.

Derivative support is parameter- and layout-specific; fixed-primitive contraction training does not imply native exponent/auxiliary-integral AD. Storage bounds and compile estimates are distinct from measured peak memory.

Molecular response execution

Molecular TDA/full TDDFT supports AO-density and DF operator actions with the shared Davidson/RPA solvers. Local restricted semilocal kernels can reuse occupied/virtual MO grid fields prepared from streamed AO blocks. Grid blocks, trial-vector batches, packed J/K workspace and solver subspaces share the response memory plan. Changing the budget changes execution blocks and caches; it does not reduce grid accuracy or truncate the occupied/virtual space.

Repeated compatible response calculations reuse compiled numerical programs while orbitals, integrals and external XC parameters remain live inputs. Global/neural XC callbacks retain their declared nonlocal dependence. Geometry objectives still require the upstream SCF, integral and grid response; supported parameter products are described in the functional and integral contracts.

For eager molecular RKS/UKS with explicit auxiliary RI, density_fit(auxbasis) prepares integrals and the reference on CPU. execution_device="gpu" then places the completed reference on GPU for supported response calculations. Repeated access reuses that resident reference; changed SCF inputs require reconvergence. This is a staged CPU/GPU calculation. Full traced RI geometry derivatives remain outside this path's supported scope.

Examples

Task	Entry points
Correlated ground states	MP2/MP3, MP2–MP6 series, water MP2–MP8 comparison, CI, CC, open-shell properties
FCI/active spaces	Ground state, core/active selection
EOM excitation/ionization/attachment	EE/IP/EA, response
GW/BSE spectra	scGW, implicit GW, molecular BSE, unrestricted BSE
Vibrational spectral effects	Fixed-phonon scGW, water BSE–vibration
Orbital localization/animation	Ethylene Boys, BODIPY, rotation animation
Wavefunction/basis derivatives	CIS probes, implicit HF → MP, counterpoise
Molecular response and properties	Direct response, oscillator-strength response, electric response
Static perturbation theory	Molecular DFPT, periodic DFPT, OFDFT response
Inverse-reference assessment	Carbon FCI and neural functional, independent PySCF reference
External-model training	External functional, density-matrix MLP, H2/FCI modes
OFDFT	Molecular, periodic, DFTpy comparison, WGC scope/results
Periodic methods/solvers	Bands, shared solvers, degenerate subspace

Comparison scripts and literature-derived illustrations are identified in their headers/guides. Follow instructions for optional dependencies, cached inputs and preceding calculations.

Validation and current scope

Validation uses independent software/Hamiltonians, reconverged finite differences, residuals, normalization, symmetry and conservation checks. Systems, basis/XC, units, tolerances, backend and code context are recorded in MP, CI, FCI, CC, EOM, solver, and OFDFT/WGC reports.

These are scoped records, not a whole-project test counter. CPU/float64 comparisons do not establish GPU coverage; finite-basis agreement does not establish model accuracy. Performance comparisons need matching workloads, settings, compilation/warmup policy and timed regions.

The TD response execution report records matched parent/candidate timings, new-object compilation counts, original-operator residuals, CPU/GPU memory metrics and neural/basis AD checks. Its reproduction commands separate cold calls, warm solves and prepared action probes. Reported gains apply to those workloads; GPU packed contraction has a measured time/memory tradeoff rather than a universal speedup.

With test dependencies and native integrals available:

JAX_PLATFORMS=cpu JAX_ENABLE_X64=1 python -m pytest -q -ra tests/mp
JAX_PLATFORMS=cpu JAX_ENABLE_X64=1 python -m pytest -q -ra tests/solvers tests/ci tests/fci tests/cc

Use python -m pytest -q -ra for the full suite and inspect dependency skips. Reference/property/derivative limits are part of each method's contract.

Documentation and contributing

Start with the public API, import migration table and migration guide. Method guides contain equations, derivative contracts, references and validation; CHANGELOG.md records changes. Use short domain imports for calculation objects and owning modules for advanced arrays.

Report bugs through GitHub Issues with a minimal example, commit/version, device/backend, dtype, geometry/units, basis/XC, convergence controls and reference result. Contributions should keep physical equations separate from shared solvers, add focused numerical regressions, and preserve upstream notices.

Repository-local .vscode/ settings and the root docs/ working directory are kept locally and excluded from GitHub synchronization. Published API and model documentation lives in the root Markdown guides and the relevant src/gradscf/ modules.

Citation

Identify the repository and commit/version used in published calculations. Cite the methods and upstream software relevant to your calculation; references and attribution are maintained in MP, CI, FCI, CC, BSE, and GW.

License and acknowledgments

Original GradSCF code uses the MIT license. Foundational DFT/TDDFT code originates from GradTDDFT. Adapted/vendored components retain their own notices and licenses: PySCF CC contractions, native PySCF/libcint, OFDFT/libKEDF, and MACE-JAX provenance.