Rotation-equivariant Graph Neural Networks for Learning Glassy Liquids Representations
======

[toc]

This repository is the official implementation of [Rotation-equivariant Graph Neural Networks for Learning Glassy Liquids Representations](https://arxiv.org/abs/2211.03226).

To be cited as (TODO: add correct bibtex entry - referrring to SciPost instead of arxiv):
> @article{pezzicoli2022se,
title={SE (3)-equivariant graph neural networks for learning glassy liquids representations},
author={Pezzicoli, Francesco Saverio and Charpiat, Guillaume and Landes, Fran{\c{c}}ois P},
journal={arXiv preprint arXiv:2211.03226},
year={2022}
}

<!-- [put image /home/flandes/downloads/SciPostsubmission/scipost resubmission/fig3_update_and_cm.pdf , here] -->

## 1. Requirements

### 1.1 Packages

Here we assume you have access to hardware with a GPU and CUDA, python3 and conda or equivalent (e.g. mamba).

We provide 3 methods:
- create a new conda env just for this code, from scratch, using our environment.yml file
- create a new conda env just for this code, from scratch
- simply install the additional packages needed in your currently existing conda env where torch is already installed

#### Check out your cuda version:
`nvcc --version`
In our case, we have cuda11.2 but will install cuda11.8 in the python packages.
Using an even more recent version should be fine, you just have to make sure your cuda on the python side is fine with the one on the nvcc side.

#### Create a new conda env just for this code, from scratch, using our environment.yml file

`conda env create -f se3gnn4glasses.yml`

`conda activate se3gnn4glasses`

If this does not work well for you, try the next method instead:

#### Create a new conda env just for this code, from scratch

First, clean up:
`conda remove -n se3gnn4glasses --all`

1. Then, create the env:
`
conda create -n se3gnn4glasses
`

2. And install torch (on GPU):
`
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
`

> For CPU-only install, do instead :
`conda install pytorch torchvision torchaudio cpuonly -c pytorch `

You can check everything is ok by running:

python -c "import torch; print(torch.__version__)
print(torch.version.cuda); import torch_geometric;"

3. Install pytorch geometric, to build GNNs:
`conda install pyg -c pyg`
You can check everything is ok by running:
` python -c "import torch; print(torch.version.cuda); import torch_geometric; print(torch_geometric.__version__)" `

4. Install e3nn
`conda install e3nn`

5. Install also scatter:
`conda install pytorch-scatter -c pyg`

#### Simply install the additional packages needed in your currently existing conda env where torch is already installed

Starting directly from step 3. or 4. in the above list.

------

## 2. Data

Our code is compatible with 3 different datasets, that originate from 3 different papers.

In chronological order:

(2020) Bapst's [dataset](https://console.cloud.google.com/storage/browser/deepmind-research-glassy-dynamics/) (click public_dataset), from the paper [Unveiling the predictive power of static structure in glassy systems](https://www.nature.com/articles/s41567-020-0842-8).

(2022) Shiba's [dataset](https://ipomoea-www.cc.u-tokyo.ac.jp/i29002/botan/public_dataset.tar.gz) and related [code]( https://github.com/h3-Open-BDEC/pyg_botan), from the paper [BOTAN: BOnd TArgeting Network for prediction of slow glassy dynamics by machine learning relative motion](https://pubs.aip.org/aip/jcp/article/158/8/084503/2868947/BOTAN-BOnd-TArgeting-Network-for-prediction-of).

(2022) Jung's dataset, from paper [Predicting dynamic heterogeneity in glass-forming liquids by physics-inspired machine learning](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.130.238202), is available in the larger Roadmap's GlassBench (seee below).

For Information: The collectively written Roadmap paper (2023), organized by G. Jung [Roadmap on machine learning glassy liquids (temporary title)](https://arxiv.org/abs/2311.14752) introduces a uniform database for training and testing models: [GlassBench](https://doi.org/10.5281/zenodo.10118191), which covers the Jung and the Shiba's data, and more (some 2D data)

In all cases, we are using either the quenched positions (IS: Inherent States) as input, or the potential energy associated to the configuration.
These original datasets only provide the so-called "thermal" positions, one needs to compute the corresponding IS.

### 2.1 Ready-to-use train files

We provide some ready-to-use files, to make it easier to reproduce our paper results.

We only provide a large number of these files for the Bapst dataset (our paper is centered on this dataset), for the training temperature T=0.44, the smallest one avaialbe (on which we focus in the paper).

These pre-processed files are in
- `./data/BapstNPT/T_044/raw`
> Note that the first time you run a network on this input, it will produce the `./data/BapstNPT/T_044/processed` folder with the appropriate files in it. It may take a while. Once the `processed/` folder is complete, the next runs will be faster (only training).

We do not provide all datasets' and temperatures' pre-processed files to avoid using too much disk storage. If you deem it useful, contact us, we can upload more data in a future version of this Zenodo repository.

We provide a couple of files (4 for each kind and temperature) to make it easier to quickly tryout transfer experiments.

### 2.2 Data pre-processing by yourself

In that case you'll need to install `hoomd`, or hack our (small) pre-processing codes, to compute the IS.

For now we have not ported our codes from hoomd v2 to hoomd v3 or v4.

conda create -n hoomd297
conda activate hoomd297
conda install -c conda-forge hoomd=2.9.7

It may be difficult to install hoomd with GPU support. However, it takes about 100 seconds to obtain 1 Inherent State even on a simple CPU, so since this step is needed only once per file, it may be simpler to use a CPU (or multiple ones..)

#### 2.2.1 Shiba dataset:

You can download this original dataset from here:
[dataset](https://ipomoea-www.cc.u-tokyo.ac.jp/i29002/botan/public_dataset.tar.gz) (it's a 60G .zip file).

Or a small sample from the github repo:
https://github.com/h3-Open-BDEC/pyg_botan

cd ./data/preprocessing/
git clone https://github.com/h3-Open-BDEC/pyg_botan
mv pyg_botan/small_data/T0.44/train/*   downloaded_data/Shiba_dataset/T0_44/train/.
conda activate hoomd297
python Fire_quench_Shiba.py  --gpu 0  --input_path  "./downloaded_data/Shiba_dataset/T0_44/train/*"
mv downloaded_data/Shiba_dataset/T0_44/train-FIRE/*   ../Shiba_dataset_torch/T_044/raw/.

Note that it is important to locate the files correctly for our code to work.

If you have GPU support with your hoomd install, of course you should then use the flag `--gpu 1` instead of `--gpu 0`

You may now even remove the original files:

rm  ./downloaded_data/Shiba_dataset/T0_44/train/*npz
cd ..
ls

#### 2.2.2 Bapst dataset:

You can download this original dataset from here: [dataset](https://console.cloud.google.com/storage/browser/deepmind-research-glassy-dynamics/) (go for public_dataset).

Then you need to run `Fire_quench_bapst.py` on these files so as to compute the IS positions, similarly to the case above.

#### 2.2.3 Jung dataset:

To download the Jung (KA2D) dataset: see GlassBench, above.

## 3. Code usage

*To use our scripts on CPU, add the option -g 0 to the calls. By default, we expect a GPU to be present.*

### 3.1 Training a model

To train the model as in the paper, since you probably use a cluster:

conda acticate se3gnn4glasses
sbatch Launch_training.sh

Or directly from a terminal:
`python main.py`
(for options, see the commands inside `Launch_training.sh`). All the hyper-parameters used for training are also in that file.

Note: to avoid mistakes during transfer-learning, the code (main.py) constrains the directeory structure quite a lot.
The code is expected to be ran from a folder containing at least 3 folders:
- `./src/` with codes inside
- `./data/` with data inside (with fixed names)
- `./expe/` will be created at the first run, and each run will output there.

### 3.2 Pre-trained Models

We provide two experiments' outputs, to avoid the cost of training again the models. See the `./expe/` folder, it contains two experiments, one for the Baspt dataset, and one for the Shiba dataset. In each case, we trained on the smallest temperature available, for 7-10h on a A100 GPU.

### 3.3 Transfer-learning

Warning: even if you use our pre-trained model, you need a large amount of RAM to be able to fine-tune the model: the forward pass incurrs the computation of numerous activations, in the current version of the code.

Again, the directory structure is constrained, to avoid mistakes. One should:
- cd to `./expe/`
- cd to the desired experiment (the training from which we want to load the model)

See the detailed commands inside `Launch_training.sh`:

cd expe/
cd "$(ls -dt */ | head -1)" ## with thanks to chatGPT

epochs=50
epochs=10

N_train=2 ## this is enough to fine tune some decoders
N_train=10 ## to improve a bit

echo -e "\nStarting the transfer learning now. We are in the directory:\n"
pwd
echo -e "\n"
## with thanks to chatGPT
for input_dir in ../../data/BapstNPT/*/; do
    input_dir="${input_dir%/}"  ## trims off the last / character
    echo -e "\n\nTransfer using input dir: $input_dir \n"
    time python3 ./src/main.py --mode transfer  --input_dir $input_dir \
    --epochs $epochs --N_train $N_train
done

### 3.4 Prediction mode (pure eval mode)

Similarly to the transfer case, we can use the model trained at a given temperature, to apply it to (test it on) others, without re-learning anything. The logic is very similar to transfer, one simply does not re-learn the decoders.

See the detailed commands inside `Launch_training.sh`:

Essentially, it is:
- cd to `./expe/`
- cd to the desired experiment (the training from which we want to load the model)
- `python ./src/main.py --mode predict --input_dir ../../data/BapstNPT/T_056 -g 0
`

### 3.5 Basic hyper-parameter changes

Typical choices one may want to investigate include:
- the network's depth and width: for each layer, how many channel of each angular order ($\ell $ value) and how many layers ? This is specified using the argument `--hidden_irreps`. Our paper uses `--hidden_irreps ['8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e', '8x0e+8x1e+8x2e+8x3e']`
- `--listTimesIdx` : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] -- the time indices to be regressed. Simply use [9] to regress only the last time index.
- `--e_pot_node_feat` : use the IS-calculated Potential Energy, or its thermal counterpart, or nothing.
- `--masked_loss`: {0,A,B}. Choose A,B if you want to mask A or B particles in the loss, 0 for no masking

Others are pretty straightforward to understand from main.py

### 4. Licence

Our licence is the [GNU GPL v3](https://www.gnu.org/licenses/gpl-3.0.en.html).

##### Credits:
- e3nn
- This readme was written using the template available here: https://github.com/paperswithcode/releasing-research-code/blob/master/templates/README.md

---------------------