testbench_tools.py

import numpy as np
import json
from itertools import product
import matplotlib.pyplot as plt

# Load configuration from JSON file
def load_config(config_path):
    with open(config_path, 'r') as f:
        return json.load(f)
    
def simulation_configuration_setup():
    config_path = "config.json"
    config = load_config(config_path)

    m_start = config["m_variants_start"]
    m_end = config["m_variants_end"]
    m_step = config["m_variants_step"]

    k_start = config["k_variants_start"]
    k_end = config["k_variants_end"]
    k_step = config["k_variants_step"]

    b_start = config["b_variants_start"]
    b_end = config["b_variants_end"]
    b_step = config["b_variants_step"]

    # Generate model variants
    ms = np.arange(m_start, m_end + m_step, m_step).tolist()
    ks = np.arange(k_start, k_end + k_step, k_step).tolist()
    bs = np.arange(b_start, b_end + b_step, b_step).tolist()

    # Generate all possible combinations of m, k, and b
    λs = [np.array(λ) for λ in product(ms, ks, bs)]

    m = config['true_m']
    k = config['true_k']
    b = config['true_b']
    dt = config["dt"]
    H = np.array(config["H"])
    Q = np.eye(H.shape[1]) * config["Q"]
    R = np.eye(H.shape[0]) * config["R"]
    x0 = np.array(config["initial_state"])
    max_time = config['max_time']
    max_steps = int(config['max_time'] / dt)
    amplitude = config['amplitude']

    return λs, m, k, b, dt, H, Q, R, x0, max_time, max_steps, amplitude

def plot_λ_hat(times, lambda_hats, true_λ):
    # Convert the list of lambda_hats (which are vectors) to a numpy array for easy indexing
    lambda_hats = np.array(lambda_hats)

    fig, ax1 = plt.subplots()

    # Plot estimated mass (m) on the primary y-axis
    ax1.set_xlabel('Time (s)')
    ax1.set_ylabel('Mass (m)', color='blue')
    ax1.plot(times, lambda_hats[:, 0], label='Estimated mass (m)', color='blue', linestyle='-')
    ax1.axhline(y=true_λ[0], color='blue', linestyle='--', label='True mass (m)')
    ax1.tick_params(axis='y', labelcolor='blue')

    # Create a secondary y-axis for spring constant (k)
    ax2 = ax1.twinx()
    ax2.set_ylabel('Spring constant (k)', color='green')
    ax2.plot(times, lambda_hats[:, 1], label='Estimated spring constant (k)', color='green', linestyle='-')
    ax2.axhline(y=true_λ[1], color='green', linestyle='--', label='True spring constant (k)')
    ax2.tick_params(axis='y', labelcolor='green')

    # Create another secondary y-axis for damping coefficient (b) (offset from the right)
    ax3 = ax1.twinx()
    ax3.spines['right'].set_position(('outward', 60))  # Offset third axis to the right
    ax3.set_ylabel('Damping coefficient (b)', color='red')
    ax3.plot(times, lambda_hats[:, 2], label='Estimated damping coefficient (b)', color='red', linestyle='-')
    ax3.axhline(y=true_λ[2], color='red', linestyle='--', label='True damping coefficient (b)')
    ax3.tick_params(axis='y', labelcolor='red')

    # Add legends for each axis
    ax1.legend(loc='upper left')
    ax2.legend(loc='upper right')
    ax3.legend(loc='lower right')

    plt.title('Parameter Estimates (m, k, b) vs Time')
    fig.tight_layout()  # To avoid overlap of labels
    plt.grid(True)
    plt.show()

def plot_heatmap(models_summary, times, λs, title):
    # Convert to string for label on plot
    λs = [f"m: {round(λ[0], 2)}, k: {round(λ[1], 2)}, b: {round(λ[2], 2)}" for λ in λs]
    plt.figure(figsize=(10, 6))
    plt.imshow(models_summary.T, aspect='auto', cmap='hot', origin='lower',
               extent=[times[0], times[-1], λs[0], λs[-1]])

    plt.colorbar(label='Model Likelihoods/Probabilities')
    plt.xlabel('Time (s)')
    plt.ylabel('Model Variants (λs)')
    plt.title(title)
    plt.show()

def mmae_simulator_simulation_and_plot(simulator, λs, true_λ, max_steps, dt):
    # Lists to store time and λ_hat values
    times = []
    lambda_hats = []
    cumulative_posteriors_summary = []
    pdvs_summary = []

    # Main simulation loop
    for step_counter in range(1, max_steps):
        λ_hat, cumulative_posteriors, pdvs = simulator.update(step_counter)
        times.append(step_counter * dt)
        lambda_hats.append(λ_hat)
        cumulative_posteriors_summary.append(cumulative_posteriors)
        pdvs_summary.append(pdvs)
        print(f"Step {step_counter}: λ_hat = {λ_hat}")

    # Convert the list of lambda_hat values (vectors) to a numpy array
    lambda_hats = np.array(lambda_hats)

    # Convert the list of model probabilities to a numpy array
    cumulative_posteriors_summary = np.array(cumulative_posteriors_summary)

    # Convert the list of model probabilities to a numpy array
    pdvs_summary = np.array(pdvs_summary)

    # np.savetxt("pdvs_summary.txt", pdvs_summary)
    # np.savetxt("cumulative_posteriors_summary.txt", cumulative_posteriors_summary)

    # Plot the estimated mass over time
    plot_λ_hat(times, lambda_hats, true_λ)

    # Plot the heatmap for likelihoods (PDVs) over time
    plot_heatmap(pdvs_summary, times, λs, title="Heatmap of Model Likelihood Over Time")

    # Plot the heatmap for model cumulative posterior probabilities over time
    plot_heatmap(cumulative_posteriors_summary, times, λs, title="Heatmap of Cumulative Posterior Probabilities Over Time")