ConformFlow

KO | EN

gitlite — search

#hacktoberfest #ai #go #android #linux #macos #java #windows #awesome #cli #pytorch #ios #openai #open-source

(★ 69)

Latent normalizing flow for protein conformational ensemble generation

ConformFlow 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# ConformFlow

**Scalable normalizing flow for protein conformational ensemble generation.**

ConformFlow is a normalizing-flow framework for sequence-conditioned
protein conformational ensemble generation. It combines a continuous
backbone latent representation with a RealNVP-style flow parameterized
by sequence-aware Transformer coupling networks. The result is a model
that supports:

- **Exact-likelihood training** on equilibrium MD trajectories and
  single structures from large protein databases.
- **Single-step parallel sampling** — typically **10–80× faster** than
  diffusion-based ensemble generators.
- **Plug-and-play conditional generation** with arbitrary differentiable
  geometric constraints (radius of gyration, end-to-end distance, RMSD
  to a native structure, contact distances, …) — no retraining required.

For full methodology and benchmarks see our paper, *ConformFlow: scalable
normalizing flow for protein conformational ensemble generation*.

## Installation

```bash
git clone https://github.com/Harrydirk41/ConformFlow.git
cd ConformFlow
pip install -r requirements.txt
```

Requirements:

- Python ≥ 3.10
- `torch`, `pytorch_lightning`, `numpy`
- `esm` — the ESM3 package (provides the frozen structure
  encoder/decoder trunks that the autoencoder is built on)
- `mdtraj` — for `.xtc` / `.pdb` trajectory I/O at inference time

## Hugging Face setup

ConformFlow relies on **ESM3** for the frozen structure encoder/decoder
trunks. Before running anything, you need to authenticate with
Hugging Face so the ESM3 weights can be downloaded on first use.
Please refer to
<https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1> and
<https://github.com/evolutionaryscale/esm> for full details on obtaining
a token.

1. **Create a Hugging Face account** at
   <https://huggingface.co/join>.
2. **Generate a token.** Go to
   <https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1> and follow
   the instructions to accept the model license and obtain an access
   token.
3. **Log in (one-time setup).** Run in your terminal and paste your
   token when prompted:

```bash
   huggingface-cli login
   ```

After this one-time setup, the ESM3 weights will be cached in
`~/.cache/huggingface/` and used automatically by every subsequent
training and inference run.

## Repository layout

```
conformflow.py          # ProteinRealNVP + AE Lightning module + NF Lightning module
custom_rotary.py        # rotary positional embeddings
esm/                    # ESM3 fork (frozen encoder/decoder + Stage-1 losses)

train_autoencoder.py    # train the continuous structure autoencoder
train_nf.py             # train the latent normalizing flow (with optional --encode)

inference_utils.py      # shared helpers (model loaders, batched O inference, XTC/PDB I/O)
infer_encode.py         # structure  →  latent embedding
infer_sample.py         # sequence   →  i.i.d. backbone ensemble
infer_pnp.py            # sequence + constraint  →  guided ensemble
```

A trained run produces two checkpoints, organised by latent dimension `M`:

```
./model/Latent_4/autoencoder/last.ckpt
./model/Latent_4/nf/last.ckpt
```

## Data format

Place your training proteins as `.npz` files under `./data/raw_data/`.
Each file represents one protein (or one MD frame) and must contain:

| key     | shape       | dtype   | description                                    |
|---------|-------------|---------|------------------------------------------------|
| `struc` | `(L, 3, 3)` | float32 | backbone atoms ordered (N, CA, C) in Å         |
| `seq`   | `(L,)`      | int64   | amino-acid indices 0…19 (20 reserved for pad)  |

## Training

### 1. Train the autoencoder

```bash
python train_autoencoder.py
```

Trains the continuous structure autoencoder (frozen ESM3 trunks +
trainable adapter MLPs) on `./data/raw_data/` using the four geometric
Stage-1 reconstruction losses. Checkpoints are written to
`./model/Latent_4/autoencoder/`.

Common overrides:

- `--latent_dim 8` — change the per-residue latent dimension `M`.
- `--batch_size 16 --epochs 200` — larger batch, longer schedule.
- `--devices 4` — multi-GPU via DDP.

### 2. Train the normalizing flow

The flow trains on latents produced by the autoencoder. First-time users
should pass `--encode` to encode `./data/raw_data/` through the trained
autoencoder into `./data/latent_data/` before training:

```bash
python train_nf.py --encode
```

Subsequent runs reuse the cached latents, so the flag can be omitted:

```bash
python train_nf.py
```

Checkpoints are written to `./model/Latent_4/nf/`.

## Inference

All three inference scripts default to the `M=4` checkpoints produced by
the training workflow. Outputs are written to
`./result/{task}/{protein_name}/`.

### (1) Encode a structure to its latent

```bash
python infer_encode.py --input native.pdb
```

Produces:

```
./result/embedding/native/
    latent.npy         # (T, L, M) per-residue continuous latent
    recon.npy          # (T, L, 3, 3) reconstructed backbone in Å
    recon.xtc          # trajectory with inferred O atoms
    recon_first.pdb    # first frame as PDB
    sequence.txt
```

### (2) Sample an ensemble from a sequence

```bash
python infer_sample.py \
    --sequence YYDPETGTWY \
    --protein_name chignolin \
    --n_samples 2000
```

Produces:

```
./result/sample/chignolin/
    sample.npy         # (N, L, 3, 3) backbone in Å
    sample.xtc         # full ensemble as trajectory
    first_frame.pdb
    sequence.txt
```

### (3) Plug-and-play guided generation

ConformFlow supports inference-time control under arbitrary
differentiable geometric constraints. Three example observables are
built in:

- `rg` — radius of gyration of Cα atoms
- `eed` — Cα end-to-end distance
- `rmsd_to_native` — Kabsch RMSD to a user-supplied PDB

```bash
# Radius of gyration target 12 Å
python infer_pnp.py \
    --sequence YYDPETGTWY \
    --protein_name chignolin \
    --feature rg --target 12.0
```

```bash
# RMSD to a native structure
python infer_pnp.py \
    --sequence YYDPETGTWY \
    --protein_name chignolin \
    --feature rmsd_to_native --target 2.5 --native_pdb native.pdb
```

Adding a new observable is a few lines: write a torch function
`f(coords: (B, L, 3, 3)) -> (B,)` and register it in the `FEATURES`
dict at the top of `infer_pnp.py`. Everything else (Langevin loop,
gradients through the decoder, XTC/PDB output) is automatic.

Outputs:

```
./result/pnp/chignolin/
    sample.npy         # (K, L, 3, 3) guided backbones
    sample.xtc
    first_frame.pdb
    stats.json         # feature-value distribution and |f − v| statistics
    sequence.txt
```

## Reproducing the paper

The default hyperparameters in this repository match those used in the
paper:

| Parameter            | Default                                            |
|----------------------|----------------------------------------------------|
| Latent dimension `M` | 4                                                  |
| Flow depth `K`       | 12 BigBlocks                                       |
| Block depth          | 2 affine coupling layers (alt-odd + alt-even)      |
| Transformer depth    | 4 layers per coupling network                      |
| Embedding dim        | 256                                                |
| Heads                | 8                                                  |
| FFN                  | SwiGLU                                             |
| Positional encoding  | Rotary (RoPE)                                      |
| Optimizer            | AdamW (lr = 1e-4, wd = 1e-5, β = (0.9, 0.95))      |
| Warmup               | 500 steps                                          |
| Precision            | bf16-mixed                                         |

## Citation

Paper coming soon!

## Contact

For questions, please open an issue or contact:

- Yikai Liu — `liu3307@purdue.edu`
- Ming Chen — `chen4116@purdue.edu`

## License

This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

ConformFlow

ConformFlow 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

ConformFlow

ConformFlow 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README