mattertune.backbones.nequip_foundation.util

Functions

nequip_model_package(ckpt_path, ...)

A suggested NequIP workflow is: 1.
mattertune.backbones.nequip_foundation.util.nequip_model_package(ckpt_path, example_atoms, output_path)[source]

A suggested NequIP workflow is: 1. Train a NequIP model and save the checkpoint (.ckpt) file. 2. Test the trained model using the checkpoint file if needed. 3. Package the trained model into a NequIP package file (.nequip.zip) 4. Compile the NequIP package file into a compiled model file (.nequip.pth/pt2)

This function packages a trained NequIP model from a checkpoint file into a NequIP package file. The implementation of this function is based on the nequip-package API in the NequIP repository, and the .nequip.zip packages produced by this function are fully compatible with subsequent nequip-compile api in nequip repo.

Some references: 1. nequip workflow: https://nequip.readthedocs.io/en/latest/guide/getting-started/workflow.html 2. example usage: TO-BE-ADDED

Parameters:

  • ckpt_path (str | Path)

  • example_atoms (Atoms)

  • output_path (str | Path)