dft.gga

PBE-form GGA exchange-correlation with an autodiff-derived potential.

The defining idea: a GGA energy is E_xc = ∫ ε(ρ, ∇ρ) dr, and its potential is the functional derivative v_xc = δE_xc/δρ = ∂ε/∂ρ - ∇·(∂ε/∂∇ρ). The gradient term is the part that is famously error-prone to hand-derive. Here we write only the energy density and obtain v_xc from mx.grad of the total energy — the autodiff machinery reconstructs the divergence term automatically, provided the density gradient is built with a differentiable (MLX-native) FFT.

import mlx_atomistic.dft.gga

Classes

PBEExchangeCorrelation

class PBEExchangeCorrelation
    def __init__(name: str = 'pbe-pz81-gga-alpha')

Alpha PBE-form GGA exchange-correlation with a PZ81 uniform baseline.

The exchange enhancement and correlation gradient terms follow the PBE form, while the uniform-gas correlation baseline reuses the PZ81 parameterization already in the package. Full production PBE uses a PW92 baseline, so this public alpha class intentionally reports an explicit alpha result name.

Parameters

Name	Type	Default	Description
name	str	'pbe-pz81-gga-alpha'

Methods

evaluate

def evaluate(density: mx.array, grid: RealSpaceGrid | None = None, *, density_floor: float = 1e-12) -> XCResult

Evaluate alpha PBE-form GGA energy density, potential, and total energy.

Parameters

Name	Type	Default	Description
density	mx.array		Electron density rho sampled on the grid.
grid	RealSpaceGrid | None	None	Real-space grid; required for GGA to evaluate the density gradient. Defaults to None.
density_floor	float	1e-12	Lower clamp applied to the density for numerical stability. Defaults to 1e-12.

Returns

• XCResult — An XCResult with the energy density, potential, and total energy.

Raises

• ValueError — If grid is None (a grid is required for GGA).

ProductionPBEExchangeCorrelation

class ProductionPBEExchangeCorrelation
    def __init__(name: str = 'pbe-pw92-gga')

PBE GGA exchange-correlation with the PW92 uniform-gas baseline.

Parameters

Name	Type	Default	Description
name	str	'pbe-pw92-gga'

Methods

evaluate

def evaluate(density: mx.array, grid: RealSpaceGrid | None = None, *, density_floor: float = 1e-12) -> XCResult

Evaluate production PBE energy density, potential, and total energy.

Parameters

Name	Type	Default	Description
density	mx.array		Electron density rho sampled on the grid.
grid	RealSpaceGrid | None	None	Real-space grid required for the GGA gradient. Defaults to None.
density_floor	float	1e-12	Lower density clamp. Defaults to 1e-12.

Returns

• XCResult — Production PBE energy density, potential, and total energy.

Raises

• ValueError — If no real-space grid is provided.

Functions

density_gradient

def density_gradient(rho: mx.array, grid: RealSpaceGrid) -> mx.array

Return ∇ρ (shape (3, *grid.shape)) via reciprocal-space differentiation.

Parameters

Name	Type	Default	Description
rho	mx.array
grid	RealSpaceGrid

Returns

• mx.array