nequip.integrations.torchsim

class nequip.integrations.torchsim.NequIPTorchSimCalc(model: Module, device: str | device = 'cpu', transforms: List[Callable] = [], atomic_numbers: Tensor | None = None, system_idx: Tensor | None = None)[source]

NequIP framework torch-sim calculator.

This torch-sim calculator is compatible with models from the NequIP framework, including NequIP and Allegro models.

The recommended way to use this calculator is with a compiled model, i.e. nequip-compile the model and load it into the calculator with NequIPTorchSimCalc.from_compiled_model(...).

Parameters:

  • model (torch.nn.Module) – a model in the NequIP framework

  • device (str or torch.device) – device for model to evaluate on, e.g. "cpu" or "cuda" (default: "cpu")

  • transforms (List[Callable]) – list of data transforms

  • atomic_numbers (torch.Tensor or None) – atomic numbers with shape [n_atoms]. If provided at initialization, cannot be provided again during forward pass

  • system_idx (torch.Tensor or None) – batch indices with shape [n_atoms] indicating which system each atom belongs to. If not provided with atomic_numbers, all atoms are assumed to be in the same system

property compute_forces: bool

Whether the model computes forces.

property compute_stress: bool

Whether the model computes stresses.

classmethod from_compiled_model(compile_path: str | Path, device: str | device = 'cpu', chemical_species_to_atom_type_map: Dict[str, str] | bool | None = None, neighborlist_backend: str = 'alchemiops', **kwargs)[source]

Build an integration calculator from a compiled model artifact.

Parameters:

  • compile_path – path to the compiled model artifact.

  • device – device where the model is loaded and evaluated.

  • chemical_species_to_atom_type_map – optional chemical species mapping override.

  • neighborlist_backend – neighbor list backend used by neighbor transforms.

  • **kwargs – forwarded to the integration class constructor.

setup_from_system_idx(atomic_numbers: Tensor, system_idx: Tensor) → None[source]

Set up internal state from atomic numbers and system indices.

Parameters:

  • atomic_numbers (torch.Tensor) – atomic numbers with shape [n_atoms].

  • system_idx (torch.Tensor) – system indices with shape [n_atoms].

forward(state: SimState) → dict[str, Tensor][source]

Compute energies, forces, and stresses.

Parameters:

state (SimState) – state object containing positions, cell, and system information.

Returns:

computed properties ("energy", "forces", "stress").

Return type:

dict[str, torch.Tensor]