dft.potentials

Toy DFT potentials and energy helpers.

import mlx_atomistic.dft.potentials

Classes

`LocalGaussianPseudopotential`

class LocalGaussianPseudopotential
    def __init__(centers: Sequence[Sequence[float]], amplitudes: Sequence[float] | float, widths: Sequence[float] | float)

Simple periodic local Gaussian pseudopotential for toy SCF examples.

Parameters

Name	Type	Default	Description
`centers`	`Sequence[Sequence[float]]`
`amplitudes`	`Sequence[float] \| float`
`widths`	`Sequence[float] \| float`

Methods

`field`

def field(grid: RealSpaceGrid) -> mx.array

Evaluate the local potential on a real-space grid.

Parameters

Name	Type	Default	Description
`grid`	`RealSpaceGrid`		Real-space grid to sample the potential on.

Returns

mx.array — The summed Gaussian local potential on the grid, shape grid.shape.

Functions

`apply_kinetic`

def apply_kinetic(orbital: mx.array, grid: RealSpaceGrid) -> mx.array

Apply the plane-wave kinetic operator -1/2 ∇² to one orbital.

Parameters

Name	Type	Default	Description
`orbital`	`mx.array`		A single real-space orbital, shape `grid.shape`.
`grid`	`RealSpaceGrid`		Real-space grid defining the FFT mesh.

Returns

mx.array — The kinetic operator applied to the orbital (real part), shape grid.shape.

`density_from_normalized_orbitals`

def density_from_normalized_orbitals(orbitals: mx.array, grid: RealSpaceGrid, *, occupations: Sequence[float]) -> mx.array

Build a density after normalizing orbitals.

Parameters

Name	Type	Description
`orbitals`	`mx.array`	Real-space orbital stack `(n_orbitals, *grid.shape)`.
`grid`	`RealSpaceGrid`	Real-space grid defining the FFT mesh.
`occupations`	`Sequence[float]`	Per-orbital occupation numbers.

Returns

mx.array — The electron density ρ = Σᵢ fᵢ |ψᵢ|² on the grid, shape grid.shape.

`electron_count`

def electron_count(density: mx.array, grid: RealSpaceGrid) -> mx.array

Integrate a density over the real-space grid.

Parameters

Name	Type	Default	Description
`density`	`mx.array`		Electron density `ρ` on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid`		Real-space grid; its `dv` weights the integral.

Returns

mx.array — The scalar integrated electron count ∫ρ dr.

`energy_decomposition`

def energy_decomposition(orbitals: mx.array, density: mx.array, local_potential: mx.array, grid: RealSpaceGrid, *, occupations: Sequence[float], density_floor: float = 1e-12, xc_functional: ExchangeCorrelationFunctional | None = None) -> dict[str, mx.array]

Return core toy-DFT energy terms.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space occupied orbital(s).
`density`	`mx.array`		Electron density `ρ` on the grid, shape `grid.shape`.
`local_potential`	`mx.array`		External local potential on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid`		Real-space grid; its `dv` weights the integrals.
`occupations`	`Sequence[float]`		Per-orbital occupation numbers.
`density_floor`	`float`	`1e-12`	Lower clamp on the density for XC stability. Defaults to `1e-12`.
`xc_functional`	`ExchangeCorrelationFunctional \| None`	`None`	Exchange-correlation functional; `None` uses `DiracExchange`. Defaults to `None`.

Returns

dict[str, mx.array] — A dict of energy terms — kinetic, local, hartree, xc, total (plus exchange/correlation when the functional exposes them).

`hartree_potential`

def hartree_potential(density: mx.array, grid: RealSpaceGrid) -> mx.array

Solve the periodic Hartree potential for ρ with the G = 0 term removed.

Parameters

Name	Type	Default	Description
`density`	`mx.array`		Electron density `ρ` on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid`		Real-space grid defining the FFT mesh.

Returns

mx.array — The real-space Hartree potential V_H, shape grid.shape (the G = 0 term is dropped against a neutralizing background).

`kinetic_energy`

def kinetic_energy(orbitals: mx.array, grid: RealSpaceGrid, *, occupations: Sequence[float]) -> mx.array

Return the occupied one-particle kinetic energy.

Parameters

Name	Type	Description
`orbitals`	`mx.array`	Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.
`grid`	`RealSpaceGrid`	Real-space grid; its `dv` weights the integral.
`occupations`	`Sequence[float]`	Per-orbital occupation numbers.

Returns

mx.array — The scalar occupied kinetic energy Σᵢ fᵢ ⟨ψᵢ|T|ψᵢ⟩.

`lda_exchange_energy_potential`

def lda_exchange_energy_potential(density: mx.array, grid: RealSpaceGrid | None = None, *, density_floor: float = 1e-12) -> tuple[mx.array, mx.array]

Return Dirac LDA exchange energy and potential for an unpolarized density.

Parameters

Name	Type	Default	Description
`density`	`mx.array`		Electron density `ρ` on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid \| None`	`None`	Real-space grid providing the integration weight `dv`; `None` uses unit weight. Defaults to `None`.
`density_floor`	`float`	`1e-12`	Lower clamp on the density for numerical stability. Defaults to `1e-12`.

Returns

tuple[mx.array, mx.array] — An (energy, potential) tuple: the scalar exchange energy and the exchange potential on the grid.

`local_pseudopotential_forces`

def local_pseudopotential_forces(density: mx.array, grid: RealSpaceGrid, pseudopotential: LocalGaussianPseudopotential) -> mx.array

Return Hellmann-Feynman forces on local Gaussian centers.

The local energy is ∫ρ(r)V_loc(r; R_I)dr. This computes -∂E/∂R_I for the Gaussian center coordinates.

Parameters

Name	Type	Description
`density`	`mx.array`	Electron density `ρ` on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid`	Real-space grid; its `dv` weights the integral.
`pseudopotential`	`LocalGaussianPseudopotential`	The `LocalGaussianPseudopotential` whose centers feel the force.

Returns

mx.array — The force on each Gaussian center, shape (n_centers, 3).