nequip.nn

class nequip.nn.pair_potential.ZBL(type_names: List[str], chemical_species: List[str], units: str, polynomial_cutoff_p: float = 6.0, per_atom_energy_field: str = 'atomic_energy', irreps_in=None)[source]

ZBL pair potential energy term.

Useful as a prior for core repulsion to mitigate molecular dynamics failure modes associated with atoms getting too close.

Parameters:

  • type_names (List[str]) – list of type names known by the model, [atom1, atom2, atom3]

  • chemical_species (List[str]) – list of chemical symbols, e.g. [C, H, O]

  • units (str) – LAMMPS units that the data is in; metal and real are presently supported – raise a GitHub issue if more is desired

  • polynomial_cutoff_p (float) – exponent used for the polynomial cutoff (default 6)

class nequip.nn.PerTypeScaleShift(type_names: List[str], field: str, out_field: str | None = None, scales: float | Dict[str, float] | None = None, shifts: float | Dict[str, float] | None = None, scales_trainable: bool = False, shifts_trainable: bool = False, irreps_in={})[source]

Scale and/or shift a predicted per-atom property based on (learnable) per-species/type parameters.

Note that scaling/shifting is always done casting into the global dtype (float64), even if model_dtype is a lower precision.

If a single scalar is provided for scales/shifts, a shortcut implementation is used. Otherwise, a more expensive implementation that assigns separate scales/shifts to each atom type is used.

If scales/shifts are trainable, the more expensive implementation that assigns separate scales/shifts to each atom type is used, even if a single scalar was provided for the initialization.

classmethod modify_PerTypeScaleShift(model, scales: float | Dict[str, float] | None = None, shifts: float | Dict[str, float] | None = None, scales_trainable: bool = False, shifts_trainable: bool = False)[source]

Modify per-type scales and shifts of a model.

The new scales and shifts should be provided as dicts. The keys must correspond to the type_names registered in the model being modified, and may not include all the possible type_names of the original model. For example, if one uses a pretrained model with 50 atom types, and seeks to only modify 3 per-atom shifts to be consistent with a fine-tuning dataset’s DFT settings, one could use

shifts:
  C: 1.23
  H: 0.12
  O: 2.13

In this case, the per-type atomic energy shifts of the original model will be used for every other atom type, except for atom types with the new shifts specified.

For more details on fine-tuning, see https://nequip.readthedocs.io/en/latest/guide/training-techniques/fine_tuning.html

Parameters:

  • scales – the new per-type atomic energy scales

  • shifts – the new per-type atomic energy shifts (e.g. isolated atom energies of a dataset used for fine-tuning)

  • scales_trainable (bool) – whether the new scales are trainable

  • shifts_trainable (bool) – whether the new shifts are trainable