Anisotropic Model Training

Functions

_train_aniso_adam

Trains an anisotropic PIP model using the Adam optimizer.
molpipx.trainig_anisotropic_model._train_aniso_adam(data: Tuple, atoms: List, f_mono: Callable, f_poly: Callable, optimizer_info: jaxtyping.PyTree) → Tuple[jaxtyping.Float, jaxtyping.PyTree][source]

Trains an anisotropic PIP model using the Adam optimizer.

Parameters:

  • data (Tuple) – A tuple containing training and validation data: ((X_tr, y_tr), (X_val, y_val)).

  • atoms (List) – A list of atom types (e.g., ['H', 'C', 'H']) used to generate the mask.

  • f_mono (Callable) – Function that returns the monomials.

  • f_poly (Callable) – Function that returns the polynomials.

  • optimizer_info (PyTree) – A dictionary or PyTree containing optimizer settings: * 'tol': Convergence tolerance. * 'Maxiters': Maximum number of epochs. * 'learning_rate': Learning rate for Adam.

Returns:

A tuple containing:

  • The final optimized parameters (l_opt[-1]).

  • The history of parameters over all epochs (l_opt).

Return type:

Tuple[Float, PyTree]