diffpes

Differentiable ARPES simulation in JAX: from tight-binding or DFT electronic structure to band structures, spectra, Fermi surfaces, and detector-level photoemission data. Gradients flow through the entire chain.

Tutorials · Guides · API reference

pip install diffpes

From bands to spectra

The graphene π bands and their simulated ARPES spectrum: the occupied band is bright. The Fermi function at 300 K cuts off the unoccupied band.

The self-energy is an explicit model parameter — the same cut at 20, 80, and 250 meV linewidth.

Stagger the two onsite energies and the Dirac crossing opens into a gap; ARPES sees the valence band.

The Fermi edge at 25, 100, and 300 K.

Across the Brillouin zone

Constant-energy slices: Dirac points grow into trigonally warped pockets, touch at the van Hove singularity, and close into a ring around Γ.

Move the Fermi level and the Fermi surface follows — hole-doped by 1 eV.

The intensity cube

The full cube I(kx, ky, E) around one Dirac point, rendered by intensity transparency over its deepest constant-energy slice.

Energy-window integrals of the cube — rings collapse onto the cone apex.

EDC and MDC stacks pulled from the same cut.

Through the detector

The same physics before and after the instrument: analyser calibration, point-spread functions, transmission, background, exposure, and counting.

Expected counts and one reproducible Poisson acquisition.

Rotate the photon polarization and difference the two acquisitions — matrix-element contrast at the detector.

First spectrum

This script is complete — a graphene π-band model, a momentum cut through the Dirac point, and the finite-temperature spectral function:

import jax.numpy as jnp
import matplotlib.pyplot as plt

from diffpes.plots import plot_arpes_spectrum
from diffpes.simul import assemble_spectral_intensity_bands_chunk
from diffpes.tightb import (
    build_kpath,
    diagonalize_tb,
    kpath_arc_length,
    kpoints_frac_to_cart,
)
from diffpes.types import (
    make_arpes_spectrum,
    make_crystal_geometry,
    make_orbital_basis,
    make_self_energy_model,
    make_tb_model,
)

# Graphene pi bands: two carbon sites, six nearest-neighbour hoppings.
a = 2.46
crystal = make_crystal_geometry(
    lattice=jnp.asarray(
        [[a, 0.0, 0.0], [a / 2, a * 3**0.5 / 2, 0.0], [0.0, 0.0, 20.0]]
    ),
    positions=jnp.asarray([[0.0, 0.0, 0.0], [1 / 3, 1 / 3, 0.0]]),
    species=("C", "C"),
)
basis = make_orbital_basis(
    atom_indices=(0, 1), n=(2, 2), l=(0, 0), m=(0, 0), labels=("pz_A", "pz_B")
)
model = make_tb_model(
    hopping_amplitudes=-2.7 * jnp.ones(6, dtype=jnp.complex128),
    onsite_energies=jnp.zeros(2),
    soc_lambdas=jnp.zeros(0),
    geometry=crystal,
    basis=basis,
    hopping_pairs=((0, 1), (0, 1), (0, 1), (1, 0), (1, 0), (1, 0)),
    hopping_cells=(
        (0, 0, 0), (-1, 0, 0), (0, -1, 0),
        (0, 0, 0), (1, 0, 0), (0, 1, 0),
    ),
    shell_index=(-1, -1),
)

# A straight momentum cut through the Dirac point at K.
path = build_kpath(
    jnp.asarray([[0.0, 0.0, 0.0], [0.5, 1.0, 0.0]]),
    crystal,
    301,
    ("Gamma", "K'"),
)
bands = diagonalize_tb(model, path.kpoints)

# Occupied spectral function: Lorentzian self-energy + Fermi cutoff at 300 K.
energies = jnp.linspace(-9.2, 1.2, 480)
intensity = assemble_spectral_intensity_bands_chunk(
    bands.eigenvalues,
    jnp.ones((path.kpoints.shape[0], energies.shape[0], 2)),
    energies,
    make_self_energy_model(gamma=0.09),
    jnp.asarray(0.0),
    300.0,
    allow_degenerate_value_only=True,
)
spectrum = make_arpes_spectrum(
    intensity,
    energies,
    kpath_arc_length(path, crystal),
    kpoints_frac_to_cart(path.kpoints, crystal),
)
plot_arpes_spectrum(spectrum, cmap="magma")
plt.show()

From DFT

A real material, straight from VASP output: the Bi₂Se₃ slab bands along M–Γ–M and their occupied spectrum at 35 K.

The near-Fermi window of the same calculation, sharpened to a 12 meV linewidth.

Slab and bulk calculations on the same M–Γ–K–M path: the slab carries states inside the bulk gap.

EDCs around Γ from the same map.

DOSCAR densities of states and the CHGCAR charge density, resolving all six quintuple layers.

diffpes.inout reads EIGENVAL, PROCAR, POSCAR, KPOINTS, OUTCAR, DOSCAR, CHGCAR, WAVECAR, Wannier90 hr.dat/tb.dat, and Cartesian hopping lists. Parsed eigenvalues and orbital projections drop into the same spectral calls as tight-binding models, so a converged VASP calculation becomes a simulated ARPES measurement.

Differentiable end to end

jax.grad, jax.vmap, and jax.jit work through the whole pipeline — crystal geometry, hoppings, self-energy, matrix elements, experiment geometry, and detector response to expected counts. Fit any of it to measured spectra by gradient descent.

Agent-runnable experiments

The automatons/ directory provides standalone experiments for automated diffpes workflows. Each file supports discovery, parameter validation, smoke execution, and a final JSON result.

Run a small example with:

JAX_PLATFORMS=cpu MPLCONFIGDIR=/tmp/dp-mpl .venv/bin/python \
  automatons/forward_bands.py --smoke --outdir /tmp/dp-bands --json

Read the experiment catalog before selecting a file. Read the agent guide for the complete execution protocol.