dft.operators

Kohn-Sham operator and tiny-grid numerical references.

import mlx_atomistic.dft.operators

Classes

`DavidsonDiagonalizer`

class DavidsonDiagonalizer
    def __init__(config: EigensolverConfig = EigensolverConfig(), preconditioner: KineticPreconditioner = KineticPreconditioner())

Small block Davidson-style eigensolver.

The dense reference path is still used for tiny grids. Above that cutoff, this applies a conservative preconditioned residual iteration that avoids building the dense Hamiltonian.

Parameters

Name	Type	Default	Description
`config`	`EigensolverConfig`	`EigensolverConfig()`
`preconditioner`	`KineticPreconditioner`	`KineticPreconditioner()`

Methods

`solve`

def solve(operator: KohnShamOperator, *, n_orbitals: int, initial_orbitals: mx.array | None = None) -> DiagonalizationResult

Solve for the lowest occupied orbitals.

Parameters

Name	Type	Default	Description
`operator`	`KohnShamOperator`		The `KohnShamOperator` to diagonalize.
`n_orbitals`	`int`		Number of occupied orbitals to solve for.
`initial_orbitals`	`mx.array \| None`	`None`	Optional starting orbital stack; `None` uses a deterministic random guess. Defaults to `None`.

Returns

DiagonalizationResult — A DiagonalizationResult for the lowest n_orbitals states (dense reference below the fallback size, preconditioned iteration above).

`DenseHamiltonianReference`

class DenseHamiltonianReference
    def __init__(operator: KohnShamOperator)

Explicit dense Hamiltonian for tiny-grid validation.

Parameters

Name	Type	Default	Description
`operator`	`KohnShamOperator`

Methods

`diagonalize`

def diagonalize(n_orbitals: int) -> DiagonalizationResult

Diagonalize the dense reference matrix.

Parameters

Name	Type	Default	Description
`n_orbitals`	`int`		Number of lowest eigenpairs to return.

Returns

DiagonalizationResult — A DiagonalizationResult with the lowest n_orbitals eigenpairs (orthonormalized).

`matrix`

def matrix() -> np.ndarray

Build the dense matrix by applying the operator to basis vectors.

Returns

np.ndarray — The Hermitized dense Hamiltonian, shape (grid.size, grid.size).

`matvec`

def matvec(orbital: mx.array) -> mx.array

Apply the dense matrix to a single orbital.

Parameters

Name	Type	Default	Description
`orbital`	`mx.array`		A single real-space orbital of shape `grid.shape`.

Returns

mx.array — H ψ as a grid-shaped array.

`DiagonalizationResult`

class DiagonalizationResult
    def __init__(eigenvalues: mx.array, orbitals: mx.array, residuals: mx.array, orthonormality_error: float, iterations: int, converged: bool, metadata: dict = dict())

Small eigenproblem result bundle.

Parameters

Name	Type	Default
`eigenvalues`	`mx.array`
`orbitals`	`mx.array`
`residuals`	`mx.array`
`orthonormality_error`	`float`
`iterations`	`int`
`converged`	`bool`
`metadata`	`dict`	`dict()`

Methods

`to_dict`

def to_dict() -> dict

Return a JSON-safe summary.

Returns

dict — The eigenvalues, residuals, orthonormality error, iteration count, convergence flag, and metadata as a JSON-serializable dict.

`EigensolverConfig`

class EigensolverConfig
    def __init__(max_iterations: int = 16, tolerance: float = 1e-06, max_subspace_size: int = 12, dense_fallback_size: int = 512)

Configuration for iterative Kohn-Sham eigensolvers.

Parameters

Name	Type	Default
`max_iterations`	`int`	`16`
`tolerance`	`float`	`1e-06`
`max_subspace_size`	`int`	`12`
`dense_fallback_size`	`int`	`512`

`KineticPreconditioner`

class KineticPreconditioner
    def __init__(shift: float = 0.5)

Simple reciprocal-space kinetic preconditioner.

Parameters

Name	Type	Default	Description
`shift`	`float`	`0.5`

Methods

`apply`

def apply(residual: mx.array, operator: KohnShamOperator) -> mx.array

Apply 1 / (0.5|G+k|² + shift) to a residual.

Parameters

Name	Type	Default	Description
`residual`	`mx.array`		A real-space residual to precondition.
`operator`	`KohnShamOperator`		The `KohnShamOperator` supplying the grid and k-point.

Returns

mx.array — The preconditioned residual, same shape as residual.

`KohnShamOperator`

class KohnShamOperator
    def __init__(grid: RealSpaceGrid, local_potential: mx.array, density: mx.array, hartree: mx.array, xc_potential: mx.array, effective_potential: mx.array, nonlocal_operator: NonlocalPseudopotentialOperator | None = None, kpoint: tuple[float, float, float] = (0.0, 0.0, 0.0))

Kohn-Sham Hamiltonian operator for a fixed density and local potential.

Parameters

Name	Type	Default
`grid`	`RealSpaceGrid`
`local_potential`	`mx.array`
`density`	`mx.array`
`hartree`	`mx.array`
`xc_potential`	`mx.array`
`effective_potential`	`mx.array`
`nonlocal_operator`	`NonlocalPseudopotentialOperator \| None`	`None`
`kpoint`	`tuple[float, float, float]`	`(0.0, 0.0, 0.0)`

Methods

`apply_hamiltonian`

def apply_hamiltonian(orbitals: mx.array) -> mx.array

Apply H = T + V_eff.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.

Returns

mx.array — H ψ (kinetic + effective local + any nonlocal term), same shape as orbitals.

`apply_hartree_xc_potential`

def apply_hartree_xc_potential(orbitals: mx.array) -> mx.array

Apply only the Hartree plus exchange-correlation potential.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.

Returns

mx.array — (V_H + V_xc) · ψ, same shape as orbitals.

`apply_kinetic`

def apply_kinetic(orbitals: mx.array) -> mx.array

Apply only the kinetic part.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.

Returns

mx.array — T ψ at this operator’s k-point, same shape as orbitals.

`apply_local_potential`

def apply_local_potential(orbitals: mx.array) -> mx.array

Apply only the external local potential.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.

Returns

mx.array — V_local · ψ, same shape as orbitals.

`from_density`

def from_density(grid: RealSpaceGrid, local_potential: mx.array, density: mx.array, *, xc_functional: ExchangeCorrelationFunctional | None = None, density_floor: float = 1e-12, nonlocal_operator: NonlocalPseudopotentialOperator | None = None, kpoint: Sequence[float] | None = None) -> KohnShamOperator

Build an operator from ρ and a local potential.

Parameters

Name	Type	Default	Description
`grid`	`RealSpaceGrid`		Real-space grid defining the FFT mesh.
`local_potential`	`mx.array`		External local (ionic) potential on the grid, shape `grid.shape`.
`density`	`mx.array`		Electron density `ρ` on the grid, shape `grid.shape`.
`xc_functional`	`ExchangeCorrelationFunctional \| None`	`None`	Exchange-correlation functional; `None` uses `LDAExchangeCorrelation`. Defaults to `None`.
`density_floor`	`float`	`1e-12`	Lower clamp on the density for XC numerical stability. Defaults to `1e-12`.
`nonlocal_operator`	`NonlocalPseudopotentialOperator \| None`	`None`	Optional nonlocal pseudopotential operator. Defaults to `None`.
`kpoint`	`Sequence[float] \| None`	`None`	Optional Bloch k-point; `None` uses Γ. Defaults to `None`.

Returns

KohnShamOperator — A KohnShamOperator with the Hartree, XC, and effective potentials precomputed.

`rayleigh_quotients`

def rayleigh_quotients(orbitals: mx.array) -> mx.array

Return <ψᵢ|H|ψᵢ>/<ψᵢ|ψᵢ> for each orbital.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital stack `(n_orbitals, *grid.shape)`.

Returns

mx.array — The per-orbital Rayleigh quotients (energy estimates), shape (n_orbitals,).

`SubspaceDiagonalizer`

class SubspaceDiagonalizer
    def __init__(max_iterations: int = 8, tolerance: float = 1e-06)

Operator-backed Rayleigh-Ritz diagonalizer for tiny prototype grids.

Parameters

Name	Type	Default	Description
`max_iterations`	`int`	`8`
`tolerance`	`float`	`1e-06`

Methods

`solve`

def solve(operator: KohnShamOperator, *, n_orbitals: int) -> DiagonalizationResult

Solve the occupied subspace.

For this milestone the implementation intentionally uses the explicit tiny-grid reference matrix as the Rayleigh-Ritz space. The public solver boundary lets us replace this with an iterative Davidson path later.

Parameters

Name	Type	Default	Description
`operator`	`KohnShamOperator`		The `KohnShamOperator` to diagonalize.
`n_orbitals`	`int`		Number of occupied orbitals to solve for.

Returns

DiagonalizationResult — A DiagonalizationResult for the lowest n_orbitals states.

Functions

`apply_hamiltonian`

def apply_hamiltonian(orbitals: mx.array, effective_potential: mx.array, grid: RealSpaceGrid) -> mx.array

Apply T + V_eff using an explicit effective potential.

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)`.
`effective_potential`	`mx.array`	Combined local effective potential on the grid, shape `grid.shape`.
`grid`	`RealSpaceGrid`	Real-space grid defining the FFT mesh.

Returns

mx.array — (T + V_eff) ψ, same shape as orbitals.

`apply_hartree_xc_potential`

def apply_hartree_xc_potential(orbitals: mx.array, hartree: mx.array, xc_potential: mx.array) -> mx.array

Apply Hartree plus exchange-correlation potential to orbitals.

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)`.
`hartree`	`mx.array`	Hartree potential on the grid, shape `grid.shape`.
`xc_potential`	`mx.array`	Exchange-correlation potential on the grid, shape `grid.shape`.

Returns

mx.array — The product (V_H + V_xc) · ψ, same shape as orbitals.

`apply_kinetic`

def apply_kinetic(orbitals: mx.array, grid: RealSpaceGrid, *, kpoint: Sequence[float] | None = None) -> mx.array

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

Parameters

Name	Type	Default	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 defining the FFT mesh.
`kpoint`	`Sequence[float] \| None`	`None`	Optional k-point 3-vector added to the reciprocal-grid G-vectors; `None` uses the Γ point. Defaults to `None`.

Returns

mx.array — The kinetic operator applied to orbitals, same shape as orbitals.

`apply_local_potential`

def apply_local_potential(orbitals: mx.array, local_potential: mx.array) -> mx.array

Apply a local multiplicative potential to orbitals.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital(s); a single orbital of shape `grid.shape` or a stack `(n_orbitals, *grid.shape)`.
`local_potential`	`mx.array`		Local potential on the grid, shape `grid.shape`.

Returns

mx.array — The elementwise product V · ψ, same shape as orbitals.

`orbital_residuals`

def orbital_residuals(orbitals: mx.array, operator: KohnShamOperator, eigenvalues: mx.array | None = None) -> mx.array

Return ||Hψᵢ - εᵢψᵢ|| for each orbital.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital stack `(n_orbitals, *grid.shape)`.
`operator`	`KohnShamOperator`		The `KohnShamOperator` supplying `H`.
`eigenvalues`	`mx.array \| None`	`None`	Optional per-orbital eigenvalues εᵢ; `None` uses the Rayleigh quotients. Defaults to `None`.

Returns

mx.array — The per-orbital residual norms, shape (n_orbitals,).

`orthonormality_error`

def orthonormality_error(orbitals: mx.array, grid: RealSpaceGrid) -> float

Return max absolute overlap error from the identity matrix.

Parameters

Name	Type	Default	Description
`orbitals`	`mx.array`		Real-space orbital stack `(n_orbitals, *grid.shape)`.
`grid`	`RealSpaceGrid`		Real-space grid; its `dv` weights the overlap integral.

Returns

float — The maximum absolute deviation of the overlap matrix ⟨ψᵢ|ψⱼ⟩ from the identity.