# TorchMD-NET

TorchMD-NET provides state-of-the-art graph neural networks and equivariant transformer neural networks potentials for learning molecular potentials. It offers an efficient and fast implementation and it is integrated in GPU-accelerated molecular dynamics code like [ACEMD](https://www.acellera.com/products/molecular-dynamics-software-gpu-acemd/) and [OpenMM](https://www.openmm.org). See the full paper at https://arxiv.org/abs/2202.02541.

## Installation

Clone the repository:
```
git clone https://github.com/torchmd/torchmd-net.git
cd torchmd-net
```

Create a Conda environment and activate it:
```
conda env create -f environment.yml
conda activate torchmd-net
```

Install TorchMD-NET into the Conda environment:
```
pip install -e .
```

## Usage


Specifying training arguments can either be done via a configuration yaml file or through command line arguments directly. An example configuration file for a TorchMD Graph Network can be found at [examples/graph-network.yaml](https://github.com/compsciencelab/torchmd-net/blob/main/examples/graph-network.yaml). For an example on how to train the network on the QM9 dataset, see [examples/train_GN_QM9.sh](https://github.com/compsciencelab/torchmd-net/blob/main/examples/train_GN_QM9.sh). GPUs can be selected by their index by listing the device IDs (coming from `nvidia-smi`) in the `CUDA_VISIBLE_DEVICES` environment variable. Otherwise, the argument `--ngpus` can be used to select the number of GPUs to train on (-1 uses all available GPUs or the ones specified in `CUDA_VISIBLE_DEVICES`).

To pre-train on PCQM4Mv2 with denoising:

```
CUDA_VISIBLE_DEVICES=0,1 python scripts/train.py --conf examples/ET-PCQM4MV2.yaml --dataset QM9 --log-dir experiments/pcq_denoising/ 
```

To fine-tune on QM9 using `checkpoints/pcq-denoising.ckpt`: 

```
CUDA_VISIBLE_DEVICES=0,1 python scripts/train.py --load-model checkpoints/pcq-denoising.ckpt --conf examples/ET-QM9.yaml --position-noise-scale 0.01 --output-model-noise VectorOutput --denoising-weight 0.1  --log-dir experiments/qm9/
```

## Credit

This repo is developed off the [code base by Thölke and De Fabritiis](https://github.com/torchmd/torchmd-net).