UniFlow

KO | EN

gitlite — search

#hacktoberfest #react #typescript #golang #android #machine-learning #linux #windows #cli #kubernetes #pytorch #security #ios #mcp

(★ 23)

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

UniFlow Latest Version Download

Download Latest Version (.zip)

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

UniFlow Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

UniFlow

UniFlow Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README