163 lines
7.7 KiB
Python
163 lines
7.7 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")
|
|
from analysis.analysis_conditions import analysis_conditions
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Helper: Filter epoch data by range and sampling step
|
|
# -----------------------------------------------------------------------------
|
|
def filter_epoch_data(epochs, theta_over_epochs, epoch_range, epoch_step):
|
|
"""
|
|
Filters epochs and corresponding theta arrays.
|
|
Only epochs between epoch_range[0] and epoch_range[1] (inclusive) are kept,
|
|
then every epoch_step element is selected.
|
|
"""
|
|
filtered_epochs = []
|
|
filtered_theta_over_epochs = []
|
|
for i, ep in enumerate(epochs):
|
|
if epoch_range[0] <= 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
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# 3D Epoch Evolution Plotting (unchanged)
|
|
# -----------------------------------------------------------------------------
|
|
def plot_3d_epoch_evolution_on_axis(ax, epochs, theta_over_epochs, desired_theta, title, time_array):
|
|
"""
|
|
Plots the evolution of theta for each epoch using the provided time_array.
|
|
"""
|
|
# Determine overall theta range
|
|
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)
|
|
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_title(title)
|
|
ax.set_zlim(desired_range_min, desired_range_max)
|
|
ax.view_init(elev=20, azim=-135)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Composite 3D Plot for Base Loss Evolution (manually assigned subplots)
|
|
# -----------------------------------------------------------------------------
|
|
def plot_composite_base_loss_evolution_for_condition(cond_results, condition_name, desired_theta,
|
|
left_list, right_list, save_path,
|
|
plot_epoch_range, plot_epoch_step):
|
|
"""
|
|
Creates a composite 3D epoch evolution plot (theta evolution) that is arranged as follows:
|
|
- Top row (spanning 2 columns): The "one" data (base) is plotted.
|
|
- Subsequent rows are manually assigned:
|
|
* Left column: uses the base loss function specified by left_list[i]
|
|
* Right column: uses the base loss function specified by right_list[i]
|
|
This lets you compare, for example, fractional loss functions on the left versus multiple ones on the right.
|
|
|
|
The plotting function used is the same as before (which plots theta evolution).
|
|
"""
|
|
total_rows = 1 + len(left_list)
|
|
fig = plt.figure(figsize=(12, 3 * total_rows))
|
|
gs = gridspec.GridSpec(total_rows, 2)
|
|
|
|
# Use the common time array from the "one" data.
|
|
time_array = cond_results["one"]["time"]
|
|
|
|
# Top row: "one" data (base)
|
|
ax_top = fig.add_subplot(gs[0, :], projection='3d')
|
|
epochs_one = cond_results["one"]["epochs"]
|
|
theta_one = cond_results["one"]["theta_over_epochs"]
|
|
epochs_one, theta_one = filter_epoch_data(epochs_one, theta_one, plot_epoch_range, plot_epoch_step)
|
|
plot_3d_epoch_evolution_on_axis(ax_top, epochs_one, theta_one, desired_theta, "one", time_array)
|
|
|
|
# Subsequent rows: assign left and right manually
|
|
for i in range(len(left_list)):
|
|
# Left subplot uses left_list[i]
|
|
left_key = left_list[i]
|
|
ax_left = fig.add_subplot(gs[i+1, 0], projection='3d')
|
|
if left_key in cond_results:
|
|
epochs_left = cond_results[left_key]["epochs"]
|
|
theta_left = cond_results[left_key]["theta_over_epochs"]
|
|
epochs_left, theta_left = filter_epoch_data(epochs_left, theta_left, plot_epoch_range, plot_epoch_step)
|
|
plot_3d_epoch_evolution_on_axis(ax_left, epochs_left, theta_left, desired_theta, left_key, time_array)
|
|
else:
|
|
ax_left.set_title(f"No data for {left_key}")
|
|
|
|
# Right subplot uses right_list[i]
|
|
right_key = right_list[i]
|
|
ax_right = fig.add_subplot(gs[i+1, 1], projection='3d')
|
|
if right_key in cond_results:
|
|
epochs_right = cond_results[right_key]["epochs"]
|
|
theta_right = cond_results[right_key]["theta_over_epochs"]
|
|
epochs_right, theta_right = filter_epoch_data(epochs_right, theta_right, plot_epoch_range, plot_epoch_step)
|
|
plot_3d_epoch_evolution_on_axis(ax_right, epochs_right, theta_right, desired_theta, right_key, time_array)
|
|
else:
|
|
ax_right.set_title(f"No data for {right_key}")
|
|
|
|
plt.tight_layout()
|
|
plt.savefig(save_path, dpi=300)
|
|
plt.close()
|
|
print(f"Saved composite base loss evolution plot to {save_path}")
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Main Loop
|
|
# -----------------------------------------------------------------------------
|
|
if __name__ == "__main__":
|
|
# Directory where condition folders are located
|
|
output_dir = "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/base_loss"
|
|
|
|
# Set epoch range and step
|
|
plot_epoch_range = (0, 100)
|
|
plot_epoch_step = 1
|
|
|
|
# Define manual assignment lists:
|
|
# Left column: fractional loss functions (e.g., one_fifth, one_fourth, one_third)
|
|
# Right column: multiple loss functions (e.g., two, three, four)
|
|
left_list = ["one_fifth", "one_fourth", "one_third", "one_half"]
|
|
right_list = ["two", "three", "four", "five"]
|
|
|
|
# Get condition folders that also exist in analysis_conditions
|
|
condition_names = analysis_conditions.keys()
|
|
|
|
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"):
|
|
key = file.replace(".json", "")
|
|
file_path = os.path.join(data_dir, file)
|
|
with open(file_path, "r") as f:
|
|
cond_results[key] = json.load(f)
|
|
|
|
plots_dir = os.path.join(output_dir, condition, "plots", "epoch_evolution")
|
|
os.makedirs(plots_dir, exist_ok=True)
|
|
|
|
save_path = os.path.join(plots_dir, "composite_base_loss_3d.png")
|
|
plot_composite_base_loss_evolution_for_condition(cond_results, condition, desired_theta,
|
|
left_list, right_list, save_path,
|
|
plot_epoch_range, plot_epoch_step)
|
|
print(f"Completed 3D base loss evolution plotting for condition: {condition}")
|