Skip to content

Commit 872c6de

Browse files
Federico-PizarroBejaranoFederico-PizarroBejarano
committed
Further cleaning plotting code
1 parent 01312a5 commit 872c6de

File tree

5 files changed

+224
-194
lines changed

5 files changed

+224
-194
lines changed

‎mm_utils/setup.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
from catkin_pkg.python_setup import generate_distutils_setup
66

7-
setup_args = generate_distutils_setup(packages=["mm_utils"], package_dir={"": "src"})
7+
setup_args = generate_distutils_setup(
8+
packages=["mm_utils", "mm_utils.plotting"], package_dir={"": "src"}
9+
)
810

911
setup(**setup_args)

‎mm_utils/src/mm_utils/plotting/__init__.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
"""
2+
Plotting utilities for mobile manipulation experiments.
3+
4+
This module provides plotting functionality organized into logical sections:
5+
- DataPlotter: Core data loading and processing
6+
- Trajectory plotting: Path and tracking visualization
7+
- MPC plotting: Controller-specific visualization
8+
- Utility functions: Logger construction
9+
"""
10+
11+
from mm_utils.plotting.core import DataPlotter, construct_logger
12+
from mm_utils.plotting.mpc import MPCPlotterMixin
13+
from mm_utils.plotting.trajectory import TrajectoryPlotterMixin
14+
15+
__all__ = [
16+
"DataPlotter",
17+
"MPCPlotterMixin",
18+
"TrajectoryPlotterMixin",
19+
"construct_logger",
20+
]

mm_utils/src/mm_utils/plotting.py renamed to mm_utils/src/mm_utils/plotting/core.py

Lines changed: 5 additions & 193 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
"""
2-
Plotting utilities for mobile manipulation experiments.
2+
Core plotting functionality for mobile manipulation experiments.
33
4-
This module provides plotting functionality organized into logical sections:
5-
- DataPlotter: Core data loading and processing
6-
- Trajectory plotting: Path and tracking visualization
7-
- Utility functions: PDF export, logger construction
4+
This module provides the main DataPlotter class and utility functions for
5+
loading and processing simulation data.
86
"""
97

108
import copy
@@ -18,6 +16,8 @@
1816
import mm_control.MPC as MPC
1917
from mm_utils import math, parsing
2018
from mm_utils.math import wrap_pi_array
19+
from mm_utils.plotting.mpc import MPCPlotterMixin
20+
from mm_utils.plotting.trajectory import TrajectoryPlotterMixin
2121

2222
# =============================================================================
2323
# UTILITY FUNCTIONS
@@ -66,194 +66,6 @@ def construct_logger(
6666
raise ValueError(f"Unrecognized folder structure in {path_to_folder}.")
6767

6868

69-
# =============================================================================
70-
# MPC PLOTTING MIXIN
71-
# =============================================================================
72-
73-
74-
class MPCPlotterMixin:
75-
"""Mixin class containing MPC-specific plotting methods."""
76-
77-
def plot_cost(self):
78-
"""Plot MPC cost function over time."""
79-
t_sim = self.data["ts"]
80-
cost_final = self.data.get("mpc_cost_finals")
81-
82-
if cost_final is None:
83-
print("No cost data found")
84-
return
85-
86-
f, ax = plt.subplots(1, 1)
87-
88-
# Convert to numpy array and handle different shapes
89-
cost_final = np.array(cost_final)
90-
if cost_final.ndim == 1:
91-
ax.plot(
92-
t_sim,
93-
cost_final,
94-
".-",
95-
label=self.data["name"],
96-
linewidth=2,
97-
markersize=8,
98-
)
99-
else:
100-
# Multi-dimensional cost - plot first dimension
101-
ax.plot(
102-
t_sim,
103-
cost_final[:, 0],
104-
".-",
105-
label=self.data["name"],
106-
linewidth=2,
107-
markersize=8,
108-
)
109-
110-
ax.set_xlabel("Time (s)")
111-
ax.set_ylabel("Cost")
112-
ax.legend()
113-
ax.set_title("MPC Cost")
114-
ax.grid(True)
115-
116-
def plot_run_time(self):
117-
"""Plot controller execution time."""
118-
t_sim = self.data["ts"]
119-
run_time = self.data.get("controller_run_time")
120-
if run_time is None:
121-
print("Ignore run time")
122-
return
123-
124-
f, ax = plt.subplots(1, 1)
125-
ax.plot(t_sim, run_time * 1000, label=self.data["name"], linewidth=2)
126-
ax.set_xlabel("Time (s)")
127-
ax.set_ylabel("run time (ms)")
128-
ax.legend()
129-
ax.set_title("Controller Run Time")
130-
131-
132-
# =============================================================================
133-
# TRAJECTORY PLOTTING MIXIN
134-
# =============================================================================
135-
136-
137-
class TrajectoryPlotterMixin:
138-
"""Mixin class containing trajectory and path plotting methods."""
139-
140-
def plot_ee_tracking(self):
141-
"""Plot end-effector position tracking."""
142-
ts = self.data["ts"]
143-
r_ew_w_ds = self.data.get("r_ew_w_ds", [])
144-
r_ew_ws = self.data.get("r_ew_ws", [])
145-
146-
if len(r_ew_w_ds) == 0 and len(r_ew_ws) == 0:
147-
return
148-
149-
_, axes = plt.subplots(1, 1, sharex=True)
150-
legend = self.data["name"]
151-
152-
if len(r_ew_w_ds) > 0:
153-
axes.plot(
154-
ts, r_ew_w_ds[:, 0], label=legend + "$x_d$", color="r", linestyle="--"
155-
)
156-
axes.plot(
157-
ts, r_ew_w_ds[:, 1], label=legend + "$y_d$", color="g", linestyle="--"
158-
)
159-
axes.plot(
160-
ts, r_ew_w_ds[:, 2], label=legend + "$z_d$", color="b", linestyle="--"
161-
)
162-
if len(r_ew_ws) > 0:
163-
axes.plot(ts, r_ew_ws[:, 0], label=legend + "$x$", color="r")
164-
axes.plot(ts, r_ew_ws[:, 1], label=legend + "$y$", color="g")
165-
axes.plot(ts, r_ew_ws[:, 2], label=legend + "$z$", color="b")
166-
axes.grid()
167-
axes.legend()
168-
axes.set_xlabel("Time (s)")
169-
axes.set_ylabel("Position (m)")
170-
axes.set_title("End effector position tracking")
171-
172-
return axes
173-
174-
def plot_base_path(self):
175-
"""Plot base path."""
176-
r_b = self.data.get("r_bw_ws", [])
177-
178-
if len(r_b) == 0:
179-
return
180-
181-
_, ax = plt.subplots(1, 1)
182-
183-
if len(r_b) > 0:
184-
r_b = np.array(r_b) # Convert to numpy array
185-
ax.plot(r_b[:, 0], r_b[:, 1], label=self.data["name"], linewidth=1)
186-
187-
ax.grid()
188-
ax.legend()
189-
ax.set_xlabel("x (m)")
190-
ax.set_ylabel("y (m)")
191-
ax.set_title("Base Path Tracking")
192-
193-
def plot_tracking_err(self):
194-
"""Plot tracking error."""
195-
ts = self.data["ts"]
196-
197-
_, ax = plt.subplots(1, 1)
198-
199-
ax.plot(
200-
ts,
201-
self.data["err_base"],
202-
label=self.data["name"]
203-
+ f" base err, rms={self.data['statistics']['err_base']['rms']:.3f}",
204-
linestyle="--",
205-
)
206-
ax.plot(
207-
ts,
208-
self.data["err_ee"],
209-
label=self.data["name"]
210-
+ f" EE err, rms={self.data['statistics']['err_ee']['rms']:.3f}",
211-
linestyle="-",
212-
)
213-
214-
ax.grid()
215-
ax.legend()
216-
ax.set_xlabel("Time (s)")
217-
ax.set_ylabel("Error (m)")
218-
ax.set_title("Tracking Error vs Time")
219-
220-
def plot_task_performance(self):
221-
"""Plot task performance metrics."""
222-
f, axes = plt.subplots(4, 1, sharex=True)
223-
224-
legend = self.data["name"]
225-
t_sim = self.data["ts"]
226-
227-
axes[0].plot(
228-
t_sim,
229-
self.data["constraints_violation"] * 100,
230-
label=f"{legend} mean={self.data['statistics']['constraints_violation']['mean']*100:.1f}%",
231-
)
232-
axes[0].set_ylabel("Constraints violation (%)")
233-
234-
axes[1].plot(
235-
t_sim,
236-
self.data["err_ee"],
237-
label=f"{legend} acc={self.data['statistics']['err_ee']['integral']:.3f}",
238-
)
239-
axes[1].set_ylabel("EE Error (m)")
240-
241-
axes[2].plot(
242-
t_sim,
243-
self.data["err_base"],
244-
label=f"{legend} acc={self.data['statistics']['err_base']['integral']:.3f}",
245-
)
246-
axes[2].set_ylabel("Base Error (m)")
247-
248-
axes[3].plot(t_sim, self.data["arm_manipulability"], label=legend)
249-
axes[3].set_ylabel("Arm Manipulability")
250-
axes[3].set_xlabel("Time (s)")
251-
252-
for ax in axes:
253-
ax.legend()
254-
ax.grid(True)
255-
256-
25769
# =============================================================================
25870
# CORE DATAPLOTTER CLASS
25971
# =============================================================================

‎mm_utils/src/mm_utils/plotting/mpc.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
"""
2+
MPC-specific plotting functionality.
3+
4+
This module provides plotting methods for Model Predictive Control (MPC) data,
5+
including cost function visualization and controller execution time.
6+
"""
7+
8+
import matplotlib.pyplot as plt
9+
import numpy as np
10+
11+
12+
class MPCPlotterMixin:
13+
"""Mixin class containing MPC-specific plotting methods."""
14+
15+
def plot_cost(self):
16+
"""Plot MPC cost function over time."""
17+
t_sim = self.data["ts"]
18+
cost_final = self.data.get("mpc_cost_finals")
19+
20+
if cost_final is None:
21+
print("No cost data found")
22+
return
23+
24+
f, ax = plt.subplots(1, 1)
25+
26+
# Convert to numpy array and handle different shapes
27+
cost_final = np.array(cost_final)
28+
if cost_final.ndim == 1:
29+
ax.plot(
30+
t_sim,
31+
cost_final,
32+
".-",
33+
label=self.data["name"],
34+
linewidth=2,
35+
markersize=8,
36+
)
37+
else:
38+
# Multi-dimensional cost - plot first dimension
39+
ax.plot(
40+
t_sim,
41+
cost_final[:, 0],
42+
".-",
43+
label=self.data["name"],
44+
linewidth=2,
45+
markersize=8,
46+
)
47+
48+
ax.set_xlabel("Time (s)")
49+
ax.set_ylabel("Cost")
50+
ax.legend()
51+
ax.set_title("MPC Cost")
52+
ax.grid(True)
53+
54+
def plot_run_time(self):
55+
"""Plot controller execution time."""
56+
t_sim = self.data["ts"]
57+
run_time = self.data.get("controller_run_time")
58+
if run_time is None:
59+
print("Ignore run time")
60+
return
61+
62+
f, ax = plt.subplots(1, 1)
63+
ax.plot(t_sim, run_time * 1000, label=self.data["name"], linewidth=2)
64+
ax.set_xlabel("Time (s)")
65+
ax.set_ylabel("run time (ms)")
66+
ax.legend()
67+
ax.set_title("Controller Run Time")

0 commit comments

Comments
 (0)