"""Time-resolved one-dimensional slices of scalar XGC mesh diagnostics.
The catalog supplies the variable's logical steps and physics times. Each
frame is reduced to a mapped radial profile, a theta-star surface trace, or
one branch of a line of sight before the next batch is read.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
import numpy as np
from .mesh_slicing import (flux_surface_profile, line_of_sight,
split_line_at_psi_minimum, surface_lines)
from xgc_analysis.mesh import PlaneData
[docs]
@dataclass
class MeshSliceTimeMap:
"""One scalar slice sampled across physical time.
Attributes
----------
coordinate : ndarray
Increasing horizontal coordinates: normalized psi or theta-star radians.
time_s, gsteps : ndarray
Physics time in seconds and matching logical gsteps.
values : ndarray
Slice values with shape ``(n_time, n_coordinate)``.
coordinate_name, statistic, detail : str
Horizontal coordinate, chosen reduction, and slice description.
"""
coordinate: np.ndarray
time_s: np.ndarray
gsteps: np.ndarray
values: np.ndarray
coordinate_name: str
statistic: str
detail: str
def _plane_values(array, mesh, vertex_axis, plane_axis, plane_index, other_indices):
"""Apply Monitor-style spatial selectors to one catalog frame."""
field = np.asarray(array, dtype=float)
if not 0 <= vertex_axis < field.ndim or field.shape[vertex_axis] != mesh.get_plane(0).n_n:
raise ValueError("Vertex axis does not match the XGC mesh")
if plane_axis is not None:
if not 0 <= plane_axis < field.ndim or field.shape[plane_axis] != mesh.nphi:
raise ValueError("Toroidal plane axis does not match the XGC mesh")
if not 0 <= plane_index < mesh.nphi:
raise ValueError("Toroidal plane index is out of range")
selection = tuple(
slice(None) if axis == vertex_axis else
plane_index if axis == plane_axis else
int(other_indices.get(axis, 0))
for axis in range(field.ndim)
)
values = np.asarray(field[selection], dtype=float)
if values.shape != (mesh.get_plane(plane_index).n_n,):
raise ValueError("Selected frame must contain one scalar per mesh vertex")
return values
def _physics_times(catalog, product_key, records):
"""Use catalog fragment times, falling back to the product's time channel."""
times = {int(record.logical_step): getattr(record.selected_fragment, "time", None)
for record in records}
missing = [step for step, value in times.items()
if value is None or not np.isfinite(float(value))]
if missing:
product = catalog.get_product(product_key)
if "time" not in product.variables:
raise ValueError(f"{product_key} has no physics time for {len(missing)} selected steps")
result = catalog.read_arrays(product_key, variables=["time"], steps=missing)
for step in missing:
raw = result.arrays.get("time", {}).get(step)
if raw is None or np.asarray(raw).size != 1:
raise ValueError(f"No scalar physics time for {product_key} gstep {step}")
times[step] = float(np.asarray(raw).reshape(-1)[0])
values = np.asarray([float(times[int(record.logical_step)]) for record in records])
if not np.all(np.isfinite(values)) or np.any(np.diff(values) <= 0):
raise ValueError("Selected physics times must be finite and strictly increasing")
return values
def _line_origin(plane, origin, r0, z0):
"""Resolve the requested sight-line origin on the selected plane."""
if origin == "axis":
return float(plane.axis_r), float(plane.axis_z)
if origin == "xpoint":
return float(plane.x_r), float(plane.x_z)
if origin == "custom" and r0 is not None and z0 is not None:
return float(r0), float(z0)
raise ValueError("Line origin must be axis, xpoint, or custom R,Z")
def _interpolate_samples(coordinate, values, grid):
"""Interpolate finite samples to a fixed grid without extrapolation."""
coordinate = np.asarray(coordinate, dtype=float)
values = np.asarray(values, dtype=float)
valid = np.isfinite(coordinate) & np.isfinite(values)
if np.count_nonzero(valid) < 2:
return np.full(grid.shape, np.nan)
order = np.argsort(coordinate[valid], kind="stable")
x = coordinate[valid][order]
y = values[valid][order]
unique, inverse = np.unique(x, return_inverse=True)
if unique.size < 2:
return np.full(grid.shape, np.nan)
if unique.size != x.size:
y = np.bincount(inverse, weights=y)/np.bincount(inverse)
return np.interp(grid, unique, y, left=np.nan, right=np.nan)
[docs]
def compute_mesh_slice_time_map(
simulation, product_key, variable, *, statistic="average", vertex_axis=0,
plane_axis=None, plane_index=0, other_indices=None, surface_index=0,
line_origin="axis", line_r0=None, line_z0=None, line_angle_deg=0.0,
line_branch="toward", line_points=500, line_end_ignore=5,
n_coordinate=256, gstep_range=None, time_range_ms=None, frame_stride=1,
last_frame_index=None, check_cancel: Callable[[], None] | None=None,
) -> MeshSliceTimeMap:
"""Reduce catalog frames to a contour-ready mesh-slice time map.
Radial averages and RMS profiles use only ``Plane.surf_map`` surfaces,
which cross the outer midplane. A theta trace uses one selected surface.
The line of sight is split at its interior minimum psi and one branch is
interpolated to a fixed normalized-psi grid. Missing spatial samples stay
NaN; time is never inferred from output-frame spacing.
Parameters
----------
simulation : Simulation
Initialized XGC Simulation with mesh, magnetic field, and catalog.
product_key, variable : str
Source product and scalar vertex variable.
statistic : {"average", "rms", "surface", "line"}
Spatial reduction to apply independently at every time step.
vertex_axis, plane_axis, plane_index, other_indices
Spatial axis indices and selectors for one frame of the variable.
surface_index : int
Zero-based surface for a theta-star time map.
line_origin, line_r0, line_z0, line_angle_deg
Sight-line origin and angle from +R for a line time map.
line_branch : {"toward", "away"}
Side of the interior psi minimum to retain.
line_points, line_end_ignore, n_coordinate : int
Sight-line samples, ignored endpoint samples, and contour-grid points.
gstep_range, time_range_ms : tuple or None
Optional inclusive limits on logical gsteps and physics time in ms.
frame_stride : int
Take every Nth selected frame after applying window limits.
last_frame_index : int or None
Optional last eligible ordinal frame; ``None`` includes all frames.
check_cancel : callable or None
Invoked between catalog batches and frame reductions.
Returns
-------
MeshSliceTimeMap
Increasing horizontal coordinate, physics times, gsteps, and values.
"""
if statistic not in {"average", "rms", "surface", "line"}:
raise ValueError("Unknown mesh time-map statistic")
if frame_stride < 1 or n_coordinate < 2:
raise ValueError("Frame stride must be positive and grid resolution at least two")
if statistic == "line" and line_branch not in {"toward", "away"}:
raise ValueError("Line branch must be toward or away from minimum psi")
catalog = simulation.catalog
records = list(catalog.available_steps(product_key, variable))
if last_frame_index is not None:
if not 0 <= last_frame_index < len(records):
raise ValueError("Selected frame is outside the available time series")
records = records[:last_frame_index + 1]
if gstep_range is not None:
records = [record for record in records
if gstep_range[0] <= record.logical_step <= gstep_range[1]]
if not records:
raise ValueError("No variable frames in the requested gstep window")
times = _physics_times(catalog, product_key, records)
if time_range_ms is not None:
keep = (times*1.0e3 >= time_range_ms[0]) & (times*1.0e3 <= time_range_ms[1])
records = [record for record, selected in zip(records, keep) if selected]
times = times[keep]
records = records[::frame_stride]
times = times[::frame_stride]
if len(records) < 2:
raise ValueError("A time map needs at least two selected time slices")
mesh = simulation.mesh
plane = mesh.get_plane(plane_index)
if statistic in {"average", "rms"}:
mapped = np.asarray(plane.surf_map, dtype=int)
coordinate = np.asarray(plane.psi_surf[mapped]/plane.x_psi, dtype=float)
order = np.argsort(coordinate, kind="stable")
mapped = mapped[order]
coordinate = coordinate[order]
coordinate_name = "psi"
detail = "Outer-midplane flux surfaces"
elif statistic == "surface":
if not 0 <= surface_index < len(plane.psi_surf):
raise ValueError("Selected surface index is outside the mesh surface range")
vertices = np.asarray(plane.get_surface_vertex_indices(surface_index), dtype=int)
theta = np.asarray(plane.theta[vertices], dtype=float)
valid = np.isfinite(theta) & (theta >= -2*np.pi)
if np.isclose(plane.psi_surf[surface_index]/plane.x_psi, 1.0, atol=1e-6):
valid &= plane.rz[vertices, 1] >= plane.x_z
theta = (theta[valid] + np.pi) % (2*np.pi) - np.pi
if theta.size < 2:
raise ValueError("Selected surface has fewer than two valid theta-star points")
coordinate = np.linspace(float(theta.min()), float(theta.max()), n_coordinate)
coordinate_name = "theta"
detail = f"Surface {surface_index}, psi_N={plane.psi_surf[surface_index]/plane.x_psi:.3f}"
else:
origin_r, origin_z = _line_origin(plane, line_origin, line_r0, line_z0)
if plane_index == 0 or mesh.is_axisymmetric:
psi_values = np.asarray(simulation.magnetic_field.psi_pd.get_data(), dtype=float)
else:
rz = plane.rz
psi_values = np.asarray(simulation.magnetic_field.eq_psi_rz_spline.ev(rz[:, 0], rz[:, 1]), dtype=float)
psi_data = PlaneData(plane, data_array=psi_values)
reference = PlaneData(plane, data_array=np.zeros(plane.n_n))
_, _, psi, _ = line_of_sight(reference, psi_data, r0=origin_r, z0=origin_z,
angle_degrees=line_angle_deg, n_points=line_points)
_, _, turning = split_line_at_psi_minimum(
psi, np.where(np.isfinite(psi), 0.0, np.nan),
ignore_end_points=line_end_ignore)
branch = slice(None, turning + 1) if line_branch == "toward" else slice(turning, None)
branch_psi = psi[branch]
finite = branch_psi[np.isfinite(branch_psi)]
if finite.size < 2 or finite.min() == finite.max():
raise ValueError("Selected sight-line branch has no finite psi interval")
coordinate = np.linspace(float(finite.min()), float(finite.max()), n_coordinate)
coordinate_name = "psi"
detail = f"Line of sight: {line_branch} minimum psi_N"
if not np.all(np.isfinite(coordinate)) or np.any(np.diff(coordinate) <= 0):
raise ValueError("Mesh slice coordinate must be finite and increasing")
steps = [int(record.logical_step) for record in records]
rows = []
for start in range(0, len(steps), 4):
if check_cancel is not None:
check_cancel()
batch = steps[start:start + 4]
result = catalog.read_arrays(product_key, variables=[variable], steps=batch)
for step in batch:
if check_cancel is not None:
check_cancel()
raw = result.arrays.get(variable, {}).get(step)
if raw is None:
raise ValueError(f"Missing {product_key}:{variable} at gstep {step}")
values = _plane_values(raw, mesh, vertex_axis, plane_axis, plane_index,
other_indices or {})
field = PlaneData(plane, data_array=values)
if statistic in {"average", "rms"}:
profile = flux_surface_profile(field, statistic)
rows.append(np.asarray(profile[mapped], dtype=float))
elif statistic == "surface":
_, _, theta, samples = surface_lines(field, [surface_index])[0]
rows.append(_interpolate_samples(theta, samples, coordinate))
else:
_, _, _, samples = line_of_sight(
field, psi_data, r0=origin_r, z0=origin_z,
angle_degrees=line_angle_deg, n_points=line_points)
rows.append(_interpolate_samples(branch_psi, samples[branch], coordinate))
return MeshSliceTimeMap(
coordinate=coordinate, time_s=np.asarray(times, dtype=float),
gsteps=np.asarray(steps, dtype=int), values=np.asarray(rows, dtype=float),
coordinate_name=coordinate_name, statistic=statistic, detail=detail,
)