Pi-Ensemble

Animation: Markov walk on Pi-Ensemble generated structures for AtSWEET13.

Pi-Ensemble is a modular framework for generating protein sequence/structure interpolations between two template structures. The package combines:

• a structure predictor that maps sequence -> structure
• an inverse folding model that maps structure -> sequence probabilities
• an interpolation algorithm that mixes information between two templates

The code is organized under pie/, with the main entrypoint in pie/run_pie.py. Scientific context for the method will be linked here once the associated manuscript and data resources are public.

Table of contents

• Implemented Models
  • Structure Prediction
  • Inverse Folding
  • Interpolation Algorithms
• Usage
  • Installation
  • Prediction
  • Output Interpretation
  • Visualization
• Citations
• License

Implemented Models

Pi-Ensemble runs pre-trained models for structure prediction and sequence design. The implemented models and sample configuration files are listed in the tables below.

Structure Prediction

Registry name	Class	Example config	Model reference
esm3	ESM3Predictor	examples/esm3.yaml	ESM3
boltz	BoltzPredictor	examples/boltz.yaml	Boltz
bioemu	BioEmuPredictor	examples/bioemu.yaml	BioEmu
ESMFold2	ESMFold2Predictor	examples/esmfold2.yaml	ESMFold2

Inverse Folding

Registry name	Class	Example config	Model reference
proteinmpnn	ProteinMPNNPredictor	examples/proteinmpnn.yaml	ProteinMPNN

Interpolation Algorithms

Registry name	Class	Example config
serial	SerialInterpolation	examples/serial.yaml
batch	BatchInterpolation	examples/batch.yaml

Usage

Please note that installation and usage configuration may vary according to the models you wish to use.

Installation

Docker

Docker is the recommended path when you want the tested dependency layout without reproducing the host setup manually. The repository ships a CUDA 12.8 image definition and Compose file:

• Dockerfile
• docker-compose.yml

The published image is zcorn/pi-ensemble:cuda12.8. Pull it with:

docker pull zcorn/pi-ensemble:cuda12.8

Or build it locally with:

docker compose build PI-Ensemble

Run Pi-Ensemble inside the container with:

docker compose run --rm PI-Ensemble run_pie test/48G7g/esmfold2_proteinmpnn_serial.yaml

The Compose setup mounts the repository at /workspace, keeps Hugging Face downloads in the host cache, persists model caches in a named Docker volume, and reserves all available NVIDIA GPUs. The main pie environment uses Python 3.12 for ESMFold2 support. BioEmu is installed in a separate Python 3.11 bioemu Conda environment inside the image to avoid dependency conflicts with Boltz, and cg2all is kept in its own Python 3.11 environment.

If you use ESM3, authenticate with Hugging Face before the first run because the model repository is gated:

huggingface-cli login

Host Installation

For host installation, use Conda. Pi-Ensemble uses separate environments because Boltz and BioEmu require incompatible dependency sets.

Create the main Pi-Ensemble environment. This environment uses Python 3.12 because Biohub ESMFold2 requires it:

conda env create -f environment.yml
conda activate pie

Clone ProteinMPNN at a known location:

git clone https://github.com/dauparas/ProteinMPNN.git /opt/ProteinMPNN

If /opt is not writable on your machine, clone elsewhere and set sequence_prediction.kwargs.pmpnn_path in your YAML config to that directory.

Create the optional BioEmu environment when you need structure_prediction.model: "bioemu":

conda env create -f environment-bioemu.yml

Create the optional cg2all environment when you enable cg2all: true:

conda env create -f environment-cg2all.yml

The provided cg2all environment is CPU-only. Pi-Ensemble always runs cg2all with --device cpu.

The host environment files mirror the container layout:

• environment.yml: main Python 3.12 pie environment with Boltz, Biohub ESM/ESMFold2, OpenMM, and the Pi-Ensemble package
• environment-bioemu.yml: separate Python 3.11 bioemu environment
• environment-cg2all.yml: separate Python 3.11 cg2all environment

ESMFold2 dependencies are installed from Biohub's GitHub ESM package rather than the generic PyPI esm package.

When running BioEmu from a host install, the default configuration assumes the auxiliary environment is named bioemu. Override bioemu_environment in the model kwargs only if you use another name.

Prediction

Pi-Ensemble is executed from a YAML configuration file through pie/run_pie.py. A run consists of three blocks:

• structure_prediction: selects the sequence -> structure model
• sequence_prediction: selects the structure -> sequence model
• interpolation: selects the interpolation algorithm and its settings

After installing the package, run:

run_pie path/to/config.yaml

Example:

run_pie test/48G7g/esm3_proteinmpnn_serial.yaml

The config schema is:

structure_prediction:
  model: "<registry name>"
  kwargs:
    ...

sequence_prediction:
  model: "<registry name>"
  kwargs:
    ...

interpolation:
  name: "<registry name>"
  kwargs:
    ...

The registry names currently available are found in the tables above and in the following files:

• pie/structure/registry.py
• pie/sequence/registry.py
• pie/algorithms/registry.py

Output Interpretation

Outputs are written under the outpath specified in the interpolation config.

For the serial algorithm, outputs are organized by interpolation weight, direction, and round:

outpath/
  round_0/
  weight_<value>/
    direction_A/
      round_1/
      round_2/
      ...
    direction_B/
      round_1/
      round_2/
      ...
  log.json

For the batch algorithm, outputs are organized by round, direction, and generated sequence index:

outpath/
  round_0/
  round_1/
    direction_A/
      sequence_000/
      sequence_001/
      ...
    direction_B/
      sequence_000/
      sequence_001/
      ...
  round_2/
  ...
  log.json

Common output artifacts include:

• predicted structure files such as structure.pdb
• model-specific auxiliary files, for example confidence JSON files
• inverse-folding outputs such as seqs/*.fa and probs/*.npz
• log.json, a lightweight session summary of the run

If enabled by the interpolation config, postprocessing outputs are also written at the end of the run:

• cg2all/ for backbone-to-all-atom conversion
• minimized/ for OpenMM-minimized structures

OpenMM minimization uses minimize_forcefield, which defaults to charmm36_2024.xml. You can override it with any OpenMM-compatible forcefield XML name, or with a list when multiple XML files are needed:

interpolation:
  kwargs:
    minimize: true
    minimize_forcefield: "amber14-all.xml"

interpolation:
  kwargs:
    minimize: true
    minimize_forcefield:
      - "amber14/protein.ff14SB.xml"
      - "amber14/tip3pfb.xml"

Visualization

Visualization scripts will be shared soon.

Citations

If you use Pi-Ensemble, please cite its associated publication as well as the models called by your config choices.

@article {Nadeem2026.08.12.744498,
	author = {Nadeem, Hassan and Kleiman, Diego E. and Zhou, Yuming and Leakey, Andrew D. B. and Shukla, Diwakar},
	title = {Pi-Ensemble: Sequence-guided generation of interpolated protein conformational ensembles},
	elocation-id = {2026.08.12.744498},
	year = {2026},
	doi = {10.64898/2026.08.12.744498},
	publisher = {Cold Spring Harbor Laboratory},
	URL = {https://www.biorxiv.org/content/early/2026/08/15/2026.08.12.744498},
	eprint = {https://www.biorxiv.org/content/early/2026/08/15/2026.08.12.744498.full.pdf},
	journal = {bioRxiv}
}

If your run used...	Also cite...
esm3	ESM3
esmfold2	ESMFold2
boltz	Boltz
bioemu	BioEmu
proteinmpnn	ProteinMPNN
cg2all	cg2all
minimize	OpenMM

License

Pi-Ensemble itself is distributed under the MIT License. Model outputs and pretrained weights may be subject to additional third-party licenses or usage restrictions, depending on which options you choose. Users are responsible for obtaining any required model weights and complying with the corresponding terms of use.

Component	License / source note
ESM3/ESMFold2	MIT License
Boltz	MIT License
BioEmu	MIT License
ProteinMPNN	MIT License
cg2all	Apache License 2.0
OpenMM	Various open licenses

Generative AI Acknowledgement

This repository was built with input from Codex.

Parent Repository

This repository was developed from the parent repository:
hnadeem2/Predicting_Interpolated_Ensemble