"""Flux-surface and volume-integrated rates from XGC f-source diagnostics."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from xgc_analysis.data_access.products.fsource_data import FsourceData
_CHANNELS = {
"fueling": ("density_change", "s^-1"),
"heating": ("energy_change", "W"),
"torque": ("momentum_change", "N m"),
}
[docs]
@dataclass
class SourceRateResult:
"""One f-source term evaluated over the selected diagnostic window.
``local`` is the source rate in each mapped flux-surface shell;
``cumulative`` is the enclosed volume integral, following the IDL
core/SOL separatrix treatment. ``total`` is the full-domain rate at each
selected time. All rates use the unit in ``units``.
Attributes
----------
psi_n, local, cumulative : ndarray
Flux-surface coordinate, shell rates, and enclosed rates.
gsteps, time_s, total : ndarray
Selected gsteps, physics times in seconds, and total-rate trace.
units, channel, source, species : str
Rate unit and selected diagnostic identifiers.
"""
psi_n: np.ndarray
local: np.ndarray
cumulative: np.ndarray
gsteps: np.ndarray
time_s: np.ndarray
total: np.ndarray
units: str
channel: str
source: str
species: str
[docs]
def available_source_terms(simulation) -> list[str]:
"""Return f-source suffixes with matching density and change channels.
Parameters
----------
simulation
Initialized XGC Simulation containing the needed geometry and diagnostics.
Returns
-------
list[str]
source suffixes present in the diagnostic.
"""
names = simulation.catalog.get_product("xgc.fsourcediag.bp").variables
prefix = "i_density_change_"
return sorted(
name[len(prefix):] for name in names
if name.startswith(prefix) and f"i_density_{name[len(prefix):]}" in names
)
def _source_step_duration_by_gstep(
steps: np.ndarray,
times: np.ndarray,
f_source_period: int,
source_steps_per_diagnostic: int | None,
) -> dict[int, float]:
"""Recover the duration of one source update from diagnostic coordinates.
The diagnostic does not accumulate increments between outputs. XGC calls
f_source once every ``f_source_period`` global simulation steps.
"""
if f_source_period <= 0:
raise ValueError("sml_f_source_period must be positive.")
gaps = np.diff(steps)
if source_steps_per_diagnostic is None:
if np.any(gaps % f_source_period):
raise ValueError("F-source diagnostic gsteps are not separated by whole source periods.")
source_counts = gaps//f_source_period
else:
if source_steps_per_diagnostic <= 0:
raise ValueError("source_steps_per_diagnostic must be positive.")
source_counts = np.full(gaps.shape, source_steps_per_diagnostic)
dt_step = np.diff(times)/source_counts
return dict(zip(steps, np.r_[dt_step[0], dt_step]))
[docs]
def compute_source_rates(
simulation,
*,
channel: str,
source: str,
species: str,
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,
) -> SourceRateResult:
"""Evaluate local, cumulative, and total rates from ``xgc.fsourcediag``.
The BP change channel is normalized to its source-specific density. Its
product with that density is a change per XGC source step. A source step's
duration is inferred from neighboring diagnostic ``time`` and ``gstep``
coordinates and ``sml_f_source_period``, including for the first selected
sample. The optional ``source_steps_per_diagnostic`` overrides the number
of source updates between output slices, matching the IDL routine's
``fsource_rate`` argument. Positive momentum
sources are converted from parallel linear momentum to toroidal torque by
multiplying by ``R*B_phi/|B|`` before spatial integration.
Parameters
----------
simulation
Initialized XGC Simulation containing the needed geometry and diagnostics.
channel
Fueling, heating, or torque channel.
source
F-source term suffix to evaluate.
species
Ion or electron species selector.
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
-------
SourceRateResult
shell, cumulative, and time-trace rates.
"""
if channel not in _CHANNELS:
raise ValueError(f"Unknown f-source channel {channel!r}.")
if species not in {"i", "e"}:
raise ValueError("F-source species must be 'i' or 'e'.")
if source not in available_source_terms(simulation):
raise ValueError(f"F-source term {source!r} is unavailable.")
records = list(simulation.catalog.available_steps("xgc.fsourcediag.bp", "time"))
if len(records) < 2:
raise ValueError("At least two f-source time samples are needed to determine the source-step duration.")
steps = np.asarray([record.logical_step for record in records], dtype=int)
times = np.asarray([record.selected_fragment.time for record in records], dtype=float)
if not np.isfinite(times).all() or not np.all(np.diff(steps) > 0) or not np.all(np.diff(times) > 0):
raise ValueError("F-source gsteps and physics times must increase monotonically.")
f_source_period = int(simulation.input_params.get("sml_param", {}).get("sml_f_source_period", 1))
dt_by_step = _source_step_duration_by_gstep(
steps, times, f_source_period, source_steps_per_diagnostic,
)
selected = np.ones(steps.size, dtype=bool)
if gstep_range is not None:
selected &= (steps >= gstep_range[0]) & (steps <= gstep_range[1])
if time_range_ms is not None:
selected &= (times*1.0e3 >= time_range_ms[0]) & (times*1.0e3 <= time_range_ms[1])
selected_steps = steps[selected]
if gstep_range is None and time_range_ms is None:
selected_steps = selected_steps[-max_default_samples:]
if selected_steps.size == 0:
raise ValueError("The selected range contains no f-source diagnostic samples.")
change_name = f"{species}_{_CHANNELS[channel][0]}_{source}"
density_name = f"{species}_density_{source}"
reader = FsourceData(
simulation.mesh,
catalog=simulation.catalog,
steps=selected_steps,
variables=[change_name, density_name],
)
plane = simulation.mesh.get_plane(0)
surfaces = np.asarray(plane.surf_map, dtype=int)
psi_surf = np.asarray(plane.psi_surf, dtype=float)[surfaces]
psi_n = psi_surf/float(plane.x_psi)
sep = int(np.argmin(np.abs(psi_n - 1.0)))
vertex_psi = np.asarray(simulation.magnetic_field.psi_pd.data, dtype=float).reshape(-1)
z = np.asarray(plane.rz, dtype=float)[:, 1]
volume_factor = int(simulation.mesh.wedge_n)*int(simulation.mesh.nphi)
vertex_volume = np.asarray(plane.node_vol_nearest, dtype=float)*volume_factor
shell_volume = np.zeros(surfaces.size, dtype=float)
for j, surface in enumerate(surfaces):
vertices = np.asarray(plane.get_surface_vertex_indices(int(surface)), dtype=int)
if j == sep:
vertices = vertices[z[vertices] >= plane.x_z]
shell_volume[j] = np.sum(vertex_volume[vertices])
moment_arm = None
if channel == "torque":
field = np.asarray(simulation.magnetic_field.bfield.data, dtype=float)
moment_arm = np.asarray(plane.rz)[:, 0]*field[:, 2]/np.linalg.norm(field, axis=1)
local_samples = []
cumulative_samples = []
totals = []
sample_steps = []
for local_step in reader.list_step_indices(change_name):
gstep = int(reader.get_step_info(local_step)["logical_step"])
change = np.asarray(reader.get_plane_data(change_name, local_step).data, dtype=float).reshape(-1)
density = np.asarray(reader.get_plane_data(density_name, local_step).data, dtype=float).reshape(-1)
rate = np.nan_to_num(change*density, nan=0.0, posinf=0.0, neginf=0.0)/dt_by_step[gstep]
if moment_arm is not None:
rate *= moment_arm
local_samples.append(np.asarray(plane.flux_avg_to_surf(rate))[surfaces]*shell_volume)
weighted = rate*vertex_volume
core = np.asarray([
np.sum(weighted[(vertex_psi <= psi) & (z >= plane.x_z)])
for psi in psi_surf[:sep+1]
])
sol = np.asarray([
core[-1] + np.sum(weighted[(vertex_psi > plane.x_psi) & (vertex_psi <= psi)])
for psi in psi_surf[sep+1:]
])
cumulative_samples.append(np.r_[core, sol])
totals.append(float(np.sum(weighted)))
sample_steps.append(gstep)
sample_steps = np.asarray(sample_steps, dtype=int)
time_by_step = dict(zip(steps, times))
return SourceRateResult(
psi_n=psi_n,
local=np.mean(local_samples, axis=0),
cumulative=np.mean(cumulative_samples, axis=0),
gsteps=sample_steps,
time_s=np.asarray([time_by_step[step] for step in sample_steps]),
total=np.asarray(totals),
units=_CHANNELS[channel][1],
channel=channel,
source=source,
species=species,
)