Skip to content

aimnet2 #

Utilities for working with AIMNet2. Mainly taken from https://github.com/openmm/openmm-ml/pull/64 and https://github.com/SimonBoothroyd/befit/blob/main/befit/utils/aimnet2.py

See discussion at https://github.com/isayevlab/AIMNet2/issues/15 re ensemble models.

See https://github.com/isayevlab/aimnetcentral/blob/47969eb3e29e34824d82a648dd756669c875ecdb/scripts/compile/compile_off.yaml for available models. May compile the ensemble models in future.

Classes:

AIMNet2PotentialImplFactory #

Bases: MLPotentialImplFactory

This is the factory that creates AIMNet2PotentialImpl objects.

AIMNet2PotentialImpl #

AIMNet2PotentialImpl(name: str)

Bases: MLPotentialImpl

This is the MLPotentialImpl implementing the AIMNet2 potential.

Source code in presto/utils/aimnet2.py 
51
52
def __init__(self, name: str):
    self.name = name

_load_local_model #

_load_local_model(
    method: str, device: TorchDevice | None = None
) -> ScriptModule

Load an AIMNet2 model from local package resources.

Source code in presto/utils/aimnet2.py 
29
30
31
32
33
34
35
36
37
38
def _load_local_model(
    method: str, device: TorchDevice | None = None
) -> torch.jit.ScriptModule:
    """Load an AIMNet2 model from local package resources."""
    filename = f"{method}.jpt"
    with resources.path("presto.models", filename) as model_path:
        model: torch.jit.ScriptModule = torch.jit.load(  # type: ignore[no-untyped-call]
            str(model_path), map_location=device
        )
        return model

_register_aimnet2_potentials #

_register_aimnet2_potentials() -> None

Register the AIMNET2 potential implementation factory.

Source code in presto/utils/aimnet2.py 
124
125
126
127
def _register_aimnet2_potentials() -> None:
    """Register the AIMNET2 potential implementation factory."""
    for model in _AVAILABLE_MODELS:
        MLPotential.registerImplFactory(model, AIMNet2PotentialImplFactory())