import os import json import numpy as np import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec import sys sys.path.append("/home/judson/Neural-Networks-in-GNC/inverted_pendulum") from analysis.analysis_conditions import analysis_conditions # --- Helper function to filter epoch data --- def filter_epoch_data(epochs, theta_over_epochs, epoch_range, epoch_step): """ Filters the list of epochs and corresponding theta values. Only epochs between epoch_range[0] and epoch_range[1] (inclusive) are kept, and then every epoch_step element is selected. """ filtered_epochs = [] filtered_theta_over_epochs = [] for i, ep in enumerate(epochs): if ep >= epoch_range[0] and ep <= epoch_range[1]: filtered_epochs.append(ep) filtered_theta_over_epochs.append(theta_over_epochs[i]) if epoch_step > 1: filtered_epochs = filtered_epochs[::epoch_step] filtered_theta_over_epochs = filtered_theta_over_epochs[::epoch_step] return filtered_epochs, filtered_theta_over_epochs # --- Updated composite plotting functions using the saved time array --- def plot_3d_epoch_evolution_on_axis(ax, epochs, theta_over_epochs, desired_theta, title, time_array): """ Plots the evolution for one epoch using the provided time_array. """ # time_array is provided from the JSON file (same for all epochs) theta_values = np.concatenate(theta_over_epochs) theta_min = np.min(theta_values) theta_max = np.max(theta_values) desired_range_min = max(theta_min, desired_theta - 1.5 * np.pi) desired_range_max = min(theta_max, desired_theta + 1.5 * np.pi) for epoch, theta_vals in reversed(list(zip(epochs, theta_over_epochs))): masked_theta_vals = np.array(theta_vals) masked_theta_vals[(masked_theta_vals < desired_range_min) | (masked_theta_vals > desired_range_max)] = np.nan ax.plot([epoch] * len(time_array), time_array, masked_theta_vals) epochs_array = np.array(epochs) # Use the last time value as the fixed time for the desired_theta reference line ax.plot(epochs_array, [time_array[-1]] * len(epochs_array), [desired_theta] * len(epochs_array), color='r', linestyle='--', linewidth=2, label='Desired Theta') ax.set_xlabel("Epoch") ax.set_ylabel("Time (s)") ax.set_zlabel("Theta (rad)") # ax.set_zscale('symlog') ax.set_title(title) ax.set_zlim(desired_range_min, desired_range_max) ax.view_init(elev=20, azim=-135) def plot_composite_loss_functions_for_condition(cond_results, condition_name, desired_theta, loss_list, save_path, plot_epoch_range, plot_epoch_step, swap_columns=False): """ Generates a composite figure for a given condition using the saved time array. Assumes each JSON file in cond_results contains keys: "epochs", "theta_over_epochs", and "time". """ total_rows = 1 + len(loss_list) # one row for "constant" + one row per loss function fig = plt.figure(figsize=(12, 3 * total_rows)) gs = gridspec.GridSpec(total_rows, 2) # Get time array from the constant loss data (assumed to be the same for all loss functions) time_array = cond_results["constant"]["time"] # Top row: "constant" spanning both columns ax_const = fig.add_subplot(gs[0, :], projection='3d') epochs_const = cond_results["constant"]["epochs"] theta_const = cond_results["constant"]["theta_over_epochs"] epochs_const, theta_const = filter_epoch_data(epochs_const, theta_const, plot_epoch_range, plot_epoch_step) plot_3d_epoch_evolution_on_axis(ax_const, epochs_const, theta_const, desired_theta, "constant", time_array) # For each loss in the provided list, plot the pair of original and mirrored for i, loss in enumerate(loss_list): if not swap_columns: # Left: original; Right: mirrored. ax_left = fig.add_subplot(gs[i+1, 0], projection='3d') if loss in cond_results: epochs_loss = cond_results[loss]["epochs"] theta_loss = cond_results[loss]["theta_over_epochs"] epochs_loss, theta_loss = filter_epoch_data(epochs_loss, theta_loss, plot_epoch_range, plot_epoch_step) plot_3d_epoch_evolution_on_axis(ax_left, epochs_loss, theta_loss, desired_theta, loss, time_array) else: ax_left.set_title(f"No data for {loss}") ax_right = fig.add_subplot(gs[i+1, 1], projection='3d') mirrored_loss = loss + "_mirrored" if mirrored_loss in cond_results: epochs_mir = cond_results[mirrored_loss]["epochs"] theta_mir = cond_results[mirrored_loss]["theta_over_epochs"] epochs_mir, theta_mir = filter_epoch_data(epochs_mir, theta_mir, plot_epoch_range, plot_epoch_step) plot_3d_epoch_evolution_on_axis(ax_right, epochs_mir, theta_mir, desired_theta, mirrored_loss, time_array) else: ax_right.set_title(f"No data for {mirrored_loss}") else: # Swap: Left: mirrored; Right: original. mirrored_loss = loss + "_mirrored" ax_left = fig.add_subplot(gs[i+1, 0], projection='3d') if mirrored_loss in cond_results: epochs_mir = cond_results[mirrored_loss]["epochs"] theta_mir = cond_results[mirrored_loss]["theta_over_epochs"] epochs_mir, theta_mir = filter_epoch_data(epochs_mir, theta_mir, plot_epoch_range, plot_epoch_step) plot_3d_epoch_evolution_on_axis(ax_left, epochs_mir, theta_mir, desired_theta, mirrored_loss, time_array) else: ax_left.set_title(f"No data for {mirrored_loss}") ax_right = fig.add_subplot(gs[i+1, 1], projection='3d') if loss in cond_results: epochs_loss = cond_results[loss]["epochs"] theta_loss = cond_results[loss]["theta_over_epochs"] epochs_loss, theta_loss = filter_epoch_data(epochs_loss, theta_loss, plot_epoch_range, plot_epoch_step) plot_3d_epoch_evolution_on_axis(ax_right, epochs_loss, theta_loss, desired_theta, loss, time_array) else: ax_right.set_title(f"No data for {loss}") plt.tight_layout() plt.savefig(save_path, dpi=300) plt.close() print(f"Saved composite plot to {save_path}") # --- Load simulation results from JSON files (one per condition) --- output_dir = "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/time_weighting" # Get all JSON files (each representing one condition) by listing condition folders condition_names = analysis_conditions.keys() # Specify the epoch range and step for plotting plot_epoch_range = (0, 100) # Only plot epochs between 0 and 100 plot_epoch_step = 1 # Use every epoch for condition in condition_names: condition_dir = os.path.join(output_dir, condition) data_dir = os.path.join(condition_dir, "data") # Load JSON files for each loss function into a dictionary cond_results = {} for file in os.listdir(data_dir): if file.endswith(".json"): loss_function = file.replace(".json", "") file_path = os.path.join(data_dir, file) with open(file_path, "r") as f: cond_results[loss_function] = json.load(f) # Determine desired_theta from the "constant" loss data (last theta value of the last epoch) epochs_const = cond_results["constant"]["epochs"] theta_const = cond_results["constant"]["theta_over_epochs"] desired_theta = analysis_conditions[condition][-1] # Create a directory for plots for this condition (if not already present) plots_dir = os.path.join(condition_dir, "plots") plots_dir = os.path.join(plots_dir, "epoch_evolution") os.makedirs(plots_dir, exist_ok=True) # Composite figure for linear, quadratic, cubic (with constant at top) loss_list1 = ["linear", "quadratic", "cubic"] composite_save_path1 = os.path.join(plots_dir, "composite_linear_quadratic_cubic.png") plot_composite_loss_functions_for_condition(cond_results, condition, desired_theta, loss_list1, composite_save_path1, plot_epoch_range, plot_epoch_step) # Composite figure for inverse, inverse_squared, inverse_cubed (mirrored version on the left) loss_list2 = ["inverse", "inverse_squared", "inverse_cubed"] composite_save_path2 = os.path.join(plots_dir, "composite_inverse.png") plot_composite_loss_functions_for_condition(cond_results, condition, desired_theta, loss_list2, composite_save_path2, plot_epoch_range, plot_epoch_step, swap_columns=True) # Composite figure for square_root and cubic_root (no swap, original on left) loss_list3 = ["square_root", "cubic_root"] composite_save_path3 = os.path.join(plots_dir, "composite_square_root_cubic_root.png") plot_composite_loss_functions_for_condition(cond_results, condition, desired_theta, loss_list3, composite_save_path3, plot_epoch_range, plot_epoch_step) print(f"Completed plotting for condition: {condition}")