179 lines
9.2 KiB
Python
179 lines
9.2 KiB
Python
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}")
|