"""Reusable mapped R-Z diagnostic plots for XGC monitor and Python workflows."""
from __future__ import annotations
import numpy as np
[docs]
def plot_mesh_time_map(
time_map, *, variable, description="", units="", color_bounds=None,
coordinate_bounds=None, time_bounds_ms=None,
):
"""Plot a mesh slice over physics time with a normalized-psi or theta-star axis.
Parameters
----------
time_map : MeshSliceTimeMap
Precomputed spatial slices and their physical times.
variable, description, units : str
Field label, optional verbose title, and display units.
color_bounds, coordinate_bounds, time_bounds_ms : tuple or None
Optional color, horizontal, and millisecond axis limits.
Returns
-------
matplotlib.figure.Figure
Contour-ready time map rendered as a pseudocolor figure.
"""
from matplotlib import pyplot as plt
coordinate = np.asarray(time_map.coordinate, dtype=float)
time_ms = np.asarray(time_map.time_s, dtype=float)*1.0e3
values = np.asarray(time_map.values, dtype=float)
if values.shape != (time_ms.size, coordinate.size):
raise ValueError("Time-map values do not match their coordinate axes")
if not np.any(np.isfinite(values)):
raise ValueError("The selected time map has no finite field values")
fig, ax = plt.subplots(figsize=(8, 5.5))
image = ax.pcolormesh(
coordinate, time_ms, values, shading="auto", cmap="viridis",
vmin=None if color_bounds is None else color_bounds[0],
vmax=None if color_bounds is None else color_bounds[1],
rasterized=True,
)
quantity = (f"RMS({variable} − ⟨{variable}⟩)" if time_map.statistic == "rms"
else f"⟨{variable}⟩" if time_map.statistic == "average" else variable)
fig.colorbar(image, ax=ax, label=f"{quantity} [{units}]" if units else quantity)
ax.set_xlabel(r"$\theta^*$ (rad)" if time_map.coordinate_name == "theta" else r"$\psi_N$")
ax.set_ylabel("Time (ms)")
if coordinate_bounds is not None:
ax.set_xlim(*coordinate_bounds)
if time_bounds_ms is not None:
ax.set_ylim(*time_bounds_ms)
fig.suptitle(f"{description or quantity} — {time_map.detail}")
fig.tight_layout()
return fig
[docs]
def plot_rz_diagnostic(
plane_data,
*,
variable: str,
description: str = "",
units: str = "",
split_average: bool = False,
view: str | None = None,
n_int: int = 300,
r_bounds=None,
z_bounds=None,
color_bounds=None,
residual_color_bounds=None,
profile_psi_bounds=None,
profile_value_bounds=None,
show_wall: bool = True,
show_mesh: bool = False,
show_flux: bool = False,
):
"""Plot a full field, mapped radial mean or RMS, or its R-Z residual.
``plane_data`` must hold one scalar per vertex. The surface average and
residual use its existing flux-surface conversion methods. Radial views
select the outer-midplane surfaces in ``Plane.surf_map``. ``view`` also
accepts ``split`` for existing Python callers; the GUI presents the
average and residual as separate choices. R-Z fields are mapped to a
bitmap grid by :meth:`Plane.plot_contour`.
Parameters
----------
plane_data
One scalar value per mesh vertex, with its Plane geometry.
variable
Display name or variable key for plot labeling.
description
Optional descriptive color-scale label.
units
Physical units for plot labeling.
split_average
Legacy switch for a combined average-and-residual view.
view
Named plot view.
n_int
Number of R-Z raster grid points per axis.
r_bounds
Optional R-axis bounds in meters.
z_bounds
Optional Z-axis bounds in meters.
color_bounds
Optional field color limits.
residual_color_bounds
Optional residual color limits.
profile_psi_bounds
Optional normalized-psi axis limits.
profile_value_bounds
Optional profile-value axis limits.
show_wall
Draw the wall outline.
show_mesh
Draw triangular mesh edges.
show_flux
Draw ordered flux-surface paths.
Returns
-------
matplotlib.figure.Figure
rendered R-Z field or surface profile.
"""
from matplotlib import pyplot as plt
plane = plane_data.plane
values = np.asarray(plane_data.get_data(), dtype=float)
if values.shape != (plane.n_n,):
raise ValueError(f"Expected one value for each of {plane.n_n} mesh vertices; got {values.shape}")
if not np.any(np.isfinite(values)):
raise ValueError("The selected field has no finite vertex values")
if not 2 <= n_int <= 2048:
raise ValueError("R-Z grid resolution must be between 2 and 2048 points per axis")
view = view or ("split" if split_average else "full")
if view not in {"full", "average", "rms", "residual", "split"}:
raise ValueError(f"Unknown R-Z diagnostic view: {view}")
if view != "full":
from .mesh_slicing import flux_surface_profile
average = flux_surface_profile(plane_data, "average")
if view == "rms":
profile = flux_surface_profile(plane_data, "rms")
if view not in {"average", "rms"}:
plane_data.flux_avg_to_surf()
plane_data.flux_avg_from_surf()
field = values - plane_data.flux_avg_plane
if view == "split":
fig, (line_ax, rz_ax) = plt.subplots(1, 2, figsize=(12, 5.5),
gridspec_kw={"width_ratios": [1, 1.3]})
elif view in {"average", "rms"}:
fig, line_ax = plt.subplots(figsize=(7, 5.5))
else:
fig, rz_ax = plt.subplots(figsize=(7, 5.5))
mapped = np.asarray(plane.surf_map, dtype=int)
if view in {"split", "average", "rms"}:
line_ax.plot(plane.psi_surf[mapped]/plane.x_psi,
profile[mapped] if view == "rms" else average[mapped], linewidth=1.7)
line_ax.set_xlabel(r"$\psi_N$")
quantity = f"RMS({variable} − ⟨{variable}⟩)" if view == "rms" else f"⟨{variable}⟩"
line_ax.set_ylabel(f"{quantity} [{units}]" if units else quantity)
line_ax.set_title("Poloidal residual RMS" if view == "rms" else "Flux-surface average")
line_ax.grid(alpha=0.25)
if profile_psi_bounds is not None:
line_ax.set_xlim(*profile_psi_bounds)
if profile_value_bounds is not None:
line_ax.set_ylim(*profile_value_bounds)
if view in {"average", "rms"}:
fig.suptitle(description or variable)
fig.tight_layout()
return fig
color_label = f"{variable} − ⟨{variable}⟩"
bounds = residual_color_bounds
else:
field = values
fig, rz_ax = plt.subplots(figsize=(7, 5.5))
color_label = variable
bounds = color_bounds
color_label = f"{color_label} [{units}]" if units else color_label
from xgc_analysis.mesh import PlaneData
field_data = plane_data if view == "full" else PlaneData(plane, data_array=field)
field_data.plot_contour(
n_int=n_int, var_name=color_label,
title="R-Z residual" if view != "full" else "R-Z field",
lower=None if bounds is None else bounds[0],
upper=None if bounds is None else bounds[1],
R_bounds=r_bounds, Z_bounds=z_bounds,
ax=rz_ax, render_as_image=True,
cmap="viridis",
)
if show_mesh:
rz_ax.triplot(plane.rz[:, 0], plane.rz[:, 1], plane.nd_connect_list,
linewidth=0.12, color="#263238", alpha=0.35)
if show_wall and hasattr(plane, "wall_nodes"):
wall = plane.rz[np.asarray(plane.wall_nodes, dtype=int)]
if len(wall) >= 2:
if not np.allclose(wall[0], wall[-1]):
wall = np.vstack((wall, wall[0]))
rz_ax.plot(wall[:, 0], wall[:, 1], color="#cc4b37", linewidth=1.1)
if show_flux:
for surface in plane.surface_plot_paths:
points = plane.rz[surface]
rz_ax.plot(points[:, 0], points[:, 1], color="#f4d03f", linewidth=0.55, alpha=0.75)
if r_bounds is not None:
rz_ax.set_xlim(*r_bounds)
if z_bounds is not None:
rz_ax.set_ylim(*z_bounds)
fig.suptitle(description or variable)
fig.tight_layout()
return fig
[docs]
def plot_mesh_slice(
plane_data, *, view, variable, description="", units="", mesh_data=None,
surface_indices=(), surface_index=0, n_theta=256, n_psi=256,
psi_bounds=(0.1, 1.05), remove_fsa=False, color_bounds=None, theta_bounds=None,
phi_bounds=None, surface_value_bounds=None, line_psi_data=None,
line_origin=None, line_r0=None, line_z0=None, line_slope=0.0,
line_angle_deg=None, line_xaxis="R", line_points=500, line_end_ignore=5,
line_xbounds=None, line_ybounds=None,
):
"""Plot a surface, theta-phi, psi-theta, or line-of-sight mesh slice.
The interpolation functions in :mod:`mesh_slicing` are also usable in
Python workflows without creating a figure.
Parameters
----------
plane_data
One scalar value per mesh vertex, with its Plane geometry.
view
Named plot view.
variable
Display name or variable key for plot labeling.
description
Optional descriptive color-scale label.
units
Physical units for plot labeling.
mesh_data
Vertex data on all toroidal planes.
surface_indices
Zero-based indices or collection of selected flux surfaces.
surface_index
Zero-based selected flux-surface index.
n_theta
Number of points on the uniform theta-star grid.
n_psi
Number of points on the uniform normalized-psi grid.
psi_bounds
Minimum and maximum normalized psi for interpolation.
remove_fsa
Subtract the flux-surface average before mapping.
color_bounds
Optional field color limits.
theta_bounds
Optional theta-star axis limits in radians.
phi_bounds
Optional toroidal-angle axis limits in radians.
surface_value_bounds
Optional surface-trace value limits.
line_psi_data
Poloidal-flux PlaneData for line-of-sight sampling.
line_origin
Magnetic-axis, X-point, or custom origin selector.
line_r0
Custom R coordinate in meters.
line_z0
Custom Z coordinate in meters.
line_slope
Legacy line direction dZ/dR.
line_angle_deg
Line direction from +R in degrees.
line_xaxis
R, Z, or normalized psi horizontal axis.
line_points
Number of sight-line samples.
line_end_ignore
Endpoint samples excluded from psi-minimum search.
line_xbounds
Optional sight-line horizontal limits.
line_ybounds
Optional sight-line value limits.
Returns
-------
matplotlib.figure.Figure
rendered mesh slice.
"""
from matplotlib import pyplot as plt
from .mesh_slicing import (line_of_sight, psi_theta_slice,
split_line_at_psi_minimum,
surface_lines, surface_phi_slice)
quantity = f"{variable} − ⟨{variable}⟩" if view == "psi_theta" and remove_fsa else variable
label = f"{quantity} [{units}]" if units else quantity
fig, ax = plt.subplots(figsize=(7.5, 5.5))
if view == "surface_lines":
if not surface_indices:
raise ValueError("Enter at least one zero-based surface index")
for index, psi_n, theta, values in surface_lines(plane_data, surface_indices):
ax.plot(theta, values, label=fr"surface {index}, $\psi_N$={psi_n:.3f}")
ax.set_xlabel(r"$\theta^*$ (rad)")
ax.set_ylabel(label)
ax.legend()
ax.grid(alpha=0.25)
if theta_bounds is not None:
ax.set_xlim(*theta_bounds)
if surface_value_bounds is not None:
ax.set_ylim(*surface_value_bounds)
elif view in {"surface_phi", "psi_theta"}:
if view == "surface_phi":
if mesh_data is None:
raise ValueError("Theta-phi view requires 3D MeshData")
theta, phi, mapped, psi_n = surface_phi_slice(mesh_data, surface_index, n_theta)
extent = (theta[0], theta[0] + 2*np.pi,
phi[0], phi[-1] + mesh_data.mesh.delta_phi)
ax.set_ylabel(r"$\phi$ (rad)")
ax.set_title(fr"Surface {surface_index}, $\psi_N$={psi_n:.3f}")
if phi_bounds is not None:
ax.set_ylim(*phi_bounds)
else:
if remove_fsa:
plane_data.flux_avg_to_surf()
plane_data.flux_avg_from_surf()
from xgc_analysis.mesh import PlaneData
residual = np.asarray(plane_data.get_data(), dtype=float) - plane_data.flux_avg_plane
plane_data = PlaneData(plane_data.plane, data_array=residual)
psi, theta, mapped, strike_bounds = psi_theta_slice(
plane_data, n_psi=n_psi, n_theta=n_theta, psi_bounds=psi_bounds,
theta_domain=(-np.pi, 2*np.pi), return_boundaries=True)
extent = (theta[0], theta[0] + n_theta*(theta[1] - theta[0]), psi[0], psi[-1])
ax.set_ylabel(r"$\psi_N$")
image = ax.imshow(mapped, origin="lower", aspect="auto", extent=extent,
cmap="viridis", interpolation="nearest",
vmin=None if color_bounds is None else color_bounds[0],
vmax=None if color_bounds is None else color_bounds[1])
fig.colorbar(image, ax=ax, label=label)
if view == "psi_theta" and strike_bounds.size:
from matplotlib import patheffects
for column, name in ((1, "SOL strike endpoint 1"), (2, "SOL strike endpoint 2")):
line, = ax.plot(strike_bounds[:, column], strike_bounds[:, 0],
color="#161616", linestyle="--", linewidth=1.4,
label=name, zorder=5)
line.set_path_effects([patheffects.withStroke(linewidth=3, foreground="white")])
ax.legend(loc="upper left", fontsize="small")
ax.set_xlabel(r"$\theta^*$ (rad)")
if theta_bounds is not None:
ax.set_xlim(*theta_bounds)
elif view == "line_of_sight":
if line_psi_data is None:
raise ValueError("Line-of-sight view requires a nodal psi PlaneData")
plane = plane_data.plane
if line_origin == "axis":
r0, z0 = plane.axis_r, plane.axis_z
elif line_origin == "xpoint":
r0, z0 = plane.x_r, plane.x_z
elif line_origin == "custom":
if line_r0 is None or line_z0 is None:
raise ValueError("Custom line origin requires both R and Z")
r0, z0 = line_r0, line_z0
elif line_origin is None:
r0 = plane.axis_r if line_r0 is None else line_r0
z0 = plane.axis_z if line_z0 is None else line_z0
else:
raise ValueError("Line origin must be axis, xpoint, or custom")
r, z, psi, values = line_of_sight(
plane_data, line_psi_data, r0=r0, z0=z0,
slope=line_slope, angle_degrees=line_angle_deg, n_points=line_points)
coordinate = {"R": r, "Z": z, "psi": psi}[line_xaxis]
xlabel = {"R": "R (m)", "Z": "Z (m)", "psi": r"$\psi_N$"}[line_xaxis]
if line_xaxis == "psi":
toward, away, _ = split_line_at_psi_minimum(
psi, values, ignore_end_points=line_end_ignore)
ax.plot(*toward, linewidth=1.6, color="tab:blue", label="Toward minimum ψ_N")
ax.plot(*away, linewidth=1.6, color="tab:orange", label="Away from minimum ψ_N")
ax.legend()
else:
ax.plot(coordinate, values, linewidth=1.6)
ax.set_xlabel(xlabel)
ax.set_ylabel(label)
ax.grid(alpha=0.25)
if line_xbounds is not None:
ax.set_xlim(*line_xbounds)
if line_ybounds is not None:
ax.set_ylim(*line_ybounds)
else:
plt.close(fig)
raise ValueError(f"Unknown mesh slice view: {view}")
fig.suptitle(description or variable)
fig.tight_layout()
return fig