vampnet_plot.py

"""
Kuntal Ghosh
Early 2026

===============
Plot-only script. Reads .dat files produced by vampnet_compute.py and writes
PNG figures to the current directory. Also, these .dat are GNUplot compatible, 
so feel free to use whatever I want

No trajectory I/O, no VAMPNet, no MDAnalysis dependency. Edit plot styling
here and re-run without redoing the heavy computation.
"""

from itertools import combinations

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.colors import LinearSegmentedColormap


state_color_palette = ["#e41a1c", "#377eb8", "#4daf4a", "#ff7f00", "#984ea3"]

# Condition labels (use Unicode minus sign U+2212)
LABEL_ALL   = "All trajectories"
LABEL_PLUS  = "+CPSF6"
LABEL_MINUS = "\u2212CPSF6"

# Condition colors for the population bar chart
COLOR_PLUS  = "steelblue"
COLOR_MINUS = "coral"

# Heatmap colormap (matches VAMP plot script)
HOT_WHITE = LinearSegmentedColormap.from_list(
    "hot_white",
    ["white", "#FFDD00", "#FF6600", "#CC0000", "#660000"],
    N=512,
)


# ============================================================================
# Helpers
# ============================================================================
def parse_header(filename):
    """Return dict of key=value pairs from '#' header lines."""
    meta = {}
    with open(filename) as f:
        for line in f:
            if not line.startswith("#"):
                break
            for tok in line.lstrip("#").strip().split():
                if "=" in tok:
                    k, v = tok.split("=", 1)
                    meta[k] = v
    return meta


def load_active_states(filename="vampnet_active_states.dat"):
    """Returns the list of active state names: ['A', 'B', ...]."""
    names = []
    with open(filename) as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            names.append(line.split()[0])
    return names


def load_training_curves(filename="vampnet_training_curve.dat"):
    meta  = parse_header(filename)
    n_ens = int(meta["n_ensemble"])
    best  = int(meta["best_member_index"])
    data  = np.loadtxt(filename)
    epochs = data[:, 0]
    train  = data[:, 1:1+n_ens]              # columns 1..n_ens
    val    = data[:, 1+n_ens:1+2*n_ens]      # columns n_ens+1..2n_ens
    return epochs, train, val, best, meta


def load_heatmap_long(filename):
    """Long-format heatmap .dat -> (density_2d, xlim, ylim)."""
    meta = parse_header(filename)
    nx   = int(meta["nbins_x"])
    ny   = int(meta["nbins_y"])
    xmin = float(meta["xmin"]); xmax = float(meta["xmax"])
    ymin = float(meta["ymin"]); ymax = float(meta["ymax"])
    data = np.loadtxt(filename)
    density = data[:, 2].reshape(nx, ny)
    return density, (xmin, xmax), (ymin, ymax)


def load_state_populations(filename="vampnet_state_populations.dat"):
#   These state populations should ideally correspond to what I see from
#   the RMSD, R_g and 1/R plots (and visually too)
    names      = []
    pops_plus  = []
    pops_minus = []
    with open(filename) as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            parts = line.split()
            names.append(parts[0])
            pops_plus.append(float(parts[1]))
            pops_minus.append(float(parts[2]))
    return names, np.array(pops_plus), np.array(pops_minus)


def load_per_state_hist(filename):
    """Returns (bin_centers, [hist_state_0, hist_state_1, ...])."""
    data = np.loadtxt(filename)
    bin_centers = data[:, 0]
    hists = [data[:, k+1] for k in range(data.shape[1] - 1)]
    return bin_centers, hists


# ============================================================================
# Pre-flight: load active state list
# ============================================================================
ACTIVE_NAMES  = load_active_states()
N_ACTIVE      = len(ACTIVE_NAMES)
ACTIVE_COLORS = state_color_palette[:N_ACTIVE]
print(f"Detected {N_ACTIVE} active states: {ACTIVE_NAMES}")


# ============================================================================
# Figure 1: Training/validation curves
# ============================================================================
print("Generating vampnet_training_curve.png")
epochs, train_arr, val_arr, best_idx, train_meta = load_training_curves()
n_ens = train_arr.shape[1]

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
for arr, label, color, ax_loss, ax_score in [
    (train_arr, "Train", "steelblue",  axes[0], axes[1]),
    (val_arr,   "Val",   "darkorange", axes[0], axes[1]),
]:
    mean = arr.mean(axis=1)
    std  = arr.std(axis=1)
    for i in range(n_ens):
        ax_loss.plot(epochs, -arr[:, i],  color=color, lw=0.6, alpha=0.3)
        ax_score.plot(epochs,  arr[:, i], color=color, lw=0.6, alpha=0.3)
    ax_loss.plot(epochs, -mean, color=color, lw=2.0, label=f"{label} mean")
    ax_loss.fill_between(epochs, -(mean+std), -(mean-std),
                         color=color, alpha=0.15)
    ax_score.plot(epochs, mean, color=color, lw=2.0, label=f"{label} mean")
    ax_score.fill_between(epochs, mean-std, mean+std,
                          color=color, alpha=0.15)

axes[0].set_xlabel("Epoch", fontsize=13)
axes[0].set_ylabel("Loss (\u2212VAMP2)", fontsize=13)
axes[0].set_title("Training Loss", fontsize=13, fontweight="bold")
axes[0].legend(fontsize=11)
axes[0].grid(linestyle="--", alpha=0.4)

axes[1].set_xlabel("Epoch", fontsize=13)
axes[1].set_ylabel("VAMP-2 score", fontsize=13)
axes[1].set_title("VAMP-2 Score", fontsize=13, fontweight="bold")
axes[1].grid(linestyle="--", alpha=0.4)

best_val_score = val_arr[:, best_idx-1].max()
axes[1].axhline(best_val_score, color="green", lw=1.5, linestyle="--",
                label=f"Best member {best_idx} ({best_val_score:.4f})")
axes[1].legend(fontsize=10)

plt.suptitle(f"VAMPNet Training | ensemble of {n_ens} members",
             fontsize=13, fontweight="bold")
plt.tight_layout()
plt.savefig("vampnet_training_curve.png", dpi=150, bbox_inches="tight")
plt.close(fig)


# ============================================================================
# Figure 2: Pairwise heatmaps P(state_i) vs P(state_j)
# ============================================================================
print("Generating vampnet_heatmaps.png")
pairs = list(combinations(range(N_ACTIVE), 2))
n_pairs = len(pairs)

fig, axes = plt.subplots(n_pairs, 3,
                         figsize=(19, 6 * n_pairs),
                         squeeze=False)

condition_files = [
    ("all_trajectories", LABEL_ALL),
    ("plus_cpsf6",       LABEL_PLUS),
    ("minus_cpsf6",      LABEL_MINUS),
]

for row, (i, j) in enumerate(pairs):
    si, sj = ACTIVE_NAMES[i], ACTIVE_NAMES[j]
    for col, (tag, title) in enumerate(condition_files):
        ax = axes[row, col]
        fn = f"vampnet_heatmap_{si}-{sj}_{tag}.dat"
        density, xlim, ylim = load_heatmap_long(fn)
        im = ax.imshow(
            density.T,
            origin="lower",
            extent=[xlim[0], xlim[1], ylim[0], ylim[1]],
            aspect="auto",
            cmap=HOT_WHITE,
            vmin=0, vmax=1,
        )
        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label("Normalized frequency", fontsize=10)
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.set_xlabel(f"P(State {si})", fontsize=12)
        ax.set_ylabel(f"P(State {sj})", fontsize=12)
        if row == 0:
            ax.set_title(title, fontsize=14, fontweight="bold")
        ax.text(0.02, 0.97, f"{si} vs {sj}",
                transform=ax.transAxes, fontsize=11,
                va="top", ha="left",
                bbox=dict(facecolor="white", edgecolor="gray",
                          alpha=0.7, boxstyle="round,pad=0.2"))

plt.suptitle("VAMPNet Pairwise State Probability Landscapes — Normalized Frequency",
             fontsize=15, fontweight="bold", y=1.005)
plt.tight_layout()
plt.savefig("vampnet_heatmaps.png", dpi=150, bbox_inches="tight")
plt.close(fig)


# ============================================================================
# Figure 3: State populations
# ============================================================================
print("Generating vampnet_state_populations.png")
names, pop_plus, pop_minus = load_state_populations()
fig, ax = plt.subplots(figsize=(8, 6))
x = np.arange(len(names))
w = 0.35
ax.bar(x - w/2, pop_plus,  w, label=LABEL_PLUS,
       color=COLOR_PLUS, edgecolor="k", linewidth=0.7)
ax.bar(x + w/2, pop_minus, w, label=LABEL_MINUS,
       color=COLOR_MINUS, edgecolor="k", linewidth=0.7)
for k, name in enumerate(names):
    ax.text(k - w/2, pop_plus[k]  + 0.005, f"{pop_plus[k]:.2f}",
            ha="center", va="bottom", fontsize=11,
            color=COLOR_PLUS, fontweight="bold")
    ax.text(k + w/2, pop_minus[k] + 0.005, f"{pop_minus[k]:.2f}",
            ha="center", va="bottom", fontsize=11,
            color=COLOR_MINUS, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels(names, fontsize=13)
ax.set_xlabel("State", fontsize=13)
ax.set_ylabel("Population fraction", fontsize=13)
ax.set_title("State Populations", fontsize=13)
ax.legend(fontsize=12, loc="upper right")
ax.set_ylim(0, 1.0)
ax.grid(axis="y", linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("vampnet_state_populations.png", dpi=150, bbox_inches="tight")
plt.close(fig)


# ============================================================================
# Helper for per-state population-weighted histogram figures
# ============================================================================
def per_state_hist_figure(plus_file, minus_file, xlabel, suptitle, png_name,
                          xrange=None):
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    for ax, (fn, title) in zip(axes, [(plus_file,  LABEL_PLUS),
                                      (minus_file, LABEL_MINUS)]):
        bin_centers, hists = load_per_state_hist(fn)
        n_states = len(hists)
        for k in range(n_states):
            ax.plot(bin_centers, hists[k],
                    color=ACTIVE_COLORS[k], linewidth=2.0)
            ax.fill_between(bin_centers, hists[k],
                            color=ACTIVE_COLORS[k], alpha=0.25)
        ax.set_xlabel(xlabel, fontsize=15)
        ax.set_ylabel("Probability density", fontsize=15)
        ax.set_title(title, fontsize=16)
        ax.tick_params(axis='both', labelsize=13)
        if xrange is not None:
            ax.set_xlim(*xrange)
        else:
            ax.set_xlim(bin_centers.min(), bin_centers.max())
        ax.set_ylim(bottom=0)
    legend_elements = [
        mpatches.Patch(facecolor=ACTIVE_COLORS[k], alpha=0.6,
                       label=f"State {ACTIVE_NAMES[k]}")
        for k in range(N_ACTIVE)
    ]
    fig.legend(handles=legend_elements, loc="lower center",
               ncol=N_ACTIVE, fontsize=13, frameon=True,
               bbox_to_anchor=(0.5, -0.06))
    plt.suptitle(suptitle, fontsize=16, y=1.02)
    plt.tight_layout()
    plt.savefig(png_name, dpi=150, bbox_inches="tight")
    plt.close(fig)


# ============================================================================
# Figure 4: Tilt distributions per state
# ============================================================================
print("Generating vampnet_tilt_distributions.png")
per_state_hist_figure(
    "vampnet_tilt_per_state_plus_cpsf6.dat",
    "vampnet_tilt_per_state_minus_cpsf6.dat",
    xlabel="Mean tilt angle (\u00b0)",
    suptitle=f"Tilt Angle Distributions per State\n{LABEL_PLUS} vs {LABEL_MINUS}",
    png_name="vampnet_tilt_distributions.png",
    xrange=(0, 40),
)


# ============================================================================
# Figure 5: Curvature distributions per state
# ============================================================================
print("Generating vampnet_curvature_hist_per_state.png")
per_state_hist_figure(
    "vampnet_curvature_per_state_plus_cpsf6.dat",
    "vampnet_curvature_per_state_minus_cpsf6.dat",
    xlabel="Curvature 1/R (\u00c5\u207b\u00b9)",
    suptitle=f"Curvature (1/R) Distributions per State\n{LABEL_PLUS} vs {LABEL_MINUS}",
    png_name="vampnet_curvature_hist_per_state.png",
)


# ============================================================================
# Figure 6: R_g distributions per state
# ============================================================================
print("Generating vampnet_rg_hist_per_state.png")
per_state_hist_figure(
    "vampnet_rg_per_state_plus_cpsf6.dat",
    "vampnet_rg_per_state_minus_cpsf6.dat",
    xlabel="$R_g$ (\u00c5)",
    suptitle=f"$R_g$ Distributions per State\n{LABEL_PLUS} vs {LABEL_MINUS}",
    png_name="vampnet_rg_hist_per_state.png",
)