import numpy as np
from matplotlib.tri import Triangulation
__all__ = ["compute_fieldline_mapping"]
# ------------------------------------------------------------------
# Vector‑ised helper for RK‑4 integration (all vertices at once)
# ------------------------------------------------------------------
def _rk4_step_batch(rhs, R, Z, h):
"""Advance *R, Z* by one classical RK‑4 step of size *h*. Arrays are
updated in‑place. *rhs(phi, R, Z)* returns (dR/dφ, dZ/dφ)."""
k1_R, k1_Z = rhs(R, Z)
k2_R, k2_Z = rhs(R + 0.5 * h * k1_R, Z + 0.5 * h * k1_Z)
k3_R, k3_Z = rhs(R + 0.5 * h * k2_R, Z + 0.5 * h * k2_Z)
k4_R, k4_Z = rhs(R + h * k3_R, Z + h * k3_Z)
R += h * (k1_R + 2 * k2_R + 2 * k3_R + k4_R) / 6.0
Z += h * (k1_Z + 2 * k2_Z + 2 * k3_Z + k4_Z) / 6.0
# ------------------------------------------------------------------
# Barycentric weights – explicit, vectorised
# ------------------------------------------------------------------
def _bary_weights_batch(RZ_tri, P):
"""Return barycentric coordinates for *m* points at once.
Parameters
----------
RZ_tri : (m, 3, 2) float64 – triangle vertices
P : (m, 2) float64 – query points
"""
R1, Z1 = RZ_tri[:, 0, 0], RZ_tri[:, 0, 1]
R2, Z2 = RZ_tri[:, 1, 0], RZ_tri[:, 1, 1]
R3, Z3 = RZ_tri[:, 2, 0], RZ_tri[:, 2, 1]
R, Z = P[:, 0], P[:, 1]
denom = (Z2 - Z3) * (R1 - R3) + (R3 - R2) * (Z1 - Z3)
w0 = ((Z2 - Z3) * (R - R3) + (R3 - R2) * (Z - Z3)) / denom
w1 = ((Z3 - Z1) * (R - R3) + (R1 - R3) * (Z - Z3)) / denom
w2 = 1.0 - w0 - w1
return np.column_stack((w0, w1, w2))
# ------------------------------------------------------------------
# Robust helper to obtain a *hash*‑grid TriFinder regardless of mpl version
# ------------------------------------------------------------------
def _get_hash_trifinder(tri_obj):
try:
return tri_obj.get_trifinder(kind="hash") # mpl ≥ 3.7
except TypeError:
try:
return tri_obj.get_trifinder("hash") # mpl ≤ 3.6
except TypeError:
return tri_obj.get_trifinder() # fallback trapezoid
# ------------------------------------------------------------------
# Main routine
# ------------------------------------------------------------------
[docs]
def compute_fieldline_mapping(
mesh,
magnetic_field,
*,
tor_turns: float = 1.0,
rk4_substeps: int = 2,
n_int: int = 250,
direction: str = "forward",
):
"""Trace mesh vertices along **B** and return interpolation meta‑data.
Parameters
----------
tor_turns : float, optional
How many *toroidal turns* (2π each) to follow. Fractional values
are allowed; the actual number of steps is rounded to the nearest
multiple of stored planes so the trace always lands on a plane.
direction : {'forward', 'backward'}
Sign of dφ. *forward* ⇒ increasing φ, *backward* ⇒ decreasing.
"""
if direction not in ("forward", "backward"):
raise ValueError("direction must be 'forward' or 'backward'")
planes = mesh.planes
nphi = mesh.nphi
delta_phi = mesh.delta_phi
wedge_angle = mesh.wedge_angle
wedge_n = int(round(2.0 * np.pi / wedge_angle))
steps_per_turn = wedge_n * nphi
n_steps_total = max(1, int(round(tor_turns * steps_per_turn)))
sign = 1.0 if direction == "forward" else -1.0
h = abs(delta_phi) / rk4_substeps
h_signed = sign * h
# ------------------------------------------------------------------
# Build TriFinders and cached arrays
# ------------------------------------------------------------------
trifinders, connects, coords = [], [], []
for p in planes:
tri_obj = getattr(p, "triObj", p.get_triangulation_data(n_int=n_int).triObj)
trifinders.append(_get_hash_trifinder(tri_obj))
connects.append(p.nd_connect_list)
coords.append(p.rz)
ref_plane = planes[0]
n_vert = ref_plane.n_n
# global vertex positions (updated in‑place)
R_all = ref_plane.rz[:, 0].copy()
Z_all = ref_plane.rz[:, 1].copy()
tri_index = np.full((n_vert, n_steps_total), -1, dtype=np.int32)
bary_weights = np.full((n_vert, n_steps_total, 3), np.nan)
plane_index = np.empty(n_steps_total, dtype=np.int32)
# RHS function for RK‑4 (vectorised)
def rhs(R, Z):
BR, BZ, Bphi = magnetic_field.compute_background_field(R, Z)
Bphi = np.where(np.abs(Bphi) < 1e-12, 1e-12 * np.sign(Bphi), Bphi)
return R * BR / Bphi, R * BZ / Bphi
for s in range(n_steps_total):
# integrate all vertices to the next stored plane
for _ in range(rk4_substeps):
_rk4_step_batch(rhs, R_all, Z_all, h_signed)
tgt_plane_idx = 0 if mesh.is_axisymmetric else (s + 1) % nphi
plane_index[s] = tgt_plane_idx
tri_ids = trifinders[tgt_plane_idx](R_all, Z_all)
valid = tri_ids >= 0
tri_index[valid, s] = tri_ids[valid]
if np.any(valid):
verts = connects[tgt_plane_idx][tri_ids[valid]]
if verts.max() >= coords[tgt_plane_idx].shape[0]:
verts = verts - 1
RZ_tri = coords[tgt_plane_idx][verts]
bary_weights[valid, s] = _bary_weights_batch(
RZ_tri, np.column_stack((R_all[valid], Z_all[valid]))
)
return dict(
triangle_index=tri_index,
bary_weights=bary_weights,
plane_index=plane_index,
delta_phi=delta_phi,
direction=direction,
)