import os
import warnings
import numpy as np
from matplotlib.tri import Triangulation
from xgc_analysis.accessor_mixin import ArrayAccessorMixin
from xgc_analysis.mesh_data import MeshData
from xgc_analysis.plane_data import PlaneData
from xgc_analysis.read_bp_file import ReadBPFile
[docs]
class FieldData(ArrayAccessorMixin):
"""
Reader for XGC ``xgc.2d.XXXXX.bp`` and ``xgc.3d.XXXXX.bp`` field files.
Data are stored in ``self.data`` with the structure::
self.data[var_name][file_step_index] = PlaneData | MeshData | scalar | np.ndarray
``FieldData`` inherits :class:`ArrayAccessorMixin`, which provides
``get_array(var_name)`` for converting stored values to plain NumPy arrays
for plotting and analysis.
"""
def __init__(self, mesh, work_dir="./", file_indices=None, is_axisymmetric=False, split_n0_turb=False):
"""
Initialize the field-data reader.
Parameters
----------
mesh : Mesh
Mesh object used to construct :class:`PlaneData` / :class:`MeshData`
wrappers.
work_dir : str, optional
Directory containing ``xgc.2d.XXXXX.bp`` / ``xgc.3d.XXXXX.bp`` files.
file_indices : list[int], optional
List of file indices to read. If omitted, reads ``[0]``.
is_axisymmetric : bool, optional
If ``True``, reads ``xgc.2d`` files. Otherwise reads ``xgc.3d`` files.
split_n0_turb : bool, optional
Reserved for future post-processing; currently stored but not used.
"""
self.mesh = mesh
self.work_dir = work_dir
self.file_indices = file_indices if file_indices is not None else [0]
self.is_axisymmetric = is_axisymmetric
self.split_n0_turb = split_n0_turb
self.data = {}
if is_axisymmetric:
self.required_vars = [
"depsidt", "dethetadt", "dpot", "eden",
"epsi", "etheta", "iden", "pot0", "time"
]
self._read_2d_files(mesh)
else:
self.required_vars = [
"aparh", "apars", "dBphi", "dBpsi", "dBtheta",
"dpot", "eden", "ejpar", "epara", "epara2",
"epsi", "etheta", "iden", "ijpar", "phi_right",
"pot0", "time"
]
self._read_3d_files(mesh)
def _read_2d_files(self, mesh):
"""
Read ``xgc.2d.XXXXX.bp`` files into ``self.data``.
2D field arrays are wrapped as :class:`PlaneData` objects.
Scalar values (notably ``time``) are stored directly.
"""
for i, idx in enumerate(self.file_indices):
fname = os.path.join(self.work_dir, f"xgc.2d.{idx:05d}.bp")
file_data = ReadBPFile(fname)
for step, variables in file_data.items():
for var in self.required_vars:
val = variables.get(var, 0.0)
if var == "time":
self.data.setdefault(var, {})[i] = val
continue
if isinstance(val, np.ndarray):
if val.ndim == 2 and val.shape[0] == 1:
val = val[0]
val = PlaneData(
plane=mesh.get_plane(0),
data_array=val,
n_components=1,
dtype=val.dtype,
)
self.data.setdefault(var, {})[i] = val
def _read_3d_files(self, mesh):
"""
Read ``xgc.3d.XXXXX.bp`` files into ``self.data``.
3D mesh fields are wrapped as :class:`MeshData` when truly toroidally
varying (shape ``(nphi, n_n)``). Axisymmetric arrays (``(n_n,)`` or
``(1, n_n)``) are stored as :class:`PlaneData` to avoid duplicating data
across planes. Scalars and special arrays (e.g. ``phi_right``) are stored
directly.
"""
for i, idx in enumerate(self.file_indices):
fname = os.path.join(self.work_dir, f"xgc.3d.{idx:05d}.bp")
file_data = ReadBPFile(fname)
for step, variables in file_data.items():
for var in self.required_vars:
val = variables.get(var, 0.0)
if var == "time":
self.data.setdefault(var, {})[i] = val
continue
if var == "phi_right":
self.data.setdefault(var, {})[i] = np.squeeze(val)
continue
if isinstance(val, np.ndarray):
if val.ndim == 2 and val.shape[0] == mesh.nphi:
# val shape is (nphi, n_n)
val = MeshData(
mesh,
data_array=val,
n_components=1,
dtype=val.dtype,
mesh_is_axisym=mesh.is_axisymmetric
)
elif val.ndim == 2 and val.shape[0] == 1:
val = PlaneData(
plane=mesh.get_plane(0),
data_array=val[0],
n_components=1,
dtype=val.dtype
)
elif val.ndim == 1:
val = PlaneData(
plane=mesh.get_plane(0),
data_array=val,
n_components=1,
dtype=val.dtype,
)
self.data.setdefault(var, {})[i] = val
# --- Typed accessors ---
[docs]
def get_field(self, var_name, step_index=0):
"""Return the raw stored field item for a variable/time selection."""
return self.get_item(var_name, step_index)
[docs]
def get_plane_data(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as :class:`PlaneData`."""
return self.get_as(var_name, step_index, PlaneData)
[docs]
def get_mesh_data(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as :class:`MeshData`."""
return self.get_as(var_name, step_index, MeshData)
[docs]
def get_scalar(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as a scalar numeric value."""
return self.get_as(var_name, step_index, (int, float, np.integer, np.floating))
[docs]
def export_vtu(self, mesh):
"""
Export loaded field data to ``.vtu`` files for visualization.
Parameters
----------
mesh : Mesh
Mesh object providing geometry/connectivity for the exported grid.
"""
import pyvista as pv
# Ensure the output directory for VTK files exists
vtu_dir = os.path.join(self.work_dir, "vtus")
if not os.path.exists(vtu_dir):
os.makedirs(vtu_dir)
# Warn if multiple file indices are used, as this may use significant memory
if len(self.file_indices) > 1:
warnings.warn(
"`export_vtu()` is called with multiple file_indices. "
"This will work, but may require a large amount of memory if you plan to generate many VTK files. "
"Consider instantiating separate FieldData objects for individual VTK file indices. ",
UserWarning
)
# Generate triangle connectivity (could be replaced with TriangulationData from plane.py)
triObj = Triangulation(mesh.rz[:,0], mesh.rz[:,1], mesh.nd_connect_list)
# Prepare 3D points for VTK: 2D poloidal points with a zero Z component (VTK expects 3D coordinates)
points_vtu = np.column_stack([mesh.rz[:,0], mesh.rz[:,1], np.zeros_like(mesh.rz[:,0])])
# Prepare cells for VTK
num_cells = triObj.triangles.shape[0]
cell_types = np.full(len(triObj.triangles), pv.CellType.TRIANGLE, dtype=np.uint8)
cell_array = np.hstack([
np.full((num_cells,1), 3), # triangle has 3 vertices
triObj.triangles
]).astype(np.int64)
# Choose filename format based on dimensionality
dim = "2d" if self.is_axisymmetric else "3d"
for i, idx in enumerate(self.file_indices):
fname = os.path.join(vtu_dir, f"xgc.{dim}.{idx:05d}.vtu")
# Create unstructured grid
grid = pv.UnstructuredGrid(cell_array, cell_types, points_vtu)
# -----------------------------------------------------
# Dataset 1: Point data on XGC poloidal grid
# expected shape: [nindices, nphi, nnode]
# -----------------------------------------------------
for var in self.required_vars:
if var in ("time", "phi_right"): # TODO: handle these more robustly
continue
datum = self.data[var][i]
if isinstance(datum, MeshData):
data_array = datum.get_data()[0]
elif isinstance(datum, PlaneData):
data_array = datum.get_data()
else:
continue
grid.point_data[var] = data_array
# -----------------------------------------------------
# Dataset 2: Global field data
# expected shape: [nindices]
# -----------------------------------------------------
grid.field_data["time"] = np.array([self.get_array("time")[i]], dtype=np.float64)
# -----------------------------------------------------
# Dataset 3: Additional data not present in the original xgc.3d.#####.bp file
# Manually added to support post-processing and diagnostics
# -----------------------------------------------------
grid.field_data["step"] = np.array([idx], dtype=np.int32)
# Save to .vtu file
grid.save(fname)