Update README.md

README.md

@@ -1,3 +1,6 @@
+---
+license: cc-by-nc-sa-2.0
+---
 <div align='center'>
 <h1> MolCRAFT Series for Drug Design </h1>
 
@@ -68,49 +71,111 @@ The MolCRAFT series addresses critical challenges in generative models for SBDD,
     * Achieves 95.9% PoseBusters passing rate on CrossDock with significantly improved molecular geometries.
 
 ---
+# MolCRAFT
+Official implementation of ICML 2024 ["MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space"](https://arxiv.org/abs/2404.12141).
 
-## ⚙️ Installation & Usage
+🎉 Our demo is now available [here](http://61.241.63.126:8000). Welcome to have a try!
 
-Please refer to the `README.md` file within each project's subdirectory for specific instructions on installation, dependencies (docker recommended), and how to run the code.
+![](../asset/molcraft_framework.png)
 
----
+## Environment
+It is highly recommended to install via docker if a Linux server with NVIDIA GPU is available.
 
-## 📊 Datasets and Benchmarks
+Otherwise, you might check [README for env](docker/README.md) for further details of docker or conda setup.
 
-Our models are evaluated on standard benchmarks in the field, such as:
+### Prerequisite
+A docker with `nvidia-container-runtime` enabled on your Linux system is required.
 
-* **CrossDocked2020**: Used for evaluating binding affinity, molecular validity, and optimization success rates.
-* **PoseBusters V2**: Used for assessing the quality of generated molecular poses.
+> [!TIP]
+> - This repo provides an easy-to-use script to install docker and nvidia-container-runtime, in `./docker` run `sudo ./setup_docker_for_host.sh` to set up your host machine.
+> - For details, please refer to the [install guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
 
-Details about the specific datasets used for training and evaluation can be found in the respective publications and project READMEs.
 
----
+### Install via Docker
+We highly recommend you to set up the environment via docker, since all you need to do is a simple `make` command.
+```bash
+cd ./docker
+make
+```
 
-## 🤝 Contributing
 
-We welcome contributions to the MolCRAFT series! If you are interested in contributing, please feel free to fork the repository, make your changes, and submit a pull request. You can also open an issue if you find any bugs or have suggestions for improvements.
+-----
+## Data
+We use the same data as [TargetDiff](https://github.com/guanjq/targetdiff/tree/main?tab=readme-ov-file#data). Data used for training / evaluating the model should be put in the `data` folder by default, and accessible in the [data](https://drive.google.com/drive/folders/1j21cc7-97TedKh_El5E34yI8o5ckI7eK?usp=share_link) Google Drive folder.
+
+To train the model from scratch, download the lmdb file and split file into data folder:
+* `crossdocked_v1.1_rmsd1.0_pocket10_processed_final.lmdb`
+* `crossdocked_pocket10_pose_split.pt`
+
+To evaluate the model on the test set, download _and_ unzip the `test_set.zip` into data folder. It includes the original PDB files that will be used in Vina Docking.
+
 
 ---
+## Training
+Run `make -f scripts.mk` (without the need for data preparation), or alternatively (with data folder correctly configured),
+```bash
+python train_bfn.py --exp_name ${EXP_NAME} --revision ${REVISION}
+```
 
-## 📝 Citation
+where the default values should be set the same as:
+```bash
+python train_bfn.py --sigma1_coord 0.03 --beta1 1.5 --lr 5e-4 --time_emb_dim 1 --epochs 15 --max_grad_norm Q --destination_prediction True --use_discrete_t True --num_samples 10 --sampling_strategy end_back_pmf
+```
 
-If you use any of the methods or code from this repository in your research, please cite the respective papers:
+### Testing
+For quick evaluation of the official checkpoint, refer to `make evaluate` in `scripts.mk`:
+```bash
+python train_bfn.py --test_only --no_wandb --ckpt_path ./checkpoints/${CKPT_NAME}
+```
 
-```bibtex
-@article{qiu2025piloting,
-  title={Piloting Structure-Based Drug Design via Modality-Specific Optimal Schedule},
-  author={Qiu, Keyue and Song, Yuxuan and Fan, Zhehuan and Liu, Peidong and Zhang, Zhe and Zheng, Mingyue and Zhou, Hao and Ma, Wei-Ying},
-  journal={ICML 2025},
-  year={2025}
-}
+### Debugging
+For quick debugging training process, run `make debug -f scripts.mk`:
+```bash
+python train_bfn.py --no_wandb --debug --epochs 1
+```
 
-@article{qiu2025empower,
-  title={Empower Structure-Based Molecule Optimization with Gradient Guidance},
-  author={Qiu, Keyue and Song, Yuxuan and Yu, Jie and Ma, Hongbo and Cao, Ziyao and Zhang, Zhilong and Wu, Yushuai and Zheng, Mingyue and Zhou, Hao and Ma, Wei-Ying},
-  journal={ICML 2025},
-  year={2025}
-}
+## Sampling
+We provide the pretrained MolCRAFT checkpoint [here](https://drive.google.com/file/d/1TcUQM7Lw1klH2wOVBu20cTsvBTcC1WKu/view?usp=share_link). 
+
+
+### Sampling for pockets in the testset
+Run `make evaluate -f scripts.mk`, or alternatively,
+```bash
+python train_bfn.py --config_file configs/default.yaml --exp_name ${EXP_NAME} --revision ${REVISION} --test_only --num_samples ${NUM_MOLS_PER_POCKET} --sample_steps 100
+```
 
+The output molecules `vina_docked.pt` for all 100 test pockets will be saved in `./logs/${USER}_bfn_sbdd/${EXP_NAME}/${REVISION}/test_outputs/${TIMESTAMP}` folders.
+
+### Sampling from pdb file
+To sample from a whole protein pdb file, we need the corresponding reference ligand to clip the protein pocket (a 10A region around the reference position).
+
+Below is an example that stores the generated 10 molecules under `output` folder. The configurations are managed in the ``call()`` function of ``sample_for_pocket.py``.
+
+```bash
+python sample_for_pocket.py ${PDB_PATH} ${SDF_PATH}
+```
+
+## Evaluation
+### Evaluating molecules
+For binding affinity (Vina Score / Min / Dock) and molecular properties (QED, SA), it is calculated upon sampling.
+
+For PoseCheck (strain energy, clashes) and other conformational results (bond length, bond angle, torsion angle, RMSD), please refer to `test` folder.
+
+### Evaluating meta files
+We provide samples for all SBDD baselines in the [sample](https://drive.google.com/drive/folders/1A3Mthm9ksbfUnMCe5T2noGsiEV1RfChH?usp=sharing) Google Drive folder.
+
+You may download the `all_samples.tar.gz` and then `tar xzvf all_samples.tar.gz`, which extracts all the pt files into `samples` folder for evaluation.
+
+<!-- ## Demo
+### Host our web app demo locally
+
+With ``gradio`` and ``gradio_molecule3d`` installed, you can simply run ``python app.py`` to open the demo locally. Port mapping has been set in Makefile if you are using docker. You should also forward this port if you run the docker in an ssh server. We will share a permanent demo link later.
+
+Great thanks to @duerrsimon for his kind support in resolving rendering issues! -->
+
+## Citation
+
+```
 @article{qu2024molcraft,
   title={MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space},
   author={Qu, Yanru and Qiu, Keyue and Song, Yuxuan and Gong, Jingjing and Han, Jiawei and Zheng, Mingyue and Zhou, Hao and Ma, Wei-Ying},
@@ -124,8 +189,4 @@ If you use any of the methods or code from this repository in your research, ple
   journal={ICLR 2024},
   year={2024}
 }
-```
-
-## 📄 License 
-
-The project is licensed under the terms of the CC-BY-NC-SA license. See [LICENSE](https://github.com/algomole/MolCRAFT/blob/main/LICENSE) for more details.
+```