import re
import sqlite3
import warnings
import zlib
from pathlib import Path
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 open_catalog
[docs]
class Simulation:
REQUIRED_CATALOG_PRODUCTS = ("xgc.mesh.bp", "xgc.equil.bp", "xgc.bfield.bp")
def __init__(
self,
directories=None,
is_stellarator=False,
sim_is_axisymmetric=False,
catalog=None,
initialize=True,
):
"""
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, and to locate a local ``input`` file
when one is available. If omitted and ``catalog`` has a
``root_dir`` attribute, that root is used as the local fallback
directory. 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``,
and ``xgc.bfield.bp``. 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.
Raises
------
RuntimeError
If ``initialize`` is True and no catalog advertises the required
static products.
Notes
-----
Static BP products are resolved through the catalog so directory and
campaign backends share the same initialization path. The text
``input`` file is read from an embedded campaign text dataset named
``input`` when available, otherwise from a local fallback directory.
"""
self.catalog = catalog
self.is_stellarator = is_stellarator
self.sim_is_axisymmetric = sim_is_axisymmetric
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.input_params = {}
self.species = []
self.mesh = None
self.magnetic_field = None
self.velocity_grid = None
return
if self.catalog is None:
self.catalog = self._open_required_catalog(directories)
self._require_catalog_products(self.catalog, self.REQUIRED_CATALOG_PRODUCTS)
self.data_directory = self._catalog_data_directory(self.catalog, directories)
# Parse the input file to get the simulation parameters
self.input_params = self._read_input_params(directories)
# Set up species list
self.species = self._initialize_species()
# 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.
self.mesh = Mesh(
is_axisymmetric=(not is_stellarator),
data_dir=self.data_directory,
catalog=self.catalog,
)
# 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.
self.magnetic_field = MagneticField(
plane_instance=self.mesh.get_plane(0),
data_dir=self.data_directory,
catalog=self.catalog,
)
# Velocity-space grid metadata (used by distribution-function readers).
# Not all analysis directories contain xgc.f0.mesh.bp, so fail softly.
try:
self.velocity_grid = VelocityGrid(work_dir=self.data_directory, catalog=self.catalog)
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,
)
@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 and to locate a local
``input`` file.
"""
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)
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 _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, directories):
"""
Read XGC namelist parameters from campaign or local ``input`` data.
Parameters
----------
directories : iterable[str]
Local directories to search when the active catalog does not expose
embedded input text.
Returns
-------
dict
Parsed namelist data. An empty dictionary is returned when no input
text is available.
"""
input_text = self._read_campaign_input_text()
if input_text is not None:
return self._parse_namelist_text(input_text)
for directory in directories:
input_file = Path(directory) / "input"
if input_file.exists():
return self._parse_namelist_file(input_file)
warnings.warn(
"No XGC input text found in the campaign catalog or local fallback "
"directories, so Simulation.input_params is empty.",
RuntimeWarning,
stacklevel=2,
)
return {}
def _read_campaign_input_text(self):
"""
Return embedded campaign text for the ``input`` dataset when available.
Returns
-------
str or None
Decoded input-file text, or ``None`` when the active catalog is not
campaign-backed or does not store an embedded ``input`` text file.
Notes
-----
HPC-Campaign stores archive metadata in SQLite. ADIOS FileReader is
used for BP products, but text products are not currently surfaced by
the catalog discovery layer, so ``Simulation`` reads this one required
text dataset directly from the campaign archive.
"""
campaign_path = getattr(self.catalog, "campaign_path", None)
if campaign_path is None:
return None
campaign_path = Path(campaign_path)
if not campaign_path.exists():
return None
try:
with sqlite3.connect(campaign_path) as connection:
row = connection.execute(
"""
select file.compression, file.data
from dataset
join replica on dataset.rowid = replica.datasetid
join repfiles on replica.rowid = repfiles.replicaid
join file on repfiles.fileid = file.rowid
where dataset.name = ? and dataset.fileformat = ?
order by replica.modtime desc, file.modtime desc
limit 1
""",
("input", "TEXT"),
).fetchone()
except sqlite3.Error:
return None
if row is None:
return None
compression, data = row
if data is None:
return None
raw = bytes(data)
if int(compression or 0) == 1:
raw = zlib.decompress(raw)
return raw.decode("utf-8")
# def _parse_namelist_file(self,filename):
# namelists = {}
# current_namelist = None
#
# with open(filename, 'r') as f:
# lines = f.readlines()
# for line in lines:
# line = line.split("!", 1)[0].strip() # Remove comment part and trim whitespace
# if not line:
# continue
# if line.startswith("&"):
# current_nml = line[1:].strip().lower()
# namelists[current_nml] = {}
# elif line.startswith("/"):
# current_nml = None
# elif current_nml:
# keyval = line.split('=')
# if len(keyval) == 2:
# key = keyval[0].strip().lower()
# val = keyval[1].strip()
# namelists[current_nml][key] = val
# return namelists
def _parse_namelist_file(self, filepath):
"""
Parse one local XGC Fortran namelist input file.
Parameters
----------
filepath : str or pathlib.Path
File to read.
Returns
-------
dict
Parsed namelist mapping.
"""
with open(filepath, 'r') as f:
return self._parse_namelist_lines(f.readlines())
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.")