import numpy as np
from pathlib import PurePosixPath
from dataclasses import dataclass, replace
from typing import Dict, List, Any
from xgc_analysis.constants import echarge, atomic_mass_unit
from xgc_analysis.species.periodic_table import get_element_name_by_mass # You’ll need to implement or import this.
from .prescribed_profiles import (
build_deposition_shapes,
build_prescribed_profile,
indexed_value,
read_input_package_profile,
read_profile_text,
)
[docs]
class Species:
"""One XGC species and its input-defined radial functions.
``initial_profiles`` retains the existing file-profile API.
``prescribed_profiles`` describes all input-defined radial functions,
including analytic choices and particle source rates. Heat/torque
deposition weights are stored separately in ``deposition_shapes``.
Mass and charge are retained in both XGC input units and SI units.
"""
def __init__(self, sim, index):
"""
Initialize one species from its input-parameter list entry.
Parameters
----------
sim : xgc_analysis.simulation.Simulation
Initialized simulation with input parameters and catalog text.
index : int
Zero-based species index in the XGC input parameter lists.
"""
ptl_param = sim.input_params.get("ptl_param", {})
# Extract parameters from Simulation input
mass_au = ptl_param["ptl_mass_au"][index] # mass in atomic units
charge_eu = ptl_param["ptl_charge_eu"][index] # charge in units of e
marker_type = ptl_param["ptl_marker_type"][index]
dynamic_f0 = ptl_param["ptl_dynamic_f0"][index]
# Derived quantities
self.mass_au = mass_au
self.index = index
self.mass_kg = mass_au * atomic_mass_unit
self.charge_eu = charge_eu
self.charge_C = charge_eu * echarge
self.name = get_element_name_by_mass(mass_au)
self.species_type = "None" if marker_type == 0 else marker_type
self.dynamic_f0 = dynamic_f0
# ----------------------------------------------------------------------
# 0. A tiny container that keeps one profile together
# ----------------------------------------------------------------------
@dataclass
class InitialProfile:
"""One initial species profile sampled against normalized poloidal flux."""
psi_n: np.ndarray
value: np.ndarray
type: int | None = None # used only for flow
# ----------------------------------------------------------------------
# 1. Scan the eq_param namelist once, then build per-species dictionaries
# ----------------------------------------------------------------------
eq_param: Dict[str, Any] = sim.input_params.get("eq_param", {})
sml_param: Dict[str, Any] = sim.input_params.get("sml_param", {})
# Every key we care about in one place
SHAPE_KEYS = [
"eq_dens_shape", "eq_temp_shape", "eq_flow_shape",
"eq_dens_diff_shape", "eq_pinch_v_shape",
"eq_flow_diff_shape", "eq_t_diff_shape",
]
FILE_KEY_TO_LABEL = {
"eq_dens_file": "density",
"eq_temp_file": "temperature",
"eq_flow_file": "flow",
"eq_dens_diff_file": "ptl_diffusivity",
"eq_pinch_v_file": "ptl_pinch_velocity",
"eq_flow_diff_file": "momentum_diffusivity",
"eq_t_diff_file": "heat_conductivity",
}
# ----------------------------------------------------------------------
# Compute the manifest-relative directory that holds profile files.
# ----------------------------------------------------------------------
# The catalog owns all text reads. The namelist still defines the
# relative path convention, but Species only uses it to build candidate
# artifact names such as "input_dir/profile.dat".
# ----------------------------------------------------------------------
input_dir_raw = sml_param.get("sml_input_file_dir", "")
if isinstance(input_dir_raw, (list, tuple)):
input_dir_raw = input_dir_raw[0] if input_dir_raw else ""
input_dir = _clean_artifact_path(str(input_dir_raw).strip().strip('"').strip("'"))
input_format = str(sml_param.get("sml_input_file_format", "text")).strip().strip('"').strip("'").lower()
package_file = str(sml_param.get("sml_input_package_file", "xgc.input.bp")).strip().strip('"').strip("'")
package_key = _clean_artifact_path(f"{input_dir}/{package_file}")
i_sp = index # for readability
self.initial_profiles: Dict[str, InitialProfile] = {}
for shape_key in SHAPE_KEYS:
if shape_key not in eq_param:
continue
shapes = eq_param[shape_key]
shapes = shapes if isinstance(shapes, (list, tuple, np.ndarray)) else [shapes]
if i_sp >= len(shapes):
continue # parameter missing or too short
if shapes[i_sp] != -1:
continue # profile not file‑based
file_key = shape_key.replace("_shape", "_file")
label = FILE_KEY_TO_LABEL[file_key]
if input_format == "adios" and label in {"density", "temperature", "flow"}:
continue
file_paths = eq_param.get(file_key, [])
file_paths = file_paths if isinstance(file_paths, (list, tuple, np.ndarray)) else [file_paths]
f_name = str(file_paths[i_sp]).strip() if i_sp < len(file_paths) else ""
if not f_name:
# "_shape == -1" but filename missing → silently skip, per spec
continue
artifact_path = _resolve_catalog_text_artifact(sim, input_dir, f_name)
psi_n, vals = read_profile_text(sim.catalog.read_text(artifact_path))
self.initial_profiles[label] = InitialProfile(
psi_n=psi_n,
value=vals,
type=int(indexed_value(eq_param, "eq_flow_type", i_sp, 2))
if label == "flow" and "eq_flow_type" in eq_param else None
)
def read_catalog_profile(filename: str) -> tuple[np.ndarray, np.ndarray]:
"""Resolve one input profile filename through the active catalog."""
artifact = _resolve_catalog_text_artifact(sim, input_dir, filename)
return read_profile_text(sim.catalog.read_text(artifact))
families = (
("eq_dens", "density", "m^-3", (0.0, 0.0, 0.0), (1.0e20, 0.0, 0.0)),
("eq_temp", "temperature", "eV", (0.0, 0.0, 0.0), (1.0e3, 0.0, 0.0)),
("eq_flow", "flow", "rad/s", (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)),
("eq_dens_diff", "ptl_diffusivity", "m^2/s", (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)),
("eq_pinch_v", "ptl_pinch_velocity", "m/s", (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)),
("eq_flow_diff", "momentum_diffusivity", "m^2/s", (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)),
("eq_t_diff", "heat_conductivity", "m^2/s", (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)),
)
self.prescribed_profiles = {}
base_flow_type = int(indexed_value(eq_param, "eq_flow_type", index, 2))
self.flow_types = {"flow": base_flow_type}
for prefix, label, unit, default_x, default_v in families:
if label == "flow" and base_flow_type == 0:
unit = "m/s"
package_group = f"/profiles/{label}/species{index}"
from_package = input_format == "adios" and label in {"density", "temperature", "flow"}
file_reader = (
lambda filename, group=package_group: read_input_package_profile(
sim.catalog, package_key, group)
) if from_package else read_catalog_profile
self.prescribed_profiles[label] = build_prescribed_profile(
params=eq_param, prefix=prefix, name=label, unit=unit,
index=index, defaults_x=default_x, defaults_v=default_v,
read_file=file_reader,
source_override=f"{package_key}:{package_group}" if from_package else None,
)
# XGC inherits f-grid/marker temperature and flow shape parameters
# from the corresponding equilibrium profiles, then applies overrides.
for prefix, label, parent_label in (
("eq_fg_temp", "fgrid_temperature", "temperature"),
("eq_fg_flow", "fgrid_flow", "flow"),
("eq_mk_temp", "marker_temperature", "temperature"),
("eq_mk_flow", "marker_flow", "flow"),
):
parent = self.prescribed_profiles[parent_label]
base_prefix = "eq_temp" if parent_label == "temperature" else "eq_flow"
base_filename = str(indexed_value(eq_param, f"{base_prefix}_file", index, "example.dat"))
from_package = input_format == "adios"
package_group = f"/profiles/{label}/species{index}"
file_reader = (
lambda filename, group=package_group: read_input_package_profile(
sim.catalog, package_key, group)
) if from_package else read_catalog_profile
inherit_package_profile = (
from_package
and parent.shape == -1
and int(indexed_value(eq_param, f"{prefix}_shape", index, parent.shape)) == -1
and not str(indexed_value(eq_param, f"{prefix}_file", index, "")).strip()
)
if inherit_package_profile:
self.prescribed_profiles[label] = replace(parent, name=label)
else:
self.prescribed_profiles[label] = build_prescribed_profile(
params=eq_param, prefix=prefix, name=label, unit=parent.unit,
index=index, inherited=parent, read_file=file_reader,
default_file=parent.source if from_package else base_filename,
source_override=f"{package_key}:{package_group}" if from_package else None,
)
if parent_label == "flow":
flow_type = int(indexed_value(eq_param, f"{prefix}_type", index, base_flow_type))
self.flow_types[label] = flow_type
self.prescribed_profiles[label] = replace(
self.prescribed_profiles[label],
unit="m/s" if flow_type == 0 else "rad/s",
)
src_param = sim.input_params.get("src_param", {})
self.prescribed_profiles["particle_source_rate"] = build_prescribed_profile(
params=src_param, prefix="src_particle_source", name="particle_source_rate",
unit="m^-3 s^-1", index=index, read_file=read_catalog_profile,
)
self.deposition_shapes = build_deposition_shapes(src_param, index)
self._profile_decay_options = {
"outpsi_n": float(sml_param.get("sml_outpsi", 1.05)),
"outer_factor": float(eq_param.get("eq_out_decay_factor", 0.1)),
"outer_width_n": float(eq_param.get("eq_out_decay_width", 0.03)),
"private_factor": float(eq_param.get("eq_priv_flux_decay_factor", 0.05)),
"private_width_n": float(eq_param.get("eq_priv_flux_decay_width", 0.015)),
}
species_count = len(ptl_param.get("ptl_mass_au", ()))
responder = int(eq_param.get("eq_den_rsp_index", 0)) if species_count > 2 else -1
self.density_is_quasineutral_responder = index == responder
[docs]
def evaluate_xgc_profile(
self, name: str, psi_n, *, region: str,
xpoint_psi_n=None, upper_xpoint_psi_n=None,
):
"""Evaluate an input profile with XGC's region-dependent decay.
Parameters
----------
name : str
Key in ``prescribed_profiles``.
psi_n : array_like
Normalized poloidal flux.
region : {'core_edge', 'private_lower', 'private_upper'}
Region already classified from the magnetic equilibrium and R/Z.
xpoint_psi_n : float or None, optional
Lower X-point flux divided by the run's psi normalization.
Required for the lower private-flux region; this need not be 1
for a 3D equilibrium.
upper_xpoint_psi_n : float or None, optional
Required for an upper private-flux region.
Returns
-------
ndarray
Profile values with XGC's configured exponential decay.
Notes
-----
For a quasineutrality responder, XGC constructs the density from the
other species; this method cannot reconstruct it from one species'
input parameters. Runtime region selection and particle-source masks
also require the caller's geometry context.
"""
if name == "density" and self.density_is_quasineutral_responder:
raise RuntimeError("quasineutrality-adjusted density requires the other species and XGC sampling")
if region == "private_lower" and xpoint_psi_n is None:
raise ValueError("lower private-flux evaluation requires xpoint_psi_n")
return self.prescribed_profiles[name].evaluate_with_region(
psi_n, region=region, xpoint_psi_n=xpoint_psi_n,
upper_xpoint_psi_n=upper_xpoint_psi_n,
**self._profile_decay_options,
)
def __repr__(self):
"""Summarize this species and its loaded profile state for debugging."""
return (f"Species(name={self.name}, type={self.species_type}, mass={self.mass_au:.2f} amu, "
f"{self.mass_kg:.2e} kg, charge={self.charge_eu} e, {self.charge_C:.2e} C, "
f"dynamic_f0={self.dynamic_f0})")
# --------------------------------------------------------------
# Vectorised interpolation of an initial profile (simplified)
# --------------------------------------------------------------
[docs]
def interpolate_profile(
self,
name: str,
psi_new: np.ndarray,
*,
fill: str | float | None = "edge",
) -> np.ndarray:
"""
Interpolate one of this species’ initial profiles onto a new ψ grid.
Parameters
----------
name : str
Profile key in ``self.initial_profiles``
(“density”, “temperature”, “flow”, …).
psi_new : ndarray
Target ψ‑grid (1‑D NumPy array).
fill : {"edge", None, float}, default "edge"
How to handle ψ outside the original range:
• "edge" – use nearest edge value (default)
• None – raise ValueError if extrapolation needed
• float – fill with that scalar
Returns
-------
ndarray
Interpolated profile values at ``psi_new``.
"""
# ---------- fetch the stored profile ----------
# ---------- return zeros if the profile is absent ----------
if name not in self.initial_profiles:
print(f"Unknown profile '{name}'. ")
print(f"Available: {list(self.initial_profiles)}")
return np.zeros_like(psi_new, dtype=float)
# ----------------------------------------------------------------
prof = self.initial_profiles[name]
x, y = prof.psi_n, prof.value
if x[0] > x[-1]: # ensure ascending order for np.interp
x, y = x[::-1], y[::-1]
# ---------- decide on fill strategy ----------
if fill == "edge":
left, right = y[0], y[-1]
elif fill is None:
if (psi_new < x[0]).any() or (psi_new > x[-1]).any():
raise ValueError(
f"psi_new extends outside [{x[0]}, {x[-1]}] and "
f"fill=None was requested."
)
left = right = np.nan # never used because no extrapolation
else: # float
left = right = float(fill)
# ---------- vectorised interpolation ----------
return np.interp(psi_new, x, y, left=left, right=right)
def _clean_artifact_path(path: str) -> str:
"""
Normalize an input-file path token to a manifest-relative POSIX path.
Parameters
----------
path : str
Raw path token from the XGC input namelist.
Returns
-------
str
Relative POSIX path without leading ``./``. Empty strings and ``"."``
normalize to ``""``.
"""
clean = path.strip().strip('"').strip("'")
if clean in {"", "."}:
return ""
clean = PurePosixPath(clean).as_posix()
while clean.startswith("./"):
clean = clean[2:]
return clean.lstrip("/")
def _resolve_catalog_text_artifact(sim, input_dir: str, filename: str) -> str:
"""
Resolve a file-based profile reference against catalog text artifacts.
Parameters
----------
sim : xgc_analysis.simulation.Simulation
Simulation object holding the active catalog.
input_dir : str
Manifest-relative directory from ``sml_input_file_dir``.
filename : str
Profile filename from ``eq_*_file``.
Returns
-------
str
Catalog text artifact path.
Raises
------
RuntimeError
If the simulation does not expose a catalog.
FileNotFoundError
If no candidate artifact is present in the catalog.
"""
catalog = getattr(sim, "catalog", None)
if catalog is None:
raise RuntimeError("File-based species profiles require Simulation.catalog text artifacts.")
name = _clean_artifact_path(filename)
candidates = []
if input_dir:
candidates.append(_clean_artifact_path(f"{input_dir}/{name}"))
candidates.append(name)
basename = PurePosixPath(name).name
if basename not in candidates:
candidates.append(basename)
for candidate in candidates:
if candidate and catalog.has_text(candidate):
return candidate
raise FileNotFoundError(
"Initial-profile file is not present as a catalog text artifact. "
f"Tried: {', '.join(candidate for candidate in candidates if candidate)}"
)
[docs]
def build_run_prescribed_profiles(sim) -> dict:
"""Build active run-wide profiles that are not owned by one species.
Parameters
----------
sim : xgc_analysis.simulation.Simulation
Simulation with parsed input and a catalog of input text artifacts.
Returns
-------
dict[str, PrescribedProfile]
Currently contains radiation ``effective_charge`` only when
``rad_use_zeff_profile`` is enabled.
"""
if not sim.input_params.get("rad_param", {}).get("rad_use_zeff_profile", False):
return {}
eq_param = sim.input_params.get("eq_param", {})
sml_param = sim.input_params.get("sml_param", {})
input_dir_raw = sml_param.get("sml_input_file_dir", "")
if isinstance(input_dir_raw, (list, tuple)):
input_dir_raw = input_dir_raw[0] if input_dir_raw else ""
input_dir = _clean_artifact_path(str(input_dir_raw))
def read_catalog_profile(filename):
"""Load a species profile array from the catalog-backed input product."""
artifact = _resolve_catalog_text_artifact(sim, input_dir, filename)
return read_profile_text(sim.catalog.read_text(artifact))
return {
"effective_charge": build_prescribed_profile(
params=eq_param, prefix="eq_zeff", name="effective_charge",
unit="dimensionless", index=0,
defaults_x=(1.0, 0.1, 0.0), defaults_v=(1.0, 1.0, 1.0),
default_file="zeff.dat", read_file=read_catalog_profile,
)
}