146 lines
7.4 KiB
Python
146 lines
7.4 KiB
Python
import os
|
|
import json
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.gridspec as gridspec
|
|
|
|
# --- 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
|
|
|
|
# --- Composite plotting functions (with z-scale set to linear) ---
|
|
def plot_3d_epoch_evolution_on_axis(ax, epochs, theta_over_epochs, desired_theta, title, num_steps, dt):
|
|
time_steps = np.arange(num_steps) * dt
|
|
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_steps), time_steps, masked_theta_vals)
|
|
|
|
epochs_array = np.array(epochs)
|
|
ax.plot(epochs_array, [time_steps.max()] * 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, num_steps, dt, loss_list, save_path, plot_epoch_range, plot_epoch_step, swap_columns=False):
|
|
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)
|
|
|
|
# Top row: "constant" spanning both columns
|
|
ax_const = fig.add_subplot(gs[0, :], projection='3d')
|
|
epochs_const, theta_const = cond_results["constant"]
|
|
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", num_steps, dt)
|
|
|
|
# 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, theta_loss = cond_results[loss]
|
|
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, num_steps, dt)
|
|
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, theta_mir = cond_results[mirrored_loss]
|
|
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, num_steps, dt)
|
|
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, theta_mir = cond_results[mirrored_loss]
|
|
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, num_steps, dt)
|
|
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, theta_loss = cond_results[loss]
|
|
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, num_steps, dt)
|
|
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_weighting2"
|
|
# Get all JSON files (each representing one condition)
|
|
condition_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
|
|
|
|
# Define simulation parameters (must match those used in run_tests.py)
|
|
dt = 0.02
|
|
num_steps = 500
|
|
|
|
# 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 (change this if you want to sample)
|
|
|
|
# For each condition file, load its data and create the composite plots
|
|
for file in condition_files:
|
|
condition = file.replace(".json", "")
|
|
condition_file = os.path.join(output_dir, file)
|
|
with open(condition_file, "r") as f:
|
|
cond_results = json.load(f)
|
|
|
|
# Determine desired_theta from constant loss data (last theta value of the last epoch)
|
|
epochs_const, theta_const = cond_results["constant"]
|
|
desired_theta = theta_const[-1][-1]
|
|
|
|
# Create a directory for this condition's plots
|
|
condition_path = os.path.join(output_dir, condition)
|
|
os.makedirs(condition_path, exist_ok=True)
|
|
|
|
# Composite figure for linear, quadratic, cubic (with constant at top)
|
|
loss_list1 = ["linear", "quadratic", "cubic"]
|
|
composite_save_path1 = os.path.join(condition_path, "composite_epoch_evolution_linear_quadratic_cubic.png")
|
|
plot_composite_loss_functions_for_condition(cond_results, condition, desired_theta, num_steps, dt, 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(condition_path, "composite_epoch_evolution_inverse_losses.png")
|
|
plot_composite_loss_functions_for_condition(cond_results, condition, desired_theta, num_steps, dt, loss_list2, composite_save_path2, plot_epoch_range, plot_epoch_step, swap_columns=True)
|
|
|
|
print(f"Completed plotting for condition: {condition}")
|