vampnet.py

"""
Kuntal Ghosh
Early 2026

vampnet_compute.py
==================
VAMPNet-based conformational analysis of
CA hexamer assemblies, comparing +CPSF6 vs -CPSF6 conditions.
(I have two systems: one with 6 NUP153 chains, the other with 3 NUP153 and 3 CPSF6 chains)

This script:
  1. Loads 6 trajectories (3 +CPSF6 + 3 -CPSF6)
  2. Kabsch-aligns CA atoms per frame against a reference (analogous to aligning on VMD)
  3. Computes inter-hexamer plane-normal angles, folds to acute, encodes as (sin, cos)
  4. Trains a VAMPNet ensemble (N_ENSEMBLE members), selects best by val VAMP-2 score
  5. Projects all features through best lobe -> soft state probabilities
  6. Filters unpopulated states (active-state set), relabels A, B, C, ...
  7. Computes per-frame curvature (1/R, sphere-fit to all CAs): (am looking into this: might give some insights into the handover mechanism)
  8. Computes per-frame R_g of contact-residue cluster
  9. All plotting done separately (a little annoying to have to run this multiple times for plotting changes)

"""

import warnings
warnings.filterwarnings("ignore")

from itertools import combinations

import numpy as np
import MDAnalysis as mda
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from scipy.stats import chi2_contingency
from scipy.optimize import least_squares
from scipy.ndimage import gaussian_filter
from deeptime.decomposition.deep import VAMPNet
from deeptime.util.data import TrajectoryDataset


# ============================================================================
# Configuration
# ============================================================================
CPSF6_TOP  = "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_cpsf6_nup153/CPSF6_fresh/ca_cpsf6_nup153.pdb"
NUP153_TOP = "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_nup153/main_run/first_frame.pdb"

# Note that all these trajectories comprise a hexamer-of-hexamer patch with the NUP153 and CPSF6 chains
# Refer to my slides for details on positioning the IDP chains

cpsf6_trajs = [
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_cpsf6_nup153/CPSF6_fresh/replica_0/traj.xtc",
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_cpsf6_nup153/CPSF6_fresh/replica_1/traj.xtc",
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_cpsf6_nup153/CPSF6_fresh/replica_2/traj.xtc",
]
nup153_trajs = [
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_nup153/main_run/traj.xtc",
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_nup153/replica_1/traj.xtc",
    "/project2/gavoth/kuntalg/HIV_NUP/mesh/ca_nup153/replica_3/traj.xtc",
]

REFERENCE_PDB = CPSF6_TOP   # reference structure for Kabsch alignment (any toplogy works, as the hexamer patch base is same)

# Time / VAMPNet hyperparameters
FRAME_TIME   = 1.0     # ns per frame
LAG          = 20      # VAMPNet lag in frames
BATCH_SIZE   = 512
N_EPOCHS     = 500
LR           = 3e-4
L2           = 1e-3
N_STATES     = 5       # raw output dim; active states filtered later (this seems to be the best)
N_ENSEMBLE   = 3
SEED         = 36
MIN_FRAMES   = 10      # raw states with fewer combined frames are dropped

# Hexamer chain definitions (6 hexamers x 6 chains = 36 chains)
hexamer1 = ['A', 'B', 'C', 'D', 'E', 'F']
hexamer2 = ['G', 'H', 'I', 'J', 'K', 'L']
hexamer3 = ['M', 'N', 'O', 'P', 'Q', 'R']
hexamer4 = ['S', 'T', 'U', 'V', 'W', 'X']
hexamer5 = ['Y', 'Z', 'a', 'b', 'c', 'd']
hexamer6 = ['e', 'f', 'g', 'h', 'i', 'j']
all_hexamers       = [hexamer1, hexamer2, hexamer3, hexamer4, hexamer5, hexamer6]
ALL_HEXAMER_CHAINS = [ch for hx in all_hexamers for ch in hx]

# Inter-hexamer pairs to compute angles between (7 pairs -> 14 features)
hexamer_pairs = [
    (hexamer1, hexamer2),
    (hexamer1, hexamer3),
    (hexamer2, hexamer4),
    (hexamer3, hexamer4),
    (hexamer3, hexamer5),
    (hexamer4, hexamer6),
    (hexamer5, hexamer6),
]
N_PAIRS = len(hexamer_pairs)

# Selection for per-frame R_g (contact-residue cluster across 3 hexamers)
#RG_SELECTION = "(chainID Q or chainID Y or chainID g) and resid 212 and name CA"
RG_SELECTION = "(chainID Q or chainID Y or chainID g) and resid 212 213 and name CA"
# Note: this is meant to reflect the interactions between the RRR residues and GLU 212/213
# Looks pretty cool!

# Heatmap (2D VAMPNet pairwise) parameters
HEATMAP_NBINS = 100
HEATMAP_SIGMA = 1.5

# Per-state 1D histogram parameters (population-weighted)
HIST_NBINS     = 40
TILT_HIST_XMIN = 0.0
TILT_HIST_XMAX = 90.0

torch.manual_seed(SEED)
np.random.seed(SEED)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {DEVICE}")


# ============================================================================
# Build alignment selection
# ============================================================================
_chain_sel = " or ".join(f"segid {ch}" for ch in ALL_HEXAMER_CHAINS)
ALIGN_SEL  = f"name CA and ({_chain_sel})"


# ============================================================================
# Geometry utilities
# ============================================================================

# Each hexamer is roughly planar: so fitting a plane, getting the normal vector
# to that plane and then computing the angles between them

def fit_plane_normal(points):
    """SVD plane fit to (6,3) hexamer COMs -> unit normal vector."""
    c = points - points.mean(axis=0)
    _, _, vh = np.linalg.svd(c)
    n = vh[-1]
    return n / np.linalg.norm(n)

# Briggs' paper suggests that the hexamer-hexamer tilt angles would be somewhere
# between 0 and ~45-50 degrees (which I'm getting here)
# This folds the hexamer-hexamer angles to [0,90]
def acute_angle_between_vectors(v1, v2):
    """Acute angle (0 to pi/2) between two vectors. Uses |dot| to fold
    away the arbitrary sign of plane normals."""
    v1 = v1 / np.linalg.norm(v1)
    v2 = v2 / np.linalg.norm(v2)
    return float(np.arccos(np.clip(np.abs(np.dot(v1, v2)), 0.0, 1.0)))

# Tested and verified: yields the same thing as aligning on VMD
def kabsch_rotation(mobile_c, ref_c):
    """3x3 rotation minimising RMSD between two centered point sets."""
    H        = mobile_c.T @ ref_c
    U, _, Vt = np.linalg.svd(H)
    d        = np.linalg.det(Vt.T @ U.T)
    D        = np.diag([1.0, 1.0, d])
    return Vt.T @ D @ U.T


def _sphere_residuals(params, coords):
    cx, cy, cz, R = params
    dists = np.sqrt((coords[:, 0] - cx)**2 +
                    (coords[:, 1] - cy)**2 +
                    (coords[:, 2] - cz)**2)
    return dists - R


# Okay, so this is to fit a sphere to the WHOLE patch (not just one hexamer)
def fit_sphere_radius(coords):
    centroid = coords.mean(axis=0)
    R0       = np.linalg.norm(coords - centroid, axis=1).mean()
    x0       = np.array([centroid[0], centroid[1], centroid[2], R0])
    result   = least_squares(_sphere_residuals, x0, args=(coords,), method='lm')
    return abs(result.x[3])


# ============================================================================
# Load reference
# ============================================================================
ref_u        = mda.Universe(REFERENCE_PDB)
ref_ag       = ref_u.select_atoms(ALIGN_SEL)
ref_com      = ref_ag.positions.mean(axis=0).copy()
ref_centered = (ref_ag.positions - ref_com).copy()

# ============================================================================
# Build (topology, trajectory) list
# ============================================================================
traj_paths = (
    [(CPSF6_TOP,  t) for t in cpsf6_trajs] +
    [(NUP153_TOP, t) for t in nup153_trajs]
)

# ============================================================================
# Feature extraction: acute-folded inter-hexamer angles -> (sin, cos)
# Also returns aligned CAs per trajectory.
# ============================================================================
def compute_features_and_geometry(traj_path):
    topology, trajectory = traj_path
    print(f"  Loading: {trajectory}")
    u = mda.Universe(topology, trajectory)

    segids  = {seg.segid for seg in u.segments}
    missing = [ch for ch in ALL_HEXAMER_CHAINS if ch not in segids]
    if missing:
        raise ValueError(f"Chains {missing} not found in {trajectory}. "
                         f"Available segids: {sorted(segids)}")

    n_frames  = len(u.trajectory)
    mobile_ag = u.select_atoms(ALIGN_SEL)
    if len(mobile_ag) != len(ref_ag):
        raise ValueError(f"Atom count mismatch: mobile={len(mobile_ag)} "
                         f"reference={len(ref_ag)}")

    mobile_index_map = {global_idx: local_idx
                        for local_idx, global_idx in enumerate(mobile_ag.indices)}
    chain_indices_local = [
        [np.array([mobile_index_map[i]
                   for i in u.select_atoms(f"segid {ch} and name CA").indices])
         for ch in hx]
        for hx in all_hexamers
    ]
    chain_masses = [
        [u.select_atoms(f"segid {ch} and name CA").masses for ch in hx]
        for hx in all_hexamers
    ]

    hexamer_chain_coms = np.zeros((n_frames, len(all_hexamers), 6, 3))
    aligned_ca_all     = np.zeros((n_frames, len(mobile_ag), 3), dtype=np.float32)

    for i, _ in enumerate(u.trajectory):
        mob_pos = mobile_ag.positions.copy()
        mob_com = mob_pos.mean(axis=0)
        mob_c   = mob_pos - mob_com
        R       = kabsch_rotation(mob_c, ref_centered)
        aligned = mob_c @ R.T
        aligned_ca_all[i] = aligned

        for h, (hx_idx, hx_masses) in enumerate(
                zip(chain_indices_local, chain_masses)):
            for c, (idx, masses) in enumerate(zip(hx_idx, hx_masses)):
                pos_c = aligned[idx]
                hexamer_chain_coms[i, h, c] = (
                    (pos_c * masses[:, None]).sum(axis=0) / masses.sum()
                )

    normals = np.zeros((n_frames, len(all_hexamers), 3))
    for i in range(n_frames):
        for h in range(len(all_hexamers)):
            normals[i, h] = fit_plane_normal(hexamer_chain_coms[i, h])

#   Each angle contributes two features: sin and consine of the angle
    feat_blocks = []
    for hx1, hx2 in hexamer_pairs:
        idx1 = all_hexamers.index(hx1)
        idx2 = all_hexamers.index(hx2)
        ang  = np.array([
            acute_angle_between_vectors(normals[i, idx1], normals[i, idx2])
            for i in range(n_frames)
        ])
        feat_blocks.append(np.column_stack([np.sin(ang), np.cos(ang)]))
    features = np.hstack(feat_blocks)

    print(f"    Frames: {n_frames}, features shape: {features.shape}")
    return features, aligned_ca_all


print("\n" + "=" * 60)
print("Feature extraction (Kabsch alignment + acute-folded angles)")
print("=" * 60)
all_features_list = []
aligned_ca_list   = []
for p in traj_paths:
    feats, aligned = compute_features_and_geometry(p)
    all_features_list.append(feats)
    aligned_ca_list.append(aligned)

all_features_f32 = [f.astype(np.float32) for f in all_features_list]
INPUT_DIM        = all_features_f32[0].shape[1]
print(f"\nInput dimensionality: {INPUT_DIM}")


# ============================================================================
# Dataset construction
# ============================================================================
dataset = TrajectoryDataset.from_trajectories(lagtime=LAG, data=all_features_f32)
# This is what constructs the {x_t, x_(t+tau)} pairs: critical for VAMP
n_total = len(dataset)
n_train = int(0.8 * n_total)
n_val   = n_total - n_train
train_data, val_data = torch.utils.data.random_split(
    dataset, [n_train, n_val],
    generator=torch.Generator().manual_seed(SEED)
)
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE,
                          shuffle=True, drop_last=True)
val_loader   = DataLoader(val_data,   batch_size=BATCH_SIZE,
                          shuffle=False, drop_last=False)
val_x0 = torch.cat([b[0] for b in val_loader], dim=0).to(DEVICE)
val_x1 = torch.cat([b[1] for b in val_loader], dim=0).to(DEVICE)
print(f"Dataset: {n_total} pairs | Train: {n_train} | Val: {n_val}")


# ============================================================================
# Architecture
# ============================================================================
class VAMPEncoder(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, output_dim),
            nn.Softmax(dim=-1),
        )

    def forward(self, x):
        return self.network(x)


# ============================================================================
# Ensemble training
# ============================================================================
print(f"\n{'='*60}")
print(f"VAMPNet ensemble training | device={DEVICE} | "
      f"states={N_STATES} | ensemble={N_ENSEMBLE}")
print(f"{'='*60}")

trained_lobes    = []
all_train_curves = []
all_val_curves   = []
best_val_scores  = []

# VAMPNet maximizes the VAMP-2 score
for member in range(N_ENSEMBLE):
    print(f"\n  [Ensemble member {member+1}/{N_ENSEMBLE}]")
    print(f"  {'Epoch':>6}  {'Train VAMP-2':>12}  {'Val VAMP-2':>10}  {'LR':>10}")
    print(f"  {'-'*44}")

    lobe    = VAMPEncoder(INPUT_DIM, N_STATES).to(DEVICE)
    vampnet = VAMPNet(lobe=lobe, optimizer="Adam",
                      learning_rate=LR, device=DEVICE, score_method="VAMP2")
    for pg in vampnet.optimizer.param_groups:
        pg["weight_decay"] = L2

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        vampnet.optimizer, mode="max", patience=20, factor=0.5, min_lr=1e-6,
    )

    best_val   = -np.inf
    best_state = None
    train_eps  = []
    val_eps    = []

    for epoch in range(1, N_EPOCHS + 1):
        vampnet.fit(train_loader)
        tr = float(np.mean(vampnet.train_scores[-len(train_loader):]))
        vl = float(vampnet.validate((val_x0, val_x1)))
        train_eps.append(tr)
        val_eps.append(vl)

        scheduler.step(vl)
        current_lr = vampnet.optimizer.param_groups[0]["lr"]

        if vl > best_val:
            best_val   = vl
            best_state = {k: v.cpu().clone()
                          for k, v in lobe.state_dict().items()}
        if epoch % 50 == 0 or epoch == 1:
            tag = " <-- best" if vl == best_val else ""
            print(f"  {epoch:>6}  {tr:>12.4f}  {vl:>10.4f}  "
                  f"{current_lr:>10.2e}{tag}")

    trained_lobes.append(best_state)
    all_train_curves.append(train_eps)
    all_val_curves.append(val_eps)
    best_val_scores.append(best_val)
    print(f"  Best val VAMP-2: {best_val:.4f}")

print(f"\nEnsemble training complete.")
print(f"Val VAMP-2 per member: {[f'{s:.4f}' for s in best_val_scores]}")
print(f"Mean +/- std: {np.mean(best_val_scores):.4f} +/- {np.std(best_val_scores):.4f}")
for i, sd in enumerate(trained_lobes):
    torch.save(sd, f"vampnet_lobe_member{i+1}.pt")
print("Lobe weights saved.")


# ============================================================================
# Project all frames using best-scoring lobe
# ============================================================================
best_idx = int(np.argmax(best_val_scores))
print(f"\nBest member: {best_idx+1} (val VAMP-2 = {best_val_scores[best_idx]:.4f})")

best_lobe = VAMPEncoder(INPUT_DIM, N_STATES).to(DEVICE)
best_lobe.load_state_dict(trained_lobes[best_idx])
best_lobe.eval()

vampnet_best = VAMPNet(lobe=best_lobe, device=DEVICE, score_method="VAMP2")
model        = vampnet_best.fetch_model()

all_soft   = [model.transform(feat) for feat in all_features_f32]
all_dtrajs = [np.argmax(p, axis=1) for p in all_soft]

all_state_ids_raw = np.concatenate(all_dtrajs,     axis=0)
ids_cpsf6_raw     = np.concatenate(all_dtrajs[:3], axis=0)
ids_bare_raw      = np.concatenate(all_dtrajs[3:], axis=0)
soft_all_raw      = np.concatenate(all_soft,       axis=0)
soft_cpsf6_raw    = np.concatenate(all_soft[:3],   axis=0)
soft_bare_raw     = np.concatenate(all_soft[3:],   axis=0)

print(f"\nRaw state populations:")
for k in range(N_STATES):
    n = (all_state_ids_raw == k).sum()
    print(f"  Raw state {k}: {n} ({100*n/len(all_state_ids_raw):.1f}%)")


# ============================================================================
# Filter unpopulated states and relabel alphabetically
# ============================================================================
active_raw = [k for k in range(N_STATES)
              if (ids_cpsf6_raw == k).sum() + (ids_bare_raw == k).sum() >= MIN_FRAMES]
dropped    = [k for k in range(N_STATES) if k not in active_raw]
N_ACTIVE   = len(active_raw)
ACTIVE_NAMES = [chr(65 + i) for i in range(N_ACTIVE)]   # "A","B","C",...
print(f"\nActive raw states : {active_raw}")
print(f"Dropped raw states: {dropped}")
print(f"N_ACTIVE = {N_ACTIVE}, relabeled as {ACTIVE_NAMES}")

raw_to_new = {raw: new for new, raw in enumerate(active_raw)}

def remap_dtraj(dtraj):
    new = np.full_like(dtraj, -1)
    for raw, new_idx in raw_to_new.items():
        new[dtraj == raw] = new_idx
    return new

all_dtrajs_active = [remap_dtraj(d) for d in all_dtrajs]

soft_all_active   = soft_all_raw[:,   active_raw]
soft_cpsf6_active = soft_cpsf6_raw[:, active_raw]
soft_bare_active  = soft_bare_raw[:,  active_raw]

ids_all_active   = np.concatenate(all_dtrajs_active,     axis=0)
ids_cpsf6_active = np.concatenate(all_dtrajs_active[:3], axis=0)
ids_bare_active  = np.concatenate(all_dtrajs_active[3:], axis=0)

mask_all_valid   = ids_all_active   >= 0
mask_cpsf6_valid = ids_cpsf6_active >= 0
mask_bare_valid  = ids_bare_active  >= 0

print(f"\nFrames after filtering:")
print(f"  All     : {mask_all_valid.sum()} / {len(mask_all_valid)}")
print(f"  +CPSF6  : {mask_cpsf6_valid.sum()} / {len(mask_cpsf6_valid)}")
print(f"  -CPSF6  : {mask_bare_valid.sum()} / {len(mask_bare_valid)}")


# ============================================================================
# Recover tilt angles per frame (acute-folded; averaged over 7 pairs)
# ============================================================================
def recover_tilt_angles(features):
    angles = np.zeros((features.shape[0], N_PAIRS))
    for k in range(N_PAIRS):
        angles[:, k] = np.degrees(np.arccos(
            np.clip(features[:, 2*k + 1], -1.0, 1.0)
        ))
    return angles

all_tilt_mean = [recover_tilt_angles(f).mean(axis=1) for f in all_features_list]
tilt_all_cat  = np.concatenate(all_tilt_mean,     axis=0)
tilt_cpsf6    = np.concatenate(all_tilt_mean[:3], axis=0)
tilt_bare     = np.concatenate(all_tilt_mean[3:], axis=0)


# ============================================================================
# Curvature per frame (sphere fit to all aligned CAs)
# ============================================================================
print("\n" + "=" * 60)
print("Computing per-frame curvature (1/R)")
print("=" * 60)

# Briggs' reports hexamer-hexamer tilt -> correlated VAMPNet with this/curvature
def curvature_per_frame(aligned_ca):
    n = aligned_ca.shape[0]
    out = np.zeros(n)
    for i in range(n):
        out[i] = 1.0 / fit_sphere_radius(aligned_ca[i])
    return out

curv_per_traj = []
for i, aligned in enumerate(aligned_ca_list):
    print(f"  Replica {i}: {aligned.shape[0]} frames")
    curv_per_traj.append(curvature_per_frame(aligned))

curv_all_cat = np.concatenate(curv_per_traj,     axis=0)
curv_cpsf6   = np.concatenate(curv_per_traj[:3], axis=0)
curv_bare    = np.concatenate(curv_per_traj[3:], axis=0)


# ============================================================================
# Per-frame R_g of contact-residue cluster
# ============================================================================
print("\n" + "=" * 60)
print("Computing per-frame R_g (contact-residue cluster)")
print("=" * 60)

# R_g evolution shows the RRR motif of NUP153 interacting with E212/E213
# Check if VAMPNet reflect similar curvature patterns
# DO NOT choose CA residues other than E212/E213

def rg_per_frame(traj_path):
    topology, trajectory = traj_path
    print(f"  R_g: {trajectory}")
    u = mda.Universe(topology, trajectory)
    ag = u.select_atoms(RG_SELECTION)
    if ag.n_atoms == 0:
        raise ValueError(f"Empty R_g selection in {trajectory}. "
                         f"Selection: {RG_SELECTION}")
    n_frames = len(u.trajectory)
    rg = np.zeros(n_frames)
    for i, _ in enumerate(u.trajectory):
        rg[i] = ag.radius_of_gyration()
    return rg

rg_per_traj = [rg_per_frame(p) for p in traj_paths]
rg_all_cat  = np.concatenate(rg_per_traj,     axis=0)
rg_cpsf6    = np.concatenate(rg_per_traj[:3], axis=0)
rg_bare     = np.concatenate(rg_per_traj[3:], axis=0)


# ============================================================================
# Writers
# ============================================================================
def write_dat(filename, header_lines, columns, fmt="%.6g"):
    if isinstance(columns, list):
        data = np.column_stack(columns)
    else:
        data = columns
    with open(filename, "w") as f:
        for line in header_lines:
            f.write(f"# {line}\n")
        np.savetxt(f, data, fmt=fmt)
    print(f"  wrote {filename}")


# ---------------------------------------------------------------------------
# 1. Training curves (per ensemble member, per epoch)
# ---------------------------------------------------------------------------
train_arr = np.array(all_train_curves)
val_arr   = np.array(all_val_curves)
n_epochs  = train_arr.shape[1]
epoch_arr = np.arange(1, n_epochs + 1)

cols = [epoch_arr]
hdr  = ["lag_time_frames=" + str(LAG),
        f"frame_time_ns={FRAME_TIME}",
        f"n_ensemble={N_ENSEMBLE}",
        f"n_states_raw={N_STATES}",
        f"best_member_index={best_idx+1}",
        f"best_val_vamp2={best_val_scores[best_idx]:.6g}",
        f"per_member_best_val_vamp2: "
        + " ".join(f"{i+1}={s:.6g}" for i, s in enumerate(best_val_scores)),
        "columns: epoch  " +
        "  ".join(f"train_m{i+1}" for i in range(N_ENSEMBLE)) +
        "  " + "  ".join(f"val_m{i+1}"   for i in range(N_ENSEMBLE))]
for i in range(N_ENSEMBLE):
    cols.append(train_arr[i])
for i in range(N_ENSEMBLE):
    cols.append(val_arr[i])

write_dat("vampnet_training_curve.dat", hdr, cols)

# ---------------------------------------------------------------------------
# 2. Pairwise heatmaps: P(state_i) vs P(state_j), all 3 conditions
# ---------------------------------------------------------------------------
pairs = list(combinations(range(N_ACTIVE), 2))

def heatmap_long(soft_data, i, j, xlim, ylim):
    H, xedges, yedges = np.histogram2d(
        soft_data[:, i], soft_data[:, j],
        bins=HEATMAP_NBINS, range=[xlim, ylim],
    )
    H = gaussian_filter(H, sigma=HEATMAP_SIGMA)
    H = H / H.max()
    H = H ** 0.5    # gamma-correct for readability (matches original code)
    xcen = 0.5 * (xedges[:-1] + xedges[1:])
    ycen = 0.5 * (yedges[:-1] + yedges[1:])
    XX, YY = np.meshgrid(xcen, ycen, indexing='ij')
    return XX.ravel(), YY.ravel(), H.ravel()

for (i, j) in pairs:
    pad   = 0.02
    xi_min = max(soft_all_active[:, i].min() - pad, 0.0)
    xi_max = min(soft_all_active[:, i].max() + pad, 1.0)
    xj_min = max(soft_all_active[:, j].min() - pad, 0.0)
    xj_max = min(soft_all_active[:, j].max() + pad, 1.0)
    xlim   = (xi_min, xi_max)
    ylim   = (xj_min, xj_max)

    for tag, soft_data in [("all_trajectories", soft_all_active),
                           ("plus_cpsf6",       soft_cpsf6_active),
                           ("minus_cpsf6",      soft_bare_active)]:
        xc, yc, dens = heatmap_long(soft_data, i, j, xlim, ylim)
        si = ACTIVE_NAMES[i]; sj = ACTIVE_NAMES[j]
        write_dat(
            f"vampnet_heatmap_{si}-{sj}_{tag}.dat",
            [
                f"state_i={si} state_j={sj}",
                f"nbins_x={HEATMAP_NBINS} nbins_y={HEATMAP_NBINS}",
                f"xmin={xlim[0]:.6g} xmax={xlim[1]:.6g} "
                f"ymin={ylim[0]:.6g} ymax={ylim[1]:.6g}",
                f"sigma_smooth={HEATMAP_SIGMA}",
                "gamma=0.5 (sqrt applied for readability)",
                f"condition={tag}",
                "columns: P_state_i  P_state_j  density",
            ],
            [xc, yc, dens],
        )

with open("vampnet_active_states.dat", "w") as f:
    f.write(f"# n_active={N_ACTIVE}\n")
    f.write("# active state names mapped from raw model output index\n")
    f.write("# columns: new_name  raw_index\n")
    for new_idx, raw_idx in enumerate(active_raw):
        f.write(f"{ACTIVE_NAMES[new_idx]}  {raw_idx}\n")
print("  wrote vampnet_active_states.dat")

# ---------------------------------------------------------------------------
# 3. State populations + chi2 test
# ---------------------------------------------------------------------------
valid_c = ids_cpsf6_active >= 0
valid_b = ids_bare_active  >= 0
ids_c_v = ids_cpsf6_active[valid_c]
ids_b_v = ids_bare_active[valid_b]

pop_cpsf6    = np.array([(ids_c_v == i).mean() for i in range(N_ACTIVE)])
pop_bare     = np.array([(ids_b_v == i).mean() for i in range(N_ACTIVE)])
counts_cpsf6 = np.array([(ids_c_v == i).sum()  for i in range(N_ACTIVE)])
counts_bare  = np.array([(ids_b_v == i).sum()  for i in range(N_ACTIVE)])

counts_cpsf6_safe = counts_cpsf6 + 1
counts_bare_safe  = counts_bare  + 1
chi2, pval, dof, _ = chi2_contingency(
    np.vstack([counts_cpsf6_safe, counts_bare_safe])
)
print(f"  chi2={chi2:.3f}, dof={dof}, p={pval:.3e}")

with open("vampnet_state_populations.dat", "w") as f:
    f.write(f"# chi2={chi2:.6g}  dof={dof}  p={pval:.6g}  "
            f"(pseudocount +1 per cell to guard against zeros)\n")
    f.write("# columns: state  pop_plus_cpsf6  pop_minus_cpsf6  "
            "count_plus_cpsf6  count_minus_cpsf6\n")
    for i in range(N_ACTIVE):
        f.write(f"{ACTIVE_NAMES[i]}  {pop_cpsf6[i]:.6g}  {pop_bare[i]:.6g}  "
                f"{counts_cpsf6[i]}  {counts_bare[i]}\n")
print("  wrote vampnet_state_populations.dat")

# ---------------------------------------------------------------------------
# 4. Per-state population-weighted histograms: tilt, curvature, R_g
# ---------------------------------------------------------------------------
def write_per_state_pophist(filename, observable, state_ids,
                            xmin, xmax, n_bins, header_extra):
    """
    For each active state k, compute (density=True) histogram on shared bin grid,
    multiply by population weight (n_k / n_total_valid).
    Sum over all states integrates to 1.
    Only frames with state_ids >= 0 contribute.
    """
    valid_mask = state_ids >= 0
    obs_v      = observable[valid_mask]
    ids_v      = state_ids[valid_mask]
    n_total    = len(ids_v)

    edges       = np.linspace(xmin, xmax, n_bins + 1)
    bin_centers = 0.5 * (edges[:-1] + edges[1:])
    bin_width   = edges[1] - edges[0]

    cols = [bin_centers]
    weights = []
    for k in range(N_ACTIVE):
        mask = (ids_v == k)
        n_k  = mask.sum()
        if n_k > 0 and n_total > 0:
            weight_k = n_k / n_total
            hist, _  = np.histogram(obs_v[mask], bins=edges, density=True)
            hist_w   = hist * weight_k
        else:
            weight_k = 0.0
            hist_w   = np.zeros_like(bin_centers)
        weights.append(weight_k)
        cols.append(hist_w)

    header = header_extra + [
        f"n_bins={n_bins}",
        f"xmin={xmin:.6g} xmax={xmax:.6g} bin_width={bin_width:.6g}",
        f"n_total_valid_frames={n_total}",
        "population_weights: " + " ".join(
            f"{ACTIVE_NAMES[k]}={weights[k]:.6g}" for k in range(N_ACTIVE)
        ),
        "each column 'density_X' integrates to that state's population fraction",
        "columns: bin_center  " +
        "  ".join(f"density_{ACTIVE_NAMES[k]}" for k in range(N_ACTIVE)),
    ]
    write_dat(filename, header, cols)

# Tilt — fixed range [0, 90] after acute fold
write_per_state_pophist(
    "vampnet_tilt_per_state_plus_cpsf6.dat",
    tilt_cpsf6, ids_cpsf6_active,
    TILT_HIST_XMIN, TILT_HIST_XMAX, HIST_NBINS,
    ["observable=mean_tilt_deg (acute)"],
)
write_per_state_pophist(
    "vampnet_tilt_per_state_minus_cpsf6.dat",
    tilt_bare, ids_bare_active,
    TILT_HIST_XMIN, TILT_HIST_XMAX, HIST_NBINS,
    ["observable=mean_tilt_deg (acute)"],
)

# Curvature — shared range
curv_min = min(curv_cpsf6.min(), curv_bare.min())
curv_max = max(curv_cpsf6.max(), curv_bare.max())
curv_pad = 0.02 * (curv_max - curv_min)
curv_xmin, curv_xmax = curv_min - curv_pad, curv_max + curv_pad
write_per_state_pophist(
    "vampnet_curvature_per_state_plus_cpsf6.dat",
    curv_cpsf6, ids_cpsf6_active,
    curv_xmin, curv_xmax, HIST_NBINS,
    ["observable=curvature_inv_angstrom"],
)
write_per_state_pophist(
    "vampnet_curvature_per_state_minus_cpsf6.dat",
    curv_bare, ids_bare_active,
    curv_xmin, curv_xmax, HIST_NBINS,
    ["observable=curvature_inv_angstrom"],
)

# R_g — shared range
rg_min = min(rg_cpsf6.min(), rg_bare.min())
rg_max = max(rg_cpsf6.max(), rg_bare.max())
rg_pad = 0.02 * (rg_max - rg_min)
rg_xmin, rg_xmax = rg_min - rg_pad, rg_max + rg_pad
write_per_state_pophist(
    "vampnet_rg_per_state_plus_cpsf6.dat",
    rg_cpsf6, ids_cpsf6_active,
    rg_xmin, rg_xmax, HIST_NBINS,
    ["observable=Rg_contact_residues_angstrom",
     f"selection={RG_SELECTION}"],
)
write_per_state_pophist(
    "vampnet_rg_per_state_minus_cpsf6.dat",
    rg_bare, ids_bare_active,
    rg_xmin, rg_xmax, HIST_NBINS,
    ["observable=Rg_contact_residues_angstrom",
     f"selection={RG_SELECTION}"],
)

# ---------------------------------------------------------------------------
# 5. Summary table — tilt, curvature, R_g (mean / std / population-weighted)
# ---------------------------------------------------------------------------
def overall_and_pw(obs_c, obs_b, ids_c, ids_b, pop_c, pop_b):
    vc = ids_c >= 0
    vb = ids_b >= 0
    obs_c_v, obs_b_v = obs_c[vc], obs_b[vb]
    ids_c_v, ids_b_v = ids_c[vc], ids_b[vb]
    mean_c, std_c = obs_c_v.mean(), obs_c_v.std()
    mean_b, std_b = obs_b_v.mean(), obs_b_v.std()
    pw_c = sum(pop_c[k] * obs_c_v[ids_c_v == k].mean() for k in range(N_ACTIVE)
               if (ids_c_v == k).any())
    pw_b = sum(pop_b[k] * obs_b_v[ids_b_v == k].mean() for k in range(N_ACTIVE)
               if (ids_b_v == k).any())
    return mean_c, std_c, mean_b, std_b, pw_c, pw_b

stats_lines = []
for name, obs_c, obs_b in [
        ("mean_tilt_deg",        tilt_cpsf6, tilt_bare),
        ("curvature_inv_A",      curv_cpsf6, curv_bare),
        ("Rg_contact_A",         rg_cpsf6,   rg_bare),
    ]:
    mc, sc, mb, sb, pwc, pwb = overall_and_pw(
        obs_c, obs_b, ids_cpsf6_active, ids_bare_active, pop_cpsf6, pop_bare)
    stats_lines.append((name, "mean",              mc,  mb,  mc - mb))
    stats_lines.append((name, "std",               sc,  sb,  sc - sb))
    stats_lines.append((name, "pop_weighted_mean", pwc, pwb, pwc - pwb))

with open("vampnet_summary.dat", "w") as f:
    f.write(f"# n_valid_frames_plus_cpsf6={int(valid_c.sum())}\n")
    f.write(f"# n_valid_frames_minus_cpsf6={int(valid_b.sum())}\n")
    f.write(f"# n_active_states={N_ACTIVE}\n")
    f.write("# columns: observable  metric  plus_cpsf6  minus_cpsf6  "
            "delta(plus-minus)\n")
    for obs, met, c, b, d in stats_lines:
        f.write(f"{obs:24s}  {met:18s}  {c:14.6g}  {b:14.6g}  {d:+14.6g}\n")
print("  wrote vampnet_summary.dat")

print("\n" + "=" * 60)
print(f"Done. N_ACTIVE={N_ACTIVE}, {len(pairs)} state pairs, "
      f"{3*len(pairs) + 10} .dat files written.")