from_compiled_model¶

torch_sim.models.nequip_framework.from_compiled_model(compile_path, device='cpu')[source]¶

Load a compiled NequIP model from a file.

Loads a compiled NequIP model from a file and extracts the necessary metadata for using it in TorchSim. The model must have been compiled using nequip-compile.

Parameters:

compile_path (str) – Path to the compiled model file. The file should have been created using nequip-compile.
device (str | torch.device) – Device to load the model on. Can be either a string like ‘cpu’ or ‘cuda’, or a torch.device object. Defaults to ‘cpu’.

Returns:

A tuple containing:

The loaded NequIP model as a torch.nn.Module
A tuple with:
- r_max (float): Cutoff radius used by the model
- type_names (list[str]): List of chemical symbols supported by the model

Return type:

tuple[Module, tuple[float, list[str]]]

Example

>>> model, (r_max, type_names) = from_compiled_model("model.pth", device="cuda")
>>> print(f"Model cutoff: {r_max:.2f}")
>>> print(f"Supported elements: {type_names}")

References

For model compilation please refer to the NequIP documentation:: https://nequip.readthedocs.io/en/latest/guide/getting-started/workflow.html#compilation