Inference

Extension packages can perform inference using NequIP framework models through existing integrations or by extending them with custom functionality.

Integrations

The NequIP framework provides pre-built integrations that extension packages can leverage:

ASE Integration

The NequIPCalculator provides ASE calculator functionality. Extension packages can:

  • Subclass NequIPCalculator to add custom behavior

  • Override the save_extra_outputs() method to process additional model outputs

  • Extend the supported properties by modifying implemented_properties

See the ASE integration documentation for usage details.

LAMMPS ML-IAP Integration

The NequIPLAMMPSMLIAPWrapper is the LAMMPS ML-IAP integration module. Extension packages that implement basic interatomic potential models in a compliant manner can be used with this wrapper.

Custom Integrations

For other simulation packages or custom workflows, extension packages can build upon the underlying model loading and inference utilities. The NequIP framework provides two main model loading functions for inference:

nequip.model.saved_models.load_utils.load_saved_model(input_path, compile_mode: str = 'eager', model_key: str = 'sole_model', return_data_dict: bool = False)[source]

Load a saved model from checkpoint, package, or nequip.net.

This function can load models from:

  • Checkpoint files (.ckpt): saved during training runs

  • Package files (.nequip.zip): created with nequip-package

  • nequip.net models: using model ID format nequip.net:group-name/model-name:version from nequip.net

Parameters:

  • input_path – path to the model checkpoint or package file, or nequip.net model ID (format: nequip.net:group-name/model-name:version)

  • compile_mode (str) – compile mode for the model (default: "eager")

  • model_key (str) – key to select the model from ModuleDict (default: "sole_model")

  • return_data_dict (bool) – if True, also return the data dict for compilation (default: False)

Returns:

the loaded model, or (model, data) tuple if return_data_dict=True

Return type:

torch.nn.Module or tuple

Use this function for loading trained models from:

nequip.model.inference_models.load_compiled_model(compile_path: str, device: str | device, input_keys: List[str] | None = None, output_keys: List[str] | None = None) → Tuple[Module, dict][source]

Load a compiled model from either TorchScript or AOTInductor format.

This function can load compiled models created with nequip-compile:

  • TorchScript models (.nequip.pth): legacy compiled format

  • AOT Inductor models (.nequip.pt2): modern compiled format with better performance

Parameters:

  • compile_path – path to compiled model file (.nequip.pth or .nequip.pt2)

  • device – the device to use

  • input_keys – optional input field names for AOTInductor models (for .nequip.pt2)

  • output_keys – optional output field names for AOTInductor models (for .nequip.pt2)

Returns:

(model, metadata) with model prepared for inference

Return type:

tuple

Use this function for loading compiled models created with nequip-compile, including:

  • TorchScript models (.nequip.ts)

  • AOT Inductor compiled models (.nequip.pt2)

AOT Inductor Compilation Targets

The NequIP framework provides predefined compilation targets for common use cases:

  • ase: Optimized for ASE calculator integration with standard ASE outputs

  • batch: Supports batched inference with additional batch dimension handling

Extension packages can register custom AOT Inductor input/output specifications using the following function:

nequip.scripts._compile_utils.register_compile_targets(target_dict: Dict[str, List[str] | Callable]) → None[source]

Register compile targets for AOT compilation.

The intended clients of this function are NequIP extension packages to register their custom compilation targets.

Parameters:

target_dict – dict containing keys input, output, batch_map_settings, data_settings

Contact

For questions about developing inference capabilities in extension packages: