mace_fitting

Contents

mace_fitting#

autoplex.fitting.common.utils.mace_fitting(db_dir, path_to_default_hyperparameters=MLIP_RSS_DEFAULTS_FILE_PATH, device='cuda', ref_energy_name='REF_energy', ref_force_name='REF_forces', ref_virial_name='REF_virial', use_defaults=True, fit_kwargs=None)[source]#

Perform the MACE potential fitting.

This function sets up and executes a python script to perform MACE fitting using specified parameters and input data located in the provided directory. It handles the input/output of atomic configurations, sets up the NequIP model, and calculates training and testing errors after fitting.

Parameters:

  • db_dir (Path) – directory containing the training and testing data files.

  • path_to_default_hyperparameters (str or Path.) – Path to mlip-rss-defaults.json.

  • device (str) – specify device to use cuda or cpu.

  • ref_energy_name (str, optional) – Reference energy name.

  • ref_force_name (str, optional) – Reference force name.

  • ref_virial_name (str, optional) – Reference virial name.

  • fit_kwargs (dict.) – optional dictionary with parameters for mace fitting with keys same as mlip-rss-defaults.json.

Keyword Arguments:

  • model (str) – type of model to be trained

  • config_type_weights (str) – weights of config types

  • hidden_irreps (str) – control the model size

  • r_max (float) – cutoff radius controls the locality of the model

  • batch_size (int) – batch size (note that batch size cannot be larger than the size of training datasets)

  • start_swa (str) – if the keyword –swa is enabled, the energy weight of the loss is increased for the last ~20% of the training epochs (from –start_swa epochs)

  • correlation (int) – correlation order corresponds to the order that MACE induces at each layer

  • loss (str) – loss functions

  • default_dtype (str) – type of float to use, e.g. float32 and float64

Returns:

A dictionary containing train_error, test_error, and the path to the fitted MLIP.

Return type:

dict[str, float]

Contents