Example extension package for nequip

This repository serves as an example and template for writing extension packages for the nequip framework for (equivariant) machine learning on atomic systems. Extension packages to nequip can implement new models, new model components and wrappers that can be mixed with or added to existing models,

For an interactive tutorial please see the last section the Allegro tutorial Colab.

Contents

Model customization

The nequip framework, for which Allegro is an extension package, makes it easy to modify or extend models while preserving compatability with all the nequip-* tools and LAMMPS pair styles.

To illustrate this, we demonstrate a simple modification to an Allegro model that adds random noise to the predicted pairwise energies before they are summed into per-atom and total energies. To do this, we define two kinds of nequip extensions:

1. A module, which includes code we want to put in our model. Modules are just PyTorch torch.nn.Modules that include some extra information for nequip about the irreps of the data they expect to input and output. See nequip_example_extension.nn.AddNoiseModule.
2. A model builder, a function that takes the config (and optionally a model returned by previous model builders) and returns a new version of the model. The model builder is responsible in this case for adding our new module to an existing Allegro model. See nequip_example_extension.model.AddNoiseToPairEnergies.

Extensions like this can be used from the existing YAML config files you are familar with; see configs/minimal_allegro_with_custom_module.yaml.

Custom loss functions and metrics

See loss_terms.py and metrics.py. Please not that custom loss functions are NOT necessary if you want to use a conventional loss on a new prediction / field / property. Just like the existing loss functions/metrics can be used transparently on both energy and forces according to the specification in nequip's full.yaml, they can extend to other fields by giving the correct name. All that is required is for that field to be present both in the model's predictions and the dataset. Rather, custom loss functions / metrics are for different functional forms of the loss/metric, regardless of what field they are applied to, or for custom metrics/losses that depend on multiple fields and ignore the key passed through from the YAML config.

FAQ

What about LAMMPS?

If your model modules / additions are fully TorchScript compatible, then your modified/extended models can be used in the corresponding LAMMPS plugin for MD (pair_nequip for the general case, and pair_allegro for Allegro models.)

Extension package vs pull request?

We welcome pull requests and contributions to the nequip and allegro repositories and generally smaller new features of general interest should be pull requests on those repositories. If you aren't sure what is most appropriate for your idea, please always feel free to reach out via GitHub Discussions or email at albym[at]seas[dot]harvard[dot]edu. See also CONTRIBUTING.md.