Research data for "Extending the Range of Graph Neural Networks with Global Encodings"

Package Installation

Please, refer to the GitHub repository of RANGE for up-to-date instructions on how to install the codebase.

File Reconstruction

To reconstruct the .zip files from the split files, run in the terminal:

zip -s 0 input.zip --out output.zip

Model Training

Inside the training directory, data are organized in the following way:

- 📂 dataset directory
  - 📂 baseline, range or non-regularized range
    - 📂 model type (SchNet, Schnet+RANGE ...)
      - 📂 baseline cutoff and seed
          - 📄 best_model.ckpt
          - 📄 config.yaml

EwaldMP results are self contained within the EwaldMP directory.

The configuration file used for training is provided together with the best model checkpoint (i.e. minimum validation loss).
To train a new model with the same hyperparameters, first update the config.yaml file, selecting the correct path for the dataset splits and the directory where the dataset will be automatically downloaded and processed:

...
# Dataset configuration example in the .yaml file 
data:
  dataset:
    class_path: rangemp.datasets.AQMgasDataset
    init_args:
      root: YOURPATH/RANGE_review/datasets/AQM/
  log_dir: default_root_dir
  splits: YOURPATH/RANGE_review/datasets/AQM/selection_n_atoms_30_train_0.7_val_0.15_seed_4272389.npz
...

A progress bar can be visualized changing the corresponding flag (do this only if you are not redirecting into a text file).
After that, the model can be trained using

python YOURPATH/RANGE_review/mlcg/scripts/mlcg-train.py fit --config config.yaml

for baseline, Ewald, and non-regularized RANGE models, or

python YOURPATH/RANGE_review/mlcg/scripts/mlcg-train-regularized.py fit --config config.yaml

for regularized RANGE models.

If EwaldMP model are tested remember to export the pythonpath

export PYTHONPATH=/YOURPATH/RANGE_review/training/EwaldMP/ewaldmp

By default, training is done on GPU. It is also possible to use CPU by setting

...
  accelerator: 'cpu'
  # strategy: 'ddp'
...

in the config.yaml file.

The trained model can be extracted from a checkpoint using

from mlcg.utils import extract_model_from_checkpoint

PATH = "SOME_PATH.ckpt"
model_extracted = extract_model_from_checkpoint(PATH, None)

After the model is trained, it is possible to test it using the script provided in the metrics_evaluation directory by modifying the example_setup.yaml file appropriately:

python compute_metrics.py --config example_setup.yaml

In case of memory overload, consider to properly set batch_size and inference_batch_size in accordance with your hardware.

Model Simulation

The initial configurations used to perform the simulations, together with the corresponding simulation trajectories from the manuscript, are provided for the buckyball catcher, DHA, and double-walled nanotube datasets.
Before running any simulations, process the model checkpoints saved during training. These checkpoints can be found in the corresponding training folders.
Be sure to select the trained model corresponding to the system you want to simulate.

Use the mlcg-combine_model.py script as follows:

python YOURPATH/RANGE_review/mlcg/scripts/mlcg-combine_model.py --ckpt YOURMODEL.ckpt --prior YOURPATH/RANGE_review/simulations/empty_prior.pt --out combined_model.pt

In all simulations reported in the manuscript, no prior energy term is used.
The object in YOURPATH/RANGE_review/simulations/empty_prior.pt is an empty PyTorch dictionary that is only used to correctly process the model for simulation, without adding any additional prior energy contribution.

Simulations are performed using the mlcg-nvt_langevin.py script. An input YAML file is provided (please be sure to properly modify the paths inside):

python YOURPATH/RANGE_review/mlcg/scripts/mlcg-nvt_langevin.py --config sim_config.yaml

The simulation output consists of .npy files, each containing a segment of independent trajectories (10 for DHA, 5 for buckyball catcher and 2 for double-walled nanotube). The shape of the .npy files is (n_trajectories, n_frames, n_atoms, 3). To reconstruct the full trajectories, all coordinate files should be loaded and concatenated along the first axis.

Neural P3M

All Neural P3M experiment were conducted using the original code provided by the autors.
A separate environment need to be created to reproduce the test, following the instruction provided in the P3M repository.
Training configurations and best model checkpoints are reported for different datasets in training_p3m.