import re
import warnings
from pathlib import Path
from time import perf_counter
from .mesh import Mesh
from .plane import Plane
from .magnetic_field import MagneticField
from .species import Species
from .velocity_grid import VelocityGrid
from .catalog import build_static_buffer, open_catalog
[docs]
class Simulation:
REQUIRED_CATALOG_PRODUCTS = ("xgc.mesh.bp", "xgc.equil.bp", "xgc.bfield.bp")
REQUIRED_CATALOG_TEXTS = ("input",)
STATIC_BUFFER_REQUESTS = {
"xgc.mesh.bp": (
"delta_phi",
"wedge_angle",
"n_n",
"n_t",
"nsurf",
"surf_maxlen",
"grid_nwall",
"rz",
"nd_connect_list",
"grid_wall_nodes",
"tr_area",
"node_vol",
"node_vol_nearest",
"node_vol_ff0",
"node_vol_ff1",
"theta",
"region",
"surf_len",
"surf_idx",
"rmin",
"rmaj",
"epsilon",
"m_max_surf",
"psi_surf",
"qsafety",
"trapped",
),
"xgc.equil.bp": (
"bp_sign",
"bt_sign",
"eq_I",
"eq_axis_b",
"eq_axis_r",
"eq_axis_z",
"eq_max_r",
"eq_max_z",
"eq_min_r",
"eq_min_z",
"eq_mpsi",
"eq_mr",
"eq_mz",
"eq_psi_grid",
"eq_psi_rz",
"eq_x_psi",
"eq_x_r",
"eq_x_z",
),
"xgc.bfield.bp": (
"bfield",
"n_n",
"psi",
"jpar_bg",
"jpar_bg_fs_avg",
),
"xgc.current_drive.bp": ("jpar_bg", "jpar_bg_fs_avg"),
"xgc.volumes.bp": ("diag_1d_vol",),
"xgc.grad_rz.bp": (
"basis",
"n_r",
"m_r",
"w_r",
"nelement_r",
"eindex_r",
"value_r",
"n_z",
"m_z",
"w_z",
"nelement_z",
"eindex_z",
"value_z",
),
"xgc.cnv_to_surf.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
),
"xgc.cnv_from_surf.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
),
"xgc.ff_1dp_fwd.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
"dl_par",
),
"xgc.ff_1dp_rev.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
"dl_par",
),
"xgc.ff_hdp_fwd.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
"dl_par",
),
"xgc.ff_hdp_rev.bp": (
"ncols",
"nrows",
"width",
"nelement",
"eindex",
"value",
"dl_par",
),
"xgc.f0.mesh.bp": (
"f0_nmu",
"f0_nvp",
"f0_smu_max",
"f0_vp_max",
"f0_dsmu",
"f0_dvp",
"f0_fg_T_ev",
"f0_T_ev",
"f0_den",
"f0_flow",
"gradpsi",
"nb_curl_nb",
"v_curv",
"v_gradb",
"f0_grid_vol_vonly",
),
}
OPTIONAL_STATIC_BUFFER_PRODUCTS = {
"xgc.current_drive.bp",
"xgc.grad_rz.bp",
"xgc.cnv_to_surf.bp",
"xgc.cnv_from_surf.bp",
"xgc.ff_1dp_fwd.bp",
"xgc.ff_1dp_rev.bp",
"xgc.ff_hdp_fwd.bp",
"xgc.ff_hdp_rev.bp",
"xgc.f0.mesh.bp",
}
def __init__(
self,
directories=None,
is_stellarator=False,
sim_is_axisymmetric=False,
catalog=None,
initialize=True,
static_buffer=None,
):
"""
Initializes a Simulation instance for the gyrokinetic code XGC.
Parameters
----------
directories : str or list[str] or None, optional
Directory or directories used to build a directory-backed catalog
when ``catalog`` is omitted. If omitted and ``catalog`` has a
``root_dir`` attribute, that root is used as the directory anchor.
Otherwise the current directory is used.
is_stellarator : bool, optional
True when the simulation is for a stellarator. The default is
False, corresponding to a tokamak setup.
sim_is_axisymmetric : bool, optional
True when the simulation itself is axisymmetric.
catalog : object or None, optional
Pre-built dataset catalog. Initialization requires a catalog that
advertises the static products ``xgc.mesh.bp``, ``xgc.equil.bp``,
``xgc.bfield.bp``, and a text artifact named ``input``. If
omitted, directory-backed catalogs are tried for the candidate
directories.
initialize : bool, optional
If True, construct mesh, magnetic-field, velocity-grid, input, and
species state as before. If False, only directory/catalog state is
stored and heavy analysis members are set to ``None`` or empty
containers.
static_buffer : mapping[str, mapping[str, object]] or None, optional
Optional reusable static-product buffer keyed by product and
variable name. When omitted, the constructor builds a buffer for
the core Simulation products once and shares it across downstream
analysis objects.
Raises
------
RuntimeError
If ``initialize`` is True and no catalog advertises the required
static products.
Notes
-----
Static BP products and required text inputs are resolved through the
catalog so directory and campaign backends share the same
initialization path. Direct local file fallback reads are disabled.
"""
self.catalog = catalog
self.is_stellarator = is_stellarator
self.sim_is_axisymmetric = sim_is_axisymmetric
self.construction_timing = {}
directories = self._normalize_directories(directories, catalog)
if not initialize:
if catalog is not None and hasattr(catalog, "root_dir"):
self.data_directory = str(catalog.root_dir)
else:
self.data_directory = str(directories[0])
self.static_buffer = static_buffer
self.input_params = {}
self.species = []
self.mesh = None
self.magnetic_field = None
self.velocity_grid = None
return
total_start = perf_counter()
catalog_open_seconds = 0.0
if self.catalog is None:
open_start = perf_counter()
self.catalog = self._open_required_catalog(directories)
catalog_open_seconds = perf_counter() - open_start
validation_start = perf_counter()
self._require_catalog_products(self.catalog, self.REQUIRED_CATALOG_PRODUCTS)
self._require_catalog_texts(self.catalog, self.REQUIRED_CATALOG_TEXTS)
self.data_directory = self._catalog_data_directory(self.catalog, directories)
validation_seconds = perf_counter() - validation_start
# Parse the input file to get the simulation parameters
input_start = perf_counter()
self.input_params = self._read_input_params()
input_seconds = perf_counter() - input_start
# Set up species list
species_start = perf_counter()
self.species = self._initialize_species()
species_seconds = perf_counter() - species_start
static_buffer_timing = {
"plan_s": 0.0,
"read_s": 0.0,
"materialize_s": 0.0,
"total_s": 0.0,
"planned_variable_reads": 0,
"record_count": 0,
"product_count": len(static_buffer or {}),
"source": "provided" if static_buffer is not None else "built",
}
if static_buffer is not None:
self.static_buffer = static_buffer
else:
buffer_start = perf_counter()
self.static_buffer, static_buffer_timing = build_static_buffer(
self.catalog,
self.STATIC_BUFFER_REQUESTS,
optional_products=self.OPTIONAL_STATIC_BUFFER_PRODUCTS,
return_timing=True,
)
static_buffer_timing["total_s"] = perf_counter() - buffer_start
# Create the Mesh and MagneticField instances.
# It is assumed that the Mesh class is defined in mesh.py and accepts a filename and a flag for axisymmetry.
mesh_start = perf_counter()
self.mesh = Mesh(
is_axisymmetric=(not is_stellarator),
data_dir=self.data_directory,
catalog=self.catalog,
static_buffer=self.static_buffer,
)
mesh_seconds = perf_counter() - mesh_start
# MagneticField is assumed to be defined in magnetic_field.py.
# Its constructor accepts a plane instance (e.g., the first plane from the mesh) and the equil and bfield file paths.
magnetic_field_start = perf_counter()
self.magnetic_field = MagneticField(
plane_instance=self.mesh.get_plane(0),
data_dir=self.data_directory,
catalog=self.catalog,
static_buffer=self.static_buffer,
)
magnetic_field_seconds = perf_counter() - magnetic_field_start
# Velocity-space grid metadata (used by distribution-function readers).
# Not all analysis directories contain xgc.f0.mesh.bp, so fail softly.
velocity_grid_start = perf_counter()
try:
self.velocity_grid = VelocityGrid(
work_dir=self.data_directory,
catalog=self.catalog,
static_buffer=self.static_buffer,
)
except Exception as exc:
self.velocity_grid = None
warnings.warn(
f"Could not initialize VelocityGrid from '{self.data_directory}/xgc.f0.mesh.bp': {exc}",
RuntimeWarning,
stacklevel=2,
)
velocity_grid_seconds = perf_counter() - velocity_grid_start
object_construction_seconds = (
mesh_seconds + magnetic_field_seconds + velocity_grid_seconds
)
self.construction_timing = {
"total_s": perf_counter() - total_start,
"catalog_open_s": catalog_open_seconds,
"catalog_validation_s": validation_seconds,
"input_parse_s": input_seconds,
"species_init_s": species_seconds,
"object_construction_s": object_construction_seconds,
"mesh_construct_s": mesh_seconds,
"magnetic_field_construct_s": magnetic_field_seconds,
"velocity_grid_construct_s": velocity_grid_seconds,
"static_buffer": static_buffer_timing,
}
@staticmethod
def _normalize_directories(directories, catalog):
"""
Normalize constructor directory input to a list of strings.
Parameters
----------
directories : str, pathlib.Path, iterable, or None
User-provided local directory candidates.
catalog : object or None
Optional catalog whose ``root_dir`` supplies the fallback directory
when ``directories`` is omitted.
Returns
-------
list[str]
Candidate local directories. These directories are only used to
build a directory catalog when needed.
"""
if directories is None:
if catalog is not None and hasattr(catalog, "root_dir"):
directories = [str(catalog.root_dir)]
else:
directories = ["./"]
if isinstance(directories, (str, Path)):
return [str(directories)]
return [str(directory) for directory in directories]
@classmethod
def _open_required_catalog(cls, directories):
"""
Build a directory-backed catalog from the first complete candidate.
Parameters
----------
directories : iterable[str]
Candidate local directories.
Returns
-------
xgc_analysis.catalog.SimulationCatalog
Directory-backed catalog advertising the required static products.
Raises
------
RuntimeError
If no candidate directory produces a complete catalog.
"""
errors = []
for directory in directories:
try:
catalog = open_catalog(directory)
cls._require_catalog_products(catalog, cls.REQUIRED_CATALOG_PRODUCTS)
cls._require_catalog_texts(catalog, cls.REQUIRED_CATALOG_TEXTS)
return catalog
except Exception as exc:
errors.append(f"{directory}: {exc}")
details = "; ".join(errors) if errors else "no directories were provided"
raise RuntimeError(
"Simulation requires a catalog with static products "
f"{', '.join(cls.REQUIRED_CATALOG_PRODUCTS)}. Could not build one from "
f"the candidate directories: {details}"
)
@staticmethod
def _require_catalog_products(catalog, product_keys):
"""
Verify that the catalog advertises required products.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog to validate.
product_keys : iterable[str]
Product keys required for simulation initialization.
Raises
------
RuntimeError
If the catalog is missing one or more products.
"""
products = getattr(catalog, "products", {})
missing = [key for key in product_keys if key not in products]
if missing:
raise RuntimeError(
"Simulation catalog is missing required product(s): "
+ ", ".join(missing)
)
@staticmethod
def _require_catalog_texts(catalog, artifact_paths):
"""
Verify that the catalog advertises required text artifacts.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog to validate.
artifact_paths : iterable[str]
Manifest-relative text artifact paths required for simulation
initialization.
Raises
------
RuntimeError
If the catalog does not expose one or more required text artifacts.
"""
missing = []
for artifact_path in artifact_paths:
has_text = getattr(catalog, "has_text", None)
if has_text is None or not has_text(artifact_path):
missing.append(artifact_path)
if missing:
raise RuntimeError(
"Simulation catalog is missing required text artifact(s): "
+ ", ".join(missing)
)
@staticmethod
def _catalog_data_directory(catalog, directories):
"""
Return a string directory anchor for compatibility fields.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Active catalog.
directories : list[str]
Local fallback directories.
Returns
-------
str
Catalog root directory when available, otherwise the first local
candidate.
"""
if catalog is not None and hasattr(catalog, "root_dir"):
return str(catalog.root_dir)
return str(directories[0])
def _read_input_params(self):
"""
Read XGC namelist parameters from catalog text artifact ``input``.
Returns
-------
dict
Parsed namelist data.
"""
return self._parse_namelist_text(self.catalog.read_text("input"))
def _parse_namelist_file(self, filepath):
"""
Disabled local namelist-file parser.
Parameters
----------
filepath : str or pathlib.Path
Local file that would have been read by the legacy path.
Returns
-------
dict
Never returned.
Notes
-----
``Simulation`` initialization reads input through
:meth:`SimulationCatalog.read_text`. Call
:meth:`_parse_namelist_text` with catalog-provided text instead.
"""
raise RuntimeError(
f"Direct namelist file reads are disabled for '{filepath}'. "
"Read text through SimulationCatalog.read_text('input') instead."
)
def _parse_namelist_text(self, text):
"""
Parse XGC Fortran namelist text.
Parameters
----------
text : str
Full input-file content.
Returns
-------
dict
Parsed namelist mapping.
"""
return self._parse_namelist_lines(text.splitlines())
def _parse_namelist_lines(self, lines):
"""
Parse XGC Fortran namelist lines.
Parameters
----------
lines : iterable[str]
Lines from a local or campaign-stored ``input`` file.
Returns
-------
dict
Parsed namelist mapping with lowercase namelist and variable names.
"""
namelists = {}
current_nml = None
for line in lines:
# Strip comments
line = line.split("!", 1)[0].strip()
if not line:
continue
# Start or end of namelist
if line.startswith("&"):
current_nml = line[1:].strip().lower()
namelists[current_nml] = {}
elif line.startswith("/"):
current_nml = None
elif current_nml:
if "=" in line:
key, val_str = line.split("=", 1)
key = key.strip().lower()
val_str = val_str.strip()
# Split by space (unless quoted string)
raw_values = re.findall(r"'[^']*'|[^ ]+", val_str)
# Convert values
values = [self.convert_fortran_value(v) for v in raw_values]
# Collapse to scalar if length 1
namelists[current_nml][key] = values if len(values) > 1 else values[0]
return namelists
[docs]
def convert_fortran_value(self,val):
val = val.strip()
# Boolean
if val.lower() in [".true.", "true", "t"]:
return True
if val.lower() in [".false.", "false", "f"]:
return False
# String
if val.startswith("'") and val.endswith("'"):
return val.strip("'")
# Fortran float with D exponent
try:
return float(val.replace('D', 'E'))
except ValueError:
pass
# Default fallback
return val
def _initialize_species(self):
"""
Initialize Species objects based on input parameters.
"""
ptl_param = self.input_params.get("ptl_param", {})
n_species = len(ptl_param.get("ptl_mass_au", []))
return [Species(self, i) for i in range(n_species)]
# ------------------------------------------------------------------------------
# Example usage:
# ------------------------------------------------------------------------------
if __name__ == "__main__":
# Create a Simulation instance searching in the current directory (or a list of directories)
sim = Simulation(directories=["./"], is_stellarator=False, sim_is_axisymmetric=True)
print("Simulation data loaded from directory:", sim.data_directory)
print("Mesh and MagneticField instances have been set up.")