You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
107 lines
4.0 KiB
107 lines
4.0 KiB
#!/usr/bin/env python3
|
|
import argparse
|
|
import os
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
|
|
def plot_quiver_heads(
|
|
df,
|
|
outdir,
|
|
normalize=True,
|
|
subsample=1,
|
|
eps=0.1, # base length scale (increase for longer arrows)
|
|
headlength=30, headwidth=12, headaxislength=12,
|
|
cmap="jet", width=0.015,
|
|
length_mode="scaled" # "scaled" -> length ∝ |E|; "fixed" -> same length
|
|
):
|
|
os.makedirs(outdir, exist_ok=True)
|
|
|
|
if subsample > 1:
|
|
df = df.iloc[::subsample, :].copy()
|
|
|
|
for p in sorted(df['port_idx'].unique()):
|
|
sub = df[df['port_idx'] == p]
|
|
if sub.empty:
|
|
continue
|
|
|
|
x = sub['x'].to_numpy()
|
|
y = sub['y'].to_numpy()
|
|
u = sub['Et_x'].to_numpy()
|
|
v = sub['Et_y'].to_numpy()
|
|
|
|
mag = np.hypot(u, v) # color by magnitude
|
|
|
|
# Normalize direction only (optional)
|
|
if normalize:
|
|
nz = mag > 0
|
|
u[nz] /= mag[nz]
|
|
v[nz] /= mag[nz]
|
|
|
|
# Determine drawn length
|
|
if length_mode == "scaled":
|
|
# Length scales with |E| relative to max, multiplied by eps
|
|
maxmag = np.max(mag) if np.any(np.isfinite(mag)) else 1.0
|
|
sf = eps * (mag / maxmag if maxmag > 0 else 0.0)
|
|
u_draw = u * sf
|
|
v_draw = v * sf
|
|
else:
|
|
# Fixed tiny shafts for heads-only look; increase eps to see longer arrows
|
|
u_draw = u * eps
|
|
v_draw = v * eps
|
|
|
|
fig, ax = plt.subplots(figsize=(6, 6))
|
|
ax.quiver(
|
|
x, y, u_draw, v_draw, mag, # mag provides the color
|
|
cmap=cmap,
|
|
angles="xy", scale_units="xy", scale=1.0,
|
|
pivot="tip",
|
|
headlength=headlength, headwidth=headwidth, headaxislength=headaxislength,
|
|
width=width
|
|
)
|
|
|
|
ax.set_title(f"Port {int(p)} field quiver (heads only, colored)")
|
|
ax.set_xlabel("x")
|
|
ax.set_ylabel("y")
|
|
ax.set_aspect("equal", adjustable="box")
|
|
ax.grid(True, ls="--", alpha=0.4)
|
|
|
|
outpath = os.path.join(outdir, f"port_{int(p)}_quiver_heads_colored.png")
|
|
fig.savefig(outpath, dpi=220, bbox_inches="tight")
|
|
plt.close(fig)
|
|
print(f"Saved {outpath}")
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Plot per-port quiver (arrow heads only) of Et")
|
|
ap.add_argument("--csv", required=True, help="CSV with columns: port_idx,x,y,Et_x,Et_y")
|
|
ap.add_argument("--outdir", default="figs", help="Output directory")
|
|
ap.add_argument("--normalize", action="store_true", help="Normalize vectors (direction only)")
|
|
ap.add_argument("--subsample", type=int, default=1, help="Use every Nth row")
|
|
ap.add_argument("--eps", type=float, default=0.1, help="Base length scale. Increase for longer arrows")
|
|
ap.add_argument("--headlength", type=float, default=65, help="Arrow headlength (points)")
|
|
ap.add_argument("--headwidth", type=float, default=24, help="Arrow headwidth (points)")
|
|
ap.add_argument("--headaxislength", type=float, default=12, help="Arrow headaxislength (points)")
|
|
ap.add_argument("--cmap", type=str, default="jet", help="Colormap for magnitudes")
|
|
ap.add_argument("--width", type=float, default=0.015, help="Arrow line width (thickness)")
|
|
ap.add_argument("--length-mode", choices=["fixed","scaled"], default="scaled",
|
|
help="fixed: constant arrow length = eps; scaled: length ∝ |E| * eps")
|
|
args = ap.parse_args()
|
|
|
|
df = pd.read_csv(args.csv)
|
|
required = {"port_idx", "x", "y", "Et_x", "Et_y"}
|
|
if not required.issubset(df.columns):
|
|
raise ValueError(f"CSV missing required columns: {sorted(required)}")
|
|
|
|
plot_quiver_heads(
|
|
df, args.outdir,
|
|
normalize=args.normalize,
|
|
subsample=args.subsample,
|
|
eps=args.eps,
|
|
headlength=args.headlength, headwidth=args.headwidth, headaxislength=args.headaxislength,
|
|
cmap=args.cmap, width=args.width,
|
|
length_mode=args.length_mode
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|