"""Input-defined XGC species profiles and heat/torque deposition shapes.
The analytic formulas follow ``XGC_core/cpp/plasma/profile.hpp``. File shapes
follow ``profile_setup.F90``: a not-a-knot spline is sampled on 1000 equally
spaced flux points and those samples are linearly interpolated at query points.
Values here are the prescribed radial functions before XGC's spatial-region
selection, edge/private-flux decay, and source normalization.
"""
from __future__ import annotations
from dataclasses import dataclass
from io import StringIO
from typing import Callable, Mapping
import numpy as np
from scipy.interpolate import CubicSpline
SHAPES = {-1: "file", 0: "constant", 1: "tanh", 2: "linear", 5: "shape5"}
[docs]
def indexed_value(params: Mapping, key: str, index: int, default):
"""Read one zero-based per-species namelist entry with its XGC default.
Parameters
----------
params : mapping
Parsed ``eq_param`` or ``src_param`` values.
key : str
Per-species input name.
index : int
Zero-based species index.
default : object
Value used only when the input key is absent.
Returns
-------
object
Selected value or ``default``. A scalar input belongs to index 0.
Raises
------
ValueError
If a present per-species input has no entry for this species, as in
XGC's ``NamelistReader::get``.
"""
if key not in params:
return default
value = params[key]
if isinstance(value, (list, tuple, np.ndarray)):
if index >= len(value):
raise ValueError(f"{key} has no value for species index {index}")
return value[index]
if index > 0:
raise ValueError(f"{key} has no value for species index {index}")
return value
[docs]
def read_profile_text(text: str) -> tuple[np.ndarray, np.ndarray]:
"""Read XGC's count/psi-value/-1 profile text without reordering points.
Parameters
----------
text : str
Catalog text artifact in XGC's one-dimensional profile format.
Returns
-------
tuple[ndarray, ndarray]
Normalized poloidal flux and corresponding profile values.
"""
fp = StringIO(text)
count = int(fp.readline().strip().lstrip("#").strip())
psi, value = np.loadtxt([next(fp) for _ in range(count)], unpack=True)
fp.readline() # Existing readers consume but do not validate the sentinel.
return np.atleast_1d(psi), np.atleast_1d(value)
[docs]
@dataclass(frozen=True)
class PrescribedProfile:
"""One per-species or run-wide XGC ``Eq::Profile`` input definition.
``x`` contains normalized poloidal-flux coordinates; ``v`` contains
values in ``unit``. ``evaluate_raw`` returns the shape before XGC applies
edge/private-flux decay or other geometry-dependent adjustments.
"""
name: str
shape: int
x: tuple[float, float, float, float, float]
v: tuple[float, float, float, float]
unit: str
source: str
psi_n: np.ndarray | None = None
values: np.ndarray | None = None
unavailable_reason: str | None = None
[docs]
def evaluate_raw(self, psi_n):
"""Evaluate the XGC radial formula at normalized poloidal flux.
Parameters
----------
psi_n : array_like
Normalized poloidal flux, including possible edge/SOL points.
Returns
-------
ndarray
Input-defined profile values in ``unit``; no region decay is
applied. File shapes use XGC's spline-then-linear interpolation.
Raises
------
ValueError
If the selected analytic shape has unusable parameters.
FileNotFoundError
If the selected file shape has no catalog text artifact.
"""
if self.unavailable_reason:
raise FileNotFoundError(self.unavailable_reason)
psi = np.asarray(psi_n, dtype=float)
x1, x2, x3, _, _ = self.x
v1, v2, v3, _ = self.v
if self.shape == 0:
return np.full_like(psi, v1, dtype=float)
if self.shape == 1:
if x2 == 0:
raise ValueError(f"{self.name}: tanh width x2 must be nonzero")
return 0.5*(v1 + v2) + 0.5*(v1 - v2)*np.tanh(2*(x1 - psi)/x2)
if self.shape == 2:
if x2 == 0:
raise ValueError(f"{self.name}: linear interval x2 must be nonzero")
return v1 + (v2 - v1)*(psi - x1)/x2
if self.shape == 5:
if x2 == 0 or x1 <= 0 or x3 < 0 or x3 == x1:
raise ValueError(f"{self.name}: invalid shape-5 center, width, or core coordinate")
height = 0.5*(v1 - v2)
width_inverse = 4/x2
slope = (v3 - v1)/(np.sqrt(x3) - np.sqrt(x1))/(-(height + 1.0e-99)*width_inverse*np.sqrt(x1))
psi_safe = np.where(np.abs(psi) < 1.0e-3, 1.0e-3, psi)
z = width_inverse*(x1 - np.sqrt(x1*psi_safe))
tanh_z = np.tanh(z)
return height*(tanh_z + 0.5*z*slope*(1.0 + tanh_z)) + 0.5*(v1 + v2)
if self.shape == -1:
if self.psi_n is None or self.values is None:
raise FileNotFoundError(f"{self.name}: file profile has no catalog values")
spline = CubicSpline(self.psi_n, self.values, bc_type="not-a-knot")
grid = np.linspace(self.psi_n[0], self.psi_n[-1], 1000)
samples = spline(grid)
return np.interp(psi, grid, samples, left=samples[0], right=samples[-1])
raise ValueError(f"{self.name}: unsupported XGC profile shape {self.shape}")
[docs]
def evaluate_with_region(
self,
psi_n,
*,
region: str,
outpsi_n: float = 1.05,
outer_factor: float = 0.1,
outer_width_n: float = 0.03,
xpoint_psi_n: float = 1.0,
private_factor: float = 0.05,
private_width_n: float = 0.015,
upper_xpoint_psi_n: float | None = None,
):
"""Apply XGC's profile decay for an explicitly identified region.
Parameters
----------
psi_n : array_like
Normalized poloidal flux.
region : {'core_edge', 'private_lower', 'private_upper'}
Region classification supplied by the caller. XGC determines this
from R/Z and the magnetic equilibrium; psi alone is insufficient.
outpsi_n, outer_factor, outer_width_n : float, optional
``sml_outpsi`` and ``eq_out_decay_*`` values in normalized psi.
xpoint_psi_n, private_factor, private_width_n : float, optional
Lower X-point normalized flux and ``eq_priv_flux_decay_*`` values.
upper_xpoint_psi_n : float or None, optional
Upper X-point normalized flux; required for ``private_upper``.
Returns
-------
ndarray
Profile values with the same edge/private exponential decay as
``Eq::Profile::value`` in the caller-designated region.
"""
psi = np.asarray(psi_n, dtype=float)
if region == "core_edge":
if outer_width_n <= 0:
raise ValueError("outer profile decay width must be positive")
boundary = self.evaluate_raw(outpsi_n)
decayed = boundary*((1-outer_factor)*np.exp(-np.abs(psi-outpsi_n)/outer_width_n) + outer_factor)
return np.where(psi <= outpsi_n, self.evaluate_raw(psi), decayed)
if region in {"private_lower", "private_upper"}:
if private_width_n <= 0:
raise ValueError("private-flux profile decay width must be positive")
if region == "private_upper":
if upper_xpoint_psi_n is None:
raise ValueError("upper private-flux evaluation requires upper_xpoint_psi_n")
boundary_psi = upper_xpoint_psi_n
else:
boundary_psi = xpoint_psi_n
boundary = self.evaluate_raw(boundary_psi)
return boundary*((1-private_factor)*np.exp(-np.abs(psi-boundary_psi)/private_width_n) + private_factor)
raise ValueError(f"unknown XGC region {region!r}; R/Z classification is required")
[docs]
@dataclass(frozen=True)
class DepositionShape:
"""Unnormalized heat/torque deposition weight for one source region.
The spatial weight is separate from whole-torus power (W), torque (N m),
volume normalization, and source timing. Special mode 1 includes a
``-Z`` distance factor in metres; it is not a dimensionless weight.
Runtime region and wall masks used by XGC are not available from a
one-dimensional input shape and are deliberately excluded here.
"""
region: int
pin: float
pout: float
decay_width: float
shape: int
special_mode: int
r_begin: float
z_begin: float
radius: float
[docs]
def evaluate_weight(self, psi_n, *, r=None, z=None):
"""Evaluate XGC's input-defined deposition weight.
Parameters
----------
psi_n : array_like
Normalized poloidal flux at the query points.
r, z : array_like or None, optional
Cylindrical coordinates in metres, required for special modes 2
and 1 respectively.
Returns
-------
ndarray
Spatial weighting factor before source amplitude and volume
normalization. Runtime region/wall masks remain the caller's job.
"""
psi = np.asarray(psi_n, dtype=float)
if self.special_mode == 2:
if r is None or z is None:
raise ValueError("circular deposition shape requires r and z")
return np.asarray((np.asarray(r)-self.r_begin)**2 + (np.asarray(z)-self.z_begin)**2 < self.radius**2, dtype=float)
width = self.decay_width
if self.shape == 2:
if width <= 0:
raise ValueError("Gaussian deposition width must be positive")
inner, outer = self.pin - 3.5*width, self.pin + 3.5*width
weight = np.exp(-0.5*((psi-self.pin)/width)**2)
else:
inner, outer = self.pin - width, self.pout + width
if width > 0:
weight = np.where(psi < self.pin, (psi-inner)/width,
np.where(psi < self.pout, 1.0, (outer-psi)/width))
else:
weight = np.where((psi > self.pin) & (psi < self.pout), 1.0, 0.0)
weight = np.where((psi > inner) & (psi < outer), weight, 0.0)
if self.special_mode == 1:
if z is None:
raise ValueError("divertor/SOL deposition shape requires z")
weight = weight*np.maximum(self.z_begin-np.asarray(z), 0.0)
return weight
[docs]
def build_prescribed_profile(
*,
params: Mapping,
prefix: str,
name: str,
unit: str,
index: int,
defaults_x: tuple[float, float, float] = (0.0, 0.0, 0.0),
defaults_v: tuple[float, float, float] = (0.0, 0.0, 0.0),
inherited: PrescribedProfile | None = None,
read_file: Callable[[str], tuple[np.ndarray, np.ndarray]] | None = None,
default_file: str = "example.dat",
source_override: str | None = None,
) -> PrescribedProfile:
"""Build a profile definition from XGC's per-species namelist defaults.
Parameters
----------
params : mapping
One parsed XGC namelist, usually ``eq_param`` or ``src_param``.
prefix : str
Input key stem, for example ``eq_temp``.
name : str
Public quantity name attached to ``Species.prescribed_profiles``.
unit : str
Units of the raw input-defined function.
index : int
Zero-based species index.
defaults_x, defaults_v : tuple[float, float, float]
XGC source-code defaults for the first three coordinate/value inputs.
inherited : PrescribedProfile or None
Base temperature/flow definition inherited by fg/mk profiles.
read_file : callable or None
Catalog text resolver returning normalized psi and values.
default_file : str
XGC source-code fallback filename for file shapes.
source_override : str or None
Catalog package/group name when a file shape comes from ADIOS input.
Returns
-------
PrescribedProfile
Shape definition; a missing file is reported when evaluated.
"""
parent_x = inherited.x if inherited is not None else (*defaults_x, 1.0, 0.0)
parent_v = inherited.v if inherited is not None else (*defaults_v, 0.0)
shape = int(indexed_value(params, f"{prefix}_shape", index,
inherited.shape if inherited is not None else 0))
if shape not in SHAPES:
raise ValueError(f"{prefix}: unsupported XGC shape {shape}")
x = tuple(float(indexed_value(params, f"{prefix}_x{i+1}", index, parent_x[i])) for i in range(5))
v = tuple(float(indexed_value(params, f"{prefix}_v{i+1}", index, parent_v[i])) for i in range(4))
file_value = indexed_value(params, f"{prefix}_file", index, default_file)
filename = str(file_value).strip().strip('"').strip("'")
if inherited is not None and not filename:
filename = default_file
source = (source_override or filename) if shape == -1 else f"{prefix}_shape={shape}"
psi_n = values = None
reason = None
if shape == -1:
if not filename:
reason = f"{name}: XGC file shape has no filename"
elif inherited is not None and filename == inherited.source and inherited.shape == -1:
psi_n, values = inherited.psi_n, inherited.values
reason = inherited.unavailable_reason
source = inherited.source
elif read_file is None:
reason = f"{name}: no catalog text reader was supplied"
else:
try:
psi_n, values = read_file(filename)
except FileNotFoundError as exc:
reason = str(exc)
return PrescribedProfile(name, shape, x, v, unit, source, psi_n, values, reason)
[docs]
def build_deposition_shapes(src_params: Mapping, index: int) -> dict[int, DepositionShape]:
"""Read one species' active heat/torque region weights from ``src_param``.
Parameters
----------
src_params : mapping
Parsed ``src_param`` namelist.
index : int
Zero-based species index.
Returns
-------
dict[int, DepositionShape]
One-based active region number to unnormalized shape definition.
"""
count = int(indexed_value(src_params, "src_narea", index, 1 if index == 1 else 0))
if not 0 <= count <= 4:
raise ValueError("src_narea must be between 0 and 4")
shapes = {}
for region in range(1, count+1):
suffix = str(region)
def param(stem, default):
"""Read one numeric analytic-shape parameter with its XGC default."""
return indexed_value(src_params, f"src_{stem}{suffix}", index, default)
shapes[region] = DepositionShape(
region=region,
pin=float(param("pin", 1.04 if region == 2 else 0.0)),
pout=float(param("pout", 1.05 if region == 2 else (0.01 if region == 1 else -1.0))),
decay_width=float(param("decay_width", 0.01)),
shape=int(param("ishape", 1)),
special_mode=int(param("special_mode", 0)),
r_begin=float(param("r_begin", 0.0)),
z_begin=float(param("z_begin", 0.0)),
radius=float(param("radius", 0.1)),
)
return shapes