torch_sim.models.nequip_framework

Wrapper for NequIP-Allegro models in TorchSim.

This module provides a TorchSim wrapper of the NequIP-Allegro models for computing energies, forces, and stresses for atomistic systems. It integrates the NequIP-Allegro models with TorchSim’s simulation framework, handling batched computations for multiple systems simultaneously.

The implementation supports various features including:

  • Computing energies, forces, and stresses

  • Batched calculations for multiple systems

References

Functions

from_compiled_model

Load a compiled NequIP model from a file.

Classes

ChemicalSpeciesToAtomTypeMapper

Maps atomic numbers to model-specific atom type indices.

NequIPFrameworkModel

NequIP model for energy, force and stress calculations.