NEQ ATM Switching for Relative Binding Free Energies

Minimal, standalone reproduction of the nonequilibrium (NEQ) Alchemical Transfer Method (ATM) relative binding free energy workflow. This repository accompanies the publication and demonstrates end-to-end:

1. NEQ switching — equilibrium sampling at endstates and an intermediate, followed by nonequilibrium driving in four directions
2. Free energy estimation — BAR (Bennett Acceptance Ratio) with bootstrap uncertainty and convergence diagnostics
3. Benchmark correlation — computed vs. experimental DDG scatter plot for the TYK2 congeneric series
4. Paper figures — data and scripts to regenerate Figure 2 (multi-system correlation, nonchimeric switch) and Figure 3 (forward/backward NEQ work distributions)

Disclaimer: This is research/demonstration code provided for reproducibility of the accompanying publication. It is provided "as is", without warranty of any kind, and is not intended for production use. See LICENSE and NOTICE.

Security note (pickle files): The switching pipeline reads and writes Python pickle files (e.g. nex_switch_store.pkl). Loading a pickle file executes arbitrary code — only load .pkl files you generated yourself or obtained from a trusted source. See SECURITY.md for details.

Installation

OpenMM and openmmtools are distributed via conda-forge. Create the environment from the provided environment.yml, then install the local nex package with poetry:

micromamba create -f environment.yml
micromamba activate nex-atm
poetry install

Or equivalently with conda:

conda env create -f environment.yml
conda activate nex-atm
poetry install

Example Data

The data/ directory contains inputs for one TYK2 edge (jmc_23 → ejm_55) and benchmark results for all 24 TYK2 edges:

Path	Description
data/tyk2_jmc_23_ejm_55/atm_system.xml	ATM-equipped OpenMM System (~34 MB)
data/tyk2_jmc_23_ejm_55/equilibrated.xml	Equilibrated OpenMM State (~10 MB)
data/tyk2_jmc_23_ejm_55/equilibrated.pdb	Topology PDB (~6 MB)
data/tyk2_jmc_23_ejm_55/nex_switch_store.pkl	Production work data (100 snapshots x 3 replicas)
data/tyk2_results/TYK2_collated.csv	Computed + experimental DDGs for all 24 TYK2 edges

TYK2_collated.csv columns

This file is the input to plot_correlation.py. The columns it uses are:

Column	Description
Ligand_1, Ligand_2	Edge endpoints
calculated_DDG, error	Computed DDG and uncertainty for the edge (a single run, using the production parameters below). This is what plot_correlation.py reads, and it matches the column name written by compute_free_energy.py.
experimental_ddG	Experimental DDG

The file also carries per-replicate data from the paper (calculated_DDG_0/1/2, error_0/1/2, and per-leg bootstrap/overlap diagnostics) for the three independent runs. These are kept for reference; the plot uses the single calculated_DDG/error columns.

Note: the paper's Figure 1 reports the mean over the three independent runs with the standard error of the mean, whereas plot_correlation.py here uses a single run (calculated_DDG), so the correlation metrics differ slightly from the published values.

Running a NEQ Switch

The switching script performs equilibrium warmup, snapshot collection, and nonequilibrium switching for a single edge. It requires a GPU (CUDA) by default.

Demo mode (smoke test, ~minutes on GPU):

python scripts/run_neq_switching.py \
    --system data/tyk2_jmc_23_ejm_55/atm_system.xml \
    --state data/tyk2_jmc_23_ejm_55/equilibrated.xml \
    --topology data/tyk2_jmc_23_ejm_55/equilibrated.pdb \
    --output-dir results/jmc_23_ejm_55 \
    --demo

Demo mode uses reduced parameters (1 snapshot, 5k warmup steps, 5k switching steps) to verify the machinery runs end-to-end. The resulting free energy will not be meaningful.

Production parameters (used to generate published results):

Parameter	Value
warmup_steps	625,000
snapshot_interval	125,000
num_snapshots	100
num_replicas	3
num_steps_neq_switch	50,000

These are the script defaults (without --demo). A production run takes several hours on a modern GPU.

Computing Free Energies

Given a nex_switch_store.pkl from the switching script (or the provided production data), compute the relative binding free energy:

python scripts/compute_free_energy.py \
    --work-pkl data/tyk2_jmc_23_ejm_55/nex_switch_store.pkl \
    --output-dir results/jmc_23_ejm_55 \
    --ligand-1 jmc_23 --ligand-2 ejm_55

This runs BAR on both legs (ligand 1 ↔ intermediate, intermediate ↔ ligand 2), performs bootstrap BAR with a KS normality test, generates work distribution histograms, and writes a CSV with the DDG and diagnostics.

Plotting the Correlation

Generate a computed vs. experimental DDG scatter plot for the full TYK2 benchmark:

python scripts/plot_correlation.py \
    --data data/tyk2_results/TYK2_collated.csv \
    --output correlation_plot.pdf

The plot includes error bars, a y = x reference line, a ±1 kcal/mol shaded band, and annotation of RMSE, MUE, R², Kendall τ, and Spearman ρ.

The input CSV needs calculated_DDG, error, and experimental_ddG columns. compute_free_energy.py writes calculated_DDG and error per edge; the experimental values are not produced by the pipeline, so you add the experimental_ddG column by hand (collating the per-edge outputs and joining experimental data) before plotting.

Reproducing Figure 2 (multi-system correlation, nonchimeric)

Figure 2 compares predicted vs. experimental ΔΔG across three congeneric series — CDK2, TYK2, JNK1 — computed with the nonchimeric (direct, one-step) NEQ ATM switch (skipping the alchemical intermediate). The per-edge data for each system lives in data/fig2/:

Path	Description
data/fig2/CDK2_collated.csv	CDK2 — 25 edges
data/fig2/TYK2_collated.csv	TYK2 — 24 edges
data/fig2/JNK1_collated.csv	JNK1 — 30 edges

Each row holds the raw triplicate predictions (calculated_DDG_0/1/2, error_0/1/2), per-replicate overlap flags (left_isoverlapping_i, right_isoverlapping_i), and experimental_ddG. The plotting script computes every reported statistic directly from that raw data — nothing is pre-baked: the per-edge prediction is the overlap-filtered mean ± SEM over the three replicates, and each panel reports RMSE and Pearson r with a hierarchical bootstrap 95% CI (resampling edges).

python scripts/plot_fig2_correlation.py \
    --data-dir data/fig2 \
    --output fig2_correlation.pdf \
    --seed 0

Expected output (seed 0):

System	RMSE_exp (kcal/mol)	Pearson r (95% CI)	edges
CDK2	1.84 ± 0.12	0.58 [0.41, 0.75]	25
TYK2	0.89 ± 0.05	0.62 [0.39, 0.80]	24
JNK1	0.79 ± 0.08	0.60 [0.31, 0.82]	30

Reproducing Figure 3 (nonequilibrium work distributions)

Figure 3 is a 3×2 grid for a single CDK2 edge (receptor_30_31) contrasting chimeric switching (through the alchemical intermediate, rows 1–2) with the nonchimeric direct switch (row 3). The left column shows protocol-work vs. λ trajectories; the right column shows forward (blue) / backward (orange, sign-flipped) final-work histograms.

The per-edge work data is distilled into one compact archive:

Path	Description
data/fig3/cdk2_receptor_30_31_works.npz	Work-vs-λ trajectories + final works for all six protocols: chimeric (1↔int, int↔2) and nonchimeric (1↔2)

python scripts/plot_fig3_work_distributions.py \
    --npz data/fig3/cdk2_receptor_30_31_works.npz \
    --output fig3_work_distributions.pdf

The chimeric legs give tight, well-overlapping work (~10–100 kcal/mol); the nonchimeric direct switch gives broad, poorly overlapping work (~±1000 kcal/mol) — the motivation for the chimeric protocol.

Data provenance

The Figure 2 CSVs are the per-system collated outputs of the nonchimeric switching + BAR analysis pipeline. The Figure 3 .npz is a distilled form of the full per-edge simulation output (the raw work is ~21 MB across ~2400 small per-snapshot files for this single edge); an offline extraction step reduces it to the trajectories and final works the figure needs, stored raw in kT (the plotting script applies the kT → kcal/mol factor). Work values, replicate structure, and overlap filtering are preserved so the figures are byte-faithful to the published versions. See data/README.md for full data provenance and licensing.

Method Overview

The Alchemical Transfer Method (ATM) computes relative binding free energies by transferring a ligand between two binding poses via a shared intermediate state. The nonequilibrium protocol:

1. Equilibrium sampling at three states: endstate 1 (ligand 1 bound), endstate 2 (ligand 2 bound), and a chimeric intermediate (λ = 0.5)
2. Snapshot collection from each equilibrium ensemble
3. Nonequilibrium switching in four directions:
   • 1 → intermediate (forward positive)
   • intermediate → 1 (backward positive)
   • intermediate → 2 (forward negative)
   • 2 → intermediate (backward negative)
4. BAR estimation of ΔG for each leg from the accumulated protocol work; ΔΔG = ΔG(left leg) + ΔG(right leg)

The AlchemicalNonequilibriumLangevinIntegrator from openmmtools drives the lambda schedule, and protocol work is extracted from the integrator at the end of each switch. The ATMForce in OpenMM handles the alchemical coupling.

Citation

If you use this code, please cite:

@article{pitman2026nonequilibrium,
  title={Nonequilibrium Chimeric Switching (NEX) Stabilizes Binding Free Energy Calculations Across Chemical Space},
  author={Pitman, Mary and Sood, Amogh and Rufa, Dominic and Huntington, Lee},
  journal={ChemRxiv},
  year={2026},
  doi={10.26434/chemrxiv.15001585/v2},
  url={https://doi.org/10.26434/chemrxiv.15001585/v2}
}

Contributing & Security

Contributions are welcome — see CONTRIBUTING.md and our CODE_OF_CONDUCT.md. To report a security issue, see SECURITY.md.

License

Source code is licensed under the Apache License 2.0. See LICENSE and NOTICE.

The benchmark data under data/ is derived from third-party sources and may carry separate licensing and attribution requirements — see data/README.md.