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:
the force acting on the heavy atom; and
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:
bibtex
1@article{mirarchi2024amaro,
2title = {AMARO: All Heavy-Atom Transferable Neural Network Potentials of Protein Thermodynamics},
3author = {Mirarchi, Antonio and Pel{\'a}ez, Ra{\'u}l P. and Simeon, Guillem and De Fabritiis, Gianni},
4journal = {Journal of Chemical Theory and Computation},
5volume = {20},
6number = {22},
7pages = {9871--9878},
8year = {2024},
9publisher = {ACS Publications}
10}
