"""Wall-polygon arclength utilities."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
[docs]
@dataclass(frozen=True)
class ZeroCrossing:
"""A crossing of the line Z=0 by the wall curve."""
s: float
r: float
segment_index: int
[docs]
class WallCurve:
"""
Closed wall polygon with arclength coordinate.
Parameters
----------
points_rz : np.ndarray
Wall vertices as shape ``(n, 2)`` in clockwise order. The first vertex
should not be repeated at the end.
verify_clockwise : bool
Validate clockwise orientation and raise ``ValueError`` if violated.
auto_reverse : bool
If ``True`` and orientation is counter-clockwise, reverse order instead
of raising.
"""
def __init__(
self,
points_rz: np.ndarray,
*,
verify_clockwise: bool = True,
auto_reverse: bool = False,
):
pts = np.asarray(points_rz, dtype=float)
if pts.ndim != 2 or pts.shape[1] != 2:
raise ValueError("points_rz must be an array of shape (n, 2).")
if pts.shape[0] < 3:
raise ValueError("A wall polygon needs at least 3 vertices.")
if np.allclose(pts[0], pts[-1]):
pts = pts[:-1]
area_signed = self._signed_area(pts)
is_clockwise = area_signed < 0.0
if verify_clockwise and not is_clockwise:
if auto_reverse:
pts = pts[::-1].copy()
else:
raise ValueError(
"Wall polygon is not clockwise. Pass auto_reverse=True to fix automatically."
)
self._points = pts
self._n = pts.shape[0]
self._seg_len = self._segment_lengths(pts)
self._s_vertex = np.zeros(self._n, dtype=float)
self._s_vertex[1:] = np.cumsum(self._seg_len[:-1])
self._L = float(np.sum(self._seg_len))
self._origin_s = 0.0
@staticmethod
def _signed_area(points: np.ndarray) -> float:
x = points[:, 0]
y = points[:, 1]
return 0.5 * float(np.sum(x * np.roll(y, -1) - np.roll(x, -1) * y))
@staticmethod
def _segment_lengths(points: np.ndarray) -> np.ndarray:
p_next = np.roll(points, -1, axis=0)
return np.sqrt(np.sum((p_next - points) ** 2, axis=1))
@property
def points(self) -> np.ndarray:
"""Wall vertices ``(R,Z)`` in the stored order."""
return self._points
@property
def n_vertices(self) -> int:
"""Number of wall vertices."""
return self._n
@property
def total_length(self) -> float:
"""Total closed-curve length."""
return self._L
@property
def s_vertex(self) -> np.ndarray:
"""Arclength at each vertex (origin-shifted, modulo total length)."""
return np.mod(self._s_vertex - self._origin_s, self._L)
[docs]
def set_origin_at_inboard_midplane(self, r_axis: float, *, z_tol: float = 1.0e-10) -> float:
"""
Set s=0 at the intersection with Z=0 and R<R_axis.
Returns
-------
float
The selected crossing arclength in the original coordinate.
"""
crossings = self.find_z0_crossings(z_tol=z_tol)
inboard = [c for c in crossings if c.r < r_axis]
if not inboard:
raise ValueError("No wall crossing with Z=0 and R<R_axis found.")
chosen = max(inboard, key=lambda c: c.r)
self._origin_s = chosen.s
return chosen.s
[docs]
def find_z0_crossings(self, *, z_tol: float = 1.0e-10) -> list[ZeroCrossing]:
"""Find intersections of the closed curve with Z=0."""
out: list[ZeroCrossing] = []
p = self._points
p_next = np.roll(p, -1, axis=0)
for i in range(self._n):
z1 = p[i, 1]
z2 = p_next[i, 1]
if abs(z1) <= z_tol and abs(z2) <= z_tol:
continue
if abs(z1) <= z_tol:
t = 0.0
elif abs(z2) <= z_tol:
t = 1.0
elif z1 * z2 < 0.0:
t = -z1 / (z2 - z1)
else:
continue
t = float(np.clip(t, 0.0, 1.0))
r = float(p[i, 0] + t * (p_next[i, 0] - p[i, 0]))
s = float(self._s_vertex[i] + t * self._seg_len[i])
out.append(ZeroCrossing(s=s, r=r, segment_index=i))
return out
[docs]
def circular_distance(self, s1: float, s2: float) -> float:
"""Shortest arclength distance between two arclength values."""
d = abs((s2 - s1) % self._L)
return float(min(d, self._L - d))
[docs]
def shortest_interval(self, s1: float, s2: float) -> tuple[float, float, bool]:
"""
Return shortest directed interval from s1 to s2.
Returns
-------
(s_start, s_end, wraps)
Interval on [0,L) where ``wraps=True`` means it crosses L->0.
"""
a = float(s1 % self._L)
b = float(s2 % self._L)
d_fwd = (b - a) % self._L
d_bwd = (a - b) % self._L
if d_fwd <= d_bwd:
return a, b, bool(b < a)
return b, a, bool(a < b)