dft.optimization

Fixed-cell DFT geometry optimization workflows.

import mlx_atomistic.dft.optimization

Classes

GeometryOptimizationConfig

class GeometryOptimizationConfig
    def __init__(max_steps: int = 25, force_tolerance: float = 0.001, energy_tolerance: float = 1e-06, initial_step_size: float = 0.08, max_step: float = 0.25, line_search_shrink: float = 0.5, line_search_min_step: float = 0.0001, max_line_search_iterations: int = 8, optimizer: GeometryOptimizer = 'lbfgs', scf_config: SCFConfig | None = None, reuse_scf_state: bool = True, relaxation_mode: GeometryRelaxationMode = 'ions', stress_tolerance: float = 0.001, cell_step_size: float = 0.02)

Configuration for fixed-cell ion-position relaxation.

Parameters

Name	Type	Default	Description
max_steps	int	25
force_tolerance	float	0.001
energy_tolerance	float	1e-06
initial_step_size	float	0.08
max_step	float	0.25
line_search_shrink	float	0.5
line_search_min_step	float	0.0001
max_line_search_iterations	int	8
optimizer	GeometryOptimizer	'lbfgs'
scf_config	SCFConfig | None	None
reuse_scf_state	bool	True
relaxation_mode	GeometryRelaxationMode	'ions'
stress_tolerance	float	0.001
cell_step_size	float	0.02

GeometryOptimizationRecord

class GeometryOptimizationRecord
    def __init__(positions: np.ndarray, forces: np.ndarray, energies: np.ndarray, max_forces: np.ndarray, statuses: tuple[str, ...], metadata: dict[str, Any], history: list[dict[str, Any]])

Loaded NPZ geometry-optimization history.

Parameters

Name	Type	Default	Description
positions	np.ndarray
forces	np.ndarray
energies	np.ndarray
max_forces	np.ndarray
statuses	tuple[str, ...]
metadata	dict[str, Any]
history	list[dict[str, Any]]

Methods

to_dict

def to_dict() -> dict[str, Any]

Return a JSON-safe summary.

Returns

• dict[str, Any]

GeometryOptimizationResult

class GeometryOptimizationResult
    def __init__(status: GeometryStatus, convergence_reason: str, initial_system: DFTSystem, final_system: DFTSystem, final_scf: SCFResult | None, steps: tuple[GeometryOptimizationStep, ...], config: GeometryOptimizationConfig, elapsed_ms: float)

Result bundle for a fixed-cell DFT geometry optimization.

Parameters

Name	Type	Default	Description
status	GeometryStatus
convergence_reason	str
initial_system	DFTSystem
final_system	DFTSystem
final_scf	SCFResult | None
steps	tuple[GeometryOptimizationStep, ...]
config	GeometryOptimizationConfig
elapsed_ms	float

Properties

• converged bool — Whether the geometry optimization met a convergence criterion.
• final_energy float | None — Final total energy, if an SCF result is available.
• final_forces np.ndarray | None — Final ion-center forces, if available.
• final_max_force float | None — Final maximum per-center force norm.
• final_positions np.ndarray — Final ion-center positions.

Methods

to_dict

def to_dict() -> dict[str, Any]

Return a JSON-safe summary without dense SCF arrays.

Returns

• dict[str, Any]

GeometryOptimizationStep

class GeometryOptimizationStep
    def __init__(index: int, energy: float, energy_delta: float, max_force: float, rms_force: float, force_norm: float, step_norm: float, accepted_step_size: float, line_search_iterations: int, scf_status: str, scf_iterations: int, scf_residual: float, electron_count: float, timing_summary: dict[str, float], positions: np.ndarray, forces: np.ndarray, status: str = 'accepted', stress_norm: float | None = None, rejected_reason: str | None = None)

One accepted geometry-optimization step.

Parameters

Name	Type	Default	Description
index	int
energy	float
energy_delta	float
max_force	float
rms_force	float
force_norm	float
step_norm	float
accepted_step_size	float
line_search_iterations	int
scf_status	str
scf_iterations	int
scf_residual	float
electron_count	float
timing_summary	dict[str, float]
positions	np.ndarray
forces	np.ndarray
status	str	'accepted'
stress_norm	float | None	None
rejected_reason	str | None	None

Methods

to_dict

def to_dict() -> dict[str, Any]

Return a JSON-safe step summary.

Returns

• dict[str, Any]

Functions

geometry_demo_system

def geometry_demo_system(name: str, *, grid_shape: Sequence[int] = (4, 4, 4)) -> DFTSystem

Build one of the compact DFT geometry-optimization demo systems.

Parameters

Name	Type	Default	Description
name	str
grid_shape	Sequence[int]	(4, 4, 4)

Returns

• DFTSystem

load_geometry_optimization

def load_geometry_optimization(path: str | Path) -> GeometryOptimizationRecord

Load a compressed geometry-optimization history.

Parameters

Name	Type	Default	Description
path	str | Path

Returns

• GeometryOptimizationRecord

optimize_geometry

def optimize_geometry(system: DFTSystem, *, config: GeometryOptimizationConfig | None = None, xc_functional: ExchangeCorrelationFunctional | None = None) -> GeometryOptimizationResult

Relax ion positions in a fixed orthorhombic cell using SCF forces.

Parameters

Name	Type	Default	Description
system	DFTSystem
config	GeometryOptimizationConfig | None	None
xc_functional	ExchangeCorrelationFunctional | None	None

Returns

• GeometryOptimizationResult

save_geometry_optimization

def save_geometry_optimization(path: str | Path, result: GeometryOptimizationResult, *, metadata: dict[str, Any] | None = None) -> None

Save a geometry-optimization history to compressed NPZ plus JSON metadata.

Parameters

Name	Type	Default	Description
path	str | Path
result	GeometryOptimizationResult
metadata	dict[str, Any] | None	None

Returns

• None