import numpy as np
from .plane import Plane
from xgc_analysis.data_access.catalog import has_static_product, read_static_variables
from scipy.sparse import csr_matrix as _CSR
from typing import List
def _mapping_to_csr_list(mapping: dict,
nd_connect: np.ndarray,
vertex_count: int) -> List[_CSR]:
"""
Convert ``compute_fieldline_mapping`` output into a list of CSR matrices,
one per toroidal step. For an axisymmetric mesh every matrix has shape
``(n_vert, n_vert)`` and projects data forward by Δφ.
Each matrix Pₛ fulfils
f(φ₀ + (s+1)Δφ) = Pₛ @ f(φ₀)
"""
tri = mapping["triangle_index"] # (n_vert, n_steps)
w = mapping["bary_weights"] # (n_vert, n_steps, 3)
n_v, n_steps = tri.shape
# auto-detect 1-based connectivity
conn = nd_connect.copy()
if conn.max() >= vertex_count:
conn -= 1
mats = []
for s in range(n_steps):
tri_s = tri[:, s]
valid = tri_s >= 0
if not np.any(valid):
mats.append(_CSR((vertex_count, vertex_count)))
continue
rows = np.repeat(np.nonzero(valid)[0], 3)
cols = conn[tri_s[valid]].reshape(-1)
data = w[valid, s, :].reshape(-1)
mats.append(_CSR((data, (rows, cols)),
shape=(vertex_count, vertex_count)))
return mats
###############################################################################
# The "mesh" class: an overarching container for one or more planes.
###############################################################################
[docs]
class Mesh:
"""
A general mesh class that consists of an array of "plane" objects.
The mesh data for all planes comes from the same BP file ("xgc.mesh.bp"). The file also
contains the variables "delta_phi" and "wedge_angle". The number of wedges is computed as:
wedge_n = 2*pi / wedge_angle,
and the number of planes is computed as:
nphi = wedge_angle / delta_phi.
If the input flag is_axisymmetric is True, then the mesh is axisymmetric and only one plane is set up.
Otherwise, nphi planes are created (all identical in this version).
"""
def __init__(
self,
is_axisymmetric,
data_dir=".",
catalog=None,
source_reader=None,
static_buffer=None,
):
"""
Initialize mesh planes and toroidal wedge metadata.
Parameters
----------
is_axisymmetric : bool
If True, construct one reusable axisymmetric plane. Otherwise
construct one plane per toroidal mesh plane.
data_dir : str, optional
Directory containing static mesh BP products for legacy local reads.
catalog : xgc_analysis.data_access.catalog.SimulationCatalog or None, optional
Catalog used to read static mesh metadata. Direct local path
fallback reads are disabled.
source_reader : callable or None, optional
Optional static-product read backend passed through to catalog
reads.
static_buffer : mapping[str, mapping[str, object]] or None, optional
Optional reusable static-product buffer shared with downstream
``Plane`` instances.
"""
# Read the wedge information from the BP file.
if has_static_product(catalog, "xgc.mesh.bp", ("delta_phi", "wedge_angle")):
mesh_values = read_static_variables(
catalog,
"xgc.mesh.bp",
("delta_phi", "wedge_angle"),
source_reader=source_reader,
static_buffer=static_buffer,
)
self.delta_phi = float(np.asarray(mesh_values["delta_phi"]).item())
self.wedge_angle = float(np.asarray(mesh_values["wedge_angle"]).item())
else:
raise RuntimeError(
"Mesh requires catalog product 'xgc.mesh.bp' with "
"'delta_phi' and 'wedge_angle'; direct local BP fallback is disabled."
)
# Compute the number of wedges and number of planes.
# Here, we round to the nearest integer.
self.wedge_n = int(round(2 * np.pi / self.wedge_angle))
self.nphi = int(round(self.wedge_angle / self.delta_phi))
self.is_axisymmetric = is_axisymmetric
if is_axisymmetric:
# For axisymmetric meshes, all planes are identical, so only one plane is created.
self.planes = [
Plane(
data_dir=data_dir,
catalog=catalog,
source_reader=source_reader,
static_buffer=static_buffer,
)
]
self.n_planes = 1
else:
# For non-axisymmetric meshes, set up an array of nphi planes.
self.planes = [
Plane(
data_dir=data_dir,
catalog=catalog,
source_reader=source_reader,
static_buffer=static_buffer,
)
for _ in range(self.nphi)
]
self.n_planes = self.nphi
[docs]
def setup_fieldline_mapping(self, magnetic_field,
*, direction="forward",
rk4_substeps: int = 2, n_int: int = 250):
"""
Trace all mesh vertices along B and build one CSR propagator matrix
for every stored toroidal step. The objects are cached as
self.fl_mapping – raw dict from compute_fieldline_mapping
self.fl_csr_matrices – list[CSR] (length = n_steps)
Call this after construction if the mapping was not required earlier.
``magnetic_field`` supplies the background field. ``direction`` sets
tracing orientation, ``rk4_substeps`` sets integration subdivisions,
and ``n_int`` sets toroidal steps. Updates ``fl_mapping`` and
``fl_csr_matrices`` in place. Only axisymmetric meshes are supported.
"""
from .fieldline_mapping import compute_fieldline_mapping
if (not self.is_axisymmetric):
raise ValueError("Field-following mapping is currently only available for axisymmetric meshes.")
# --- run tracer ----------------------------------------------------
self.fl_mapping = compute_fieldline_mapping(
self,
magnetic_field,
rk4_substeps=rk4_substeps,
n_int=n_int,
direction=direction,
)
# --- assemble sparse matrices -------------------------------------
ref_plane = self.planes[0] # axisymmetric assumption for now
self.fl_csr_matrices = _mapping_to_csr_list(
self.fl_mapping,
ref_plane.nd_connect_list,
ref_plane.n_n,
)
@property
def wall_rz(self):
"""Return wall node coordinates per plane as a list of arrays."""
return [plane.wall_rz for plane in self.planes]
[docs]
def get_wall_curve(
self,
*,
plane_index: int = 0,
verify_clockwise: bool = True,
auto_reverse: bool = False,
set_inboard_origin: bool = True,
):
"""
Build a :class:`WallCurve` from a selected plane's wall polygon.
``plane_index`` selects the plane. The three boolean options control
orientation checks, reversal, and arclength origin as documented by
:meth:`Plane.get_wall_curve`. Return the new wall curve.
"""
plane = self.get_plane(plane_index)
return plane.get_wall_curve(
verify_clockwise=verify_clockwise,
auto_reverse=auto_reverse,
set_inboard_origin=set_inboard_origin,
)
@property
def wall_phi(self):
"""Return toroidal angle for each plane in radians."""
if self.nphi <= 1:
return np.array([0.0])
return np.arange(self.nphi, dtype=float) * self.delta_phi
[docs]
def get_plane(self, index=0):
"""
Return the plane at ``index``; ignore the index for axisymmetric meshes.
Raise ``IndexError`` when a non-axisymmetric index is out of range.
"""
if self.is_axisymmetric:
return self.planes[0]
else:
if index < 0 or index >= len(self.planes):
raise IndexError("Plane index out of range.")
return self.planes[index]
[docs]
def get_vertex_counts(self):
"""
Returns a list of the number of vertices on each of the mesh's planes.
For an axisymmetric mesh (is_axisymmetric=True), the same number of vertices (from the
single plane stored) is repeated for each plane (nphi times).
For a non-axisymmetric mesh, the function returns the number of vertices for each plane.
"""
if self.is_axisymmetric:
return [self.planes[0].n_n] * self.nphi
else:
return [p.n_n for p in self.planes]
# ---------------- global (3‑D) diagnostics -----------------
[docs]
def flux_avg_to_surf(self, meshdata: np.ndarray, *, fsa: str = "volume") -> np.ndarray:
"""Flux‑surface average (⟨f⟩_ψ) of a 3‑D field, averaged over φ.
``meshdata`` has toroidal plane as axis zero; ``fsa`` selects the
plane-level weighting scheme. Axisymmetric data may omit that axis.
Returns
-------
ndarray
Shape (nsurf,) for scalar, or (nsurf, n_components) for vector data.
"""
meshdata = np.asarray(meshdata)
if len(self.planes) == 1 and meshdata.shape[0] != 1:
meshdata = meshdata[None, ...]
per_plane = [pl.flux_avg_to_surf(meshdata[i], fsa=fsa)
for i, pl in enumerate(self.planes)]
return np.mean(per_plane, axis=0)
[docs]
def flux_avg_from_surf(self, flux_avg: np.ndarray, *, fsa: str = "volume") -> np.ndarray:
"""Project surface values onto vertices and average over stored planes.
``flux_avg`` contains surface samples; ``fsa`` selects the weighting
scheme. Return a representative vertex array, not a ``PlaneData``.
"""
per_plane = [pl.flux_avg_from_surf(flux_avg, fsa=fsa) for pl in self.planes]
return np.mean(per_plane, axis=0)
[docs]
def flux_avg_to_from_surf(self, meshdata: np.ndarray, *, fsa: str = "volume") -> np.ndarray:
"""Average ``meshdata`` over flux surfaces and project to vertices.
``fsa`` selects the plane-level weighting scheme. Return the
toroidally averaged projection on reference plane zero.
"""
meshdata = np.asarray(meshdata)
if len(self.planes) == 1 and meshdata.shape[0] != 1:
meshdata = meshdata[None, ...]
per_plane = [pl.flux_avg_to_surf(meshdata[i], fsa=fsa)
for i, pl in enumerate(self.planes)]
surf_mean = np.mean(per_plane, axis=0)
return self.planes[0].flux_avg_from_surf(surf_mean, fsa=fsa)
[docs]
def flux_avg_to_surf_norm(self, meshdata: np.ndarray, *, fsa: str = "volume") -> np.ndarray:
"""Average the vector norm of ``meshdata`` over surfaces and planes.
``fsa`` selects the plane-level weighting scheme. Return one value
per flux surface.
"""
meshdata = np.asarray(meshdata)
if len(self.planes) == 1 and meshdata.shape[0] != 1:
meshdata = meshdata[None, ...]
per_plane = [pl.flux_avg_to_surf_norm(meshdata[i], fsa=fsa)
for i, pl in enumerate(self.planes)]
return np.mean(per_plane, axis=0)
# ------------ toroidal / field‑line helpers --------------
[docs]
def toroidal_average(self, meshdata: np.ndarray) -> np.ndarray:
"""Return the arithmetic mean of ``meshdata`` along toroidal axis zero.
The output retains its vertex and optional component axes.
"""
return np.mean(meshdata, axis=0)
[docs]
def ff_map_real2ff(self, meshdata: np.ndarray) -> np.ndarray:
"""
Maps mesh data from the regular representation of data on
planes (points with the same index on different planes are
connected toroidally) to field-following representation (vertices
with the same index are connected along the magnetic field).
``meshdata`` has shape ``(nphi, node[, component])``. Return a NumPy
array whose leading dimension selects left and right mapped data.
The mesh must have multiple toroidal planes and loaded interpolation
matrices; otherwise the function raises ``ValueError`` or
``RuntimeError``.
"""
if (self.nphi == 1 or meshdata.ndim == 1):
raise ValueError("Data must not be axisymmetric for field-aligned mapping.")
if (meshdata.shape[0] != self.nphi):
raise ValueError("Input data has wrong shape.")
required = ("ff_hdp_fwd", "ff_hdp_rev")
missing = [a for a in required
if getattr(self.planes[0], a, None) is None]
if missing:
raise RuntimeError("Parallel-derivative data missing in plane 0: "
+ ", ".join(missing))
# Not yet compatible with Stellarator
if (meshdata.ndim == 3):
# Make the vector components the first dimension
n_components = meshdata.shape[2]
data = np.transpose(meshdata,(2,0,1))
else:
# Add artificial extra dimension
n_components = 1
data = meshdata[np.newaxis,:]
mapped_data = np.zeros((2,n_components,self.nphi,self.planes[0].n_n))
for i_comp in range(n_components):
for i in range(self.nphi):
if (i == 0):
im = self.nphi-1
else:
im = i-1
l_plane_data = data[i_comp,im,:]
r_plane_data = data[i_comp,i, :]
if (self.is_axisymmetric):
l_plane_obj = self.planes[0]
r_plane_obj = l_plane_obj
else:
l_plane_obj = self.planes[im]
r_plane_obj = self.planes[i ]
mapped_data[1,i_comp,i,:] = r_plane_obj.ff_hdp_fwd.multiply_left(r_plane_data)
mapped_data[0,i_comp,i,:] = l_plane_obj.ff_hdp_rev.multiply_left(l_plane_data)
if (n_components > 1):
# Transpose back to original ordering
mapped_data = np.transpose(mapped_data,(0,2,3,1))
else:
# Remove dummy dimension
mapped_data = np.squeeze(mapped_data)
return mapped_data
[docs]
def apply_fieldline_mapping(self, meshdata, *, ref_plane: int = 0):
"""
Pull the data of *every* plane onto the vertices of *ref_plane*
along magnetic-field lines.
``meshdata`` has toroidal plane as axis zero; ``ref_plane`` selects
the reference mesh plane. A field-line mapping must already exist.
Returns
-------
np.ndarray
If the stored field is scalar:
shape (n_steps, n_vert)
If the stored field is vector:
shape (n_steps, n_vert, n_components)
Axis-0 walks along the field line (φ = φ₀ + (s+1)Δφ);
Axis-1 enumerates the vertices of the reference plane.
"""
# ----------------------------------------------------------
# 0. sanity checks
# ----------------------------------------------------------
if not hasattr(self, "fl_csr_matrices") or not hasattr(self, "fl_mapping"):
raise RuntimeError("Call mesh.setup_fieldline_mapping(...) first.")
if (self.nphi == 1 or meshdata.ndim == 1):
raise ValueError("Data must not be axisymmetric.")
if (meshdata.shape[0] != self.nphi):
raise ValueError("Input data has wrong shape.")
if (meshdata.ndim == 3):
# Make the vector components the first dimension
n_comp = meshdata.shape[2]
else:
# Add artificial extra dimension
n_comp = 1
P_list = self.fl_csr_matrices # list[CSR]
plane_idx = self.fl_mapping["plane_index"] # (n_steps,)
n_steps = len(P_list)
n_vert = P_list[0].shape[0] # rows = reference vertices
# ----------------------------------------------------------
# 1. allocate output
# ----------------------------------------------------------
if n_comp == 1:
out = np.empty((n_steps, n_vert), dtype=meshdata.dtype)
else:
out = np.empty((n_steps, n_vert, n_comp), dtype=meshdata.dtype)
# ----------------------------------------------------------
# 2. loop over toroidal steps
# ----------------------------------------------------------
for s in range(n_steps):
data_src = meshdata[int(plane_idx[s])] # (n_src, [n_comp])
if n_comp == 1: # -------- scalar
out[s] = P_list[s] @ data_src
else: # -------- vector
for c in range(n_comp):
out[s, :, c] = P_list[s] @ data_src[:, c]
return out
# ---------------- Gradient -----------------
[docs]
def evaluate_scalar_gradients(self, field_2d: np.ndarray) -> np.ndarray:
"""
Apply per-plane gradient operators to a scalar field on the full 3-D mesh.
Parameters
----------
field_2d : np.ndarray, shape (nphi, n_n)
`field_2d[i_phi, i_node]` holds the scalar value at node *i_node*
on poloidal plane *i_phi*.
Returns
-------
grads : np.ndarray, shape (nphi, n_n, 2)
`grads[i_phi]` is the `(n_n, 2)` array returned by the
corresponding plane's `evaluate_scalar_gradients`.
"""
if (self.planes[0].gradient_r_psi is None or self.planes[0].gradient_z_theta is None):
raise RuntimeError("Gradient matrices are unavailable; cannot evaluate "
"in-plane gradients on this mesh.")
if field_2d.shape != (self.nphi, self.planes[0].n_n):
raise ValueError(
f"Expected shape ({self.nphi}, {self.planes[0].n_n}), got {field_2d.shape}"
)
out = np.empty((self.nphi, self.planes[0].n_n, 2), dtype=field_2d.dtype)
for i_phi, (plane, field_slice) in enumerate(zip(self.planes, field_2d)):
out[i_phi] = plane.evaluate_scalar_gradients(field_slice)
return out
# ------------------------------------------------------------------
# 2-nd–order finite-difference derivative parallel to B
# ------------------------------------------------------------------
[docs]
def evaluate_grad_par(self, field_2d: np.ndarray, bt_sign: int) -> np.ndarray:
"""
Compute b·∇(field) on an axisymmetric mesh.
Parameters
----------
field_2d : np.ndarray, shape (nphi, n_n)
Scalar field values on the full mesh
(first index = poloidal plane, second = node).
bt_sign : int, sign of the toroidal magnetic field ( +1 or -1 ).
Returns
-------
np.ndarray
Same shape as *field_2d* with the parallel derivative.
Notes
-----
Uses the centred, 2-nd-order accurate stencil
b·∇f ≈ sgn * [ -l_r/(l_l l_tot) f_{i-1}
+ (l_r-l_l)/(l_l l_r) f_i
+ l_l/(l_r l_tot) f_{i+1} ]
where the forward / reverse projection matrices *ff_1dp_fwd* and
*ff_1dp_rev* as well as the parallel arc-lengths *dl_par_1dp_fwd*
and *dl_par_1dp_rev* are taken from ``planes[0]`` (valid only when
``self.is_axisymmetric`` is *True*).
"""
# -------- sanity checks -------------------------------------------------
if not getattr(self, "is_axisymmetric", False):
raise RuntimeError("evaluate_grad_par currently supports only "
"axisymmetric meshes (self.is_axisymmetric must be True).")
required = ("ff_1dp_fwd", "ff_1dp_rev",
"dl_par_1dp_fwd", "dl_par_1dp_rev")
missing = [a for a in required
if getattr(self.planes[0], a, None) is None]
if missing:
raise RuntimeError("Parallel-derivative data missing in plane 0: "
+ ", ".join(missing))
nphi, n_n = field_2d.shape
if (nphi != self.nphi) or (n_n != self.planes[0].n_n):
raise ValueError(f"field_2d has shape {field_2d.shape}, "
f"expected ({self.nphi}, {self.planes[0].n_n}).")
plane0 = self.planes[0]
required = ("ff_1dp_fwd", "ff_1dp_rev",
"dl_par_1dp_fwd", "dl_par_1dp_rev", "bfield")
if not all(hasattr(plane0, attr) for attr in required):
missing = [a for a in required if not hasattr(plane0, a)]
raise AttributeError("Plane is missing required attributes: "
+ ", ".join(missing))
# -------- pre-compute constant coefficients -----------------------------
# forward / reverse arc-length arrays
l_r = np.asarray(plane0.dl_par_1dp_fwd)
l_l = np.asarray(plane0.dl_par_1dp_rev)
l_tot = l_r + l_l
coeff_l = -l_r / (l_l * l_tot) # shape (n_n,)
coeff_0 = (l_r - l_l) / (l_l * l_r)
coeff_r = l_l / (l_r * l_tot)
# overall sign (same on all nodes for an axisymmetric mesh)
sgn = int(bt_sign)
# -------- main loop over φ planes ---------------------------------------
out = np.empty_like(field_2d)
fwd_mat = plane0.ff_1dp_fwd # CSR (n_n × n_n)
rev_mat = plane0.ff_1dp_rev
for iphi in range(nphi):
iphi_l = iphi - 1 if iphi > 0 else nphi - 1
iphi_r = (iphi + 1) % nphi
# project neighbouring planes onto the current one
f_l = rev_mat.dot(field_2d[iphi_l]) # (n_n,)
f_r = fwd_mat.dot(field_2d[iphi_r])
out[iphi] = sgn * (
coeff_l * f_l +
coeff_0 * field_2d[iphi] +
coeff_r * f_r
)
return out
###############################################################################
# Example usage of the overall mesh class.
###############################################################################
if __name__ == "__main__":
# Create an axisymmetric mesh (only one plane).
axis_mesh = Mesh(is_axisymmetric=True)
print("Axisymmetric mesh:")
print("Number of planes:", len(axis_mesh.planes))
axis_mesh.get_plane().plot_surfaces()
# Create a non-axisymmetric mesh (multiple planes).
nonaxis_mesh = Mesh(is_axisymmetric=False)
print("Non-axisymmetric mesh:")
print("Number of planes:", len(nonaxis_mesh.planes))