"""Compare the XGC and TRANSP gradients of flux-surface averaged potential."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import h5py
import numpy as np
from xgc_analysis.analysis.geometry import interp_radius_at_psi, midplane_radius_maps
[docs]
@dataclass
class PotentialGradientComparison:
"""Profiles of ``-d<phi>/d(r/a)`` in volts on the XGC surface grid.
Attributes
----------
psi_n, r_over_a : ndarray
Normalized poloidal flux and outboard-midplane minor-radius coordinate.
xgc, transp : ndarray or None
XGC gradient and optional TRANSP comparison, both in volts.
gstep_range, time_range_s : tuple
Averaged gsteps and optional physics-time range in seconds.
transp_path, field_product : str or None
Reference path and XGC potential product used.
"""
psi_n: np.ndarray
r_over_a: np.ndarray
xgc: np.ndarray
transp: np.ndarray | None
gstep_range: tuple[int, int]
time_range_s: tuple[float, float] | None
transp_path: str | None
field_product: str
def _node_profile(value, n_nodes: int, product: str) -> np.ndarray:
"""Reduce a field sample to vertices, averaging its toroidal axis if present."""
field = np.squeeze(np.asarray(value, dtype=float))
if field.shape == (n_nodes,):
return field
if field.ndim == 2 and field.shape[1] == n_nodes:
return field.mean(axis=0)
if field.ndim == 2 and field.shape[0] == n_nodes:
return field.mean(axis=1)
raise ValueError(f"{product}:pot0 has shape {field.shape}; expected {n_nodes} vertices with an optional toroidal axis.")
[docs]
def compare_potential_gradient(
simulation,
*,
transp_path: str | Path | None = None,
gstep_range: tuple[int, int] | None = None,
time_range_ms: tuple[float, float] | None = None,
max_default_samples: int = 30,
) -> PotentialGradientComparison:
"""Average XGC ``pot0`` and compare ``-d<phi>/d(r/a)`` with TRANSP ``VRPOT``.
The XGC field is averaged over toroidal planes when necessary, then over
each flux surface and the selected diagnostic times. The separatrix FSA
excludes its legs below the X-point. Both derivatives use
the outboard midplane ``r/a`` coordinate of the XGC mesh. The result has
units of volts, rather than the V/m units of a physical radial field.
Parameters
----------
simulation
Initialized XGC Simulation containing the needed geometry and diagnostics.
transp_path
Optional TRANSP HDF5 reference path.
gstep_range
Optional inclusive gstep bounds.
time_range_ms
Optional inclusive physics-time bounds in milliseconds.
max_default_samples
Number of latest samples used when no window is supplied.
Returns
-------
PotentialGradientComparison
XGC and optional TRANSP gradient profiles.
"""
catalog = simulation.catalog
product = next((key for key in ("xgc.2d.bp", "xgc.3d.bp")
if key in catalog.products and "pot0" in catalog.get_product(key).variables), None)
if product is None:
raise FileNotFoundError("xgc.2d.bp:pot0 or xgc.3d.bp:pot0 is required.")
records = list(catalog.available_steps(product, "pot0"))
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(f"The selected range contains no {product}:pot0 samples.")
steps = [int(record.logical_step) for record in records]
times = {int(record.logical_step): record.selected_fragment.time for record in records}
if time_range_ms is not None and any(value is None for value in times.values()):
if "time" not in catalog.get_product(product).variables:
raise ValueError(f"{product} has no time coordinate for a time-window selection.")
read_times = catalog.read_arrays(product, ["time"], steps=steps).arrays["time"]
times = {step: float(np.asarray(read_times[step]).reshape(-1)[0]) for step in steps}
if time_range_ms is not None:
records = [record for record in records
if times[int(record.logical_step)] is not None
and time_range_ms[0] <= 1.0e3*times[int(record.logical_step)] <= time_range_ms[1]]
elif gstep_range is None:
records = records[-max_default_samples:]
if not records:
raise ValueError(f"The selected averaging window has no {product}:pot0 samples.")
steps = [int(record.logical_step) for record in records]
plane = simulation.mesh.get_plane(0)
surfaces = np.asarray(plane.surf_map, dtype=int)
psi_n = np.asarray(plane.psi_surf, dtype=float)[surfaces]/float(plane.x_psi)
# The outer-midplane radius can fold on far-SOL surfaces; the IDL plot
# covers only the near-SOL region around the separatrix.
near_sol = np.isfinite(psi_n) & (psi_n <= 1.05)
surfaces, psi_n = surfaces[near_sol], psi_n[near_sol]
psi_mid, r_out, _ = midplane_radius_maps(plane)
r_sep = interp_radius_at_psi(psi_mid, r_out, 1.0, "outer")
good_mid = np.isfinite(psi_mid) & np.isfinite(r_out)
r_profile = np.interp(psi_n, psi_mid[good_mid], r_out[good_mid])
r_over_a = (r_profile - float(plane.axis_r))/(r_sep - float(plane.axis_r))
if len(psi_n) and np.isclose(psi_n[0], 0.0):
r_over_a[0] = 0.0
if len(r_over_a) < 2 or not np.all(np.diff(r_over_a) > 0):
raise ValueError("The XGC outer-midplane r/a surface map must be strictly increasing.")
read = catalog.read_arrays(product, ["pot0"], steps=steps).arrays["pot0"]
mean_potential = np.mean([
plane.flux_avg_to_surf(
_node_profile(read[step], plane.n_n, product),
exclude_separatrix_legs=True,
)[surfaces]
for step in steps
], axis=0)
xgc = -np.gradient(mean_potential, r_over_a, edge_order=2 if len(r_over_a) > 2 else 1)
transp = None
reference_path = None
if transp_path is not None:
reference_path = str(Path(transp_path).expanduser().resolve())
with h5py.File(reference_path) as file:
if "PSIN" not in file or "VRPOT" not in file:
raise KeyError("TRANSP HDF5 file requires PSIN and VRPOT datasets.")
tr_psi = np.asarray(file["PSIN"], dtype=float).reshape(-1)
tr_potential = np.asarray(file["VRPOT"], dtype=float).reshape(-1)
if tr_psi.shape != tr_potential.shape or len(tr_psi) < 2:
raise ValueError("TRANSP PSIN and VRPOT must have matching profiles of at least two points.")
order = np.argsort(tr_psi)
tr_psi, tr_potential = tr_psi[order], tr_potential[order]
tr_r_over_a = np.interp(tr_psi, psi_n, r_over_a)
if not np.all(np.diff(tr_r_over_a) > 0):
raise ValueError("TRANSP PSIN does not map to a strictly increasing XGC r/a grid.")
tr_gradient = -np.gradient(tr_potential, tr_r_over_a, edge_order=2 if len(tr_psi) > 2 else 1)
transp = np.interp(psi_n, tr_psi, tr_gradient, left=np.nan, right=np.nan)
transp[psi_n > 1.0] = np.nan
selected_times = [float(times[step]) for step in steps if times[step] is not None]
return PotentialGradientComparison(
psi_n=psi_n,
r_over_a=r_over_a,
xgc=xgc,
transp=transp,
gstep_range=(min(steps), max(steps)),
time_range_s=(min(selected_times), max(selected_times)) if selected_times else None,
transp_path=reference_path,
field_product=product,
)