diff --git a/pyproject.toml b/pyproject.toml index d4e46ea..28fd39b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ dependencies = [ # sys.exit(main()) so every successful run exited 1) — that broke # agent error-detection on every csvtool call. "csvkit>=2.2.0", - "h5py>=3.16.0", + "h5py>=3.16.0" ] [project.scripts] @@ -70,6 +70,9 @@ dev = [ "ruff>=0.5.0", ] +[project.optional-dependencies] +fix-certificates = ["truststore>=0.10.4"] + [build-system] requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" diff --git a/src/dsagt/__init__.py b/src/dsagt/__init__.py index 5ad0529..6d1b160 100644 --- a/src/dsagt/__init__.py +++ b/src/dsagt/__init__.py @@ -22,6 +22,20 @@ # sentence-transformers' tokenizer is used after a fork (e.g. under # pytest-xdist or DataLoader workers). _os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") +# uv bundles its own OpenSSL that does not read the macOS Security keychain. +# On networks with an SSL-intercepting proxy (e.g. Zscaler), the corporate +# root CA lives only in the keychain, so HTTPS downloads (HuggingFace model +# weights, arXiv PDFs) fail with "unable to get local issuer certificate". +# truststore patches ssl.SSLContext to use the OS-native trust store +# (macOS Security framework / Windows Certificate Store), making the +# installed corporate CA visible to httpx, requests, and urllib3. +# inject_into_ssl() must run before any ssl.SSLContext is constructed. +try: + import truststore as _truststore + _truststore.inject_into_ssl() + del _truststore +except ImportError: + pass del _os, _default_threads from dsagt.registry import ToolRegistry diff --git a/use_cases/tokamak_stability/AGENTS.md b/use_cases/tokamak_stability/AGENTS.md new file mode 100644 index 0000000..1d44a78 --- /dev/null +++ b/use_cases/tokamak_stability/AGENTS.md @@ -0,0 +1,8 @@ +# Agent Guide - tokamak_stability + + +## Tools + +This subdirectory provides python modules containing functions for interacting with HDF5 data files created by the M3D-C1 finite-element simulation code. Before using these functions, and in particular before creating CLI tools based on these functions, read the agent skill file at `skills/m3dc1-skill/SKILL.md` for calling conventions, input/output formats, and examples. + +The three python modules containing useful functions are: `m3dc1_tools.py`, `m3dc1_plots.py`, and `hdf5.py`. diff --git a/use_cases/tokamak_stability/README.md b/use_cases/tokamak_stability/README.md new file mode 100644 index 0000000..e8c5c00 --- /dev/null +++ b/use_cases/tokamak_stability/README.md @@ -0,0 +1,117 @@ +# Fusion energy use case + +Here we're going to use dsagt to investigate the stability properties of a tokamak configuration. We'll be looking at linear MHD simulation data produced by the [M3D-C1](https://sites.google.com/pppl.gov/m3d-c1) unstructured-mesh finite-element code. The session below has been tested with Claude Code using Sonnet 4.6. + + +## Dependencies + +In addition a standard dsagt installation you'll also need to build and install the fusion-io library from [https://github.com/nferraro/fusion-io](https://github.com/nferraro/fusion-io). The top commit of the main branch will work. Use these environment variables to point to the fusion-io installation: + +``` +export FIO_INSTALL_DIR=/path/to/your/fusion-io/install/ +export PYTHONPATH=$FIO_INSTALL_DIR/lib:$PYTHONPATH +export DYLD_LIBRARY_PATH=$FIO_INSTALL_DIR/lib:$DYLD_LIBRARY_PATH +``` + +This use case also comes with a wrapper python module called `m3dc1`, which you may want to add to your PYTHONPATH: + +``` +export PYTHONPATH=/path/to/dsagt/use_cases/tokamak_stability:$PYTHONPATH +``` + +## Example data + +You can download a [demonstration dataset](https://drive.google.com/file/d/1ZghND-G2SInuovLqrg-DVPyECEASPTSq/view?usp=sharing), consisting of a single M3D-C1 simulation output. +This dataset is courtesy of Alvaro Sanchez-Villar (asvillar@pppl.gov). Untar the .tar.gz file somewhere convenient. + + +## Getting started + + +With dsagt installed, activate the virtual environment from the repository root: + +``` +source .venv/bin/activate +``` + +Now create a new dsagt project and associated directory, here called `fusion-use-case`: + +``` +dsagt init fusion-use-case --agent claude --location ~/data/ +``` + +Here we're using Claude Code; see the dsagt [documentation](https://ai-modcon.github.io/dsagt/) for guidelines for using other agents. The created project directory will be a subdirectory of `~/data/`; omitting this with place your project directories in `~/dsagt-projects/`. + +We can start MLflow in the background: + +``` +dsagt mlflow fusion-use-case +``` + +This isn't strictly necessary but will allow us to more accurately recreate our work in a shell script later on. Copy and paste the MLflow export block printed by this command into this current shell. For example the block may be: + +``` +export CLAUDE_CODE_ENABLE_TELEMETRY=1 +export CLAUDE_CODE_ENHANCED_TELEMETRY_BETA=1 +export OTEL_LOG_TOOL_DETAILS=1 +export OTEL_LOG_USER_PROMPTS=1 +export OTEL_TRACES_EXPORTER=otlp +export OTEL_LOG_RAW_API_BODIES=file:/path/to/data/fusion-use-case/api_bodies +``` + +Now enter the project directory and start the agent: + +``` +cd ~/data/fusion-use-case +claude +``` + + +## An example session + +Inside the agent, enter these prompts one at a time, replacing the placeholder directory paths with the corresponding paths on your system: + +- There are three python files in the `/path/to/your/dsagt/use_cases/tokamak_stability` directory containing functions for dealing with HDF5 datasets produced by the M3D-C1 code: `hdf5.py`, `m3dc1_tools.py`, and `m3dc1_plots.py`. Register these functions as tools and print a summary here. Read the `AGENTS.md` file in that directory first. + +Creating and registering the tools may take several minutes. + +- Using your tools, tell me about the data in the `/path/to/your/m3dc1_data` directory. + +- What are the Miller parameters for this configuration? + +- What's the safety factor? + +- Make plots of the t=1 fields of the electron temperature, all components of the current density, and the perturbations of the density and magnetic flux. + +Plots will go in the `/path/to/fusion-use-case/plots/` subdirectory of your project directory by default. Instruct the agent if you prefer an alternative location. + +- Create plots of the standard poloidal spectra and the kinetic energy trace. + +- Save the poloidal spectral data for the pressure field in an HDF5 file `pressure_spectrum.h5`. + +New data products will go to a `processed_data/` subdirectory of your project directory by default. + +- Extract the electron temperature and electron density data at t=1 and place them in an HDF5 file `electrons.h5`. + +- Using the MLflow traces in the `mlflow/` directory, create a shell script that recreates this session's tool calls on a general data directory that is set at the top of the script. Save as `dsagt_session_script.sh`. + + +Now exit the agent session as usual and stop the MLflow server: + +``` +dsagt stop fusion-use-case +``` + +The project name (here "fusion-use-case") is registered in `~/dsagt-projects/projects.yaml` (default location). You can unregister the project name with + +``` +dsagt rm fusion-use-case +``` + +You will be asked whether you want to delete the project directory. + + + + + + diff --git a/use_cases/tokamak_stability/hdf5.py b/use_cases/tokamak_stability/hdf5.py index 3955b9a..6f55d0e 100644 --- a/use_cases/tokamak_stability/hdf5.py +++ b/use_cases/tokamak_stability/hdf5.py @@ -3,6 +3,10 @@ Functions: - ``list_h5_files(directory, recursive=False)``: return list of .h5 files in directory. - ``list_h5_variables(file_path)``: return list of dataset paths inside an HDF5 file. +- ``read_h5_dataset(file_path, dataset_path)``: read one dataset as a NumPy array. +- ``read_h5_attrs(file_path, group_path="")``: read HDF5 attributes as a plain dict. +- ``repackage_h5(output_path, sources, ...)``: copy a subset of datasets to a new file. +- ``json_to_h5(output_path, source, ...)``: write JSON data to an HDF5 file. Example: >>> list_h5_files('data') @@ -10,13 +14,20 @@ >>> list_h5_variables('data/run1.h5') ['group1/dset1', 'group2/sub/dset2'] + + >>> read_h5_dataset('data/run1.h5', 'scalars/time') + array([0., 1., 2., ...]) + + >>> read_h5_attrs('data/run1.h5', 'equilibrium') + {'version': 45, 'nspace': 2, 'ntimestep': 0, 'time': 0.0} """ from __future__ import annotations from pathlib import Path -from typing import List +from typing import Any, Dict, List import h5py +import numpy as np def list_h5_files(directory: str | Path, recursive: bool = False) -> List[str]: @@ -72,10 +83,62 @@ def visitor(name, obj): __all__ = [ 'list_h5_files', 'list_h5_variables', + 'read_h5_dataset', + 'read_h5_attrs', 'repackage_h5', + 'json_to_h5', ] +def read_h5_dataset(file_path: str | Path, dataset_path: str) -> np.ndarray: + """Read one dataset from an HDF5 file and return it as a NumPy array. + + Args: + file_path: Path to the HDF5 file. + dataset_path: Internal HDF5 path to the dataset (e.g. ``"scalars/E_K3"``). + + Returns: + NumPy array containing the dataset values. + + Raises: + KeyError: If ``dataset_path`` is not found in the file. + OSError: If the file cannot be opened. + """ + with h5py.File(Path(file_path), 'r') as f: + if dataset_path not in f: + raise KeyError(f"Dataset '{dataset_path}' not found in {file_path}") + return f[dataset_path][()] + + +def read_h5_attrs(file_path: str | Path, group_path: str = "") -> Dict[str, Any]: + """Read HDF5 attributes of a group or dataset as a plain Python dict. + + Args: + file_path: Path to the HDF5 file. + group_path: Internal HDF5 path of the group or dataset whose attributes + are to be read. Empty string (default) reads root-level attributes. + + Returns: + Dict mapping attribute name to value. NumPy scalar types are cast to + Python ``int`` or ``float`` for easy serialisation. + + Raises: + KeyError: If ``group_path`` is not empty and is not found in the file. + OSError: If the file cannot be opened. + """ + result: Dict[str, Any] = {} + with h5py.File(Path(file_path), 'r') as f: + obj = f[group_path] if group_path else f + for key, val in obj.attrs.items(): + if hasattr(val, 'item'): + result[key] = val.item() + elif isinstance(val, np.ndarray): + result[key] = val.tolist() + else: + result[key] = val + return result + + def repackage_h5( output_path: str | Path, sources: List[str | Path], @@ -204,3 +267,160 @@ def repackage_h5( return written +# --------------------------------------------------------------------------- +# json_to_h5 helpers +# --------------------------------------------------------------------------- + +def _try_numeric_array(lst: list) -> "np.ndarray | None": + """Return a numeric ndarray if lst is a uniform or rectangular numeric list, else None.""" + try: + arr = np.array(lst) + except (ValueError, TypeError): + return None + return arr if arr.dtype.kind in ('i', 'u', 'f', 'b') else None + + +def _write_json_node(parent: "h5py.Group", name: str, value: Any) -> List[str]: + """Recursively write one JSON value into an HDF5 group. + + Conversion rules: + None → omitted (returns []) + dict → group; keys become child names + numeric list / rectangular list-of-lists → dataset + other list → numbered subgroups ("0", "1", …) + bool → scalar uint8 dataset (checked before int) + int / float → scalar dataset + str → scalar string dataset + """ + if value is None: + return [] + + if isinstance(value, dict): + grp = parent.require_group(name) + written: List[str] = [] + for k, v in value.items(): + written.extend(_write_json_node(grp, str(k), v)) + return written + + if isinstance(value, list): + arr = _try_numeric_array(value) + if arr is not None: + ds = parent.create_dataset(name, data=arr) + return [ds.name] + grp = parent.require_group(name) + written = [] + for i, item in enumerate(value): + written.extend(_write_json_node(grp, str(i), item)) + return written + + # bool must be tested before int (bool is a subclass of int in Python) + if isinstance(value, bool): + ds = parent.create_dataset(name, data=np.uint8(value)) + return [ds.name] + + if isinstance(value, (int, float)): + ds = parent.create_dataset(name, data=value) + return [ds.name] + + if isinstance(value, str): + ds = parent.create_dataset(name, data=value) + return [ds.name] + + return [] # unrecognised type: skip + + +def json_to_h5( + output_path: str | Path, + source: "str | Path | dict | list", + overwrite: bool = False, +) -> List[str]: + """Write JSON data to an HDF5 file, mapping structure to groups and datasets. + + The ``source`` argument is resolved in this order: + + 1. ``dict`` or ``list`` — used directly as parsed JSON data. + 2. ``str`` or ``Path`` pointing to an existing file — read and parsed as JSON. + 3. ``str`` beginning with ``{`` or ``[`` — parsed as an inline JSON string. + + Conversion rules applied recursively: + + ================== ===================================================== + JSON type HDF5 result + ================== ===================================================== + ``dict`` group; each key becomes a child name + uniform / rect list dataset (converted to ``np.ndarray``) + jagged / mixed list numbered subgroups (``"0"``, ``"1"``, …) + list of dicts numbered subgroups + ``int`` / ``float`` scalar dataset + ``bool`` scalar ``uint8`` dataset + ``str`` scalar string dataset + ``None`` omitted silently + ================== ===================================================== + + A top-level ``list`` is wrapped in a group named ``"data"``. + + Args: + output_path: Destination HDF5 file to create. + source: JSON data as a parsed dict/list, a path to a ``.json`` + file, or an inline JSON string. + overwrite: If ``True``, overwrite an existing ``output_path``. + + Returns: + List of HDF5 dataset paths written into the output file. + + Raises: + FileExistsError: If ``output_path`` exists and ``overwrite=False``. + ValueError: If ``source`` is a string that is neither a valid + file path nor valid JSON. + + Examples:: + + # From a JSON file written by a tool + json_to_h5("spectrum.h5", "spectrum_output.json") + + # From an in-memory dict (e.g. output of compute_poloidal_spectrum) + m, psi, spec = compute_poloidal_spectrum(...) + json_to_h5("spectrum.h5", {"m_modes": m.tolist(), + "psi_norm": psi.tolist(), + "spectrum": spec.tolist()}) + + # From an inline JSON string + json_to_h5("out.h5", '{"time": [0.0, 1.0], "ke": [1e-13, 9e-8]}') + """ + import json as _json + + outp = Path(output_path) + if outp.exists() and not overwrite: + raise FileExistsError(f"Output file exists: {outp}") + outp.parent.mkdir(parents=True, exist_ok=True) + + if isinstance(source, (dict, list)): + data = source + elif isinstance(source, (str, Path)): + src_path = Path(source) + if src_path.exists(): + with src_path.open() as fh: + data = _json.load(fh) + else: + try: + data = _json.loads(str(source)) + except _json.JSONDecodeError as exc: + raise ValueError( + f"source is neither an existing file nor valid JSON: {exc}" + ) from exc + else: + raise TypeError( + f"source must be dict, list, str, or Path; got {type(source).__name__}" + ) + + written: List[str] = [] + with h5py.File(outp, 'w') as f: + if isinstance(data, dict): + for k, v in data.items(): + written.extend(_write_json_node(f, str(k), v)) + else: + written.extend(_write_json_node(f, "data", data)) + + return written + + diff --git a/use_cases/tokamak_stability/m3dc1/__init__.py b/use_cases/tokamak_stability/m3dc1/__init__.py new file mode 100644 index 0000000..dc35ed5 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/__init__.py @@ -0,0 +1,1133 @@ +"""m3dc1 — Python interface to M3D-C1 simulation data via fusion-io. + +Public API +---------- +eval_field(name, R, phi, Z, coord="scalar", sim=..., time=..., quiet=True) +get_time_of_slice(time_idx, filename="C1.h5", units="mks", quiet=True) +get_timetrace(name, sim=..., units="m3dc1", growth=False, renorm=False, quiet=True) +flux_average(field, sim=..., fcoords="pest", points=200) +get_shape(sim, res=250) +flux_coordinates(sim=..., fcoords="pest", phit=0.0, points=200) +eigenfunction(sim=[fc_obj, sim_lin], field="p", coord="scalar", ...) + +Plotting (Priority 2) — require matplotlib: +plot_field, plot_shape, plot_mesh, plot_signal, +plot_time_trace_fast, plot_flux_average, plot_line, plot_field_vs_phi, +plot_mag_probes, plot_gfile, tpf, run_trace, plot_poincare, poincare_movie + +Submodule aliases (for `from m3dc1.eval_field import eval_field` etc.): + m3dc1.eval_field, m3dc1.get_time_of_slice, m3dc1.get_timetrace +""" + +import warnings +from pathlib import Path + +import h5py +import numpy as np + +from . import _neo_input as _ni + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +_KE_ALIASES = { + "ke": "E_K3", + "ke_ion": "E_K3", +} + +# Fields pre-computed by write_neo_input (map caller name → nc variable) +_NEO_FSA_FIELDS = { + "q": "q", + "ne": "ne0", + "te": "Te0", + "ni": "ni0", + "ti": "Ti0", +} + +# Mapping from typedict short name to ftype for fields we know about +_VECTOR_FIELDS = {"j", "v", "B", "A", "E", "gradA"} + +_FIELD_LATEX = { + "j": "J", "B": "B", "A": "A", "E": "E", "v": "v", "gradA": r"\nabla A", + "p": "p", "psi": r"\psi", "ne": "n_e", "te": "T_e", "ni": "n_i", "ti": "T_i", + "q": "q", +} +_COORD_SUB = {"R": "R", "phi": r"\phi", "Z": "Z"} + + +def _field_label(field_name, coord, ftype): + """Return a LaTeX colorbar label for a field/coord combination.""" + base = _FIELD_LATEX.get(field_name) + if base is None: + return field_name + sub = _COORD_SUB.get(coord) + if sub is not None: + return rf"${base}_{{{sub}}}$" + if ftype == "vector": + return rf"$|{base}|$" + return rf"${base}$" + + +def _ftype_of(field_name, sim): + """Return 'vector' or 'scalar' for a field name.""" + if field_name in sim.typedict: + return sim.typedict[field_name][1] + return "scalar" + + +# --------------------------------------------------------------------------- +# eval_field +# --------------------------------------------------------------------------- + +def eval_field(field_name, R, phi, Z, coord="scalar", sim=None, time=None, + quiet=True): + """Evaluate a fusion-io field at arbitrary (R, phi, Z) coordinates. + + Parameters + ---------- + field_name : str + Field name accepted by fpy.sim_data (e.g. 'psi', 'B', 'p'). + R, phi, Z : array_like + Evaluation coordinates (cylindrical, SI). All must be broadcastable. + coord : str + 'scalar' — return the scalar value (or magnitude for vector fields). + 'R', 'phi', 'Z' — return that cylindrical component (vector fields). + 'vector' — return (3, *shape) array of (R, phi, Z) components. + sim : fpy.sim_data + Open simulation object. + time : int or None + Timeslice to evaluate at. If None, uses sim.timeslice. + quiet : bool + Suppress informational output. + + Returns + ------- + np.ndarray + Shape matches R (or (3, *R.shape) for coord='vector'). + Out-of-domain points are NaN. + """ + if sim is None: + raise ValueError("eval_field: sim must be provided") + + if time is not None: + sim.set_timeslice(int(time)) + + R = np.asarray(R, dtype=float) + phi = np.asarray(phi, dtype=float) + Z = np.asarray(Z, dtype=float) + R, phi, Z = np.broadcast_arrays(R, phi, Z) + shape = R.shape + R_flat = R.ravel() + phi_flat = phi.ravel() + Z_flat = Z.ravel() + n = len(R_flat) + + ftype = _ftype_of(field_name, sim) + fld = sim.get_field(field_name, time=None) + + if ftype == "vector" or coord in ("R", "phi", "Z", "vector"): + out = np.full((3, n), np.nan) + for i in range(n): + result = fld.evaluate((float(R_flat[i]), float(phi_flat[i]), + float(Z_flat[i]))) + if result[0] is not None: + out[0, i] = result[0] + if len(result) > 1 and result[1] is not None: + out[1, i] = result[1] + if len(result) > 2 and result[2] is not None: + out[2, i] = result[2] + + if coord == "vector": + return out.reshape((3,) + shape) + comp_map = {"R": 0, "phi": 1, "Z": 2} + if coord in comp_map: + return out[comp_map[coord]].reshape(shape) + # coord == 'scalar' for a vector field → return magnitude. + # Use nansum for the sum but restore NaN where ALL components are NaN + # (out-of-domain points), so pcolormesh autoscale excludes them. + mag = np.sqrt(np.nansum(out ** 2, axis=0)) + mag[np.all(np.isnan(out), axis=0)] = np.nan + return mag.reshape(shape) + + else: + out = np.full(n, np.nan) + for i in range(n): + result = fld.evaluate((float(R_flat[i]), float(phi_flat[i]), + float(Z_flat[i]))) + if result[0] is not None: + out[i] = float(result[0]) + return out.reshape(shape) + + +# --------------------------------------------------------------------------- +# get_time_of_slice +# --------------------------------------------------------------------------- + +def get_time_of_slice(time_idx, filename="C1.h5", units="mks", quiet=True): + """Return the simulation time of a snapshot. + + Parameters + ---------- + time_idx : int + Snapshot index (the NNN in time_NNN.h5). + filename : str or Path + Path to C1.h5. + units : str + 'alfven' or 'm3dc1' — return Alfvén times (raw). + 'mks' or 's' — convert to physical seconds. + quiet : bool + Suppress warnings. + + Returns + ------- + float + Simulation time in the requested units. Returns nan on failure. + """ + try: + with h5py.File(str(filename), "r") as f: + key = f"time_{int(time_idx):03d}" + if key not in f: + if not quiet: + print(f"get_time_of_slice: group {key!r} not found in {filename}") + return float("nan") + t_alfven = float(f[key].attrs["time"]) + + if units in ("alfven", "m3dc1"): + return t_alfven + + # MKS conversion: tau_A = l0 / v_A (compute in CGS, result in seconds) + try: + n0 = float(f.attrs["n0_norm"]) # cm^-3 + l0 = float(f.attrs["l0_norm"]) # cm + b0 = float(f.attrs["b0_norm"]) # Gauss + ion_mass = float(f.attrs["ion_mass"]) # proton masses + m_p_cgs = 1.6726e-24 # grams + v_A_cgs = b0 / np.sqrt(4 * np.pi * ion_mass * m_p_cgs * n0) + tau_A = l0 / v_A_cgs # seconds + return t_alfven * tau_A + except KeyError as exc: + if not quiet: + warnings.warn( + f"get_time_of_slice: cannot compute MKS time " + f"(missing attr {exc}); returning Alfvén time.", + stacklevel=2, + ) + return t_alfven + + except Exception as exc: + if not quiet: + print(f"get_time_of_slice: error reading {filename}: {exc}") + return float("nan") + + +# --------------------------------------------------------------------------- +# get_timetrace +# --------------------------------------------------------------------------- + +def get_timetrace(name, sim=None, units="m3dc1", growth=False, renorm=False, + quiet=True): + """Read a scalar time trace from C1.h5. + + Parameters + ---------- + name : str + Scalar name ('ke' aliases to 'E_K3'; otherwise must be in C1.h5/scalars/). + sim : fpy.sim_data + Open simulation object. + units : str + Time units: 'm3dc1'/'alfven' (Alfvén times) or 'mks'/'s' (seconds). + growth : bool + If True, return instantaneous growth rate gamma = 0.5 * d(ln values)/dt + instead of the raw values. + renorm : bool + If True, divide values by the first nonzero value before returning. + quiet : bool + Suppress output. + + Returns + ------- + (time, values, label, units_str) : tuple of (ndarray, ndarray, str, str) + """ + if sim is None: + raise ValueError("get_timetrace: sim must be provided") + + raw_name = _KE_ALIASES.get(name.lower(), name) + try: + time = np.asarray(sim._all_traces["time"], dtype=float) + values = np.asarray(sim._all_traces[raw_name], dtype=float) + except KeyError: + raise KeyError(f"get_timetrace: trace {raw_name!r} not found in C1.h5/scalars/") + + # Trim to equal length if needed + n = min(len(time), len(values)) + time = time[:n] + values = values[:n] + + # Remove NaNs + valid = ~np.isnan(values) + time = time[valid] + values = values[valid] + + label = raw_name + units_str = units + + if renorm: + nonzero = values != 0 + if nonzero.any(): + values = values / values[nonzero][0] + + if growth: + # gamma = 0.5 * d(ln |values|) / dt + log_vals = np.log(np.abs(values)) + dt = np.diff(time) + gamma = 0.5 * np.diff(log_vals) / np.where(dt != 0, dt, np.nan) + time = 0.5 * (time[:-1] + time[1:]) + values = gamma + label = f"gamma({raw_name})" + units_str = "1/tau_A" + + if units in ("mks", "s") and not growth: + # Scale time axis only; values are dimensionless or in code units + # Try to get the Alfvén time from the file + try: + with h5py.File(sim.filename, "r") as f: + n0 = float(f.attrs["n0_norm"]) + l0 = float(f.attrs["l0_norm"]) + b0 = float(f.attrs["b0_norm"]) + ion_mass_val = float(f.attrs["ion_mass"]) + m_p_cgs = 1.6726e-24 + v_A_cgs = b0 / np.sqrt(4 * np.pi * ion_mass_val * m_p_cgs * n0) + tau_A = l0 / v_A_cgs + time = time * tau_A + units_str = "s" + except Exception: + if not quiet: + warnings.warn("get_timetrace: cannot convert time to MKS; " + "returning Alfvén times.", stacklevel=2) + + return time, values, label, units_str + + +# --------------------------------------------------------------------------- +# flux_average +# --------------------------------------------------------------------------- + +def flux_average(field, sim=None, fcoords="pest", points=200, ntheta=300, nphi=4): + """Compute the flux-surface average of a field. + + Parameters + ---------- + field : str + Field name. For 'q', 'ne', 'te', 'ni', 'ti' the pre-computed values + from write_neo_input are returned directly. All other fields are + evaluated on the 3-D flux surface grid and averaged with the Jacobian. + sim : fpy.sim_data + Equilibrium simulation object. + fcoords : str + Flux coordinate system passed to write_neo_input (currently unused; + reserved for future coordinate choices). + points : int + Number of radial grid points (nr). + ntheta : int + Poloidal grid points per surface for the FSA integration. + nphi : int + Toroidal grid points for the FSA integration. Use nphi=1 for a + strictly axisymmetric equilibrium (exact); nphi>=4 if the field + has toroidal variation. + + Returns + ------- + (psi_norm, profile) : (ndarray, ndarray), each of shape (nr,) + """ + if sim is None: + raise ValueError("flux_average: sim must be provided") + + c1h5 = Path(sim.filename) + timeslice = sim.timeslice + nc_path, tmpdir_obj = _ni.run_write_neo_input( + c1h5, timeslice, + psi_start=0.01, psi_end=0.99, + nr=points, ntheta=ntheta, nphi=nphi, + ) + try: + neo = _ni.read_neo_input(nc_path) + finally: + if tmpdir_obj is not None: + tmpdir_obj.cleanup() + + psi_norm = neo["psi_norm"] + + # Direct lookup for pre-computed averages + neo_key = _NEO_FSA_FIELDS.get(field.lower()) + if neo_key is not None and neo_key in neo: + return psi_norm, neo[neo_key] + + # General case: evaluate field on 3-D grid and integrate + R_grid = neo["R"] # (ntheta, nphi, nr) + Z_grid = neo["Z"] + Jac = neo["Jac"] + Phi = neo["Phi"] # (nphi,) toroidal angles + + ntheta, nphi, nr = R_grid.shape + theta = np.linspace(0, 2 * np.pi, ntheta, endpoint=False) + + # Evaluate field at every grid point + R_flat = R_grid.ravel() + Z_flat = Z_grid.ravel() + phi_3d = np.broadcast_to( + Phi[np.newaxis, :, np.newaxis], (ntheta, nphi, nr) + ).copy() + phi_flat = phi_3d.ravel() + + f_flat = eval_field(field, R_flat, phi_flat, Z_flat, + coord="scalar", sim=sim, time=timeslice, quiet=True) + f_grid = f_flat.reshape(ntheta, nphi, nr) + + profile = np.empty(nr) + for s in range(nr): + profile[s] = _ni.flux_surface_average( + f_grid[:, :, s], Jac[:, :, s], theta, Phi + ) + + return psi_norm, profile + + +# --------------------------------------------------------------------------- +# get_shape +# --------------------------------------------------------------------------- + +def get_shape(sim, res=250): + """Compute Miller geometry parameters of the last closed flux surface. + + Parameters + ---------- + sim : fpy.sim_data + Equilibrium simulation object. + res : int + Grid resolution per axis for psi evaluation (res × res points). + + Returns + ------- + dict with keys 'R0', 'a', 'kappa', 'delta' (all in metres). + Returns {} on failure. + """ + try: + import matplotlib # noqa: PLC0415 + matplotlib.use("Agg") + import matplotlib.pyplot as plt # noqa: PLC0415 + + # Read LCFS psi value + scalars = sim._all_attrs["scalars"] + psi_lcfs = float(np.asarray(scalars["psi_lcfs"])[0]) + + # Get mesh bounding box + mesh = sim.get_mesh(quiet=True) + R_mesh = mesh.elements[:, 4] + Z_mesh = mesh.elements[:, 5] + R_min, R_max = R_mesh.min(), R_mesh.max() + Z_min, Z_max = Z_mesh.min(), Z_mesh.max() + + # Regular grid for psi evaluation + R_1d = np.linspace(R_min, R_max, res) + Z_1d = np.linspace(Z_min, Z_max, res) + R_2d, Z_2d = np.meshgrid(R_1d, Z_1d) + phi_2d = np.zeros_like(R_2d) + + psi_2d = eval_field("psi", R_2d, phi_2d, Z_2d, + coord="scalar", sim=sim, + time=sim.timeslice, quiet=True) + + # Extract LCFS contour (collect paths before closing figure) + fig, ax = plt.subplots() + cs = ax.contour(R_2d, Z_2d, psi_2d, levels=[psi_lcfs]) + try: + paths = list(cs.get_paths()) # matplotlib >= 3.8 + except AttributeError: + paths = [] + for coll in cs.collections: # matplotlib < 3.8 + paths.extend(coll.get_paths()) + plt.close(fig) + if not paths: + return {} + longest = max(paths, key=lambda p: len(p.vertices)) + R_lcfs = longest.vertices[:, 0] + Z_lcfs = longest.vertices[:, 1] + + R0 = (R_lcfs.max() + R_lcfs.min()) / 2.0 + a = (R_lcfs.max() - R_lcfs.min()) / 2.0 + kappa = (Z_lcfs.max() - Z_lcfs.min()) / (2.0 * a) + R_at_Zmax = R_lcfs[np.argmax(Z_lcfs)] + delta = (R0 - R_at_Zmax) / a + + return {"R0": float(R0), "a": float(a), + "kappa": float(kappa), "delta": float(delta)} + + except Exception as exc: + warnings.warn(f"get_shape failed: {exc}", stacklevel=2) + return {} + + +# --------------------------------------------------------------------------- +# flux_coordinates +# --------------------------------------------------------------------------- + +class _FCSummary: + """Holds the flux-coordinate summary accessible via flux_coordinates().fc.""" + def __init__(self, neo_data): + self.psi_norm = neo_data["psi_norm"] # (nr,) + self.q = neo_data["q"] # (nr,) + self.rpath = neo_data["R"][:, 0, :] # (ntheta, nr) at phi plane 0 + self.zpath = neo_data["Z"][:, 0, :] # (ntheta, nr) + + +class _FluxCoordsResult: + """Return value of flux_coordinates(). + + Attributes + ---------- + fc : _FCSummary + Exposes .psi_norm, .q, .rpath, .zpath + """ + def __init__(self, neo_data, sim, phit, tmpdir_obj): + self.fc = _FCSummary(neo_data) + self._neo = neo_data + self._sim = sim + self._phit = float(phit) + self._tmpdir_obj = tmpdir_obj # keep temp dir alive + + +def flux_coordinates(sim=None, fcoords="pest", phit=0.0, points=200, ntheta=128): + """Compute flux coordinates using write_neo_input. + + Parameters + ---------- + sim : fpy.sim_data + Equilibrium simulation object. + fcoords : str + Flux coordinate system (reserved; currently 'pest' only). + phit : float + Toroidal angle in radians for the poloidal cross-section. + points : int + Number of radial grid points (nr). + ntheta : int + Number of poloidal grid points per flux surface. Controls the highest + resolvable poloidal mode number (max_m = ntheta//2) and the cost of + subsequent eigenfunction() calls. 128 resolves up to m=64, which + exceeds typical M3D-C1 mesh resolution. + + Returns + ------- + _FluxCoordsResult + Object with .fc.psi_norm, passed as sim[0] to eigenfunction(). + """ + if sim is None: + raise ValueError("flux_coordinates: sim must be provided") + + c1h5 = Path(sim.filename) + timeslice = sim.timeslice + nc_path, tmpdir_obj = _ni.run_write_neo_input( + c1h5, timeslice, + psi_start=0.01, psi_end=0.99, + nr=points, ntheta=ntheta, nphi=1, + ) + neo = _ni.read_neo_input(nc_path) + return _FluxCoordsResult(neo, sim, phit, tmpdir_obj) + + +# --------------------------------------------------------------------------- +# eigenfunction +# --------------------------------------------------------------------------- + +def eigenfunction(sim=None, field="p", coord="scalar", fcoords="pest", + points=200, makeplot=False, fourier=True, full_fft=False, + norm_to_unity=True, quiet=True): + """Compute the poloidal mode spectrum of a perturbed field. + + Parameters + ---------- + sim : [_FluxCoordsResult, fpy.sim_data] + sim[0] — flux coordinate result from flux_coordinates(). + sim[1] — perturbed simulation object (linear timeslice). + field : str + Field to decompose (e.g. 'p', 'B'). + coord : str + 'scalar', 'R', 'phi', or 'Z'. + fcoords : str + Flux coordinate system (reserved). + points : int + Number of radial points (inherited from flux_coordinates). + makeplot : bool + If True, plot amplitude vs psi_norm. + fourier : bool + If True, return Fourier spectrum in poloidal mode number. + If False, return the raw (theta, nr) field. + full_fft : bool + If True, return full two-sided FFT (using np.fft.fft + fftshift). + If False, return one-sided amplitude (using np.fft.rfft). + norm_to_unity : bool + If True, divide the spectrum by its global maximum. + quiet : bool + Suppress output. + + Returns + ------- + np.ndarray + Shape (ntheta//2 + 1, nr) for fourier=True, full_fft=False. + Shape (ntheta, nr) for fourier=True, full_fft=True. + Shape (ntheta, nr) for fourier=False. + """ + if sim is None or not hasattr(sim, "__len__") or len(sim) < 2: + raise ValueError( + "eigenfunction: sim must be [flux_coords_result, sim_data_linear]" + ) + + fc_result = sim[0] + sim_lin = sim[1] + + # Poloidal cross-section at the stored phi plane (index 0 in neo nphi axis) + R_grid = fc_result._neo["R"][:, 0, :] # (ntheta, nr) + Z_grid = fc_result._neo["Z"][:, 0, :] # (ntheta, nr) + ntheta, nr = R_grid.shape + phi_arr = np.full_like(R_grid, fc_result._phit) + + f_flat = eval_field( + field, + R_grid.ravel(), phi_arr.ravel(), Z_grid.ravel(), + coord=coord, sim=sim_lin, + time=sim_lin.timeslice, quiet=quiet, + ) + f_grid = f_flat.reshape(ntheta, nr) + + if not fourier: + spec = f_grid + elif full_fft: + spec = np.abs(np.fft.fftshift(np.fft.fft(f_grid, axis=0), axes=0)) + else: + spec = np.abs(np.fft.rfft(f_grid, axis=0)) + + if norm_to_unity: + peak = spec.max() + if peak > 0: + spec = spec / peak + + if makeplot: + try: + import matplotlib.pyplot as plt # noqa: PLC0415 + psin = fc_result.fc.psi_norm + n_modes = spec.shape[0] + fig, ax = plt.subplots() + for m in range(min(n_modes, 20)): + ax.plot(psin, spec[m, :], label=f"m={m}") + ax.set_xlabel(r"$\psi_\mathrm{norm}$") + ax.set_ylabel("Mode amplitude") + ax.set_title(f"Eigenfunction: {field} ({coord})") + ax.legend(fontsize="small", ncol=2) + plt.tight_layout() + if plt.isinteractive(): + plt.show() + except Exception as exc: + if not quiet: + warnings.warn(f"eigenfunction makeplot failed: {exc}", stacklevel=2) + + return spec + + +# --------------------------------------------------------------------------- +# Priority 2 — Plotting functions +# --------------------------------------------------------------------------- + +def _require_sim(filename, time, filetype="m3dc1"): + """Open a sim_data object from a filename.""" + try: + import sys + import os + fio_lib = os.path.join(os.path.dirname(os.path.dirname(__file__))) + if fio_lib not in sys.path: + sys.path.insert(0, fio_lib) + from fpy import sim_data # noqa: PLC0415 + except ImportError: + raise ImportError("fpy.sim_data not importable; check your fusion-io installation.") + return sim_data(filename=str(filename), filetype=filetype, time=time) + + +def plot_field(field, filename, time=0, coord="scalar", phi=0.0, points=250, + tor_av=1, units="mks", cmap="inferno", mesh=False, bound=False, lcfs=False, + coils=False, save=False, savedir="./", quiet=True): + """Plot a field on a regular (R, Z) grid as a filled contour map. + + Evaluates *field* on a *points* × *points* grid at fixed toroidal angle + *phi*. If *tor_av* > 1, averages over that many equally-spaced phi values. + """ + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, time) + mesh_obj = sim.get_mesh(quiet=True) + R_mesh = mesh_obj.elements[:, 4] + Z_mesh = mesh_obj.elements[:, 5] + R_1d = np.linspace(R_mesh.min(), R_mesh.max(), points) + Z_1d = np.linspace(Z_mesh.min(), Z_mesh.max(), points) + R_2d, Z_2d = np.meshgrid(R_1d, Z_1d) + + if tor_av > 1: + phis = np.linspace(0, 2 * np.pi, tor_av, endpoint=False) + f_sum = np.zeros_like(R_2d) + for p in phis: + f_sum += eval_field(field, R_2d, np.full_like(R_2d, p), Z_2d, + coord=coord, sim=sim, time=time, quiet=True) + f_2d = f_sum / tor_av + else: + f_2d = eval_field(field, R_2d, np.full_like(R_2d, phi), Z_2d, + coord=coord, sim=sim, time=time, quiet=True) + + pos_extreme = max(float(np.nanmax(f_2d)), 0.0) + neg_extreme = abs(min(float(np.nanmin(f_2d)), 0.0)) + larger = max(pos_extreme, neg_extreme) + smaller = min(pos_extreme, neg_extreme) + if larger > 0 and smaller / larger > 0.05: + pcm_kwargs = dict(cmap="RdBu_r", vmin=-larger, vmax=larger, shading="auto") + else: + pcm_kwargs = dict(cmap=cmap, shading="auto") + + ftype = _ftype_of(field, sim) + label = _field_label(field, coord, ftype) + + fig, ax = plt.subplots(figsize=(5, 7)) + pcm = ax.pcolormesh(R_2d, Z_2d, f_2d, **pcm_kwargs) + plt.colorbar(pcm, ax=ax, label=label) + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + + if lcfs: + try: + psi_lcfs = float(np.asarray(sim._all_attrs["scalars"]["psi_lcfs"])[0]) + psi_2d = eval_field("psi", R_2d, np.zeros_like(R_2d), Z_2d, + coord="scalar", sim=sim, time=time, quiet=True) + ax.contour(R_2d, Z_2d, psi_2d, levels=[psi_lcfs], colors="w") + except Exception: + pass + + ax.set_title(f"{label} t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"{field}_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_shape(filename, time=0, points=200, save=False, savedir="./", + quiet=True): + """Plot psi contours (flux surfaces).""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, time) + mesh_obj = sim.get_mesh(quiet=True) + R_mesh = mesh_obj.elements[:, 4] + Z_mesh = mesh_obj.elements[:, 5] + R_1d = np.linspace(R_mesh.min(), R_mesh.max(), points) + Z_1d = np.linspace(Z_mesh.min(), Z_mesh.max(), points) + R_2d, Z_2d = np.meshgrid(R_1d, Z_1d) + psi_2d = eval_field("psi", R_2d, np.zeros_like(R_2d), Z_2d, + coord="scalar", sim=sim, time=time, quiet=True) + + fig, ax = plt.subplots(figsize=(5, 7)) + ax.contour(R_2d, Z_2d, psi_2d, levels=20) + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title(f"Flux surfaces t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"shape_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_mesh(filename, time=0, boundary=False, save=False, savedir="./", + quiet=True): + """Plot the M3D-C1 triangular mesh.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + import matplotlib.tri as tri # noqa: PLC0415 + + # Read mesh geometry directly from HDF5; no fpy needed for vertex positions. + # Columns 0-2 of mesh/elements are not integer node indices (they are + # floating-point values), so connectivity cannot be derived from them. + # Instead, deduplicate per-element (R, Z) from columns 4-5 to get unique + # vertices and let matplotlib compute a Delaunay triangulation. + with h5py.File(filename, "r") as _f: + if "equilibrium/mesh/elements" in _f: + _el = _f["equilibrium/mesh/elements"][:] + elif "mesh/elements" in _f: + _el = _f["mesh/elements"][:] + else: + raise KeyError(f"No mesh/elements group found in {filename}") + + _rz = np.unique(np.column_stack([_el[:, 4], _el[:, 5]]), axis=0) + R = _rz[:, 0].astype(float) + Z = _rz[:, 1].astype(float) + triang = tri.Triangulation(R, Z) + + fig, ax = plt.subplots() + ax.triplot(triang, lw=0.3, color="k") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title("Mesh") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / "mesh.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_signal(filename, signame, save=False, savedir="./", quiet=True): + """Plot a diagnostic signal (probe/flux loop) time trace.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, 0) + sig = sim.get_signal(signame) + fig, ax = plt.subplots() + time = np.arange(len(sig.sigvalues)) + ax.plot(time, sig.sigvalues) + ax.set_xlabel("Time step") + ax.set_ylabel(signame) + ax.set_title(signame) + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"signal_{signame}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_time_trace_fast(trace, filename, units="m3dc1", save=False, + savedir="./", quiet=True): + """Plot log(E_K3) vs time, optionally with fitted growth rate.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, 0) + time, values, label, units_str = get_timetrace( + trace, sim=sim, units=units, quiet=quiet + ) + fig, ax = plt.subplots() + ax.semilogy(time, np.abs(values)) + ax.set_xlabel(f"t ({units_str})") + ax.set_ylabel(label) + ax.set_title(f"{label} time trace") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"timetrace_{trace}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_flux_average(field, filename, time=0, fcoords="pest", points=200, + save=False, savedir="./", quiet=True): + """Plot the flux-surface average of a field vs psi_norm.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, time) + psin, profile = flux_average(field, sim=sim, fcoords=fcoords, points=points) + fig, ax = plt.subplots() + ax.plot(psin, profile) + ax.set_xlabel(r"$\psi_\mathrm{norm}$") + ax.set_ylabel(field) + ax.set_title(f"FSA of {field} t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"fsa_{field}_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_line(field, filename, time=0, angle=0.0, Zoff=0.0, coord="scalar", + dist_from_magax=False, points=200, save=False, savedir="./", + quiet=True): + """Plot a field along a radial line in (R, Z) at fixed phi=0.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, time) + mesh_obj = sim.get_mesh(quiet=True) + R_max = mesh_obj.elements[:, 4].max() + R_min = mesh_obj.elements[:, 4].min() + theta_rad = np.deg2rad(angle) + R_line = np.linspace(R_min, R_max, points) + Z_line = np.full_like(R_line, Zoff) + R_line * np.tan(theta_rad) + phi_line = np.zeros_like(R_line) + f_line = eval_field(field, R_line, phi_line, Z_line, + coord=coord, sim=sim, time=time, quiet=True) + + x_axis = R_line + xlabel = "R (m)" + if dist_from_magax: + try: + R_mag = float(np.asarray(sim._all_attrs["scalars"]["xmag"])[0]) + x_axis = np.abs(R_line - R_mag) + xlabel = "|R - R_mag| (m)" + except Exception: + pass + + fig, ax = plt.subplots() + ax.plot(x_axis, f_line) + ax.set_xlabel(xlabel) + ax.set_ylabel(f"{field} ({coord})") + ax.set_title(f"Radial profile: {field} t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"line_{field}_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_field_vs_phi(field, filename, time=0, R=None, Z=None, phi_res=64, + coord="scalar", save=False, savedir="./", quiet=True): + """Plot a field at fixed (R, Z) as a function of toroidal angle phi.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + sim = _require_sim(filename, time) + if R is None or Z is None: + R_mag = float(np.asarray(sim._all_attrs["scalars"]["xmag"])[0]) + Z_mag = float(np.asarray(sim._all_attrs["scalars"]["zmag"])[0]) + R = R_mag + Z = Z_mag + phis = np.linspace(0, 2 * np.pi, phi_res, endpoint=False) + R_arr = np.full_like(phis, R) + Z_arr = np.full_like(phis, Z) + f_arr = eval_field(field, R_arr, phis, Z_arr, + coord=coord, sim=sim, time=time, quiet=True) + fig, ax = plt.subplots() + ax.plot(np.rad2deg(phis), f_arr) + ax.set_xlabel("phi (degrees)") + ax.set_ylabel(f"{field} ({coord})") + ax.set_title(f"{field} vs phi at R={R:.3f}, Z={Z:.3f} t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"vs_phi_{field}_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def tpf(field, sim, filename, units="m3dc1", res=64): + """Return (time_value, toroidal_peaking_factor) for the midplane. + + The toroidal peaking factor is max(|f|) / mean(|f|) evaluated toroidally + at the outboard midplane. + """ + R_mag = float(np.asarray(sim._all_attrs["scalars"]["xmag"])[0]) + Z_mag = float(np.asarray(sim._all_attrs["scalars"]["zmag"])[0]) + phis = np.linspace(0, 2 * np.pi, res, endpoint=False) + R_arr = np.full_like(phis, R_mag) + Z_arr = np.full_like(phis, Z_mag) + f_arr = eval_field(field, R_arr, phis, Z_arr, + coord="scalar", sim=sim, time=sim.timeslice, quiet=True) + f_abs = np.abs(f_arr) + mean_f = f_abs.mean() + pf = f_abs.max() / mean_f if mean_f > 0 else float("nan") + + time_arr = np.asarray(sim._all_traces["time"]) + t = float(time_arr[sim.timeslice]) if sim.timeslice < len(time_arr) else float("nan") + return t, pf + + +def run_trace(filename, time=0, nparticles=1000, nsteps=500, verbose=False): + """Run the fusion-io field-line tracer and return the output path.""" + import os as _os # noqa: PLC0415 + import shutil # noqa: PLC0415 + import subprocess as sp # noqa: PLC0415 + + exe = shutil.which("trace") + fio_dir = _os.environ.get("FIO_INSTALL_DIR") + if exe is None and fio_dir: + candidate = Path(fio_dir) / "bin" / "trace" + if candidate.is_file(): + exe = str(candidate) + if exe is None: + raise RuntimeError( + "trace executable not found. Set $FIO_INSTALL_DIR or add trace to $PATH." + ) + cmd = [exe, "-m3dc1", str(filename), str(time), + "-nparticles", str(nparticles), "-nsteps", str(nsteps)] + sp.run(cmd, check=True, capture_output=not verbose) + return Path("poincare.dat") + + +def plot_poincare(poincare_file="poincare.dat", save=False, savedir="./"): + """Plot a Poincaré section from a trace output file.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + data = np.loadtxt(str(poincare_file)) + fig, ax = plt.subplots() + ax.scatter(data[:, 0], data[:, 1], s=0.5, c="k") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title("Poincaré section") + plt.tight_layout() + if save: + plt.savefig(Path(savedir) / "poincare.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def poincare_movie(filename, time_indices=None, nparticles=500, nsteps=300, + savedir="./poincare_frames/", verbose=False): + """Generate Poincaré sections for a sequence of timeslices.""" + import os as _os # noqa: PLC0415 + _os.makedirs(savedir, exist_ok=True) + with h5py.File(str(filename), "r") as f: + ntime = int(f.attrs.get("ntime", 0)) + if time_indices is None: + time_indices = list(range(ntime)) + for t in time_indices: + poincare_path = run_trace(filename, time=t, nparticles=nparticles, + nsteps=nsteps, verbose=verbose) + plot_poincare(poincare_path, save=True, savedir=savedir) + + +def plot_mag_probes(filename, save=False, savedir="./", quiet=True): + """Plot magnetic probe (R, Z) positions from coil.dat.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + coil_dat = Path(filename).parent / "coil.dat" + if not coil_dat.exists(): + warnings.warn(f"plot_mag_probes: {coil_dat} not found", stacklevel=2) + return + R_list, Z_list = [], [] + with open(coil_dat) as fh: + for line in fh: + parts = line.split() + if len(parts) >= 2: + try: + R_list.append(float(parts[0])) + Z_list.append(float(parts[1])) + except ValueError: + pass + fig, ax = plt.subplots() + ax.scatter(R_list, Z_list, marker="x") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title("Magnetic probe positions") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / "mag_probes.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def _parse_geqdsk_grid(filename): + """Parse an EFIT GEQDSK file and return a dict of equilibrium arrays.""" + with open(str(filename)) as fh: + header = fh.readline() + parts = header.split() + nw = int(parts[-2]) + nh = int(parts[-1]) + + def _read_floats(n): + vals = [] + while len(vals) < n: + line = fh.readline() + vals.extend([float(x) for x in line.split()]) + return np.array(vals[:n]) + + efit_data = {} + scalars = _read_floats(20) + keys = ["rdim", "zdim", "rcentr", "rleft", "zmid", + "rmaxis", "zmaxis", "simag", "sibry", "bcentr", + "current", "simag2", "xdum", "rmaxis2", "xdum2", + "zmaxis2", "xdum3", "sibry2", "xdum4", "xdum5"] + for k, v in zip(keys, scalars): + efit_data[k] = v + + efit_data["fpol"] = _read_floats(nw) + efit_data["pres"] = _read_floats(nw) + efit_data["ffprim"] = _read_floats(nw) + efit_data["pprime"] = _read_floats(nw) + efit_data["psirz"] = _read_floats(nw * nh).reshape(nh, nw) + efit_data["qpsi"] = _read_floats(nw) + nbbbs = int(fh.readline().split()[0]) + bbbs = _read_floats(2 * nbbbs) + efit_data["rbbbs"] = bbbs[::2] + efit_data["zbbbs"] = bbbs[1::2] + efit_data["nw"] = nw + efit_data["nh"] = nh + efit_data["rdim"] = efit_data["rdim"] + return efit_data + + +def plot_gfile(gfile_path, save=False, savedir="./", quiet=True): + """Parse a GEQDSK file and plot psi contours and LCFS.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + + try: + g = _parse_geqdsk_grid(gfile_path) + except Exception as exc: + warnings.warn(f"plot_gfile: could not parse {gfile_path}: {exc}", stacklevel=2) + return + + nw, nh = g["nw"], g["nh"] + rdim, zdim = g["rdim"], g["zdim"] + rleft, zmid = g["rleft"], g["zmid"] + R_1d = np.linspace(rleft, rleft + rdim, nw) + Z_1d = np.linspace(zmid - zdim / 2, zmid + zdim / 2, nh) + R_2d, Z_2d = np.meshgrid(R_1d, Z_1d) + + fig, ax = plt.subplots() + ax.contour(R_2d, Z_2d, g["psirz"], levels=20) + ax.plot(g["rbbbs"], g["zbbbs"], "r-", lw=2, label="LCFS") + ax.plot(g["rmaxis"], g["zmaxis"], "r+", ms=10) + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title(str(gfile_path)) + ax.legend() + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / "gfile.png", dpi=150) + if plt.isinteractive(): + plt.show() + + +def plot_vector_field(field, filename, time=0, points=30, phi=0.0, + quiet=True): + """Plot a vector field using mayavi (requires mayavi installation).""" + try: + from mayavi import mlab # noqa: F401, PLC0415 + except ImportError: + raise ImportError( + "plot_vector_field requires mayavi. " + "Install it with: conda install -c conda-forge mayavi" + ) + raise NotImplementedError( + "plot_vector_field mayavi rendering is not yet implemented." + ) + + +# --------------------------------------------------------------------------- +# Public API list +# --------------------------------------------------------------------------- + +__all__ = [ + "eval_field", + "get_time_of_slice", + "get_timetrace", + "flux_average", + "get_shape", + "flux_coordinates", + "eigenfunction", + "plot_field", + "plot_shape", + "plot_mesh", + "plot_signal", + "plot_time_trace_fast", + "plot_flux_average", + "plot_line", + "plot_field_vs_phi", + "tpf", + "run_trace", + "plot_poincare", + "poincare_movie", + "plot_mag_probes", + "plot_gfile", + "plot_vector_field", +] diff --git a/use_cases/tokamak_stability/m3dc1/_neo_input.py b/use_cases/tokamak_stability/m3dc1/_neo_input.py new file mode 100644 index 0000000..f1a1e14 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/_neo_input.py @@ -0,0 +1,139 @@ +"""Private module: run write_neo_input subprocess and read neo_input.nc.""" + +import os +import shutil +import subprocess +import tempfile +from pathlib import Path + +import numpy as np + + +def _find_write_neo_input(): + fio_dir = os.environ.get("FIO_INSTALL_DIR") + if fio_dir: + candidate = Path(fio_dir) / "bin" / "write_neo_input" + if candidate.is_file(): + return str(candidate) + found = shutil.which("write_neo_input") + if found: + return found + raise RuntimeError( + "write_neo_input executable not found. " + "Set $FIO_INSTALL_DIR to the fusion-io install prefix " + "or add write_neo_input to $PATH." + ) + + +def run_write_neo_input( + c1h5_path, + timeslice, + psi_start=0.01, + psi_end=0.99, + nr=50, + ntheta=400, + nphi=1, + dl=0.01, + workdir=None, + verbose=False, +): + """Run write_neo_input and return (nc_path, tmpdir_obj). + + tmpdir_obj is a TemporaryDirectory instance if workdir was None (keep a + reference to prevent early cleanup), or None if workdir was provided. + nc_path is the absolute Path to neo_input.nc inside workdir. + """ + exe = _find_write_neo_input() + c1h5_path = Path(c1h5_path).resolve() + + tmpdir_obj = None + if workdir is None: + tmpdir_obj = tempfile.TemporaryDirectory() + workdir = Path(tmpdir_obj.name) + else: + workdir = Path(workdir) + workdir.mkdir(parents=True, exist_ok=True) + + cmd = [ + exe, + "-m3dc1", str(c1h5_path), str(timeslice), + "-psi_start", str(psi_start), + "-psi_end", str(psi_end), + "-nr", str(nr), + "-ntheta", str(ntheta), + "-nphi", str(nphi), + "-dl", str(dl), + "-bootstrap", "0", + ] + + subprocess.run( + cmd, + cwd=str(workdir), + check=True, + capture_output=not verbose, + ) + + nc_path = workdir / "neo_input.nc" + if not nc_path.exists(): + raise RuntimeError( + f"write_neo_input completed but {nc_path} was not created." + ) + return nc_path, tmpdir_obj + + +def read_neo_input(nc_path): + """Read neo_input.nc; return a plain dict of NumPy arrays. + + Tries netCDF4 first, falls back to scipy.io.netcdf_file. + Adds derived 'psi_norm' = (psi - psi_0) / (psi_1 - psi_0). + """ + nc_path = str(nc_path) + + _3d_vars = ["R", "Z", "Jac", "B_R", "B_Phi", "B_Z", "p_3d", "ne_3d", "ni_3d"] + _1d_vars = ["q", "psi", "ne0", "Te0", "ni0", "Ti0", "b2_fa", "bmag_fa", "Phi"] + _opt_1d = ["bx_fa"] + + try: + import netCDF4 # noqa: PLC0415 + + with netCDF4.Dataset(nc_path, "r") as ds: + data = {v: np.array(ds.variables[v][:]) for v in _3d_vars + _1d_vars} + for v in _opt_1d: + if v in ds.variables: + data[v] = np.array(ds.variables[v][:]) + data["ion_mass"] = float(ds.ion_mass) + data["psi_0"] = float(ds.psi_0) + data["psi_1"] = float(ds.psi_1) + + except ImportError: + from scipy.io import netcdf_file # noqa: PLC0415 + + with netcdf_file(nc_path, "r", mmap=False) as ds: + data = {v: np.array(ds.variables[v].data) for v in _3d_vars + _1d_vars} + for v in _opt_1d: + if v in ds.variables: + data[v] = np.array(ds.variables[v].data) + data["ion_mass"] = float(ds.ion_mass) + data["psi_0"] = float(ds.psi_0) + data["psi_1"] = float(ds.psi_1) + + psi_0 = data["psi_0"] + psi_1 = data["psi_1"] + data["psi_norm"] = (data["psi"] - psi_0) / (psi_1 - psi_0) + return data + + +def flux_surface_average(quantity, jacobian, theta, phi): + """Flux-surface average: = integral(f*J dtheta dphi) / integral(J dtheta dphi). + + quantity, jacobian: shape (ntheta, nphi) for a single flux surface. + theta: 1-D array of length ntheta. + phi: 1-D array of length nphi. + """ + try: + _trapz = np.trapezoid # numpy >= 2.0 + except AttributeError: + _trapz = np.trapz # numpy < 2.0 + denom = _trapz(_trapz(jacobian, x=theta, axis=0), x=phi, axis=0) + numer = _trapz(_trapz(quantity * jacobian, x=theta, axis=0), x=phi, axis=0) + return numer / denom diff --git a/use_cases/tokamak_stability/m3dc1/eval_field.py b/use_cases/tokamak_stability/m3dc1/eval_field.py new file mode 100644 index 0000000..bfcf851 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/eval_field.py @@ -0,0 +1,4 @@ +"""Re-export for `from m3dc1.eval_field import eval_field`.""" +from m3dc1 import eval_field # noqa: F401 + +__all__ = ["eval_field"] diff --git a/use_cases/tokamak_stability/m3dc1/get_time_of_slice.py b/use_cases/tokamak_stability/m3dc1/get_time_of_slice.py new file mode 100644 index 0000000..6723013 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/get_time_of_slice.py @@ -0,0 +1,4 @@ +"""Re-export for `from m3dc1.get_time_of_slice import get_time_of_slice`.""" +from m3dc1 import get_time_of_slice # noqa: F401 + +__all__ = ["get_time_of_slice"] diff --git a/use_cases/tokamak_stability/m3dc1/get_timetrace.py b/use_cases/tokamak_stability/m3dc1/get_timetrace.py new file mode 100644 index 0000000..4c66251 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/get_timetrace.py @@ -0,0 +1,4 @@ +"""Re-export for `from m3dc1.get_timetrace import get_timetrace`.""" +from m3dc1 import get_timetrace # noqa: F401 + +__all__ = ["get_timetrace"] diff --git a/use_cases/tokamak_stability/m3dc1/plot_field_basic.py b/use_cases/tokamak_stability/m3dc1/plot_field_basic.py new file mode 100644 index 0000000..e4f81c6 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/plot_field_basic.py @@ -0,0 +1,9 @@ +"""Submodule: `import m3dc1.plot_field_basic` — exposes mutable shortlbl/label globals.""" +from m3dc1 import eval_field, plot_field # noqa: F401 + +# These module-level globals are set by calling code (e.g. plotsm3dc1.py) +# before invoking the plot functions. +shortlbl: dict = {} +label: dict = {} + +__all__ = ["plot_field", "eval_field", "shortlbl", "label"] diff --git a/use_cases/tokamak_stability/m3dc1/plot_field_mesh.py b/use_cases/tokamak_stability/m3dc1/plot_field_mesh.py new file mode 100644 index 0000000..f8621ce --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/plot_field_mesh.py @@ -0,0 +1,67 @@ +"""Submodule: `import m3dc1.plot_field_mesh` — exposes mutable shortlbl/label globals. + +plot_field_mesh evaluates the field at mesh vertices and renders with +matplotlib.tri.Triangulation (avoids interpolation artefacts near element edges). +""" +from pathlib import Path +import warnings +import numpy as np +from m3dc1 import eval_field # noqa: F401 + +# Mutable globals set by calling code before invoking plot functions. +shortlbl: dict = {} +label: dict = {} + +__all__ = ["plot_field_mesh", "eval_field", "shortlbl", "label"] + + +def plot_field_mesh(field, filename, time=0, coord="scalar", phi=0.0, + save=False, savedir="./", quiet=True): + """Plot a field evaluated at mesh vertices using Triangulation rendering.""" + import matplotlib.pyplot as plt # noqa: PLC0415 + import matplotlib.tri as tri # noqa: PLC0415 + import sys + import os + + fio_lib = os.path.join(os.path.dirname(os.path.dirname(__file__))) + if fio_lib not in sys.path: + sys.path.insert(0, fio_lib) + + try: + from fpy import sim_data # noqa: PLC0415 + except ImportError: + raise ImportError("fpy.sim_data not importable; check your fusion-io installation.") + + sim = sim_data(filename=str(filename), time=time) + mesh_obj = sim.get_mesh(quiet=True) + el = mesh_obj.elements + + # Columns 0-2 of mesh/elements are floating-point values, not integer node + # indices, so they cannot be used as triangle connectivity. Evaluate the + # field at all per-element (R, Z) positions, then deduplicate to obtain + # unique vertices and a valid Delaunay triangulation. + R_all = el[:, 4].astype(float) + Z_all = el[:, 5].astype(float) + phi_all = np.full_like(R_all, phi) + + f_all = eval_field(field, R_all, phi_all, Z_all, + coord=coord, sim=sim, time=time, quiet=quiet) + + _, idx = np.unique(np.column_stack([R_all, Z_all]), axis=0, return_index=True) + R_unique = R_all[idx] + Z_unique = Z_all[idx] + f_unique = f_all[idx] + + triang = tri.Triangulation(R_unique, Z_unique) + fig, ax = plt.subplots(figsize=(5, 7)) + tcf = ax.tripcolor(triang, f_unique, shading="gouraud") + plt.colorbar(tcf, ax=ax, label=label.get(field, field)) + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.set_title(f"{shortlbl.get(field, field)} t={time}") + plt.tight_layout() + if save: + fig.savefig(Path(savedir) / f"mesh_{field}_{time:03d}.png", dpi=150) + if plt.isinteractive(): + plt.show() diff --git a/use_cases/tokamak_stability/m3dc1/tests/__init__.py b/use_cases/tokamak_stability/m3dc1/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/use_cases/tokamak_stability/m3dc1/tests/test_m3dc1.py b/use_cases/tokamak_stability/m3dc1/tests/test_m3dc1.py new file mode 100644 index 0000000..e94e457 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/tests/test_m3dc1.py @@ -0,0 +1,289 @@ +"""Unit tests for the m3dc1 package. + +Tests are split into: + - Unit tests (no data, no write_neo_input): mock or analytic only. + - Integration tests (require sparc_1425 data and write_neo_input on PATH): + marked with @pytest.mark.integration. + +Run unit tests only: + pytest fusion_io/m3dc1/tests/test_m3dc1.py -m "not integration" + +Run all tests (requires data): + pytest fusion_io/m3dc1/tests/test_m3dc1.py +""" + +import numpy as np +import pytest +from pathlib import Path +from unittest.mock import MagicMock, patch + +# Fixtures +SPARC_DIR = Path("/Users/kparfrey/data/asv/run1/sparc_1425") +HAS_DATA = SPARC_DIR.is_dir() + +try: + import fio_py # noqa: F401 + HAS_FIO_PY = True +except ImportError: + HAS_FIO_PY = False + +integration = pytest.mark.skipif( + not (HAS_DATA and HAS_FIO_PY), + reason="Requires sparc_1425 data and fio_py extension (fusion-io must be built)", +) + + +# --------------------------------------------------------------------------- +# Unit tests — eval_field +# --------------------------------------------------------------------------- + +def _make_mock_sim(ftype="scalar", return_val=42.0): + """Create a minimal fpy.sim_data mock.""" + sim = MagicMock() + sim.typedict = { + "psi": ("psi", "scalar", None, "composite"), + "B": ("magnetic field", "vector", None, "simple"), + "p": ("total pressure", "scalar", None, "simple"), + } + fld = MagicMock() + fld.ftype = ftype + if ftype == "scalar": + fld.evaluate.return_value = (return_val,) + else: + fld.evaluate.return_value = (1.0, 2.0, 3.0) + sim.get_field.return_value = fld + sim.timeslice = 0 + return sim + + +def test_eval_field_scalar_shape(): + from m3dc1 import eval_field + sim = _make_mock_sim(ftype="scalar", return_val=5.0) + R = np.ones((3, 4)) + phi = np.zeros((3, 4)) + Z = np.zeros((3, 4)) + out = eval_field("psi", R, phi, Z, coord="scalar", sim=sim, quiet=True) + assert out.shape == (3, 4) + assert np.all(out == 5.0) + + +def test_eval_field_vector_shape(): + from m3dc1 import eval_field + sim = _make_mock_sim(ftype="vector") + R = np.ones(10) + phi = np.zeros(10) + Z = np.zeros(10) + out = eval_field("B", R, phi, Z, coord="vector", sim=sim, quiet=True) + assert out.shape == (3, 10) + + +def test_eval_field_vector_component_shape(): + from m3dc1 import eval_field + sim = _make_mock_sim(ftype="vector") + R = np.ones(5) + phi = np.zeros(5) + Z = np.zeros(5) + out = eval_field("B", R, phi, Z, coord="R", sim=sim, quiet=True) + assert out.shape == (5,) + assert np.all(out == 1.0) + + +def test_eval_field_nan_on_none(): + from m3dc1 import eval_field + sim = _make_mock_sim(ftype="scalar") + fld = MagicMock() + fld.ftype = "scalar" + # Simulate out-of-domain: fpy returns (None,) + fld.evaluate.return_value = (None,) + sim.get_field.return_value = fld + R = np.array([1.0]) + out = eval_field("psi", R, R * 0, R * 0, coord="scalar", sim=sim, quiet=True) + assert np.isnan(out[0]) + + +# --------------------------------------------------------------------------- +# Unit tests — _neo_input.flux_surface_average +# --------------------------------------------------------------------------- + +def test_fsa_uniform(): + """FSA of f=1 with Jac=1 over [0,2pi]x[0,2pi] should be 1.""" + from m3dc1._neo_input import flux_surface_average + ntheta, nphi = 64, 4 + theta = np.linspace(0, 2 * np.pi, ntheta, endpoint=False) + phi = np.linspace(0, 2 * np.pi, nphi, endpoint=False) + q = np.ones((ntheta, nphi)) + j = np.ones((ntheta, nphi)) + result = flux_surface_average(q, j, theta, phi) + assert abs(result - 1.0) < 1e-10 + + +def test_fsa_linear(): + """FSA of cos(theta) ≈ 0 (trapezoidal error bounded by 1/ntheta for open interval).""" + from m3dc1._neo_input import flux_surface_average + ntheta, nphi = 512, 4 + theta = np.linspace(0, 2 * np.pi, ntheta, endpoint=False) + phi = np.linspace(0, 2 * np.pi, nphi, endpoint=False) + th2d = theta[:, np.newaxis] * np.ones((1, nphi)) + q = np.cos(th2d) + j = np.ones((ntheta, nphi)) + result = flux_surface_average(q, j, theta, phi) + # Trapezoidal rule on open interval [0, 2π) has O(1/N) error for periodic f + assert abs(result) < 2.0 / ntheta + + +# --------------------------------------------------------------------------- +# Unit tests — get_time_of_slice (mocked h5py) +# --------------------------------------------------------------------------- + +def test_get_time_of_slice_alfven(): + """Returns float in Alfvén units from h5 mock.""" + from m3dc1 import get_time_of_slice + mock_h5 = MagicMock() + mock_h5.__enter__ = lambda s: s + mock_h5.__exit__ = MagicMock(return_value=False) + mock_h5.__contains__ = lambda s, k: k == "time_001" + mock_h5.__getitem__ = lambda s, k: MagicMock(attrs={"time": 500.0}) + + with patch("h5py.File", return_value=mock_h5): + t = get_time_of_slice(1, filename="C1.h5", units="alfven", quiet=True) + assert t == pytest.approx(500.0) + + +# --------------------------------------------------------------------------- +# Unit tests — get_timetrace (mocked sim) +# --------------------------------------------------------------------------- + +def test_get_timetrace_ke_alias(): + """'ke' is aliased to E_K3.""" + from m3dc1 import get_timetrace + sim = MagicMock() + time = np.linspace(0, 100, 101) + ke = np.exp(0.01 * time) + sim._all_traces = {"time": time, "E_K3": ke} + sim.filename = "C1.h5" + t, vals, label, units_str = get_timetrace("ke", sim=sim, quiet=True) + assert label == "E_K3" + assert len(vals) == len(t) + + +def test_get_timetrace_growth(): + """Growth mode: gamma of exp(2*gamma0*t) = gamma0.""" + from m3dc1 import get_timetrace + gamma0 = 0.005 + sim = MagicMock() + t_arr = np.linspace(0, 1000, 1001) + ke_arr = np.exp(2 * gamma0 * t_arr) + sim._all_traces = {"time": t_arr, "E_K3": ke_arr} + sim.filename = "C1.h5" + t, vals, label, units_str = get_timetrace( + "ke", sim=sim, growth=True, quiet=True + ) + # Should recover gamma0 (to within trapezoidal discretisation error) + assert np.allclose(vals, gamma0, rtol=1e-4) + + +# --------------------------------------------------------------------------- +# Unit tests — get_shape +# --------------------------------------------------------------------------- + +def test_get_shape_returns_dict(): + """get_shape returns dict with correct keys even on simple mock data.""" + # We can't call get_shape without a real sim and matplotlib, so just + # verify the key set when provided a working mock. + from m3dc1 import get_shape + import matplotlib + matplotlib.use("Agg") + + sim = MagicMock() + # Use plain dicts so dict["key"] access works without MagicMock magic-method subtleties + sim._all_attrs = { + "scalars": { + "psi_lcfs": np.array([0.25]), + "xmag": np.array([1.8]), + "zmag": np.array([0.0]), + } + } + + # Mesh: circle of radius 0.5 centred at R=1.8 + n = 100 + angles = np.linspace(0, 2 * np.pi, n, endpoint=False) + R_c = 1.8 + 0.5 * np.cos(angles) + Z_c = 0.0 + 0.5 * np.sin(angles) + elements = np.zeros((n, 6)) + elements[:, 4] = R_c + elements[:, 5] = Z_c + mesh_obj = MagicMock() + mesh_obj.elements = elements + sim.get_mesh.return_value = mesh_obj + sim.timeslice = 0 + + # Stub eval_field to return circular psi (centre at R=1.8, Z=0) + with patch("m3dc1.eval_field") as mock_ef: + def _psi_stub(field, R, phi, Z, **kw): + return (R - 1.8) ** 2 + Z ** 2 + mock_ef.side_effect = _psi_stub + result = get_shape(sim, res=50) + + assert set(result.keys()) == {"R0", "a", "kappa", "delta"} + + +# --------------------------------------------------------------------------- +# Integration tests +# --------------------------------------------------------------------------- + +@integration +def test_eval_field_psi_sparc1425(): + """psi increases monotonically from axis to LCFS.""" + from fpy import sim_data + from m3dc1 import eval_field + sim = sim_data(filename=str(SPARC_DIR / "C1.h5"), time=0) + R_mag = float(np.asarray(sim._all_attrs["scalars"]["xmag"])[0]) + Z_mag = float(np.asarray(sim._all_attrs["scalars"]["zmag"])[0]) + R_line = np.linspace(R_mag, R_mag + 0.5, 20) + psi = eval_field("psi", R_line, np.zeros(20), np.full(20, Z_mag), + sim=sim, time=0, quiet=True) + assert np.all(np.diff(psi[~np.isnan(psi)]) > 0) + + +@integration +def test_get_time_of_slice_mks(): + """Returns a positive finite time in seconds for sparc_1425 snapshot 1.""" + from m3dc1 import get_time_of_slice + t = get_time_of_slice(1, filename=str(SPARC_DIR / "C1.h5"), + units="mks", quiet=True) + assert np.isfinite(t) and t > 0 + + +@integration +def test_get_shape_sparc1425(): + """R0 ≈ 1.85 m, kappa > 1, 0 < delta < 1.""" + from fpy import sim_data + from m3dc1 import get_shape + sim = sim_data(filename=str(SPARC_DIR / "C1.h5"), time=0) + shape = get_shape(sim, res=100) + assert shape, "get_shape returned empty dict" + assert abs(shape["R0"] - 1.855) < 0.1, f"R0={shape['R0']}" + assert shape["kappa"] > 1, f"kappa={shape['kappa']}" + assert 0 < shape["delta"] < 1, f"delta={shape['delta']}" + + +@integration +def test_flux_average_q_sparc1425(): + """q-profile is monotonically increasing for sparc_1425 equilibrium.""" + from fpy import sim_data + from m3dc1 import flux_average + sim = sim_data(filename=str(SPARC_DIR / "C1.h5"), time=0) + psin, q = flux_average("q", sim=sim, points=30) + assert np.all(np.diff(q) > 0), f"q not monotone: {q}" + + +@integration +def test_flux_coordinates_psi_norm_sparc1425(): + """fc.psi_norm spans [0.01, 0.99] for sparc_1425.""" + from fpy import sim_data + from m3dc1 import flux_coordinates + sim = sim_data(filename=str(SPARC_DIR / "C1.h5"), time=0) + fc_obj = flux_coordinates(sim=sim, points=30) + psin = fc_obj.fc.psi_norm + assert psin[0] < 0.05 + assert psin[-1] > 0.95 diff --git a/use_cases/tokamak_stability/m3dc1/time_trace_fast.py b/use_cases/tokamak_stability/m3dc1/time_trace_fast.py new file mode 100644 index 0000000..717dda4 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/time_trace_fast.py @@ -0,0 +1,4 @@ +"""Submodule alias: `import m3dc1.time_trace_fast as ttf`.""" +from m3dc1 import plot_time_trace_fast, get_timetrace # noqa: F401 + +__all__ = ["plot_time_trace_fast", "get_timetrace"] diff --git a/use_cases/tokamak_stability/m3dc1/user_guide.md b/use_cases/tokamak_stability/m3dc1/user_guide.md new file mode 100644 index 0000000..a20c929 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1/user_guide.md @@ -0,0 +1,301 @@ +# m3dc1 Python Module — User Guide + +The `m3dc1` package provides a Python interface to M3D-C1 simulation data. +It wraps `fpy.py` (the fusion-io Python bindings) and adds flux-surface +analysis functions that rely on the `write_neo_input` C++ executable. + +--- + +## Setup + +Three environment variables are required: + +```bash +export FIO_INSTALL_DIR=/path/to/fusion-io/install +export PYTHONPATH=$FIO_INSTALL_DIR/lib:$PYTHONPATH +export DYLD_LIBRARY_PATH=$FIO_INSTALL_DIR/lib:$DYLD_LIBRARY_PATH # macOS +# export LD_LIBRARY_PATH=$FIO_INSTALL_DIR/lib:$LD_LIBRARY_PATH # Linux +``` + +`PYTHONPATH` makes `fio_py`, `fpy`, and `m3dc1` importable. `DYLD_LIBRARY_PATH` +(or `LD_LIBRARY_PATH` on Linux) lets the dynamic linker find `libm3dc1.so` and +`libfusionio.so` at runtime. `FIO_INSTALL_DIR` tells `m3dc1` where to find the +`write_neo_input` and `trace` executables; if you prefer, you can instead add +`$FIO_INSTALL_DIR/bin` to `PATH`. + +The source directory `fusion_io/` can be used in place of the install tree by +pointing `PYTHONPATH` at it directly — no build step is needed for the pure-Python +parts of `m3dc1`. + +--- + +## Import patterns + +All of the following import styles are supported: + +```python +import m3dc1 as m1 # main namespace +from m3dc1.eval_field import eval_field # submodule alias +from m3dc1.get_time_of_slice import get_time_of_slice +from m3dc1.get_timetrace import get_timetrace +import m3dc1.time_trace_fast as ttf # time trace submodule +``` + +--- + +## Opening a simulation + +`m3dc1` functions take a `fpy.sim_data` object rather than a filename. Open +one with: + +```python +from fpy import sim_data + +sim_eq = sim_data("path/to/C1.h5", time=0) # equilibrium / timeslice 0 +sim_lin = sim_data("path/to/C1.h5", time=1) # linear perturbation, slice 1 +``` + +The `time` argument sets the active timeslice. Use `sim.set_timeslice(n)` to +change it later. Useful attributes: + +| Attribute | Description | +|------------------------|------------------------------------------------| +| `sim.filename` | Absolute path to `C1.h5` | +| `sim.timeslice` | Currently active timeslice index | +| `sim.ntime` | Total number of timeslices | +| `sim.ntor` | Toroidal mode number | +| `sim._all_traces` | h5py group for `C1.h5/scalars/` (time traces) | +| `sim._all_attrs` | Open h5py file handle for direct HDF5 access | + +--- + +## Core functions + +### `eval_field` + +Evaluate a field at arbitrary (R, φ, Z) coordinates. + +```python +from m3dc1.eval_field import eval_field + +psi = eval_field("psi", R, phi, Z, sim=sim_eq, time=0) +BR = eval_field("B", R, phi, Z, coord="R", sim=sim_eq) +Bvec = eval_field("B", R, phi, Z, coord="vector", sim=sim_eq) +``` + +`R`, `phi`, `Z` are NumPy arrays (or scalars) in SI units and cylindrical +coordinates. They are broadcast against each other so mixed shapes work. + +| `coord` | Return shape | Meaning | +|------------|--------------------|----------------------------------------| +| `"scalar"` | `R.shape` | Scalar value, or magnitude for vectors | +| `"R"` | `R.shape` | R-component of a vector field | +| `"phi"` | `R.shape` | φ-component of a vector field | +| `"Z"` | `R.shape` | Z-component of a vector field | +| `"vector"` | `(3, *R.shape)` | All three cylindrical components | + +Out-of-domain points return `NaN`. Field names accepted by `sim.typedict`: +`"psi"`, `"B"`, `"p"`, `"j"`, `"v"`, `"ne"`, `"ni"`, `"te"`, `"ti"`, +`"pe"`, `"pi"`, `"A"`, `"E"`, `"kprad_rad"`. + +**Performance note:** `eval_field` calls the C extension once per grid point +in a Python loop. For large arrays this is the main bottleneck. Prefer the +pre-computed values from `flux_average` where possible. + +--- + +### `get_time_of_slice` + +Return the simulation time of a snapshot. + +```python +from m3dc1.get_time_of_slice import get_time_of_slice + +t_alfven = get_time_of_slice(1, filename="C1.h5", units="alfven") +t_s = get_time_of_slice(1, filename="C1.h5", units="mks") +``` + +| `units` | Returns | +|----------------------|-----------------------------| +| `"alfven"`, `"m3dc1"` | Alfvén times (code units) | +| `"mks"`, `"s"` | Physical seconds | + +The MKS conversion reads `n0_norm`, `l0_norm`, `b0_norm`, and `ion_mass` from +the `C1.h5` global attributes. Returns `nan` if those attributes are absent +(older file versions). + +--- + +### `get_timetrace` + +Read a scalar time trace. + +```python +from m3dc1.get_timetrace import get_timetrace + +t, ke, label, units_str = get_timetrace("ke", sim=sim_eq) +t, gamma, _, _ = get_timetrace("ke", sim=sim_eq, growth=True) +``` + +`"ke"` is an alias for `"E_K3"` (total kinetic energy). Any name present in +`C1.h5/scalars/` is valid. + +| Parameter | Effect | +|-------------|------------------------------------------------------------| +| `growth` | Return γ(t) = 0.5 · d(ln\|values\|)/dt instead of values | +| `renorm` | Divide values by the first non-zero value | +| `units` | `"m3dc1"` (Alfvén) or `"mks"` for the time axis | + +Returns `(time, values, label, units_str)`. + +--- + +## Flux-surface functions + +These functions require `write_neo_input` to be installed and findable (see +Setup above). Each call launches the executable as a subprocess, which +re-opens `C1.h5`, loads the mesh, and traces flux surfaces. **Expect a few +seconds of startup cost per call** even before the tracing begins. + +### `flux_average` + +Compute the flux-surface average ⟨f⟩ of a field as a function of ψ_norm. + +```python +psin, q_prof = m1.flux_average("q", sim=sim_eq) +psin, p_prof = m1.flux_average("p", sim=sim_eq) +psin, ne_prof = m1.flux_average("ne", sim=sim_eq) +``` + +For the fields below, `write_neo_input` pre-computes the average internally +and the result is read directly — no secondary `eval_field` loop: + +| Field arg | neo_input.nc variable | +|-----------|-----------------------| +| `"q"` | `q` | +| `"ne"` | `ne0` | +| `"te"` | `Te0` | +| `"ni"` | `ni0` | +| `"ti"` | `Ti0` | + +All other fields are evaluated on the full 3-D flux surface grid with +`eval_field` (slow for large grids) and integrated with the Jacobian. + +Key parameters: + +| Parameter | Default | Effect | +|-----------|---------|--------------------------------------------------| +| `points` | 200 | Number of radial surfaces (nr) | +| `ntheta` | 128 | Poloidal points per surface | +| `nphi` | 4 | Toroidal planes; use 1 for axisymmetric fields | + +--- + +### `get_shape` + +Compute Miller geometry parameters of the last closed flux surface. + +```python +shape = m1.get_shape(sim_eq) +# {"R0": 1.855, "a": 0.570, "kappa": 1.85, "delta": 0.42} +``` + +Uses `matplotlib.contour` on a regular psi grid — no `write_neo_input` +needed. The `res` parameter (default 250) sets the grid resolution per axis. +Returns `{}` on failure. + +| Key | Description | +|----------|----------------------------------------------| +| `"R0"` | Major radius of LCFS midpoint (m) | +| `"a"` | Minor radius (m) | +| `"kappa"`| Elongation | +| `"delta"`| Triangularity | + +--- + +### `flux_coordinates` and `eigenfunction` + +These two functions are designed to be used together. `flux_coordinates` +runs `write_neo_input` once and caches the flux surface grid; `eigenfunction` +uses that grid to decompose a perturbed field into poloidal modes. + +```python +# Step 1 — compute flux coordinates from the equilibrium (runs write_neo_input) +fc = m1.flux_coordinates(sim=sim_eq, points=100, ntheta=128) + +# fc.fc.psi_norm → normalised flux grid, shape (nr,) +# fc.fc.q → safety factor profile, shape (nr,) + +# Step 2 — decompose a perturbed field (runs eval_field loop) +spec = m1.eigenfunction( + sim=[fc, sim_lin], + field="p", coord="scalar", + fourier=True, full_fft=False, + norm_to_unity=True, +) +# spec shape: (ntheta//2 + 1, nr) — poloidal mode amplitude vs psi_norm +``` + +Reuse the same `fc` object for multiple `eigenfunction` calls on different +fields or components — `write_neo_input` only runs once. + +`flux_coordinates` parameters: + +| Parameter | Default | Effect | +|-----------|---------|-----------------------------------------------------------| +| `points` | 200 | Number of radial surfaces (nr) | +| `ntheta` | 128 | Poloidal points; sets max resolvable mode (m ≤ ntheta//2) | +| `phit` | 0.0 | Toroidal angle (radians) of the cross-section | + +`eigenfunction` parameters: + +| Parameter | Default | Effect | +|-----------------|------------|---------------------------------------------------| +| `field` | `"p"` | Field to decompose | +| `coord` | `"scalar"` | Component: `"scalar"`, `"R"`, `"phi"`, `"Z"` | +| `fourier` | `True` | Return Fourier spectrum; `False` returns raw grid | +| `full_fft` | `False` | One-sided amplitude vs full two-sided FFT | +| `norm_to_unity` | `True` | Divide spectrum by its global peak | +| `makeplot` | `False` | Show amplitude vs ψ_norm plot | + +--- + +## Performance guide + +| Function | Mechanism | Typical cost | +|------------------------------------|----------------------------|-----------------------| +| `get_time_of_slice` | h5py read | Milliseconds | +| `get_timetrace` | h5py read | Milliseconds | +| `flux_average` (q, ne, te, ni, ti) | Subprocess + nc read | Seconds | +| `flux_average` (other fields) | Subprocess + eval_field loop | Minutes | +| `get_shape` | eval_field on 2-D grid | Seconds (res=250) | +| `flux_coordinates` | Subprocess | Seconds–minutes | +| `eigenfunction` | eval_field loop (ntheta×nr) | Minutes | + +The subprocess cost is dominated by: +1. Process launch and HDF5 file re-open (~1–2 s fixed overhead) +2. Flux surface tracing: scales as `nr × ntheta` + +The `eval_field` loop cost scales as `ntheta × nr` Python→C extension calls. +Reduce both by lowering `points` and `ntheta`. For the pre-computed fields +in `flux_average`, only the subprocess cost matters. + +--- + +## Plotting functions + +All plotting functions follow the same pattern — they open a new `sim_data` +internally from a filename, evaluate the requested quantity, and display with +matplotlib. + +```python +m1.plot_field("p", "C1.h5", time=1, lcfs=True) +m1.plot_shape("C1.h5", time=0) +m1.plot_time_trace_fast("ke", "C1.h5") +m1.plot_flux_average("q", "C1.h5", time=0) +m1.plot_gfile("sparc.geqdsk") +``` + +Pass `save=True` and `savedir="./output/"` to write PNG files instead of +displaying interactively. `plot_vector_field` requires mayavi and raises +`ImportError` if it is not installed. diff --git a/use_cases/tokamak_stability/m3dc1_plots.py b/use_cases/tokamak_stability/m3dc1_plots.py new file mode 100644 index 0000000..dd5a358 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1_plots.py @@ -0,0 +1,1348 @@ +"""Plotting library for M3D-C1 simulation output. + +Each public function produces exactly one figure, saves it to a caller-specified +path, closes the figure, and returns the output Path. No plt.show() calls are +made; the module is headless-safe. + +Functions +--------- +Category A — pure matplotlib, data via m3dc1_tools.py (no m3dc1 plotting): + plot_kinetic_energy Semi-log KE vs time with optional γ annotation + plot_growth_rate_vs_time Instantaneous γτ_A(t) trace + plot_flux_average_profiles Multi-panel radial flux-averaged profiles + plot_safety_factor q(ψ_norm) with reference lines and q95 + plot_poloidal_spectrum 2-D heatmap of m-spectrum vs ψ_norm + plot_standard_spectra 4-panel spectra: p, B_R, B_Z, B_φ + plot_perturbed_field_map δf filled contour on R,Z plane + plot_geqdsk Flux contours from a GEQDSK file + plot_geqdsk_compare 3-panel GEQDSK vs C1.h5 psi comparison + plot_tpf_vs_time Total poloidal flux trace over time snapshots + +Category B — m3dc1 library wrappers (figure capture / tempdir pattern): + plot_field 2-D field contour on R,Z plane via m3.plot_field + plot_mesh Triangular mesh via m3.plot_mesh + plot_flux_surface_shape Flux surface shape via m3.plot_shape + plot_field_mesh Field on mesh via m3dc1.plot_field_mesh + plot_field_vs_phi Field vs φ at fixed R or Z + plot_flux_average_m3 Flux average via m3.plot_flux_average + plot_line Field along a line via m3.plot_line + plot_eigenfunction Eigenfunction via m3.eigenfunction + plot_poincare Poincaré plot via m3.plot_poincare + plot_signal Diagnostic signal via m3.plot_signal + plot_time_trace Time trace via m3.plot_time_trace_fast + +Convenience wrappers: + plot_stability_summary KE + growth rate + spectra + field map + plot_equilibrium_overview Profiles + q + mesh + plot_case_summary Calls both wrappers above +""" +from __future__ import annotations + +import contextlib +import math +import os +import re +import shutil +import tempfile +from pathlib import Path + +# Set a headless default before importing pyplot; respected if user has already +# set MPLBACKEND or imported matplotlib with a different backend. +os.environ.setdefault("MPLBACKEND", "Agg") + +import matplotlib.pyplot as plt +from matplotlib.colors import LogNorm +import numpy as np + +from m3dc1_tools import ( + compute_flux_average_profiles, + compute_growth_rate, + compute_ke_growth_trace, + compute_perturbed_fields, + compute_poloidal_spectrum, + compute_q95, + compute_standard_spectra, + make_evaluation_grid, + read_case_metadata, + read_mesh_vertices, +) + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + +@contextlib.contextmanager +def _in_case_dir(path: Path): + """Temporarily change CWD to path; restore on exit even if an exception is raised.""" + old = os.getcwd() + os.chdir(path) + try: + yield + finally: + os.chdir(old) + + +def _import_m3dc1(): + """Return the m3dc1 module or raise ImportError with a helpful message.""" + try: + import m3dc1 as m3 + return m3 + except ImportError as exc: + raise ImportError( + "m3dc1 Python package not found. Add its location to sys.path " + "before calling Category B plot functions." + ) from exc + + +def _save_first_new_fig(before: set, output_path: Path, dpi: int) -> None: + """Save the first figure created since `before` was recorded; close all new figures.""" + after = set(plt.get_fignums()) + new = sorted(after - before) + if not new: + raise RuntimeError("m3dc1 library call produced no new figure") + fig = plt.figure(new[0]) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + for n in new: + plt.close(plt.figure(n)) + + +def _parse_geqdsk(gfile_path: Path) -> dict: + """Parse a GEQDSK equilibrium file; return a dict of grid arrays and scalars.""" + number_re = re.compile(r"[+-]?(?:\d+\.\d*|\.\d+|\d+)(?:[Ee][+-]?\d+)?") + with open(gfile_path) as f: + header = f.readline() + parts = header.split() + if len(parts) < 2: + raise ValueError(f"Invalid GEQDSK header in {gfile_path}") + nw, nh = int(parts[-2]), int(parts[-1]) + nums = [float(x) for line in f for x in number_re.findall(line)] + + idx = 0 + + def take(n: int) -> list: + nonlocal idx + out = nums[idx: idx + n] + idx += n + return out + + rdim, zdim, _rcentr, rleft, zmid = take(5) + rmaxis, zmaxis, simag, sibry, _bcentr = take(5) + current = take(5)[0] + take(nw) # fpol + take(nw) # pres + take(nw) # ffprim + take(nw) # pprim + psirz_flat = take(nw * nh) + take(nw) # qpsi + + _nbbbs_limitr = take(2) + nbbbs, limitr = int(_nbbbs_limitr[0]), int(_nbbbs_limitr[1]) + + rbbbs, zbbbs, rlim, zlim = [], [], [], [] + if nbbbs > 0: + bdry = take(2 * nbbbs) + rbbbs, zbbbs = bdry[0::2], bdry[1::2] + if limitr > 0: + lim = take(2 * limitr) + rlim, zlim = lim[0::2], lim[1::2] + + rg = np.linspace(rleft, rleft + rdim, nw) + zg = np.linspace(zmid - zdim / 2.0, zmid + zdim / 2.0, nh) + psirz = np.reshape(psirz_flat, (nh, nw)) + psirzn = (psirz - simag) / (sibry - simag) if sibry != simag else np.zeros_like(psirz) + + return { + "nw": nw, "nh": nh, + "rg": rg, "zg": zg, + "psirz": psirz, "psirzn": psirzn, + "simag": simag, "sibry": sibry, + "rmaxis": rmaxis, "zmaxis": zmaxis, + "current": current, + "rbbbs": rbbbs, "zbbbs": zbbbs, + "rlim": rlim, "zlim": zlim, + } + + +# ============================================================================ +# CATEGORY A — PURE MATPLOTLIB +# Data is loaded and computed via m3dc1_tools.py functions. +# No m3dc1 library plotting routines are called in this section. +# ============================================================================ + +def plot_kinetic_energy( + case_dir: str | Path, + output_path: str | Path, + annotate_growth_rate: bool = True, + dpi: int = 150, +) -> Path: + """Semi-log plot of total kinetic energy vs time. + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + annotate_growth_rate: If True, annotate with mean γτ_A. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + time, ke = compute_ke_growth_trace(case_dir) + mask = ke > 0 + if mask.sum() < 2: + raise ValueError("Fewer than 2 non-zero KE points in the trace.") + + fig, ax = plt.subplots(figsize=(7, 4)) + ax.semilogy(time[mask], ke[mask], lw=1.5) + ax.set_xlabel("time (τ_A)") + ax.set_ylabel("E_K3 (normalised)") + ax.set_title(f"Kinetic Energy — {Path(case_dir).name}") + ax.grid(True, which="both", alpha=0.4) + + if annotate_growth_rate: + gamma = compute_growth_rate(case_dir) + if math.isfinite(gamma): + ax.annotate( + f"γτ_A = {gamma:.3f}", + xy=(0.05, 0.95), xycoords="axes fraction", + va="top", ha="left", fontsize=9, color="C1", + ) + + fig.tight_layout() + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_growth_rate_vs_time( + case_dir: str | Path, + output_path: str | Path, + smooth_window: int = 5, + dpi: int = 150, +) -> Path: + """Instantaneous growth rate γτ_A(t) = 0.5 · d(ln E_K3)/dt. + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + smooth_window: Running-average window width for the smoothed overlay. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + time, ke = compute_ke_growth_trace(case_dir) + mask = ke > 0 + if mask.sum() < 3: + raise ValueError("Fewer than 3 non-zero KE points; cannot compute growth rate trace.") + + t = time[mask] + k = ke[mask] + dt = np.diff(t) + dt[dt == 0] = np.nan + gamma = 0.5 * np.diff(np.log(k)) / dt + t_mid = 0.5 * (t[:-1] + t[1:]) + + w = min(smooth_window, len(gamma)) + gamma_smooth = np.convolve(gamma, np.ones(w) / w, mode="same") + + fig, ax = plt.subplots(figsize=(7, 4)) + ax.plot(t_mid, gamma, alpha=0.35, lw=0.8, color="C0") + ax.plot(t_mid, gamma_smooth, lw=1.5, color="C0", label=f"smoothed (w={w})") + ax.axhline(0, color="k", lw=0.6, ls="--") + ax.set_xlabel("time (τ_A)") + ax.set_ylabel("γτ_A") + ax.set_title(f"Instantaneous Growth Rate — {Path(case_dir).name}") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.4) + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_flux_average_profiles( + case_dir: str | Path, + output_path: str | Path, + fields: list[str] | None = None, + fcoords: str = "pest", + points: int = 200, + dpi: int = 150, +) -> Path: + """Multi-panel flux-surface-averaged radial profiles vs ψ_norm. + + Uses the equilibrium state (time=-1) from C1.h5. Requires m3dc1 + fpy. + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + fields: Field names to plot; default ``["p", "j", "ne", "q"]``. + fcoords: Flux coordinate system (``"pest"`` or ``"equal_arc"``). + points: Radial grid resolution. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + profiles = compute_flux_average_profiles( + case_dir, fields=fields, fcoords=fcoords, points=points + ) + if not profiles: + raise ValueError("No profiles computed; check case_dir and field names.") + + _labels = {"p": "Pressure", "j": "Current density", "ne": "Electron density nₑ", "q": "Safety factor q"} + n = len(profiles) + ncols = min(n, 2) + nrows = math.ceil(n / ncols) + fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows), squeeze=False) + + for ax, (name, (psin, profile)) in zip(axes.flat, profiles.items()): + ax.plot(psin, profile, lw=1.5) + ax.set_xlabel("ψ_norm") + ax.set_ylabel(_labels.get(name, name)) + ax.set_xlim(0, 1) + ax.grid(True, alpha=0.4) + + for ax in list(axes.flat)[n:]: + ax.set_visible(False) + + fig.suptitle(f"Flux-averaged profiles — {Path(case_dir).name}") + fig.tight_layout() + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_safety_factor( + case_dir: str | Path, + output_path: str | Path, + fcoords: str = "pest", + points: int = 200, + dpi: int = 150, +) -> Path: + """Safety factor q(ψ_norm) with q=1,2,3 reference lines and q95 annotation. + + Uses the equilibrium state (time=-1). Requires m3dc1 + fpy. + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + fcoords: Flux coordinate system. + points: Radial grid resolution. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + profiles = compute_flux_average_profiles( + case_dir, fields=["q"], fcoords=fcoords, points=points + ) + if "q" not in profiles: + raise ValueError("q profile not available for this case.") + + psin, q = profiles["q"] + q95 = compute_q95(psin, q) + + fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True) + ax.plot(psin, q, lw=1.8) + for q_ref in (1, 2, 3): + ax.axhline(q_ref, color="gray", lw=0.8, ls="--", alpha=0.7) + ax.text(0.99, q_ref + 0.05, f"q={q_ref}", ha="right", va="bottom", + fontsize=8, color="gray") + if math.isfinite(q95): + ax.axvline(0.95, color="C1", lw=0.8, ls=":", alpha=0.8) + ax.annotate( + f"q₉₅ = {q95:.2f}", + xy=(0.95, q95), xytext=(0.78, q95 + 0.4), + arrowprops=dict(arrowstyle="->", lw=0.8, color="C1"), + fontsize=8, color="C1", + ) + ax.set_xlabel("ψ_norm") + ax.set_ylabel("q") + ax.set_xlim(0, 1) + ax.set_title(f"Safety Factor — {Path(case_dir).name}") + ax.grid(True, alpha=0.4) + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def _norm_and_m_max( + m_modes: np.ndarray, spectrum: np.ndarray, min_amplitude: float +) -> tuple[np.ndarray, int]: + """Normalize spectrum to its peak and return (norm_spec, m_max). + + m_max is the outermost |m| where any psi_norm value exceeds min_amplitude, + giving a symmetric mode range [-m_max, m_max] that contains all visible structure. + """ + peak = float(np.nanmax(spectrum)) + norm_spec = spectrum / peak if peak > 0 else spectrum.copy() + visible = norm_spec.max(axis=1) >= min_amplitude + m_max = int(np.max(np.abs(m_modes[visible]))) if visible.any() else int(np.max(np.abs(m_modes))) + return norm_spec, m_max + + +def plot_poloidal_spectrum( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + coord: str = "scalar", + fcoords: str = "pest", + points: int = 200, + min_amplitude: float = 1e-5, + dpi: int = 150, +) -> Path: + """2-D heatmap of poloidal mode spectrum: normalized amplitude[m, ψ_norm]. + + The mode range is chosen automatically: all modes where the normalized + amplitude exceeds ``min_amplitude`` at any ψ_norm value are shown, with the + range kept symmetric around m=0. Color scale is logarithmic from + ``min_amplitude`` to 1. + + Requires m3dc1 + fpy. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index to analyse. + field: Field name, e.g. ``"p"``, ``"B"``. + output_path: Destination PNG file. + coord: Field component (``"scalar"`` for scalars; ``"R"``, ``"Z"``, ``"phi"`` for vectors). + fcoords: Flux coordinate system. + points: Radial grid resolution. + min_amplitude: Normalized amplitude threshold; modes below this at all ψ_norm + values are excluded, and it sets the colorbar minimum. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m_modes, psi_norm, spectrum = compute_poloidal_spectrum( + case_dir, time_idx, field, coord=coord, fcoords=fcoords, points=points + ) + + norm_spec, m_max = _norm_and_m_max(m_modes, spectrum, min_amplitude) + mask = (m_modes >= -m_max) & (m_modes <= m_max) + + fig, ax = plt.subplots(figsize=(7, 5)) + im = ax.pcolormesh( + psi_norm, m_modes[mask], norm_spec[mask, :], + cmap="inferno", shading="auto", + norm=LogNorm(vmin=min_amplitude, vmax=1), + ) + fig.colorbar(im, ax=ax, label="Normalized Amplitude") + ax.set_xlabel(r"$\psi_{\rm{norm}}$") + ax.set_ylabel("poloidal mode $m$") + ax.set_title(f"Poloidal Spectrum — {field} (t={time_idx}) — {Path(case_dir).name}") + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_standard_spectra( + case_dir: str | Path, + time_idx: int, + output_path: str | Path, + fcoords: str = "pest", + points: int = 200, + min_amplitude: float = 1e-5, + dpi: int = 150, +) -> Path: + """2×2 panel figure: poloidal spectra for p, B_R, B_Z, B_φ. + + All panels share the same mode range, determined by the widest visible range + across all four fields (modes where normalized amplitude >= ``min_amplitude`` + at any ψ_norm). Color scale is logarithmic from ``min_amplitude`` to 1, + independently normalized per panel. + + Requires m3dc1 + fpy. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + output_path: Destination PNG file. + fcoords: Flux coordinate system. + points: Radial grid resolution. + min_amplitude: Normalized amplitude threshold; modes below this at all ψ_norm + values are excluded, and it sets the colorbar minimum for all panels. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + spectra = compute_standard_spectra(case_dir, time_idx, fcoords=fcoords, points=points) + if not spectra: + raise ValueError("No spectra computed; check case_dir and time_idx.") + + _panel_labels = {"p": r"Pressure ($p$)", "br": r"$B_R$", "bz": r"$B_Z$", "bphi": r"$B_\phi$"} + keys = [k for k in ("p", "br", "bz", "bphi") if k in spectra] + n = len(keys) + ncols = min(n, 2) + nrows = math.ceil(n / ncols) + + # First pass: normalize each panel and find the widest visible mode range. + prepped: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {} + m_max_global = 0 + for key in keys: + m_modes, psi_norm, spectrum = spectra[key] + norm_spec, m_max = _norm_and_m_max(m_modes, spectrum, min_amplitude) + m_max_global = max(m_max_global, m_max) + prepped[key] = (m_modes, psi_norm, norm_spec) + + # Second pass: draw all panels with the shared mode range. + fig, axes = plt.subplots(nrows, ncols, figsize=(6 * ncols, 4 * nrows), squeeze=False) + + for ax, key in zip(axes.flat, keys): + m_modes, psi_norm, norm_spec = prepped[key] + mask = (m_modes >= -m_max_global) & (m_modes <= m_max_global) + im = ax.pcolormesh( + psi_norm, m_modes[mask], norm_spec[mask, :], + cmap="inferno", shading="auto", + norm=LogNorm(vmin=min_amplitude, vmax=1), + ) + fig.colorbar(im, ax=ax, label="Normalized Amplitude") + ax.set_xlabel(r"$\psi_{\rm{norm}}$") + ax.set_ylabel("poloidal mode $m$") + ax.set_title(_panel_labels.get(key, key)) + + for ax in list(axes.flat)[n:]: + ax.set_visible(False) + + fig.suptitle(f"Poloidal Spectra — t={time_idx} — {Path(case_dir).name}") + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_perturbed_field_map( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + mode: str = "grid", + grid_res: int = 200, + phi: float = 0.0, + dpi: int = 150, +) -> Path: + """Filled contour of a perturbed field δf on the R,Z plane. + + Requires m3dc1 + fpy. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index for the perturbed state. + field: Scalar field name, e.g. ``"psi"``, ``"p"``, ``"ne"``. + output_path: Destination PNG file. + mode: Evaluation grid: ``"grid"`` (regular Cartesian) or ``"mesh"`` (mesh vertices). + grid_res: Points per axis when ``mode="grid"``. + phi: Toroidal angle in radians. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + case_dir = Path(case_dir) + R_verts, Z_verts = read_mesh_vertices(case_dir / "C1.h5") + R, Z, phi_arr = make_evaluation_grid(R_verts, Z_verts, mode=mode, grid_res=grid_res, phi=phi) + fields_data = compute_perturbed_fields(case_dir, time_idx, R, Z, phi_arr, fields=[field]) + + if field not in fields_data: + raise ValueError(f"Field '{field}' not available. Got: {list(fields_data.keys())}") + + delta_f = fields_data[field] + vmax = float(np.nanquantile(np.abs(delta_f), 0.99)) or 1.0 + + fig, ax = plt.subplots(figsize=(5, 7)) + im = ax.pcolormesh(R, Z, delta_f, cmap="RdBu_r", shading="auto", vmin=-vmax, vmax=vmax) + fig.colorbar(im, ax=ax, label=f"δ{field}") + ax.set_aspect("equal") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_title(f"Perturbed {field} — t={time_idx}, φ={phi:.2f} rad\n{case_dir.name}") + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_geqdsk( + gfile_path: str | Path, + output_path: str | Path, + dpi: int = 150, +) -> Path: + """Flux surface contours and boundary from a GEQDSK equilibrium file. + + No m3dc1 or fpy dependency. + + Args: + gfile_path: Path to the GEQDSK file (typically named ``geqdsk``). + output_path: Destination PNG file. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + geo = _parse_geqdsk(Path(gfile_path)) + psirzn = geo["psirzn"] + + fig, ax = plt.subplots(figsize=(5, 7)) + ax.contour(geo["rg"], geo["zg"], psirzn, + levels=np.linspace(0.1, 0.9, 9), linewidths=0.7, colors="C0") + out_max = float(np.nanmax(psirzn)) + if out_max > 1.05: + ax.contour(geo["rg"], geo["zg"], psirzn, + levels=np.linspace(1.05, out_max, 20), linewidths=0.7, colors="C0") + ax.plot(geo["rmaxis"], geo["zmaxis"], lw=0, marker="+", markersize=10, color="C0") + if geo["rbbbs"]: + ax.plot(geo["rbbbs"], geo["zbbbs"], color="m", lw=1.2, label="LCFS") + if geo["rlim"]: + ax.plot(geo["rlim"], geo["zlim"], color="C1", lw=1.2, label="limiter") + ax.set_aspect("equal") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_title(f"{Path(gfile_path).name} (I = {geo['current']:.3e} A)") + ax.grid(True, alpha=0.4) + if geo["rbbbs"] or geo["rlim"]: + ax.legend(fontsize=8) + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_geqdsk_compare( + case_dir: str | Path, + output_path: str | Path, + gfile: str = "geqdsk", + dpi: int = 150, +) -> Path: + """3-panel comparison: GEQDSK ψ_norm, C1.h5 ψ_norm, and their difference. + + Requires m3dc1 + fpy (evaluates ψ on the GEQDSK grid from C1.h5). + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + gfile: GEQDSK filename relative to case_dir (default ``"geqdsk"``). + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + import fpy + m3 = _import_m3dc1() + + case_dir = Path(case_dir) + gfile_path = case_dir / gfile + if not gfile_path.exists(): + raise FileNotFoundError(f"GEQDSK file not found: {gfile_path}") + + geo = _parse_geqdsk(gfile_path) + rg, zg, psirzn_geq = geo["rg"], geo["zg"], geo["psirzn"] + + c1h5 = str(case_dir / "C1.h5") + with _in_case_dir(case_dir): + sim = fpy.sim_data(c1h5, time=0) + psi_axis = float(sim.get_time_trace("psimin").values[0]) + psi_lcfs = float(sim.get_time_trace("psi_lcfs").values[0]) + + R, Z = np.meshgrid(rg, zg) + phi_grid = np.zeros_like(R) + with _in_case_dir(case_dir): + psi_c1 = m3.eval_field("psi", R, phi_grid, Z, coord="scalar", + sim=sim, time=0, quiet=True) + if psi_lcfs != psi_axis: + psin_c1 = (psi_c1 - psi_axis) / (psi_lcfs - psi_axis) + else: + psin_c1 = np.zeros_like(psi_c1) + + diff = psin_c1 - psirzn_geq + levels = np.linspace(0.05, 1.2, 18) + + fig, axs = plt.subplots(1, 3, figsize=(14, 6), sharey=True) + axs[0].contour(rg, zg, psirzn_geq, levels=levels, colors="r", linewidths=0.7) + axs[0].contour(rg, zg, psirzn_geq, levels=[1.0], colors="k", linewidths=2) + axs[0].set_title("GEQDSK ψ_norm") + + axs[1].contour(rg, zg, psin_c1, levels=levels, colors="C0", linewidths=0.7) + axs[1].contour(rg, zg, psin_c1, levels=[1.0], colors="k", linewidths=2) + axs[1].set_title("C1.h5 ψ_norm") + + im = axs[2].contourf(rg, zg, diff, levels=30, cmap="coolwarm") + axs[2].contour(rg, zg, psin_c1, levels=[1.0], colors="k", linewidths=1.5) + axs[2].set_title("ψ_norm difference (C1 − GEQDSK)") + fig.colorbar(im, ax=axs[2], shrink=0.8) + + for ax in axs: + ax.set_aspect("equal") + ax.set_xlabel("R (m)") + ax.grid(True, alpha=0.3) + axs[0].set_ylabel("Z (m)") + fig.suptitle(f"GEQDSK comparison — {case_dir.name}") + fig.tight_layout() + + output_path = Path(output_path) + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +def plot_tpf_vs_time( + case_dir: str | Path, + field: str, + output_path: str | Path, + ts_list: list[int] | None = None, + units: str = "mks", + points: int = 250, + dpi: int = 150, +) -> Path: + """Total poloidal flux of a field as a function of time snapshot. + + Iterates over time snapshots, calling m3.tpf for each. Requires m3dc1 + fpy. + A sidecar .dat file with columns (ts, time, tpf) is written alongside output_path. + + Args: + case_dir: M3D-C1 case directory. + field: Field name, e.g. ``"prad"``. + output_path: Destination PNG file. + ts_list: Time-slice indices to include; None means all available snapshots. + units: ``"mks"`` or ``"alfven"`` — affects axis label only. + points: Spatial resolution passed to m3.tpf. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + import fpy + m3 = _import_m3dc1() + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + + if ts_list is None: + from m3dc1_tools import list_time_snapshots + ts_list = list_time_snapshots(case_dir) + if not ts_list: + raise ValueError("No time snapshots found.") + + time_vals, tpf_vals = [], [] + with _in_case_dir(case_dir): + for ts in ts_list: + try: + sim = fpy.sim_data(filename=c1h5, time=ts) + t, pf = m3.tpf(field=field, sim=sim, filename=c1h5, units=units, res=points) + time_vals.append(t) + tpf_vals.append(pf) + except Exception as exc: + print(f"WARNING: tpf failed for ts={ts}: {exc}") + + if not time_vals: + raise RuntimeError("TPF computation failed for all time slices.") + + output_path = Path(output_path) + dat_path = output_path.with_suffix(".dat") + with open(dat_path, "w") as f: + f.write("# ts time tpf\n") + for ts, t, pf in zip(ts_list, time_vals, tpf_vals): + f.write(f"{ts} {t} {pf}\n") + + time_label = "time (s)" if units == "mks" else "time (τ_A)" + fig, ax = plt.subplots(figsize=(7, 4)) + ax.plot(time_vals, tpf_vals, lw=1.5) + ax.set_xlabel(time_label) + ax.set_ylabel(f"TPF ({field})") + ax.set_title(f"Total Poloidal Flux — {field} — {case_dir.name}") + ax.grid(True, alpha=0.4) + fig.tight_layout() + + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path + + +# ============================================================================ +# CATEGORY B — M3DC1 LIBRARY WRAPPERS +# These functions call m3dc1's own plotting routines and redirect their output +# to a caller-specified path using either the figure-capture pattern or a +# temporary directory. All require m3dc1 (and usually fpy) to be importable. +# ============================================================================ + +def plot_field( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + coord: str = "scalar", + phi: float = 0.0, + points: int = 250, + tor_av: int = 1, + units: str = "mks", + cmap: str = "inferno", + mesh: bool = False, + bound: bool = False, + lcfs: bool = False, + dpi: int = 150, +) -> Path: + """2-D field contour on the R,Z plane via m3dc1.plot_field. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + field: Field name, e.g. ``"psi"``, ``"te"``, ``"ne"``, ``"jphi"``. + output_path: Destination PNG file. + coord: Field component (``"scalar"``, ``"R"``, ``"Z"``, ``"phi"``). + phi: Toroidal angle for the slice (radians). + points: R/Z grid resolution. + tor_av: Toroidal average over N planes. + units: ``"mks"`` or ``"alfven"``. + mesh: Overlay mesh triangulation. + bound: Overlay boundary only (no full mesh). + lcfs: Overlay last closed flux surface. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_field( + field=field, + filename=c1h5, + time=time_idx, + coord=coord, + phi=phi, + points=points, + tor_av=tor_av, + units=units, + cmap=cmap, + mesh=mesh, + bound=bound, + lcfs=lcfs, + save=False, + ) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_mesh( + case_dir: str | Path, + output_path: str | Path, + boundary: bool = False, + dpi: int = 150, +) -> Path: + """Triangular mesh visualisation via m3dc1.plot_mesh. + + Args: + case_dir: M3D-C1 case directory. + output_path: Destination PNG file. + boundary: If True, show boundary outline only (no interior mesh). + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_mesh(filename=c1h5, save=False, boundary=boundary) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_flux_surface_shape( + case_dir: str | Path, + time_idx: int, + output_path: str | Path, + points: int = 200, + dpi: int = 150, +) -> Path: + """Flux surface shape on the R,Z plane via m3dc1.plot_shape. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + output_path: Destination PNG file. + points: Grid resolution for psi evaluation. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_shape(filename=c1h5, time=time_idx, points=points) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_field_mesh( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + coord: str = "scalar", + phi: float = 0.0, + dpi: int = 150, +) -> Path: + """Field plotted on the mesh triangulation via m3dc1.plot_field_mesh. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + field: Field name. + output_path: Destination PNG file. + coord: Field component. + phi: Toroidal angle in radians. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + import importlib + _import_m3dc1() + pfm = importlib.import_module("m3dc1.plot_field_mesh") + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + pfm.plot_field_mesh(field=field, filename=c1h5, time=time_idx, coord=coord, phi=phi) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_field_vs_phi( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + R: float | None = None, + Z: float | None = None, + coord: str = "scalar", + phi_res: int = 64, + dpi: int = 150, +) -> Path: + """Field vs toroidal angle φ at a fixed (R, Z) point. + + R and Z default to the magnetic axis when not supplied. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + field: Field name. + output_path: Destination PNG file. + R: Radial coordinate in metres (None → magnetic axis R). + Z: Vertical coordinate in metres (None → magnetic axis Z). + coord: Field component. + phi_res: Number of toroidal angle points. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_field_vs_phi( + field=field, + filename=c1h5, + time=time_idx, + R=R, + Z=Z, + coord=coord, + phi_res=phi_res, + save=False, + ) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_flux_average_m3( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + fcoords: str = "pest", + points: int = 200, + dpi: int = 150, +) -> Path: + """Flux-averaged profile via m3dc1.plot_flux_average. + + The ``_m3`` suffix distinguishes this from the pure-matplotlib + ``plot_flux_average_profiles`` in Category A. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + field: Field name, e.g. ``"p"``. + output_path: Destination PNG file. + fcoords: Flux coordinate system. + points: Radial grid resolution. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_flux_average(field=field, filename=c1h5, time=time_idx, fcoords=fcoords, points=points) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_line( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + coord: str = "scalar", + angle: float = 0.0, + zoff: float = 0.0, + dist_from_magax: bool = False, + dpi: int = 150, +) -> Path: + """Field profile along a radial line at fixed phi=0 via m3dc1.plot_line. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index. + field: Field name. + output_path: Destination PNG file. + coord: Field component. + angle: Poloidal angle of the line in degrees. + zoff: Z offset of the line origin (metres). + dist_from_magax: If True, use distance from magnetic axis as x-axis. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_line( + field=field, + coord=coord, + angle=angle, + Zoff=zoff, + dist_from_magax=dist_from_magax, + filename=c1h5, + time=time_idx, + ) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_eigenfunction( + case_dir: str | Path, + time_idx: int, + field: str, + output_path: str | Path, + coord: str = "scalar", + phit: float = 0.0, + fcoords: str = "pest", + points: int = 200, + dpi: int = 150, +) -> Path: + """Eigenfunction and poloidal spectrum via m3dc1.eigenfunction. + + Builds the [flux_coords_result, linear_sim] pair that m3dc1.eigenfunction + requires: runs flux_coordinates on the equilibrium (time=0) and opens a + second sim object for the requested time snapshot. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index for the perturbed state. + field: Field name, e.g. ``"p"``. + output_path: Destination PNG file. + coord: Field component. + phit: Toroidal angle in radians for the poloidal cross-section. + fcoords: Flux coordinate system. + points: Radial grid resolution. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + with _in_case_dir(case_dir): + eq_sim = m3._require_sim(c1h5, 0) + fc_result = m3.flux_coordinates(sim=eq_sim, fcoords=fcoords, phit=phit, points=points) + lin_sim = m3._require_sim(c1h5, time_idx) + m3.eigenfunction( + sim=[fc_result, lin_sim], + field=field, + coord=coord, + fcoords=fcoords, + points=points, + makeplot=True, + ) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_poincare( + case_dir: str | Path, + time_idx: int, + output_path: str | Path, + dpi: int = 150, +) -> Path: + """Poincaré section via m3dc1.plot_poincare. + + Poincaré data must already exist in case_dir (generated by m3.run_trace or + equivalent). m3dc1 saves the figure to a directory internally; this + function redirects it to a temporary directory then moves the result to + output_path. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index to plot. + output_path: Destination PNG file. + dpi: Not applied (m3dc1 controls its own dpi). + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + output_path = Path(output_path) + + with tempfile.TemporaryDirectory() as tmpdir: + tmpdir_path = Path(tmpdir) + with _in_case_dir(case_dir): + m3.plot_poincare( + time=time_idx, + savefig=True, + savepath=str(tmpdir_path) + os.sep, + ) + plt.close("all") + new_files = sorted(tmpdir_path.iterdir()) + if not new_files: + raise RuntimeError("m3.plot_poincare produced no output file.") + shutil.copy2(new_files[0], output_path) + + return output_path + + +def plot_signal( + case_dir: str | Path, + signal: str, + output_path: str | Path, + pspec: bool = False, + pts_per_probe: int = 1, + dpi: int = 150, +) -> Path: + """Diagnostic signal time traces via m3dc1.plot_signal. + + Args: + case_dir: M3D-C1 case directory. + signal: Signal type, e.g. ``"mag_probes"``. + output_path: Destination PNG file. + pspec: If True, also plot the power spectrum. + pts_per_probe: Points per probe. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_signal( + signal=signal, + filename=c1h5, + pspec=pspec, + pts_per_probe=pts_per_probe, + ) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +def plot_time_trace( + case_dir: str | Path, + trace: str, + output_path: str | Path, + units: str = "mks", + dpi: int = 150, +) -> Path: + """Global scalar time trace via m3dc1.plot_time_trace_fast. + + Args: + case_dir: M3D-C1 case directory. + trace: Trace name, e.g. ``"ke"``. + output_path: Destination PNG file. + units: ``"mks"`` or ``"alfven"``. + dpi: Figure resolution. + + Returns: + Path to the saved figure. + """ + m3 = _import_m3dc1() + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + output_path = Path(output_path) + + before = set(plt.get_fignums()) + m3.plot_time_trace_fast(trace=trace, filename=c1h5, units=units, save=False) + _save_first_new_fig(before, output_path, dpi) + return output_path + + +# ============================================================================ +# CONVENIENCE WRAPPERS +# These call individual functions above to produce standard sets of figures. +# Each returns a list of Paths for all files created. +# ============================================================================ + +def plot_stability_summary( + case_dir: str | Path, + time_idx: int, + output_dir: str | Path, + prefix: str = "", + dpi: int = 150, +) -> list[Path]: + """Standard linear stability figure set for one time snapshot. + + Produces: KE trace, growth rate trace, standard spectra, p spectrum, δpsi map. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index for spectral and field-map figures. + output_dir: Directory where figures are written. + prefix: Optional filename prefix for all output files. + dpi: Figure resolution. + + Returns: + List of Paths created (skips any figure that raises an exception). + """ + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + results = [] + + for name, func, kwargs in [ + ("ke.png", plot_kinetic_energy, {"case_dir": case_dir}), + ("growth_rate.png", plot_growth_rate_vs_time, {"case_dir": case_dir}), + ]: + try: + p = output_dir / f"{prefix}{name}" + results.append(func(output_path=p, dpi=dpi, **kwargs)) + except Exception as exc: + print(f"WARNING: {name} failed: {exc}") + + for name, func, kwargs in [ + ("spectra.png", plot_standard_spectra, {"case_dir": case_dir, "time_idx": time_idx}), + ("spectrum_p.png", plot_poloidal_spectrum, {"case_dir": case_dir, "time_idx": time_idx, "field": "p"}), + ("field_psi.png", plot_perturbed_field_map, {"case_dir": case_dir, "time_idx": time_idx, "field": "psi"}), + ]: + try: + p = output_dir / f"{prefix}{name}" + results.append(func(output_path=p, dpi=dpi, **kwargs)) + except Exception as exc: + print(f"WARNING: {name} failed: {exc}") + + return results + + +def plot_equilibrium_overview( + case_dir: str | Path, + output_dir: str | Path, + prefix: str = "", + dpi: int = 150, +) -> list[Path]: + """Equilibrium figure set: flux-averaged profiles, q profile, and mesh. + + Args: + case_dir: M3D-C1 case directory. + output_dir: Directory where figures are written. + prefix: Optional filename prefix. + dpi: Figure resolution. + + Returns: + List of Paths created. + """ + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + results = [] + + for name, func in [ + ("profiles.png", plot_flux_average_profiles), + ("q_profile.png", plot_safety_factor), + ("mesh.png", plot_mesh), + ]: + try: + p = output_dir / f"{prefix}{name}" + results.append(func(case_dir=case_dir, output_path=p, dpi=dpi)) + except Exception as exc: + print(f"WARNING: {name} failed: {exc}") + + return results + + +def plot_case_summary( + case_dir: str | Path, + time_idx: int, + output_dir: str | Path, + prefix: str = "", + dpi: int = 150, +) -> list[Path]: + """Complete figure set: stability summary + equilibrium overview. + + Args: + case_dir: M3D-C1 case directory. + time_idx: Snapshot index for spectral and field-map figures. + output_dir: Directory where figures are written. + prefix: Optional filename prefix. + dpi: Figure resolution. + + Returns: + Combined list of Paths created by both wrappers. + """ + results = plot_stability_summary(case_dir, time_idx, output_dir, prefix=prefix, dpi=dpi) + results += plot_equilibrium_overview(case_dir, output_dir, prefix=prefix, dpi=dpi) + return results diff --git a/use_cases/tokamak_stability/m3dc1_tools.py b/use_cases/tokamak_stability/m3dc1_tools.py new file mode 100644 index 0000000..ed0c579 --- /dev/null +++ b/use_cases/tokamak_stability/m3dc1_tools.py @@ -0,0 +1,874 @@ +"""Standalone post-processing tools for M3D-C1 simulation output. + +Each function does one computation and returns plain Python / NumPy data. +None of the public functions retain file handles or simulation objects between +calls, making them safe to use from an agentic pipeline. + +Functions that require the ``m3dc1`` / ``fpy`` libraries are marked in the +docstring. All other functions depend only on ``h5py`` and ``numpy``. + +Functions +--------- +Inspection (h5py only): + read_c1input Parse C1input namelist + list_time_snapshots Discover time_NNN.h5 snapshot files + read_snapshot_time Physical time of a snapshot + read_scalar_traces Global time-trace scalars from C1.h5 + read_case_metadata Combined summary of parameters and geometry + +Mesh / evaluation grid (h5py / numpy only): + read_mesh_vertices Unique (R, Z) vertex positions + make_evaluation_grid Build (R, Z, phi) arrays for eval_field() + +Growth rate (h5py / numpy only): + compute_ke_growth_trace Kinetic energy vs time from scalars + compute_growth_rate Mean linear growth rate (1/tau_A) + +Equilibrium profiles (requires m3dc1 + fpy): + compute_flux_average_profiles Radial flux-surface-averaged profiles + compute_q95 Safety factor at psi_norm = 0.95 + compute_miller_geometry R0, a, kappa, delta of the LCFS + +Perturbed fields (requires m3dc1 + fpy): + compute_perturbed_fields Perturbed fields at arbitrary (R, Z, phi) points + +Field evaluation on a grid (requires m3dc1 + fpy): + evaluate_field_on_grid One or more fields evaluated on the mesh-vertex or Cartesian grid + +Spectral analysis (requires m3dc1 + fpy): + compute_poloidal_spectrum Poloidal m-spectrum of a single field + compute_standard_spectra Spectra for p, B_R, B_Z, B_phi +""" +from __future__ import annotations + +import contextlib +import os +import warnings +from pathlib import Path + +import h5py +import numpy as np + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + +@contextlib.contextmanager +def _in_case_dir(case_dir: Path): + """Temporarily change the working directory to ``case_dir``. + + The m3dc1 library reads auxiliary files (geqdsk, equilibrium data) using + relative paths from the working directory. This context manager provides + that guarantee for m3dc1-dependent functions while leaving the global cwd + restored on exit, even if an exception is raised. + """ + old = os.getcwd() + os.chdir(case_dir) + try: + yield + finally: + os.chdir(old) + + +# --------------------------------------------------------------------------- +# Inspection tools (no m3dc1 dependency) +# --------------------------------------------------------------------------- + +def read_c1input(case_dir: str | Path) -> dict: + """Parse the C1input namelist and return simulation parameters as a dict. + + Reads the Fortran namelist file ``C1input`` in the case directory and + extracts the parameters most relevant to post-processing. Any parameter + absent from the file is returned with value ``None``. + + Args: + case_dir: Path to the M3D-C1 case directory containing ``C1input``. + + Returns: + Dict with keys: ``"ntor"`` (int), ``"pscale"`` (float), + ``"batemanscale"`` (float), ``"dt"`` (float), ``"ntimemax"`` (int), + ``"ntimepr"`` (int), ``"linear"`` (int), ``"numvar"`` (int), + ``"ion_mass"`` (float), ``"zeff"`` (float). + """ + int_keys = {"ntor", "ntimemax", "ntimepr", "linear", "numvar"} + float_keys = {"pscale", "batemanscale", "dt", "ion_mass", "zeff"} + all_keys = int_keys | float_keys + params: dict = {k: None for k in all_keys} + + c1input = Path(case_dir) / "C1input" + if not c1input.exists(): + return params + + with c1input.open("r") as fh: + for line in fh: + line = line.split("!")[0].strip() + if not line or "=" not in line: + continue + key_part, _, val_part = line.partition("=") + key = key_part.strip().lower() + val_str = val_part.strip().split()[0].rstrip(",") if val_part.strip() else None + if key not in all_keys or val_str is None: + continue + try: + if key in int_keys: + params[key] = int(float(val_str)) + else: + params[key] = float(val_str) + except ValueError: + pass + + return params + + +def list_time_snapshots(case_dir: str | Path) -> list[int]: + """Return sorted list of time snapshot indices present in the case directory. + + Scans for files matching the pattern ``time_NNN.h5`` where NNN is an + integer. Does not open the files. + + Args: + case_dir: Path to the M3D-C1 case directory. + + Returns: + Sorted list of integer indices, e.g. ``[0, 1]`` for ``time_000.h5`` + and ``time_001.h5``. Empty list if no matching files are found. + """ + indices = [] + for path in sorted(Path(case_dir).glob("time_*.h5")): + stem = path.stem # e.g. "time_001" + parts = stem.split("_", maxsplit=1) + if len(parts) == 2 and parts[1].isdigit(): + indices.append(int(parts[1])) + return sorted(set(indices)) + + +def read_snapshot_time( + case_dir: str | Path, + time_idx: int, + units: str = "alfven", +) -> float: + """Return the simulation time corresponding to a snapshot index. + + Reads the ``time`` attribute stored directly on the snapshot group in + ``time_NNN.h5`` (or the equivalent group inside ``C1.h5``). No m3dc1 + library needed for Alfvén-time results. + + Args: + case_dir: Path to the M3D-C1 case directory. + time_idx: Snapshot index (the NNN in ``time_NNN.h5``). + units: ``"alfven"`` (default) — return simulation time in Alfvén + times directly from the file attribute. ``"s"`` or + ``"mks"`` — return physical seconds via + ``m3dc1.get_time_of_slice``; falls back to Alfvén times + with a warning if m3dc1 is not installed. + + Returns: + Simulation time as a float. Returns ``nan`` if the snapshot is not + found. + """ + case_dir = Path(case_dir) + snap_file = case_dir / f"time_{time_idx:03d}.h5" + + alfven_time = float("nan") + if snap_file.exists(): + with h5py.File(snap_file, "r") as f: + alfven_time = float(f.attrs.get("time", float("nan"))) + else: + c1h5 = case_dir / "C1.h5" + group_name = f"time_{time_idx:03d}" + if c1h5.exists(): + with h5py.File(c1h5, "r") as f: + if group_name in f: + alfven_time = float(f[group_name].attrs.get("time", float("nan"))) + + if units in ("s", "mks"): + try: + from m3dc1 import get_time_of_slice + return float(get_time_of_slice( + time_idx, + filename=str(case_dir / "C1.h5"), + units="mks", + quiet=True, + )) + except ImportError: + warnings.warn( + f"m3dc1 not installed; returning Alfvén time {alfven_time:.4g} " + f"for time_idx={time_idx}", + RuntimeWarning, + stacklevel=2, + ) + + return alfven_time + + +def read_scalar_traces( + case_dir: str | Path, + names: list[str] | None = None, +) -> dict[str, np.ndarray]: + """Read global time-trace scalars from ``C1.h5/scalars/``. + + Each scalar is a 1-D array of length ``ntimestep + 1`` (one entry per + recorded timestep, including the initial state at t = 0). The ``"time"`` + scalar gives the simulation time in Alfvén times. + + Args: + case_dir: Path to the M3D-C1 case directory. + names: List of scalar names to read, e.g. ``["E_K3", "time"]``. + If ``None``, all available scalars are returned. + + Returns: + Dict mapping scalar name to a 1-D NumPy array. + """ + c1h5 = Path(case_dir) / "C1.h5" + result: dict[str, np.ndarray] = {} + with h5py.File(c1h5, "r") as f: + scalars = f["scalars"] + keys = list(names) if names is not None else list(scalars.keys()) + for k in keys: + if k in scalars: + result[k] = scalars[k][:] + return result + + +def read_case_metadata(case_dir: str | Path) -> dict: + """Return a structured summary of a case's parameters and equilibrium state. + + Combines ``read_c1input``, ``list_time_snapshots``, and key equilibrium + geometry scalars from ``C1.h5`` into a single dict. Useful as the first + call when an agent encounters an unfamiliar case. + + Args: + case_dir: Path to the M3D-C1 case directory. + + Returns: + Dict with keys: + + ``"params"`` + Output of :func:`read_c1input`. + ``"snapshots"`` + Output of :func:`list_time_snapshots`. + ``"final_time"`` + Simulation time of the last snapshot in Alfvén times (``nan`` if + no snapshots found). + ``"R_mag"`` + R coordinate of the magnetic axis (metres). + ``"Z_mag"`` + Z coordinate of the magnetic axis (metres). + ``"R_xpoint"`` + R coordinate of the primary X-point (metres). + ``"Z_xpoint"`` + Z coordinate of the primary X-point (metres). + ``"psi_min"`` + ψ at the magnetic axis (internal units). + ``"psi_lcfs"`` + ψ at the last closed flux surface (internal units). + """ + case_dir = Path(case_dir) + result: dict = {} + result["params"] = read_c1input(case_dir) + result["snapshots"] = list_time_snapshots(case_dir) + + if result["snapshots"]: + result["final_time"] = read_snapshot_time(case_dir, result["snapshots"][-1]) + else: + result["final_time"] = float("nan") + + c1h5 = case_dir / "C1.h5" + if c1h5.exists(): + with h5py.File(c1h5, "r") as f: + sc = f["scalars"] + result["R_mag"] = float(sc["xmag"][0]) + result["Z_mag"] = float(sc["zmag"][0]) + result["R_xpoint"] = float(sc["xnull"][0]) + result["Z_xpoint"] = float(sc["znull"][0]) + result["psi_min"] = float(sc["psimin"][0]) + result["psi_lcfs"] = float(sc["psi_lcfs"][0]) + else: + for key in ("R_mag", "Z_mag", "R_xpoint", "Z_xpoint", "psi_min", "psi_lcfs"): + result[key] = float("nan") + + return result + + +# --------------------------------------------------------------------------- +# Mesh and evaluation grid (no m3dc1 dependency) +# --------------------------------------------------------------------------- + +def read_mesh_vertices(c1h5_path: str | Path) -> tuple[np.ndarray, np.ndarray]: + """Extract unique mesh vertex (R, Z) coordinates from an M3D-C1 HDF5 file. + + M3D-C1 stores the reference vertex (R, Z) of each triangular finite element + in columns 4 and 5 of ``mesh/elements``. Deduplicating these gives the set + of unique vertex positions in the poloidal cross-section. + + Args: + c1h5_path: Path to a ``C1.h5``, ``equilibrium.h5``, or + ``time_NNN.h5`` file. + + Returns: + ``(R, Z)`` — a pair of 1-D float32 arrays, each of length + ``n_unique_vertices``, sorted lexicographically by (R, Z). + """ + path = Path(c1h5_path) + with h5py.File(path, "r") as f: + if "equilibrium/mesh/elements" in f: + elements = f["equilibrium/mesh/elements"][:] + elif "mesh/elements" in f: + elements = f["mesh/elements"][:] + else: + raise KeyError(f"No mesh/elements group found in {c1h5_path}") + + rz = np.unique(np.c_[elements[:, 4], elements[:, 5]], axis=0) + return rz[:, 0].astype(np.float32), rz[:, 1].astype(np.float32) + + +def make_evaluation_grid( + R_mesh: np.ndarray, + Z_mesh: np.ndarray, + mode: str = "mesh", + grid_res: int = 200, + phi: float = 0.0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build (R, Z, phi) evaluation point arrays for field interpolation. + + The returned arrays are suitable for passing directly to ``eval_field()`` + from the ``m3dc1`` library. Note that ``eval_field`` uses the argument + order ``(field_name, R, phi, Z, ...)``, so pass ``phi_arr`` before ``Z``. + + Args: + R_mesh: 1-D array of mesh vertex R coordinates (from + :func:`read_mesh_vertices`). + Z_mesh: 1-D array of mesh vertex Z coordinates. + mode: ``"mesh"`` — use mesh vertices directly; R and Z are 1-D. + ``"grid"`` — build a regular ``grid_res × grid_res`` + Cartesian grid spanning the mesh bounding box; R and Z are + 2-D. + grid_res: Number of points per axis for ``mode="grid"``. + phi: Toroidal angle in radians applied to all evaluation points. + + Returns: + ``(R, Z, phi_arr)`` where all three arrays have the same shape: + + - ``mode="mesh"``: 1-D, length ``n_vertices``. + - ``mode="grid"``: 2-D, shape ``(grid_res, grid_res)``. + """ + if mode == "grid": + r_lin = np.linspace(float(np.nanmin(R_mesh)), float(np.nanmax(R_mesh)), grid_res) + z_lin = np.linspace(float(np.nanmin(Z_mesh)), float(np.nanmax(Z_mesh)), grid_res) + R, Z = np.meshgrid(r_lin, z_lin) + elif mode == "mesh": + R = np.asarray(R_mesh) + Z = np.asarray(Z_mesh) + else: + raise ValueError(f"mode must be 'mesh' or 'grid', got {mode!r}") + + phi_arr = np.full_like(R, phi, dtype=float) + return R, Z, phi_arr + + +# --------------------------------------------------------------------------- +# Growth rate (h5py / numpy only — reads scalars directly from C1.h5) +# --------------------------------------------------------------------------- + +def compute_ke_growth_trace( + case_dir: str | Path, +) -> tuple[np.ndarray, np.ndarray]: + """Return total kinetic energy as a function of time from ``C1.h5/scalars/``. + + Uses the ``E_K3`` scalar (total 3-D kinetic energy). No m3dc1 dependency. + + Args: + case_dir: Path to the M3D-C1 case directory. + + Returns: + ``(time, ke)`` — two 1-D float arrays of length ``ntimestep + 1``. + ``time`` is in Alfvén times; ``ke`` is in internal (normalised) units. + """ + traces = read_scalar_traces(case_dir, names=["time", "E_K3"]) + return traces["time"], traces["E_K3"] + + +def compute_growth_rate( + case_dir: str | Path, + time_idx: int | None = None, +) -> float: + """Estimate the linear growth rate from the kinetic energy time trace. + + Computes the mean instantaneous growth rate + + γ(t) = 0.5 · d(ln E_K3) / dt + + over the trace, optionally truncated at the timestep corresponding to + snapshot ``time_idx``. Returns the mean of all instantaneous γ values, + in units of 1/τ_A. + + Args: + case_dir: Path to the M3D-C1 case directory. + time_idx: Snapshot index (the NNN in ``time_NNN.h5``). If given, the + trace is truncated at the ``ntimestep`` value stored in that + snapshot's HDF5 attributes before computing the rate. If + ``None``, the full available trace is used. + + Returns: + Mean growth rate in 1/τ_A. Returns ``nan`` if fewer than two non-zero + kinetic energy values are found. + """ + time, ke = compute_ke_growth_trace(case_dir) + + if time_idx is not None: + snap_file = Path(case_dir) / f"time_{time_idx:03d}.h5" + if snap_file.exists(): + with h5py.File(snap_file, "r") as f: + ntimestep = f.attrs.get("ntimestep") + if ntimestep is not None: + mask = time <= float(ntimestep) + time = time[mask] + ke = ke[mask] + + nonzero = ke > 0 + if nonzero.sum() < 2: + return float("nan") + + first = int(np.argmax(nonzero)) + ke_nz = ke[first:] + time_nz = time[first:] + dt = np.diff(time_nz) + dt[dt == 0] = np.nan + gamma = 0.5 * np.diff(np.log(ke_nz)) / dt + return float(np.nanmean(gamma)) + + +# --------------------------------------------------------------------------- +# Equilibrium profiles (requires m3dc1 + fpy) +# --------------------------------------------------------------------------- + +def compute_flux_average_profiles( + case_dir: str | Path, + fields: list[str] | None = None, + fcoords: str = "pest", + points: int = 200, +) -> dict[str, tuple[np.ndarray, np.ndarray]]: + """Compute flux-surface-averaged radial profiles for equilibrium quantities. + + Requires the ``m3dc1`` and ``fpy`` libraries. Uses the equilibrium + (``time=-1``) state from ``C1.h5``. + + Args: + case_dir: Path to the M3D-C1 case directory. + fields: Field names to compute. Default: ``["p", "j", "ne", "q"]``. + Any name accepted by ``m3dc1.flux_average()`` is valid. + fcoords: Flux coordinate system: ``"pest"`` (default) or + ``"equal_arc"``. + points: Number of radial grid points in the output profiles. + + Returns: + Dict mapping field name to ``(psi_norm, profile)`` — two 1-D float64 + arrays of length ``points``. ``psi_norm`` runs from 0 (magnetic axis) + to 1 (LCFS). Fields that cannot be computed are omitted. + """ + import fpy + import m3dc1 as m1 + + if fields is None: + fields = ["p", "j", "ne", "q"] + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + result: dict[str, tuple[np.ndarray, np.ndarray]] = {} + + with _in_case_dir(case_dir): + sim = fpy.sim_data(c1h5, time=-1) + for field in fields: + try: + psin, profile = m1.flux_average(field, sim=sim, fcoords=fcoords, points=points) + if psin is not None and profile is not None: + result[field] = (np.asarray(psin, dtype=float), np.asarray(profile, dtype=float)) + except Exception: + pass + + return result + + +def compute_q95(psin: np.ndarray, q_profile: np.ndarray) -> float: + """Interpolate the safety factor at normalised flux psi_norm = 0.95. + + Args: + psin: 1-D array of normalised poloidal flux values spanning + 0 (magnetic axis) to at least 0.95. + q_profile: 1-D array of safety factor values at each ``psin`` point. + + Returns: + q95 as a float. Returns ``nan`` if the profile does not cover + psi_norm = 0.95. + """ + psin = np.asarray(psin, dtype=float) + q_profile = np.asarray(q_profile, dtype=float) + if len(psin) == 0 or psin.max() < 0.95: + return float("nan") + return float(np.interp(0.95, psin, q_profile)) + + +def compute_miller_geometry( + case_dir: str | Path, + res: int = 250, +) -> dict[str, float]: + """Compute Miller geometry parameters of the last closed flux surface. + + Requires the ``m3dc1`` library. Uses ``m3dc1.get_shape()`` on the + equilibrium state. + + Args: + case_dir: Path to the M3D-C1 case directory. + res: Poloidal resolution used to trace the LCFS shape. + + Returns: + Dict with keys ``"R0"`` (major radius, m), ``"a"`` (minor radius, m), + ``"kappa"`` (elongation), ``"delta"`` (triangularity). Returns an + empty dict if the computation fails. + """ + import fpy + import m3dc1 as m1 + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + + with _in_case_dir(case_dir): + sim = fpy.sim_data(c1h5, time=-1) + try: + shape = m1.get_shape(sim, res=res) + return {k: float(shape[k]) for k in ("R0", "a", "kappa", "delta")} + except Exception: + return {} + + +# --------------------------------------------------------------------------- +# Perturbed fields (requires m3dc1 + fpy) +# --------------------------------------------------------------------------- + +_DEFAULT_SKIP = frozenset({"alpha", "gradA", "kprad_rad"}) +_VECTOR_FIELDS = frozenset({"B", "E", "A", "j", "v"}) + + +def compute_perturbed_fields( + case_dir: str | Path, + time_idx: int, + R: np.ndarray, + Z: np.ndarray, + phi: np.ndarray, + fields: str | list[str] = "all", + skip_fields: list[str] | None = None, +) -> dict[str, np.ndarray]: + """Evaluate perturbed fields at given (R, Z, phi) points for one snapshot. + + Computes Δf = f(time_idx) − f(equilibrium) by calling ``eval_field()`` at + both the equilibrium and the perturbed state and subtracting. Vector fields + (B, E, A, j, v) are decomposed into R, PHI, and Z components stored under + keys such as ``"BR"``, ``"BPHI"``, ``"BZ"``. + + Note: ``eval_field`` expects the argument order ``(name, R, phi, Z, ...)``, + which is handled internally. + + Args: + case_dir: Path to the M3D-C1 case directory. + time_idx: Snapshot index to use for the perturbed state. + R: Array of R evaluation coordinates (any broadcastable shape). + Z: Array of Z evaluation coordinates (same shape as R). + phi: Array of toroidal angles in radians (same shape as R). + fields: ``"all"`` to compute every available field, or a list of + specific field names. + skip_fields: Field names to exclude even when ``fields="all"``. + Default: ``{"alpha", "gradA", "kprad_rad"}``. + + Returns: + Dict mapping field name → NumPy array with the same shape as R. Fields + that fail evaluation are omitted (a warning is printed). + """ + import fpy + from m3dc1 import eval_field + + skip = set(skip_fields) if skip_fields is not None else set(_DEFAULT_SKIP) + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + sim_eq = fpy.sim_data(c1h5, time=-1) + sim_lin = fpy.sim_data(c1h5, time=time_idx) + + available = set(sim_lin.typedict.keys()) + out: dict[str, np.ndarray] = {} + + def _add_scalar(name: str) -> None: + try: + eq_val = eval_field(name, R, phi, Z, coord="scalar", + sim=sim_eq, time=sim_eq.timeslice, quiet=True) + lin_val = eval_field(name, R, phi, Z, coord="scalar", + sim=sim_lin, time=sim_lin.timeslice, quiet=True) + out[name] = lin_val - eq_val + except Exception as exc: + print(f"WARNING: skipping scalar '{name}': {exc}") + + def _add_vector(base: str) -> None: + try: + eq_vec = eval_field(base, R, phi, Z, coord="vector", + sim=sim_eq, time=sim_eq.timeslice, quiet=True) + lin_vec = eval_field(base, R, phi, Z, coord="vector", + sim=sim_lin, time=sim_lin.timeslice, quiet=True) + delta = lin_vec - eq_vec + tag = base.upper() + out[f"{tag}R"] = delta[0] + out[f"{tag}PHI"] = delta[1] + out[f"{tag}Z"] = delta[2] + except Exception as exc: + print(f"WARNING: skipping vector '{base}': {exc}") + + if fields == "all": + for name in sorted(available - _VECTOR_FIELDS): + if name not in skip: + _add_scalar(name) + for base in sorted(_VECTOR_FIELDS & available): + if base not in skip: + _add_vector(base) + else: + for name in fields: + if name in skip: + continue + if name in _VECTOR_FIELDS: + _add_vector(name) + else: + _add_scalar(name) + + return out + + +# --------------------------------------------------------------------------- +# Field evaluation on a grid (requires m3dc1 + fpy) +# --------------------------------------------------------------------------- + +def evaluate_field_on_grid( + case_dir: str | Path, + fields: str | list[str], + time_idx: int = -1, + coord: str = "scalar", + mode: str = "mesh", + grid_res: int = 200, + phi: float = 0.0, + output_path: str | Path | None = None, +) -> dict[str, np.ndarray]: + """Evaluate one or more fields on an (R, Z) grid and optionally save to HDF5. + + Builds the evaluation grid from the mesh vertices or a regular Cartesian + grid, opens the simulation at the requested timeslice, and evaluates each + field using ``eval_field``. + + Requires the ``m3dc1`` and ``fpy`` libraries. + + Args: + case_dir: Path to the M3D-C1 case directory. + fields: Field name or list of field names, e.g. ``"te"`` or + ``["te", "ne", "p"]``. + time_idx: Timeslice index. ``-1`` selects the equilibrium state; + any non-negative integer matches the NNN in + ``time_NNN.h5``. + coord: Component or mode for field evaluation: + + - ``"scalar"`` — scalar value, or magnitude for vector + fields. + - ``"R"``, ``"phi"``, ``"Z"`` — a specific cylindrical + component (vector fields only). + - ``"all"`` — for scalar fields behaves like + ``"scalar"``; for vector fields (B, E, A, j, v) + evaluates all three components and stores them as + separate keys, e.g. ``"BR"``, ``"BPHI"``, ``"BZ"``. + + mode: Grid type: + + - ``"mesh"`` — unique mesh vertex positions; output + arrays are 1-D. + - ``"grid"`` — regular ``grid_res × grid_res`` Cartesian + grid spanning the mesh bounding box; output arrays are + 2-D. + + grid_res: Points per axis for ``mode="grid"``. + phi: Toroidal angle in radians applied to all evaluation + points. + output_path: If given, write results to an HDF5 file at this path. + The file will contain datasets ``"R"``, ``"Z"``, and one + dataset per field or component key. + + Returns: + Dict with keys ``"R"``, ``"Z"``, and one key per evaluated field or + component. Array shapes match the grid: 1-D for ``mode="mesh"``, 2-D + ``(grid_res, grid_res)`` for ``mode="grid"``. + """ + import fpy + from m3dc1 import eval_field + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + + if isinstance(fields, str): + fields = [fields] + + R_verts, Z_verts = read_mesh_vertices(c1h5) + R, Z, phi_arr = make_evaluation_grid( + R_verts, Z_verts, mode=mode, grid_res=grid_res, phi=phi + ) + + sim = fpy.sim_data(c1h5, time=time_idx) + result: dict[str, np.ndarray] = {"R": R, "Z": Z} + + for field in fields: + is_vector = field in _VECTOR_FIELDS + if coord == "all" and is_vector: + tag = field.upper() + for comp, key in (("R", f"{tag}R"), ("phi", f"{tag}PHI"), ("Z", f"{tag}Z")): + result[key] = eval_field( + field, R, phi_arr, Z, + coord=comp, sim=sim, time=sim.timeslice, quiet=True, + ) + else: + effective_coord = "scalar" if (coord == "all" and not is_vector) else coord + result[field] = eval_field( + field, R, phi_arr, Z, + coord=effective_coord, sim=sim, time=sim.timeslice, quiet=True, + ) + + if output_path is not None: + with h5py.File(str(output_path), "w") as f: + f.create_dataset("R", data=R) + f.create_dataset("Z", data=Z) + for key, arr in result.items(): + if key not in ("R", "Z"): + f.create_dataset(key, data=arr) + + return result + + +# --------------------------------------------------------------------------- +# Spectral analysis (requires m3dc1 + fpy) +# --------------------------------------------------------------------------- + +def compute_poloidal_spectrum( + case_dir: str | Path, + time_idx: int, + field: str, + coord: str = "scalar", + fcoords: str = "pest", + points: int = 200, + full_fft: bool = False, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Compute the poloidal mode spectrum of a field at a given time snapshot. + + Decomposes the field along the poloidal direction using + ``m3dc1.eigenfunction()`` with Fourier=True. The result gives the amplitude + of each poloidal harmonic m as a function of normalised flux psi_norm. + + Requires the ``m3dc1`` and ``fpy`` libraries. + + Args: + case_dir: Path to the M3D-C1 case directory. + time_idx: Snapshot index to analyse. + field: Field name: ``"p"`` (pressure), ``"B"`` (magnetic field), etc. + coord: Component: ``"scalar"`` for scalar fields; ``"R"``, ``"Z"``, + or ``"phi"`` for vector components. + fcoords: Flux coordinate system: ``"pest"`` (default) or + ``"equal_arc"``. + points: Radial resolution (number of psi_norm points). + full_fft: If ``True``, return the full two-sided FFT with all m modes. + If ``False`` (default), return a symmetric spectrum mirrored + to include both positive and negative m. + + Returns: + ``(m_modes, psi_norm, spectrum)`` where: + + - ``m_modes``: 1-D int array of poloidal mode numbers. + - ``psi_norm``: 1-D float array of length ``points``. + - ``spectrum``: 2-D float array, shape ``(len(m_modes), points)``. + ``spectrum[i, j]`` is the amplitude of mode ``m_modes[i]`` at + ``psi_norm[j]``. + """ + import fpy + import m3dc1 as m1 + + case_dir = Path(case_dir) + c1h5 = str(case_dir / "C1.h5") + + with _in_case_dir(case_dir): + sim_eq = fpy.sim_data(c1h5, time=-1) + sim_lin = fpy.sim_data(c1h5, time=time_idx) + + sim_eq_fc = m1.flux_coordinates( + sim=sim_eq, fcoords=fcoords, phit=0.0, points=points + ) + psi_norm = np.asarray(sim_eq_fc.fc.psi_norm, dtype=float) + + spec = m1.eigenfunction( + sim=[sim_eq_fc, sim_lin], + field=field, + coord=coord, + fcoords=fcoords, + points=points, + makeplot=False, + fourier=True, + full_fft=full_fft, + norm_to_unity=True, + quiet=True, + ) + + spec = np.asarray(spec, dtype=float) + + if full_fft: + m_modes = (np.fft.fftshift(np.fft.fftfreq(points, d=1.0)) * points).astype(int) + return m_modes, psi_norm, spec + + m_max = spec.shape[0] - 1 + m_modes = np.arange(-m_max, m_max + 1, dtype=int) + spec_full = np.concatenate([spec[1:][::-1], spec], axis=0) + return m_modes, psi_norm, spec_full + + +_STANDARD_SPEC_MAP = { + "p": ("p", "scalar"), + "br": ("B", "R"), + "bz": ("B", "Z"), + "bphi": ("B", "phi"), +} + + +def compute_standard_spectra( + case_dir: str | Path, + time_idx: int, + fcoords: str = "pest", + points: int = 200, + full_fft: bool = False, +) -> dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]]: + """Compute poloidal spectra for the standard set of fields: p, B_R, B_Z, B_phi. + + Calls :func:`compute_poloidal_spectrum` for each of the four default + spectrum fields used in ``postprocess_ndarray.py``. Any field that fails is + omitted from the result without raising an exception. + + Requires the ``m3dc1`` and ``fpy`` libraries. + + Args: + case_dir: Path to the M3D-C1 case directory. + time_idx: Snapshot index to analyse. + fcoords: Flux coordinate system. + points: Radial resolution. + full_fft: If ``True``, use the full two-sided FFT. + + Returns: + Dict with keys ``"p"``, ``"br"``, ``"bz"``, ``"bphi"``. Each value is + a ``(m_modes, psi_norm, spectrum)`` tuple as returned by + :func:`compute_poloidal_spectrum`. Keys for fields that fail are absent. + """ + result: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {} + for key, (field, coord) in _STANDARD_SPEC_MAP.items(): + try: + result[key] = compute_poloidal_spectrum( + case_dir, time_idx, field, coord=coord, + fcoords=fcoords, points=points, full_fft=full_fft, + ) + except Exception as exc: + print(f"WARNING: spectrum failed for '{key}': {exc}") + return result diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md b/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md new file mode 100644 index 0000000..dbd6e05 --- /dev/null +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/SKILL.md @@ -0,0 +1,80 @@ +--- +name: m3dc1-skill +description: Use this skill when working with the python modules `m3dc1_tools.py`, `m3dc1_plots.py`, `hdf5.py` and the tools created from the functions within. Triggers when working with M3D-C1 simulation data or repackaging general HDF5 files. +--- + +# m3dc1-skill + + +## Creating CLI tools from M3D-C1 python modules + +Always read the relevant API guide before creating CLI tools --- see the `Reference material` section below. + + +IMPORTANT NOTE FOR AGENTS GENERATING NEW TOOLS FROM m3dc1_tools.py +------------------------------------------------------------------ +The functions marked "requires m3dc1 + fpy" in the initial comment +in `m3dc1_tools.py` call into compiled C/Fortran code that writes diagnostic +messages (e.g. "deleting simulation object", "period = 6.28318") directly +to the OS-level stdout file descriptor. This output bypasses Python's +sys.stdout entirely and cannot be suppressed or redirected from Python. + +Consequence: any CLI tool that (1) calls one of these functions and (2) prints +its JSON result to stdout will produce contaminated output that cannot be parsed +as JSON when captured via shell redirect. + +Required pattern for any new CLI tool wrapping these functions: + - Accept an --output-json FILE argument. + - Write the JSON result to that file instead of printing to stdout. + - Document in the tool spec that the agent MUST use --output-json and read + from the file rather than capturing stdout. + +Functions that ARE affected: + `compute_flux_average_profiles`, `compute_q95`, `compute_miller_geometry`, + `compute_perturbed_fields`, `compute_poloidal_spectrum`, `compute_standard_spectra` + `evaluate_field_on_grid` + +Functions NOT affected (h5py / numpy only — no compiled stdout writes): + `read_c1input`, `list_time_snapshots`, `read_snapshot_time`, `read_scalar_traces`, + `read_case_metadata`, `read_mesh_vertices`, `make_evaluation_grid`, + `compute_ke_growth_trace`, `compute_growth_rate`, `compute_q95` + + +## When to use these tools + +Always use the M3D-C1 tools when processing or investigating datasets created by this code. Only create custom tools when your needs cannot be met by any of the existing tools. Be watchful for synonyms of alternative expressions from those used in the tool descriptions; for example, use of `repackage_hdf5` should also be triggered by calls to repack or reorganize HDF5 data, or to create a new file to hold a new dataset etc. + +When using any tools that write a temporary JSON file (to circumvent the corruption of stdout by fpy as described above) that temporary file should be deleted as soon as it is no longer needed to connect tool calls. Do not leave unnecessary JSON files in the filesystem. + +If the user's meaning or intent is unclear always ask a clarifying question --- do not silently fail or say that a required tool does not exist. + + +## Note on how M3D-C1 field data is stored + +M3D-C1 stores fields using the coefficients of the basis functions rather than the pointwise field values themselves. The field values can be found using the `eval_field` function in the `m3dc1` python module; the `eval_field_on_grid` function in `m3dc1_tools.py` evaluates field values across a spatial grid and either returns a dict or writes an HDF5 file. Use these functions or tools derived from them to evaluate fields. If the user asks for field data (e.g. temperature, density, pressure, magnetic field or current density components) to be repackaged, repacked, extracted, moved, rewritten etc. assume that they are interested in these real evaluated field values rather than the raw coefficient values. Only repackage raw coefficients if the user specifically requests this. If the user's intent is unclear always ask a clarifying question. + + +## Where to store products + +Unless explicitly requested by the user, products you create (e.g. plots or new files containing repacked or processed data) should be placed in the active dsagt project directory or its subdirectories. In particular, do not place any new files in the data directory in which the source M3D-C1 data files are located unless requested by the user. By default, place plots in a `plots/` subdirectory and new files containing data in a `processed_data/` subdirectory, creating them if necessary. + +If the user expresses a preference for new file locations attempt to maintain consistency with this choice for the remainder of the session unless new instructions are given; for example, if the user asks for a plot to go in a `new_images/` directory place subsequent plots in the same directory unless directed otherwise. If the user's intentions are unclear (e.g. plots have been placed in more the one directory in the same session) ask for clarification before proceeding. + + +## Generating shell scripts for session recreation + +When asked to make a shell script always check the user's default shell. Create the script for the default shell unless the user requests a different shell, in which case ask for clarification. In particular, Mac users may informally ask for a "bash script" even if they have zsh set as their default shell. This may be a problem as environment variables required for finding libraries such as fusion-io may exist only in the default shell's dot files. + + +## Reference material + +ALWAYS read the relevant API guide before attempting to create CLI tools from one of the python modules. + +- `references/m3dc1_output_guide.md`: explains M3D-C1's native output format in detail, including how the field coefficients are stored and evaluated. + +- `references/m3dc1_tools_api.md`: API and usage guide to the functions in `m3dc1_tools.py`. These are mostly wrappers of the `m3dc1` python module. + +- `references/m3dc1_plots_api.md`: API and usage guide for plotting using the functions in `m3dc1_plots.py`. + +- `references/hdf5_api.md`: API and usage guide for the basic HDF5 functions in `hdf5.py`. + diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/references/hdf5_api.md b/use_cases/tokamak_stability/skills/m3dc1-skill/references/hdf5_api.md new file mode 100644 index 0000000..e446e99 --- /dev/null +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/references/hdf5_api.md @@ -0,0 +1,206 @@ +# hdf5.py — API Reference + +General-purpose helpers for inspecting and repackaging HDF5 files using h5py. +All functions live in `hdf5.py`. No M3D-C1-specific knowledge is assumed; +these utilities work with any HDF5 file. + +```python +from hdf5 import list_h5_files, list_h5_variables, read_h5_dataset, read_h5_attrs, repackage_h5, json_to_h5 +``` + +--- + +## `list_h5_files(directory, recursive=False)` + +Return a sorted list of `.h5` file paths in a directory. + +| Parameter | Type | Description | +|-------------|------|----------------------------------------------------| +| `directory` | path | Directory to search | +| `recursive` | bool | If `True`, descend into subdirectories (default `False`) | + +Returns `list[str]`. Returns `[]` if the directory does not exist. + +```python +list_h5_files("m3dc1_data") +# ['m3dc1_data/C1.h5', 'm3dc1_data/equilibrium.h5', +# 'm3dc1_data/time_000.h5', 'm3dc1_data/time_001.h5'] +``` + +--- + +## `list_h5_variables(file_path)` + +Return a sorted list of all dataset paths inside an HDF5 file. + +| Parameter | Type | Description | +|-------------|------|----------------------| +| `file_path` | path | Path to the HDF5 file | + +Returns `list[str]` of internal HDF5 paths (e.g. `"group/subgroup/dataset"`). +Returns `[]` if the file cannot be opened or contains no datasets. + +```python +list_h5_variables("m3dc1_data/C1.h5") +# ['equilibrium/fields/B', 'equilibrium/fields/psi', ..., 'scalars/E_K3', ...] +``` + +--- + +## `read_h5_dataset(file_path, dataset_path)` + +Read a single dataset from an HDF5 file and return it as a NumPy array. + +| Parameter | Type | Description | +|----------------|------|----------------------------------------------------------| +| `file_path` | path | Path to the HDF5 file | +| `dataset_path` | str | Internal HDF5 path to the dataset, e.g. `"scalars/E_K3"` | + +Returns `np.ndarray`. + +Raises `KeyError` if `dataset_path` is not found; raises `OSError` if the +file cannot be opened. + +```python +ke = read_h5_dataset("m3dc1_data/C1.h5", "scalars/E_K3") +# array of kinetic energy vs timestep +``` + +--- + +## `read_h5_attrs(file_path, group_path="")` + +Read HDF5 attributes of a group or dataset as a plain Python dict. + +| Parameter | Type | Description | +|--------------|------|---------------------------------------------------------------------| +| `file_path` | path | Path to the HDF5 file | +| `group_path` | str | Internal HDF5 path of the group/dataset. Default `""` reads root attributes. | + +Returns `dict[str, Any]`. NumPy scalar types are cast to Python `int` or +`float` for easy serialisation. + +Raises `KeyError` if `group_path` is non-empty and not found; raises `OSError` +if the file cannot be opened. + +```python +read_h5_attrs("m3dc1_data/C1.h5", "equilibrium") +# {'version': 45, 'nspace': 2, 'ntimestep': 0, 'time': 0.0} + +read_h5_attrs("m3dc1_data/C1.h5") # root attributes +# {'n0_norm': 2e20, 'b0_norm': 6.5, 'l0_norm': 0.57, ...} +``` + +--- + +## `repackage_h5(output_path, sources, variables=None, selection=None, overwrite=False)` + +Create a new HDF5 file containing a subset of datasets copied from one or more +source files. + +| Parameter | Type | Description | +|---------------|------------|-----------------------------------------------------------------------------| +| `output_path` | path | Destination HDF5 file to create | +| `sources` | list[path] | Source HDF5 files to read from | +| `variables` | list[str] \| None | Dataset paths to copy from every source. `None` copies all datasets. | +| `selection` | dict \| None | Per-source override: maps source path → list of dataset paths. Takes precedence over `variables` for that source. | +| `overwrite` | bool | If `True`, overwrite an existing `output_path` (default `False`) | + +Returns `list[str]` of dataset paths written into the output file. + +Raises `FileExistsError` if `output_path` exists and `overwrite=False`. + +**Grouping behaviour** — when multiple sources are provided, datasets are +placed under a group named after the source file stem to avoid collisions: + +| Scenario | Destination path | +|---|---| +| Single source | `dataset_path` (top-level, no grouping) | +| Multiple sources, same directory | `{stem}/{dataset_path}` | +| Multiple sources, different directories | `{dir}/{stem}/{dataset_path}` | +| Multiple sources, different grandparent directories | `{granddir}/{dir}/{stem}/{dataset_path}` | + +Dataset attributes and compression settings are preserved from the source. + +```python +# Copy two datasets from a single file +repackage_h5( + "subset.h5", + sources=["m3dc1_data/C1.h5"], + variables=["scalars/E_K3", "scalars/time"], +) + +# Take different variables from two different files +repackage_h5( + "combined.h5", + sources=["m3dc1_data/C1.h5", "m3dc1_data/equilibrium.h5"], + selection={ + "m3dc1_data/C1.h5": ["scalars/E_K3", "scalars/time"], + "m3dc1_data/equilibrium.h5": ["fields/psi", "fields/B"], + }, +) +# Result groups: C1/scalars/E_K3, C1/scalars/time, +# equilibrium/fields/psi, equilibrium/fields/B +``` + +--- + +## `json_to_h5(output_path, source, overwrite=False)` + +Write arbitrarily nested JSON data to an HDF5 file. + +The `source` argument is resolved in order: + +1. `dict` or `list` — used directly as parsed data. +2. `str` or `Path` pointing to an existing file — read and parsed as JSON. +3. `str` beginning with `{` or `[` — parsed as an inline JSON string. + +**Conversion rules** (applied recursively): + +| JSON type | HDF5 result | +|---|---| +| `dict` | group; each key becomes a child name | +| uniform or rectangular list | dataset (converted to `np.ndarray`) | +| jagged, mixed, or list of dicts | numbered subgroups (`"0"`, `"1"`, …) | +| `int` / `float` | scalar dataset | +| `bool` | scalar `uint8` dataset | +| `str` | scalar string dataset | +| `None` | omitted silently | + +A top-level `list` is wrapped in a group named `"data"`. + +| Parameter | Type | Description | +|---------------|------|-------------| +| `output_path` | path | Destination HDF5 file to create | +| `source` | dict, list, str, or Path | JSON data — see above | +| `overwrite` | bool | If `True`, overwrite an existing `output_path` (default `False`) | + +Returns `list[str]` of dataset paths written. + +Raises `FileExistsError` if `output_path` exists and `overwrite=False`. +Raises `ValueError` if `source` is a string that is neither a valid file nor valid JSON. + +```python +# From a JSON file written by a pipeline tool +json_to_h5("spectrum.h5", "spectrum_output.json") + +# From an in-memory dict +m, psi, spec = compute_poloidal_spectrum(...) +json_to_h5("spectrum.h5", { + "m_modes": m.tolist(), + "psi_norm": psi.tolist(), + "spectrum": spec.tolist(), +}) +# Writes datasets: /m_modes, /psi_norm, /spectrum + +# From an inline JSON string +json_to_h5("out.h5", '{"time": [0.0, 1.0], "ke": [1e-13, 9e-8]}') + +# Nested dict — compute_standard_spectra result +spectra = compute_standard_spectra(...) +json_to_h5("spectra.h5", { + k: {"m_modes": m.tolist(), "psi_norm": p.tolist(), "spectrum": s.tolist()} + for k, (m, p, s) in spectra.items() +}) +# Writes groups /p, /br, /bz, /bphi, each containing three datasets +``` diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_output_guide.md b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_output_guide.md new file mode 100644 index 0000000..09777dc --- /dev/null +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_output_guide.md @@ -0,0 +1,652 @@ +# M3D-C1 Output File Guide + +This guide documents the structure of HDF5 output files produced by the M3D-C1 +extended-MHD simulation code, based on direct inspection of the example dataset at +`/Users/kparfrey/data/asv/run1/m3dc1_data/` and the reference scripts in `existing/`. + + +## 1. Files in a Case Directory + +A typical case directory contains the following files: + +``` +C1.h5 Main simulation output (HDF5) +equilibrium.h5 Standalone copy of the equilibrium state (HDF5) +time_000.h5 Standalone snapshot at the first saved timestep (HDF5) +time_001.h5 Standalone snapshot at the last saved timestep (HDF5) +C1input Fortran namelist input parameters (ASCII) +geqdsk Grad-Shafranov equilibrium in GEQDSK format (ASCII) +``` + +The `time_NNN.h5` files contain the same data as the corresponding `time_NNN/` groups +inside `C1.h5`. Similarly, `equilibrium.h5` is byte-for-byte identical to the +`equilibrium/` group in `C1.h5`. The standalone files are written as a convenience so +that individual snapshots can be loaded without opening the full aggregated file. + + +## 2. Simulation Units and Coordinates + +M3D-C1 uses an internal normalization where the time unit is the Alfvén time τ_A. +Physical quantities stored in the HDF5 files are in these internal (normalized) units +unless the post-processing library performs a unit conversion. + +The spatial coordinate system is cylindrical: `(R, φ, Z)` where R is the major radius +in metres, φ is the toroidal angle in radians, and Z is the vertical coordinate in +metres. For the example SPARC case: + +- R range: 1.25 – 2.55 m +- Z range: −1.25 – 1.25 m +- Magnetic axis: R ≈ 1.855 m, Z ≈ 0.0 m +- Lower X-point: R ≈ 1.492 m, Z ≈ −1.043 m +- ψ at magnetic axis (ψ_min): −1.438 (internal units) +- ψ at LCFS (ψ_LCFS): +0.0235 (internal units) + + +## 3. The C1.h5 File + +`C1.h5` is the master output file. It contains three top-level groups: + +``` +C1.h5 +├── equilibrium/ n=0 equilibrium state (stored once) +│ ├── mesh/ +│ └── fields/ +├── time_NNN/ Perturbed state at each saved timestep (one group per snapshot) +│ ├── mesh/ +│ └── fields/ +└── scalars/ Time-history traces of global diagnostics +``` + +Top-level attributes on each time group and the equilibrium group: + +``` +Attribute Type Meaning +────────────────────────────────────────────────────────────────────────────── +version int32 File format version (example: 45) +nspace int32 Spatial dimensionality (2 = 2D poloidal mesh) +ntimestep int32 Timestep number of this snapshot (0 for equilibrium) +time float64 Simulation time of this snapshot in Alfvén times +iwall_regions int32 Flag: whether wall-region elements are present +``` + + +## 4. Mesh Structure + +Both `equilibrium/mesh/` and every `time_NNN/mesh/` contain the same two datasets +describing the finite-element mesh on the poloidal cross-section. The mesh does not +change during the simulation. + +### 4.1 `mesh/elements` — shape (nelms, 8) + +Each row describes one triangular finite element. For the example case, +nelms = 18 555. + +``` +Column Meaning +────────────────────────────────────────────────────────────────────────────── +0 ΔR component of the first edge vector (metres; may be negative) +1 ΔZ component of the first edge vector (metres; may be negative) +2 Length of the first edge: √(col0² + col1²) (metres; always positive) +3 Orientation angle α of the local coordinate frame (radians; −π to π) +4 R coordinate of the element reference vertex (metres) +5 Z coordinate of the element reference vertex (metres) +6 Zone/region flag (integer encoded as float: 0, 1, 2, or 4) +7 Constant equal to 1.0 (padding / normalisation flag) +``` + +Columns 0–3 define the affine mapping from a reference triangle to the physical +element, which is used internally by the FEM basis functions. + +Zone codes (column 6): + +``` +Zone Count Meaning +────────────────────────────────────────────────────────────────────────────── +0 18 246 Interior plasma region (bulk of the mesh) +1 1 Magnetic axis element +2 286 Scrape-off layer / wall-adjacent region +4 22 Special boundary elements +``` + +Mesh group attributes: + +``` +Attribute Type Value (example) Meaning +──────────────────────────────────────────────────────────────────────────────── +nelms int32 18 555 Number of triangular elements +nplanes int32 1 Number of toroidal planes (1 = 2D) +nperiods int32 1 Toroidal periodicity +period float64 6.2832… Toroidal period (2π for full torus) +height float64 2.4996 Poloidal height of the mesh (m) +width float64 1.2999 Radial width of the mesh (m) +3D int32 0 Flag: 0 = 2D mesh, 1 = 3D mesh +ifull_torus int32 0 Flag: full-torus (1) or wedge (0) +version int32 45 Format version +``` + +### 4.2 `mesh/adjacency` — shape (nelms, 3) + +Each row lists the three element indices (0-based) of the neighbours sharing an edge +with this element. A value of −1 indicates a boundary with no neighbour on that side. + +``` +dtype: int32 +range: −1 to nelms − 1 +``` + + +## 5. Field Arrays + +### 5.1 Dataset shape and layout + +Every field dataset inside `fields/` has shape **(nelms, 20)**. + +- **First axis** (size nelms = 18 555): element index, matching the row order in + `mesh/elements` and `mesh/adjacency`. +- **Second axis** (size 20): the 20 finite-element expansion coefficients stored per + element. M3D-C1 uses high-order C1-continuous triangular elements; these 20 values + are the degrees of freedom of the FEM basis within each element. They are not + directly interpretable as pointwise field values — use the `m3dc1` Python library + (specifically `eval_field()`) to evaluate the field at arbitrary (R, φ, Z) points. + +All field datasets use dtype `float32`. + +### 5.2 Real and imaginary parts: the `_i` suffix + +Because M3D-C1 performs a Fourier decomposition in the toroidal direction, perturbed +fields are stored as complex amplitudes. The real and imaginary parts are stored as +separate datasets using the naming convention: + +``` +field real part of the n = ntor Fourier component +field_i imaginary part of the n = ntor Fourier component +``` + +The equilibrium (`n = 0`) fields have no meaningful imaginary part (they are +axisymmetric), but `field_i` datasets are still written for consistency. + +The full 3D field at a given toroidal angle φ is reconstructed as: + +``` +f_total(R, φ, Z) = f_eq(R, Z) + f_real(R, Z) · cos(n·φ) + − f_imag(R, Z) · sin(n·φ) +``` + +where `n = ntor` (the toroidal mode number from `C1input`). + +### 5.3 Equilibrium vs. perturbed fields + +The `equilibrium/fields/` group stores the **axisymmetric (n = 0) equilibrium state**. +Each `time_NNN/fields/` group stores the **perturbed (n = ntor) component only** — +not the total field. The post-processing workflow explicitly separates the two: + +```python +f_pert(R, Z) = eval_field(name, R, φ, Z, sim=sim_lin) - eval_field(name, R, φ, Z, sim=sim_eq) +``` + +In the example case (ntor = 9, linear run), `time_000/fields/psi` is effectively zero +(no initial psi perturbation), while `time_001/fields/psi` shows the fully grown n = 9 +instability after 1000 Alfvén times. + + +## 6. Field Catalogue + +### 6.1 Fields present only in the equilibrium group + +These transport coefficient fields are time-independent and are only written once: + +``` +Name Meaning +────────────────────────────────────────────────────────────────────────────── +eta Resistivity +eta_J Resistivity × J (Ohmic heating term) +kappa Isotropic thermal conductivity +kappar Field-aligned thermal conductivity +visc Isotropic viscosity +visc_c Compressional viscosity +denm Neutral/minority species density +``` + +### 6.2 Fields present in all groups (equilibrium and all time snapshots) + +#### Magnetic / electromagnetic + +``` +Name Long name / description +────────────────────────────────────────────────────────────────────────────── +psi Poloidal magnetic flux (stream function) +psi_i Imaginary part of psi perturbation +f Toroidal field function F = R·Bφ +f_i Imaginary part of f perturbation +fp Derivative of f with respect to ψ +fp_i Imaginary part of fp perturbation +phi Electrostatic / stream-function potential +phi_i Imaginary part of phi +I Parallel current +I_i Imaginary part of I +E_R Radial electric field component +E_R_i Imaginary part of E_R +E_Z Vertical electric field component +E_Z_i Imaginary part of E_Z +E_PHI Toroidal electric field component +E_PHI_i Imaginary part of E_PHI +E_par Parallel electric field +E_par_i Imaginary part of E_par +jphi Toroidal current density +jphi_i Imaginary part of jphi +``` + +#### Kinetic / thermodynamic + +``` +Name Long name / description +────────────────────────────────────────────────────────────────────────────── +P Total pressure (ions + electrons) +P_i Imaginary part of P +Pe Electron pressure +Pe_i Imaginary part of Pe +V Flow velocity (or poloidal flux related) +V_i Imaginary part of V +te Electron temperature +te_i Imaginary part of te +ti Ion temperature +ti_i Imaginary part of ti +ne Electron number density +ne_i Imaginary part of ne +den Total mass density +den_i Imaginary part of den +chi Flux-surface label (related to poloidal flux) +chi_i Imaginary part of chi +zeff Effective ion charge Z_eff +zeff_i Imaginary part of zeff +``` + +#### Transport / source terms + +``` +Name Long name / description +────────────────────────────────────────────────────────────────────────────── +bdotgradp B · ∇p (parallel pressure gradient, drive term) +bdotgradp_i Imaginary part of bdotgradp +bdotgradt B · ∇T (parallel temperature gradient) +bdotgradt_i Imaginary part of bdotgradt +torque_em Electromagnetic torque +torque_em_i Imaginary part of torque_em +torque_ntv Neoclassical toroidal viscosity torque +torque_ntv_i Imaginary part of torque_ntv +``` + +#### Classification / geometry flags + +``` +Name Meaning +────────────────────────────────────────────────────────────────────────────── +magnetic_region 0.0 = outside separatrix, 1.0 = inside LCFS +mesh_zone Zone code, matching elements column 6 +wall_dist Distance to the nearest wall surface (metres) +``` + + +## 7. Scalars: Global Time Traces + +`scalars/` contains 97 datasets, each of shape **(ntimestep + 1,)** = **(1001,)** for +this case. The values are stored in the same internal (Alfvén-unit) normalisation as +the field arrays. + +### 7.1 Time coordinate + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +time Simulation time at each output step (Alfvén times; 0.0 – 1000.0) +dt Time-step size used at each step +``` + +The `scalars/` group itself carries `@ntimestep = 1000`, reflecting the total number of +steps taken. + +### 7.2 Energy diagnostics + +All energies are in internal (normalised) units. + +``` +Name Meaning +────────────────────────────────────────────────────────────────────────────── +E_K3 Total 3D kinetic energy +E_K3D 3D kinetic energy (D-species component) +E_K3H 3D kinetic energy (H-species component) +E_KP Poloidal kinetic energy +E_KPD Poloidal kinetic energy (D-species) +E_KPH Poloidal kinetic energy (H-species) +E_KT Toroidal kinetic energy +E_KTD Toroidal kinetic energy (D-species) +E_KTH Toroidal kinetic energy (H-species) +E_MP Perturbed magnetic energy +E_MPC Perturbed magnetic energy (compressional) +E_MPD Perturbed magnetic energy (D-species) +E_MPH Perturbed magnetic energy (H-species) +E_MPV Perturbed magnetic energy (vacuum region) +E_MT Total magnetic energy (equilibrium + perturbation) +E_MTC Total magnetic energy (compressional) +E_MTD Total magnetic energy (D-species) +E_MTH Total magnetic energy (H-species) +E_MTV Total magnetic energy (vacuum) +E_P Total pressure energy +E_PD Pressure energy (density component) +E_PE Electron pressure energy +E_PH Pressure energy (H-component) +E_grav Gravitational potential energy +W_M Integrated perturbed magnetic energy (used for growth-rate diagnostics) +W_P Integrated perturbed pressure energy +``` + +In the example linear run, `E_K3` grows from ~1.4 × 10⁻¹³ at t = 1 to ~9.7 × 10⁻⁸ at +t = 1000, giving an estimated growth rate γ ≈ 0.0067 τ_A⁻¹ for the n = 9 mode. + +### 7.3 Geometry and equilibrium state + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +xmag R coordinate of the magnetic axis (metres) +zmag Z coordinate of the magnetic axis (metres) +xnull R coordinate of the primary X-point (metres) +znull Z coordinate of the primary X-point (metres) +xnull2 R coordinate of the secondary X-point (metres) +znull2 Z coordinate of the secondary X-point (metres) +psimin ψ at the magnetic axis (internal units; most negative value) +psi_lcfs ψ at the last closed flux surface +psi0 Reference ψ value (axis or vacuum) +``` + +For the example case: magnetic axis at (1.855, ≈0) m; X-point at (1.492, −1.043) m. + +### 7.4 Integral conservation quantities + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +toroidal_current Total toroidal plasma current +toroidal_current_p Toroidal current (plasma region only) +toroidal_current_w Toroidal current (wall/SOL region) +toroidal_flux Total toroidal magnetic flux +toroidal_flux_p Toroidal flux (plasma region) +angular_momentum Total angular momentum +angular_momentum_p Angular momentum (plasma region) +helicity Magnetic helicity +volume Total plasma volume +volume_p Plasma volume (inside LCFS) +volume_pd Plasma volume (D-species) +area Plasma cross-sectional area +area_p Cross-sectional area (plasma region) +particle_number Total particle number (ions) +particle_number_p Particle number (plasma region) +electron_number Total electron count +circulation Poloidal flow circulation +loop_voltage Inductive loop voltage +``` + +### 7.5 Radiation and particle sources + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +radiation Total radiated power +brem_rad Bremsstrahlung radiation +line_rad Line radiation +reck_rad Recombination radiation (K channel) +recp_rad Recombination radiation (P channel) +kprad_n KPRAD impurity density +kprad_n0 KPRAD initial impurity density +kprad_dt KPRAD time derivative +Particle_source Particle source rate +Particle_Flux_convective Convective particle flux across LCFS +Particle_Flux_diffusive Diffusive particle flux across LCFS +ion_loss Ion particle loss rate +runaways Runaway electron count +temax Maximum electron temperature +``` + +### 7.6 Power and energy fluxes + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +power_injected Total injected power +Flux_kinetic Kinetic energy flux across LCFS +Flux_poynting Poynting flux +Flux_pressure Pressure-driven energy flux +Flux_thermal Thermal conduction flux +Parallel_viscous_heating Heating from parallel viscous dissipation +``` + +### 7.7 Torques + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +Torque_em Electromagnetic (J × B) torque +Torque_visc Viscous torque +Torque_ntv Neoclassical toroidal viscosity torque +Torque_com Combined torque +Torque_gyro Gyroviscous torque +Torque_parvisc Parallel viscosity torque +Torque_sol Scrape-off-layer torque +``` + +### 7.8 Wall forces and Fourier diagnostics + +``` +Name Description +────────────────────────────────────────────────────────────────────────────── +Wall_Force_n0_x n=0 wall force, x-component +Wall_Force_n0_x_halo n=0 wall force (halo current), x-component +Wall_Force_n0_y n=0 wall force, y-component +Wall_Force_n0_z n=0 wall force, z-component +Wall_Force_n0_z_halo n=0 wall force (halo current), z-component +Wall_Force_n1_x n=1 wall force, x-component +Wall_Force_n1_y n=1 wall force, y-component +IP_co Plasma current, cos-component of n=ntor Fourier mode +IP_sn Plasma current, sin-component of n=ntor Fourier mode +M_IZ Edge / wall current amplitude +M_IZ_co Edge current, cos-component +M_IZ_sn Edge current, sin-component +i_control%err_i Current control integrator error +i_control%err_p_old Current control proportional error (previous step) +n_control%err_i Density control integrator error +n_control%err_p_old Density control proportional error (previous step) +Ave_P Volume-averaged pressure +``` + + +## 8. The Standalone Snapshot Files + +### `equilibrium.h5` + +Contains exactly the same data as `C1.h5/equilibrium/`. The top-level structure is: + +``` +equilibrium.h5 +├── fields/ (same 56 datasets as C1.h5/equilibrium/fields/) +└── mesh/ (elements, adjacency) +``` + +Top-level attributes: same as the `equilibrium` group in `C1.h5`. + +### `time_NNN.h5` + +Contains exactly the same data as `C1.h5/time_NNN/`. The top-level structure is: + +``` +time_NNN.h5 +├── fields/ (50 datasets: same as C1.h5/time_NNN/fields/) +└── mesh/ (elements, adjacency) +``` + +Top-level attributes: same as the `time_NNN` group in `C1.h5`, including `ntimestep` +and `time` reflecting the actual simulation time of this snapshot. + +The files are named sequentially (`time_000.h5`, `time_001.h5`, …) corresponding to +the order in which snapshots were saved; the file index does not necessarily equal the +timestep number. The actual timestep number and simulation time are read from the +`ntimestep` and `time` attributes. + + +## 9. The `C1input` Parameter File + +`C1input` is an ASCII Fortran namelist file (`&inputnl … /`) that controls the +simulation. Key parameters relevant to post-processing: + +``` +Parameter Type Example value Meaning +───────────────────────────────────────────────────────────────────────────────────── +ntor int 9 Toroidal mode number n of the perturbation +pscale float 0.8150 Pressure normalisation scale factor +batemanscale float 1.0634 Bateman transport scaling factor +dt float 1.0 Time step (Alfvén times) +ntimemax int 1000 Total number of time steps +ntimepr int 1000 Steps between snapshot outputs +linear int 1 0 = nonlinear run, 1 = linear run +numvar int 3 Physics model: 1=2-field, 2=4-field, 3=6-field +idens int 1 1 = evolve density equation +ion_mass float 2.0 Ion mass in atomic mass units (2 = deuterium) +zeff float 1.0 Effective charge number Z_eff +itor int 1 1 = toroidal geometry +``` + +`pscale` and `batemanscale` are used to rescale pressure and transport from the GEQDSK +equilibrium to the M3D-C1 normalisation. `ntor` determines which Fourier component is +stored in the `field` / `field_i` datasets in the time-slice groups. + + +## 10. How to Access Data + +A set of standalone post-processing tools is provided in `m3dc1_tools.py`. +These are designed for use in an agentic pipeline and cover the most common +operations. See `m3dc1_tools_api.md` in this directory for the full API reference. + +### Recommended starting point — case summary + +```python +from m3dc1_tools import read_case_metadata + +meta = read_case_metadata("m3dc1_data") +# meta["params"]["ntor"] → 9 +# meta["snapshots"] → [0, 1] +# meta["R_mag"] → 1.855 m +``` + +### Listing and reading scalar time traces + +```python +from m3dc1_tools import read_scalar_traces + +traces = read_scalar_traces("m3dc1_data", names=["time", "E_K3"]) +t = traces["time"] # Alfvén times, shape (1001,) +ke = traces["E_K3"] # kinetic energy +``` + +### Reading mesh vertices + +```python +from m3dc1_tools import read_mesh_vertices + +R, Z = read_mesh_vertices("m3dc1_data/C1.h5") +# R, Z are 1-D float32 arrays of unique mesh vertex positions +``` + +### Building evaluation grids for field interpolation + +```python +from m3dc1_tools import read_mesh_vertices, make_evaluation_grid + +R_v, Z_v = read_mesh_vertices("m3dc1_data/C1.h5") +R, Z, phi = make_evaluation_grid(R_v, Z_v, mode="mesh") +# or use a regular 200×200 Cartesian grid: +R, Z, phi = make_evaluation_grid(R_v, Z_v, mode="grid", grid_res=200) +``` + +### Growth rate + +```python +from m3dc1_tools import compute_growth_rate + +gamma = compute_growth_rate("m3dc1_data") # → ~0.0067 τ_A⁻¹ +``` + +### Listing all HDF5 variables + +```python +from hdf5 import list_h5_variables + +variables = list_h5_variables("m3dc1_data/C1.h5") +``` + +### Reading a single HDF5 dataset directly + +```python +from hdf5 import read_h5_dataset, read_h5_attrs + +psi = read_h5_dataset("m3dc1_data/C1.h5", "equilibrium/fields/psi") +# psi shape: (18555, 20) — FEM coefficients, not pointwise values + +attrs = read_h5_attrs("m3dc1_data/C1.h5", "equilibrium") +# {"version": 45, "nspace": 2, "ntimestep": 0, "time": 0.0} +``` + +### Evaluating a field at arbitrary (R, Z) points (requires m3dc1) + +The FEM coefficient arrays must be evaluated through the `m3dc1` library; +they cannot be interpolated manually. Use: + +```python +import fpy +from m3dc1.eval_field import eval_field + +sim = fpy.sim_data("m3dc1_data/C1.h5", time=-1) # time=-1 = equilibrium +# eval_field argument order: (name, R, phi, Z, ...) +psi_values = eval_field("psi", R, phi, Z, coord="scalar", sim=sim, time=sim.timeslice) +``` + +For perturbed-field evaluation across a full snapshot, use the higher-level +wrapper: + +```python +from m3dc1_tools import compute_perturbed_fields + +perts = compute_perturbed_fields("m3dc1_data", time_idx=1, R=R, Z=Z, phi=phi) +# perts["psi"] → perturbed psi at each (R, Z, phi) point +``` + + +## 11. Notes on Field Magnitudes and Units + +The fields are stored in M3D-C1's internal normalisation. Absolute values are only +meaningful in context: + +- **Equilibrium `psi`**: ranges ≈ −480 000 to +452 000 (internal flux units) +- **Equilibrium `te`**: ranges ≈ ±7.9 × 10⁶ (normalised temperature; includes FEM + coefficients, not directly in eV) +- **Perturbed `psi` (time_001)**: max |ψ_pert| ≈ 1.8 × 10⁶ after 1000 τ_A of growth +- **Energies (scalars)**: in normalised units; E_K3 grows from ~10⁻¹³ to ~10⁻⁷ over + 1000 τ_A + +The `m3dc1` Python library handles conversion to physical units (MKS) when `units="mks"` +is passed to evaluation functions. + + +## 12. Summary Table of HDF5 Datasets + +``` +Path Shape dtype Description +─────────────────────────────────────────────────────────────────────────────────── +equilibrium/mesh/elements (18555, 8) float32 Element geometry +equilibrium/mesh/adjacency (18555, 3) int32 Element neighbours +equilibrium/fields/ (18555, 20) float32 Equilibrium field (56 fields) +time_NNN/mesh/elements (18555, 8) float32 Same mesh (unchanged) +time_NNN/mesh/adjacency (18555, 3) int32 Same adjacency +time_NNN/fields/ (18555, 20) float32 Perturbed field (50 fields) +time_NNN/fields/_i (18555, 20) float32 Imaginary part (50 fields) +scalars/ (1001,) float32 Global time traces (97 traces) +scalars/time (1001,) float32 Time axis (Alfvén times) +``` diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_plots_api.md b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_plots_api.md new file mode 100644 index 0000000..1602a5f --- /dev/null +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_plots_api.md @@ -0,0 +1,441 @@ +# m3dc1_plots.py — API Reference + +Plotting library for M3D-C1 simulation output. All public functions produce +exactly one figure, save it to a caller-specified path, close the figure, and +return the output `Path`. No `plt.show()` calls are made; the module is +headless-safe. + +--- + +## Architecture + +Functions are divided into two categories and one layer of convenience wrappers, +all clearly labelled with section headers in the source file. + +### Category A — Pure matplotlib + +Data is loaded and computed by calling functions from `m3dc1_tools.py`. No +m3dc1 library plotting routines are called. These functions can be used +anywhere h5py and numpy are available; those that call fpy-dependent +`m3dc1_tools` functions will emit C/Fortran-level diagnostics to stdout (see +[stdout note](#stdout-note)). + +### Category B — m3dc1 library wrappers + +These wrap plotting routines from the `m3dc1` Python package. The module uses +one of two patterns to redirect output to a caller-supplied path: + +- **Figure-capture pattern** (most functions): records `plt.get_fignums()` + before and after the m3dc1 call, saves the first new figure to `output_path`, + closes all new figures. +- **Tempdir pattern** (`plot_poincare`): calls the m3dc1 function with its own + save mechanism directed to a temporary directory, then moves the resulting + file to `output_path`. + +All Category B functions raise `ImportError` with a helpful message if `m3dc1` +is not importable. + +### Stdout note + +Category A functions that call fpy-dependent `m3dc1_tools` functions +(`plot_flux_average_profiles`, `plot_safety_factor`, `plot_poloidal_spectrum`, +`plot_standard_spectra`, `plot_perturbed_field_map`, `plot_geqdsk_compare`) will +emit C/Fortran-level diagnostic messages (e.g. `"deleting simulation object"`) +directly to the OS-level stdout. This is harmless in library use — Python +output is not captured. It becomes an issue only when capturing stdout in a +subprocess; in that case write results to a file rather than relying on stdout. + +--- + +## Category A Functions + +### `plot_kinetic_energy` + +```python +plot_kinetic_energy( + case_dir, output_path, + annotate_growth_rate=True, dpi=150 +) -> Path +``` + +Semi-log plot of total kinetic energy (E_K3) vs time in Alfvén units. +Optionally annotates with the mean linear growth rate γτ_A computed by +`compute_growth_rate`. + +--- + +### `plot_growth_rate_vs_time` + +```python +plot_growth_rate_vs_time( + case_dir, output_path, + smooth_window=5, dpi=150 +) -> Path +``` + +Instantaneous growth rate γτ_A(t) = 0.5 · d(ln E_K3)/dt. Shows both the raw +trace (faint) and a running-average smoothed overlay. + +--- + +### `plot_flux_average_profiles` + +```python +plot_flux_average_profiles( + case_dir, output_path, + fields=None, fcoords="pest", points=200, dpi=150 +) -> Path +``` + +Multi-panel figure with one subplot per field vs ψ_norm. Uses the equilibrium +state (time = −1). Default fields: `["p", "j", "ne", "q"]`. Requires m3dc1 + fpy. + +--- + +### `plot_safety_factor` + +```python +plot_safety_factor( + case_dir, output_path, + fcoords="pest", points=200, dpi=150 +) -> Path +``` + +q(ψ_norm) with horizontal reference lines at q = 1, 2, 3 and a q₉₅ annotation. +Uses the equilibrium state. Requires m3dc1 + fpy. + +--- + +### `plot_poloidal_spectrum` + +```python +plot_poloidal_spectrum( + case_dir, time_idx, field, output_path, + coord="scalar", fcoords="pest", points=200, n_modes=20, dpi=150 +) -> Path +``` + +2-D heatmap: spectrum amplitude as a function of poloidal mode number m (y-axis) +and ψ_norm (x-axis). Shows ±`n_modes` harmonics around m = 0. Requires m3dc1 + fpy. + +--- + +### `plot_standard_spectra` + +```python +plot_standard_spectra( + case_dir, time_idx, output_path, + fcoords="pest", points=200, n_modes=20, dpi=150 +) -> Path +``` + +2 × 2 panel figure: poloidal spectra for p, B_R, B_Z, B_φ. Requires m3dc1 + fpy. + +--- + +### `plot_perturbed_field_map` + +```python +plot_perturbed_field_map( + case_dir, time_idx, field, output_path, + mode="grid", grid_res=200, phi=0.0, dpi=150 +) -> Path +``` + +Filled contour of δf = f(time_idx) − f(equilibrium) on the R,Z plane at a +given toroidal angle φ. Uses a diverging colormap (RdBu_r) symmetric about +zero. Requires m3dc1 + fpy. + +| Parameter | Description | +|-----------|-------------| +| `field` | Scalar field name: `"psi"`, `"p"`, `"ne"`, etc. | +| `mode` | `"grid"` (regular Cartesian grid) or `"mesh"` (mesh vertices) | +| `grid_res`| Points per axis for `mode="grid"` | +| `phi` | Toroidal angle in radians | + +--- + +### `plot_geqdsk` + +```python +plot_geqdsk(gfile_path, output_path, dpi=150) -> Path +``` + +Flux surface contours and plasma/limiter boundary from a GEQDSK equilibrium +file. No m3dc1 or fpy dependency. + +--- + +### `plot_geqdsk_compare` + +```python +plot_geqdsk_compare( + case_dir, output_path, + gfile="geqdsk", dpi=150 +) -> Path +``` + +3-panel figure: GEQDSK ψ_norm (left), C1.h5 ψ_norm (centre), difference +(right). Requires m3dc1 + fpy to evaluate ψ from C1.h5 on the GEQDSK grid. + +--- + +### `plot_tpf_vs_time` + +```python +plot_tpf_vs_time( + case_dir, field, output_path, + ts_list=None, units="mks", points=250, dpi=150 +) -> Path +``` + +Total poloidal flux of a field as a function of time snapshot. Iterates over +`ts_list` (default: all available snapshots), calling `m3.tpf` for each. +Writes a sidecar `.dat` file (columns: `ts time tpf`) alongside `output_path`. +Requires m3dc1 + fpy. + +--- + +## Category B Functions + +All Category B functions call `_import_m3dc1()` at entry and raise `ImportError` +if m3dc1 is not available. + +### `plot_field` + +```python +plot_field( + case_dir, time_idx, field, output_path, + coord="scalar", phi=0.0, points=250, tor_av=1, + units="mks", mesh=False, bound=False, lcfs=False, dpi=150 +) -> Path +``` + +2-D filled-contour field plot on the R,Z plane via `m3.plot_field`. Uses the +figure-capture pattern. + +The colormap and normalisation are chosen automatically from the data: + +- If both positive and negative values are present and neither extreme is less + than 5 % of the other, uses **RdBu_r** with symmetric limits + `vmin = −vmax = max(|data|)`. +- Otherwise uses **inferno** with natural data limits (suited to + single-sign fields such as pressure or density). + +The colorbar and figure title use a LaTeX label derived from the field name and +component, e.g. `coord="phi"` on field `"j"` produces $J_{\phi}$. + +--- + +### `plot_mesh` + +```python +plot_mesh(case_dir, output_path, boundary=False, dpi=150) -> Path +``` + +Triangular mesh visualisation via `m3.plot_mesh`. + +--- + +### `plot_flux_surface_shape` + +```python +plot_flux_surface_shape( + case_dir, time_idx, output_path, + points=200, dpi=150 +) -> Path +``` + +Flux surface psi contours on the R,Z plane via `m3.plot_shape`. + +--- + +### `plot_field_mesh` + +```python +plot_field_mesh( + case_dir, time_idx, field, output_path, + coord="scalar", phi=0.0, dpi=150 +) -> Path +``` + +Field plotted on the mesh triangulation via `m3dc1.plot_field_mesh`. + +--- + +### `plot_field_vs_phi` + +```python +plot_field_vs_phi( + case_dir, time_idx, field, output_path, + R=None, Z=None, coord="scalar", phi_res=64, dpi=150 +) -> Path +``` + +Field vs toroidal angle φ at a fixed (R, Z) point. `R` and `Z` default to the +magnetic axis when not supplied. + +--- + +### `plot_flux_average_m3` + +```python +plot_flux_average_m3( + case_dir, time_idx, field, output_path, + fcoords="pest", points=200, dpi=150 +) -> Path +``` + +Flux-averaged profile via `m3.plot_flux_average`. The `_m3` suffix +distinguishes this from the pure-matplotlib `plot_flux_average_profiles` in +Category A. + +--- + +### `plot_line` + +```python +plot_line( + case_dir, time_idx, field, output_path, + coord="scalar", angle=0.0, zoff=0.0, + dist_from_magax=False, dpi=150 +) -> Path +``` + +Field profile along a radial line at fixed phi=0 via `m3.plot_line`. + +--- + +### `plot_eigenfunction` + +```python +plot_eigenfunction( + case_dir, time_idx, field, output_path, + coord="scalar", phit=0.0, fcoords="pest", + points=200, dpi=150 +) -> Path +``` + +Eigenfunction and poloidal spectrum via `m3.eigenfunction`. Internally runs +`flux_coordinates` on the equilibrium (time=0) then evaluates the perturbed +field at `time_idx`, producing the `[fc_result, lin_sim]` pair that +`eigenfunction` requires. + +--- + +### `plot_poincare` + +```python +plot_poincare(case_dir, time_idx, output_path, dpi=150) -> Path +``` + +Poincaré section via `m3.plot_poincare`. Requires pre-computed Poincaré data +in `case_dir` (generated by `m3.run_trace` or equivalent). Uses the tempdir +pattern. + +--- + +### `plot_signal` + +```python +plot_signal( + case_dir, signal, output_path, + pspec=False, pts_per_probe=1, dpi=150 +) -> Path +``` + +Diagnostic signal time traces via `m3.plot_signal`. + +--- + +### `plot_time_trace` + +```python +plot_time_trace( + case_dir, trace, output_path, + units="mks", dpi=150 +) -> Path +``` + +Global scalar time trace via `m3.plot_time_trace_fast`. + +--- + +## Convenience Wrappers + +### `plot_stability_summary` + +```python +plot_stability_summary( + case_dir, time_idx, output_dir, + prefix="", dpi=150 +) -> list[Path] +``` + +Produces the standard linear stability figure set: + +| Filename | Function called | +|----------|----------------| +| `{prefix}ke.png` | `plot_kinetic_energy` | +| `{prefix}growth_rate.png` | `plot_growth_rate_vs_time` | +| `{prefix}spectra.png` | `plot_standard_spectra` | +| `{prefix}spectrum_p.png` | `plot_poloidal_spectrum(field="p")` | +| `{prefix}field_psi.png` | `plot_perturbed_field_map(field="psi")` | + +Individual failures are caught and printed as warnings; successful outputs are +collected in the returned list. + +--- + +### `plot_equilibrium_overview` + +```python +plot_equilibrium_overview( + case_dir, output_dir, + prefix="", dpi=150 +) -> list[Path] +``` + +Equilibrium figure set: + +| Filename | Function called | +|----------|----------------| +| `{prefix}profiles.png` | `plot_flux_average_profiles` | +| `{prefix}q_profile.png` | `plot_safety_factor` | +| `{prefix}mesh.png` | `plot_mesh` | + +--- + +### `plot_case_summary` + +```python +plot_case_summary( + case_dir, time_idx, output_dir, + prefix="", dpi=150 +) -> list[Path] +``` + +Calls `plot_stability_summary` then `plot_equilibrium_overview` and returns +the combined list of Paths. + +--- + +## Usage Examples + +```python +import sys +sys.path.insert(0, "/path/to/tokamak_stability") +import m3dc1_plots + +case = "/data/runs/m3dc1_data" + +# Quick KE check (no m3dc1 needed) +m3dc1_plots.plot_kinetic_energy(case, "ke.png") + +# Safety factor (requires m3dc1 + fpy) +m3dc1_plots.plot_safety_factor(case, "q_profile.png") + +# Full stability figure set for snapshot 1 +paths = m3dc1_plots.plot_stability_summary(case, time_idx=1, output_dir="figs/") +print("Created:", paths) +``` diff --git a/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_tools_api.md b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_tools_api.md new file mode 100644 index 0000000..18f94ae --- /dev/null +++ b/use_cases/tokamak_stability/skills/m3dc1-skill/references/m3dc1_tools_api.md @@ -0,0 +1,388 @@ +# m3dc1_tools — API Reference + +For plotting functions that consume the outputs of these tools, see +[m3dc1_plots_api.md](m3dc1_plots_api.md). + +All functions live in `m3dc1_tools.py` in the `use_cases/tokamak_stability/` +directory. Import them directly: + +```python +from m3dc1_tools import read_case_metadata, compute_growth_rate, ... +``` + +Functions are grouped by dependency level. The first three groups require only +`h5py` and `numpy`. The final two groups additionally require the `m3dc1` and +`fpy` libraries. + + +--- + +## Case Inspection + +### `read_c1input(case_dir)` + +Parse the `C1input` Fortran namelist and return the key simulation parameters +as a plain dict. + +| Parameter | Type | Description | +|------------|--------|----------------------------------------------| +| `case_dir` | path | M3D-C1 case directory containing `C1input` | + +Returns `dict` with keys: `ntor` (int), `pscale`, `batemanscale`, `dt` +(float), `ntimemax`, `ntimepr`, `linear`, `numvar` (int), `ion_mass`, `zeff` +(float). Any absent parameter is `None`. + +```python +params = read_c1input("m3dc1_data") +# {'ntor': 9, 'pscale': 0.815, 'linear': 1, ...} +``` + +--- + +### `list_time_snapshots(case_dir)` + +Return a sorted list of integer indices for the `time_NNN.h5` snapshot files +present in the case directory. + +| Parameter | Type | Description | +|------------|------|---------------------------| +| `case_dir` | path | M3D-C1 case directory | + +Returns `list[int]`, e.g. `[0, 1]`. Returns `[]` if no snapshots are found. + +```python +snaps = list_time_snapshots("m3dc1_data") +# [0, 1] +``` + +--- + +### `read_snapshot_time(case_dir, time_idx, units="alfven")` + +Return the simulation time of a specific snapshot. + +| Parameter | Type | Description | +|------------|------|----------------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `time_idx` | int | Snapshot file index (the NNN in `time_NNN.h5`) | +| `units` | str | `"alfven"` (default) or `"s"`/`"mks"` for SI seconds | + +Returns `float`. `"s"` requires `m3dc1`; falls back to Alfvén times with a +warning if unavailable. Returns `nan` when the snapshot is not found. + +```python +t = read_snapshot_time("m3dc1_data", 1) # → 1000.0 (Alfvén times) +t_s = read_snapshot_time("m3dc1_data", 1, "mks") # → physical seconds +``` + +--- + +### `read_scalar_traces(case_dir, names=None)` + +Read global time-trace diagnostics from `C1.h5/scalars/`. + +| Parameter | Type | Description | +|------------|-----------------|----------------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `names` | list[str] | None | Scalar names to read. `None` returns all (~97 traces). | + +Returns `dict[str, np.ndarray]` — each array has length `ntimestep + 1`. The +`"time"` key gives Alfvén times. + +```python +traces = read_scalar_traces("m3dc1_data", names=["time", "E_K3", "W_M"]) +t = traces["time"] # Alfvén times, shape (1001,) +ke = traces["E_K3"] # kinetic energy, shape (1001,) +``` + +--- + +### `read_case_metadata(case_dir)` + +Return a combined summary of a case's parameters and equilibrium geometry. +Useful as the first call when an agent encounters an unfamiliar case. + +| Parameter | Type | Description | +|------------|------|-----------------------| +| `case_dir` | path | M3D-C1 case directory | + +Returns `dict` with keys: `params`, `snapshots`, `final_time`, `R_mag`, +`Z_mag`, `R_xpoint`, `Z_xpoint`, `psi_min`, `psi_lcfs`. + +```python +meta = read_case_metadata("m3dc1_data") +# meta["params"]["ntor"] → 9 +# meta["R_mag"] → 1.855 m +# meta["snapshots"] → [0, 1] +``` + +--- + +## Mesh and Evaluation Grid + +### `read_mesh_vertices(c1h5_path)` + +Extract the unique (R, Z) mesh vertex positions from the `mesh/elements` +array in an M3D-C1 HDF5 file. + +| Parameter | Type | Description | +|--------------|------|------------------------------------------------------| +| `c1h5_path` | path | Path to `C1.h5`, `equilibrium.h5`, or `time_NNN.h5` | + +Returns `(R, Z)` — two 1-D `float32` arrays sorted lexicographically. + +```python +R, Z = read_mesh_vertices("m3dc1_data/C1.h5") +# R shape: (n_vertices,) R ∈ [1.25, 2.55] m +``` + +--- + +### `make_evaluation_grid(R_mesh, Z_mesh, mode="mesh", grid_res=200, phi=0.0)` + +Build `(R, Z, phi)` arrays ready for passing to `eval_field()`. + +| Parameter | Type | Description | +|------------|---------|--------------------------------------------------------------| +| `R_mesh` | ndarray | Mesh vertex R coordinates from `read_mesh_vertices` | +| `Z_mesh` | ndarray | Mesh vertex Z coordinates | +| `mode` | str | `"mesh"` (use vertices) or `"grid"` (regular Cartesian grid) | +| `grid_res` | int | Grid points per axis for `mode="grid"` | +| `phi` | float | Toroidal angle in radians applied to all points | + +Returns `(R, Z, phi_arr)` with identical shapes. **Note:** `eval_field` takes +arguments as `(name, R, phi, Z, ...)` — pass `phi_arr` before `Z`. + +```python +R, Z = read_mesh_vertices("m3dc1_data/C1.h5") +R_eval, Z_eval, phi_eval = make_evaluation_grid(R, Z, mode="mesh") +# or: R_eval, Z_eval, phi_eval = make_evaluation_grid(R, Z, mode="grid", grid_res=200) +``` + +--- + +## Growth Rate + +### `compute_ke_growth_trace(case_dir)` + +Return the total kinetic energy `E_K3` and corresponding time array directly +from `C1.h5/scalars/`. No m3dc1 dependency. + +| Parameter | Type | Description | +|------------|------|-----------------------| +| `case_dir` | path | M3D-C1 case directory | + +Returns `(time, ke)` — two 1-D float arrays of length `ntimestep + 1`. + +```python +t, ke = compute_ke_growth_trace("m3dc1_data") +# ke[-1] / ke[1] ≈ 7 × 10⁵ (mode amplitude growth over 1000 τ_A) +``` + +--- + +### `compute_growth_rate(case_dir, time_idx=None)` + +Estimate the linear growth rate γ = 0.5 · d(ln E_K3)/dt, averaged over the +kinetic energy trace. No m3dc1 dependency. + +| Parameter | Type | Description | +|------------|----------|------------------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `time_idx` | int|None | Snapshot index to truncate the trace at. `None` = full. | + +Returns `float` in units of 1/τ_A. Returns `nan` if insufficient data. + +```python +gamma = compute_growth_rate("m3dc1_data") +# → 0.0067 τ_A⁻¹ (n=9 mode growth rate) +``` + +--- + +## Equilibrium Profiles *(requires m3dc1 + fpy)* + +### `compute_flux_average_profiles(case_dir, fields=None, fcoords="pest", points=200)` + +Compute flux-surface-averaged radial profiles for equilibrium quantities. + +| Parameter | Type | Description | +|------------|-----------|-------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `fields` | list|None | Field names. Default: `["p", "j", "ne", "q"]` | +| `fcoords` | str | Flux coordinate system: `"pest"` or `"equal_arc"` | +| `points` | int | Number of radial grid points | + +Returns `dict[str, (psi_norm, profile)]` where each value is a pair of 1-D +`float64` arrays of length `points`. Fields that fail are omitted. + +```python +profiles = compute_flux_average_profiles("m3dc1_data") +psin, q = profiles["q"] # safety factor profile vs psi_norm +``` + +--- + +### `compute_q95(psin, q_profile)` + +Interpolate the safety factor at psi_norm = 0.95. + +| Parameter | Type | Description | +|-------------|---------|-----------------------------------------------| +| `psin` | ndarray | Normalised flux values (0 to ≥ 0.95) | +| `q_profile` | ndarray | Safety factor at each psin | + +Returns `float`. Returns `nan` if the profile does not reach psi_norm = 0.95. + +```python +psin, q = profiles["q"] +q95 = compute_q95(psin, q) # → e.g. 3.2 +``` + +--- + +### `compute_miller_geometry(case_dir, res=250)` + +Compute the Miller geometry parameters of the last closed flux surface. + +| Parameter | Type | Description | +|------------|------|---------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `res` | int | Poloidal resolution for LCFS tracing | + +Returns `dict` with keys `"R0"` (m), `"a"` (m), `"kappa"`, `"delta"`. +Returns `{}` on failure. + +```python +shape = compute_miller_geometry("m3dc1_data") +# {"R0": 1.85, "a": 0.57, "kappa": 1.85, "delta": 0.42} +``` + +--- + +## Perturbed Fields *(requires m3dc1 + fpy)* + +### `compute_perturbed_fields(case_dir, time_idx, R, Z, phi, fields="all", skip_fields=None)` + +Evaluate perturbed fields (Δf = f_linear − f_equilibrium) at arbitrary +(R, Z, φ) points for one time snapshot. + +| Parameter | Type | Description | +|---------------|-------------|---------------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `time_idx` | int | Snapshot index | +| `R` | ndarray | R evaluation coordinates | +| `Z` | ndarray | Z evaluation coordinates (same shape as R) | +| `phi` | ndarray | Toroidal angles in radians (same shape as R) | +| `fields` | str or list | `"all"` or list of field names | +| `skip_fields` | list|None | Fields to skip. Default: `{"alpha", "gradA", "kprad_rad"}` | + +Returns `dict[str, ndarray]`. Vector fields (B, E, A, j, v) produce keys like +`"BR"`, `"BPHI"`, `"BZ"`. Shapes match R. + +```python +R, Z = read_mesh_vertices("m3dc1_data/C1.h5") +R, Z, phi = make_evaluation_grid(R, Z) +perts = compute_perturbed_fields("m3dc1_data", 1, R, Z, phi, fields=["psi", "B"]) +# perts["psi"] → 1-D array, same length as R +# perts["BR"] → R-component of perturbed B +``` + +--- + +## Field Evaluation on a Grid *(requires m3dc1 + fpy)* + +### `evaluate_field_on_grid(case_dir, fields, time_idx=-1, coord="scalar", mode="mesh", grid_res=200, phi=0.0, output_path=None)` + +Evaluate one or more fields on an (R, Z) grid derived from the mesh vertices or a regular +Cartesian grid, and optionally save the result to an HDF5 file. + +| Parameter | Type | Description | +|---------------|-------------------|-------------| +| `case_dir` | path | M3D-C1 case directory | +| `fields` | str or list[str] | Field name(s), e.g. `"te"` or `["te", "ne", "p"]` | +| `time_idx` | int | `-1` for equilibrium; non-negative integer for a snapshot | +| `coord` | str | `"scalar"` (value or magnitude), `"R"` / `"phi"` / `"Z"` (vector component), or `"all"` (see below) | +| `mode` | str | `"mesh"` (unique mesh vertices, 1-D output) or `"grid"` (regular grid, 2-D output) | +| `grid_res` | int | Points per axis for `mode="grid"` | +| `phi` | float | Toroidal angle in radians | +| `output_path` | path \| None | If given, write results to an HDF5 file at this path | + +**`coord="all"`** behaviour: + +- Scalar fields (e.g. `te`, `p`, `ne`): evaluated as `"scalar"`, stored under the field name. +- Vector fields (`B`, `E`, `A`, `j`, `v`): all three cylindrical components are evaluated and + stored as separate keys — e.g. `"BR"`, `"BPHI"`, `"BZ"` for `field="B"`. + +Returns `dict[str, np.ndarray]` with keys `"R"`, `"Z"`, and one key per field or component. +Array shapes are 1-D for `mode="mesh"` or 2-D `(grid_res, grid_res)` for `mode="grid"`. + +```python +# Equilibrium temperature and density on the mesh vertex grid, saved to file +result = evaluate_field_on_grid( + "m3dc1_data", ["te", "ne", "p"], + time_idx=-1, mode="mesh", + output_path="equilibrium_scalars.h5", +) +# result["te"] → 1-D array at each mesh vertex + +# All components of B on a 200×200 Cartesian grid at snapshot 1 +result = evaluate_field_on_grid( + "m3dc1_data", "B", + time_idx=1, coord="all", mode="grid", grid_res=200, + output_path="B_components.h5", +) +# result["BR"], result["BPHI"], result["BZ"] → each (200, 200) +``` + +--- + +## Spectral Analysis *(requires m3dc1 + fpy)* + +### `compute_poloidal_spectrum(case_dir, time_idx, field, coord="scalar", fcoords="pest", points=200, full_fft=False)` + +Compute the poloidal mode (m) spectrum of a field at a given snapshot. + +| Parameter | Type | Description | +|-------------|------|----------------------------------------------------------| +| `case_dir` | path | M3D-C1 case directory | +| `time_idx` | int | Snapshot index | +| `field` | str | Field name: `"p"`, `"B"`, etc. | +| `coord` | str | Component: `"scalar"`, `"R"`, `"Z"`, `"phi"` | +| `fcoords` | str | Flux coordinate system | +| `points` | int | Radial resolution | +| `full_fft` | bool | `True` = full two-sided FFT; `False` = mirrored spectrum | + +Returns `(m_modes, psi_norm, spectrum)`: +- `m_modes`: 1-D int array of poloidal mode numbers +- `psi_norm`: 1-D float array of length `points` +- `spectrum`: 2-D float array, shape `(len(m_modes), points)` + +```python +m, psi, spec = compute_poloidal_spectrum("m3dc1_data", 1, "p", points=200) +# spec[i, j] = amplitude of mode m[i] at psi_norm[j] +``` + +--- + +### `compute_standard_spectra(case_dir, time_idx, fcoords="pest", points=200, full_fft=False)` + +Compute poloidal spectra for the default set of four fields: pressure (p), +and the R, Z, and toroidal components of B. + +| Parameter | Type | Description | +|------------|------|-----------------------| +| `case_dir` | path | M3D-C1 case directory | +| `time_idx` | int | Snapshot index | +| `fcoords` | str | Flux coordinate system | +| `points` | int | Radial resolution | +| `full_fft` | bool | Full two-sided FFT | + +Returns `dict` with keys `"p"`, `"br"`, `"bz"`, `"bphi"`. Each value is a +`(m_modes, psi_norm, spectrum)` tuple. Failed fields are absent. + +```python +spectra = compute_standard_spectra("m3dc1_data", 1, points=200) +m, psi, spec_p = spectra["p"] +``` + diff --git a/use_cases/tokamak_stability/test/conftest.py b/use_cases/tokamak_stability/test/conftest.py new file mode 100644 index 0000000..9196008 --- /dev/null +++ b/use_cases/tokamak_stability/test/conftest.py @@ -0,0 +1,6 @@ +import sys +from pathlib import Path + +# Make the use-case directory importable so tests can do +# "from hdf5 import ..." and "from m3dc1_tools import ..." +sys.path.insert(0, str(Path(__file__).parent.parent)) diff --git a/use_cases/tokamak_stability/test_hdf5_tools.py b/use_cases/tokamak_stability/test/test_hdf5_tools.py similarity index 71% rename from use_cases/tokamak_stability/test_hdf5_tools.py rename to use_cases/tokamak_stability/test/test_hdf5_tools.py index ee4bdec..a94dee4 100644 --- a/use_cases/tokamak_stability/test_hdf5_tools.py +++ b/use_cases/tokamak_stability/test/test_hdf5_tools.py @@ -1,6 +1,7 @@ import h5py +import numpy as np from pathlib import Path -from dsagt.tools.hdf5 import list_h5_files, list_h5_variables, repackage_h5 +from hdf5 import list_h5_files, list_h5_variables, read_h5_dataset, read_h5_attrs, repackage_h5 def _make_file(path: Path, datasets: dict): @@ -180,3 +181,82 @@ def test_repackage_attributes_preserved(tmp_path): with h5py.File(out, 'r') as f: assert f["dset"].attrs.get("units") == "m" + + +# --------------------------------------------------------------------------- +# read_h5_dataset +# --------------------------------------------------------------------------- + +def test_read_h5_dataset_basic(tmp_path): + src = tmp_path / "src.h5" + data = np.array([1.0, 2.0, 3.0]) + _make_file(src, {"mydata": data}) + result = read_h5_dataset(src, "mydata") + np.testing.assert_array_equal(result, data) + + +def test_read_h5_dataset_nested(tmp_path): + src = tmp_path / "src.h5" + data = np.array([[1, 2], [3, 4]]) + _make_file(src, {"group/nested": data}) + result = read_h5_dataset(src, "group/nested") + np.testing.assert_array_equal(result, data) + + +def test_read_h5_dataset_missing_key_raises(tmp_path): + src = tmp_path / "src.h5" + _make_file(src, {"existing": [1]}) + import pytest + with pytest.raises(KeyError): + read_h5_dataset(src, "nonexistent") + + +def test_read_h5_dataset_preserves_dtype(tmp_path): + src = tmp_path / "src.h5" + data = np.array([1, 2, 3], dtype=np.float32) + _make_file(src, {"vals": data}) + result = read_h5_dataset(src, "vals") + assert result.dtype == np.float32 + + +# --------------------------------------------------------------------------- +# read_h5_attrs +# --------------------------------------------------------------------------- + +def test_read_h5_attrs_root(tmp_path): + src = tmp_path / "src.h5" + with h5py.File(src, "w") as f: + f.attrs["version"] = np.int32(45) + f.attrs["time"] = np.float64(1.5) + result = read_h5_attrs(src) + assert result["version"] == 45 + assert isinstance(result["version"], int) + assert abs(result["time"] - 1.5) < 1e-10 + assert isinstance(result["time"], float) + + +def test_read_h5_attrs_group(tmp_path): + src = tmp_path / "src.h5" + with h5py.File(src, "w") as f: + grp = f.create_group("mygroup") + grp.attrs["label"] = "hello" + grp.attrs["count"] = np.int32(7) + result = read_h5_attrs(src, "mygroup") + assert result["label"] == "hello" + assert result["count"] == 7 + + +def test_read_h5_attrs_empty(tmp_path): + src = tmp_path / "src.h5" + with h5py.File(src, "w") as f: + f.create_group("empty") + result = read_h5_attrs(src, "empty") + assert result == {} + + +def test_read_h5_attrs_missing_group_raises(tmp_path): + src = tmp_path / "src.h5" + _make_file(src, {"x": [1]}) + import pytest + with pytest.raises(KeyError): + read_h5_attrs(src, "nosuchgroup") diff --git a/use_cases/tokamak_stability/test/test_m3dc1_plots.py b/use_cases/tokamak_stability/test/test_m3dc1_plots.py new file mode 100644 index 0000000..ff2716d --- /dev/null +++ b/use_cases/tokamak_stability/test/test_m3dc1_plots.py @@ -0,0 +1,304 @@ +"""Tests for m3dc1_plots.py. + +Unit tests run without real data or m3dc1. Integration tests (marked +``integration``) require real data at the path in CASE_DIR and a working +m3dc1 + fpy installation. +""" +import math +import textwrap +from pathlib import Path + +import numpy as np +import pytest + +import m3dc1_plots +from m3dc1_plots import _parse_geqdsk, _save_first_new_fig + +CASE_DIR = Path("/Users/kparfrey/data/asv/run1/sparc_1425") + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def minimal_geqdsk(tmp_path) -> Path: + """Write a tiny but structurally valid GEQDSK file.""" + nw, nh = 4, 5 + header = f" EQ 0 0 0 {nw} {nh}\n" + # rdim zdim rcentr rleft zmid + scalars_1 = " 1.0 2.0 1.5 0.5 0.0\n" + # rmaxis zmaxis simag sibry bcentr + scalars_2 = " 1.5 0.0 0.1 0.9 2.0\n" + # current + 4 padding zeros (first 5 of the next block) + scalars_3 = " 1.0e6 0.0 0.0 0.0 0.0\n" + + def _row(vals): + return " ".join(f"{v: .6e}" for v in vals) + "\n" + + fpol = [1.0] * nw + pres = [0.0] * nw + ffprim = [0.0] * nw + pprim = [0.0] * nw + psirz = list(np.linspace(0.1, 0.9, nw * nh)) + qpsi = [2.0] * nw + + nbbbs, limitr = 3, 2 + rbbbs = [1.2, 1.5, 1.8] + zbbbs = [-0.5, 0.5, 0.0] + rlim = [1.0, 2.0] + zlim = [-1.0, 1.0] + + lines = [header, scalars_1, scalars_2, scalars_3] + + def _write_1d(lst): + chunk = 5 + for i in range(0, len(lst), chunk): + lines.append(_row(lst[i: i + chunk])) + + _write_1d(fpol) + _write_1d(pres) + _write_1d(ffprim) + _write_1d(pprim) + _write_1d(psirz) + _write_1d(qpsi) + lines.append(f" {nbbbs} {limitr}\n") + bdry = [] + for r, z in zip(rbbbs, zbbbs): + bdry += [r, z] + _write_1d(bdry) + lim = [] + for r, z in zip(rlim, zlim): + lim += [r, z] + _write_1d(lim) + + gfile = tmp_path / "geqdsk" + gfile.write_text("".join(lines)) + return gfile + + +@pytest.fixture +def fake_case_dir(tmp_path) -> Path: + """Create a minimal case directory with a synthetic C1.h5 for h5py-only tests.""" + import h5py + case_dir = tmp_path / "sparc_test" + case_dir.mkdir() + + n_steps = 20 + time = np.linspace(0, 10, n_steps + 1) + ke = 1e-6 * np.exp(0.5 * time) + + with h5py.File(case_dir / "C1.h5", "w") as f: + sc = f.create_group("scalars") + sc.create_dataset("time", data=time) + sc.create_dataset("E_K3", data=ke) + sc.create_dataset("xmag", data=[1.5]) + sc.create_dataset("zmag", data=[0.0]) + sc.create_dataset("xnull", data=[1.3]) + sc.create_dataset("znull", data=[-1.2]) + sc.create_dataset("psimin", data=[0.1]) + sc.create_dataset("psi_lcfs", data=[0.9]) + + return case_dir + + +# --------------------------------------------------------------------------- +# Unit tests — _parse_geqdsk helper +# --------------------------------------------------------------------------- + +class TestParseGeqdsk: + def test_basic_structure(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + assert "rg" in geo and "zg" in geo + assert "psirzn" in geo + assert "rmaxis" in geo + + def test_grid_shape(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + nw, nh = 4, 5 + assert len(geo["rg"]) == nw + assert len(geo["zg"]) == nh + assert geo["psirz"].shape == (nh, nw) + + def test_psirzn_normalisation(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + # psirzn should span 0..1 given simag=0.1, sibry=0.9, psirz=linspace(0.1,0.9) + assert float(np.nanmin(geo["psirzn"])) == pytest.approx(0.0, abs=0.01) + assert float(np.nanmax(geo["psirzn"])) == pytest.approx(1.0, abs=0.01) + + def test_boundary_parsed(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + assert len(geo["rbbbs"]) == 3 + assert len(geo["zbbbs"]) == 3 + + def test_limiter_parsed(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + assert len(geo["rlim"]) == 2 + + def test_z_grid_centred(self, minimal_geqdsk): + geo = _parse_geqdsk(minimal_geqdsk) + # zmid=0, zdim=2 → zg should span -1..1 + assert geo["zg"][0] == pytest.approx(-1.0, abs=1e-9) + assert geo["zg"][-1] == pytest.approx(1.0, abs=1e-9) + + +# --------------------------------------------------------------------------- +# Unit tests — _save_first_new_fig helper +# --------------------------------------------------------------------------- + +class TestSaveFirstNewFig: + def test_saves_file(self, tmp_path): + import matplotlib.pyplot as plt + before = set(plt.get_fignums()) + fig = plt.figure() + out = tmp_path / "fig.png" + _save_first_new_fig(before, out, dpi=72) + assert out.exists() + assert out.stat().st_size > 0 + + def test_closes_all_new_figs(self, tmp_path): + import matplotlib.pyplot as plt + before = set(plt.get_fignums()) + plt.figure() + plt.figure() + out = tmp_path / "fig.png" + _save_first_new_fig(before, out, dpi=72) + after = set(plt.get_fignums()) + assert after == before + + def test_raises_if_no_new_fig(self, tmp_path): + import matplotlib.pyplot as plt + before = set(plt.get_fignums()) + out = tmp_path / "fig.png" + with pytest.raises(RuntimeError, match="no new figure"): + _save_first_new_fig(before, out, dpi=72) + + +# --------------------------------------------------------------------------- +# Unit tests — Category A functions with synthetic data +# --------------------------------------------------------------------------- + +class TestPlotKineticEnergy: + def test_creates_file(self, fake_case_dir, tmp_path): + out = tmp_path / "ke.png" + result = m3dc1_plots.plot_kinetic_energy(fake_case_dir, out) + assert result == out + assert out.exists() + assert out.stat().st_size > 100 + + def test_no_annotation_option(self, fake_case_dir, tmp_path): + out = tmp_path / "ke_no_ann.png" + m3dc1_plots.plot_kinetic_energy(fake_case_dir, out, annotate_growth_rate=False) + assert out.exists() + + def test_raises_on_all_zero_ke(self, tmp_path): + import h5py + case_dir = tmp_path / "zero_ke" + case_dir.mkdir() + with h5py.File(case_dir / "C1.h5", "w") as f: + sc = f.create_group("scalars") + sc.create_dataset("time", data=np.zeros(5)) + sc.create_dataset("E_K3", data=np.zeros(5)) + with pytest.raises(ValueError, match="non-zero"): + m3dc1_plots.plot_kinetic_energy(case_dir, tmp_path / "ke.png") + + +class TestPlotGrowthRateVsTime: + def test_creates_file(self, fake_case_dir, tmp_path): + out = tmp_path / "gr.png" + result = m3dc1_plots.plot_growth_rate_vs_time(fake_case_dir, out) + assert result == out + assert out.exists() + + +class TestPlotGeqdsk: + def test_creates_file(self, minimal_geqdsk, tmp_path): + out = tmp_path / "geqdsk.png" + result = m3dc1_plots.plot_geqdsk(minimal_geqdsk, out) + assert result == out + assert out.exists() + assert out.stat().st_size > 100 + + def test_returns_path_object(self, minimal_geqdsk, tmp_path): + out = tmp_path / "geqdsk.png" + result = m3dc1_plots.plot_geqdsk(minimal_geqdsk, out) + assert isinstance(result, Path) + + +# --------------------------------------------------------------------------- +# Integration tests — require real data + m3dc1/fpy +# --------------------------------------------------------------------------- + +pytestmark_integration = pytest.mark.integration + + +@pytest.mark.integration +class TestIntegrationCategoryA: + def test_plot_kinetic_energy(self, tmp_path): + out = tmp_path / "ke.png" + m3dc1_plots.plot_kinetic_energy(CASE_DIR, out) + assert out.exists() + + def test_plot_growth_rate_vs_time(self, tmp_path): + out = tmp_path / "gr.png" + m3dc1_plots.plot_growth_rate_vs_time(CASE_DIR, out) + assert out.exists() + + def test_plot_flux_average_profiles(self, tmp_path): + out = tmp_path / "profiles.png" + m3dc1_plots.plot_flux_average_profiles(CASE_DIR, out) + assert out.exists() + + def test_plot_safety_factor(self, tmp_path): + out = tmp_path / "q.png" + m3dc1_plots.plot_safety_factor(CASE_DIR, out) + assert out.exists() + + def test_plot_poloidal_spectrum(self, tmp_path): + out = tmp_path / "spec.png" + m3dc1_plots.plot_poloidal_spectrum(CASE_DIR, 1, "p", out) + assert out.exists() + + def test_plot_standard_spectra(self, tmp_path): + out = tmp_path / "spectra.png" + m3dc1_plots.plot_standard_spectra(CASE_DIR, 1, out) + assert out.exists() + + def test_plot_perturbed_field_map(self, tmp_path): + out = tmp_path / "field_psi.png" + m3dc1_plots.plot_perturbed_field_map(CASE_DIR, 1, "psi", out) + assert out.exists() + + def test_plot_geqdsk_uses_case_dir_gfile(self, tmp_path): + out = tmp_path / "geqdsk.png" + gfile = CASE_DIR / "geqdsk" + if not gfile.exists(): + pytest.skip("geqdsk file not present in case dir") + m3dc1_plots.plot_geqdsk(gfile, out) + assert out.exists() + + +@pytest.mark.integration +class TestIntegrationWrappers: + def test_plot_stability_summary(self, tmp_path): + paths = m3dc1_plots.plot_stability_summary(CASE_DIR, 1, tmp_path) + assert len(paths) >= 2 + for p in paths: + assert p.exists() + + def test_plot_equilibrium_overview(self, tmp_path): + paths = m3dc1_plots.plot_equilibrium_overview(CASE_DIR, tmp_path) + assert len(paths) >= 1 + for p in paths: + assert p.exists() + + def test_plot_case_summary_returns_combined(self, tmp_path): + paths = m3dc1_plots.plot_case_summary(CASE_DIR, 1, tmp_path) + assert len(paths) >= 2 + + def test_prefix_applied(self, tmp_path): + paths = m3dc1_plots.plot_stability_summary( + CASE_DIR, 1, tmp_path, prefix="run1_" + ) + for p in paths: + assert p.name.startswith("run1_") diff --git a/use_cases/tokamak_stability/test/test_m3dc1_tools.py b/use_cases/tokamak_stability/test/test_m3dc1_tools.py new file mode 100644 index 0000000..a67dd2f --- /dev/null +++ b/use_cases/tokamak_stability/test/test_m3dc1_tools.py @@ -0,0 +1,494 @@ +"""Tests for m3dc1_tools.py. + +Unit tests use temporary fixtures and run without any external data or +libraries. Integration tests require the example dataset at DATA_DIR and are +automatically skipped when it is absent. Tests for m3dc1/fpy-dependent +functions are additionally skipped when those libraries are not installed. +""" +import pytest +import numpy as np +import h5py +from pathlib import Path + +from m3dc1_tools import ( + read_c1input, + list_time_snapshots, + read_snapshot_time, + read_scalar_traces, + read_case_metadata, + read_mesh_vertices, + make_evaluation_grid, + compute_ke_growth_trace, + compute_growth_rate, + compute_q95, +) + +DATA_DIR = Path("/Users/kparfrey/data/asv/run1/sparc_1425") + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def sparc_case(): + if not DATA_DIR.exists(): + pytest.skip(f"Example data not available at {DATA_DIR}") + return DATA_DIR + + +@pytest.fixture +def require_m3dc1(): + try: + import m3dc1 # noqa: F401 + import fpy # noqa: F401 + except ImportError: + pytest.skip("m3dc1/fpy not installed") + + +@pytest.fixture +def minimal_c1input(tmp_path): + """Write a minimal C1input with all tracked parameters.""" + content = """&inputnl +\tntimemax = 500\t! steps +\tntimepr = 250\t! output interval +\tdt = 0.5\t! time step +\tlinear = 1\t! linear run +\tnumvar = 3\t! 6-field +\tion_mass = 2.\t! deuterium +\tzeff = 1.\t! Z_eff +pscale = 0.815 +batemanscale = 1.063 +ntor = 9 + / +""" + (tmp_path / "C1input").write_text(content) + return tmp_path + + +@pytest.fixture +def case_with_snapshots(tmp_path): + """Case directory containing dummy time_NNN.h5 snapshot files.""" + for idx in (0, 2, 5): + fname = tmp_path / f"time_{idx:03d}.h5" + with h5py.File(fname, "w") as f: + f.attrs["ntimestep"] = idx * 100 + f.attrs["time"] = float(idx * 100) + return tmp_path + + +# --------------------------------------------------------------------------- +# read_c1input — unit tests +# --------------------------------------------------------------------------- + +def test_read_c1input_all_params(minimal_c1input): + params = read_c1input(minimal_c1input) + assert params["ntor"] == 9 + assert params["ntimemax"] == 500 + assert params["ntimepr"] == 250 + assert abs(params["dt"] - 0.5) < 1e-9 + assert params["linear"] == 1 + assert params["numvar"] == 3 + assert abs(params["ion_mass"] - 2.0) < 1e-9 + assert abs(params["zeff"] - 1.0) < 1e-9 + assert abs(params["pscale"] - 0.815) < 1e-6 + assert abs(params["batemanscale"] - 1.063) < 1e-6 + + +def test_read_c1input_missing_file(tmp_path): + params = read_c1input(tmp_path) + assert all(v is None for v in params.values()) + + +def test_read_c1input_partial_params(tmp_path): + (tmp_path / "C1input").write_text("&inputnl\n\tntor = 3\n /\n") + params = read_c1input(tmp_path) + assert params["ntor"] == 3 + assert params["pscale"] is None + assert params["linear"] is None + + +def test_read_c1input_returns_int_for_int_keys(minimal_c1input): + params = read_c1input(minimal_c1input) + assert isinstance(params["ntor"], int) + assert isinstance(params["linear"], int) + assert isinstance(params["numvar"], int) + + +def test_read_c1input_returns_float_for_float_keys(minimal_c1input): + params = read_c1input(minimal_c1input) + assert isinstance(params["pscale"], float) + assert isinstance(params["batemanscale"], float) + + +def test_read_c1input_comment_lines_ignored(tmp_path): + content = "! ntor = 999\n&inputnl\nntor = 5\n /\n" + (tmp_path / "C1input").write_text(content) + params = read_c1input(tmp_path) + assert params["ntor"] == 5 + + +# --------------------------------------------------------------------------- +# list_time_snapshots — unit tests +# --------------------------------------------------------------------------- + +def test_list_time_snapshots_finds_files(case_with_snapshots): + result = list_time_snapshots(case_with_snapshots) + assert result == [0, 2, 5] + + +def test_list_time_snapshots_empty_dir(tmp_path): + assert list_time_snapshots(tmp_path) == [] + + +def test_list_time_snapshots_ignores_non_numeric(tmp_path): + with h5py.File(tmp_path / "time_abc.h5", "w"): + pass + with h5py.File(tmp_path / "time_001.h5", "w") as f: + f.attrs["time"] = 1.0 + result = list_time_snapshots(tmp_path) + assert result == [1] + + +def test_list_time_snapshots_deduplicated(tmp_path): + """Only the file indices matter; no duplicates.""" + for name in ("time_003.h5", "time_003.h5"): + with h5py.File(tmp_path / name, "w"): + pass + assert list_time_snapshots(tmp_path) == [3] + + +# --------------------------------------------------------------------------- +# read_snapshot_time — unit tests +# --------------------------------------------------------------------------- + +def test_read_snapshot_time_from_file(case_with_snapshots): + t = read_snapshot_time(case_with_snapshots, 2) + assert abs(t - 200.0) < 1e-6 + + +def test_read_snapshot_time_missing_returns_nan(tmp_path): + t = read_snapshot_time(tmp_path, 99) + assert np.isnan(t) + + +def test_read_snapshot_time_from_c1h5(tmp_path): + """Falls back to reading from C1.h5/time_NNN group.""" + with h5py.File(tmp_path / "C1.h5", "w") as f: + grp = f.create_group("time_001") + grp.attrs["time"] = 42.0 + t = read_snapshot_time(tmp_path, 1) + assert abs(t - 42.0) < 1e-6 + + +# --------------------------------------------------------------------------- +# make_evaluation_grid — unit tests +# --------------------------------------------------------------------------- + +def _sample_mesh(): + rng = np.random.default_rng(0) + R = rng.uniform(1.0, 2.0, 50).astype(np.float32) + Z = rng.uniform(-1.0, 1.0, 50).astype(np.float32) + return R, Z + + +def test_make_evaluation_grid_mesh_mode(): + R_m, Z_m = _sample_mesh() + R, Z, phi_arr = make_evaluation_grid(R_m, Z_m, mode="mesh") + assert R.shape == R_m.shape + assert Z.shape == Z_m.shape + assert phi_arr.shape == R_m.shape + np.testing.assert_array_equal(R, R_m) + + +def test_make_evaluation_grid_grid_mode(): + R_m, Z_m = _sample_mesh() + R, Z, phi_arr = make_evaluation_grid(R_m, Z_m, mode="grid", grid_res=10) + assert R.shape == (10, 10) + assert Z.shape == (10, 10) + assert phi_arr.shape == (10, 10) + + +def test_make_evaluation_grid_phi_constant(): + R_m, Z_m = _sample_mesh() + _, _, phi_arr = make_evaluation_grid(R_m, Z_m, phi=1.23) + assert np.all(phi_arr == 1.23) + + +def test_make_evaluation_grid_bounding_box(): + R_m, Z_m = _sample_mesh() + R, Z, _ = make_evaluation_grid(R_m, Z_m, mode="grid", grid_res=20) + assert R.min() >= R_m.min() - 1e-6 + assert R.max() <= R_m.max() + 1e-6 + assert Z.min() >= Z_m.min() - 1e-6 + assert Z.max() <= Z_m.max() + 1e-6 + + +def test_make_evaluation_grid_invalid_mode(): + R_m, Z_m = _sample_mesh() + with pytest.raises(ValueError): + make_evaluation_grid(R_m, Z_m, mode="invalid") + + +# --------------------------------------------------------------------------- +# compute_q95 — unit tests +# --------------------------------------------------------------------------- + +def test_compute_q95_analytical(): + psin = np.linspace(0, 1, 11) + q = 1 + 2 * psin # q=1 at axis, q=3 at edge → q95 = 1 + 2*0.95 = 2.9 + assert abs(compute_q95(psin, q) - 2.9) < 1e-10 + + +def test_compute_q95_below_range_returns_nan(): + psin = np.linspace(0, 0.9, 10) # does not reach 0.95 + q = np.ones(10) + assert np.isnan(compute_q95(psin, q)) + + +def test_compute_q95_empty_returns_nan(): + assert np.isnan(compute_q95(np.array([]), np.array([]))) + + +def test_compute_q95_at_boundary(): + psin = np.array([0.0, 0.95, 1.0]) + q = np.array([1.0, 2.5, 3.0]) + assert abs(compute_q95(psin, q) - 2.5) < 1e-10 + + +# --------------------------------------------------------------------------- +# compute_growth_rate — unit tests +# --------------------------------------------------------------------------- + +def _make_c1h5_with_scalars(tmp_path, ke_values, time_values=None): + if time_values is None: + time_values = np.arange(len(ke_values), dtype=float) + with h5py.File(tmp_path / "C1.h5", "w") as f: + sc = f.create_group("scalars") + sc.create_dataset("E_K3", data=np.array(ke_values, dtype=np.float32)) + sc.create_dataset("time", data=np.array(time_values, dtype=np.float32)) + return tmp_path + + +def test_compute_growth_rate_nan_on_zeros(tmp_path): + _make_c1h5_with_scalars(tmp_path, [0.0, 0.0, 0.0]) + assert np.isnan(compute_growth_rate(tmp_path)) + + +def test_compute_growth_rate_nan_on_single_nonzero(tmp_path): + _make_c1h5_with_scalars(tmp_path, [0.0, 1.0]) + assert np.isnan(compute_growth_rate(tmp_path)) + + +def test_compute_growth_rate_exponential_growth(tmp_path): + """For E_K3 = exp(2*gamma*t), the recovered rate should equal gamma.""" + gamma_true = 0.05 + t = np.arange(0, 101, dtype=float) + ke = np.exp(2 * gamma_true * t) + ke[0] = 0.0 # initial timestep often zero + _make_c1h5_with_scalars(tmp_path, ke, t) + gamma_est = compute_growth_rate(tmp_path) + assert abs(gamma_est - gamma_true) < 1e-4 + + +def test_compute_ke_growth_trace_shape(tmp_path): + ke = np.linspace(1e-14, 1e-7, 1001) + t = np.arange(1001, dtype=float) + _make_c1h5_with_scalars(tmp_path, ke, t) + time_out, ke_out = compute_ke_growth_trace(tmp_path) + assert len(time_out) == 1001 + assert len(ke_out) == 1001 + + +# --------------------------------------------------------------------------- +# Integration tests — require real data at DATA_DIR +# --------------------------------------------------------------------------- + +def test_read_c1input_sparc1425(sparc_case): + params = read_c1input(sparc_case) + assert params["ntor"] == 9 + assert params["linear"] == 1 + assert params["numvar"] == 3 + assert abs(params["pscale"] - 0.815) < 0.01 + assert abs(params["ion_mass"] - 2.0) < 0.01 + + +def test_list_time_snapshots_sparc1425(sparc_case): + snaps = list_time_snapshots(sparc_case) + assert snaps == [0, 1] + + +def test_read_snapshot_time_alfven(sparc_case): + t = read_snapshot_time(sparc_case, 0, units="alfven") + assert abs(t - 0.0) < 1e-6 + t1 = read_snapshot_time(sparc_case, 1, units="alfven") + assert abs(t1 - 1000.0) < 1.0 + + +def test_read_scalar_traces_subset(sparc_case): + traces = read_scalar_traces(sparc_case, names=["E_K3", "time"]) + assert "E_K3" in traces + assert "time" in traces + assert len(traces["time"]) == 1001 + assert traces["E_K3"][-1] > traces["E_K3"][1] # mode grows + + +def test_read_scalar_traces_all(sparc_case): + traces = read_scalar_traces(sparc_case) + assert len(traces) > 50 + assert "xmag" in traces + + +def test_read_case_metadata_sparc1425(sparc_case): + meta = read_case_metadata(sparc_case) + assert meta["params"]["ntor"] == 9 + assert meta["snapshots"] == [0, 1] + assert abs(meta["final_time"] - 1000.0) < 1.0 + assert abs(meta["R_mag"] - 1.855) < 0.01 + assert abs(meta["Z_mag"]) < 0.01 + assert abs(meta["R_xpoint"] - 1.492) < 0.01 + assert meta["psi_min"] < 0 # negative at axis + assert meta["psi_lcfs"] > meta["psi_min"] + + +def test_read_mesh_vertices_sparc1425(sparc_case): + R, Z = read_mesh_vertices(sparc_case / "C1.h5") + assert R.shape == Z.shape + assert len(R) > 1000 + assert len(R) < 18555 # fewer unique vertices than elements + assert R.min() >= 1.24 + assert R.max() <= 2.56 + assert Z.min() >= -1.26 + assert Z.max() <= 1.26 + + +def test_read_mesh_vertices_equilibrium_h5(sparc_case): + R1, Z1 = read_mesh_vertices(sparc_case / "C1.h5") + R2, Z2 = read_mesh_vertices(sparc_case / "equilibrium.h5") + np.testing.assert_array_equal(R1, R2) + np.testing.assert_array_equal(Z1, Z2) + + +def test_compute_ke_growth_trace_sparc1425(sparc_case): + time, ke = compute_ke_growth_trace(sparc_case) + assert len(time) == 1001 + assert len(ke) == 1001 + assert ke[-1] > ke[1] + + +def test_compute_growth_rate_sparc1425(sparc_case): + gamma = compute_growth_rate(sparc_case) + assert np.isfinite(gamma) + assert 0.001 < gamma < 0.1 # reasonable range for a linear MHD run + + +def test_compute_growth_rate_time_idx_sparc1425(sparc_case): + gamma_full = compute_growth_rate(sparc_case) + gamma_t1 = compute_growth_rate(sparc_case, time_idx=1) + # With all data, both should be close (time_idx=1 → ntimestep=1000 = full trace) + assert abs(gamma_full - gamma_t1) < 1e-6 + + +# --------------------------------------------------------------------------- +# Integration tests — additionally require m3dc1 + fpy +# --------------------------------------------------------------------------- + +def test_compute_flux_average_profiles_sparc1425(sparc_case, require_m3dc1): + from m3dc1_tools import compute_flux_average_profiles + profiles = compute_flux_average_profiles(sparc_case) + assert "p" in profiles + assert "q" in profiles + psin, q = profiles["q"] + assert abs(psin[0]) < 0.05 + assert abs(psin[-1] - 1.0) < 0.05 + assert q.min() > 0.5 + + +def test_compute_q95_sparc1425(sparc_case, require_m3dc1): + from m3dc1_tools import compute_flux_average_profiles + profiles = compute_flux_average_profiles(sparc_case) + psin, q = profiles["q"] + q95 = compute_q95(psin, q) + assert np.isfinite(q95) + assert q95 > 1.0 + + +def test_compute_miller_geometry_sparc1425(sparc_case, require_m3dc1): + from m3dc1_tools import compute_miller_geometry + shape = compute_miller_geometry(sparc_case) + assert set(shape.keys()) == {"R0", "a", "kappa", "delta"} + assert 1.5 < shape["R0"] < 2.5 + assert shape["a"] > 0 + assert shape["kappa"] > 1.0 + + +def test_compute_perturbed_fields_psi(sparc_case, require_m3dc1): + from m3dc1_tools import compute_perturbed_fields + R, Z = read_mesh_vertices(sparc_case / "C1.h5") + R, Z, phi = make_evaluation_grid(R, Z, mode="mesh") + fields = compute_perturbed_fields(sparc_case, 1, R, Z, phi, fields=["psi"]) + assert "psi" in fields + assert fields["psi"].shape == R.shape + assert np.abs(fields["psi"]).max() > 0 + + +def test_compute_perturbed_fields_zero_at_t0(sparc_case, require_m3dc1): + from m3dc1_tools import compute_perturbed_fields + R, Z = read_mesh_vertices(sparc_case / "C1.h5") + R, Z, phi = make_evaluation_grid(R, Z, mode="mesh") + fields = compute_perturbed_fields(sparc_case, 0, R, Z, phi, fields=["psi"]) + assert "psi" in fields + assert np.abs(fields["psi"]).max() < 1e-3 + + +def test_compute_perturbed_fields_vector_B(sparc_case, require_m3dc1): + from m3dc1_tools import compute_perturbed_fields + R, Z = read_mesh_vertices(sparc_case / "C1.h5") + R, Z, phi = make_evaluation_grid(R, Z, mode="mesh") + fields = compute_perturbed_fields(sparc_case, 1, R, Z, phi, fields=["B"]) + assert "BR" in fields + assert "BPHI" in fields + assert "BZ" in fields + + +def test_compute_perturbed_fields_skip_respected(sparc_case, require_m3dc1): + from m3dc1_tools import compute_perturbed_fields + R, Z = read_mesh_vertices(sparc_case / "C1.h5") + R, Z, phi = make_evaluation_grid(R, Z, mode="mesh") + # psi should be absent when skipped + fields = compute_perturbed_fields( + sparc_case, 1, R, Z, phi, fields=["psi", "te"], skip_fields=["psi"] + ) + assert "psi" not in fields + assert "te" in fields + + +def test_compute_poloidal_spectrum_shape(sparc_case, require_m3dc1): + from m3dc1_tools import compute_poloidal_spectrum + m_modes, psi_norm, spec = compute_poloidal_spectrum(sparc_case, 1, "p", points=50) + assert m_modes.dtype == int or np.issubdtype(m_modes.dtype, np.integer) + assert len(psi_norm) == 50 + assert spec.shape == (len(m_modes), 50) + assert psi_norm[0] >= 0.0 + assert psi_norm[-1] <= 1.0 + 1e-6 + + +def test_compute_poloidal_spectrum_full_fft(sparc_case, require_m3dc1): + from m3dc1_tools import compute_poloidal_spectrum + m_half, _, spec_half = compute_poloidal_spectrum(sparc_case, 1, "p", points=50) + m_full, _, spec_full = compute_poloidal_spectrum( + sparc_case, 1, "p", points=50, full_fft=True + ) + # full FFT should have exactly `points` m-modes (not mirrored) + assert len(m_full) == 50 + # half-spectrum is mirrored so has 2*m_max+1 modes + assert len(m_half) > 50 + + +def test_compute_standard_spectra_keys(sparc_case, require_m3dc1): + from m3dc1_tools import compute_standard_spectra + result = compute_standard_spectra(sparc_case, 1, points=50) + assert set(result.keys()) == {"p", "br", "bz", "bphi"} + for key, (m, psi, spec) in result.items(): + assert spec.ndim == 2 + assert spec.shape[1] == 50