"""Slices of scalar data on XGC poloidal and toroidal triangular meshes.
Surface indices are the zero-based indices of ``Plane.surf_idx``. Angles are
the straight field-line angle stored in ``Plane.theta`` and the physical
toroidal plane angle ``plane_index * Mesh.delta_phi``.
"""
from __future__ import annotations
import numpy as np
from matplotlib.tri import LinearTriInterpolator
from scipy.sparse import csr_matrix
[docs]
def flux_surface_overlay_paths(plane):
"""Return ordered vertex paths for all R-Z flux-surface overlays.
Valid straight field-line angles set the order. XGC assigns ``theta=-10``
to SOL vertices, so those paths are traced through mesh edges shared by
vertices of the same surface. Separate paths are returned separately to
avoid lines across disconnected surface pieces.
Parameters
----------
plane
Mesh Plane whose flux-surface paths are requested.
Returns
-------
list[np.ndarray]
ordered vertex paths, split at disconnected pieces.
"""
adjacency = None
for surface_index in range(plane.nsurf):
vertices = np.asarray(plane.get_surface_vertex_indices(surface_index), dtype=int)
if vertices.size < 2:
continue
theta = np.asarray(plane.theta[vertices], dtype=float)
if np.all(np.isfinite(theta) & (theta >= -2*np.pi)):
ordered = vertices[np.argsort(theta, kind="stable")]
psi_n = float(plane.psi_surf[surface_index]/plane.x_psi)
if psi_n < 1:
ordered = np.append(ordered, ordered[0])
yield ordered
continue
if adjacency is None:
triangles = np.asarray(plane.nd_connect_list, dtype=int)
edges = np.concatenate((triangles[:, [0, 1]], triangles[:, [1, 2]],
triangles[:, [2, 0]]))
src = np.concatenate((edges[:, 0], edges[:, 1]))
dst = np.concatenate((edges[:, 1], edges[:, 0]))
adjacency = csr_matrix((np.ones(src.size, dtype=bool), (src, dst)),
shape=(plane.n_n, plane.n_n))
graph = adjacency[vertices][:, vertices]
degree = np.diff(graph.indptr)
remaining = set(range(vertices.size))
while remaining:
endpoints = [i for i in remaining if degree[i] <= 1]
current = min(endpoints) if endpoints else min(remaining)
path = []
while current in remaining:
remaining.remove(current)
path.append(current)
neighbors = graph.indices[graph.indptr[current]:graph.indptr[current + 1]]
following = [int(i) for i in neighbors if i in remaining]
if not following or (len(following) > 1 and len(path) > 1):
break
current = min(following)
if len(path) > 1:
tail_neighbors = graph.indices[graph.indptr[path[-1]]:graph.indptr[path[-1] + 1]]
if path[0] in tail_neighbors:
path.append(path[0])
yield vertices[path]
def _surface_samples(plane, values, surface_index, *, closed_loop=True):
"""Return finite theta/value pairs for one mesh surface.
The separatrix legs are excluded when ``closed_loop`` is true. Open SOL
surfaces are otherwise kept; their samples must not be wrapped across the
theta seam.
"""
if not 0 <= surface_index < plane.nsurf:
raise ValueError(f"Surface index must be between 0 and {plane.nsurf - 1}")
vertices = np.asarray(plane.get_surface_vertex_indices(surface_index), dtype=int)
psi_n = float(plane.psi_surf[surface_index]/plane.x_psi)
if closed_loop and np.isclose(psi_n, 1.0, atol=1e-6):
vertices = vertices[plane.rz[vertices, 1] >= plane.x_z]
theta = np.asarray(plane.theta[vertices], dtype=float)
samples = np.asarray(values[vertices], dtype=float)
valid = np.isfinite(theta) & np.isfinite(samples) & (theta >= -2*np.pi)
theta, samples = theta[valid], samples[valid]
if theta.size < 2:
raise ValueError(f"Surface {surface_index} has fewer than two valid samples")
theta = (theta + np.pi) % (2*np.pi) - np.pi
order = np.argsort(theta)
theta, samples = theta[order], samples[order]
unique, inverse = np.unique(theta, return_inverse=True)
if unique.size != theta.size:
totals = np.bincount(inverse, weights=samples)
samples = totals/np.bincount(inverse)
theta = unique
return theta, samples, psi_n
[docs]
def surface_lines(plane_data, surface_indices):
"""Return ``(index, psi_N, theta_star, values)`` for selected surfaces.
Parameters
----------
plane_data
One scalar value per mesh vertex, with its Plane geometry.
surface_indices
Zero-based indices or collection of selected flux surfaces.
Returns
-------
list[tuple]
surface index, psi, theta, and values for each selection.
"""
plane = plane_data.plane
values = np.asarray(plane_data.get_data(), dtype=float)
if values.shape != (plane.n_n,):
raise ValueError("A surface slice requires one scalar per mesh vertex")
result = []
for index in surface_indices:
theta, samples, psi_n = _surface_samples(plane, values, index)
result.append((index, psi_n, theta, samples))
return result
def _periodic_surface(plane, values, surface_index, theta_grid):
"""Interpolate a closed surface periodically onto a uniform theta grid."""
theta, samples, psi_n = _surface_samples(plane, values, surface_index)
if psi_n > 1.0 + 1e-6:
raise ValueError("A periodic theta slice requires a closed surface (psi_N <= 1)")
extended_theta = np.concatenate((theta[-1:] - 2*np.pi, theta, theta[:1] + 2*np.pi))
extended_values = np.concatenate((samples[-1:], samples, samples[:1]))
return np.interp(theta_grid, extended_theta, extended_values), psi_n
def _open_sol_surface(plane, values, surface_index, theta_grid):
"""Interpolate an open, outer-midplane SOL surface without joining its strikes.
XGC assigns theta* from 0 to 2π along the strike-to-strike path. The
maximum-R vertex locates the outer midplane; it becomes theta*=0 on the
plotted grid. The shifted coordinate stays on an open interval: wrapping
it modulo 2π would incorrectly populate angles beyond a strike point.
"""
vertices = np.asarray(plane.get_surface_vertex_indices(surface_index), dtype=int)
raw_theta = np.asarray(plane.theta[vertices], dtype=float)
samples = np.asarray(values[vertices], dtype=float)
valid = np.isfinite(raw_theta) & np.isfinite(samples) & (raw_theta >= 0) & (raw_theta <= 2*np.pi + 1e-6)
vertices, raw_theta, samples = vertices[valid], raw_theta[valid], samples[valid]
if raw_theta.size < 3 or np.ptp(raw_theta) < 1.5*np.pi:
raise ValueError(f"SOL surface {surface_index} has no full strike-to-strike theta range")
origin = raw_theta[np.argmax(plane.rz[vertices, 0])]
order = np.argsort(raw_theta)
raw_theta, samples = raw_theta[order], samples[order]
unique, inverse = np.unique(raw_theta, return_inverse=True)
if unique.size != raw_theta.size:
samples = np.bincount(inverse, weights=samples)/np.bincount(inverse)
raw_theta = unique
query = theta_grid + origin
row = np.interp(query, raw_theta, samples, left=np.nan, right=np.nan)
return row, (raw_theta[0] - origin, raw_theta[-1] - origin)
[docs]
def surface_phi_slice(mesh_data, surface_index, n_theta=256):
"""Interpolate all toroidal planes at one closed surface to common theta*.
Parameters
----------
mesh_data
Vertex data on all toroidal planes.
surface_index
Zero-based selected flux-surface index.
n_theta
Number of points on the uniform theta-star grid.
Returns
-------
tuple
theta grid, toroidal angles, sampled values, and surface psi.
"""
mesh = mesh_data.mesh
if mesh.nphi < 2:
raise ValueError("Theta-phi view requires at least two toroidal planes")
if not 16 <= n_theta <= 2048:
raise ValueError("Theta resolution must be between 16 and 2048")
theta = np.linspace(-np.pi, np.pi, n_theta, endpoint=False)
rows = []
psi_n = None
for i in range(mesh.nphi):
plane_data = mesh_data.get_plane(i)
row, psi_n = _periodic_surface(
plane_data.plane, np.asarray(plane_data.get_data(), dtype=float),
surface_index, theta)
rows.append(row)
phi = np.arange(mesh.nphi, dtype=float)*mesh.delta_phi
return theta, phi, np.asarray(rows), psi_n
[docs]
def psi_theta_slice(plane_data, *, n_psi=256, n_theta=256, psi_bounds=(0.1, 1.05),
theta_domain=(-np.pi, np.pi), return_boundaries=False):
"""Map outer-midplane surfaces to a uniform ``(psi_N, theta*)`` grid.
Interpolation follows ``mesh2uni.pro``: periodic interpolation along each
core surface, then linear interpolation in normalized poloidal flux. The
separatrix uses only vertices with defined theta*. Open SOL surfaces use
their strike-to-strike theta* with zero shifted to the outer midplane.
Surfaces with unset theta* are skipped. ``theta_domain`` can span more
than one turn to show both strike endpoints; closed surfaces occupy only
the single [-π, π] turn. With ``return_boundaries=True`` the fourth return value is
an array of ``(psi_N, first endpoint, second endpoint)`` SOL coordinates.
Parameters
----------
plane_data
One scalar value per mesh vertex, with its Plane geometry.
n_psi
Number of points on the uniform normalized-psi grid.
n_theta
Number of points on the uniform theta-star grid.
psi_bounds
Minimum and maximum normalized psi for interpolation.
theta_domain
Angular interval for the output grid, in radians.
return_boundaries
Also return open-SOL strike-point angle boundaries.
Returns
-------
tuple
psi grid, theta grid, mapped values, and optional SOL boundaries.
"""
if not 16 <= n_psi <= 2048 or not 16 <= n_theta <= 2048:
raise ValueError("Psi and theta resolutions must be between 16 and 2048")
lo, hi = psi_bounds
if not 0 <= lo < hi:
raise ValueError("Psi-theta bounds must satisfy 0 <= min < max")
theta_lo, theta_hi = theta_domain
if not np.isfinite(theta_lo) or not np.isfinite(theta_hi) or theta_lo >= theta_hi:
raise ValueError("Theta domain must have finite increasing bounds")
plane = plane_data.plane
values = np.asarray(plane_data.get_data(), dtype=float)
if values.shape != (plane.n_n,):
raise ValueError("Psi-theta view requires one scalar per mesh vertex")
theta = np.linspace(theta_lo, theta_hi, n_theta, endpoint=False)
rows, radial, boundaries = [], [], []
for index in np.asarray(plane.surf_map, dtype=int):
psi_n = float(plane.psi_surf[index]/plane.x_psi)
if psi_n < 0:
continue
try:
if psi_n <= 1 + 1e-6:
row, _ = _periodic_surface(plane, values, int(index), theta)
row[(theta < -np.pi) | (theta >= np.pi)] = np.nan
else:
row, (first, second) = _open_sol_surface(plane, values, int(index), theta)
boundaries.append((psi_n, first, second))
except ValueError:
continue
rows.append(row)
radial.append(psi_n)
if len(radial) < 2:
raise ValueError("Psi-theta mapping requires at least two mapped surfaces with defined theta")
order = np.argsort(radial)
radial = np.asarray(radial)[order]
rows = np.asarray(rows)[order]
radial, unique_index = np.unique(radial, return_index=True)
rows = rows[unique_index]
psi = np.linspace(lo, hi, n_psi)
mapped = np.empty((n_psi, n_theta), dtype=float)
for j in range(n_theta):
mapped[:, j] = np.interp(psi, radial, rows[:, j], left=np.nan, right=np.nan)
# Undefined SOL angles, including those beyond strike points, are zero.
sol = psi > 1.0
mapped[sol] = np.where(np.isfinite(mapped[sol]), mapped[sol], 0.0)
if return_boundaries:
return psi, theta, mapped, np.asarray(sorted(boundaries), dtype=float).reshape(-1, 3)
return psi, theta, mapped
[docs]
def line_of_sight(plane_data, psi_plane_data, *, r0, z0, slope=0.0,
angle_degrees=None, n_points=500):
"""Sample field and normalized psi along a line through ``(r0,z0)``.
``angle_degrees`` measures the line direction from +R toward +Z and must
lie in [-90, 90]. The legacy ``slope=dZ/dR`` argument remains available
when no angle is supplied. Samples outside the triangulated XGC domain
are returned as NaN.
Parameters
----------
plane_data
One scalar value per mesh vertex, with its Plane geometry.
psi_plane_data
Node-centered poloidal flux on the same Plane.
r0
R coordinate of the sight-line origin in meters.
z0
Z coordinate of the sight-line origin in meters.
slope
Legacy dZ/dR direction when angle_degrees is omitted.
angle_degrees
Sight-line angle from +R, within [-90, 90] degrees.
n_points
Number of samples along the sight line.
Returns
-------
tuple[np.ndarray, ...]
R, Z, normalized psi, and sampled field.
"""
if not 2 <= n_points <= 10000:
raise ValueError("Line resolution must be between 2 and 10000")
if angle_degrees is not None:
if not np.isfinite(angle_degrees) or not -90 <= angle_degrees <= 90:
raise ValueError("Line angle must be between -90 and 90 degrees")
slope = np.inf if abs(angle_degrees) == 90 else np.tan(np.deg2rad(angle_degrees))
plane = plane_data.plane
rz = np.asarray(plane.rz, dtype=float)
rmin, rmax = np.min(rz[:, 0]), np.max(rz[:, 0])
zmin, zmax = np.min(rz[:, 1]), np.max(rz[:, 1])
if np.isinf(slope):
if not rmin <= r0 <= rmax:
raise ValueError("Vertical line does not cross the mesh R range")
r = np.full(n_points, r0)
z = np.linspace(zmin, zmax, n_points)
else:
if not np.isfinite(slope):
raise ValueError("Line slope must be finite or inf")
low, high = rmin-r0, rmax-r0
if slope != 0:
zl, zh = sorted(((zmin-z0)/slope, (zmax-z0)/slope))
low, high = max(low, zl), min(high, zh)
elif not zmin <= z0 <= zmax:
raise ValueError("Horizontal line does not cross the mesh Z range")
if low >= high:
raise ValueError("Line of sight does not cross the mesh bounds")
r = r0 + np.linspace(low, high, n_points)
z = z0 + slope*(r-r0)
tri = plane.get_triangulation_data(n_int=2).triObj
values = np.asarray(plane_data.get_data(), dtype=float)
psi_values = np.asarray(psi_plane_data.get_data(), dtype=float)
if values.shape != (plane.n_n,) or psi_values.shape != (plane.n_n,):
raise ValueError("Line of sight requires scalar field and psi at each vertex")
data = np.ma.filled(LinearTriInterpolator(tri, values)(r, z), np.nan)
psi = np.ma.filled(LinearTriInterpolator(tri, psi_values)(r, z), np.nan)/plane.x_psi
return r, z, np.asarray(psi), np.asarray(data)
[docs]
def split_line_at_psi_minimum(psi, values, *, ignore_end_points=5):
"""Split a sampled sight line into its two branches around minimum ψ_N.
Nonfinite samples do not take part in finding the minimum. A configurable
number of finite samples at each end are also excluded from that search,
to avoid dead boundary points. The returned slices retain their original
order and share the turning point.
Parameters
----------
psi
Normalized psi samples along a sight line.
values
Field samples aligned with psi.
ignore_end_points
Finite points ignored at each end when finding minimum psi.
Returns
-------
tuple
two (psi, value) branches and the turning index.
"""
psi = np.asarray(psi, dtype=float)
values = np.asarray(values, dtype=float)
if psi.shape != values.shape or psi.ndim != 1:
raise ValueError("Psi and line values must be matching one-dimensional arrays")
if ignore_end_points < 0:
raise ValueError("Ignored end samples must be nonnegative")
finite = np.flatnonzero(np.isfinite(psi) & np.isfinite(values))
if finite.size < 3:
raise ValueError("At least three finite line samples are needed to split psi branches")
ignore = min(int(ignore_end_points), (finite.size - 3)//2)
candidates = finite[ignore:finite.size - ignore] if ignore else finite
turning = int(candidates[np.argmin(psi[candidates])])
return (psi[:turning + 1], values[:turning + 1]), (psi[turning:], values[turning:]), turning
[docs]
def flux_surface_profile(plane_data, statistic="average", *, fsa="volume"):
"""Return a radial mean or poloidal residual RMS on all flux surfaces.
The RMS back-projects the surface mean to vertices, then computes
``sqrt(<(X - <X>)²>)`` with the same flux-surface weighting. This avoids
cancellation between two large moments when the variation is small.
Parameters
----------
plane_data : PlaneData
Scalar data at every vertex of one poloidal plane.
statistic : {"average", "rms"}
Flux-surface average or poloidal residual RMS.
fsa : str
Plane flux-surface averaging convention.
Returns
-------
ndarray
One value per flux surface in the input field's units.
"""
if statistic not in {"average", "rms"}:
raise ValueError("statistic must be 'average' or 'rms'")
plane = plane_data.plane
values = np.asarray(plane_data.get_data(), dtype=float)
if values.shape != (plane.n_n,):
raise ValueError("A flux-surface profile requires one scalar per vertex")
mean = np.asarray(plane.flux_avg_to_surf(values, fsa=fsa), dtype=float)
if statistic == "average":
return mean
mean_at_vertices = np.asarray(plane.flux_avg_from_surf(mean, fsa=fsa), dtype=float)
residual = values - mean_at_vertices
variance = np.asarray(plane.flux_avg_to_surf(residual*residual, fsa=fsa), dtype=float)
return np.sqrt(np.maximum(variance, 0.0))