AMARO: All Heavy-Atom Transferable Neural Network Potential

This repository hosts the trained checkpoint for AMARO v1.0.

All-atom molecular simulations provide detailed insight into macromolecular phenomena, but their computational cost limits the exploration of complex biological processes.

AMARO, Advanced Machine-learning Atomic Representation Omni-force-field, is a neural network potential that combines the O(3)-equivariant message-passing architecture TensorNet with a coarse-graining map that excludes hydrogen atoms.

AMARO demonstrates that coarse-grained neural network potentials can be trained without explicit prior-energy terms while retaining stable protein dynamics, scalability, and transferability.

Representation

AMARO uses a no-hydrogen, no-water mapping. Each retained bead corresponds to one protein heavy atom.

For each retained heavy atom, the reference force is constructed as the sum of:

1. the force acting on the heavy atom; and
2. the forces acting on hydrogen atoms constrained to that heavy atom.

This mapping reduces the number of degrees of freedom while retaining a detailed representation of protein geometry.

Important: z Values Are AMARO Bead Types

The checkpoint does not interpret z as conventional atomic numbers alone.

Each bead is assigned one of 12 learned embedding types, determined by:

• the identity of the heavy atom; and
• the number of hydrogen atoms aggregated to it.

This representation distinguishes chemically different environments and electronic hybridizations that would otherwise share the same element.

Inputs must therefore be prepared using the same AMARO mapping and bead-type assignment used during training. Passing ordinary atomic numbers without applying the AMARO remapping will produce invalid predictions.

Model Inputs and Outputs

Inputs

• z: one-dimensional tensor containing AMARO bead-type, with shape (N,)
• pos: Cartesian coordinates in Å, with shape (N, 3)
• batch: optional system-assignment tensor with shape (N,)
• box: optional periodic box vectors, where supported by the installed TorchMD-Net version

Outputs

• energy: learned effective potential for each system
• forces: negative gradient of the effective potential with respect to bead positions

The model uses units of:

kcal/mol/Å for forces and Å for coordinates 

Because the model was trained using force labels without energy labels, its absolute energy reference is not physically calibrated.

Citation

Please cite the following publication when using this checkpoint:

@article{mirarchi2024amaro,
title = {AMARO: All Heavy-Atom Transferable Neural Network Potentials of Protein Thermodynamics},
author = {Mirarchi, Antonio and Pel{\'a}ez, Ra{\'u}l P. and Simeon, Guillem and De Fabritiis, Gianni},
journal = {Journal of Chemical Theory and Computation},
volume = {20},
number = {22},
pages = {9871--9878},
year = {2024},
publisher = {ACS Publications}
}