ConformFlow

KO | EN

gitlite — search

#python #android #golang #docker #vue #ontology #go #cli #nodejs #mongodb #ai #cpp

★ 69 Open GitHub ↗

Latent normalizing flow for protein conformational ensemble generation

// 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

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README