upload mofdiff source code.

Author: kyonofx

README.md

@@ -1,33 +1,195 @@
-# Project
+# MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design
 
-> This repo has been populated by an initial template to help get you started. Please
-> make sure to update the content to build a great experience for community-building.
+`mofdiff` is a diffusion model for generating coarse-grained MOF structures. This codebase also contains the code for deconstructing/reconstructing the all-atom MOF structures to train MOFDiff and assemble CG structures generated by MOFDiff.
 
-As the maintainer of this project, please make a few updates:
+[paper](https://arxiv.org/abs/2310.10732) | [data and pretained models](https://zenodo.org/uploads/10467288)
 
-- Improving this README.MD file to provide a great experience
-- Updating SUPPORT.MD with content about this project's support experience
-- Understanding the security reporting process in SECURITY.MD
-- Remove this section from the README
+If you find this code useful, please consider referencing our paper:
 
-## Contributing
+```
+@article{fu2023mofdiff,
+  title={MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design},
+  author={Fu, Xiang and Xie, Tian and Rosen, Andrew S and Jaakkola, Tommi and Smith, Jake},
+  journal={arXiv preprint arXiv:2310.10732},
+  year={2023}
+}
+```
 
-This project welcomes contributions and suggestions.  Most contributions require you to agree to a
-Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
-the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
+## Table of Contents
 
-When you submit a pull request, a CLA bot will automatically determine whether you need to provide
-a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions
-provided by the bot. You will only need to do this once across all repos using our CLA.
+- [Installation](#installation)
+- [Dowlnload data](#download-data)
+- [Training](#training)
+- [Generating MOF structures](#generating-mof-structures)
+- [Assemble all-atom MOFs](#assemble-all-atom-mofs)
+- [Relax MOFs](#relax-mofs)
 
-This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
-For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or
-contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
+## Installation
 
-## Trademarks
+We recommend using [mamba](https://mamba.readthedocs.io/en/latest/) (much faster than conda) to install the dependencies. First install `mamba` following the intructions in the [mamba repository](https://mamba.readthedocs.io/en/latest/installation/mamba-installation.html).
 
-This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft 
-trademarks or logos is subject to and must follow 
-[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general).
-Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship.
-Any use of third-party trademarks or logos are subject to those third-party's policies.
+
+Install dependencies via `mamba`:
+
+```
+mamba env create -f env.yml
+```
+
+Then install `mofdiff` as a package:
+
+```
+pip install -e .
+```
+
+We use [MOFid](https://github.com/snurr-group/mofid) for preprocessing and analysis. Install MOFid following the instruction in the [MOFid repository](https://github.com/snurr-group/mofid/blob/master/compiling.md). The generative modeling part of this codebase does not depend on MOFid.
+
+## Download data
+
+You can download the preprocessed data from [Zenodo](https://zenodo.org/uploads/10467288) (recommended). 
+
+Alternatively, you can download the `BW-DB` raw data from [Materials Cloud](https://archive.materialscloud.org/record/2018.0016/v3) and preprocess the data with the following command (assuming the data is downloaded to `${raw_path}`, this step requires MOFid):
+
+```
+python preprocessing/extract_mofid.py --df_path ${raw_path}/all_MOFs_screening_data.csv --cif_path ${raw_path}/cifs --save_path ${raw_path}/mofid
+python preprocessing/preprocess.py --dataset_path
+python preprocessing/save_to_lmdb.py
+```
+
+The preprocessing inovlves 3 steps:
+1. Extract the MOFid for all structures (CPU).
+2. Construct CG MOF data objects from MOFid deconstruction results (CPU or GPU).
+3. Save the CG MOF objects to an LMDB database (relatively fast).
+
+The entire preprocessing process for `BW-DB` may take several days (depending on the CPU/GPU resources).
+
+## Training
+
+First, configure the `.env` file to set correct paths to various directories. An [example](./.env) `.env` file is provided in the repository.
+
+### training the building block encoder
+
+Before training the diffusion model, we need to train the building block encoder. The building block encoder is a graph neural network that encodes the building blocks of MOFs. The building block encoder is trained with the following command:
+
+```
+python mofdiff/scripts/train.py --config-name=bb
+```
+
+The default output directory is `${oc.env:HYDRA_JOBS}/bb/${expname}/`. `oc.env:HYDRA_JOBS` is configured in `.env`. `expname` is configured in `configs/bb.yaml`. We use [hydra](https://hydra.cc/) for config management. All configs are stored in `configs/` You can override the default output directory with command line arguments. For example:
+
+```
+python mofdiff/scripts/train.py --config-name=bb expname=bwdb_bb_dim_64 model.latent_dim=64
+```
+
+Logging is done with [wandb](https://wandb.ai/site) by default. You need to login to wandb with `wandb login` before training. The training logs will be saved to the wandb project `mofdiff`. You can also override the wandb project with command line arguments. You can also disable wandb logging by removing the `wandb` entry in the [config](./conf/logging/default.yaml).
+
+### training coarse-grained diffusion model for MOFs
+
+The output directory where the building block encoder is saved: `bb_encoder_path` is needed for training the diffusion model. With the building block encoder trained to convergence, train the CG diffusion model with the following command:
+
+```
+python mofdiff/scripts/train.py data.bb_encoder_path=${bb_encoder_path}
+```
+
+For BW-DB, training the building block encoder takes roughly 3 days and training the diffusion model takes roughly 5 days on a single NVIDIA V100 GPU.
+
+## Generating CG MOF structures
+
+Pretrained models can be found [here](https://zenodo.org/record/10467288).
+
+With a trained CG diffusion model `${diffusion_model_path}`, generate random CG MOF structures with the following command:
+
+```
+python mofdiff/scripts/sample.py --model_path ${diffusion_model_path} --bb_cache_path ${bb_cache_path}
+```
+
+`${bb_cache_path}` is the path to the building block embedding space, saved at the beginning of CG diffusion model training. To optimize MOF structures for a property (e.g., CO2 adsorption working capacity), use the following command:
+
+```
+python mofdiff/scripts/optimize.py --model_path ${diffusion_model_path} --bb_cache_path ${bb_cache_path} --data_path ${data_path}
+```
+
+Available arguments for `sample.py` and `optimize.py` can be found in the respective files. The generated CG MOF structures will be saved in `${sample_path}=${diffusion_model_path}/${sample_tag}` as `samples.pt`.
+
+The CG structures generated with the diffusion model are not guaranteed to be realizable. We need to assemble the CG structures to recover the all-atom MOF structures. The following sections describe how to assemble the CG MOF structures, and all steps further do not require a GPU.
+
+## Assemble all-atom MOFs
+
+Assembled the CG MOF structures with the following command:
+
+```
+python mofdiff/scripts/assemble.py --input ${sample_path}/samples.pt
+```
+
+This command will assemble the CG MOF structures in `${sample_path}` and save the assembled MOFs in `${sample_path}/assembled.pt`. The cif files of the assembled MOFs will be saved in `${sample_path}/cif`. If the assembled MOFs came from property-driven optimization, the optimization arguments are saved to `${sample_path}/opt_args.json`.
+
+## Relax MOFs and compute structural properties
+
+The assembled structures may not be physically plausible. These MOF structures are relaxed uses the UFF force field with LAMMPS. LAMMPS is already installed if you have followed the installation instructions in this README. The script for relaxing the MOF structures also compute structural properties (e.g., pore volume, surface area, etc.) with [Zeo++](https://www.zeoplusplus.org/download.html) and the mofids of the generated MOFs with [MOFid](https://github.com/snurr-group/mofid/tree/master). The respective packages should be installed following the instructions in the respective repositories, and the corresponding paths should be added to `.env` before running the following command. Each step should take no more than a few minutes to complete on a single CPU. We use multiprocessing to parallelize the computation.
+
+Relax MOFs and compute structural properties with the following command:
+
+```
+python mofdiff/scripts/uff_relax.py --input ${sample_path}
+```
+
+This command will relax the assembled MOFs in `${sample_path}/cif` and save the relaxed MOFs in `${sample_path}/relaxed`. The structural properties of the relaxed MOFs will be saved in `${sample_path}/relaxed/zeo_props_relax.json`. The mofids of the relaxed MOFs will be saved in `${sample_path}/mofid`.
+
+
+## GCMC simulation for gas adsorption
+
+To run GCMC simulations, first install RASPA2 (simulation software) and eGULP (charge calculation software).
+
+RASPA2 can be installed with `pip`:
+
+```
+pip install "RASPA2==2.0.4"
+```
+
+You may need to install the following Linux dependencies first:
+
+```
+apt-get update 
+apt-get install -yq libgsl0-dev pkg-config libxrender-dev
+```
+
+Install [eGULP](https://github.com/danieleongari/egulp) following the instruction in the repository. The following commands install eGULP in `/usr/local/bin/egulp`:
+
+```
+mkdir /usr/local/bin/egulp && tar -xf egulp.tar -C /usr/local/bin/egulp
+cd /usr/local/bin/egulp/src && make && cd -
+```
+
+Then, decompress the [force field parameters](./mofdiff/gcmc/UFF-TraPPe-scaled.tar) to the RASPA directory using the following commands (assuming RASPA2 installed in `RASPA_PATH=PYTHONPATH/site-packages/RASPA2` with `pip`):
+
+```
+tar -xf UFF-TraPPe-scaled.tar -C RASPA_PATH/share/raspa/forcefield/UFF-TraPPe
+```
+
+Calculate charges for relaxed samples in `${sample_path}` with the following command:
+
+```
+python mofdiff/scripts/calculate_charges.py --input ${sample_path}
+```
+
+This command will output cif files with charge information under `${sample_path}/mepo_qeq_charges`. 
+
+
+Run GCMC simulations with the following command:
+
+
+```
+python mofdiff/scripts/gcmc_screen.py --input ${sample_path}/mepo_qeq_charges
+```
+
+The GCMC simulation results will be saved in `${sample_path}/gcmc/screening_results.json`.
+
+## Acknowledgement
+
+This codebase is based on several existing repositories:
+
+- [CDVAE](https://github.com/txie-93/cdvae)
+- [Open catalyst project](https://github.com/Open-Catalyst-Project/ocp)
+- [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric)
+- [PyTorch](https://github.com/pytorch/pytorch)
+- [Lightning](https://github.com/Lightning-AI/pytorch-lightning/)
+- [Hydra](https://github.com/facebookresearch/hydra)

conf/bb.yaml

@@ -0,0 +1,28 @@
+expname: bwdb_bb
+workdir: ${oc.env:HYDRA_JOBS}/bb_models/${expname}
+config_for: bb
+
+core:
+  version: 0.0.1
+  tags:
+    - ${now:%Y-%m-%d}
+
+hydra:
+  run:
+    dir: ${oc.env:HYDRA_JOBS}/bb/${expname}/
+
+  sweep:
+    dir: ${oc.env:HYDRA_JOBS}/bb/${expname}/
+    subdir: ${hydra.job.num}_${hydra.job.id}
+
+  job:
+    env_set:
+      WANDB_START_METHOD: thread
+      WANDB_DIR: ${oc.env:WANDB_DIR}
+
+defaults:
+  - data: bwdb_bb
+  - logging: default
+  - model: bb
+  - optim: default
+  - train: default

conf/data/bwdb_bb.yaml

@@ -0,0 +1,50 @@
+name: bwdb_bb
+root_path: ${oc.env:DATASET_DIR}
+use_type_mapper: true
+
+# max num bbs
+max_bbs: 20
+max_atoms: 200
+max_cps: 20
+
+train_max_steps: 1500000
+early_stopping_patience: 100
+patience: 10
+
+data_cache_path: ${oc.env:DATASET_DIR}
+load_cached: true
+save_cached: true
+
+datamodule:
+  _target_: mofdiff.data.datamodule.DataModule
+
+  datasets:
+    train:
+      _target_: mofdiff.data.dataset.BBDataset
+      name: ${data.name}_train
+      path: ${data.root_path}
+      max_bbs: ${data.max_bbs}
+      max_atoms: ${data.max_atoms}
+      max_cps: ${data.max_cps}
+      split_file: ${oc.env:PROJECT_ROOT}/splits/train_split.txt
+
+    val:
+      _target_: mofdiff.data.dataset.BBDataset
+      name: ${data.name}_train
+      path: ${data.root_path}
+      max_bbs: ${data.max_bbs}
+      max_atoms: ${data.max_atoms}
+      max_cps: ${data.max_cps}
+      split_file: ${oc.env:PROJECT_ROOT}/splits/val_split.txt
+
+  num_workers:
+    train: 0
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 1024
+    val: 1024
+    test: 1024
+
+data_transforms: None

conf/data/bwdb_mof.yaml

@@ -0,0 +1,66 @@
+name: bwdb
+root_path: ${oc.env:DATASET_DIR}
+prop_list: 
+  - working_capacity_vacuum_swing [mmol/g]
+  - working_capacity_temperature_swing [mmol/g]
+logmod: true
+num_targets: 2
+use_type_mapper: false
+bb_encoder_path: ???
+
+lattice_scale_method: scale_length
+
+max_bbs: 20
+max_atoms: 200
+max_cps: 20
+otf_graph: false
+
+train_max_steps: 2000000
+early_stopping_patience: 1000
+teacher_forcing_max_epoch: 300
+patience: 50
+
+data_cache_path: ${oc.env:DATASET_DIR}
+load_cached: true
+save_cached: true
+
+datamodule:
+  _target_: mofdiff.data.datamodule.DataModule
+  bb_encoder_path: ${data.bb_encoder_path}
+
+  datasets:
+    train:
+      _target_: mofdiff.data.dataset.MOFDataset
+      name: ${data.name}_train
+      path: ${data.root_path}
+      prop_list: ${data.prop_list}
+      transforms: ${data.data_transforms}
+      max_bbs: ${data.max_bbs}
+      max_atoms: ${data.max_atoms}
+      max_cps: ${data.max_cps}
+      logmod: ${data.logmod}
+      split_file: ${oc.env:PROJECT_ROOT}/splits/train_split.txt
+
+    val:
+      _target_: mofdiff.data.dataset.MOFDataset
+      name: ${data.name}_val
+      path: ${data.root_path}
+      prop_list: ${data.prop_list}
+      transforms: ${data.data_transforms}
+      max_bbs: ${data.max_bbs}
+      max_atoms: ${data.max_atoms}
+      max_cps: ${data.max_cps}
+      logmod: ${data.logmod}
+      split_file: ${oc.env:PROJECT_ROOT}/splits/val_split.txt
+
+  num_workers:
+    train: 0
+    val: 0
+    test: 0
+
+  batch_size:
+    train: 128
+    val: 128
+    test: 128
+
+data_transforms: None

conf/logging/default.yaml

@@ -0,0 +1,22 @@
+# log frequency
+val_check_interval: 3
+progress_bar_refresh_rate: 10
+
+wandb:
+  name: ${expname}
+  project: mofdiff
+  entity: null
+  log_model: True
+  mode: 'online'
+  group: ${expname}
+
+tensorboard:
+  save_dir: ${oc.env:LOG_DIR}/tensorboard
+
+wandb_watch:
+  log: 'all'
+  log_freq: 5000
+
+lr_monitor:
+  logging_interval: "step"
+  log_momentum: False