"""Reusable midplane geometry for divertor and radial transport analyses."""
from __future__ import annotations
import numpy as np
def _ordered_surface_vertices(plane, isurf: int) -> tuple[np.ndarray, bool]:
"""Return surface vertices ordered by straight-field-line poloidal angle.
The returned boolean indicates whether the surface should be closed when
constructing interpolation segments. Core surfaces are closed; open SOL or
private-flux surfaces are treated as polylines.
"""
vids = np.asarray(plane.get_surface_vertex_indices(isurf), dtype=int)
if vids.size == 0:
return vids, False
if hasattr(plane, "theta"):
theta = np.asarray(plane.theta[vids], dtype=float)
finite = np.isfinite(theta)
if np.any(finite):
vids = vids[finite][np.argsort(theta[finite])]
closed = False
try:
closed = int(np.asarray(plane.region)[int(vids[0])]) == 1
except Exception:
closed = True
return vids, closed
def _surface_midplane_crossings(plane, isurf: int, z_mid: float) -> tuple[float | None, float | None]:
"""Interpolate one flux surface to the outer and inner midplane crossings.
Each surface is traversed as a polyline in increasing ``theta``. Segment
crossings with ``Z=z_mid`` are linearly interpolated in ``(R, Z)``. The
largest crossing radius is the outboard midplane radius and the smallest
crossing radius is the inboard midplane radius.
"""
vids, closed = _ordered_surface_vertices(plane, isurf)
if vids.size < 2:
return None, None
coords = np.asarray(plane.rz[vids], dtype=float)
r = coords[:, 0]
z = coords[:, 1] - float(z_mid)
nseg = vids.size if closed else vids.size - 1
crossings = []
for i in range(nseg):
j = (i + 1) % vids.size
z0 = z[i]
z1 = z[j]
r0 = r[i]
r1 = r[j]
if not np.all(np.isfinite([z0, z1, r0, r1])):
continue
if abs(z0) <= 1.0e-12:
crossings.append(float(r0))
if z0*z1 < 0.0:
frac = -z0/(z1 - z0)
crossings.append(float(r0 + frac*(r1 - r0)))
elif abs(z1) <= 1.0e-12:
crossings.append(float(r1))
if not crossings:
return None, None
axis_r = float(plane.axis_r)
crossings = np.asarray(crossings, dtype=float)
outer = crossings[crossings >= axis_r]
inner = crossings[crossings <= axis_r]
r_outer = float(np.nanmax(outer)) if outer.size else None
r_inner = float(np.nanmin(inner)) if inner.size else None
return r_outer, r_inner
[docs]
def midplane_radius_maps(plane) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Build ``psi_N -> R_midplane`` maps for outer and inner midplanes.
For each ``Plane`` flux surface, this linearly interpolates exact
intersections with ``Z=0``. Surfaces that do not cross the outboard or
inboard midplane keep ``NaN`` for that branch. The result is sorted by
normalized flux and later used to map target-wall ``psi_N`` to
upstream/midplane radial distance from the separatrix.
Parameters
----------
plane : Plane
Poloidal mesh with flux surfaces, ``(R, Z)`` coordinates, magnetic
axis radius, and X-point flux.
Returns
-------
tuple[ndarray, ndarray, ndarray]
Sorted normalized poloidal flux, outer midplane radius, and inner
midplane radius. Radii are in metres; missing crossings are NaN.
"""
psi_vals = []
r_outer = []
r_inner = []
z_mid = 0.0
for isurf in range(int(plane.nsurf)):
outer, inner = _surface_midplane_crossings(plane, isurf, z_mid)
if outer is None and inner is None:
continue
psi_vals.append(float(plane.psi_surf[isurf]) / float(plane.x_psi))
r_outer.append(np.nan if outer is None else float(outer))
r_inner.append(np.nan if inner is None else float(inner))
psi = np.asarray(psi_vals, dtype=float)
order = np.argsort(psi)
psi = psi[order]
return psi, np.asarray(r_outer, dtype=float)[order], np.asarray(r_inner, dtype=float)[order]
[docs]
def finite_pair_count(x: np.ndarray, y: np.ndarray) -> int:
"""Count positions with finite values in both ``x`` and ``y`` arrays.
The arrays must broadcast to the same comparison shape; return that
finite-pair count as an integer.
"""
return int(np.count_nonzero(np.isfinite(x) & np.isfinite(y)))
def _unique_average(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Return sorted unique ``x`` values and averaged ``y`` values."""
order = np.argsort(x)
xs = np.asarray(x[order], dtype=float)
ys = np.asarray(y[order], dtype=float)
uniq, inv = np.unique(xs, return_inverse=True)
y_acc = np.zeros(uniq.size, dtype=float)
n_acc = np.zeros(uniq.size, dtype=float)
for i, j in enumerate(inv):
y_acc[j] += ys[i]
n_acc[j] += 1.0
return uniq, y_acc/np.maximum(n_acc, 1.0)
[docs]
def interp_radius_map(psi_mid: np.ndarray, r_mid: np.ndarray, psi_query: np.ndarray, branch: str) -> np.ndarray:
"""Interpolate midplane radii at normalized-flux query points.
``psi_mid`` and ``r_mid`` are one branch from :func:`midplane_radius_maps`;
``branch`` labels error messages. Duplicate flux points are averaged.
Query points outside the finite map return NaN. The output is in metres.
A map with fewer than two finite pairs raises ``ValueError``.
"""
finite = np.isfinite(psi_mid) & np.isfinite(r_mid)
if np.count_nonzero(finite) < 2:
raise ValueError(f"Not enough finite {branch}-midplane crossings for interpolation.")
psi, radius = _unique_average(psi_mid[finite], r_mid[finite])
return np.interp(psi_query, psi, radius, left=np.nan, right=np.nan)
[docs]
def interp_radius_at_psi(psi_mid: np.ndarray, r_mid: np.ndarray, psi_value: float, branch: str) -> float:
"""Return midplane radius in metres at one normalized-flux value.
``psi_mid`` and ``r_mid`` are one branch from :func:`midplane_radius_maps`;
``branch`` labels errors. Duplicate flux points are averaged. A map with
fewer than two finite pairs, or a query outside its bounds, raises
``ValueError``.
"""
finite = np.isfinite(psi_mid) & np.isfinite(r_mid)
psi, radius = _unique_average(psi_mid[finite], r_mid[finite])
if psi.size < 2 or psi_value < psi[0] or psi_value > psi[-1]:
raise ValueError(f"The {branch}-midplane map does not bracket psi_N={psi_value}.")
return float(np.interp(float(psi_value), psi, radius))