diff --git a/README.md b/README.md index 628be8e..25e7b0a 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # MoreRed: Molecular Relaxation by Reverse Diffusion -[MoreRed](https://arxiv.org/abs/2404.10935) is a generative diffusion model that can generate new structures or denoise arbitrarily noisy ones. Unlike previous geometry relaxation methods, which require labeled equilibrium and non-equilibrium structures, MoreRed is trained using **exclusively unlabeled equilibrium structures**. Despite this, it effectively relaxes non-equilibrium structures, achieving competitive results with **100 times less data**, exhibiting better robustness to the noise level in the input and reducing computation time during relaxation. +[MoreRed](https://arxiv.org/abs/2404.10935) is a generative diffusion model that can generate new structures or denoise arbitrarily noisy ones. Unlike previous geometry relaxation methods, which require labeled equilibrium and non-equilibrium structures, MoreRed is trained using **exclusively unlabeled equilibrium structures**. Despite this, it effectively relaxes non-equilibrium structures, achieving competitive results with **much less data**, exhibiting better robustness to the noise level in the input and reducing computation time during relaxation. @@ -20,10 +20,11 @@ MoreRed is built on top of [SchNetPack 2.0](https://github.com/atomistic-machine-learning/schnetpack/tree/master), an easily configurable and extendible library for constructing and training neural network models for atomistic systems like molecules. SchNetPack utilizes [PyTorch Lightning](https://www.pytorchlightning.ai/) for model building and [Hydra](https://hydra.cc/) for straightforward management of experimental configurations. While high level usage of the `morered` package to train and use the models described in the original paper does not require knowledge of its underlying dependencies, we recommend users familiarize themselves with Hydra to be able to customize their experimental configurations. Additionally, the tutorials and the documentation provided in [SchNetPack 2.0](https://github.com/atomistic-machine-learning/schnetpack/tree/master) can be helpful. Below, we explain how to use the `morered` package. -**_NOTE: while the current documentation in the README file and the source code should be sufficient to easily use the package, we are continually enhancing it._** +**_NOTE: while the current documentation in the README file and the source code should be sufficient to use the package easily, we will continually enhance it._** #### Content ++ [News](/README.md##News) + [Installation](/README.md##Installation) + [Training](/README.md##Training) + [Pre-trained models](/README.md##Pre-trained-models) @@ -32,6 +33,10 @@ MoreRed is built on top of [SchNetPack 2.0](https://github.com/atomistic-machine + [Tutorials](/README.md##Tutorials) + [How to cite](/README.md##How-to-cite) +## News +- Under the folder `morered`, we uploaded the final models trained on QM9 and QM7-X datasets +- We updated the notebook `denoising_tutorial.ipynb` with more details on using the trained models. + ## Installation Requirements: - python >= 3.8 @@ -55,7 +60,7 @@ pip install . ``` ## Training -The human-readable and customizable YAML configuration files under `src/morered/configs` are all you need to train and run customizable experiments with `morered`. They follow the configuration structure used in [SchNetPack 2.0](https://github.com/atomistic-machine-learning/schnetpack/tree/master). Here, we explain how to train and use the different models. Besides, under the folder `notebooks` we provide step-by-step Jupyter notebooks explaining the building blocks of MoreRed and how to use the different trained models. +The human-readable and customizable YAML configuration files under `src/morered/configs` are all you need to train and run customizable experiments with `morered`. They follow the configuration structure used in [SchNetPack 2.0](https://github.com/atomistic-machine-learning/schnetpack/tree/master). Here, we explain how to train and use the different models. Installing `morered` using pip adds the new CLI command `mrdtrain`, which can be used to train the different models by running the command: ``` @@ -101,15 +106,16 @@ mrdtrain --config-dir= experiment= More about overwriting configurations in the CLI can be found in the [SchNetPack 2.0](https://github.com/atomistic-machine-learning/schnetpack/tree/master) documentation. ## Pre-trained models -Under the folder `models`, you can find models trained on both QM9 and Qm7-X datasets. You can load the models using `torch.load()` command. Besides, the tutorial notebooks provide details on how to use the models. +Under the folder `models`, you can find the final models trained on QM9 and Qm7-X datasets until complete convergence. You can load the models using `torch.load()` command. Besides, the tutorial notebooks provide details on how to use the models. ## Molecular relaxation/generation -The notebook `notebooks/denoising_tutorial.ipynb` explains how the trained models can be used for denoising. -Under `src/morered/sampling`, you can find ready-to-use Python classes that implement the different samplers: `MoreRed-ITP`, `MoreRed-JT`, `MoreRed-AS`, `DDPM`. The same classes can be used for denoising/relaxation of noisy structures as well as for new structure generation. +The notebook `notebooks/denoising_tutorial.ipynb` explains how the trained models can be used for denoising and generation from scratch using different samplers. +Under `src/morered/sampling`, you can find ready-to-use Python classes implementing the different samplers: `MoreRed-ITP`, `MoreRed-JT`, `MoreRed-AS`, `DDPM`. The same classes can be used for denoising/relaxation of noisy structures as well as for new structure generation. ## Tutorials -Under `notebooks`, we provide different tutorials in the form of Jupyter notebooks: - - `diffusion_tutorial.ipynb`: explains how to use the diffusion models and samplers implemented in `morered`. In the notebook, the DDPM sampler is used as an example. Yet it can be replaced by any other implemented sampler class (`MoreRed-ITP`, `MoreRed-JT`, `MoreRed-AS`) in `src/morered/sampling`. +Under `notebooks`, we provide different tutorials in the form of Jupyter notebooks, that will be continually updated: + - `diffusion_tutorial.ipynb`: explains how to use the diffusion processes implemented in `morered`. + - `denoising_tutorial.ipynb`: explains how to use the trained models with the different samplers implemented in `morered` for noisy structure relaxation or generation from scratch. Under the folder `models`, we provide final models trained on QM9 and QM7-X datasets until complete convergence. ## How to cite if you use MoreRed in your research, please cite the corresponding publication: