UniFlow

KO | EN

gitlite — search

#python #java #python3 #arduino #golang #machine-learning #rust #html #flask #javascript #seismology #nodejs

★ 23 Open GitHub ↗

An internal coordinate normalizing flow that unifies protein ensemble generation and machine-learned coarse-grained force fields for molecular dynamics simulation within a single framework

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# 🧬 UniFlow: Unifying Machine-Learned Force Fields and Conformational Ensemble Generation with Scalable Normalizing Flows

[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://raw.githubusercontent.com/Harrydirk41/UniFlow/main/LICENSE)
[![Python](https://img.shields.io/badge/python-3.10+-blue.svg)]()
[![PyTorch](https://img.shields.io/badge/PyTorch-2.0+-red.svg)]()
[![Paper](https://img.shields.io/badge/arXiv-UniFlow-blue)](https://arxiv.org/abs/XXXX.XXXXX)

---

## ✨ Overview

UniFlow is a sequence-conditioned normalizing flow for protein conformational
ensemble generation. Conditioned on a native structure and its sequence, it
models protein backbone conformational ensembles using a circular-spline
coupling flow over internal coordinates. The same trained flow serves three
purposes:

- **I.i.d. sampling** — draw a backbone conformational ensemble in a single flow pass.
- **Langevin dynamics** — treat UniFlow as a machine-learned coarse-grained force
  field and run underdamped Langevin dynamics for protein systems of interest.
- **Likelihood evaluation** — score the log-likelihood of existing structures.

## 🛠 Installation

```bash
pip install numpy torch pytorch_lightning mdtraj tqdm einops pyyaml matplotlib
```

Requirements:

- Python ≥ 3.10
- `torch`, `pytorch_lightning`, `numpy` — model, training, math
- `mdtraj` — `.xtc` / `.pdb` I/O and DSSP
- `einops`, `pyyaml`, `matplotlib`, `tqdm`

Every entry point below also accepts `--config <file>.yml` (see the shipped
`*_config.yml`); YAML keys set defaults and any command-line flag overrides them.

## 🏋️ Training

### Training-data format

A **data root** is a directory of per-protein subdirectories; each protein
subdirectory holds one native `.pdb` and one or more MD `.xtc` trajectories of
that protein (`.xtc` and `.pdb` file names are arbitrary):

```
DATA_ROOT/                 # <- this path is `roots` in the config
  protein_A/
    protein_A.pdb          # native structure (sequence + N/CA/C reference)
    md_0.xtc  md_1.xtc     # MD frames of protein_A
  protein_B/
    protein_B.pdb
    md_0.xtc
  ...
```

Set your paths once in `training_config.yml`:

```yaml
roots:      [/path/to/DATA_ROOT]   # one or more data-root dirs
cache_dir:  /path/to/frame_cache   # where the preprocessed cache is written
model_dir:  uniflow_ckpts          # where checkpoints are written
```

### Prepare the frame cache (optional)

Parse each protein once into a memory-mapped backbone cache so training reads only
the frames it samples instead of re-reading every `.xtc` each step:

```bash
python prepare_data.py --config training_config.yml
```

(You can skip this — training builds the cache lazily on the first epoch.)

### Train

```bash
python train.py --config training_config.yml
```

Loss = `nll_weight · NLL + clash_weight · clash + contact_weight · contact`
(set the geometry weights to 0 for pure NLL). Optimizer is AdamW
(lr 1e-4, β = (0.9, 0.95)) with linear warmup. Checkpoints (`last.ckpt` + top-3)
are written to `model_dir`.

## 🎲 Sampling (one-shot)

Set the checkpoint and target protein in `sampling_config.yml`:

```yaml
ckpt:       uniflow_ckpts/last.ckpt   # trained checkpoint
pdb:        native.pdb                # native structure to condition on
n_samples:  1000
```

```bash
python sample.py --config sampling_config.yml
```

Draws `n_samples` conformations conditioned on the native, infers oxygen, and
writes to `sample/iid/{pdb_name}/` (set `out_dir` to override):

```
sample/iid/{pdb_name}/
    samples/    samples.xtc  first_frame.pdb  samples.npy    # N,CA,C,O
    analysis/   fes  geometry  ramachandran  secondary_structure   (.png/.pdf)
    metrics.csv  metrics.json      # per-sample + ensemble-mean row
```

Analysis figures: free-energy profiles of end-to-end / Rg / RMSD-to-native;
Cα–Cα bond, angle, torsion, and contact map; sample-vs-native Ramachandran; and
per-residue secondary-structure population. `metrics.csv` gives per-structure
`rg`, `eed`, `rmsd`, helix/sheet/coil segment counts, clash count, mean Cα–Cα
bond/angle, `logP_model`, `logP_cartesian`, and native-contact order parameters
`Q` and `f`.

## 📊 Likelihood (score existing structures)

Set the checkpoint and the protein directory to score in `likelihood_config.yml`:

```yaml
ckpt:  uniflow_ckpts/last.ckpt   # trained checkpoint
dir:   /path/to/protein_dir      # a dir with one native .pdb and >=1 .xtc to score
```

```bash
python likelihood.py --config likelihood_config.yml
```

Scores every frame of the `.xtc` in a protein directory under the flow and writes
to `sample/likelihood/{dir_name}/`:

```
sample/likelihood/{dir_name}/
    structures.npy     # (N_struc, L, 3, 3) Å
    logprob.npy        # (N_struc, 2) = [logP_model, logP_cartesian]
```

`logP_model` is the internal-coordinate flow density; `logP_cartesian` maps it to
Cartesian space via the exact bond–angle–torsion Jacobian plus a harmonic bond
prior around the native bonds.

## 🌀 Langevin (ML force field)

Set the checkpoint, protein, and dynamics options in `langevin_config.yml`:

```yaml
ckpt:       uniflow_ckpts/last.ckpt   # trained checkpoint
pdb:        native.pdb                # native structure to start from
n_walkers:  8
dt:         0.001                     # ps
ghmc:       true                      # Metropolis-adjusted, unbiased (recommended)
max_force:  0                         # no force cap (correct with GHMC)
```

```bash
python langevin.py --config langevin_config.yml
```

Runs underdamped Langevin dynamics (BAOAB) and writes the time-trajectory
(`traj.npy` + per-walker `.xtc` + `first_frame.pdb`) to `sample/md/{pdb_name}/`.
Useful options:

- `ghmc: true` + `max_force: 0` — Metropolis-adjusted, unbiased sampling
  (recommended); watch the `acc` readout and lower `dt` if acceptance is low.
- `compile: true` — `torch.compile` the energy for speed (safe with GHMC).
- `gauge_fix: true` — remove the 6 external DOF each step (exact internal marginal).

## 📬 Contact

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

## 📜 License

MIT License — see the [LICENSE](LICENSE) file for details.

UniFlow

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README