Model I/O

dipm.utils.model_io.save_model(save_path: str | PathLike, model: ForceFieldPredictor) → None

Saves a force field model to a safetensors file.

Parameters:

  • save_path – The target path to the model file. Should have extension “.safetensors”.

  • model – The force field model to save. Must be passed as type ForceFieldPredictor.

dipm.utils.model_io.load_model(load_path: str | PathLike, model_type: type[ForceModel] | None = None, dtype: str | type[Any] | dtype | SupportsDType | Any | None = None, strict: bool = True, params_filters: Callable[[DatasetInfo, ForceModelConfig, dict[str, Array], type[ForceModel]], tuple[DatasetInfo, ForceModelConfig, dict[str, Array]]] | list[Callable[[DatasetInfo, ForceModelConfig, dict[str, Array], type[ForceModel]], tuple[DatasetInfo, ForceModelConfig, dict[str, Array]]]] | None = None) → ForceFieldPredictor

Loads a model from a safetensors file and returns it wrapped as a ForceFieldPredictor.

Parameters:

  • load_path – The path to the safetensors file to load.

  • model_type (optional) – The model class that corresponds to the saved model. If you are using a known model listed in dipm.models.KNOWN_MODELS, this argument can be optional and inferred from the saved metadata. Otherwise, it must be provided.

  • dtype (optional) – The dtype in computations. If not provided, it will be the same as the the parameters dtype.

  • strict (optional) – If True, raises an error if the parameters name and shape are mismatched.

  • params_filters (optional) – A function or a list of functions that takes the dataset_info, model_config, params_raw and model_type as input and returns the filtered dataset_info, model_config and params_raw. This can be used to modify the loaded parameters or metadata before creating the model.

Returns:

The loaded model wrapped as a ForceFieldPredictor object.

class dipm.utils.params_filter.ForceHeadParamsFilter

Removes the force head if exists.

class dipm.utils.params_filter.UnseenElementsParamsFilter(elements_to_drop: set[int])

Removes the parameters of unseen elements.