PEANUT
Predicting potential Energies with Artificial NeUral neTworks — a computationally efficient approach for modeling potential energy surfaces.

---

Table of Contents

• Project description
  • Implemented models
  • Model evaluation
  • Documentation
• Installation
• Demo
• Train the model
  • Data preparation
  • Configure
  • Start training
• Results
• Tests and code structure
• Workplan
• References
• License

---

Project description

This repository builds a Deep Learning model for a Deep Learning lecture at TU Vienna. Corresponding to the assignment description: I believe my project description matches the project type ''beat the stars'' the most.

In molecular dynamics (MD) simulations, it is possible to replace classical empirical force fields (FF) with neural network potentials to predict potential energy surfaces. Such networks are usually called neural network potentials (NNP) or machine learning potentials (MLP) - to avoid confusion, I will use NNP and MLP = multilayer perceptron from now on. A combination of potential energy prediction and a so-called ''MD engine'' allows for running simulations of chemical systems such as e.g. a small solute in water, which in turn is a very useful tool to replace costly lab experiments. This method is an already well-known and established tool in computational chemistry.

Given that molecules can be modeled as graphs where atoms are considered as nodes and edges are considered between interacting atoms, the most recent and successful NNPs are designed as graph neural networks (e.g. DimeNet, MACE). These graph neural networks learn embeddings of atom types and use message passing along edges to model atom interactions based on interatomic distances and 3D arrangements. Other architecture types have also been proposed: SchNet uses a continuous-filter convolutional neural network architecture for modeling interactions. Other approaches use atom-centered symmetry functions based on distances and angles as features as fixed representation (ANI). This approach is of course less computationally expensive, however these architectures do not use representation learning.

Generally, such NNPs are trained on single point energies. Thus, their use in MD simulations is definitely an application outside of their training domain, making the task even more difficult. Also, since MD simulations are computationally highly expensive, the question of how complex a NNP’s architecture can and should be is crucial. Given that the use of NNPs in computational chemistry is still quite new, many methods that are already existing in classical MD simulations (using empirical FFs), need to be re-developed for the use of NNPs. For the development of such methods, the overall accuracy is not always the key point. Often, a functional yet not fully accurate NNP would be sufficient to test new methods.

Therefore, the goal of this project is to propose a NNP architecture that achieves an acceptable trade-off between accuracy and computational effort by using and fine-tuning well-established and tested methods to predict potential energy surfaces.

Some ideas and key features need to be used in any NNP architecture, such as a neighborlist or a differentiable cutoff function. Therefore, whenever possible, I will try to use existing implementations (for pytorch-based models, this could be for example the neighbor list implementation of NNPOps).

---

Implemented Models

This project provides implementations for four different Peanut model versions:

• PEANUT
• PEANUT dual cutoff
• PEANUT dual cutoff learned boundary
• PEANUT mini

Additionally, this repository provides an implementation for the

• SchNet model.

All peanut models use a graph neural network structure where atoms are considered as nodes and pair-wise interactions are modelled considering them as edges in the graph. The first model (PEANUT) corresponds to the sketched architecture shown in the PEANUT Documentation (except for the attention mechanism as mentioned below).

The second model (PEANUT dual cutoff) further develops the PEANUT model by using two ''edge MLPs''. The model divides all interacting atom pairs into two categories: Close and non-close neighbors. Both categories are further processed through two independent, but architecturally identical MLPs. This shall allow the model to learn the higher importance of closer neighbors in the atomic environment. Of course, the number of parameters grows. However, inference time should not be influenced by this since each pair is only processed through one of these two edge MLPs.

The third model (PEANUT dual cutoff learned boundary) is a more sophisticated version of the PEANUT dual cutoff model class. For each possible atom-type-pair, the model provides a learnable parameter to differentiate between close and non-close neighbors. The message is then computed as a weighted sum, based on the additional learnable parameter, of both messages. For this model, both the number of parameters and inference time grow.

The fourth model (PEANUT mini) was implemented to test how little information could be sufficient to learn meaningful behavior. This model will not be recommended for future experiments. However, it can serve as a baseline model to test minimal configuration performance.

For means of comparison, this repository also provides an implementation of the SchNet model. There was no hyperparameter fine-tuning done for this model. It only served as test model for the initial training pipeline.

---

Model evaluation

Even though the ultimate goal would be running MD simulations with this model, this is quite challenging. Apart from a fully trained model, a use in MD simulations would also require a full integration into an existing MD engine. And of course also several MD simulations, which take a lot of time to run. Therefore, the main goal is to test the model on single point conformations predicting single point energies only. If the model looks promising, an integration into an existing MD engine might be reconsidered. However, this would clearly be too much for the goal of this lecture. So this part remains optional, depending of this project's outcome.

Currently, the implemented loss function in the training routine considers energies (= predicted values) and forces (= negative gradient of the predicted values). This corresponds to the approach most other comparable deep learning models in this field use. Also, this covers the essential information needed to run MD simulations based on Newtons' equations of motion. Both (energies and forces) are part of the training data set. The loss function then is a weighted average of the energy loss and the force loss. During training, both MSE and MAE can be chosen for the force losses (to be set in the config file). The loss function for the forces was also part of the hyperparameter optimization part. For the energy prediction loss, MSE was implemented. The reason for this differentiation was the different levels of both values. Energy values are given per system / molecule in the dataset and correspondingly high. Forces are given as 3D vectors per atom and therefore contain rather small numbers (sometimes absolute value < 1.) For the evaluation part, MAE, MSE and RMSE were computed and logged.

---

Documentation

Some information about the PEANUT models' architectures is collected here (this is work in progress): PEANUT Documentation

1. Information about the chosen model architecture, covering four versions of the Peanut model:

   • Peanut model
     
   • Peanut dual cutoff model
     
   • Peanut dual cutoff learned boundary model
     
   • Peanut mini model
     

2. Several details on the chosen features for the representation learning part:

   • Radial features
   • Spherical features

3. Overview of the final energy prediction part.

4. Information on the chosen dataset.

5. Details on the implemented training procedure.

---

Installation

To implement this package, create an environment as follows:

git clone git@github.com:AnnaPicha/PEANUT.git
cd PEANUT

conda create -n peanut python=3.12
conda activate peanut 

# install
pip install -e .
conda install nnpops -c conda-forge

We also need to use scatter-add:

pip install torch-scatter -f https://data.pyg.org/whl/torch-2.6.0+${CUDA}.html

where ${CUDA} needs to match the pytorch version (cpu, cu126, cu128, or cu129). For more information, go to https://github.com/rusty1s/pytorch_scatter. (Unfortunately I had some issues installing this package for the cpu. The cuda version works for me.)

Then you can run a test with

pytest tests/test_building_blocks.py -s

In case you want to do the model training, you have to convert the dataset (in the repository data folder you can find the hdf5 file). To convert the file to .xyz and .pt format, you additionally need to install the extxyz package:

pip install git+https://github.com/libAtoms/extxyz

---

Demo

Once you have installed the package, you can run a small streamlit demo. First, install streamlit via pip

pip install streamlit

Then go to the streamlit folder and execute

streamlit run peanut_demo.py

In the demo, you can compute the energy and force vectors for a water and an ethane molecule using the Peanut Base model, the Peanut DualCutoff model and the Peanut DualCutoff LearnedBoundary model. The molecules are displayed in a 3D plot and the dircetions of the force vectors are shown. The coordinates of all atoms are editable. Once they are changed, energies and forces are re-evaluated. The demo can be run with 3 model types (Peanut, Peanut Dual Cutoff, Peanut Dual Cutoff Learned Boundary). Each of the three model was trained for 500 epochs.

---

Train the model

Data preparation

Go to src/peanut/data_preparation/ and execute

python prepare_data.py

This will transform the hdf5 data file in the datafolder to a torch usable .pt file. Simultaneously, a .xyz file will be created. To use fixed splits (training / validation / test dataset), also execute

python split_pubsolv.py

This will make fixed splits that are used during model training and save one single molecule to a .pt file for some simple testing in CI. If the model shall not be trained on fixed splits, this can be specified in the config.yaml file.

Configure

All settings can be set in the config.yaml file. You can either train a single model, or run a hyperparameter optimization.

Start training

To start hyperparameter optimization, go to models/hyperparameter_optimization. There is an exemplary submit script that will start several training runs for various sets of hyperparameters.

To perform one single training run, go to models/train_model. There is an exemplary submit script that will start one single training run using the settings as chosen in the config.yaml file.

Or, you can run

python ../../src/peanut/model_train/make_model.py --model_type ${model} --output_log_file ${log_file}

where model can be chosen as peanut, peanut_dualcutoff, peanut_dualcutoff_learned_boundary or peanut_mini. The log_filespecies the name of the output log file.

---

Results

Bottle necks

Most datasets generated from quantum mechanical (QM) methods contain so-called minimum-energy conformations. These are particularly stable configurations with comparatively low energy values. In general, this is not a problem for training a neural network. However, if the model is later intended to be used in a molecular dynamics (MD) simulation, it could pose challenges. During an MD simulation, the probability of the system being exactly in a minimum-energy configuration is low. As soon as the simulation begins and the atoms start moving away from a minimum, the system is already outside the training domain of the model. Therefore, the model's ability to generalize is essential for it to perform reliably in an MD simulation setting. To prevent this, dropout rates and an early stopping mechanism were employed. Apart from that, choosing a more diverse dataset should be considered.

Training results

Since the hyperparameter tuning process is very time intensive, I ran the optimization on a small subset of parameters. The chosen set of parameters is listed in the PEANUT Documentation. This of course can be extended in future experiments. All models were trained for 200 epochs on in total 9 nvidia RTX 4090 GPUs. The early stopping criterion never stopped the training early. Similarly, the validation loss (and training loss) monitored for each training run was still going down at the end of the training runs. So 200 epochs were certainly not long enough. However, this process takes a very long time and should serve as a first benchmark for the hyperparameter optimization process. For future experiments and training runs, the model(s) should be trained longer. The sets of best hyperparameters were chosen based on the lowest weighted validation loss (=weighted sum of energy and force losses).

Loss metrics

Several loss metrics were computed to monitor the training procedures. Since the target values of the model are potential energies, the error the model makes when predicting those are of course considered. However, in an MD simulation, the actually needed values are the forces that are acting on the atoms. Therefore, the model errors on the force predicition are considered as well. For both, MAE, MSE and RMSE were computed. For simplicity, only MAE are listed in the table below.

Hyperparameters

After training each of the four implemented model classes for 200 epochs. Since the Mini model only served as an initial benchmark model, the following table only shows the sets of hyperparameters for the other three Peanut model versions that delivered the most promising results:

Hyperparameter	Peanut	Peanut DualCutoff	Peanut DualCutoff LearnedBoundary
Number of radial symmetry functions	16	16	16
Depth of energy linear MLP	64	64	64
Depth of interaction (message passing) MLP	128	64	64
Loss function for forces	MAE	MAE	MAE
Use dropout in energy pass	False	False	False
Apply loss normalization (weighted)	True	True	True
Learning rate	0.001	0.001	0.0005

At the end of each training routine, the best model (best = lowest weighted combined loss) is saved and can be re-evaluated on the test set. To monitor the possibility of overfitting, the validation and training loss are plotted.

Target

All values are given in eV (electronvolt). Given that energies are the designated target value of the model, I expect the energy predictions to be more accurate than the force predictions. A mean absolute error of < 0.01 eV per atom would be desirable for energies. For forces (reported per atom per coordinate in eV/Å), a similarly small magnitude is targeted, keeping in mind that forces are obtained indirectly as gradients of the energy model. However, since forces are not directly predicted by the model, I expect higher differences in the force losses comparing the implemented architectures than in the energy values.

Datapoint example: the first molecule in the training dataset consists of 98 atoms and the energy is -345.672 eV -> average atom contribution of -3.53 eV.

The following table summarizes the best validation, training and final test losses for each PEANUT model variant. Force MAE values are reported per atom per coordinate and energy MAE values per atom. All values are given in eV. Again, the Mini model is not considered in this table:

Loss comparison (MAE)

Loss Metric	Peanut	Peanut DualCutoff	Peanut DualCutoff LearnedBoundary
Training Total Loss (MAE)	1.0578e-01	1.0246e-01	1.0128e-01
Validation Total Loss (MAE)	1.0926e-01	1.0593e-01	1.0417e-01
Test Total Loss (MAE)	1.1879e-01	1.1607e-01	1.1471e-01
Training Force Loss (MAE)	1.1346e-01	1.0984e-01	1.0871e-01
Validation Force Loss (MAE)	1.1635e-01	1.1268e-01	1.1091e-01
Test Force Loss (MAE)	9.9046e-02	9.7569e-02	9.7287e-02
Training Energy Loss (MAE)	9.7681e-03	1.0221e-02	8.4890e-03
Validation Energy Loss (MAE)	2.0619e-02	2.1554e-02	1.9870e-02
Test Energy Loss (MAE)	1.9789e-02	1.8497e-02	1.7427e-02

Notes:

• Force Loss is normalized per atom and per coordinate (x, y, z).
• Energy Loss is normalized per atom across the datapoint.
• Total Loss is a weighted sum of energy and force loss.

Visualization

After a succesful training run, the training losses and validation losses per epoch are stored in csv files and plotted. Also, energies and forces are plotted in scatter plots. The energy values are compared based on the average energy per atom in a molecule, forces are compared component-wise.

Tests and code structure

The implemented model(s) are stored in src/peanut/building_blocks/model.py. The static features that are used for the representation learning mechanism are implemented in src/peanut/building_blocks/static_features.py. The code needed to pre-process the hdf5 data format is stored in src/peanut/data_preparation. The needed steps are described in the Train the model / Data preparation section above. The training routine is implemented in src/peanut/model_train/make_model.py. The steps needed to start a model training are described in the Train the model / Start training section above.

To test the code, several tests were implemented. All tests can be found in the testsfolder:

• neighbors_test.py tests the neighborlist which is implemented in src/peanut/building_blocks/neighbor.py
• read_config.pychecks if all required hyperparameters that are needed for the training routine are provided in the config.yaml file.
• read_data_set.py checks if the dataset provides the expected features.
• test_building_blocks.py checks if all needed parts for the static feature construction are working (returning tensors with the correct shape etc.).

Locally, they can be run using pytest. Additionally, a .github/workflows/tests.yml file was added to this repository. All tests are executed using CI.

Workplan

Phase	Status	Focus	Key Tasks	Notes	Code Structure	Estimated time spent
Planning & Setup	✅ Done	Project definition & environment	Dataset selection, environment setup, initial model design	Baseline pipeline established	pyproject.toml, base config, repo layout	10 h
Core Architecture	✅ Done	Model building blocks	Node/edge features, message passing, energy readout	Fully functional model	Modular building blocks, clear module separation	40 h
Feature Refinement	✅ Done	Architectural extensions	Radial & directional bases, multi-scale edges, weighted messages	Attention not finalized	Separate building blocks	16 h
Training & Evaluation	✅ Done	Model training	Training loop, energy & force losses, MAE/MSE/RMSE	Stable training achieved	Training logic separated from model	30 h
Overfitting Control	✅ Done	Generalization & robustness	Train/val/test split, early stopping, dropout options	Validation-driven decisions	Configurable via YAML	6 h
Hyperparameter Optimization	✅ Done	Search & scaling	Choose parameters, ideal loss functions, ...	Best configs under exploration	CLI overrides + config system	20 h
Testing	✅ Done	Code quality & usability	Unit tests, logging, plots	Log results to csv files and plots	Tests per module, logging utilities	10 h
Clean up & Documentation	✅ Done	User friendliness	Run tests with CI, format code, read the docs, README	Update as project progresses	.github/workflows, read the docs	10 h
Final Review, Demo & Cleanup	✅ Done	Polish & submission	Refactoring, formatting, Write streamlit demo, CI, final results	Final step after succesful training	Repo-wide cleanup	8h

---

References

Model Architectures

1. Directional Message Passing for Molecular Graphs
   Johannes Gasteiger, Janek Groß, Stephan Günnemann, 2022. arXiv:2003.03123

2. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
   Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, Gábor Csányi, 2023. arXiv:2206.07697

3. MACE-OFF: Short-Range Transferable Machine Learning Force Fields for Organic Molecules
   Dávid Péter Kovács, J. Harry Moore, Nicholas J. Browning, Ilyes Batatia, Joshua T. Horton, Yixuan Pu, Venkat Kapil, William C. Witt, Ioan-Bogdan Magdău, Daniel J. Cole, Gábor Csányi, 2025. DOI: 10.1021/jacs.4c07099

4. SchNet: A Continuous-Filter Convolutional Neural Network for Modeling Quantum Interactions
   Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, Klaus-Robert Müller, 2017. arXiv:1706.08566

5. ANI-1: An Extensible Neural Network Potential with DFT Accuracy at Force Field Computational Cost
   J. S. Smith, O. Isayev, A. E. Roitberg, 2017. DOI: 10.1039/C6SC05720A

6. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials
   Batzner, S., Musaelian, A., Sun, L. et al., 2022. doi.org/10.1038/s41467-022-29939-5

---

Datasets

1. ANI-2x
   Extending the Applicability of the ANI Deep Learning Molecular Potential to Sulfur and Halogens
   Christian Devereux, Justin S. Smith, Kate K. Huddleston, Kipton Barros, Roman Zubatyuk, Olexandr Isayev, Adrian E. Roitberg, 2020. DOI: 10.1021/acs.jctc.0c00121

2. SPICE
   A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials
   Eastman, P., Behara, P. K., Dotson, D. L., et al., 2023. DOI: 10.1038/s41597-022-01882-6

3. PubChem
   Open Chemistry Database, National Institutes of Health.

---

License

This project is licensed under the MIT License.