"""Compare XGC prescribed transport coefficients with TRANSP profiles."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import h5py
import numpy as np
from xgc_analysis.data_access.products.diffusion_data import DiffusionCoefficientData
from xgc_analysis.analysis.geometry import interp_radius_at_psi, midplane_radius_maps
COEFFICIENTS = {
"particle_diffusivity": ("ptl_diffusivity", 0, "DIFFE", 1.0e-4, r"$D_e$ [m$^2$/s]"),
"pinch_velocity": ("ptl_pinch_velocity", 0, None, 1.0, r"$v_p$ [m/s]"),
"momentum_diffusivity": ("momentum_diffusivity", 1, "CHPHI", 1.0e-4, r"$\chi_\phi$ [m$^2$/s]"),
"ion_heat_conductivity": ("heat_conductivity", 1, "CONDI", 1.0e-4, r"$\chi_i$ [m$^2$/s]"),
"electron_heat_conductivity": ("heat_conductivity", 0, "CONDE", 1.0e-4, r"$\chi_e$ [m$^2$/s]"),
}
[docs]
@dataclass
class TransportCoefficientComparison:
"""Time-averaged XGC coefficient and optional mapped TRANSP reference.
Attributes
----------
quantity, ylabel : str
Selected coefficient key and its labeled physical units.
psi_n, r_over_a, xgc, transp : ndarray or None
Radial coordinates and XGC/optional TRANSP values on the XGC grid.
gstep_range, time_range_s : tuple
Averaged gsteps and optional physics-time range in seconds.
transp_path : str or None
Reference HDF5 path when used.
"""
quantity: str
psi_n: np.ndarray
r_over_a: np.ndarray
xgc: np.ndarray
transp: np.ndarray | None
ylabel: str
gstep_range: tuple[int, int]
time_range_s: tuple[float, float] | None
transp_path: str | None
def _catalog_time_by_gstep(catalog) -> dict[int, float]:
"""Find physics times already indexed for another XGC time-series product."""
for product in ("xgc.f2d.bp", "xgc.fsourcediag.bp", "xgc.heatdiag2.bp"):
try:
records = catalog.available_steps(product, "time")
except (KeyError, ValueError):
continue
times = {int(record.logical_step): float(record.selected_fragment.time)
for record in records if record.selected_fragment.time is not None}
if times:
return times
return {}
[docs]
def compare_transport_coefficients(
simulation,
quantity: str,
*,
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,
) -> TransportCoefficientComparison:
"""Average XGC coefficient BP steps and map a TRANSP HDF5 curve to psi_N.
TRANSP diffusivities are converted from cm^2/s to m^2/s. The supplied
``transp_path`` may be outside the XGC run; it is read only. TRANSP does
not provide a matching pinch-velocity variable in the reference HDF5.
Parameters
----------
simulation
Initialized XGC Simulation containing the needed geometry and diagnostics.
quantity
Coefficient or particle/energy balance selector.
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
-------
TransportCoefficientComparison
XGC and optional TRANSP profiles.
"""
if quantity not in COEFFICIENTS:
raise ValueError(f"Unknown transport coefficient {quantity!r}.")
name, species_index, transp_name, factor, ylabel = COEFFICIENTS[quantity]
reader = DiffusionCoefficientData(catalog=simulation.catalog)
times = _catalog_time_by_gstep(simulation.catalog)
indices = []
for index in reader.list_step_indices(name):
step = reader.get_step_info(index)["gstep"]
if step is None:
continue
if gstep_range is not None and not (gstep_range[0] <= step <= gstep_range[1]):
continue
if time_range_ms is not None:
if step not in times or not (time_range_ms[0] <= times[step]*1.0e3 <= time_range_ms[1]):
continue
indices.append(index)
if gstep_range is None and time_range_ms is None:
indices = indices[-max_default_samples:]
if not indices:
raise ValueError("The selected range contains no diffusion-coefficient BP steps.")
plane = simulation.mesh.get_plane(0)
psi_ref = np.asarray(reader.get_item("psi", indices[0]), dtype=float).reshape(-1)
curves = []
for index in indices:
psi = np.asarray(reader.get_item("psi", index), dtype=float).reshape(-1)
if psi.shape != psi_ref.shape or not np.allclose(psi, psi_ref, rtol=0, atol=1.0e-10):
raise ValueError("Diffusion-coefficient psi grids differ across selected BP steps.")
values = reader.get_coefficient(name, index)
if species_index >= values.shape[0]:
raise ValueError(f"{name} lacks the requested species index {species_index}.")
curves.append(np.asarray(values[species_index], dtype=float))
psi_n = psi_ref/float(plane.x_psi)
psi_mid, r_out, _ = midplane_radius_maps(plane)
r_sep = interp_radius_at_psi(psi_mid, r_out, 1.0, "outer")
r_profile = np.interp(psi_n, psi_mid[np.isfinite(r_out)], r_out[np.isfinite(r_out)])
r_over_a = (r_profile - plane.axis_r)/(r_sep - plane.axis_r)
reference = None
reference_path = None
if transp_path is not None and transp_name is not None:
reference_path = str(Path(transp_path).expanduser().resolve())
with h5py.File(reference_path) as file:
if "PSIN" not in file or transp_name not in file:
raise KeyError(f"TRANSP file needs PSIN and {transp_name} datasets.")
tr_psi = np.asarray(file["PSIN"], dtype=float).reshape(-1)
tr_coeff = np.asarray(file[transp_name], dtype=float).reshape(-1)
if tr_psi.shape != tr_coeff.shape:
raise ValueError(f"TRANSP PSIN and {transp_name} lengths differ.")
order = np.argsort(tr_psi)
reference = np.interp(psi_n, tr_psi[order], tr_coeff[order]*factor, left=np.nan, right=np.nan)
reference[psi_n > 1.0] = np.nan
steps = [int(reader.get_step_info(index)["gstep"]) for index in indices]
selected_times = [times[step] for step in steps if step in times]
return TransportCoefficientComparison(
quantity=quantity,
psi_n=psi_n,
r_over_a=r_over_a,
xgc=np.mean(curves, axis=0),
transp=reference,
ylabel=ylabel,
gstep_range=(min(steps), max(steps)),
time_range_s=(min(selected_times), max(selected_times)) if selected_times else None,
transp_path=reference_path,
)