250 lines
7.8 KiB
Python
250 lines
7.8 KiB
Python
"""Plot normE on a plane slice (x / y / z = const), COMSOL-style."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from matplotlib.tri import Triangulation
|
|
|
|
|
|
def load_mesh(path: Path):
|
|
lines = path.read_text().splitlines()
|
|
i = 0
|
|
|
|
def need(tag: str):
|
|
nonlocal i
|
|
if lines[i] != tag:
|
|
raise ValueError(f"expected {tag!r}, got {lines[i]!r}")
|
|
i += 1
|
|
|
|
need("NbrVertex")
|
|
nv = int(lines[i])
|
|
i += 1
|
|
need("Vertex")
|
|
V = np.array([list(map(float, lines[i + k].split())) for k in range(nv)])
|
|
i += nv
|
|
need("NbrTet")
|
|
nt = int(lines[i])
|
|
i += 1
|
|
need("Tet")
|
|
T = np.array([list(map(int, lines[i + k].split())) for k in range(nt)], int) - 1
|
|
i += nt
|
|
need("DomainOfTet")
|
|
i += nt
|
|
need("NbrEdge")
|
|
ne = int(lines[i])
|
|
i += 1
|
|
need("Edge")
|
|
i += ne
|
|
need("EdgeOfTet")
|
|
i += nt
|
|
need("NbrTri")
|
|
ntri = int(lines[i])
|
|
i += 1
|
|
need("Tri")
|
|
Tri = np.array([list(map(int, lines[i + k].split())) for k in range(ntri)], int) - 1
|
|
i += ntri
|
|
need("DomainOfTri")
|
|
dom_tri = np.array([int(lines[i + k]) for k in range(ntri)])
|
|
i += ntri
|
|
if i < len(lines) and lines[i] == "ConnOfTri":
|
|
i += 1 + ntri
|
|
return V, T, Tri, dom_tri
|
|
|
|
|
|
def load_scalar_field(path: Path) -> np.ndarray:
|
|
vals = []
|
|
for line in path.read_text().splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
if line.startswith("//"):
|
|
line = line[2:].strip()
|
|
vals.append(float(line.split()[0]))
|
|
return np.array(vals)
|
|
|
|
|
|
def axis_index(axis: str) -> int:
|
|
return {"x": 0, "y": 1, "z": 2}[axis.lower()]
|
|
|
|
|
|
def other_axes(axis: str) -> tuple[int, int]:
|
|
if axis == "x":
|
|
return 1, 2
|
|
if axis == "y":
|
|
return 0, 2
|
|
return 0, 1
|
|
|
|
|
|
def triangulate_polygon(pts: list[np.ndarray]) -> list[tuple[int, int, int]]:
|
|
if len(pts) < 3:
|
|
return []
|
|
if len(pts) == 3:
|
|
return [(0, 1, 2)]
|
|
# fan triangulation for convex quad
|
|
return [(0, 1, 2), (0, 2, 3)]
|
|
|
|
|
|
def tet_plane_poly(verts: np.ndarray, axis: int, coord: float, tol: float = 1e-9):
|
|
"""Return polygon vertices (3 or 4) where tet intersects plane, or None."""
|
|
c = verts[:, axis]
|
|
cmin, cmax = c.min(), c.max()
|
|
if coord < cmin - tol or coord > cmax + tol:
|
|
return None
|
|
|
|
if np.all(np.abs(c - coord) < 1e-6):
|
|
# face on plane: pick 3 verts on plane (rare duplicate handling)
|
|
idx = np.where(np.abs(c - coord) < 1e-6)[0]
|
|
if len(idx) >= 3:
|
|
return [verts[idx[0]], verts[idx[1]], verts[idx[2]]]
|
|
return None
|
|
|
|
edges = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
|
|
pts = []
|
|
for a, b in edges:
|
|
ca, cb = c[a], c[b]
|
|
if abs(ca - coord) < 1e-9:
|
|
pts.append(verts[a])
|
|
elif abs(cb - coord) < 1e-9:
|
|
pts.append(verts[b])
|
|
elif (ca - coord) * (cb - coord) < 0:
|
|
t = (coord - ca) / (cb - ca)
|
|
pts.append(verts[a] + t * (verts[b] - verts[a]))
|
|
|
|
if len(pts) < 3:
|
|
return None
|
|
|
|
# remove near-duplicates
|
|
uniq = []
|
|
for p in pts:
|
|
if not any(np.linalg.norm(p - q) < 1e-8 for q in uniq):
|
|
uniq.append(p)
|
|
if len(uniq) < 3:
|
|
return None
|
|
if len(uniq) > 4:
|
|
uniq = uniq[:4]
|
|
return uniq
|
|
|
|
|
|
def interp_scalar(p: np.ndarray, tet_idx: np.ndarray, V: np.ndarray, field: np.ndarray) -> float:
|
|
v = V[tet_idx]
|
|
f = field[tet_idx]
|
|
A = np.column_stack([v[1] - v[0], v[2] - v[0], v[3] - v[0]])
|
|
w123 = np.linalg.solve(A, p - v[0])
|
|
w0 = 1.0 - np.sum(w123)
|
|
w = np.array([w0, w123[0], w123[1], w123[2]])
|
|
return float(np.dot(w, f))
|
|
|
|
|
|
def build_slice_tris(V: np.ndarray, T: np.ndarray, field: np.ndarray, axis: str, coord: float):
|
|
ai, bi = other_axes(axis)
|
|
ax = axis_index(axis)
|
|
|
|
polys = []
|
|
values = []
|
|
for tet in T:
|
|
verts = V[tet]
|
|
poly = tet_plane_poly(verts, ax, coord)
|
|
if poly is None:
|
|
continue
|
|
poly_arr = np.array(poly)
|
|
local_vals = [interp_scalar(p, tet, V, field) for p in poly_arr]
|
|
tris = triangulate_polygon(poly)
|
|
for t0, t1, t2 in tris:
|
|
polys.append(poly_arr[[t0, t1, t2]])
|
|
values.append([local_vals[t0], local_vals[t1], local_vals[t2]])
|
|
|
|
if not polys:
|
|
return None, None, None
|
|
|
|
polys = np.array(polys)
|
|
values = np.array(values)
|
|
xy = polys.reshape(-1, 3)[:, [ai, bi]]
|
|
zval = values.reshape(-1)
|
|
tri = Triangulation(xy[:, 0], xy[:, 1], np.arange(len(zval)).reshape(-1, 3))
|
|
return tri, zval, (ai, bi)
|
|
|
|
|
|
def build_surface_tris(V: np.ndarray, Tri: np.ndarray, field: np.ndarray, axis: str, coord: float, tol: float):
|
|
ai, bi = other_axes(axis)
|
|
ax = axis_index(axis)
|
|
|
|
xs, ys, zs, tris = [], [], [], []
|
|
node = 0
|
|
for tri in Tri:
|
|
verts = V[tri]
|
|
if not np.all(np.abs(verts[:, ax] - coord) < tol):
|
|
continue
|
|
xs.extend(verts[:, ai])
|
|
ys.extend(verts[:, bi])
|
|
zs.extend(field[tri])
|
|
tris.append([node, node + 1, node + 2])
|
|
node += 3
|
|
|
|
if not tris:
|
|
return None, None, None
|
|
|
|
xy = np.column_stack([xs, ys])
|
|
tri = Triangulation(xy[:, 0], xy[:, 1], np.array(tris))
|
|
return tri, np.array(zs), (ai, bi)
|
|
|
|
|
|
def plot_slice(tri, zval, plane_axes, out_path: Path, coord: float, axis: str, title: str):
|
|
ai, bi = plane_axes
|
|
names = ["x", "y", "z"]
|
|
fig, ax = plt.subplots(figsize=(7.5, 6.5), dpi=150)
|
|
tpc = ax.tripcolor(tri, zval, shading="gouraud", cmap="jet", edgecolors="k", linewidth=0.15)
|
|
ax.set_aspect("equal")
|
|
ax.set_xlabel(f"{names[ai]} (m)")
|
|
ax.set_ylabel(f"{names[bi]} (m)")
|
|
ax.set_title(title)
|
|
cbar = fig.colorbar(tpc, ax=ax, fraction=0.046, pad=0.04)
|
|
cbar.set_label("|E| (V/m)")
|
|
fig.text(0.02, 0.98, f"lambda0 = 0.8 m\n{axis} = {coord} m", va="top", fontsize=9)
|
|
fig.tight_layout()
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
fig.savefig(out_path, bbox_inches="tight")
|
|
plt.close(fig)
|
|
print(f"saved: {out_path}")
|
|
print(f" |E| min={zval.min():.4f} max={zval.max():.4f} mean={zval.mean():.4f}")
|
|
|
|
|
|
def main():
|
|
root = Path(__file__).resolve().parents[1]
|
|
parser = argparse.ArgumentParser(description="Plot normE on a plane slice")
|
|
parser.add_argument("--mesh", type=Path, default=root / "build" / "Release" / "SBCmesh.dat")
|
|
parser.add_argument("--outdir", type=Path, default=root / "build" / "Release" / "OutFile")
|
|
parser.add_argument("--plane", choices=["x", "y", "z"], default="z")
|
|
parser.add_argument("--coord", type=float, default=0.5, help="plane coordinate (m)")
|
|
parser.add_argument("--mode", choices=["slice", "surface"], default="slice",
|
|
help="slice=cut volume; surface=boundary tris on plane")
|
|
parser.add_argument("--tol", type=float, default=1e-3, help="tolerance for surface mode")
|
|
parser.add_argument("--output", type=Path, default=None)
|
|
args = parser.parse_args()
|
|
|
|
V, T, Tri, _ = load_mesh(args.mesh)
|
|
normE = load_scalar_field(args.outdir / "normE")
|
|
if normE.size != V.shape[0]:
|
|
raise SystemExit(f"normE size {normE.size} != vertex count {V.shape[0]}")
|
|
|
|
if args.output is None:
|
|
args.output = args.outdir / f"normE_{args.plane}{args.coord:.3f}.png"
|
|
|
|
if args.mode == "surface":
|
|
tri, zval, plane_axes = build_surface_tris(V, Tri, normE, args.plane, args.coord, args.tol)
|
|
title = f"normE on {args.plane}={args.coord} m (surface)"
|
|
else:
|
|
tri, zval, plane_axes = build_slice_tris(V, T, normE, args.plane, args.coord)
|
|
title = f"normE on {args.plane}={args.coord} m (slice)"
|
|
|
|
if tri is None:
|
|
raise SystemExit(f"no data on plane {args.plane}={args.coord}")
|
|
|
|
plot_slice(tri, zval, plane_axes, args.output, args.coord, args.plane, title)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|