import numpy as np
from xgc_analysis.mesh import Mesh
from xgc_analysis.magnetic_field import MagneticField
from matplotlib.tri import Triangulation
from matplotlib.tri import TriFinder
__all__ = ["compute_fieldline_mapping"]
# --------------------------------------------------------------
# build a hash-grid finder that works on every mpl version
# --------------------------------------------------------------
def _hash_finder(triang):
try: # mpl ≥ 3.7 (keyword allowed)
return triang.get_trifinder(kind='hash')
except TypeError:
try: # mpl ≤ 3.6 (positional only)
return triang.get_trifinder('hash')
except TypeError: # ancient fallback
return triang.get_trifinder() # default trapezoid
def _rk4_step_batch(R, Z, h_signed, magnetic_field):
"""
One RK-4 step for *all* vertices at once.
Parameters
----------
R, Z : ndarray (n_vert,)
Current positions (modified in-place on return).
h_signed : float
Signed step Δφ (positive → forward, negative → backward).
"""
# local RHS = d(R,Z)/dφ
def rhs(Ra, Za):
BR, BZ, Bphi = magnetic_field.compute_background_field(Ra, Za)
Bphi = np.where(np.abs(Bphi) < 1e-12, 1e-12*np.sign(Bphi)+1e-12, Bphi) # avoid /0
return (Ra * BR / Bphi, Ra * BZ / Bphi)
k1_R, k1_Z = rhs(R, Z)
k2_R, k2_Z = rhs(R + 0.5*h_signed*k1_R, Z + 0.5*h_signed*k1_Z)
k3_R, k3_Z = rhs(R + 0.5*h_signed*k2_R, Z + 0.5*h_signed*k2_Z)
k4_R, k4_Z = rhs(R + h_signed*k3_R, Z + h_signed*k3_Z)
R += h_signed*(k1_R + 2*k2_R + 2*k3_R + k4_R) / 6.0
Z += h_signed*(k1_Z + 2*k2_Z + 2*k3_Z + k4_Z) / 6.0
# ------------------------------------------------------------------
# 1. helper: explicit barycentric weights, vectorised
# ------------------------------------------------------------------
def _bary_weights_batch(RZ_tri, P):
"""
Vectorised barycentric coordinates for many points at once.
Parameters
----------
RZ_tri : (m, 3, 2) float64
Triangle vertices for *m* points.
P : (m, 2) float64
Query points (R, Z).
Returns
-------
w : (m, 3) float64
Barycentric weights.
"""
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))
[docs]
def compute_fieldline_mapping(mesh, magnetic_field,
*, rk4_substeps=2, n_int=250,
direction="forward"):
"""Build a magnetic‑field‑line mapping between adjacent planes.
Parameters
----------
mesh : Mesh
Instance of the :class:`Mesh` defined in *mesh.py*.
magnetic_field : MagneticField
Instance of :class:`MagneticField` (see *magnetic_field.py*).
rk4_substeps : int, optional
Number of RK4 sub‑steps per plane‑to‑plane advance (default 2).
n_int : int, optional
Interpolation‑grid resolution used when creating a
:class:`matplotlib.tri.TriFinder` (default 250).
direction : {"forward", "backward"}, optional
Direction of tracing; *forward* increases φ, *backward* decreases φ.
Returns
-------
dict
``triangle_index`` : ``(n_vert, n_steps)`` ``int32``
Triangle hosting each intersection.
``bary_weights`` : ``(n_vert, n_steps, 3)`` ``float64``
Barycentric weights inside that triangle.
``plane_index`` : ``(n_steps,)`` ``int32``
Index of the target plane within *mesh.planes*.
``delta_phi`` : ``float``
Toroidal separation between neighbouring planes (positive).
``direction`` : ``str``
Copy of the *direction* argument.
"""
planes = mesh.planes
nphi = mesh.nphi
delta_phi = mesh.delta_phi # separation of stored planes inside one wedge
wedge_angle = mesh.wedge_angle
wedge_n = int(round(2.0 * np.pi / wedge_angle))
n_steps_total = wedge_n * nphi # full 2π circuit
sign = 1.0 if direction == "forward" else -1.0
h = abs(delta_phi) / rk4_substeps # positive magnitude of sub‑step
ref_plane = planes[0] # φ₀ = 0 reference
n_vert = ref_plane.n_n
# --------------------------------------------------------------
# TriFinders (created once per plane)
# --------------------------------------------------------------
trifinders, connects, coords = [], [], []
for p in planes:
try:
tri_obj = p.triObj
except AttributeError:
tri_obj = p.get_triangulation_data(n_int=n_int).triObj
# use the faster hash-grid finder
trifinder = _hash_finder(tri_obj)
trifinders.append(trifinder)
connects.append(p.nd_connect_list)
coords.append(p.rz)
# ------------------------------------------------------------------
# initialise arrays of *all* vertex positions
# ------------------------------------------------------------------
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_w = np.full((n_vert, n_steps_total, 3), np.nan, dtype=np.float64)
plane_idx_arr = np.empty(n_steps_total, dtype=np.int32)
h = abs(delta_phi) / rk4_substeps
sign = 1.0 if direction == "forward" else -1.0
h_signed = sign*h
# ------------------------------------------------------------------
# loop over toroidal steps *first*
# ------------------------------------------------------------------
phi_curr = 0.0
for s in range(n_steps_total):
# --- integrate all vertices to plane s+1 --------------------
for _ in range(rk4_substeps):
_rk4_step_batch(R_all, Z_all, h_signed, magnetic_field)
phi_curr += h_signed
# --- locate all intersections on the target plane ----------
#tgt_plane_idx = 0 if mesh.is_axisymmetric else (s + 1) % nphi
tgt_plane_idx = (s + 1) % nphi
plane_idx_arr[s] = tgt_plane_idx
plane_idx = 0 if mesh.is_axisymmetric else tgt_plane_idx
tri_ids = trifinders[plane_idx](R_all, Z_all) # vectorised
valid = tri_ids >= 0
tri_index[valid, s] = tri_ids[valid]
# vertices of hosting triangles (m,3,2)
verts = connects[plane_idx][tri_ids[valid]]
if verts.max() >= coords[plane_idx].shape[0]:
verts -= 1
RZ_tri = coords[plane_idx][verts]
# bary weights, vectorised
bary_w[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_w,
plane_index=plane_idx_arr,
delta_phi=mesh.delta_phi,
direction=direction)