Sequence context decodes multi-state conformational heterogeneity from crystallographic B-factors

This repository contains the code for the B-Factor Corrector (BFC), a fine-tuned protein language model that treats the recovery of conformational heterogeneity as a sequence-context-aware translation task. By leveraging deep contextual embeddings, BFC filters out crystal-induced rigidity to reveal the multi-state conformational signal underlying raw crystallographic B-factors.

Overview

X-ray crystallography provides the majority of protein structures, yet the B-factors associated with these coordinates are heavily influenced by crystal packing and refinement protocols, limiting their use for quantifying solution-state dynamics. BFC is a fine-tuned ESM-2 encoder with a lightweight regression head that takes a residue sequence together with its normalized crystallographic B-factors and recovers the per-residue multi-state conformational heterogeneity (ensemble RMSF) that no single crystal structure captures.

Repository structure

.
├── model/                     # The BFC model: train and predict
│   ├── train.py               # Fine-tune ESM-2 + regression head (the BFC model)
│   ├── train_frozen.py        # Frozen-ESM ablation (only the head is trained)
│   └── predict.py             # Single-PDB inference entry point
│
└── experiments/               # Validation experiments reported in the paper
    ├── atlas_md/               # Comparison against ATLAS MD ensembles
    ├── docking/                # Flexible-residue docking on PDBbind
    ├── antibody/               # Antibody interface refinement in OpenMM
    └── sae/                    # Sparse-autoencoder analysis

Only the core code is included here. Figure-generation scripts, intermediate data products, and exploratory analyses are omitted; the experiment scripts produce the underlying numbers, and figures can be regenerated from them.

Model weights

The trained weights are distributed separately via Google Drive:

• BFC (main model, required for prediction and all experiments): bfc.pth
• SAE (sparse autoencoder over the BFC representation, used for the dictionary analysis in Fig. 8; optional): bfc_sae.pth

Download bfc.pth and pass its path to the scripts via --model_path (prediction) or --bfc_weights (experiments).

Installation

BFC depends on PyTorch and Hugging Face Transformers. A minimal environment:

pip install torch transformers biopython numpy tqdm

The ESM-2 backbone (facebook/esm2_t6_8M_UR50D) is downloaded automatically from Hugging Face on first use.

Usage

Predict conformational heterogeneity for a single PDB chain

python model/predict.py \
    --model_path /path/to/bfc_weights.pth \
    --pdb 1AKI.pdb \
    --chain A \
    --output 1AKI_A_prediction.json

The output JSON lists, per residue, the input B-factor and the BFC-corrected per-residue conformational heterogeneity (ensemble-RMSF profile).

Train the model

Prepare your own training and validation sets as JSON-lines files (one chain per line, with the keys Sequence, Norm_B_factor, and Norm_RMSF) following the data description in the paper, then:

python model/train.py \
    --train_file train.jsonl \
    --val_file validation.jsonl \
    --output_dir bfc_checkpoints

The frozen-ESM ablation is trained the same way with model/train_frozen.py.

Experiments

The scripts under experiments/ reproduce the validation analyses. They rely on external resources that are not redistributed here:

• ATLAS MD — molecular-dynamics trajectories from the ATLAS database.
• Docking — the PDBbind refined set and the Uni-Dock engine (install separately).
• Antibody — ESMFold for structure prediction and OpenMM for the refinement MD.

Paths to weights and data are passed as command-line arguments or environment variables in each script.

Google Colab

Run a prediction in the browser (no local setup) with BFC_prediction_colab.ipynb:

The notebook downloads the BFC weights, runs inference on an example crystal structure, and plots the recovered profile.