Linear Model Training

Functions

training	Computes optimal linear parameters for a PIP model using energy data.
training_w_gradients	Simple training function for PIP models with Forces.

molpipx.training_linear_model.training(model_pip: Callable, X_tr: jaxtyping.Float.(jaxtyping.Array, '...'), y_tr: jaxtyping.Float.(jaxtyping.Array, '...')) -> jaxtyping.Float.(jaxtyping.Array, '...')

Computes optimal linear parameters for a PIP model using energy data.

Solves the linear system of equations PIP * theta = Energy using least squares to find the optimal coefficients (theta) for the polynomial expansion.

Warning

Geometries must be provided in Bohr Units.

Parameters:

• model_pip (Callable) – A Flax module instance initialized to compute PIP vectors.

• X_tr (Array) – Training geometries with shape (Batch, N_atoms, 3).

• y_tr (Array) – Training energies with shape (Batch, 1).

Returns:

The optimal linear parameters (theta) for the model.

Return type:

Array

molpipx.training_linear_model.training_w_gradients(model_pip: Callable, X_tr: jaxtyping.Float.(jaxtyping.Array, '...'), F_tr: jaxtyping.Float.(jaxtyping.Array, '...'), y_tr: jaxtyping.Float.(jaxtyping.Array, '...')) -> jaxtyping.Float.(jaxtyping.Array, '...')

Simple training function for PIP models with Forces. Warning: Geometries must be in Bohr Units and Forces in Ha/Bohr Units

Parameters:

• model_pip (Callable) – A Flax module instance initialized to compute PIP vectors.

• X_tr (Array) – Training geometries with shape (Batch, N_atoms, 3).

• F_tr (Array) – Training forces with shape (Batch, N_atoms, 3).

• y_tr (Array) – Training energies with shape (Batch, 1).

Returns:

The optimal linear parameters (theta) for the model.

Return type:

Array