41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import torch
|
|
from math import pi
|
|
import matplotlib.pyplot as plt
|
|
import sys
|
|
|
|
sys.path.append("/home/judson/Neural-Networks-in-GNC/inverted_pendulum")
|
|
from training.time_weighting_functions import weight_functions
|
|
|
|
# Time grid
|
|
t_start, t_end, t_points = 0, 10, 1000
|
|
t_span = torch.linspace(t_start, t_end, t_points)
|
|
|
|
plt.figure(figsize=(10, 6))
|
|
|
|
# Dictionaries to separate non-mirrored and mirrored functions and to store colors
|
|
colors = {}
|
|
|
|
# First plot non-mirrored functions and record their colors
|
|
for name, weight_fn in weight_functions.items():
|
|
if "mirrored" not in name:
|
|
weights = weight_fn(t_span, t_max=t_end, min_val=0.01)
|
|
line, = plt.plot(t_span.numpy(), weights.numpy(), label=name, linestyle="-")
|
|
colors[name] = line.get_color()
|
|
|
|
# Now plot mirrored functions with the same color as their non-mirrored counterpart
|
|
for name, weight_fn in weight_functions.items():
|
|
if "mirrored" in name:
|
|
weights = weight_fn(t_span, t_max=t_end, min_val=0.01)
|
|
# Derive base name; assuming the mirrored function is named as "base_mirrored"
|
|
base_name = name.replace("_mirrored", "")
|
|
color = colors.get(base_name, None)
|
|
plt.plot(t_span.numpy(), weights.numpy(), label=name, linestyle="-.", color=color)
|
|
|
|
plt.axvline(x=5, label="Mirror Line", color='black', linestyle='--')
|
|
plt.xlabel(r"Time of Simulation ($t_i$)")
|
|
plt.ylabel("Weight")
|
|
plt.title("Time Weighting Functions and Their Mirror")
|
|
plt.legend()
|
|
plt.grid(True)
|
|
plt.savefig("time_weighting_functions.png")
|