Using ANN Potentials with the PyTorch Implementation

Note

Inference as described here makes use of PyTorch. Make sure to install core torch support as described in Installation & Set-up. The descriptor-backed workflows shown here also rely on the matching torch-scatter and torch-cluster wheels.

Note

Alternative: For inference using ænet’s Fortran-based tools, see Using ANN Potentials with ænet’s Fortran Binaries.

Warning

The PyTorch implementation is primarily intended for training neural network potentials. For production inference workflows, the Fortran-based implementation, described in Using ANN Potentials with ænet’s Fortran Binaries, can provide 20–50x faster performance depending on the hardware and system size. Only for very large model sizes and when using GPUs for inference, the PyTorch implementation can be competitive in terms of speed.

Example notebooks

For a longer workflow-oriented walkthrough, including the maintained TiO2 saved-model example, batched file-backed inference, and optional GPU execution, see example-06-torch-inference.ipynb.

Additional notebooks are available in the notebooks directory within the repository.

Loading a Trained Model

Use compact page-level examples for the core API, and prefer the notebook linked above for larger file collections, performance tuning, and GPU-backed workflows.

from aenet.mlip import PredictionConfig
from aenet.torch_training import TorchANNPotential

pot = TorchANNPotential.from_file("outputs/trained_model.pt")

results = pot.predict(
    ["structure.xsf"],
    eval_forces=True,
    config=PredictionConfig(
        print_atomic_energies=True,
        timing=True,
    ),
)

print(f"Total energy: {results.total_energy[0]:.4f} eV")
print(f"Cohesive energy: {results.cohesive_energy[0]:.4f} eV")
print(results.forces[0].shape)
print(results.atom_energies[0].shape)

The unified API returns a PredictOut object containing:

  • total_energy: Total energies per structure

  • cohesive_energy: Cohesive energies per structure

  • forces: Atomic forces (if eval_forces=True)

  • atom_energies: Per-atom energies (if config.print_atomic_energies=True)

  • coords: Atomic coordinates

  • atom_types: Atomic species

  • timing: Per-structure timing breakdown (if config.timing=True)

When timing=True, the per-structure timing data is available through results.timing.

This API is identical to the Fortran-based inference API, enabling seamless interoperability between training in PyTorch and inference in Fortran.

Dataset-Backed Inference

The PyTorch backend also provides a torch-only optimization path for dataset-backed inference. The notebook linked above remains the maintained home for the longer file-backed TiO2 workflow.

from aenet.geometry import AtomicStructure
from aenet.mlip import PredictionConfig
from aenet.torch_training.dataset import CachedStructureDataset

structures = [AtomicStructure.from_file("structure.xsf")]

ds = CachedStructureDataset(
    structures=structures,
    descriptor=pot.descriptor,
    show_progress=False,
)

results = pot.predict_dataset(
    ds,
    config=PredictionConfig(batch_size=32)
)

This is useful when the dataset already stores precomputed features. For example, CachedStructureDataset can reuse cached feature tensors and avoid featurizing the structures again during energy-only inference. When eval_forces=True is requested, predict_dataset() falls back to materialized structures and calls the ordinary force-prediction path, because force evaluation requires coordinate derivatives rather than only cached energy features.

Notes:

  • predict_dataset() is available only for TorchANNPotential.

  • eval_forces=True requires a dataset that can expose raw structures via get_structure(), structures, or a supported Subset wrapper.

  • The return type is still PredictOut.