Learning Non-Local Molecular Interactions via Equivariant Local Representations and Charge Equilibration

Overview

The supplementary code is organized in the following way:

examples/ -- Scripts to train models and perform experiments
  | 4GHDNNP/ -- Benchmark systems
    | models/ -- Trained models
    | 4GHDNNP*.py -- Training scripts for models
    | Benchmark*.ipynb -- Scripts to compute errors and run benchmarks
    | data.py  -- Utilities to load reference data
    | eval_utils.py -- Utilities to compute errors
    | train_utils.py -- Utilities to setup models and optimizations
  | OE62 -- Scripts for experiments on OE62 dataset
    | models/ -- Trained models
    | Benchmark.ipynb -- Evaluation of performed benchmarks
    | eval_utils.py -- Utilities to compute errors
    | evaluate.py -- Compute timings for a model
    | oe62_predict.py -- Predict energies and forces with a model
    | oe62_train.py -- Training scripts for models
    | train_utils.py -- Utilities to setup models and optimizations
  | spice: -- Scripts for experiments on SPICE dataset
    | models/ -- Trained models
    | simulations/ -- Trajectories simulated with models
    | spice_simulate.py -- Script to simulate trajectories with a model
    | train.py -- Training script for models
    | train_utils -- Utilities to setup models, simuations, and optimization
external/ -- External libraries
  | chemtrain/ -- Chemtrain package used to train models
    | chemtrain/ -- Source code of chemtrain
    | jax_md_mod/ -- JAX-MD modifications
      | model/ -- Source code for DimeNet++ model
      | custom_electrostatics.py -- Charge equilibration implementation
      ...
    ...
  | chemutils/ -- Datasets and models
    | chemutils/ -- Source code of chemutils
      | datasets/ -- Tools to load and process dataset
        | spice.py --  SPICE dataset
        ...
      | models/ -- Model implementations
        | allegroQeq/ -- Modified allegro model and CELLI
        | mace/ -- Modified MACE model
        ...
      ...
requirements.txt -- Requirements to run examples
...

Install required software

First, install the required packages for the project:

pip install -r requirements.txt

The packages chemtrain and chemutils can be installed via pip:

pip install -e "external/chemtrain[all]"
pip install -e "external/chemutils"

To install chemtrain with GPU support through JAX, please refer to the JAX Installation Instructions.

Citation

If you use this work or the accompanying code, please cite:

@ARTICLE{Fuchs2025-ii,
  title    = "Learning non-local molecular interactions via equivariant local
              representations and charge equilibration",
  author   = "Fuchs, Paul and Sanocki, Micha{\l} and Zavadlav, Julija",
  journal  = "npj Computational Materials",
  volume   =  11,
  number   =  1,
  pages    = "287",
  month    =  sep,
  year     =  2025
}