init

Author: chengtan9907

.gitignore

@@ -0,0 +1,11 @@
+.conda
+.vscode
+.ipynb*
+__pycache__
+
+results/
+
+core.*
+
+.DS_Store
+checkpoint.pt

README.md

@@ -0,0 +1,152 @@
+# CryoFold: One-shot Prediction For Cryo-EM Structure Determination
+
+CryoFold is a deep learning framework for automating the determination of three-dimensional atomic structures from high-resolution cryo-electron microscopy (Cryo-EM) density maps. It addresses the limitations of existing AI-based methods by providing an end-to-end solution that integrates training and inference into a single streamlined pipeline. CryoFold combines 3D and sequence Transformers for feature extraction and employs an equivariant graph neural network to build accurate atomic structures from density maps.
+
+## Table of Contents
+- [Background](#background)
+- [Features](#features)
+- [Installation](#installation)
+- [Quick Start](#quick-start)
+- [Usage](#usage)
+  - [Command-Line Arguments](#command-line-arguments)
+  - [Running the Example](#running-the-example)
+  - [Using Custom Data](#using-custom-data)
+- [Tutorial](#tutorial)
+- [References](#references)
+- [Contact](#contact)
+- [License](#license)
+
+## Background
+
+Cryo-electron microscopy (Cryo-EM) has revolutionized structural biology by enabling the visualization of complex biological molecules at near-atomic resolution. The technique generates **high-resolution density maps** that offer insights into the molecular structures of proteins, viruses, and other biomolecular assemblies. However, **interpreting these density maps to derive accurate atomic models** remains a challenging and labor-intensive task, often requiring expert knowledge and manual interventions.
+
+Existing AI-based methods for automating Cryo-EM structure determination face several limitations:
+1. **Multi-stage processing**: Current approaches often involve separate stages for feature extraction, sequence alignment, and structure prediction, leading to inefficiencies and discontinuities.
+2. **Alignment bias**: Techniques such as **Hidden Markov Models (HMMs)** or **Traveling Salesman Problem (TSP) solvers** introduce bias when aligning predicted atomic coordinates with the protein sequence.
+3. **Poor generalization**: Due to the limited size of available datasets, many methods struggle to generalize well to complex or previously unseen test cases.
+
+CryoFold addresses these challenges by providing a **fully integrated, end-to-end solution** that performs **one-shot inference** with minimal manual intervention, enabling faster and more accurate structure determination.
+
+## Features
+
+- **🚀 End-to-End Training and Inference**: Simplifies the process by seamlessly integrating training and inference into a single, unified framework, eliminating the need for multi-stage processing.
+- **⚡ Fast and Accurate**: Achieves a **400% improvement in TM-score** over Cryo2Struct while reducing inference time by a factor of **1,000**.
+
+For more details on the performance and benchmarking, please refer to our paper.
+
+## Installation
+
+To get started with CryoFold, follow these steps:
+
+1. **Clone the repository**:
+
+   ```bash
+   git clone https://github.com/A4Bio/CryoFold.git
+   cd CryoFold
+   ```
+
+2. **Create and activate the conda environment**:
+
+    ```bash
+    conda env create -f environment.yml
+    conda activate cryofold
+    ```
+
+3. **Download the Pretrained Model**:
+
+    We provide a pretrained model for CryoFold. [Download it here]() and place it in the pretrained_models directory.
+
+
+## Quick Start
+
+To quickly try out CryoFold using an example dataset, run the following command:
+
+```
+bash run_example.sh
+```
+
+This script runs the `inference.py` script with sample data provided in the `examples` folder. It uses a sample density map and a ground truth PDB file for evaluation.
+
+We also provide an example tutorial in `quick_start.ipynb`.
+
+## Usage
+
+### Command-line Arguments
+
+The `inference.py` script supports several command-line arguments:
+
+| Argument                 | Description                                             | Default                             |
+|--------------------------|---------------------------------------------------------|-------------------------------------|
+| `--density_map_path`     | Path to the input density map directory (required).     | None                                |
+| `--pdb_path`             | Path to the ground truth PDB file (optional).           | None                                |
+| `--model_path`           | Path to the pretrained model checkpoint.                | `pretrained_model/checkpoint.pt`    |
+| `--output_dir`           | Directory to save the output PDB file.                  | `results`                           |
+| `--device`               | Device to run the model on (`cpu` or `cuda`).           | `cuda`                              |
+| `--verbose`              | Enable verbose output for debugging.                    | Disabled                            |
+
+### Running the Example
+
+You can run the example directly from the command line:
+
+```bash
+python inference.py --density_map_path examples/density_map --pdb_path examples/5uz7.pdb
+```
+
+### Using Custom Data
+
+To use CryoFold with your own data, you need to provide a Cryo-EM density map and, optionally, a PDB file for evaluating the predicted structure. For example:
+
+```bash
+python inference.py --density_map_path /path/to/your/density_map --pdb_path /path/to/your/ground_truth.pdb --output_dir /path/to/save/results --device cuda
+```
+
+## Tutorial
+
+### 1. Preprocessing Density Maps:
+
+To normalize your density maps, run:
+
+	# Normalize you density maps
+	$ bash run_data_preparation.bash examples/
+
+After preprocessing, the directory structure should look like:
+
+The organization of the downloaded models should look like:
+```text 
+CryoFold
+├── examples
+│   ├── density_map
+│   │   ├── map.map
+│   │   ├── seq_chain_info.json
+│   │   └── normed_map.mrc
+```
+
+### 2. Running Inference:
+
+	python inference.py --density_map_path examples/density_map --pdb_path examples/5uz7.pdb 
+
+After inference, the output will be saved in the specified output directory:
+
+```text 
+CryoFold
+├── results
+│   └── output.pdb
+```
+
+## References:
+
+For a complete description of the method, see:
+
+
+
+## Contact
+
+Please submit any bug reports, feature requests, or general usage feedback as a github issue or discussion.
+
+- Jue Wang ([email protected])
+- Cheng Tan ([email protected])
+- Zhangyang Gao ([email protected])
+
+## License
+
+This project is licensed under the MIT License. See the LICENSE file for details.

environment.yml

@@ -0,0 +1,29 @@
+name: cryofold
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - python=3.9
+  - numpy
+  - pandas
+  - scipy
+  - matplotlib
+  - pytorch=2.0.1
+  - torchvision
+  - torchaudio
+  - cudatoolkit=11.7
+  - scikit-learn
+  - pip
+  - pip:
+    - transformers==4.24.0
+    - tqdm
+    - jupyterlab
+    - einops
+    - torch_scatter
+    - torch_geometric
+    - torch_cluster
+    - fair-esm
+    - mrcfile
+    - biopython
+prefix: /root/anaconda3/envs/cryofold