import numpy as np
from xgc_analysis.plane_data import PlaneData
[docs]
class MeshData:
"""
Stores field data for a mesh as a list of PlaneData instances (one per plane).
The constructor accepts:
- vertex_counts: a list of integers indicating the number of vertices on each plane.
- field_type: 'scalar' or 'vector'.
- n_components: for vector fields (2 or 3); for scalar fields this should be 1.
- dtype: the NumPy data type for the stored data.
- mesh_is_axisym (bool): True if the mesh is axisymmetric (i.e. all planes have the same number of vertices
and the same (R,Z) coordinates), False otherwise.
The data for each plane is stored as a separate PlaneData instance.
"""
def __init__(self, mesh, data_array=None, field_type='scalar', n_components=1, dtype=np.float64, mesh_is_axisym=False):
self.mesh = mesh # geometry reference
self.nphi = mesh.nphi
self.vertex_counts = [p.n_n for p in mesh.planes] # list of ints, one per plane.
self.n_components = n_components
if (self.n_components > 1):
self.field_type = 'vector'
else:
self.field_type = 'scalar'
self.dtype = dtype
self.mesh_is_axisym = mesh_is_axisym
# Create one PlaneData instance per plane.
if (self.mesh_is_axisym):
if (data_array is not None):
self.planes = [PlaneData(mesh.planes[0], n_components=n_components, dtype=dtype, \
data_array=data_array[i]) for i in range(self.nphi)]
else:
self.planes = [PlaneData(mesh.planes[0], n_components=n_components, dtype=dtype) for i in range(self.nphi)]
else:
if (data_array is not None):
self.planes = [PlaneData(mesh.planes[i], n_components=n_components, dtype=dtype, \
data_array=data_array[i]) for i in range(self.nphi)]
else:
self.planes = [PlaneData(mesh.planes[i], n_components=n_components, dtype=dtype) for i in range(self.nphi)]
# Result placeholders
self.flux_avg_psi = np.empty((self.mesh.planes[0].nsurf)) # np.ndarray
if (self.n_components > 1):
self.flux_avg_norm_psi = np.empty((self.mesh.planes[0].nsurf)) # np.ndarray
self.flux_avg_plane = None # PlaneData
[docs]
def get_data(self):
"""
Returns the data from all planes as a 2D NumPy array with shape (nphi, n_vertices).
"""
return np.array([plane.get_data() for plane in self.planes])
[docs]
def get_plane(self, plane_index):
"""
Returns the PlaneData instance for the specified plane.
"""
if plane_index < 0 or plane_index >= self.nphi:
raise IndexError("Plane index out of range")
return self.planes[plane_index]
# ------------------------------------------------------------------
# Wrappers around Mesh helpers – store results locally
# ------------------------------------------------------------------
[docs]
def flux_avg_to_surf(self):
"""Compute the flux-surface average and store in `self.flux_avg_psi`."""
self.flux_avg_psi = self.mesh.flux_avg_to_surf(self.get_data())
[docs]
def flux_avg_from_surf(self):
"""Project the flux-surface average data back to vertices and package into PlaneData."""
vertex_arr = self.mesh.flux_avg_from_surf(self.flux_avg_psi)
self.flux_avg_plane = PlaneData(self.mesh.planes[0], data_array=vertex_arr,
n_components=vertex_arr.shape[1] if vertex_arr.ndim == 2 else 1)
[docs]
def flux_avg_to_from_surf(self):
"""Compute the flux-surface average of `self.get_data()` and project back to the plane in `self.flux_avg_plane`."""
vertex_arr = self.mesh.flux_avg_to_from_surf(self.get_data())
self.flux_avg_plane = PlaneData(self.mesh.planes[0], data_array=vertex_arr,
n_components=vertex_arr.shape[1] if vertex_arr.ndim == 2 else 1)
[docs]
def flux_avg_to_surf_norm(self):
"""Flux‑surface average of vector‑norm, averaged over the toroidal direction."""
if self.data.ndim == 2:
raise ValueError("flux_avg_to_surf_norm only for vector data.")
self.flux_avg_norm_psi = self.mesh.flux_avg_to_surf_norm(self.get_data())
[docs]
def toroidal_average(self):
"""Return toroidal average as PlaneData (also cached on Mesh)."""
if not self.mesh_is_axisym:
raise RuntimeError("average_plane is only valid for axisymmetric meshes.")
vertex_arr = self.mesh.toroidal_average(self.get_data())
pd = PlaneData(self.mesh.planes[0], data_array=vertex_arr,
n_components=vertex_arr.shape[1] if vertex_arr.ndim == 2 else 1)
self.toroidal_avg_data = pd
[docs]
def ff_map_real2ff(self):
"""
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).
The output format is list of MeshData objects of the same type, one
for the left, one for the right plane.
Returns:
MeshData (vector-type): Field on a mesh mapped to field-aligned representation.
"""
if (self.nphi == 1):
raise ValueError("Data must not be axisymmetric for field-aligned mapping.")
mapped_data = self.mesh.ff_map_real2ff(self.get_data())
return [MeshData(self.mesh,data_array=mapped_data[i],n_components=self.n_components,mesh_is_axisym=self.mesh.is_axisymmetric) for i in range(2)]
[docs]
def apply_fieldline_mapping(self, *, ref_plane: int = 0):
"""
Pull the data of *every* plane onto the vertices of *ref_plane*
along magnetic-field lines.
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.
"""
out = self.mesh.apply_fieldline_mapping(self.get_data(),ref_plane = ref_plane)
return out
[docs]
def evaluate_scalar_gradients(self) -> "MeshData":
"""
Return (R,Z)/(psi,theta)-gradients as a new MeshData object (vector field).
Returns
-------
MeshData # field_type = 'vector'
"""
if getattr(self, "field_type", "scalar") != "scalar":
raise TypeError("Gradient evaluation is only defined for scalar MeshData.")
grad_arr = self.mesh.evaluate_scalar_gradients(self.get_data())
return MeshData(
mesh=self.mesh,
data=grad_arr,
field_type="vector"
)
[docs]
def evaluate_grad_par(self, bt_sign: int) -> "MeshData":
"""
Parallel derivative b·∇(f) returned as a new *scalar* MeshData.
Parameters
----------
bt_sign : int
Sign of the toroidal magnetic field (±1).
Returns
-------
MeshData # field_type = 'scalar'
"""
if getattr(self, "field_type", "scalar") != "scalar":
raise TypeError("Parallel derivatives are defined only for scalar MeshData.")
grad_par = self.mesh.evaluate_grad_par(self.get_data(), bt_sign)
return MeshData(
mesh=self.mesh,
data=grad_par,
field_type="scalar"
)
[docs]
def evaluate_full_gradients(self, bt_sign: int) -> "MeshData":
"""
Return a 3-component gradient vector field as a new MeshData:
grads[..., 0] = ∂f/∂ψ
grads[..., 1] = ∂f/∂θ
grads[..., 2] = b·∇f
Parameters
----------
bt_sign : int
Sign of the toroidal magnetic field Bᵗ (±1).
Returns
-------
MeshData # field_type = 'vector'
"""
if getattr(self, "field_type", "scalar") != "scalar":
raise TypeError(
"Full-gradient evaluation is only defined for scalar MeshData."
)
# --- in-plane gradients (shape: nphi × n_n × 2)
grad_plane = self.mesh.evaluate_scalar_gradients(self.get_data())
# --- parallel gradient (shape: nphi × n_n)
grad_par = self.mesh.evaluate_grad_par(self.get_data(), bt_sign)
# --- stack into a 3-component vector ---------------------------
grad_full = np.concatenate(
(grad_plane, grad_par[..., np.newaxis]), axis=-1
) # shape: nphi × n_n × 3
return MeshData(
mesh=self.mesh,
data=grad_full,
field_type="vector"
)
###############################################################################
# Example usage:
###############################################################################
if __name__ == "__main__":
# Suppose we have a mesh with three planes.
# For example, in a non-axisymmetric mesh, the number of vertices per plane may differ.
vertex_counts = [100, 105, 98]
# Create a MeshData instance for a scalar field.
scalar_mesh = MeshData(vertex_counts, field_type='scalar', n_components=1, dtype=np.float64, mesh_is_axisym=False)
# Fill each plane with example scalar data.
for i, pd in enumerate(scalar_mesh.planes):
pd.data = np.linspace(0, 1, vertex_counts[i])
# Retrieve and print data for plane 1.
plane1_scalar = scalar_mesh.get_plane(1)
print("Scalar field for plane 1, shape:", plane1_scalar.data.shape)
# Create a MeshData instance for a vector field with 3 components.
vector_mesh = MeshData(vertex_counts, field_type='vector', n_components=3, dtype=np.float64, mesh_is_axisym=False)
# Fill each plane with random vector data.
for i, pd in enumerate(vector_mesh.planes):
pd.data = np.random.rand(vertex_counts[i], 3)
# Retrieve vector data for plane 1.
plane1_vector = vector_mesh.get_plane(1)
print("Vector field for plane 1, shape:", plane1_vector.data.shape)
# Simulate the case where plane 1's vector data is stored in transposed shape (3, n).
vector_mesh.planes[1].data = vector_mesh.planes[1].data.T
print("After simulating transposed storage, plane 1 shape:", vector_mesh.planes[1].data.shape)
# Extract the second component (index 1) from the vector field across all planes.
scalar_component_mesh = vector_mesh.extract_component(1)
for i, pd in enumerate(scalar_component_mesh.planes):
print(f"Extracted scalar component for plane {i}, shape: {pd.data.shape}")