243 lines
13 KiB
Python
243 lines
13 KiB
Python
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}") |