jax_hf — JAX Hartree–Fock on k‑grids

jax_hf provides two JAX-jitted solvers for the Hartree–Fock free-energy minimisation problem on 2D k-meshes:

• Direct minimisation (primary): preconditioned Riemannian CG on Stiefel × capped simplex, eigen-free inner loop, Cayley retraction, one Fock build per iteration. An optional trust-region Newton variant (optimizer="newton") converges superlinearly (a handful of outer steps) on stiff problems, and deflation (solve_deflated) finds distinct self-consistent solutions on non-convex landscapes.
• Reference SCF (baseline / fallback): standard Roothaan iteration with linear mixing, Pulay DIIS, or ODA (optimal damping) acceleration, plus an optional trust_radius on the density step.

Exchange and Hartree can both be included, and the exchange kernel may be layer-resolved. See examples/ for density-scan scripts on a bilayer graphene model.

v2.0.0 note: This release is a clean-slate rewrite. The entire public API has changed relative to the deprecated v1.x line (which was already a skeleton in v1.1.0). See MIGRATION.md for the migration guide.

Install

pip install jax-hf

Minimal example

import jax.numpy as jnp
import jax_hf

# Build a HartreeFockKernel: precomputes the FFT of the interaction kernel,
# the Hartree matrix, etc., ready for JIT.
kernel = jax_hf.HartreeFockKernel(
    weights=weights,          # (nk1, nk2) k-point weights
    hamiltonian=hamiltonian,  # (nk1, nk2, nb, nb) single-particle Hamiltonian
    coulomb_q=coulomb_q,      # (nk1, nk2, 1, 1) scalar or (nk1, nk2, nb, nb) layer-resolved
    T=0.1,
    include_hartree=False,    # set True for Hartree; also pass reference_density + hartree_matrix
    include_exchange=True,
)

# Solve (direct minimisation, default)
result = jax_hf.solve(kernel, P0=jnp.zeros_like(hamiltonian), n_electrons=N)
print(result.energy, result.converged, result.n_iter)
# result.density, result.fock, result.Q, result.p, result.mu, result.history

# Or use SCF as a fallback baseline
result_scf = jax_hf.solve_scf(kernel, P0=jnp.zeros_like(hamiltonian), n_electrons=N)

Config

Both solvers take a Config dataclass with sensible defaults:

jax_hf.SolverConfig(max_iter=200, tol_E=1e-7, optimizer="cg", ...)   # or optimizer="newton"
jax_hf.SCFConfig(max_iter=200, mixing=0.3, acceleration="linear", ...)  # or "diis" / "oda"

project_fn lets you enforce symmetry constraints (spin, valley, time reversal, spatial) on the density and Fock at every iteration. See jax_hf.symmetry.make_project_fn.

Newton, deflation, and SCF acceleration

from jax_hf import SolverConfig, SCFConfig, solve, solve_scf, solve_deflated

# Trust-region Newton: far fewer Fock builds than CG on stiff problems.
# (A second-order method — enable x64 / build the kernel in float64.)
r = solve(kernel, P0, N, config=SolverConfig(optimizer="newton", tol_grad=1e-6))

# Deflation: find *distinct* HF solutions (e.g. competing broken-symmetry
# phases).  Returns them sorted by energy; `best` is the ground-state candidate.
res = solve_deflated(kernel, P0, N, n_solutions=4)
print(res.energies, res.best.energy)

# SCF with Pulay DIIS (typically ~10x fewer iterations than linear mixing):
r = solve_scf(kernel, P0, N, config=SCFConfig(acceleration="diis"))

Public API

Name	Purpose
HartreeFockKernel	Problem + precomputed arrays
solve (alias solve_direct_minimization), SolverConfig, SolveResult	Primary solver (CG + trust-region Newton)
solve_deflated, DeflatedResult	Find distinct HF solutions (deflated Newton)
solve_scf, SCFConfig, SCFResult	Reference SCF solver (linear / DIIS / ODA)
build_fock, hf_energy, free_energy, occupation_entropy	HF objective building blocks
solve_continuation, ContinuationResult, resample_kgrid	Coarse → fine multigrid driver + k-grid resampler

Lower-level modules (jax_hf.utils, jax_hf.symmetry, jax_hf.linalg, jax_hf.fock) expose the individual pieces for users who need them.

Coarse → fine continuation

For large fine grids, solve_continuation runs a cheap coarse solve first and uses its density to seed the fine solve. The two stages can mix and match direct minimisation and SCF:

import jax_hf
from jax_hf import SCFConfig, SolverConfig

coarse = jax_hf.HartreeFockKernel(weights_c, h_c, Vq_c, T=0.1)
fine   = jax_hf.HartreeFockKernel(weights_f, h_f, Vq_f, T=0.1)

result = jax_hf.solve_continuation(
    coarse, fine, P0_coarse=jnp.zeros_like(h_c),
    n_electrons_coarse=N, n_electrons_fine=N,
    coarse_config=SCFConfig(max_iter=50, mixing=0.5),   # robust coarse
    fine_config=SolverConfig(max_iter=200, tol_E=1e-8), # fast fine
)
# result.coarse, result.fine (each a SolveResult or SCFResult)
# result.P0_fine (resampled coarse density used to seed the fine solve)

The driver is intentionally algorithm-agnostic: it resamples the coarse density onto the fine grid via resample_kgrid and hands off. Callers that need physics-aware seeding (reference-density interpolation, self-energy seeds, filling-consistent electron counts across grids) should construct both kernels themselves.

Examples

• examples/multilayer_graphene_density_scan.py — PM/SVP density scan for bilayer graphene, direct minimisation, Fock only
• examples/multilayer_graphene_density_scan_extended.py — adds spin-polarised and "SVP flipped" branches (4 total)
• examples/multilayer_graphene_density_scan_hartree.py — same four branches with layer-resolved Coulomb and Hartree included
• examples/multilayer_graphene_reference_scf_scan.py — SCF baseline scan for side-by-side comparison

Running tests

pytest tests/

The bilayer regression tests (tests/test_bilayer_regression.py) require contimod and contimod_graphene and will be skipped otherwise.