ZBL short-range potential

class metatrain.utils.additive.zbl.ZBL(hypers: Dict, dataset_info: DatasetInfo)[source]

Bases: ModelInterface

A simple model for short-range repulsive interactions.

The implementation here is equivalent to its LAMMPS counterpart, where we set the inner cutoff to 0 and the outer cutoff to the sum of the covalent radii of the two atoms as tabulated in ASE. Covalent radii that are not available in ASE are set to 0.2 Å (and a warning is issued).

Parameters:

  • hypers (Dict) – A dictionary of model hyperparameters. This parameter is ignored and is only present to be consistent with the general model API.

  • dataset_info (DatasetInfo) – An object containing information about the dataset, including target quantities and atomic types.

restart(dataset_info: DatasetInfo, model_hypers: Mapping[str, Any] | None = None) → ZBL[source]

Update a model to restart training, potentially with different dataset and/or targets.

This function is called whenever training restarts, with the same or a different dataset. It enables transfer learning (changing the targets), and fine-tuning (same targets, different datasets)

Parameters:

  • dataset_info (DatasetInfo) – Information about the new dataset, including the targets that will be used for training.

  • model_hypers (Mapping[str, Any] | None) – The new hyperparameters for the model.

Returns:

The updated model, or a new instance of the model, that is able to handle the new dataset.

Return type:

ZBL

remove_output(target_name: str) → None[source]

Remove a previously registered output target.

Parameters:

target_name (str) – Name of the target to remove.

Return type:

None

supported_outputs() → Dict[str, ModelOutput][source]

Get the outputs currently supported by this model.

This will likely be the same outputs that are set as this model capabilities in ModelInterface.export().

Returns:

A dictionary of the supported outputs by this model.

Return type:

Dict[str, ModelOutput]

forward(systems: List[System], outputs: Dict[str, ModelOutput], selected_atoms: Labels | None = None) → Dict[str, TensorMap][source]

Compute the energies of a system solely based on a ZBL repulsive potential.

Parameters:

  • systems (List[System]) – List of systems to calculate the ZBL energy.

  • outputs (Dict[str, ModelOutput]) – Dictionary containing the model outputs.

  • selected_atoms (Labels | None) – Optional selection of atoms for which to compute the predictions.

Returns:

A dictionary with the computed predictions for each system.

Raises:

ValueError – If the outputs contain unsupported keys.

Return type:

Dict[str, TensorMap]

get_pairwise_zbl(zi: Tensor, zj: Tensor, rij: Tensor) → Tensor[source]

Ziegler-Biersack-Littmark (ZBL) potential.

Inputs are the atomic numbers (zi, zj) of the two atoms of interest and their distance rij.

Parameters:

  • zi (Tensor) – Atomic number of atom i.

  • zj (Tensor) – Atomic number of atom j.

  • rij (Tensor) – Distance between atom i and atom j.

Returns:

The ZBL potential energy between atom i and atom j.

Return type:

Tensor

requested_neighbor_lists() → List[NeighborListOptions][source]
Return type:

List[NeighborListOptions]

static is_valid_target(target_name: str, target_info: TargetInfo) → bool[source]

Finds if a TargetInfo object is compatible with the ZBL model.

Parameters:

  • target_name (str) – The name of the target to be checked.

  • target_info (TargetInfo) – The TargetInfo object to be checked.

Returns:

True if the target is compatible with the ZBL model, False otherwise.

Return type:

bool

get_checkpoint(best_model_state_dict: dict[str, Any] | None = None) → dict[str, Any][source]

Get the checkpoint of the model. This should contain all the information needed by load_checkpoint to recreate the same model instance.

Parameters:

best_model_state_dict (dict[str, Any] | None) – Optional state dictionary of the best model (if available). If provided, it should likely be included in the checkpoint to allow for resuming training or evaluation from the best model state.

Returns:

The model’s checkpoint.

Return type:

dict[str, Any]

classmethod load_checkpoint(checkpoint: dict[str, Any], context: Literal['restart', 'finetune', 'export']) → ZBL[source]

Create a model from a checkpoint (i.e. state dictionary).

Parameters:

  • checkpoint (dict[str, Any]) – Checkpoint’s state dictionary.

  • context (Literal['restart', 'finetune', 'export']) – Context in which to load the model. Not used.

Returns:

An instance of the model.

Return type:

ZBL

classmethod upgrade_checkpoint(checkpoint: dict, version: int | None = None) → dict[source]

Upgrade the checkpoint to the current version of the model.

Parameters:

  • checkpoint (dict) – Checkpoint’s state dictionary.

  • version (int | None) – Version to which the checkpoint should be upgraded. If None, the checkpoint will be upgraded to the current version of the model.

Returns:

The upgraded checkpoint.

Return type:

dict

export(metadata: ModelMetadata | None = None) → AtomisticModel[source]

Turn this model into an metatomic.torch.AtomisticModel, containing the model itself, its capabilities and its metadata.

Parameters:

metadata (ModelMetadata | None) – Additional metadata to add to the model, as specified by the user.

Returns:

An instance of metatomic.torch.AtomisticModel.

Return type:

AtomisticModel