405 lines
15 KiB
Python
Executable File
405 lines
15 KiB
Python
Executable File
"""
|
||
Per-trial figure generation for isogrid optimization.
|
||
|
||
Saves 4 PNG figures per sandbox per trial into the trial folder:
|
||
{sandbox_id}_density.png — density field heatmap η(x,y) [after Brain]
|
||
{sandbox_id}_mesh.png — Gmsh triangulation overlaid on density [after Brain]
|
||
{sandbox_id}_ribs.png — final rib profile (pockets) [after Brain]
|
||
{sandbox_id}_stress.png — Von Mises stress field (per-element) [after NX solve]
|
||
|
||
The stress figure overlays FEA results onto the rib pattern so you can see
|
||
which triangles/pockets are over/under-stressed — the key diagnostic for tuning
|
||
density field parameters (η₀, α, β, R₀, s_min, s_max).
|
||
|
||
These PNGs are NEVER deleted by the retention policy — full history is preserved.
|
||
|
||
Usage:
|
||
from plot_trial import plot_trial_figures, plot_stress_figures
|
||
|
||
# After Brain (density, mesh, ribs):
|
||
plot_trial_figures(sb_data, trial_dir)
|
||
|
||
# After NX solve (stress):
|
||
stress_fields = {
|
||
"sandbox_1": {"nodes_xy": [...], "stress_values": [...], "n_elements": N},
|
||
"sandbox_2": {...},
|
||
}
|
||
plot_stress_figures(sb_data, stress_fields, trial_dir, sigma_allow=100.6)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import matplotlib
|
||
matplotlib.use("Agg") # headless — required when NX session may own the display
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.patches as mpatches
|
||
import numpy as np
|
||
|
||
# Project root on path (run_optimization.py sets sys.path already)
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
|
||
|
||
from optimization_engine.isogrid.density_field import evaluate_density_grid
|
||
|
||
|
||
# ─── Resolution for density grid (mm). 5mm is fast enough for plotting. ────
|
||
_DENSITY_RESOLUTION = 5.0
|
||
_DPI = 150
|
||
|
||
|
||
def plot_trial_figures(sb_data: list[dict], trial_dir: Path) -> list[Path]:
|
||
"""
|
||
Generate and save all figures for one trial.
|
||
|
||
Parameters
|
||
----------
|
||
sb_data : list of dicts, one per sandbox.
|
||
Each dict must have keys:
|
||
sandbox_id, geometry, params, triangulation, pockets, ribbed_plate
|
||
trial_dir : Path
|
||
The trial folder where PNGs will be saved.
|
||
|
||
Returns
|
||
-------
|
||
List of Path objects for the files that were written.
|
||
"""
|
||
written: list[Path] = []
|
||
for sbd in sb_data:
|
||
try:
|
||
sb_id = sbd["sandbox_id"]
|
||
geom = sbd["geometry"]
|
||
params = sbd["params"]
|
||
tri = sbd["triangulation"]
|
||
pockets = sbd["pockets"]
|
||
plate = sbd["ribbed_plate"]
|
||
|
||
written.append(_plot_density(geom, params, trial_dir / f"{sb_id}_density.png"))
|
||
written.append(_plot_mesh(geom, params, tri, trial_dir / f"{sb_id}_mesh.png"))
|
||
written.append(_plot_ribs(geom, pockets, plate, trial_dir / f"{sb_id}_ribs.png"))
|
||
except Exception as exc:
|
||
print(f" [Plot] WARNING: could not save figures for {sbd.get('sandbox_id', '?')}: {exc}")
|
||
|
||
return [p for p in written if p is not None]
|
||
|
||
|
||
# =============================================================================
|
||
# Figure 1 — Density heatmap
|
||
# =============================================================================
|
||
|
||
def _plot_density(geometry: dict, params: dict, out_path: Path) -> Path | None:
|
||
"""Save density field heatmap with boundary and hole outlines."""
|
||
try:
|
||
X, Y, eta = evaluate_density_grid(geometry, params, resolution=_DENSITY_RESOLUTION)
|
||
except Exception:
|
||
return None
|
||
|
||
fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)
|
||
|
||
m = ax.pcolormesh(X, Y, eta, shading="auto", cmap="viridis", vmin=0.0, vmax=1.0)
|
||
fig.colorbar(m, ax=ax, label="Density η (0 = sparse · 1 = dense)", shrink=0.8)
|
||
|
||
# Outer boundary
|
||
_draw_boundary(ax, geometry["outer_boundary"], color="white", lw=1.5, alpha=0.85)
|
||
|
||
# Holes
|
||
for hole in geometry.get("holes", []):
|
||
_draw_hole(ax, hole, color="#ff6b6b", lw=1.2)
|
||
|
||
ax.set_aspect("equal")
|
||
sb_id = geometry.get("sandbox_id", "?")
|
||
ax.set_title(
|
||
f"Density Field — {sb_id}\n"
|
||
f"η₀={params['eta_0']:.2f} α={params['alpha']:.2f} β={params['beta']:.2f} "
|
||
f"γ_s={params['gamma_stress']:.2f} R₀={params['R_0']:.0f} R_e={params['R_edge']:.0f}",
|
||
fontsize=8,
|
||
)
|
||
ax.set_xlabel("x (mm)", fontsize=8)
|
||
ax.set_ylabel("y (mm)", fontsize=8)
|
||
ax.tick_params(labelsize=7)
|
||
fig.tight_layout()
|
||
fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
|
||
plt.close(fig)
|
||
return out_path
|
||
|
||
|
||
# =============================================================================
|
||
# Figure 2 — Gmsh triangulation overlay
|
||
# =============================================================================
|
||
|
||
def _plot_mesh(geometry: dict, params: dict, triangulation: dict, out_path: Path) -> Path | None:
|
||
"""Save triangulation overlaid on a translucent density background."""
|
||
vertices = triangulation.get("vertices")
|
||
triangles = triangulation.get("triangles")
|
||
if vertices is None or len(vertices) == 0:
|
||
return None
|
||
|
||
try:
|
||
X, Y, eta = evaluate_density_grid(geometry, params, resolution=_DENSITY_RESOLUTION)
|
||
except Exception:
|
||
eta = None
|
||
|
||
fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)
|
||
|
||
# Density background (translucent)
|
||
if eta is not None:
|
||
ax.pcolormesh(X, Y, eta, shading="auto", cmap="viridis",
|
||
vmin=0.0, vmax=1.0, alpha=0.35)
|
||
|
||
# Triangle edges
|
||
if triangles is not None and len(triangles) > 0:
|
||
ax.triplot(
|
||
vertices[:, 0], vertices[:, 1], triangles,
|
||
"k-", lw=0.35, alpha=0.75,
|
||
)
|
||
|
||
# Outer boundary
|
||
_draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5)
|
||
|
||
# Holes (keepout rings)
|
||
for hole in geometry.get("holes", []):
|
||
_draw_hole(ax, hole, color="#4488ff", lw=1.2)
|
||
# Keepout ring (d_keep × hole_radius)
|
||
d_keep = params.get("d_keep", 1.2)
|
||
r_hole = hole.get("radius", 0) or hole.get("diameter", 0) / 2.0
|
||
if r_hole > 0:
|
||
keepout = plt.Circle(
|
||
hole["center"], r_hole * (1.0 + d_keep),
|
||
color="#4488ff", fill=False, lw=0.8, ls="--", alpha=0.5,
|
||
)
|
||
ax.add_patch(keepout)
|
||
|
||
ax.set_aspect("equal")
|
||
n_tri = len(triangles) if triangles is not None else 0
|
||
n_pts = len(vertices)
|
||
sb_id = geometry.get("sandbox_id", "?")
|
||
ax.set_title(
|
||
f"Triangulation — {sb_id} ({n_tri} triangles · {n_pts} vertices)\n"
|
||
f"s_min={params['s_min']:.1f} mm s_max={params['s_max']:.1f} mm",
|
||
fontsize=8,
|
||
)
|
||
ax.set_xlabel("x (mm)", fontsize=8)
|
||
ax.set_ylabel("y (mm)", fontsize=8)
|
||
ax.tick_params(labelsize=7)
|
||
fig.tight_layout()
|
||
fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
|
||
plt.close(fig)
|
||
return out_path
|
||
|
||
|
||
# =============================================================================
|
||
# Figure 3 — Final rib profile
|
||
# =============================================================================
|
||
|
||
def _plot_ribs(geometry: dict, pockets: list, ribbed_plate, out_path: Path) -> Path | None:
|
||
"""Save final rib pattern — pockets (material removed) + rib plate outline."""
|
||
fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)
|
||
|
||
# Outer boundary filled (light grey = material to start with)
|
||
outer = np.array(geometry["outer_boundary"])
|
||
ax.fill(outer[:, 0], outer[:, 1], color="#e8e8e8", zorder=0)
|
||
_draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5, zorder=3)
|
||
|
||
# Pockets (material removed = pink/salmon)
|
||
# pockets is list[dict] from generate_pockets() — each dict has a 'polyline' key
|
||
for pocket in pockets:
|
||
try:
|
||
polyline = pocket.get("polyline", [])
|
||
if len(polyline) < 3:
|
||
continue
|
||
coords = np.array(polyline)
|
||
patch = mpatches.Polygon(coords, closed=True,
|
||
facecolor="#ffaaaa", edgecolor="#cc4444",
|
||
lw=0.5, alpha=0.85, zorder=2)
|
||
ax.add_patch(patch)
|
||
except Exception:
|
||
pass
|
||
|
||
# Bolt holes (not pocketed — solid keep zones)
|
||
for hole in geometry.get("holes", []):
|
||
_draw_hole(ax, hole, color="#2255cc", lw=1.2, zorder=4)
|
||
|
||
ax.set_aspect("equal")
|
||
sb_id = geometry.get("sandbox_id", "?")
|
||
n_pockets = len(pockets)
|
||
ax.set_title(
|
||
f"Rib Profile — {sb_id} ({n_pockets} pockets)\n"
|
||
f"pink = material removed · grey = rib material · blue = bolt holes",
|
||
fontsize=8,
|
||
)
|
||
ax.set_xlabel("x (mm)", fontsize=8)
|
||
ax.set_ylabel("y (mm)", fontsize=8)
|
||
ax.tick_params(labelsize=7)
|
||
|
||
# Auto-fit to sandbox bounds
|
||
boundary = np.array(geometry["outer_boundary"])
|
||
margin = 5.0
|
||
ax.set_xlim(boundary[:, 0].min() - margin, boundary[:, 0].max() + margin)
|
||
ax.set_ylim(boundary[:, 1].min() - margin, boundary[:, 1].max() + margin)
|
||
|
||
fig.tight_layout()
|
||
fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
|
||
plt.close(fig)
|
||
return out_path
|
||
|
||
|
||
# =============================================================================
|
||
# Geometry helpers
|
||
# =============================================================================
|
||
|
||
def _draw_boundary(ax, outer_boundary, color, lw, alpha=1.0, zorder=2):
|
||
"""Draw a closed polygon boundary."""
|
||
pts = np.array(outer_boundary)
|
||
x = np.append(pts[:, 0], pts[0, 0])
|
||
y = np.append(pts[:, 1], pts[0, 1])
|
||
ax.plot(x, y, color=color, lw=lw, alpha=alpha, zorder=zorder)
|
||
|
||
|
||
def _draw_hole(ax, hole: dict, color, lw, zorder=3):
|
||
"""Draw a circular hole outline.
|
||
|
||
Note: geometry_schema normalizes inner boundaries to dicts with key
|
||
'diameter' (not 'radius'), so we use diameter/2.
|
||
"""
|
||
cx, cy = hole["center"]
|
||
# Normalized hole dicts have 'diameter'; raw dicts may have 'radius'
|
||
r = hole.get("radius", 0) or hole.get("diameter", 0) / 2.0
|
||
if r > 0:
|
||
circle = plt.Circle((cx, cy), r, color=color, fill=False, lw=lw, zorder=zorder)
|
||
ax.add_patch(circle)
|
||
|
||
|
||
# =============================================================================
|
||
# Figure 4 — Von Mises stress field (post-NX-solve)
|
||
# =============================================================================
|
||
|
||
def plot_stress_figures(
|
||
sb_data: list[dict],
|
||
stress_fields: dict,
|
||
trial_dir: Path,
|
||
sigma_allow: float = 100.6,
|
||
) -> list[Path]:
|
||
"""
|
||
Generate and save stress heatmap figures for one trial.
|
||
|
||
Must be called AFTER the NX solve (stress_fields comes from
|
||
extract_sandbox_stress_field()).
|
||
|
||
Parameters
|
||
----------
|
||
sb_data : list of dicts (same as plot_trial_figures)
|
||
stress_fields : dict keyed by sandbox_id
|
||
Each value: {"nodes_xy": [...], "stress_values": [...], "n_elements": int}
|
||
trial_dir : Path where PNGs are saved
|
||
sigma_allow : allowable stress in MPa — shown as reference line on colorbar
|
||
"""
|
||
written: list[Path] = []
|
||
for sbd in sb_data:
|
||
sb_id = sbd["sandbox_id"]
|
||
sf = stress_fields.get(sb_id, {})
|
||
if not sf.get("n_elements", 0):
|
||
continue
|
||
try:
|
||
p = _plot_stress(
|
||
geometry=sbd["geometry"],
|
||
pockets=sbd["pockets"],
|
||
stress_field=sf,
|
||
sigma_allow=sigma_allow,
|
||
out_path=trial_dir / f"{sb_id}_stress.png",
|
||
)
|
||
if p:
|
||
written.append(p)
|
||
except Exception as exc:
|
||
print(f" [Plot] WARNING: stress figure failed for {sb_id}: {exc}")
|
||
return written
|
||
|
||
|
||
def _plot_stress(
|
||
geometry: dict,
|
||
pockets: list,
|
||
stress_field: dict,
|
||
sigma_allow: float,
|
||
out_path: Path,
|
||
) -> Path | None:
|
||
"""
|
||
Von Mises stress heatmap overlaid with rib pocket outlines.
|
||
|
||
Shows which triangles/pockets are over-stressed vs under-stressed.
|
||
White pocket outlines make the rib pattern visible against the stress field.
|
||
"""
|
||
nodes_xy = np.array(stress_field.get("nodes_xy", []))
|
||
stress_vals = np.array(stress_field.get("stress_values", []))
|
||
|
||
if len(nodes_xy) < 3:
|
||
return None
|
||
|
||
fig, ax = plt.subplots(figsize=(9, 5), dpi=_DPI)
|
||
|
||
# Stress field — tricontourf when enough points, scatter otherwise
|
||
cmap = "RdYlGn_r" # green = low stress, red = high/overloaded
|
||
vmax = max(float(np.max(stress_vals)), sigma_allow * 1.05)
|
||
|
||
if len(nodes_xy) >= 6:
|
||
from matplotlib.tri import Triangulation, LinearTriInterpolator
|
||
try:
|
||
triang = Triangulation(nodes_xy[:, 0], nodes_xy[:, 1])
|
||
tc = ax.tricontourf(triang, stress_vals, levels=20,
|
||
cmap=cmap, vmin=0, vmax=vmax)
|
||
cb = fig.colorbar(tc, ax=ax, label="Von Mises (MPa)", shrink=0.8)
|
||
except Exception:
|
||
sc = ax.scatter(nodes_xy[:, 0], nodes_xy[:, 1], c=stress_vals,
|
||
cmap=cmap, s=10, vmin=0, vmax=vmax, alpha=0.85)
|
||
cb = fig.colorbar(sc, ax=ax, label="Von Mises (MPa)", shrink=0.8)
|
||
else:
|
||
sc = ax.scatter(nodes_xy[:, 0], nodes_xy[:, 1], c=stress_vals,
|
||
cmap=cmap, s=10, vmin=0, vmax=vmax, alpha=0.85)
|
||
cb = fig.colorbar(sc, ax=ax, label="Von Mises (MPa)", shrink=0.8)
|
||
|
||
# Mark σ_allow on colorbar
|
||
cb.ax.axhline(y=sigma_allow, color="black", lw=1.2, ls="--")
|
||
cb.ax.text(1.05, sigma_allow / vmax, f"σ_allow\n{sigma_allow:.0f}",
|
||
transform=cb.ax.transAxes, fontsize=6, va="center", ha="left")
|
||
|
||
# Rib pocket outlines (white) — so we can visually correlate stress with pockets
|
||
for pocket in pockets:
|
||
polyline = pocket.get("polyline", [])
|
||
if len(polyline) >= 3:
|
||
coords = np.array(polyline)
|
||
patch = mpatches.Polygon(coords, closed=True,
|
||
facecolor="none", edgecolor="white",
|
||
lw=0.6, alpha=0.75, zorder=3)
|
||
ax.add_patch(patch)
|
||
|
||
# Outer boundary + holes
|
||
_draw_boundary(ax, geometry["outer_boundary"], color="#00cc66", lw=1.5, zorder=4)
|
||
for hole in geometry.get("holes", []):
|
||
_draw_hole(ax, hole, color="#4488ff", lw=1.2, zorder=4)
|
||
|
||
ax.set_aspect("equal")
|
||
sb_id = geometry.get("sandbox_id", "?")
|
||
n_el = stress_field.get("n_elements", 0)
|
||
max_s = float(np.max(stress_vals))
|
||
feasible = max_s <= sigma_allow
|
||
|
||
status = "OK" if feasible else f"OVER by {max_s - sigma_allow:.1f} MPa"
|
||
ax.set_title(
|
||
f"Von Mises Stress — {sb_id} ({n_el} elements) [{status}]\n"
|
||
f"max = {max_s:.1f} MPa · σ_allow = {sigma_allow:.0f} MPa "
|
||
f"· dashed line = limit",
|
||
fontsize=8,
|
||
)
|
||
ax.set_xlabel("x (mm)", fontsize=8)
|
||
ax.set_ylabel("y (mm)", fontsize=8)
|
||
ax.tick_params(labelsize=7)
|
||
|
||
boundary = np.array(geometry["outer_boundary"])
|
||
margin = 5.0
|
||
ax.set_xlim(boundary[:, 0].min() - margin, boundary[:, 0].max() + margin)
|
||
ax.set_ylim(boundary[:, 1].min() - margin, boundary[:, 1].max() + margin)
|
||
|
||
fig.tight_layout()
|
||
fig.savefig(out_path, dpi=_DPI, bbox_inches="tight")
|
||
plt.close(fig)
|
||
return out_path
|