"""Compare turbulence-model, flux-derived neoclassical, and input coefficients."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from xgc_analysis.data_access.products.diffusion_data import DiffusionCoefficientData
from .neoclassical import compute_neoclassical_transport
from .coefficient_comparison import COEFFICIENTS, _catalog_time_by_gstep
COEFFICIENT_KEYS = (
"particle_diffusivity",
"ion_heat_conductivity",
"electron_heat_conductivity",
"momentum_diffusivity",
)
NEO_PROFILE_KEYS = {
"particle_diffusivity": "D_eff",
"ion_heat_conductivity": "chi_i",
"electron_heat_conductivity": "chi_e",
}
INPUT_PROFILE_KEYS = {
"particle_diffusivity": "ptl_diffusivity",
"ion_heat_conductivity": "heat_conductivity",
"electron_heat_conductivity": "heat_conductivity",
"momentum_diffusivity": "momentum_diffusivity",
}
[docs]
@dataclass
class CombinedTransportResult:
"""Coefficient families mapped onto one normalized-psi grid.
Attributes
----------
psi_n : ndarray
Common normalized-poloidal-flux coordinate.
curves : dict[str, dict[str, ndarray]]
Selected source group and coefficient family, with values in m²/s.
gstep_range, time_range_s : tuple
Averaged diagnostic gsteps and optional physics-time range in seconds.
"""
psi_n: np.ndarray
curves: dict[str, dict[str, np.ndarray]]
gstep_range: tuple[int, int]
time_range_s: tuple[float, float] | None
def _selected_indices(reader, catalog, gstep_range, time_range_ms, max_default_samples):
"""Select coefficient snapshots using catalog physics time when requested."""
times = _catalog_time_by_gstep(catalog)
indices = []
for index in reader.list_step_indices("ptl_diffusivity"):
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.")
return indices, times
[docs]
def compute_combined_transport_coefficients(
simulation,
*,
groups: tuple[str, ...] = ("turbulence", "neoclassical", "anomalous"),
gstep_range: tuple[int, int] | None = None,
time_range_ms: tuple[float, float] | None = None,
max_default_samples: int = 30,
smoothing_sigma: float = 2.0,
denominator_smoothing_sigma: float | None = None,
moment_product: str = "xgc.f2d.bp",
) -> CombinedTransportResult:
"""Collect all requested coefficient groups for a standalone ``Simulation``.
``turbulence`` is the time-averaged diffusion-model BP output.
``anomalous`` is the configured input profile held by ``Species`` and is
not time averaged. ``neoclassical`` is inferred from F2D/F3D radial
particle and heat fluxes using the existing neoclassical backend. No
flux-derived ion momentum diffusivity is currently calculated there.
Parameters
----------
simulation
Initialized XGC Simulation containing the needed geometry and diagnostics.
groups
Subset of turbulence, neoclassical, and anomalous groups.
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.
smoothing_sigma
Gaussian width in flux-surface indices.
denominator_smoothing_sigma
Separate Gaussian width for density and temperature denominators.
moment_product
F2D or axisymmetric F3D moment product.
Returns
-------
CombinedTransportResult
selected coefficient groups on a common psi grid.
"""
selected_groups = tuple(dict.fromkeys(groups))
unknown = set(selected_groups) - {"turbulence", "neoclassical", "anomalous"}
if unknown:
raise ValueError(f"Unknown coefficient group(s): {sorted(unknown)}")
if not selected_groups:
raise ValueError("Select at least one coefficient group.")
reader = DiffusionCoefficientData(catalog=simulation.catalog)
indices, times = _selected_indices(reader, simulation.catalog, gstep_range, time_range_ms, max_default_samples)
steps = [int(reader.get_step_info(index)["gstep"]) for index in indices]
plane = simulation.mesh.get_plane(0)
psi_ref = np.asarray(reader.get_item("psi", indices[0]), dtype=float).reshape(-1)
psi_n = psi_ref/float(plane.x_psi)
curves: dict[str, dict[str, np.ndarray]] = {}
if "turbulence" in selected_groups:
snapshot = {}
for key in COEFFICIENT_KEYS:
name, species_index, _, _, _ = COEFFICIENTS[key]
samples = []
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 species index {species_index}.")
samples.append(np.asarray(values[species_index], dtype=float))
snapshot[key] = np.mean(samples, axis=0)
curves["turbulence"] = snapshot
if "anomalous" in selected_groups:
electrons = [sp for sp in simulation.species if sp.charge_eu < 0]
ions = [sp for sp in simulation.species if sp.charge_eu > 0]
if not electrons or not ions:
raise ValueError("Anomalous input coefficients require electron and ion Species.")
input_curves = {}
for key in COEFFICIENT_KEYS:
species = electrons[0] if key in {"particle_diffusivity", "electron_heat_conductivity"} else ions[0]
profile = species.initial_profiles.get(INPUT_PROFILE_KEYS[key])
if profile is None:
raise ValueError(f"{key} has no configured input profile in Species.")
input_curves[key] = np.interp(psi_n, profile.psi_n, profile.value,
left=np.nan, right=np.nan)
curves["anomalous"] = input_curves
if "neoclassical" in selected_groups:
neo = compute_neoclassical_transport(
simulation, gstep_window=(steps[0], steps[-1]),
smoothing_sigma=smoothing_sigma,
denominator_smoothing_sigma=denominator_smoothing_sigma,
moment_product=moment_product,
)
neo_psi = np.asarray(neo.profiles["psi_pol_norm"], dtype=float)
curves["neoclassical"] = {
key: np.interp(psi_n, neo_psi, np.asarray(neo.profiles[name], dtype=float),
left=np.nan, right=np.nan)
for key, name in NEO_PROFILE_KEYS.items()
}
selected_times = [times[step] for step in steps if step in times]
return CombinedTransportResult(
psi_n=psi_n, curves=curves, gstep_range=(steps[0], steps[-1]),
time_range_s=(min(selected_times), max(selected_times)) if selected_times else None,
)