atom_ref

Atomic reference module from [M-OFDFT].

class AtomRef(t_z: Tensor | ndarray, t_global: float, g_z: Tensor | ndarray)

Atomic reference module from [M-OFDFT].

__init__(t_z: Tensor | ndarray, t_global: float, g_z: Tensor | ndarray) → None

Initialize the atomic reference module. Alternatively use from_dataset_statistics().

Parameters:

• t_z – Kinetic energy per atomic species. Shape (n_atom_types,).

• t_global – Global bias.

• g_z – Mean gradient vector, concatenated over atomic species. Shape (basis,).

forward(atom_ind: Tensor, coeffs: Tensor, basis_function_ind: Tensor, atom_batch: Tensor | None, coeffs_batch: Tensor | None = None) → Tensor

Calculates the energy for one molecule, or a batch of molecules, according to a linear fit, e.g. computed by DatasetStatistics.

Parameters:

• atom_ind – Atom types in the molecule (indexing the atomic numbers present in the basis). Shape (n_atoms,).

• coeffs – Coefficients of the molecule. Shape (n_basis,).

• basis_function_ind – Indices of the basis functions in the molecule. See mldft.ml.data.components.of_data.OFData.basis_function_ind. Shape (n_basis,).

• atom_batch – Batch index for the atoms. Shape (n_atoms,), or None if not batched.

• coeffs_batch – Batch index for the coefficients. Shape (n_basis,), or None if not batched.

classmethod from_dataset_statistics(dataset_statistics, basis_info: BasisInfo = None, weigher_key: str = 'has_energy_label', scalar_only=False) → AtomRef

Initialize using a DatasetStatistics object.

Parameters:

• dataset_statistics – Dataset statistics object.

• basis_info – Basis information object. Required if scalar_only is True, to infer which features are scalars.

• weigher_key (str) – Selects the sample weigher for the dataset statistics.

• scalar_only (bool) – If True, only the scalar features are used. Defaults to False.

sample_forward(sample: OFData) → Tensor

Calculate the expected energy for a sample (or batch) with a linear fit, e.g. computed by DatasetStatistics.

Parameters:

sample – OFData object containing the molecule.

class SimpleQuadraticAtomRef(ground_state_coeff_mean: Tensor, factor: float)

__init__(ground_state_coeff_mean: Tensor, factor: float)

Parameters:

• ground_state_coeff_mean – The mean of the ground state coefficients per atom type.

• factor – The factor to multiply the quadratic term with.

forward(coeffs: Tensor, basis_function_ind: Tensor, coeffs_batch: Tensor | None) → Tensor

Compute an isotropic quadratic energy around the mean ground state coeffs.

Parameters:

• coeffs – The coefficients of the molecule. Shape (n_basis,).

• basis_function_ind – Indices of the basis functions in the molecule. See mldft.ml.data.components.of_data.OFData.basis_function_ind. Shape (n_basis,).

• coeffs_batch – Batch index for the coefficients. Shape (n_basis,), or None if not batched.

sample_forward(sample: OFData) → Tensor

Compute the energy for a sample.

Parameters:

sample – The sample to compute the energy for.