import os import json import torch 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/analysis") from analysis_conditions import analysis_conditions sys.path.append("/home/judson/Neural-Networks-in-GNC/inverted_pendulum") from training.time_weighting_functions import weight_functions # ----------------------------------------------------------------------------- # Helper: Replicate the training loss (standard loss) # Loss = mean( w(t) * (theta(t) - theta_desired)^2 ) # ----------------------------------------------------------------------------- def replicate_training_loss(theta_array, time_array, desired_theta, weight_fn_name): t_tensor = torch.tensor(time_array, dtype=torch.float32) theta_tensor = torch.tensor(theta_array, dtype=torch.float32) desired_tensor = torch.full_like(theta_tensor, desired_theta) w_fn = weight_functions[weight_fn_name] w = w_fn(t_tensor, t_max=t_tensor[-1], min_val=0.01) loss_val = torch.mean(w * (theta_tensor - desired_tensor)**2) return loss_val.item() # ----------------------------------------------------------------------------- # Helper: Filter epoch data by range and sampling step # ----------------------------------------------------------------------------- def filter_epoch_data(epochs, theta_over_epochs, epoch_range, epoch_step): 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 # ----------------------------------------------------------------------------- # Composite plotting function for 2-column loss convergence plots (semilogy scale) # ----------------------------------------------------------------------------- def plot_composite_loss_convergence_for_condition(cond_results, condition_name, desired_theta, loss_list, save_path, plot_epoch_range, plot_epoch_step, swap_columns=False, normalization_mode="raw"): """ Creates a composite 2D loss convergence plot on a semilogy scale with two columns per row. Layout: - Top row (spanning 2 columns): Plots the loss computed on the constant network’s trajectories using constant weighting. - Each subsequent row corresponds to one loss function in loss_list: Left subplot shows the "original" (or swapped) version. Right subplot shows the mirrored (or swapped) version. - Each row's left and right subplots share the same y–axis limits. Loss is computed as: loss = mean( weight_fn(t) * (theta(t) - theta_desired)^2 ). normalization_mode: - "raw": Plot raw loss values. - "norm_const": Plot only the constant curve, normalized by the final constant loss from the top row. - "norm_both": Plot both curves, each normalized by its own final value. The x–axis is plotted in linear scale; the y–axis uses semilogy. """ total_rows = 1 + len(loss_list) fig = plt.figure(figsize=(12, 3 * total_rows)) gs = gridspec.GridSpec(total_rows, 2) # --- Top row: Constant network loss (using constant weighting) --- ax_const = fig.add_subplot(gs[0, :]) epochs_const = cond_results["constant"]["epochs"] theta_const = cond_results["constant"]["theta_over_epochs"] time_const = cond_results["constant"]["time"] epochs_const, theta_const = filter_epoch_data(epochs_const, theta_const, plot_epoch_range, plot_epoch_step) const_losses = [replicate_training_loss(theta_arr, time_const, desired_theta, "constant") for theta_arr in theta_const] epochs_const_plot = np.array(epochs_const) norm_const_final = const_losses[-1] if const_losses[-1] != 0 else 1.0 if normalization_mode == "raw": y_const = const_losses title_const = "Unnormalized Loss" elif normalization_mode == "norm_const": y_const = [val / norm_const_final for val in const_losses] title_const = "Loss Normalized by Final Base Loss" elif normalization_mode == "norm_both": y_const = [val / norm_const_final for val in const_losses] title_const = "Loss Normalized by Final Loss" else: y_const = const_losses title_const = "Loss Convergence" ax_const.semilogy(epochs_const_plot, y_const, label="Constant Weighting", color="black", linestyle="--") ax_const.set_title(f"Initial Condition: {condition_name} – {title_const}") ax_const.set_ylabel("Loss") ax_const.legend(fontsize=10) # --- Subsequent rows: one row per loss in loss_list --- for i, base_loss in enumerate(loss_list): if not swap_columns: left_key = base_loss right_key = base_loss + "_mirrored" else: left_key = base_loss + "_mirrored" right_key = base_loss # Left subplot: ax_left = fig.add_subplot(gs[i+1, 0]) if left_key in cond_results: epochs_left = cond_results[left_key]["epochs"] theta_left = cond_results[left_key]["theta_over_epochs"] time_left = cond_results[left_key]["time"] epochs_left, theta_left = filter_epoch_data(epochs_left, theta_left, plot_epoch_range, plot_epoch_step) losses_left_const = [replicate_training_loss(theta_arr, time_left, desired_theta, "constant") for theta_arr in theta_left] losses_left_method = [replicate_training_loss(theta_arr, time_left, desired_theta, left_key) for theta_arr in theta_left] epochs_left_plot = np.array(epochs_left) if normalization_mode == "raw": y_left_const = losses_left_const y_left_method = losses_left_method label_const = "Constant Weighting" label_method = f"{left_key} Weighting" elif normalization_mode == "norm_const": y_left_const = [val / norm_const_final for val in losses_left_const] y_left_method = None label_const = "Constant Weighting (Norm)" elif normalization_mode == "norm_both": norm_left_const = losses_left_const[-1] if losses_left_const[-1] != 0 else 1.0 norm_left_method = losses_left_method[-1] if losses_left_method[-1] != 0 else 1.0 y_left_const = [val / norm_left_const for val in losses_left_const] y_left_method = [val / norm_left_method for val in losses_left_method] label_const = "Constant Weighting (Norm)" label_method = f"{left_key} Weighting (Norm)" ax_left.semilogy(epochs_left_plot, y_left_const, label=label_const, color="black", linestyle="--") if y_left_method is not None: ax_left.semilogy(epochs_left_plot, y_left_method, label=label_method, color="blue") ax_left.set_ylabel("Loss") ax_left.set_title(f"{left_key} Loss Convergence") ax_left.legend(fontsize=10) else: ax_left.set_title(f"No Data for {left_key}") # Right subplot: ax_right = fig.add_subplot(gs[i+1, 1]) if right_key in cond_results: epochs_right = cond_results[right_key]["epochs"] theta_right = cond_results[right_key]["theta_over_epochs"] time_right = cond_results[right_key]["time"] epochs_right, theta_right = filter_epoch_data(epochs_right, theta_right, plot_epoch_range, plot_epoch_step) losses_right_const = [replicate_training_loss(theta_arr, time_right, desired_theta, "constant") for theta_arr in theta_right] losses_right_method = [replicate_training_loss(theta_arr, time_right, desired_theta, right_key) for theta_arr in theta_right] epochs_right_plot = np.array(epochs_right) if normalization_mode == "raw": y_right_const = losses_right_const y_right_method = losses_right_method label_const = "Constant Weighting" label_method = f"{right_key} Weighting" elif normalization_mode == "norm_const": y_right_const = [val / norm_const_final for val in losses_right_const] y_right_method = None label_const = "Constant Weighting (Norm)" elif normalization_mode == "norm_both": norm_right_const = losses_right_const[-1] if losses_right_const[-1] != 0 else 1.0 norm_right_method = losses_right_method[-1] if losses_right_method[-1] != 0 else 1.0 y_right_const = [val / norm_right_const for val in losses_right_const] y_right_method = [val / norm_right_method for val in losses_right_method] label_const = "Constant Weighting (Norm)" label_method = f"{right_key} Weighting (Norm)" ax_right.semilogy(epochs_right_plot, y_right_const, label=label_const, color="black", linestyle="--") if y_right_method is not None: ax_right.semilogy(epochs_right_plot, y_right_method, label=label_method, color="green") ax_right.set_title(f"{right_key} Loss Convergence") ax_right.legend(fontsize=10) else: ax_right.set_title(f"No Data for {right_key}") # Force left and right subplots in this row to share the same y–axis limits. left_ylim = ax_left.get_ylim() right_ylim = ax_right.get_ylim() common_ylim = (min(left_ylim[0], right_ylim[0]), max(left_ylim[1], right_ylim[1])) ax_left.set_ylim(common_ylim) ax_right.set_ylim(common_ylim) for ax in fig.get_axes(): ax.set_xlabel("Epoch") plt.tight_layout() plt.savefig(save_path, dpi=300) plt.close() print(f"Saved composite loss convergence plot to {save_path}") # ----------------------------------------------------------------------------- # Main plotting loop # ----------------------------------------------------------------------------- if __name__ == "__main__": output_dir = "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/time_weighting" # Settings for convergence plots plot_epoch_range = (0, 25) plot_epoch_step = 1 # Define groups as before group1 = ["linear", "quadratic", "cubic"] group2 = ["inverse", "inverse_squared", "inverse_cubed"] group3 = ["square_root", "cubic_root"] condition_names = [name for name in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, name)) and name in analysis_conditions] for condition in condition_names: print(f"Processing condition: {condition}") desired_theta = analysis_conditions[condition][-1] data_dir = os.path.join(output_dir, condition, "data") 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) plots_dir = os.path.join(output_dir, condition, "plots", "loss_convergence", f"{plot_epoch_range[1]}_epochs") os.makedirs(plots_dir, exist_ok=True) # For each group, produce 3 plots (raw, norm_const, norm_both) in 2-column layout. for group, base_name in zip([group1, group2, group3], ["loss_convergence_linear_quadratic_cubic", "loss_convergence_inverse_losses", "loss_convergence_root_losses"]): for mode in ["raw", "norm_const", "norm_both"]: save_path = os.path.join(plots_dir, f"{base_name}_{mode}.png") swap_flag = True if group == group2 else False plot_composite_loss_convergence_for_condition(cond_results, condition, desired_theta, group, save_path, plot_epoch_range, plot_epoch_step, swap_columns=swap_flag, normalization_mode=mode) print(f"Completed loss convergence plotting for condition: {condition}")