dft.plane_wave

Compact cutoff-projected plane-wave bases for periodic DFT.

import mlx_atomistic.dft.plane_wave

Classes

PlaneWaveBasis

class PlaneWaveBasis
    def __init__(grid: RealSpaceGrid, cutoff_hartree: float, kpoint_cartesian: Sequence[float] = (0.0, 0.0, 0.0), *, reciprocal_grid: ReciprocalGrid | None = None, lane_label: str = 'lane:0')

Compact plane-wave basis with a public full-grid compatibility facade.

Runtime-owned coefficients use canonical ascending FFT-index order with shape (vectors, active_count). Public methods continue to accept and return full FFT grids, materializing them only at the call boundary.

Parameters

Name	Type	Default	Description
grid	RealSpaceGrid		Uniform orthorhombic real-space grid.
cutoff_hartree	float		Kinetic energy cutoff in Hartree.
kpoint_cartesian	Sequence[float]	(0.0, 0.0, 0.0)	Bloch k-point in inverse bohr. Defaults to Gamma.
reciprocal_grid	ReciprocalGrid | None	None	Optional shared reciprocal descriptor. Compatible bases created for one SCF calculation should share this object.
lane_label	str	'lane:0'	Stable label used to distinguish equal-basis runtime lanes.

Properties

• active_count int — Number of active plane waves.
• active_flat_indices mx.array — Fresh canonical ascending flat FFT indices.
• active_integer_g mx.array — Fresh ordered exact integer G coordinates.
• active_kinetic_energies mx.array — Fresh kinetic energies in compact order.
• active_shifted_vectors mx.array — Fresh shifted reciprocal vectors in compact order.
• basis_fingerprint str — Complete compact basis fingerprint.
• kinetic_energies mx.array — Fresh full-grid kinetic energies for compatibility callers.
• lane_id str — Immutable identity of this runtime lane.
• mask mx.array — Fresh full-grid cutoff mask for compatibility callers.
• order_fingerprint str — Canonical active-order fingerprint.
• reciprocal_grid ReciprocalGrid — Shared reciprocal descriptor for this basis.
• reciprocal_vectors mx.array — Shared full-grid unshifted reciprocal vectors.
• shifted_vectors mx.array — Fresh full-grid shifted vectors for compatibility callers.
• volume float — Cell volume in bohr cubed.

Methods

apply_kinetic

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

Apply 0.5 |G+k|^2 in compact coefficient space.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.

Returns

• mx.array — Fresh full-grid kinetic action with exact inactive zeros.

apply_local

def apply_local(coefficients: mx.array, potential: mx.array) -> mx.array

Apply a periodic local potential through one batched FFT pair.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.
potential	mx.array		Real local potential with shape grid.shape.

Returns

• mx.array — Fresh full-grid local-potential action.

Raises

• ValueError — If the potential shape does not match the grid.

coefficient_norms

def coefficient_norms(coefficients: mx.array) -> mx.array

Return coefficient-space norms for one orbital or a stack.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.

Returns

• mx.array — One norm per orbital.

from_reduced_kpoint

def from_reduced_kpoint(grid: RealSpaceGrid, cutoff_hartree: float, reduced_kpoint: Sequence[float], *, reciprocal_grid: ReciprocalGrid | None = None, lane_label: str = 'lane:0') -> PlaneWaveBasis

Build a basis from fractional reciprocal coordinates.

Parameters

Name	Type	Default	Description
grid	RealSpaceGrid		Uniform orthorhombic real-space grid.
cutoff_hartree	float		Kinetic energy cutoff in Hartree.
reduced_kpoint	Sequence[float]		Fractional coordinates along reciprocal cell axes.
reciprocal_grid	ReciprocalGrid | None	None	Optional shared reciprocal descriptor.
lane_label	str	'lane:0'	Stable runtime lane label.

Returns

• PlaneWaveBasis — A basis whose Cartesian k-point is 2*pi*k_i/L_i.

normalize

def normalize(coefficients: mx.array) -> mx.array

Normalize each orbital in compact coefficient space.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.

Returns

• mx.array — Fresh full-grid unit-norm coefficients.

Raises

• ValueError — If any orbital has zero norm.

orthonormalize

def orthonormalize(coefficients: mx.array) -> mx.array

Orthonormalize an admitted orbital stack in compact space.

Parameters

Name	Type	Default	Description
coefficients	mx.array		Orbital stack with shape (n, *grid.shape).

Returns

• mx.array — Fresh orthonormal full-grid coefficients.

Raises

• ValueError — If the stack is empty, too wide, or rank deficient.

overlap_matrix

def overlap_matrix(coefficients: mx.array) -> mx.array

Return the orbital overlap matrix in compact coefficient space.

Parameters

Name	Type	Default	Description
coefficients	mx.array		Orbital stack with shape (n, *grid.shape).

Returns

• mx.array — Complex Hermitian overlap matrix with shape (n, n).

project

def project(coefficients: mx.array) -> mx.array

Zero coefficients outside the admitted kinetic cutoff.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.

Returns

• mx.array — A fresh full-grid value with every inactive entry exactly zero.

real_norms

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

Return real-space integral norms for one orbital or a stack.

Parameters

Name	Type	Default	Description
orbitals	mx.array		One real-space orbital grid or a stack of grids.

Returns

• mx.array — One sqrt(integral |psi|^2) value per orbital.

to_coefficients

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

Transform real-space orbitals into the admitted coefficient basis.

Parameters

Name	Type	Default	Description
orbitals	mx.array		One real-space orbital grid or a stack of grids.

Returns

• mx.array — Fresh cutoff-projected reciprocal coefficients.

to_dict

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

Return JSON-safe basis metadata.

Returns

• dict[str, object] — Cutoff, k-point, FFT shape, active count, and normalization metadata.

to_real

def to_real(coefficients: mx.array) -> mx.array

Transform reciprocal coefficients to normalized real-space orbitals.

Parameters

Name	Type	Default	Description
coefficients	mx.array		One coefficient grid or a stack of grids.

Returns

• mx.array — Real-space complex orbital values with matching leading shape.