"""Map SOL flux-surface pairs to wall arclength intervals."""
from __future__ import annotations
from dataclasses import dataclass
import warnings
import numpy as np
from .wall_curve import WallCurve
_PSI_TOL = 1.0e-5
[docs]
@dataclass
class SOLVolumeWallBounds:
"""Wall-boundary information for one adjacent SOL surface pair."""
lower_surface: int
upper_surface: int
psi_lower: float
psi_upper: float
lower_wall_nodes: tuple[int, int]
upper_wall_nodes: tuple[int, int]
intervals: tuple[tuple[float, float, bool], tuple[float, float, bool]]
[docs]
class SOLWallVolumeMap:
"""Identify wall ranges that bound volumes between adjacent SOL surfaces."""
def __init__(self, plane, wall_curve: WallCurve):
self.plane = plane
self.wall_curve = wall_curve
self.bounds: list[SOLVolumeWallBounds] = []
@staticmethod
def _surface_wall_nodes(plane, surf_idx: int) -> np.ndarray:
verts = np.asarray(plane.get_surface_vertex_indices(int(surf_idx)), dtype=int)
wall_nodes = np.asarray(plane.wall_nodes, dtype=int)
wn = verts[np.isin(verts, wall_nodes)]
# Fallback: identify wall vertices on this surface by region marker.
if wn.shape[0] < 2 and hasattr(plane, "region"):
region = np.asarray(plane.region, dtype=int)
on_wall = verts[region[verts] == 100]
if on_wall.shape[0] >= 2:
wn = on_wall
return np.unique(wn)
@staticmethod
def _build_wall_s_map(plane, wall_curve: WallCurve) -> dict[int, float]:
wall_nodes = np.asarray(plane.wall_nodes, dtype=int)
if wall_nodes.shape[0] != wall_curve.n_vertices:
raise ValueError(
"Wall node count mismatch between plane.wall_nodes and WallCurve vertices."
)
return {int(node): float(s) for node, s in zip(wall_nodes, wall_curve.s_vertex)}
@staticmethod
def _pair_two_by_circular_distance(
s_a: np.ndarray,
s_b: np.ndarray,
wall_curve: WallCurve,
) -> tuple[tuple[int, int], tuple[int, int]]:
if s_a.shape[0] != 2 or s_b.shape[0] != 2:
raise ValueError("Current implementation expects exactly 2 strike points per surface.")
d00 = wall_curve.circular_distance(s_a[0], s_b[0])
d01 = wall_curve.circular_distance(s_a[0], s_b[1])
d10 = wall_curve.circular_distance(s_a[1], s_b[0])
d11 = wall_curve.circular_distance(s_a[1], s_b[1])
# two possible matchings
sum_01 = d00 + d11
sum_10 = d01 + d10
if sum_01 <= sum_10:
return (0, 0), (1, 1)
return (0, 1), (1, 0)
[docs]
def build_from_surf_map(self, *, psi_norm_min: float = 1.0, psi_norm_max: float | None = None) -> list[SOLVolumeWallBounds]:
"""
Build adjacent-surface wall bounds from ``plane.surf_map``.
Parameters
----------
psi_norm_min : float
Minimum normalized psi for surfaces to include (default: SOL only, >=1).
psi_norm_max : float | None
Optional maximum normalized psi cutoff (e.g., diff_bd_out).
"""
surf_map = np.asarray(self.plane.surf_map, dtype=int)
psi_surf = np.asarray(self.plane.psi_surf, dtype=float)
x_psi = float(self.plane.x_psi)
psi_n = psi_surf[surf_map] / x_psi
# Include separatrix in surface-pair construction for SOL volume binning.
keep = psi_n >= (psi_norm_min - _PSI_TOL)
if psi_norm_max is not None:
keep &= psi_n <= (float(psi_norm_max) + _PSI_TOL)
surf_sol = surf_map[keep]
if surf_sol.shape[0] < 2:
self.bounds = []
return self.bounds
wall_s = self._build_wall_s_map(self.plane, self.wall_curve)
out: list[SOLVolumeWallBounds] = []
for i in range(surf_sol.shape[0] - 1):
s_lo = int(surf_sol[i])
s_hi = int(surf_sol[i + 1])
wn_lo = np.unique(self._surface_wall_nodes(self.plane, s_lo))
wn_hi = np.unique(self._surface_wall_nodes(self.plane, s_hi))
if wn_lo.shape[0] != 2 or wn_hi.shape[0] != 2:
warnings.warn(
f"Skipping pair ({s_lo},{s_hi}): expected 2 wall nodes each, got "
f"{wn_lo.shape[0]} and {wn_hi.shape[0]}.",
RuntimeWarning,
)
continue
s_lo_vals = np.array([wall_s[int(n)] for n in wn_lo], dtype=float)
s_hi_vals = np.array([wall_s[int(n)] for n in wn_hi], dtype=float)
pair1, pair2 = self._pair_two_by_circular_distance(s_lo_vals, s_hi_vals, self.wall_curve)
i0, j0 = pair1
i1, j1 = pair2
int0 = self.wall_curve.shortest_interval(s_lo_vals[i0], s_hi_vals[j0])
int1 = self.wall_curve.shortest_interval(s_lo_vals[i1], s_hi_vals[j1])
out.append(
SOLVolumeWallBounds(
lower_surface=s_lo,
upper_surface=s_hi,
psi_lower=float(psi_surf[s_lo]),
psi_upper=float(psi_surf[s_hi]),
lower_wall_nodes=(int(wn_lo[i0]), int(wn_lo[i1])),
upper_wall_nodes=(int(wn_hi[j0]), int(wn_hi[j1])),
intervals=(int0, int1),
)
)
self.bounds = out
return out