Skip to content

Commit e39ac51

Browse files
author
Mohammadali Shahiri
committed
feat(pendulum-graphs): Implement manual graph display for pendulum test
ccc VGG zzz Revert "fix: removing unused variables" This reverts commit 7dd173e.
1 parent 3b844f8 commit e39ac51

File tree

4 files changed

+119
-5
lines changed

4 files changed

+119
-5
lines changed

bioptim/examples/getting_started/pendulum.py

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010
"""
1111

1212
import platform
13+
import pickle
14+
import matplotlib
15+
matplotlib.use('Qt5Agg') # Use 'Qt5Agg' for PyQt5 or 'Qt6Agg' for PyQt6
1316

1417
from bioptim import (
1518
OptimalControlProgram,
@@ -147,11 +150,27 @@ def main():
147150
# sol.graphs(show_bounds=True)
148151
sol.animate(n_frames=100)
149152

153+
states = sol.decision_states()
154+
controls = sol.decision_controls()
155+
time=sol.stepwise_time()
156+
final_time=ocp.phase_time
157+
158+
q_sol = states["q"]
159+
qdot_sol = states["qdot"]
160+
tau_sol = controls["tau"]
161+
162+
data = {
163+
"q_sol": q_sol,
164+
"qdot_sol": qdot_sol,
165+
"tau_sol": tau_sol,
166+
"time": time,
167+
"final_time": final_time,
168+
}
169+
150170
# # --- Save the solution --- #
151-
# import pickle
152-
# with open("pendulum.pkl", "wb") as file:
153-
# del sol.ocp
154-
# pickle.dump(sol, file)
171+
with open("pendulum.pkl", "wb") as file:
172+
del sol.ocp
173+
pickle.dump(data, file)
155174

156175

157176
if __name__ == "__main__":
141 KB
Loading
Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
2+
# This code is designed to read data from a pickle file and plot various solutions for a pendulum system.
3+
# It's likely part of a verification process for "Testing solution outputs to manually display graphs of pendulum #813".
4+
# The script handles the extraction and processing of the pendulum system's state (q, qdot) and control (tau) variables.
5+
# These variables are visualized for two degrees of freedom across a specified time interval.
6+
# Matplotlib is used for plotting, with adjustments made for layout to prevent overlapping of the subplots.
7+
8+
9+
import pickle
10+
import matplotlib.pyplot as plt
11+
import matplotlib
12+
import numpy as np
13+
14+
# Set the backend for Matplotlib to 'Qt5Agg'
15+
matplotlib.use('Qt5Agg') # Use 'Qt5Agg' for PyQt5 compatibility, 'Qt6Agg' if using PyQt6
16+
17+
# Load data from the pickle file
18+
with open("pendulum.pkl", "rb") as file:
19+
data = pickle.load(file)
20+
21+
# Extract and process q, qdot, and tau solutions from the data
22+
# Extracting the first and last elements of each solution for two degrees of freedom (DOF)
23+
q_sol_1 = [point[0][0] for point in data["q_sol"]]
24+
q_sol_2 = [point[-1][-1] for point in data["q_sol"]]
25+
qdot_sol_1 = [point[0][0] for point in data["qdot_sol"]]
26+
qdot_sol_2 = [point[-1][-1] for point in data["qdot_sol"]]
27+
tau_sol_1 = [point[0][0] for point in data["tau_sol"]]
28+
tau_sol_2 = [point[-1][-1] for point in data["tau_sol"]]
29+
30+
# Append the last value to tau solutions to ensure equal length with the time array
31+
tau_sol_1.append(tau_sol_1[-1])
32+
tau_sol_2.append(tau_sol_2[-1])
33+
34+
# Extract time data and convert each DM (Differential Matrix) to a NumPy array
35+
time = data["time"]
36+
numpy_array = np.concatenate([np.array(t.full()).flatten()[:-1] if i < len(time) - 1 else np.array(t.full()).flatten() for i, t in enumerate(time)])
37+
38+
# Create a time array between 0 and 1 with 400 intervals (total 401 points)
39+
Time = np.linspace(0, 1, 401)
40+
41+
# Create a figure and a set of subplots
42+
# 3 rows for q, qdot, tau and 2 columns for each DOF
43+
fig, axs = plt.subplots(3, 2, figsize=(10, 15))
44+
45+
# Plotting q solutions for both DOFs
46+
axs[0, 0].plot(Time, q_sol_1)
47+
axs[0, 0].set_title("q Solution for first DOF")
48+
axs[0, 0].set_ylabel("q")
49+
axs[0, 0].set_xlabel("Time")
50+
axs[0, 0].grid(True)
51+
52+
axs[0, 1].plot(Time, q_sol_2)
53+
axs[0, 1].set_title("q Solution for second DOF")
54+
axs[0, 1].set_ylabel("q")
55+
axs[0, 1].set_xlabel("Time")
56+
axs[0, 1].grid(True)
57+
58+
axs[1, 0].plot(Time, qdot_sol_1)
59+
axs[1, 0].set_title("qdot Solution for first DOF")
60+
axs[1, 0].set_ylabel("qdot")
61+
axs[1, 0].set_xlabel("Time")
62+
axs[1, 0].grid(True)
63+
64+
axs[1, 1].plot(Time, qdot_sol_2)
65+
axs[1, 1].set_title("qdot Solution for second DOF")
66+
axs[1, 1].set_ylabel("qdot")
67+
axs[1, 1].set_xlabel("Time")
68+
axs[1, 1].grid(True)
69+
70+
axs[2, 0].plot(Time, tau_sol_1)
71+
axs[2, 0].set_title("Tau Solution for first DOF")
72+
axs[2, 0].set_ylabel("Tau")
73+
axs[2, 0].set_xlabel("Time")
74+
axs[2, 0].grid(True)
75+
76+
axs[2, 1].plot(Time, tau_sol_2)
77+
axs[2, 1].set_title("Tau Solution for second DOF")
78+
axs[2, 1].set_ylabel("Tau")
79+
axs[2, 1].set_xlabel("Time")
80+
axs[2, 1].grid(True)
81+
82+
plt.tight_layout()
83+
plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=None, hspace=0.3)
84+
plt.savefig("pendulum_graph.jpg")
85+
86+
# Display the plot
87+
plt.show()

bioptim/limits/penalty_option.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ class PenaltyOption(OptionGeneric):
5050
If the minimization is applied to derivative of the penalty [f(t, t+1)]
5151
integration_rule: QuadratureRule
5252
The integration rule to use for the penalty
53+
transition: bool
54+
If the penalty is a transition
5355
nodes_phase: tuple[int, ...]
5456
The index of the phases when penalty is multinodes
5557
penalty_type: PenaltyType
@@ -309,7 +311,13 @@ def transform_penalty_to_stochastic(self, controller: PenaltyController, fcn, st
309311
terms, meaning that you have to declare the constraint h=0 and the "variation of h"=buffer ** 2 with
310312
is_stochastic=True independently.
311313
"""
314+
312315
# TODO: Charbie -> This is just a first implementation (x=[q, qdot]), it should then be generalized
316+
317+
nx = controller.states["q"].cx_start.shape[0]
318+
n_root = controller.model.nb_root
319+
n_joints = nx - n_root
320+
313321
if "cholesky_cov" in controller.algebraic_states.keys():
314322
l_cov_matrix = StochasticBioModel.reshape_to_cholesky_matrix(
315323
controller.algebraic_states["cholesky_cov"].cx_start, controller.model.matrix_shape_cov_cholesky
@@ -756,7 +764,7 @@ def _finish_add_target_to_plot(self, controller: PenaltyController):
756764
757765
"""
758766

759-
def plot_function(node_idx):
767+
def plot_function(t0, phases_dt, node_idx, x, u, p, a, penalty=None):
760768
if isinstance(node_idx, (list, tuple)):
761769
return self.target_to_plot[:, [self.node_idx.index(idx) for idx in node_idx]]
762770
else:

0 commit comments

Comments
 (0)