106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
from pathlib import Path
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
|
|
# --------------------------------------------------
|
|
# 1. Load the data
|
|
# --------------------------------------------------
|
|
file_name = "rm75b_comfort_workspace_v_minis_collision.csv"
|
|
csv_path = Path(file_name)
|
|
|
|
# The file has no column names, so header=None is important.
|
|
df_csv = pd.read_csv(
|
|
csv_path,
|
|
)
|
|
rate_res = df_csv.iloc[:, :4]
|
|
try:
|
|
rate_res_sort = rate_res.sort_values('z').reset_index(drop=True)
|
|
except:
|
|
rate_res_sort = rate_res
|
|
|
|
DECIMALS = 4
|
|
|
|
try:
|
|
x_unique = np.round(rate_res_sort['x'], DECIMALS).unique()
|
|
y_unique = np.round(rate_res_sort['y'], DECIMALS).unique()
|
|
z_unique = np.round(rate_res_sort['z'], DECIMALS).unique()
|
|
except:
|
|
x_unique = np.round(rate_res_sort.iloc[:,0], DECIMALS).unique()
|
|
y_unique = np.round(rate_res_sort.iloc[:, 1], DECIMALS).unique()
|
|
z_unique = np.round(rate_res_sort.iloc[:, 2], DECIMALS).unique()
|
|
|
|
nx, ny, nz = len(x_unique), len(y_unique), len(z_unique)
|
|
|
|
|
|
ik_rates = rate_res_sort.to_numpy()
|
|
|
|
# --------------------------------------------------
|
|
# 2. Create an output directory
|
|
# --------------------------------------------------
|
|
output_dir = Path(file_name.split(".")[0])
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
# --------------------------------------------------
|
|
# 3. Use the same colour scale for every z-plane
|
|
# --------------------------------------------------
|
|
df = rate_res_sort
|
|
value_min = df["ik_success_rate"].min()
|
|
value_max = df["ik_success_rate"].max()
|
|
|
|
# More levels give a smoother-looking contour plot.
|
|
levels = np.linspace(value_min, value_max, 51)
|
|
|
|
|
|
# --------------------------------------------------
|
|
# 4. Draw one contour plot for each z-plane
|
|
# --------------------------------------------------
|
|
for z_value, plane in df.groupby("z", sort=True):
|
|
|
|
# Rows become y-coordinates, columns become x-coordinates.
|
|
grid = plane.pivot(index="y", columns="x", values="ik_success_rate")
|
|
|
|
x = grid.columns.to_numpy()
|
|
y = grid.index.to_numpy()
|
|
ik_grid = grid.to_numpy()
|
|
|
|
X, Y = np.meshgrid(x, y)
|
|
|
|
fig, ax = plt.subplots(figsize=(7, 6))
|
|
|
|
contour = ax.contourf(
|
|
X,
|
|
Y,
|
|
ik_grid,
|
|
levels=levels,
|
|
cmap="viridis",
|
|
extend="both",
|
|
)
|
|
|
|
highlight_levels = [0.6, 0.7]
|
|
# Only plot if the levels are within the data range (optional)
|
|
if value_min <= 0.6 <= value_max or value_min <= 0.7 <= value_max:
|
|
lines = ax.contour(X, Y, ik_grid, levels=highlight_levels,
|
|
colors='red', linewidths=2, linestyles='solid')
|
|
# Optionally label the lines
|
|
ax.clabel(lines, inline=True, fontsize=10, fmt='%1.1f')
|
|
|
|
|
|
colorbar = fig.colorbar(contour, ax=ax)
|
|
colorbar.set_label("IK rate")
|
|
|
|
ax.set_title(f"IK rate at z = {z_value:.2f}")
|
|
ax.set_xlabel("x")
|
|
ax.set_ylabel("y")
|
|
ax.set_aspect("equal")
|
|
|
|
fig.tight_layout()
|
|
|
|
output_path = output_dir / f"ik_contour_z_{z_value:.2f}.png"
|
|
fig.savefig(output_path, dpi=200, bbox_inches="tight")
|
|
plt.close(fig)
|
|
|
|
print(f"Plots saved to: {output_dir.resolve()}") |