223 lines
12 KiB
Python
223 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""Plot physical-plane FEM cuts without joining material across the vacuum."""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
from matplotlib.collections import LineCollection
|
||
from matplotlib.colors import Normalize, TwoSlopeNorm
|
||
import numpy as np
|
||
|
||
|
||
@dataclass
|
||
class Slice:
|
||
directory: Path
|
||
metadata: dict[str, str]
|
||
vertices: np.ndarray
|
||
triangles: np.ndarray
|
||
surface: np.ndarray
|
||
reference_surface: np.ndarray
|
||
|
||
|
||
def table(path: Path, names: list[str]) -> np.ndarray:
|
||
with path.open(newline="") as stream:
|
||
reader = csv.DictReader(stream)
|
||
if reader.fieldnames is None or any(name not in reader.fieldnames for name in names):
|
||
raise ValueError(f"Missing required columns in {path}")
|
||
rows = [[float(row[name]) for name in names] for row in reader]
|
||
values = np.asarray(rows, dtype=float).reshape((-1, len(names)))
|
||
if not np.isfinite(values).all():
|
||
raise ValueError(f"Nonfinite coordinates or fields in {path}")
|
||
return values
|
||
|
||
|
||
def load_slice(directory: Path) -> Slice:
|
||
with (directory / "metadata.csv").open(newline="") as stream:
|
||
metadata = {row["metric"]: row["value"] for row in csv.DictReader(stream)}
|
||
axes = {"xz": (0, 2, 1), "xy": (0, 1, 2), "yz": (1, 2, 0)}
|
||
plane = metadata["plane"]
|
||
if plane not in axes or tuple(int(metadata[name]) for name in ("horizontal_axis", "vertical_axis", "normal_axis")) != axes[plane]:
|
||
raise ValueError("Slice orientation metadata is inconsistent")
|
||
if float(metadata["deformation_scale"]) != 1:
|
||
raise ValueError("This plotter requires actual deformation with scale exactly one")
|
||
tolerance, offset = float(metadata["root_tolerance"]), float(metadata["offset"])
|
||
if not np.isfinite([tolerance, offset]).all() or tolerance <= 0:
|
||
raise ValueError("Invalid physical-plane tolerance or offset")
|
||
v = table(directory / "vertices.csv", ["vertex", "element", "x", "y", "z", "rho", "h", "plane_error"])
|
||
t = table(directory / "triangles.csv", ["triangle", "element", "v0", "v1", "v2"])
|
||
if len(v) != int(metadata["vertex_count"]) or len(t) != int(metadata["triangle_count"]):
|
||
raise ValueError("Slice row counts disagree with metadata")
|
||
if not np.array_equal(v[:, 0], np.arange(len(v))) or not np.array_equal(t[:, 0], np.arange(len(t))):
|
||
raise ValueError("Slice indices must be contiguous")
|
||
connectivity = t[:, 2:].astype(int)
|
||
if not np.array_equal(connectivity, t[:, 2:]) or (connectivity.size and (connectivity.min() < 0 or connectivity.max() >= len(v))):
|
||
raise ValueError("Invalid supplied triangle connectivity")
|
||
if len(t) and not np.all(v[connectivity, 1] == t[:, 1, None]):
|
||
raise ValueError("A triangle joins different material elements")
|
||
normal = axes[plane][2]
|
||
if len(v) and (np.max(np.abs(v[:, 2 + normal] - offset)) > tolerance * 1.01 or np.max(v[:, 7]) > tolerance * 1.01):
|
||
raise ValueError("Field samples do not lie on the requested physical plane")
|
||
surfaces = []
|
||
for filename, count in (("surface.csv", "surface_segments"), ("reference_surface.csv", "reference_surface_segments")):
|
||
s = table(directory / filename, ["x0", "y0", "z0", "x1", "y1", "z1", "maximum_plane_error"])
|
||
if len(s) != int(metadata[count]):
|
||
raise ValueError("Surface row count disagrees with metadata")
|
||
if len(s) and (np.max(np.abs(s[:, [normal, 3 + normal]] - offset)) > tolerance * 1.01 or np.max(s[:, 6]) > tolerance * 1.01):
|
||
raise ValueError("An outline does not intersect the requested physical plane")
|
||
surfaces.append(s)
|
||
if int(metadata["surface_coplanar_triangles"]) or int(metadata["reference_surface_coplanar_triangles"]):
|
||
raise ValueError("A surface patch is coplanar; this is not a resolved curve intersection")
|
||
if len(t) and not len(surfaces[0]):
|
||
raise ValueError("A nonempty material slice needs its deformed boundary")
|
||
return Slice(directory, metadata, v, connectivity, *surfaces)
|
||
|
||
|
||
def color_scale(values: np.ndarray):
|
||
"""Keep the sign visible; neither abs, logarithms nor clipping changes data."""
|
||
lower, upper = float(values.min()), float(values.max())
|
||
if lower < 0 < upper:
|
||
return TwoSlopeNorm(vmin=lower, vcenter=0, vmax=upper), "coolwarm", "color scale centered at zero"
|
||
if lower == upper:
|
||
delta = max(abs(lower) * .01, 1e-12)
|
||
return Normalize(lower - delta, upper + delta), "viridis", "linear color scale"
|
||
return Normalize(lower, upper), "viridis", "linear color scale"
|
||
|
||
|
||
def rasterize_fields(data: Slice, bounds, pixels=512):
|
||
"""Evaluate each triangle's linear scalar interpolant at pixel centers.
|
||
|
||
No global point-location topology is assumed: curved-cell chords can have
|
||
gaps or overlaps. In an overlap the first supplied triangle owns the pixel.
|
||
This deterministic display rule does not turn overlaps into a union-area
|
||
diagnostic and never averages independent element traces.
|
||
"""
|
||
if not 32 <= pixels <= 2048:
|
||
raise ValueError("Raster resolution must be between 32 and 2048 pixels per axis")
|
||
u0, u1, v0, v1 = bounds
|
||
if not np.isfinite(bounds).all() or u1 <= u0 or v1 <= v0:
|
||
raise ValueError("Invalid scalar raster bounds")
|
||
du, dv = (u1 - u0) / pixels, (v1 - v0) / pixels
|
||
us = u0 + (np.arange(pixels) + .5) * du
|
||
vs = v0 + (np.arange(pixels) + .5) * dv
|
||
values = np.zeros((2, pixels, pixels))
|
||
covered = np.zeros((pixels, pixels), dtype=bool)
|
||
u, v = int(data.metadata["horizontal_axis"]), int(data.metadata["vertical_axis"])
|
||
for indices in data.triangles:
|
||
triangle = data.vertices[indices]
|
||
x, y = triangle[:, 2 + u], triangle[:, 2 + v]
|
||
left = max(0, int(np.ceil((x.min() - u0) / du - .5 - 1e-12)))
|
||
right = min(pixels - 1, int(np.floor((x.max() - u0) / du - .5 + 1e-12)))
|
||
bottom = max(0, int(np.ceil((y.min() - v0) / dv - .5 - 1e-12)))
|
||
top = min(pixels - 1, int(np.floor((y.max() - v0) / dv - .5 + 1e-12)))
|
||
if left > right or bottom > top:
|
||
continue
|
||
dx1, dx2, dy1, dy2 = x[1] - x[0], x[2] - x[0], y[1] - y[0], y[2] - y[0]
|
||
determinant = dx1 * dy2 - dy1 * dx2
|
||
if determinant == 0:
|
||
continue
|
||
dx, dy = us[None, left:right + 1] - x[0], vs[bottom:top + 1, None] - y[0]
|
||
a = (dx * dy2 - dy * dx2) / determinant
|
||
b = (dy * dx1 - dx * dy1) / determinant
|
||
ownership = covered[bottom:top + 1, left:right + 1]
|
||
inside = (a >= -2e-13) & (b >= -2e-13) & (a + b <= 1 + 2e-13) & ~ownership
|
||
for field, column in enumerate((5, 6)):
|
||
scalar = triangle[0, column] + a * (triangle[1, column] - triangle[0, column]) + b * (triangle[2, column] - triangle[0, column])
|
||
values[field, bottom:top + 1, left:right + 1][inside] = scalar[inside]
|
||
ownership[inside] = True
|
||
return np.ma.array(values, mask=np.broadcast_to(~covered, values.shape))
|
||
|
||
|
||
def render(data: Slice, *, reference_outline: bool = False, title: str | None = None,
|
||
normalizations=None, physical_bounds=None, figure=None, axes=None, heading=True, pixels=512):
|
||
m, vertices = data.metadata, data.vertices
|
||
u, v = int(m["horizontal_axis"]), int(m["vertical_axis"])
|
||
labels = ("x", "y", "z")
|
||
if figure is None:
|
||
figure, axes = plt.subplots(1, 2, figsize=(12.4, 5.6), layout="constrained")
|
||
elif axes is None or len(axes) != 2:
|
||
raise ValueError("An existing figure requires two field axes")
|
||
bounds = physical_bounds if physical_bounds is not None else (float(m["u_min"]), float(m["u_max"]), float(m["v_min"]), float(m["v_max"]))
|
||
raster = rasterize_fields(data, bounds, pixels) if len(data.triangles) else None
|
||
for field, (ax, column, name) in enumerate(zip(axes, (5, 6), ("Independent density", "Specific enthalpy"))):
|
||
if raster is not None:
|
||
values = vertices[:, column]
|
||
norm, cmap, note = color_scale(values) if normalizations is None else normalizations[field]
|
||
# Colormapping follows scalar interpolation. Gouraud shading would
|
||
# interpolate RGBA instead, changing an affine field's appearance
|
||
# under triangle refinement even with identical scalar norms.
|
||
artist = ax.imshow(raster[field], extent=bounds, origin="lower", interpolation="none", norm=norm, cmap=cmap, rasterized=True)
|
||
bar = figure.colorbar(artist, ax=ax, shrink=.82, label="model field units")
|
||
ticks = bar.get_ticks()
|
||
ticks = ticks[(ticks >= norm.vmin) & (ticks <= norm.vmax)]
|
||
bar.set_ticks(np.unique(np.r_[norm.vmin, ticks, norm.vmax]))
|
||
ax.set_title(f"{name}\nsampled min {values.min():.5g}; max {values.max():.5g}\n{note}", fontsize=10)
|
||
else:
|
||
ax.set_title(name)
|
||
ax.text(.5, .5, "No material intersection resolved", ha="center", va="center", transform=ax.transAxes)
|
||
for surface, color, linestyle, label in (
|
||
(data.surface, "#111111", "solid", "Deformed material boundary"),
|
||
(data.reference_surface if reference_outline else np.empty((0, 7)), "#656565", "dashed", "Undeformed boundary"),
|
||
):
|
||
if len(surface):
|
||
segments = np.stack((surface[:, [u, v]], surface[:, [3 + u, 3 + v]]), axis=1)
|
||
ax.add_collection(LineCollection(segments, colors=color, linewidths=1.05, linestyles=linestyle, label=label))
|
||
lower = np.array([bounds[0], bounds[2]])
|
||
upper = np.array([bounds[1], bounds[3]])
|
||
if not np.isfinite(np.r_[lower, upper]).all() or np.any(upper <= lower):
|
||
raise ValueError("Invalid slice display bounds")
|
||
padding = .04 * max(upper - lower)
|
||
ax.set_xlim(lower[0] - padding, upper[0] + padding)
|
||
ax.set_ylim(lower[1] - padding, upper[1] + padding)
|
||
ax.set_aspect("equal", adjustable="box")
|
||
ax.set_xlabel(f"{labels[u]} (model length units)")
|
||
ax.set_ylabel(f"{labels[v]} (model length units)")
|
||
ax.set_facecolor("#f3f3f3")
|
||
if len(data.surface) or (reference_outline and len(data.reference_surface)):
|
||
ax.legend(loc="upper right", fontsize=7, framealpha=.9)
|
||
description = title or data.directory.parent.name.replace("_", " ")
|
||
status = "Equilibrium not assessed by this image"
|
||
run_metadata = data.directory.parent.parent / "metadata.txt"
|
||
if data.directory.parent.name == "final_slices" and run_metadata.exists():
|
||
entries = dict(line.split("=", 1) for line in run_metadata.read_text().splitlines() if "=" in line)
|
||
if entries.get("converged") == "0":
|
||
status = "UNCONVERGED state"
|
||
elif entries.get("converged") == "1":
|
||
status = "Solver reported convergence; physical validation is separate"
|
||
if heading:
|
||
figure.suptitle(f"{description} — physical {m['plane']} plane, {labels[int(m['normal_axis'])]}={float(m['offset']):.6g}\n"
|
||
f"{status}; actual deformation; subdivisions {m['subdivisions']}; {pixels}×{pixels} scalar pixels", fontsize=11)
|
||
else:
|
||
for ax in axes:
|
||
ax.set_title(f"{description}\n{ax.get_title()}", fontsize=10)
|
||
return figure
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("directory", type=Path, help="One plane directory, e.g. final_slices/xz")
|
||
parser.add_argument("--output", type=Path, help="Output prefix; defaults to DIRECTORY/model_slice")
|
||
parser.add_argument("--reference-outline", action="store_true")
|
||
parser.add_argument("--pixels", type=int, default=512, help="Scalar sample grid per axis, 32–2048 (default 512)")
|
||
parser.add_argument("--title")
|
||
args = parser.parse_args()
|
||
output = args.output or args.directory / "model_slice"
|
||
paths = [output.with_suffix(suffix) for suffix in (".png", ".pdf")]
|
||
if any(path.exists() for path in paths):
|
||
raise FileExistsError("Refusing to overwrite an existing slice figure")
|
||
figure = render(load_slice(args.directory), reference_outline=args.reference_outline, title=args.title, pixels=args.pixels)
|
||
output.parent.mkdir(parents=True, exist_ok=True)
|
||
for path in paths:
|
||
figure.savefig(path, dpi=200)
|
||
print(path)
|
||
plt.close(figure)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|