validate.py

__all__ = ["create_validation_plots"]

import os
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['agg.path.chunksize'] = 10000
import seaborn as sns
from scipy.stats import probplot
import numpy as np
from jax.scipy.stats import norm, uniform
import jax.numpy as jnp
from scipy.stats import ks_2samp
colors = sns.color_palette()

def setup_matplotlib_config(font_size):
    """Configure matplotlib font sizes.
    
    Args:
        font_size: Base font size for all plot elements
    """
    plt.rcParams.update({
        'font.size': font_size,
        'axes.titlesize': font_size + 4,
        'axes.labelsize': font_size + 4, 
        'xtick.labelsize': font_size + 2,
        'ytick.labelsize': font_size + 2,
        'legend.fontsize': font_size + 2,
    })

def create_validation_plots(output_folder, times, obs, samples, fontsize=15):
    """Create validation plots for model predictions.
    
    Args:
        output_folder: Directory to save plots
        times: Array of time points
        obs: Observed data
        samples: Model samples (shape: time x n_samples)
        fontsize: Base font size for plots
    """
    setup_matplotlib_config(fontsize)
    dt = times[1] - times[0]

    samples_mean = jnp.mean(samples, axis=1)
    samples_median = jnp.median(samples, axis=1)
    samples_std = jnp.std(samples, axis=1)
    samples_lower = jnp.quantile(samples, 0.025, axis=1)
    samples_upper = jnp.quantile(samples, 0.975, axis=1)

    # Number of samples to plot
    num_samples = 10
    sample_paths = samples[:, :num_samples]

    # Create a parity plot
    fig, ax = plt.subplots(figsize=(6, 5))
    ax.plot(samples_mean, obs, 'o', markersize=2) 
    yy = jnp.linspace(jnp.min(obs), jnp.max(obs))
    ax.plot(yy, yy, 'r', lw=1)
    ax.set_xlabel("Predictions", fontsize=fontsize)
    ax.set_ylabel("Observations", fontsize=fontsize)
    ax.set_title("Parity Plot", fontsize=fontsize)
    ax.tick_params(axis='both', which='major', labelsize=fontsize-2)
    sns.despine(trim=True)
    out_file = os.path.join(output_folder, "parity.pdf")
    # print(f"> writing '{out_file}'")
    fig.savefig(out_file, dpi=300)

    cum_samples = jnp.cumsum(samples*dt, axis=0)
    cum_obs = jnp.cumsum(obs*dt)
    quantiles = jnp.quantile(cum_samples, jnp.array([0.025, 0.5, 0.975]), axis=1)
    cum_low, cum_median, cum_high = quantiles

    # Create a cumulative response plot over time
    fig, ax = plt.subplots(figsize=(6, 5))
    ax.plot(times, cum_obs, color=colors[0], lw=0.5, label="Measured Data")
    ax.plot(times, cum_median, '--', color=colors[1], lw=1, label="Median")
    ax.fill_between(times, cum_low, cum_high, color=colors[1], alpha=0.2, label="95% CI")
    ax.set_xlabel("Time", fontsize=fontsize)
    ax.set_ylabel("Response", fontsize=fontsize)
    # ax.set_title("Cumulative Response Plot", fontsize=fontsize)
    ax.tick_params(axis='both', which='major', labelsize=fontsize-2)
    plt.legend(loc='lower right', frameon=False, fontsize=fontsize)
    sns.despine(trim=True)
    out_file = os.path.join(output_folder, "cumulative_response.pdf")
    # print(f"> writing '{out_file}'")
    fig.savefig(out_file, dpi=300)


    ## PROBABILISTIC METRICS
    u_i = norm.cdf((obs - samples_mean) / samples_std)
    
    # Keep only the non-nan values
    u_i = u_i[~jnp.isnan(u_i)]

    # Create a historgram of u_i for the CDF
    fig, ax = plt.subplots(figsize=(6, 5))
    mean_error = jnp.mean(u_i)
    median_error = jnp.median(u_i)
    fig, ax = plt.subplots()
    ax.hist(u_i, bins=30, edgecolor='black')
    ax.axvline(mean_error, color='r', linestyle='dotted', linewidth=2, label='Mean')  
    ax.axvline(median_error, color='k', linestyle='dotted', linewidth=2, label='Median')  
    ax.legend() 
    ax.set_xlabel("Error", fontsize=fontsize)
    ax.set_ylabel("Frequency", fontsize=fontsize)
    ax.set_title(f"Model - Mean: {mean_error:.2f}, Median: {median_error:.2f}", fontsize=fontsize)
    ax.tick_params(axis='both', which='major', labelsize=fontsize-2)
    sns.despine(trim=True)
    out_file = os.path.join(output_folder, "cdf_hist.pdf")
    # print(f"> writing '{out_file}'")
    fig.savefig(out_file, dpi=300)

    # Create q-q plot for CDF of model
    fig, ax = plt.subplots(figsize=(6, 6))
    probplot(u_i, dist="uniform", sparams=(0, 1), plot=ax, fit=False)
    ax.plot(jnp.linspace(0, 1, jnp.shape(u_i)[0]), jnp.linspace(0, 1, jnp.shape(u_i)[0]), 'r', lw=1)
    ax.set_xlabel("Theoretical Quantiles")
    ax.set_ylabel("Model Quantiles")
    # set maximum number of ticks
    ax.set_xticks(jnp.linspace(0, 1, 3))
    ax.set_yticks(jnp.linspace(0, 1, 3))
    # Remove the default title
    ax.set_title(None)
    # ax.set_title("CDF Q-Q Plot", fontsize=fontsize)
    out_file = os.path.join(output_folder, "cdf_qq.pdf")
    sns.despine(trim=True)
    # print(f"> writing '{out_file}'")
    fig.savefig(out_file, dpi=400)

    # KS test
    x_values = jnp.linspace(0, 1, jnp.shape(u_i)[0])
    u_hat = uniform.cdf(x_values, loc=0, scale=1)
    try:
        ks_stat, p_value = ks_2samp(u_i, np.array(u_hat))
        print(f"KS Statistic: {ks_stat:.4f}, p-value: {p_value:.4f}")
    except ValueError as e:
        print(f"KS test failed: {e}")