Metadata-Version: 2.4
Name: jax_hf
Version: 2.2.1
Summary: A package that provides a jax-based Hartree-Fock optimization solver for simple continuum models.
Author: Tobias Wolf
Author-email: public@wolft.xyz
License-Expression: BSD-2-Clause
Project-URL: Homepage, https://github.com/skilledwolf/jax_hf
Project-URL: Source, https://github.com/skilledwolf/jax_hf
Project-URL: Issues, https://github.com/skilledwolf/jax_hf/issues
Keywords: physics,condensed matter,hartree fock,mean field,scf
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Requires-Dist: jax
Requires-Dist: optimistix
Requires-Dist: optax
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: pytest; extra == "dev"
Dynamic: license-file

## jax_hf — JAX Hartree–Fock on k‑grids

[![PyPI](https://img.shields.io/pypi/v/jax-hf.svg)](https://pypi.org/project/jax-hf/)
[![Python](https://img.shields.io/pypi/pyversions/jax-hf.svg)](https://pypi.org/project/jax-hf/)
[![Wheel](https://img.shields.io/pypi/wheel/jax-hf.svg)](https://pypi.org/project/jax-hf/#files)
[![License](https://img.shields.io/pypi/l/jax-hf.svg)](LICENSE)
[![Build](https://github.com/skilledwolf/jax_hf/actions/workflows/build-and-test.yml/badge.svg)](https://github.com/skilledwolf/jax_hf/actions/workflows/build-and-test.yml)
[![Release](https://github.com/skilledwolf/jax_hf/actions/workflows/release.yml/badge.svg)](https://github.com/skilledwolf/jax_hf/actions/workflows/release.yml)

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`](MIGRATION.md) for the migration guide.

### Install

```bash
pip install jax-hf
```

### Minimal example

```python
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:

```python
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

```python
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:

```python
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

```bash
pytest tests/
```

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