FeNNol 0.1.dev1

FeNNol: Force-field-enhanced Neural Network optimized library

FeNNol: Force-field-enhanced Neural Networks optimized library

FeNNol is a library for building, training and running neural network potentials for molecular simulations. It is based on the JAX library and is designed to be fast and flexible.

FeNNol's documentation is available here and the article describing the library at https://doi.org/10.1063/5.0217688

Active Learning tutorial in this Colab notebook

Table of Contents

• Pre-trained models
• Installation
  • From PyPI
  • Latest version from Github repo
  • Optional dependencies
• Examples
• Citation
• License

Pre-trained models

Pre-trained models can be downloaded from the Pre-trained Models Collection. Available models include:

• FeNNix-Bio1
• ANI1x, ANI1ccx, ANI2x
• MACE-OFF23
• MACE-MP

After downloading the model (for example fennix-bio1S.fnx), you can load it and start predicting energies and forces with:

import fennol
import numpy as np

model = fennol.load("fennix-bio1S.fnx")
species = np.array([8, 1, 1])
coordinates = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])

energy, forces, output_dict = model.energy_and_forces(species=species, coordinates=coordinates, unit="ev")
print("Energy:", energy,"ev")
print("Forces:", forces,"ev/Å")

Alternatively, you can use the provided ASE calculator and run a short molecular dynamics simulation:

from fennol.ase import FENNIXCalculator
from ase import Atoms
from ase.md.verlet import VelocityVerlet    
import ase.units as units

atoms = Atoms(symbols=["O", "H", "H"], positions=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
calculator = FENNIXCalculator(model="fennix-bio1S.fnx", gpu_preprocessing=True)
atoms.calc = calculator
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
print("Energy:", energy,"ev")
print("Forces:", forces,"ev/Å")

dyn = VelocityVerlet(atoms, timestep=1.0*units.fs) 
print("Starting MD simulation. ")
for step in range(10):
    dyn.run(10)
    print(f"Step {10*(step+1)}, Potential energy: {atoms.get_potential_energy()} ev")

For better performance, we recommend using fennol_md, FeNNol's native MD engine. (see the examples/md directory for instructions and example input files)

Installation

From PyPI

# CPU version
pip install fennol

# GPU version
pip install "fennol[cuda]"

Latest version from Github repo

You can start with a fresh environment, for example using venv:

python -m venv fennol
source fennol/bin/activate

The first step is to install jax (see details at: https://jax.readthedocs.io/en/latest/installation.html). For example, to install the latest version using pip:

# CPU version
pip install -U jax

# GPU version
pip install -U "jax[cuda12]"

Then, you can clone the repo and install FeNNol using pip:

git clone https://github.com/FeNNol-tools/FeNNol.git
cd FeNNol
pip install .

Optional dependencies

• Some modules require e3nn-jax (https://github.com/e3nn/e3nn-jax) which can be installed with:

pip install --upgrade e3nn-jax

• The provided training script requires pytorch (at least the cpu version) for dataloaders:

pip install torch --index-url https://download.pytorch.org/whl/cpu

• For the Deep-HP interface, cffi and pycuda are required:

pip install cffi pycuda

Examples

To learn how to train a FeNNol model, you can check the examples in the examples/training directory. The README.md file in that directory contains instructions on how to train a model on the aspirin revMD17 dataset.

To learn how to run molecular dynamics simulations with FeNNol models, you can check the examples in the examples/md directory. The README.md file in that directory contains instructions on how to run simulations with the provided ANI-2x model.

Citation

Please cite this paper if you use the library.

T. Plé, O. Adjoua, L. Lagardère and J-P. Piquemal. FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials. J. Chem. Phys. 161, 042502 (2024)

@article{ple2024fennol,
    author = {Plé, Thomas and Adjoua, Olivier and Lagardère, Louis and Piquemal, Jean-Philip},
    title = {FeNNol: An efficient and flexible library for building force-field-enhanced neural network potentials},
    journal = {The Journal of Chemical Physics},
    volume = {161},
    number = {4},
    pages = {042502},
    year = {2024},
    month = {07},
    doi = {10.1063/5.0217688},
    url = {https://doi.org/10.1063/5.0217688},
}

License

This project is licensed under the terms of the GNU LGPLv3 license. See LICENSE for additional details.