Neural Force Field

The Neural Force Field (NFF) code is an API based on SchNet [1-4], DimeNet [5], PaiNN [6-7], DANN [8], CHGNet [9], and MACE [10,11]. It provides an interface to train and evaluate neural networks (NNs) for force fields. It can also be used as a property predictor that uses both 3D geometries and 2D graph information [12]. NFF also allows the usage of NN ensembles to perform uncertainty quantification and adversarial sampling of geometries. The complete tutorials on how to perform such analysis is available at the Atomistic Adversarial Attacks repository, and the theory behind this differentiable sampling strategy is available at our paper [13].

This code repository is developed in the Learning Matter Lab (led by Prof. Rafael Gomez-Bombarelli) at MIT.

Conda environment

We highly recommend creating a conda environment to run the code. To do that, use the following command to create the nff conda environment:

conda upgrade conda
conda env create -f environment.yml

To ensure that the nff environment is accessible through Jupyter, add the the nff display name:

python -m ipykernel install --user --name nff --display-name "Python [conda env:nff"]

Installation

If you would like to use NFF as a package (no development), you can do so by running the following command in the root directory of the repository:

pip install .

Use the -e flag if you want to develop NFF:

pip install -e '.[dev]'

Usage

For those just getting started, we recommend referring to the wiki or advanced users can just jump right in. A Jupyter notebook interface or a command-line interface can be used with NFF. There is also a high-throughput option for using NFF, where you use HTVS to run simulations or perform calculations with NFF.

A set of pre-trained models can be found in models. To take take an under-the-hood look at the architecture of all the models available with NFF, go to nff/nn/models, and see the underlying modules with supporting functions in nff/nn/modules.

Jupyter notebooks interface

Please refer to the tutorials to see how to set up a Jupyter notebook interface with NFF.

Command-line interface

Force field

The simplest way to use the nff package is to use the premade scripts (in the scripts) folder. As an example, to train a SchNet model with the default parameters using the example dataset (ethanol geometries) from the command line, run the command

nff_train.py train schnet tutorials/data/dataset.pth.tar $HOME/train_model --device cuda:0

This will use 60% of the dataset for training, 20% for validation and 20% for testing. The training will happen on the device cuda:0. Results of training, checkpoints and hyperparameters will be saved on the path $HOME/train_model.

Property predictor

NFF also contains modules that predict properties from 3D geometries of conformers. These include the SchNet model, expanded to include multiple conformers, as well as the ChemProp3D (CP3D) model, which also includes graph information. A series of scripts for these modules can be found in scripts/cp3d. An in-depth discussion of how to use these scripts can be found in scripts/cp3d/README.md.

HTVS interface

Please refer to the wiki to see how to use NFF through HTVS.

Adversarial Attacks

NFF allows the usage of NN ensembles to perform uncertainty quantification and adversarial sampling of geometries. The complete tutorials on how to perform such analysis is available at the Atomistic Adversarial Attacks repository, and the theory behind this differentiable sampling strategy is available at our paper [13].

References

• [1] K.T. Schütt. F. Arbabzadah. S. Chmiela, K.-R. Müller, A. Tkatchenko.
  Quantum-chemical insights from deep tensor neural networks. Nature Communications 8. 13890 (2017)
  10.1038/ncomms13890

• [2] K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller.
  SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in Neural Information Processing Systems 30, pp. 992-1002 (2017) link

• [3] K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller.
  SchNet - a deep learning architecture for molecules and materials. The Journal of Chemical Physics 148(24), 241722 (2018) 10.1063/1.5019779

• [4] K.T. Schütt, P. Kessel, M. Gastegger, K. Nicoli, A. Tkatchenko, K.-R. Müller. SchNetPack: A Deep Learning Toolbox For Atomistic Systems. J. Chem. Theory Comput. 15(1), 448-455 (2019). 10.1021/acs.jctc.8b00908

• [5] J. Klicpera, G. Janek, S. Günnemann. Directional message passing for molecular graphs. ICLR (2020). URL.

• [6] K. T. Schütt, O. T. Unke, M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. arXiv preprint, 2021. arXiv:2102.03150

• [7] S. Axelrod, E. Shakhnovich, R. Gómez-Bombarelli. Thermal half-lives of azobenzene derivatives: virtual screening based on intersystem crossing using a machine learning potential. arXiv preprint (2022). arXiv:2207.11592.

• [8] S. Axelrod, E. Shakhnovich, R. Gómez-Bombarelli. Excited state non-adiabatic dynamics of large photoswitchable molecules using a chemically transferable machine learning potential. Nat. Commun. 13, 3440 (2022). URL

• [9] B. Deng, P. Zhong, K. Jun, J. Riebesell, K. Han, C. J. Bartel, G. Ceder. CHGNet as a pretrained universal neural network potential for charge-informed atomstic modelling. Nat. Mach. Intell. 5, 1031 (2023). 10.1038/s42256-023-00716-3

• [10] I. Batatia, D. P. Kovács, G. N. C. Simm, C. Ortner, G. Csányi. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. Advances in Neural Information Processing Systems 35, pp. 11423–11436 (2022). link

• [11] I. Batatia et al. A foundation model for atomistic materials chemistry. arXiv preprint, 2023. 2401.00096

• [12] S. Axelrod and R. Gomez-Bombarelli. Molecular machine learning with conformer ensembles. arXiv preprint (2020). arXiv:2012.08452.

• [13] D. Schwalbe-Koda, A.R. Tan, and R. Gomez-Bombarelli. Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks. Nat. Commun. 12, 5104 (2021). URL.