"""XGC particle and energy balances with optional TRANSP source profiles."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import h5py
import numpy as np
from xgc_analysis.analysis.transport.neoclassical import compute_neoclassical_transport
from .diagnostics import available_source_terms, compute_source_rates
_BALANCES = {
"particle": ("fueling", "e", "ptl_flux_e", ("SCEW", "SCEV", "SBE"), "particles/s"),
"ion_energy": ("heating", "i", "enflux_i_gc", ("PBI",), "MW"),
"electron_energy": ("heating", "e", "enflux_e", ("PBE", "POH"), "MW"),
}
_PHYSICAL_SOURCES = ("neutral", "heat_torque", "collision", "particle_source")
[docs]
@dataclass
class BalanceComparison:
"""Transport and physical-source rates enclosed by mapped flux surfaces.
Attributes
----------
quantity, units : str
Balance channel and units of all curves.
psi_n, xgc_transport, xgc_sources, transp_sources : ndarray or None
Surface coordinate and enclosed XGC/optional TRANSP rates.
source_terms : tuple[str, ...]
XGC physical-source components included in the sum.
gstep_range, time_range_s : tuple
Averaged gsteps and physics-time range in seconds.
transp_path : str or None
Reference HDF5 path when used.
"""
quantity: str
psi_n: np.ndarray
xgc_transport: np.ndarray
xgc_sources: np.ndarray
transp_sources: np.ndarray | None
units: str
source_terms: tuple[str, ...]
gstep_range: tuple[int, int]
time_range_s: tuple[float, float]
transp_path: str | None
[docs]
@dataclass
class ParticleFuelingComparison:
"""Separate enclosed electron fueling rates from XGC and TRANSP.
Attributes
----------
psi_n : ndarray
Surface coordinate shared by all curves.
xgc_neutral, xgc_beam, transp_neutral, transp_beam : ndarray or None
Cumulative particle rates in particles/s; TRANSP curves are optional.
gstep_range, time_range_s : tuple
Averaged gsteps and physics-time range in seconds.
transp_path : str or None
Reference HDF5 path when used.
"""
psi_n: np.ndarray
xgc_neutral: np.ndarray
xgc_beam: np.ndarray
transp_neutral: np.ndarray | None
transp_beam: np.ndarray | None
gstep_range: tuple[int, int]
time_range_s: tuple[float, float]
transp_path: str | None
[docs]
def read_transp_particle_fueling(path: str | Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Integrate TRANSP wall-plus-volume neutrals and beam over DVOL.
Parameters
----------
path
TRANSP HDF5 input path.
Returns
-------
tuple[np.ndarray, np.ndarray, np.ndarray]
psi, cumulative neutral, and beam rates.
"""
names = ("PSIN", "DVOL", "SCEW", "SCEV", "SBE")
with h5py.File(Path(path).expanduser(), "r") as file:
missing = [name for name in names if name not in file]
if missing:
raise KeyError(f"TRANSP file lacks {', '.join(missing)}.")
arrays = {name: np.asarray(file[name], dtype=float).reshape(-1) for name in names}
psi = arrays["PSIN"]
if any(values.shape != psi.shape for values in arrays.values()):
raise ValueError("TRANSP source arrays must all have the same length.")
if not np.isfinite(psi).all() or not np.all(np.diff(psi) > 0):
raise ValueError("TRANSP PSIN must increase monotonically.")
return (psi, np.cumsum((arrays["SCEW"] + arrays["SCEV"])*arrays["DVOL"]),
np.cumsum(arrays["SBE"]*arrays["DVOL"]))
[docs]
def compare_particle_fueling_sources(
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,
source_steps_per_diagnostic: int | None = None,
) -> ParticleFuelingComparison:
"""Compare separate XGC neutral and prescribed beam fueling with TRANSP.
XGC ``particle_source`` represents the prescribed beam source for runs
configured with a beam source profile. All curves are cumulative rates
inside the flux surface, in particles/s.
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.
source_steps_per_diagnostic
Optional override of source updates per diagnostic frame.
Returns
-------
ParticleFuelingComparison
separate XGC and optional TRANSP curves.
"""
results = {
source: compute_source_rates(
simulation, channel="fueling", source=source, species="e",
gstep_range=gstep_range, time_range_ms=time_range_ms,
max_default_samples=max_default_samples,
source_steps_per_diagnostic=source_steps_per_diagnostic,
)
for source in ("neutral", "particle_source")
}
neutral = results["neutral"]
beam = results["particle_source"]
if not np.array_equal(neutral.gsteps, beam.gsteps) or not np.array_equal(neutral.psi_n, beam.psi_n):
raise ValueError("XGC neutral and particle-source diagnostics use different steps or surfaces.")
tr_neutral = tr_beam = None
resolved_path = None
if transp_path is not None:
resolved_path = str(Path(transp_path).expanduser().resolve())
tr_psi, tr_neutral_raw, tr_beam_raw = read_transp_particle_fueling(resolved_path)
tr_neutral = np.interp(neutral.psi_n, tr_psi, tr_neutral_raw, left=np.nan, right=np.nan)
tr_beam = np.interp(neutral.psi_n, tr_psi, tr_beam_raw, left=np.nan, right=np.nan)
tr_neutral[neutral.psi_n > 1.0] = np.nan
tr_beam[neutral.psi_n > 1.0] = np.nan
return ParticleFuelingComparison(
psi_n=neutral.psi_n, xgc_neutral=neutral.cumulative,
xgc_beam=beam.cumulative, transp_neutral=tr_neutral,
transp_beam=tr_beam,
gstep_range=(int(neutral.gsteps[0]), int(neutral.gsteps[-1])),
time_range_s=(float(neutral.time_s[0]), float(neutral.time_s[-1])),
transp_path=resolved_path,
)
[docs]
def read_transp_sources(path: str | Path, quantity: str) -> tuple[np.ndarray, np.ndarray]:
"""Integrate TRANSP source densities over ``DVOL`` shell volumes.
``SCEW``, ``SCEV``, and ``SBE`` use particles/cm^3/s; ``PBI``, ``PBE``, and ``POH``
use W/cm^3; ``DVOL`` uses cm^3. Energy output is converted to MW.
Parameters
----------
path
TRANSP HDF5 input path.
quantity
Coefficient or particle/energy balance selector.
Returns
-------
tuple[np.ndarray, np.ndarray]
psi and cumulative source rates.
"""
if quantity not in _BALANCES:
raise ValueError(f"Unknown balance quantity {quantity!r}.")
names = ("PSIN", "DVOL", *_BALANCES[quantity][3])
with h5py.File(Path(path).expanduser(), "r") as file:
missing = [name for name in names if name not in file]
if missing:
raise KeyError(f"TRANSP file lacks {', '.join(missing)}.")
arrays = {name: np.asarray(file[name], dtype=float).reshape(-1) for name in names}
psi = arrays["PSIN"]
if any(values.shape != psi.shape for values in arrays.values()):
raise ValueError("TRANSP source arrays must all have the same length.")
if not np.isfinite(psi).all() or not np.all(np.diff(psi) > 0):
raise ValueError("TRANSP PSIN must increase monotonically.")
density = sum(arrays[name] for name in _BALANCES[quantity][3])
cumulative = np.cumsum(density*arrays["DVOL"])
if quantity != "particle":
cumulative *= 1.0e-6
return psi, cumulative
[docs]
def compare_particle_energy_balance(
simulation,
quantity: str,
*,
transp_path: str | Path | None = None,
gstep_range: tuple[int, int] | None = None,
time_range_ms: tuple[float, float] | None = None,
smoothing_sigma: float = 2.0,
max_default_samples: int = 30,
moment_product: str = "xgc.f2d.bp",
source_steps_per_diagnostic: int | None = None,
) -> BalanceComparison:
"""Compare XGC neo-plus-diffusion transport with physical sources.
Transport follows the IDL balance: negative cumulative diffusion plus
measured nonadiabatic radial flux times flux-surface area. The XGC source
curve sums available neutral, external heat/torque, collision, and
particle-source terms. The optional TRANSP source curve sums wall-neutral,
volume-neutral, and beam particle sources, ion beam power, or electron beam
and ohmic power.
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.
smoothing_sigma
Gaussian width in flux-surface indices.
max_default_samples
Number of latest samples used when no window is supplied.
moment_product
F2D or axisymmetric F3D moment product.
source_steps_per_diagnostic
Optional override of source updates per diagnostic frame.
Returns
-------
BalanceComparison
XGC transport and source curves with optional TRANSP sources.
"""
if quantity not in _BALANCES:
raise ValueError(f"Unknown balance quantity {quantity!r}.")
channel, species, flux_name, _, units = _BALANCES[quantity]
present = set(available_source_terms(simulation))
sources = tuple(source for source in _PHYSICAL_SOURCES if source in present)
if "diffusion" not in present or not sources:
raise ValueError("XGC balance requires diffusion and physical f-source channels.")
terms = ("diffusion", *sources)
results = {
source: compute_source_rates(
simulation, channel=channel, source=source, species=species,
gstep_range=gstep_range, time_range_ms=time_range_ms,
max_default_samples=max_default_samples,
source_steps_per_diagnostic=source_steps_per_diagnostic,
)
for source in terms
}
first = results["diffusion"]
if any(not np.array_equal(item.gsteps, first.gsteps) for item in results.values()):
raise ValueError("XGC source terms have different selected diagnostic steps.")
neo = compute_neoclassical_transport(
simulation, gstep_window=(int(first.gsteps[0]), int(first.gsteps[-1])),
smoothing_sigma=smoothing_sigma, moment_product=moment_product,
)
plane = simulation.mesh.get_plane(0)
geometry = plane.compute_core_sol_geometry_blocks()
neo_psi = np.asarray(neo.profiles["psi_pol_norm"], dtype=float)
area = np.asarray(geometry["area_core_block"], dtype=float)
if area.shape != neo_psi.shape:
raise ValueError("Flux-surface area and neoclassical profile grids differ.")
flux = np.asarray(neo.profiles[flux_name], dtype=float)
core = neo_psi < 1.0
transport = np.full(first.psi_n.shape, np.nan)
if np.count_nonzero(core) >= 2:
radial_flow = np.interp(first.psi_n, neo_psi[core], (area*flux)[core], left=np.nan, right=np.nan)
transport = radial_flow - first.cumulative
physical_sources = sum(results[source].cumulative for source in sources)
if quantity != "particle":
transport *= 1.0e-6
physical_sources *= 1.0e-6
transp = None
resolved_path = None
if transp_path is not None:
resolved_path = str(Path(transp_path).expanduser().resolve())
tr_psi, tr_curve = read_transp_sources(resolved_path, quantity)
transp = np.interp(first.psi_n, tr_psi, tr_curve, left=np.nan, right=np.nan)
transp[first.psi_n > 1.0] = np.nan
return BalanceComparison(
quantity=quantity, psi_n=first.psi_n, xgc_transport=transport,
xgc_sources=physical_sources, transp_sources=transp, units=units,
source_terms=sources,
gstep_range=(int(first.gsteps[0]), int(first.gsteps[-1])),
time_range_s=(float(first.time_s[0]), float(first.time_s[-1])),
transp_path=resolved_path,
)