plotting.py

"""Plotting utilities for OU process results."""

import os
import matplotlib.pyplot as plt
import seaborn as sns
import jax.numpy as jnp

# Set up seaborn style
colors = sns.color_palette()
sns.set_context("paper")
sns.set_style("ticks")


def plot_trajectory(times, X, save_path, font_size=16, label=r"$X_t$"):
    """Plot a single trajectory over time.
    
    Args:
        times: Array of time points
        X: Trajectory values
        save_path: Path to save the plot
        font_size: Base font size
        label: Legend label for the trajectory
    """
    plt.figure(figsize=(6, 5))
    plt.plot(times, X, label=label, color=colors[0])
    plt.legend(frameon=False, fontsize=font_size-2, loc='upper left')
    plt.xlabel("Time", fontsize=font_size)
    plt.tick_params(axis='both', which='major', labelsize=font_size-2)
    sns.despine(trim=True)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_U_values(times, U_values, save_path, font_size=16):
    """Plot the hidden variable U_t over time.
    
    Args:
        times: Array of time points
        U_values: Hidden variable values
        save_path: Path to save the plot
        font_size: Base font size
    """
    plt.figure(figsize=(6, 5))
    plt.plot(times, U_values, label=r"$U_t$", color=colors[3])
    plt.legend(frameon=False, fontsize=font_size-2, loc='upper left')
    plt.xlabel("Time (s)", fontsize=font_size)
    plt.ylabel(r"$U_t$", fontsize=font_size)
    plt.tick_params(axis='both', which='major', labelsize=font_size-2)
    sns.despine(trim=True)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_mean_function(times, X, U_values, true_a, true_b, save_path, font_size=16):
    """Plot the true mean function μ(t) = a·U_t + b along with trajectory.
    
    Args:
        times: Array of time points
        X: Trajectory values
        U_values: Hidden variable values
        true_a: True slope parameter
        true_b: True intercept parameter
        save_path: Path to save the plot
        font_size: Base font size
    """
    mu_true = true_a * U_values + true_b
    plt.figure(figsize=(6, 5))
    plt.plot(times, mu_true, label=r"$\mu(t) = a \cdot U_t + b$", color=colors[1])
    plt.plot(times, X, alpha=0.7, color=colors[0], label=r"$X_t$")
    plt.xlabel("Time (s)", fontsize=font_size)
    plt.ylabel(r"$X_t$", fontsize=font_size)
    plt.legend(frameon=False, fontsize=font_size-2, loc='upper left')
    plt.tick_params(axis='both', which='major', labelsize=font_size-2)
    sns.despine(trim=True)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_training_loss(losses, save_path, font_size=16):
    """Plot training loss over epochs.
    
    Args:
        losses: List of loss values per epoch
        save_path: Path to save the plot
        font_size: Base font size
    """
    plt.figure()
    plt.plot(losses, label="Predicted")
    plt.xlabel("Epoch", fontsize=font_size)
    plt.ylabel("Loss", fontsize=font_size)
    plt.title("Training Loss", fontsize=font_size+2)
    plt.legend(frameon=False, fontsize=font_size-2)
    plt.xticks(fontsize=font_size-2)
    plt.yticks(fontsize=font_size-2)
    sns.despine(trim=True)
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_parameter_evolution(params_history, true_params, save_path, font_size=16):
    """Plot parameter convergence during training.
    
    Shows how each parameter approaches its true value over epochs.
    
    Args:
        params_history: Parameter values at each epoch (shape: n_epochs x 4)
        true_params: True parameter values (lambda, a, b, sigma)
        save_path: Path to save the plot
        font_size: Base font size
    """
    params_array = jnp.array(params_history)
    true_lambda, true_a, true_b, true_sigma = true_params
    
    lam_hist = jnp.exp(params_array[:, 0])
    a_hist = params_array[:, 1]
    b_hist = params_array[:, 2]
    sigma_hist = jnp.exp(params_array[:, 3])
    
    fig, axes = plt.subplots(4, 1, figsize=(5, 8), sharex=True)
    
    axes[0].plot(lam_hist, label="Predicted")
    axes[0].axhline(y=true_lambda, color='r', linestyle='--', label="True Value")
    axes[0].set_ylabel(r"$\lambda$", fontsize=font_size)
    axes[0].set_title("Mean Reversion Rate Evolution", fontsize=font_size+2)
    axes[0].legend(frameon=False, fontsize=font_size-2)
    axes[0].tick_params(labelsize=font_size-2)
    
    axes[1].plot(a_hist, label="Predicted")
    axes[1].axhline(y=true_a, color='r', linestyle='--', label="True Value")
    axes[1].set_ylabel(r"$a$", fontsize=font_size)
    axes[1].set_title("Mean Function Slope Evolution", fontsize=font_size+2)
    axes[1].legend(frameon=False, fontsize=font_size-2)
    axes[1].tick_params(labelsize=font_size-2)
    
    axes[2].plot(b_hist, label="Predicted")
    axes[2].axhline(y=true_b, color='r', linestyle='--', label="True Value")
    axes[2].set_ylabel(r"$b$", fontsize=font_size)
    axes[2].set_title("Mean Function Intercept Evolution", fontsize=font_size+2)
    axes[2].legend(frameon=False, fontsize=font_size-2)
    axes[2].tick_params(labelsize=font_size-2)
    
    axes[3].plot(sigma_hist, label="Predicted")
    axes[3].axhline(y=true_sigma, color='r', linestyle='--', label="True Value")
    axes[3].set_xlabel("Epoch", fontsize=font_size)
    axes[3].set_ylabel(r"$\sigma$", fontsize=font_size)
    axes[3].set_title("Volatility Evolution", fontsize=font_size+2)
    axes[3].legend(frameon=False, fontsize=font_size-2)
    axes[3].tick_params(labelsize=font_size-2)
    
    plt.tight_layout()
    sns.despine(trim=True)
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_estimated_mean_function(times, U_values, true_a, true_b, est_a, est_b, 
                                  save_path, font_size=16):
    """Compare true and estimated mean functions.
    
    Args:
        times: Array of time points
        U_values: Hidden variable values
        true_a: True slope parameter
        true_b: True intercept parameter
        est_a: Estimated slope parameter
        est_b: Estimated intercept parameter
        save_path: Path to save the plot
        font_size: Base font size
    """
    mu_true = true_a * U_values + true_b
    mu_est = est_a * U_values + est_b
    
    plt.figure(figsize=(6, 5))
    plt.plot(times, mu_true, color=colors[0], 
             label=rf"True: $\mu(t) = {true_a:.2f}\cdot U(t) + {true_b:.2f}$", linewidth=2)
    plt.plot(times, mu_est, color=colors[2], 
             label=rf"Predicted: $\mu(t) = {est_a:.2f}\cdot U(t) + {est_b:.2f}$", 
             linewidth=2, linestyle='--')
    plt.xlabel("Time (s)", fontsize=font_size)
    plt.ylabel(r"$\mu(t)$", fontsize=font_size)
    plt.title("Estimated Mean Function", fontsize=font_size+2)
    plt.xticks(fontsize=font_size-2)
    plt.yticks(fontsize=font_size-2)
    plt.legend(frameon=False, fontsize=font_size-2)
    sns.despine(trim=True)
    plt.savefig(save_path, dpi=300)
    plt.close()


def plot_predictive_distributions(times, X, true_simulations, est_simulations, 
                                   times_train, save_path, font_size=16):
    """Plot predictive distributions with 95% confidence intervals.
    
    Compares true and estimated model predictions against observed data.
    
    Args:
        times: Array of time points
        X: Observed trajectory
        true_simulations: Trajectories from true parameters (shape: n_samples x N+1)
        est_simulations: Trajectories from estimated parameters (shape: n_samples x N+1)
        times_train: Training time points (to show train/test split)
        save_path: Path to save the plot
        font_size: Base font size
    """
    # Calculate quantiles
    true_median = jnp.median(true_simulations, axis=0)
    true_lower = jnp.quantile(true_simulations, 0.025, axis=0)
    true_upper = jnp.quantile(true_simulations, 0.975, axis=0)
    
    est_median = jnp.median(est_simulations, axis=0)
    est_lower = jnp.quantile(est_simulations, 0.025, axis=0)
    est_upper = jnp.quantile(est_simulations, 0.975, axis=0)
    
    plt.figure(figsize=(8, 5))
    plt.plot(times, X, color='k', linestyle='-', alpha=0.7, label="Generated Data", linewidth=1)
    plt.plot(times, true_median, color=colors[0], linestyle='-', label="True Median", linewidth=1)
    plt.fill_between(times, true_lower, true_upper, color=colors[0], alpha=0.2, label="True 95% CI")
    plt.plot(times, est_median, color=colors[1], linestyle='-', label="Predicted Median", linewidth=1)
    plt.fill_between(times, est_lower, est_upper, color=colors[1], alpha=0.2, label="Predicted 95% CI")
    plt.axvline(x=times_train[-1], color='r', linestyle='--', label='Train/Test Split')
    plt.xlabel("Time (s)", fontsize=font_size)
    plt.ylabel(r"$X_t$", fontsize=font_size)
    plt.xticks(fontsize=font_size-2)
    plt.yticks(fontsize=font_size-2)
    plt.legend(frameon=False, fontsize=font_size-2)
    sns.despine(trim=True)
    plt.savefig(save_path, dpi=300)
    plt.close()