import numpy as np
from .plane_data import PlaneData
[docs]
class MeshData:
"""
Stores field data for a mesh as a list of PlaneData instances (one per plane).
The data for each toroidal plane is stored in a separate ``PlaneData``.
Axisymmetric geometry can be shared while the toroidal data slices differ.
``get_data()`` stacks those slices with plane index on axis 0.
"""
def __init__(self, mesh, data_array=None, field_type='scalar', n_components=1, dtype=np.float64, mesh_is_axisym=False):
"""Construct mesh-shaped scalar or vector values.
Parameters
----------
mesh : Mesh
Mesh whose planes provide node counts and geometry.
data_array : ndarray or None, optional
Initial values indexed first by toroidal plane.
field_type : str, optional
Legacy field-type argument; component count determines the stored
scalar or vector type.
n_components : int, optional
Number of components per vertex; one means scalar data.
dtype : numpy dtype, optional
Storage type for newly allocated plane values.
mesh_is_axisym : bool, optional
Reuse the mesh's first geometry plane for every data slice.
"""
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):
"""Stack values from all toroidal planes into one NumPy array.
Returns
-------
ndarray
Toroidal plane is axis 0. The remaining axes preserve each
``PlaneData`` array's shape: typically ``(n_vertices,)`` for
loaded scalars or ``(n_vertices, n_components)`` for vectors.
"""
return np.array([plane.get_data() for plane in self.planes])
[docs]
def get_plane(self, plane_index):
"""
Return the ``PlaneData`` at zero-based ``plane_index``.
Raise ``IndexError`` when the index is outside stored planes.
"""
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, *, fsa: str = "volume"):
"""Store the surface average in ``flux_avg_psi``.
``fsa`` selects the plane weighting scheme. This mutates the object
and returns no array.
"""
self.flux_avg_psi = self.mesh.flux_avg_to_surf(self.get_data(), fsa=fsa)
[docs]
def flux_avg_from_surf(self, *, fsa: str = "volume"):
"""Project stored surface values to vertices in ``flux_avg_plane``.
``fsa`` selects the plane weighting scheme. The method returns no
value and replaces the cached ``PlaneData``.
"""
vertex_arr = self.mesh.flux_avg_from_surf(self.flux_avg_psi, fsa=fsa)
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, *, fsa: str = "volume"):
"""Average this field and cache its vertex projection.
``fsa`` selects the plane weighting scheme. The projected
``PlaneData`` is stored in ``flux_avg_plane``; return ``None``.
"""
vertex_arr = self.mesh.flux_avg_to_from_surf(self.get_data(), fsa=fsa)
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, *, fsa: str = "volume"):
"""Cache the toroidally averaged vector norm by flux surface.
``fsa`` selects the plane weighting scheme. Requires vector data;
stores ``flux_avg_norm_psi`` and returns ``None``.
"""
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(), fsa=fsa)
[docs]
def toroidal_average(self):
"""Cache the plane-wise arithmetic mean in ``toroidal_avg_data``.
Returns
-------
None
The average is stored as a ``PlaneData`` on this ``MeshData``.
"""
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.
``ref_plane`` selects the reference plane passed to the mesh mapping.
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}")