diff --git a/CHANGELOG.md b/CHANGELOG.md index b996afc..13a07cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,22 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.5.1] - 2026-07-28 + +### Added + +- Add synthetic JINTRAC radiation values dataset support and related regression tests +- Extend radiation emitter loading to support dual emissivity sources with improved validation + +### Changed + +- Improve radiation emitter loading checks for duplicate emissivity values and core-profile grid data + +### Fixed + +- Fix grid data loading checks for radiation core profiles +- Improve error handling in radiation emitter loading workflows + ## [0.5.0] - 2026-06-24 ### Added diff --git a/README.md b/README.md index 48f0340..db46d2d 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,9 @@ pip install cherab-imas mamba install -c conda-forge cherab-imas ``` +> [!NOTE] +> Some IMAS features that rely on the memory backend require a working `imas_core` runtime. On platforms where that backend is unavailable (for example, some macOS Intel environments), a subset of tests that exercise those features may be skipped automatically, while the core package functionality remains available. + ## 📝 Documentation See the [official documentation](https://cherab.github.io/imas/) to learn more. diff --git a/docs/notebooks/plasma/4_emission.ipynb b/docs/notebooks/plasma/4_emission.ipynb index 3c192e7..4c89862 100644 --- a/docs/notebooks/plasma/4_emission.ipynb +++ b/docs/notebooks/plasma/4_emission.ipynb @@ -721,7 +721,7 @@ " title=\"Ray-traced emission spectrum integrated over wavelength bands\",\n", " ylabel=\"Radiance [W/m²/sr]\",\n", " xlocator=range(len(wavelength_bands)),\n", - " xticklabels=[f\"{band[0]} – {band[1]}\" for band in wavelength_bands],\n", + " xticklabels=[f\"{band[0]:g} – {band[1]:g}\" for band in wavelength_bands],\n", ")\n", "\n", "axs.format(\n", diff --git a/docs/notebooks/radiation/radiation_3d.ipynb b/docs/notebooks/radiation/radiation_3d.ipynb index ba7ec45..9b66f18 100644 --- a/docs/notebooks/radiation/radiation_3d.ipynb +++ b/docs/notebooks/radiation/radiation_3d.ipynb @@ -21,12 +21,8 @@ "source": [ "import numpy as np\n", "import ultraplot as uplt\n", - "from imas import DBEntry\n", "from raysect.optical import World\n", - "from rich import print as rprint\n", - "from rich.table import Table\n", "\n", - "from cherab.core.math import sample3d_grid\n", "from cherab.imas.datasets import iter_jorek\n", "from cherab.imas.emitter import load_radiation_emitter\n", "\n", diff --git a/docs/source/api.md b/docs/source/api.md index 1dbe7ea..2f2ea0e 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -1,4 +1,6 @@ -# API Reference +(api-reference)= + +# 📚 API Reference This page contains auto-generated API reference documentation. diff --git a/docs/source/contributing.md b/docs/source/contributing.md index 39b3d69..aed6fc1 100644 --- a/docs/source/contributing.md +++ b/docs/source/contributing.md @@ -1,4 +1,6 @@ -# Contributing +(contributing)= + +# 🤝 Contributing We welcome contributions to `cherab-imas`! Whether you're fixing bugs, adding new features, or improving documentation, your help is appreciated. @@ -17,7 +19,7 @@ If you encounter any issues or have suggestions for improvements, please open an ### Prerequisites -Development tasks are managed using `pixi`. If you don't have `pixi` installed, please refer to the https://pixi.sh documentation for installation instructions. +Development tasks are managed using `pixi`. If you don't have `pixi` installed, please refer to the documentation for installation instructions. Other tools like `git` can be installed globally via `pixi`: ```bash @@ -74,7 +76,7 @@ If you want to host the documentation locally, you can do so with: pixi run doc-serve ``` -The documentation will be served at http://localhost:8000. +The documentation will be served at . ::: :::{md-tab-item} lint/format diff --git a/docs/source/index.md b/docs/source/index.md index 126d8ee..b217ae5 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -66,7 +66,7 @@ examples.md ``` ```{toctree} -:caption: Reference 📖 +:caption: Reference :hidden: :maxdepth: 1 diff --git a/src/cherab/imas/datasets/__init__.py b/src/cherab/imas/datasets/__init__.py index 1bf0faf..07de0ae 100644 --- a/src/cherab/imas/datasets/__init__.py +++ b/src/cherab/imas/datasets/__init__.py @@ -46,7 +46,7 @@ the internet connectivity. """ -from ._builtin import bolometer_moc +from ._builtin import bolometer_moc, iter_jintrac_radiation_values from ._fetchers import iter_jintrac, iter_jorek, iter_solps from ._utils import clear_cache @@ -54,6 +54,7 @@ "iter_jintrac", "iter_solps", "iter_jorek", + "iter_jintrac_radiation_values", "bolometer_moc", "clear_cache", ] diff --git a/src/cherab/imas/datasets/_builtin.py b/src/cherab/imas/datasets/_builtin.py index 2e6bb24..5f56daa 100644 --- a/src/cherab/imas/datasets/_builtin.py +++ b/src/cherab/imas/datasets/_builtin.py @@ -1,6 +1,7 @@ """Provide functionality to create builtin IMAS sample datasets.""" import datetime +from pathlib import Path import numpy as np import pooch @@ -8,6 +9,9 @@ from imas import DBEntry, IDSFactory from imas.ids_defs import IDS_TIME_MODE_HOMOGENEOUS +from imas.ids_toplevel import IDSToplevel + +from ..ids.common.ggd import load_grid N_CH = 5 # Number of channels per camera N_APERTURE = 3 # Number of apertures per channel (for collimator cameras) @@ -27,6 +31,79 @@ Y_AXIS = Vector3D(0, 1, 0) +def _iter_jintrac_radiation_values_data() -> tuple[IDSToplevel, IDSToplevel]: + """Create synthetic radiation/equilibrium IDS objects for values-based emitter tests. + + Returns + ------- + tuple[IDSToplevel, IDSToplevel] + A pair containing ``(equilibrium_ids, radiation_ids)``. + + Raises + ------ + RuntimeError + If the source JINTRAC dataset does not contain required equilibrium/core/edge data. + """ + from ._fetchers import iter_jintrac + + with DBEntry(iter_jintrac(), "r") as entry: + equilibrium = entry.get("equilibrium", autoconvert=False) + core_profiles = entry.get("core_profiles", autoconvert=False) + edge_profiles = entry.get("edge_profiles", autoconvert=False) + + if not len(equilibrium.time): + raise RuntimeError("The source equilibrium IDS has no time slices.") + if not len(core_profiles.profiles_1d): + raise RuntimeError("The source core_profiles IDS has no 1D profile data.") + if not len(edge_profiles.grid_ggd): + raise RuntimeError("The source edge_profiles IDS has no GGD grid.") + + src_grid = core_profiles.profiles_1d[0].grid + rho_tor_norm = np.asarray(src_grid.rho_tor_norm, dtype=np.float64) + psi = np.asarray(src_grid.psi, dtype=np.float64) + + if rho_tor_norm.size == 0 or psi.size == 0: + raise RuntimeError("The source core_profiles grid does not contain psi/rho_tor_norm data.") + + radiation = IDSFactory(equilibrium._version).new("radiation") + radiation.ids_properties.homogeneous_time = equilibrium.ids_properties.homogeneous_time + radiation.ids_properties.comment = "Synthetic radiation IDS for CHERAB-IMAS values tests" + radiation.ids_properties.creation_date = datetime.date.today().isoformat() + radiation.time = np.asarray(equilibrium.time, dtype=np.float64) + + radiation.grid_ggd.resize(1) + radiation.grid_ggd[0] = edge_profiles.grid_ggd[0] + + radiation.process.resize(1) + process = radiation.process[0] + process.identifier.index = 901 # custom emission, referenced: https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/radiation_identifier.html + process.identifier.name = "total" + process.profiles_1d.resize(1) + process.profiles_1d[0].grid.rho_tor_norm = rho_tor_norm + process.profiles_1d[0].grid.psi = psi + + # Smooth, strictly positive core emissivity profile. + core_values = 2.0e5 * (1.0 - 0.7 * np.clip(rho_tor_norm, 0.0, 1.0) ** 1.5) + 2.0e4 + process.profiles_1d[0].electrons.emissivity = core_values + + _, subsets, subset_id = load_grid(radiation.grid_ggd[0], with_subsets=True) + subset_name = next((name for name, index in subset_id.items() if index == 5), None) + if subset_name is None: + raise RuntimeError("Unable to find GGD subset id=5 (cells) in source grid_ggd.") + + num_cells = len(subsets[subset_name]) + if num_cells == 0: + raise RuntimeError("The selected GGD subset (id=5) contains no cells.") + + edge_values = np.linspace(4.0e4, 1.0e5, num_cells, dtype=np.float64) + process.ggd.resize(1) + process.ggd[0].electrons.emissivity.resize(1) + process.ggd[0].electrons.emissivity[0].grid_subset_index = 5 + process.ggd[0].electrons.emissivity[0].values = edge_values + + return equilibrium, radiation + + def _bolo_data(): """ Create a mock bolometer IDS dataset. @@ -283,3 +360,40 @@ def bolometer_moc() -> str: entry.put(ids) return str(path) + + +def iter_jintrac_radiation_values() -> str: + """Return a synthetic radiation dataset for values-based emitter examples/tests. + + This builtin dataset is generated from the fetched ``iter_jintrac`` sample and contains: + - one ``equilibrium`` IDS (copied from ``iter_jintrac``), and + - one ``radiation`` IDS with both core-profile emissivity and GGD emissivity values. + + .. note:: + + This dataset is intended for testing and demonstration purposes only, and does not + represent any real physical scenario. + + Returns + ------- + str + Path to the synthetic dataset file. + + Examples + -------- + >>> from cherab.imas import datasets + >>> data_path = datasets.iter_jintrac_radiation_values() + >>> data_path + '.../cherab/imas/iter_jintrac_radiation_values.nc' + """ + path = Path(pooch.os_cache("cherab/imas")) / "iter_jintrac_radiation_values.nc" + + path.parent.mkdir(parents=True, exist_ok=True) + + if not path.exists(): + equilibrium, radiation = _iter_jintrac_radiation_values_data() + with DBEntry(str(path), "w", dd_version=equilibrium._version) as entry: + entry.put(equilibrium) + entry.put(radiation) + + return str(path) diff --git a/src/cherab/imas/datasets/_registry.py b/src/cherab/imas/datasets/_registry.py index a8f99d1..fcb0dff 100644 --- a/src/cherab/imas/datasets/_registry.py +++ b/src/cherab/imas/datasets/_registry.py @@ -12,5 +12,6 @@ "iter_jintrac": ["iter_scenario_53298_seq1_DD4.nc", "iter_scenario_53298_seq1_DD4_mod.nc"], "iter_solps": ["iter_scenario_123364_1.nc"], "iter_jorek": ["iter_disruption_113112_1.nc"], + "iter_jintrac_radiation_values": ["iter_jintrac_radiation_values.nc"], "bolometer_moc": ["bolometer_moc.nc"], } diff --git a/src/cherab/imas/emitter/radiation.py b/src/cherab/imas/emitter/radiation.py index e3369b3..e3791ea 100644 --- a/src/cherab/imas/emitter/radiation.py +++ b/src/cherab/imas/emitter/radiation.py @@ -17,81 +17,233 @@ # under the Licence. """Module for loading radiation emissivity from IMAS-like IDS objects and creating emitter objects.""" +from collections.abc import Callable, Collection from pathlib import Path -from typing import Literal +from typing import Any, Literal import numpy as np +from numpy.typing import NDArray from raysect.core.math import translate -from raysect.core.math.function.float import Function2D +from raysect.core.math.function.float import ( + Function2D, + Function3D, + Interpolator1DArray, +) from raysect.core.scenegraph._nodebase import _NodeBase from raysect.primitive import Cylinder, Subtract from cherab.core.math import AxisymmetricMapper from cherab.tools.emitters import RadiationFunction +from cherab.tools.equilibrium import EFITEquilibrium from imas import DBEntry +from imas.ids_struct_array import IDSStructArray from imas.ids_structure import IDSStructure +from ..ggd import UnstructGrid2DExtended from ..ggd.base_mesh import InterpolatorCacheMode from ..ids.common import get_ids_time_slice from ..ids.common.ggd import load_grid -from ..ids.radiation import load_radiation_coefficients, load_radiation_emissivity +from ..ids.common.grid_radial import GridData, get_psi_norm, load_core_grid +from ..ids.radiation import load_core_emissivity, load_ggd_emissivity from ..math import FourierBezierConstructor +from ..math.blend import blend_core_edge_functions +from ..plasma.equilibrium import load_equilibrium from ..plasma.utility import get_subset_name_index __all__ = ["load_radiation_emitter"] +def _load_emissivity_values( + processes: IDSStructArray, + process_indices: Collection[int] | None, + grid_subset_id: int, +) -> tuple[NDArray[np.float64] | None, NDArray[np.float64] | None]: + values_core = None + values_ggd = None + + for process in processes: + # Validate process + if process_indices is not None and process.identifier.index.value not in process_indices: + continue + + # Values (profiles_1d) + _values = load_core_emissivity(process).sum() + + if _values is not None: + if values_core is None: + values_core = _values + else: + values_core += _values + + # Values (GGD) + _values = load_ggd_emissivity( + process, + grid_subset_index=grid_subset_id, + field="values", + ).sum() + + if _values is not None: + if values_ggd is None: + values_ggd = _values + else: + values_ggd += _values + + return values_core, values_ggd + + +def _create_rad_func_core( + grid: GridData, + data: NDArray[np.float64], + equilibrium: EFITEquilibrium | None, + psi_interpolator: Callable[[float], float] | None, + db_args: tuple | None, + db_kwargs: dict[str, Any] | None, + time: float, + occurrence: int, +) -> tuple[AxisymmetricMapper, EFITEquilibrium]: + if equilibrium is None: + equilibrium, psi_interp = load_equilibrium( + *(db_args or ()), + time=time, + occurrence=occurrence, + with_psi_interpolator=True, + **(db_kwargs or {}), + ) + psi_interpolator = psi_interpolator or psi_interp + else: + if not isinstance(equilibrium, EFITEquilibrium): + raise ValueError("Argument equilibrium must be a EFITEquilibrium instance.") + + # Create core grid + psi_norm = get_psi_norm( + grid.psi, + equilibrium.psi_axis, + equilibrium.psi_lcfs, + grid.rho_tor_norm, + psi_interpolator, + ) + psi_norm, index = np.unique(psi_norm, return_index=True) + extrapolation_range = max(0.0, psi_norm[0], 1.0 - psi_norm[-1]) + rad_func = equilibrium.map3d( + Interpolator1DArray(psi_norm, data[index], "cubic", "nearest", extrapolation_range) + ) + + return rad_func, equilibrium + + +def _create_rad_func_ggd( + grid_ggd: IDSStructure, + data: NDArray[np.float64], + grid_subset_id: int, + **interp_kwargs, +) -> tuple[AxisymmetricMapper, dict[str, float]]: + grid, subsets, subset_id = load_grid(grid_ggd, with_subsets=True) + grid_subset_name, grid_subset_index = get_subset_name_index(subset_id, grid_subset_id) + + if not np.array_equal(subsets[grid_subset_name], np.arange(grid.num_cell, dtype=int)): + grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) + + rad_func = AxisymmetricMapper(grid.interpolator(data, **interp_kwargs)) + + return rad_func, grid.mesh_extent + + def load_radiation_emitter( *args, time: float = 0, occurrence: int = 0, - process_index: int | None = None, - ion_index: int = 0, - emissivity_index: int = 0, + args2: tuple | None = None, + kwargs2: dict[str, Any] | None = None, + time2: float | None = None, + occurrence2: int = 0, + process_index: int | Collection[int] | None = None, grid_ggd: IDSStructure | None = None, - grid_subset_id: int | str = 5, - num_toroidal: int | None = None, - phis: np.ndarray | None = None, + grid_subset_id: int = 5, + equilibrium: EFITEquilibrium | None = None, + psi_interpolator: Callable[[float], float] | None = None, + mask: Function2D | Function3D | None = None, + num_toroidal: int = 64, + phis: NDArray[np.float64] | None = None, source: Literal["auto", "values", "coefficients"] = "auto", + time_threshold: float = np.inf, step: float = 0.01, parent: _NodeBase | None = None, - time_threshold: float = np.inf, interpolator_cache: InterpolatorCacheMode = "memory", interpolator_cache_dir: str | Path | None = None, **kwargs, ) -> Subtract | Cylinder: """Load radiation emissivity and create a single radiation emitter primitive. - The grid interpolator handles cache lookup and persistence internally. + There are two sources of emissivity data in the IMAS radiation IDS: + 1. core-region emissivity and/or edge-region emissivity (GGD-based) (``values``) + 2. emissivity coefficients (JOREK GGD-based) (``coefficients``) + + In the case of (1), one tries to load both core and edge (GGD-based) emissivity values from one + IMAS query. If both are available, they are blended using a mask function. + If the second IMAS query is provided, (``args2``, ``kwargs2``, etc.), it is used to load the + missing emissivity values if one of them is not available in the first query. + + In the case of (2), one tries to load emissivity coefficients from the GGD structure and + reconstructs the emissivity based on the Fourier-Bezier method, which is tied to the JOREK + specifications. + + If ``source="auto"``, the function first tries to load emissivity values (1), and if they are + not available, it falls back to emissivity coefficients (2). If neither is available, a + ``RuntimeError`` is raised. + + For GGD-based emissivity, the grid interpolator handles cache lookup and persistence internally. Parameters ---------- *args Positional arguments passed to `imas.DBEntry`. time - Time slice to load from the IDS, by default 0. + Time slice to load from the IDS, by default 0.0. occurrence - Occurrence of the radiation IDS to load, by default 0. + Occurrence of the radiation IDS, by default 0. + args2 + Arguments passed to `imas.DBEntry` for the second emissivity. If None, the second emissivity + is not loaded, by default None. + kwargs2 + Keyword arguments passed to `imas.DBEntry` for the second emissivity. If None, the second + emissivity is not loaded, by default None. + time2 + Time slice to load for the second emissivity. By default, uses the same time as the first + emissivity. + occurrence2 + Occurrence of the radiation IDS to load for the second emissivity, by default 0. process_index - Index of the radiation process to load, by default None (loads the first process). - ion_index - Index of the ion species to load, by default 0. - emissivity_index - Index of the emissivity data to load, by default 0. + Radiation process identifier index (or indices) to load. + By default, all available processes are summed together. + Reference: https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/radiation_identifier.html + .. note:: + The emissivity value array is assumed to follow the same x-axis as the grid subset. grid_ggd - Alternative grid GGD structure to use if the radiation IDS grid is empty, by default None. + Specific grid GGD structure alternative to the one in the IDS. grid_subset_id - ID or name of the grid subset to use, by default 5 (``"Cells"``). + ID of the grid subset to use, by default 5 (= "cells") subset. + Reference: https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/ggd_subset_identifier.html + equilibrium + Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` used to map core profiles. + By default None: the equilibrium is read from the same IMAS query as the core profiles. + Ignored if the core radiation is not available. + psi_interpolator + Alternative ``psi_norm(rho_tor_norm)`` interpolator. + Used only if ``psi`` is missing in the core grid, by default None. + Obtained from the ``equilibrium`` IDS in the same IMAS query as the core profiles. + mask + Mask function used for blending: ``(1 - mask) * f_gdd + mask * f_core``. + By default, uses `~cherab.tools.equilibrium.efit.EFITEquilibrium`'s `inside_lcfs`. num_toroidal - Number of toroidal subdivisions for 3D grid extension, by default None. + Number of toroidal subdivisions for 3D grid extension, by default 64. This is used only when the grid is loaded by `.load_unstruct_grid_2d_extended`. phis Array of toroidal angles in degrees for emissivity reconstruction, by default None. This is used only when the grid is loaded by `.load_unstruct_grid_2d_extended`. source Source for emissivity data: ``"auto"`` (tries values then coefficients), ``"values"`` - (emissivity values), or ``"coefficients"`` (reconstruct from Fourier-Bezier coefficients), - by default ``"auto"``. + (blended emissivity from core profiles + (edge) GGD values), or ``"coefficients"`` + (reconstruct from Fourier-Bezier coefficients), by default ``"auto"``. step Step size for the radiation function interpolator, by default 0.01 m. parent @@ -115,107 +267,258 @@ def load_radiation_emitter( Raises ------ + ValueError + If ``source`` is not one of ``"auto"``, ``"values"``, or ``"coefficients"``. RuntimeError - If the radiation IDS or its emissivity data cannot be loaded. + If no emissivity data is available in either core or GGD radiation data. """ - with DBEntry(*args, **kwargs) as entry: - radiation_ids = get_ids_time_slice( - entry, - "radiation", - time=time, - occurrence=occurrence, - time_threshold=time_threshold, + if source not in {"auto", "values", "coefficients"}: + raise ValueError( + f"Invalid source '{source}'. Expected one of: 'auto', 'values', 'coefficients'." ) - if not len(radiation_ids.grid_ggd) and grid_ggd is None: - raise RuntimeError( - "The 'grid_ggd' AOS of the radiation IDS is empty" - " and an alternative grid_ggd structure is not provided." - ) + if process_index is None: + process_indices = None + elif isinstance(process_index, int): + process_indices = {process_index} + else: + process_indices = set(process_index) - grid_ggd_struct = grid_ggd or radiation_ids.grid_ggd[0] + # Common variables + ids = None + ids2 = None + uri: str | None = None + uri2: str | None = None + emitter = None + grid = None + primitive_name: str | None = None + radius_outer: float = 0.0 + radius_inner: float = 0.0 + height: float = 0.0 + zmin: float = 0.0 try: - grid, subsets, subset_id = load_grid( - grid_ggd_struct, - with_subsets=True, - num_toroidal=num_toroidal, - ) + with DBEntry(*args, **kwargs) as entry: + ids = get_ids_time_slice( + entry, + "radiation", + time=time, + occurrence=occurrence, + time_threshold=time_threshold, + ) + uri: str | None = entry.uri + except RuntimeError as err: + raise RuntimeError("Unable to load radiation IDS.") from err + + if args2 is not None and source != "coefficients": try: - grid_subset_name, grid_subset_index = get_subset_name_index(subset_id, grid_subset_id) - if not np.array_equal(subsets[grid_subset_name], np.arange(grid.num_cell, dtype=int)): - grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) - subset_enabled = True - except ValueError: - subset_enabled = False - grid_subset_index = None - except NotImplementedError: - subset_enabled = False - grid = load_grid(grid_ggd_struct, with_subsets=False, num_toroidal=num_toroidal) - grid_subset_index = None - - emissivity = None - values_error: Exception | None = None + with DBEntry(*args2, **(kwargs2 or {})) as entry: + ids2 = get_ids_time_slice( + entry, + "radiation", + time=time2 or time, + occurrence=occurrence2, + time_threshold=time_threshold, + ) + uri2: str | None = entry.uri + except RuntimeError as err: + raise RuntimeError("Unable to load second radiation IDS.") from err + + # ------------------------------ + # === Load emissivity values === + # ------------------------------ + # temporary variables + eq_args = args + eq_kwargs = kwargs + eq_time = time + eq_occurrence = occurrence + rad_func = None + rad_func_core = None + rad_func_ggd = None if source in {"auto", "values"}: - try: - if grid_subset_index is None and subset_enabled: - raise RuntimeError("Unable to determine grid subset index for emissivity.values.") - emissivity = load_radiation_emissivity( - radiation_ids, - process_index=process_index, - grid_subset_index=5 if grid_subset_index is None else grid_subset_index, + # ------------------------------ + # === Load emissivity values === + # ------------------------------ + values_core, values_ggd = _load_emissivity_values( + ids.process, process_indices, grid_subset_id + ) + + # Load emissivity values from the second IDS if available and needed + if values_core is None and values_ggd is None: + pass + + elif values_ggd is None or values_core is None: + if ids2 is not None: + values_core2, values_ggd2 = _load_emissivity_values( + ids2.process, process_indices, grid_subset_id + ) + + if values_core is not None and values_core2 is not None: + raise RuntimeError( + "Duplicate core emissivity values are available in both radiation IDSs." + ) + if values_ggd is not None and values_ggd2 is not None: + raise RuntimeError( + "Duplicate GGD emissivity values are available in both radiation IDSs." + ) + + if values_core is None: + values_core = values_core2 + eq_args = args2 + eq_kwargs = kwargs2 + eq_time = time2 or time + eq_occurrence = occurrence2 + if values_ggd is None: + values_ggd = values_ggd2 + uri = f"{uri} + {uri2}" + + if values_core is None and values_ggd is None and source == "values": + raise RuntimeError( + "No emissivity values are available in either core or GGD radiation data." ) - except Exception as err: - values_error = err - if source == "values": - raise - - if emissivity is None and source in {"auto", "coefficients"}: - coeff = load_radiation_coefficients( - radiation_ids, - process_index=0 if process_index is None else process_index, - ion_index=ion_index, - emissivity_index=emissivity_index, - grid_subset_index=grid_subset_index, + + # ---------------------------------- + # === Create radiation functions === + # ---------------------------------- + if values_core is not None: + # TODO: Should load grid data at the same time as emissivity? + if len(ids.process) and len(ids.process[0].profiles_1d): + grid_struct = ids.process[0].profiles_1d[0].grid + elif ids2 is not None and len(ids2.process) and len(ids2.process[0].profiles_1d): + grid_struct = ids2.process[0].profiles_1d[0].grid + else: + raise RuntimeError("No core grid is available in either radiation IDS.") + + grid_data = load_core_grid(grid_struct) + rad_func_core, equilibrium = _create_rad_func_core( + grid_data, + values_core, + equilibrium, + psi_interpolator, + eq_args, + eq_kwargs, + eq_time, + eq_occurrence, + ) + mask = mask or equilibrium.inside_lcfs + radius_inner, radius_outer = equilibrium.r_range + zmin, zmax = equilibrium.z_range + height = zmax - zmin + + if values_ggd is not None: + grid_ggd = grid_ggd or ids.grid_ggd[0] + rad_func_ggd, extent = _create_rad_func_ggd( + grid_ggd, + values_ggd, + grid_subset_id, + interpolator_cache=interpolator_cache, + interpolator_cache_dir=interpolator_cache_dir, + ) + radius_outer = extent["rmax"] + radius_inner = extent["rmin"] + height = extent["zmax"] - extent["zmin"] + zmin = extent["zmin"] + + if ( + rad_func_core is not None + and rad_func_ggd is not None + and isinstance(mask, Function2D | Function3D) + ): + rad_func = blend_core_edge_functions( + rad_func_core, + rad_func_ggd, + mask, + return3d=True, + ) + elif rad_func_core is not None: + rad_func = rad_func_core + elif rad_func_ggd is not None: + rad_func = rad_func_ggd + else: + pass + + if isinstance(rad_func, Function3D): + emitter = RadiationFunction(rad_func, step=step) + primitive_name = f"RadiationEmitter_{ids.time[0]}s, uri {uri}" + + # ------------------------------------ + # === Load emissivity coefficients === + # ------------------------------------ + + if emitter is None and source in {"auto", "coefficients"}: + # Load GGD Grid + grid = load_grid( + ids.grid_ggd[0], + with_subsets=False, + num_toroidal=num_toroidal, ) - constructor = FourierBezierConstructor(grid_ggd_struct, coefficients=coeff) + if not isinstance(grid, UnstructGrid2DExtended): + raise RuntimeError( + "Coefficient-based emissivity reconstruction requires a 2D-extended grid." + ) + coeff = None + + # Load emissivity coefficients + for process in ids.process: + # Validate process + if ( + process_indices is not None + and process.identifier.index.value not in process_indices + ): + continue + + _coeff = load_ggd_emissivity( + process, + 1, # NOTE: JOREK-specific: emissivity coefficients are always associated with nodes + field="coefficients", + ).sum() + + if _coeff is not None: + if coeff is None: + coeff = _coeff + else: + coeff += _coeff + + if coeff is None: + if source == "coefficients": + raise RuntimeError("No emissivity coefficients are available in radiation data.") + raise RuntimeError( + "Unable to load emissivity from radiation IDS:" + " no values or coefficients are available." + ) + + constructor = FourierBezierConstructor(ids.grid_ggd[0], coefficients=coeff) if phis is None: - if not hasattr(grid, "num_toroidal"): - raise RuntimeError( - "Coefficient-based emissivity reconstruction requires a 2D-extended grid with a num_toroidal attribute." - ) d_phi = 360.0 / grid.num_toroidal phis_array = np.arange(d_phi * 0.5, 360.0, d_phi, dtype=np.float64) else: phis_array = np.asarray(phis, dtype=np.float64) emissivity = constructor.average_gaussian_faces_per_toroidal(phis_array).ravel() + primitive_name = f"RadiationEmitter_{ids.time[0]}s, uri {uri}" - if emissivity is None: - if values_error is not None: - raise RuntimeError( - "Unable to load emissivity from radiation IDS using either values or coefficients." - ) from values_error - raise RuntimeError("Unable to load emissivity from radiation IDS.") - - rad_func = grid.interpolator( - emissivity, - interpolator_cache=interpolator_cache, - interpolator_cache_dir=interpolator_cache_dir, - ) - - if isinstance(rad_func, Function2D): - rad_func = AxisymmetricMapper(rad_func) + rad_func = grid.interpolator( + emissivity, + interpolator_cache=interpolator_cache, + interpolator_cache_dir=interpolator_cache_dir, + ) + # Create RadiationFunction material + emitter = RadiationFunction(rad_func, step=step) - emitter = RadiationFunction(rad_func, step=step) + # Determine primitive dimensions + radius_outer = grid.mesh_extent["rmax"] + radius_inner = grid.mesh_extent["rmin"] + height = grid.mesh_extent["zmax"] - grid.mesh_extent["zmin"] + zmin = grid.mesh_extent["zmin"] - radius_outer = grid.mesh_extent["rmax"] - radius_inner = grid.mesh_extent["rmin"] - height = grid.mesh_extent["zmax"] - grid.mesh_extent["zmin"] - zmin = grid.mesh_extent["zmin"] + # ------------------------------- + # === Create Primitive object === + # ------------------------------- + if emitter is None: + raise RuntimeError("Emitter material cannot be constructed.") if radius_inner > 0: primitive = Subtract( @@ -226,6 +529,6 @@ def load_radiation_emitter( primitive.transform = translate(0, 0, zmin) primitive.material = emitter - primitive.name = f"RadiationEmitter_{radiation_ids.time[0]}s, uri {entry.uri}" + primitive.name = primitive_name or "RadiationEmitter" return primitive diff --git a/src/cherab/imas/ids/common/__init__.py b/src/cherab/imas/ids/common/__init__.py index fd66e51..db058e1 100644 --- a/src/cherab/imas/ids/common/__init__.py +++ b/src/cherab/imas/ids/common/__init__.py @@ -17,8 +17,18 @@ # under the Licence. """Subpackage for common utilities for loading data from IMAS IDS structures.""" -from . import ggd, species +from __future__ import annotations + +from . import ggd, grid_radial, species +from ._ids_numeric import get_ids_numeric_field from ._model import solve_coronal_equilibrium from .slice import get_ids_time_slice -__all__ = ["species", "get_ids_time_slice", "ggd", "solve_coronal_equilibrium"] +__all__ = [ + "ggd", + "grid_radial", + "species", + "solve_coronal_equilibrium", + "get_ids_time_slice", + "get_ids_numeric_field", +] diff --git a/src/cherab/imas/ids/common/_ids_numeric.py b/src/cherab/imas/ids/common/_ids_numeric.py new file mode 100644 index 0000000..43f07eb --- /dev/null +++ b/src/cherab/imas/ids/common/_ids_numeric.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import numpy as np +from numpy.typing import NDArray + +from imas.ids_primitive import IDSNumericArray +from imas.ids_structure import IDSStructure + + +def get_ids_numeric_field(ids_struct: IDSStructure, name: str) -> NDArray[np.float64] | None: + """Return a numeric IDS field as a float64 NumPy array. + + Returns + ------- + `NDArray[numpy.float64]` or None + The field values, or ``None`` when the field is missing or empty. + """ + data = getattr(ids_struct, name, None) + if isinstance(data, IDSNumericArray) and len(data): + return np.asarray(data, dtype=np.float64) + return None diff --git a/src/cherab/imas/ids/common/ggd/__init__.py b/src/cherab/imas/ids/common/ggd/__init__.py index c0665ab..f6434ab 100644 --- a/src/cherab/imas/ids/common/ggd/__init__.py +++ b/src/cherab/imas/ids/common/ggd/__init__.py @@ -15,8 +15,9 @@ # # See the Licence for the specific language governing permissions and limitations # under the Licence. -"""Subpackage for loading GGD grids from IMAS IDS structures.""" +"""Subpackage for loading GGD-related data.""" +from .load_data import get_ggd_subset_data from .load_grid import load_grid from .load_unstruct_2d import load_unstruct_grid_2d from .load_unstruct_3d import load_unstruct_grid_2d_extended, load_unstruct_grid_3d @@ -26,4 +27,5 @@ "load_unstruct_grid_2d", "load_unstruct_grid_2d_extended", "load_unstruct_grid_3d", + "get_ggd_subset_data", ] diff --git a/src/cherab/imas/ids/common/ggd/load_data.py b/src/cherab/imas/ids/common/ggd/load_data.py new file mode 100644 index 0000000..865c1d5 --- /dev/null +++ b/src/cherab/imas/ids/common/ggd/load_data.py @@ -0,0 +1,73 @@ +"""Functions for loading GGD-related data.""" + +from __future__ import annotations + +from typing import Literal + +import numpy as np +from numpy.typing import NDArray + +from imas.ids_defs import EMPTY_INT +from imas.ids_primitive import IDSNumericArray +from imas.ids_structure import IDSStructArray, IDSStructure + +__all__ = ["get_ggd_subset_data"] + + +def get_ggd_subset_data( + ids_struct: IDSStructure, + name: str, + grid_subset_index: int, + field: Literal["values", "coefficients"] = "values", +) -> NDArray[np.float64] | None: + """Return ``ids_struct..`` on the requested GGD subset. + + The list of grid subset indices (GGD subset identifiers) can be seen at + https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/ggd_subset_identifier.html + + Parameters + ---------- + ids_struct + The IDS structure containing the GGD data. + Must contain a field named "grid_subset_index". + name + The name of the GGD field to retrieve (e.g., ``"electrons"``). + grid_subset_index + The index of the GGD subset to retrieve. + field + The field to retrieve from the GGD subset, by default "values". + + Returns + ------- + `NDArray[numpy.float64]` | None + If `field` is "values", returns 1D `numpy.ndarray`. + If `field` is "coefficients", returns 2D `numpy.ndarray`. + Returns None if the requested GGD subset or field is not present in the IDS structure. + + Raises + ------ + ValueError + If `field` is not one of "values" or "coefficients". + + Examples + -------- + >>> get_ggd_subset_data(ids.ggd[0].electrons, "density", 5) + + >>> get_ggd_subset_data(ids.ggd[0].ion[0], "temperature", 5, field="coefficients") + """ + if field not in {"values", "coefficients"}: + raise ValueError(f"Invalid field '{field}', must be 'values' or 'coefficients'") + + struct_arr = getattr(ids_struct, name, None) + if not isinstance(struct_arr, IDSStructArray) or not len(struct_arr): + return None + + for sub_struct in struct_arr: + index = getattr(sub_struct, "grid_subset_index", EMPTY_INT) + if index == grid_subset_index: + data = getattr(sub_struct, field, None) + if isinstance(data, IDSNumericArray) and len(data): + return np.asarray(data, dtype=np.float64) + return None + + return None diff --git a/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py b/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py index 43e0964..a2ff725 100644 --- a/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py +++ b/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py @@ -150,14 +150,13 @@ def load_unstruct_grid_2d( return grid # Reading grid subsets (2D only) - cell_subset_ids = (5, 22, 23, 24, 25, 38, 39, 40) + CELL_SUBSET_IDS = {5, 22, 23, 24, 25, 38, 39, 40} subsets = {} subset_id = {} for subset in grid_ggd.grid_subset: + subset_index = subset.identifier.index.value dimension_is_2d = subset.dimension == DIMENSION.FACE + 1 # C to Fortran indexing - known_subset_id = ( - subset.dimension != EMPTY_INT and subset.identifier.index in cell_subset_ids - ) + known_subset_id = subset.dimension != EMPTY_INT and subset_index in CELL_SUBSET_IDS if (dimension_is_2d or known_subset_id) and len(subset.element): name = str(subset.identifier.name) indices = np.empty(len(subset.element), dtype=np.int32) @@ -170,6 +169,6 @@ def load_unstruct_grid_2d( break indices[i] = element.object[0].index.value subsets[name] = indices - 1 # Fortran to C indexing - subset_id[name] = subset.identifier.index.value + subset_id[name] = subset_index return grid, subsets, subset_id diff --git a/src/cherab/imas/ids/common/grid_radial.py b/src/cherab/imas/ids/common/grid_radial.py new file mode 100644 index 0000000..63710f7 --- /dev/null +++ b/src/cherab/imas/ids/common/grid_radial.py @@ -0,0 +1,101 @@ +"""Module for loading core grid properties and calculating normalized poloidal flux.""" + +from collections.abc import Callable +from dataclasses import dataclass, fields + +import numpy as np +from numpy.typing import NDArray + +from imas.ids_structure import IDSStructure + +from ._ids_numeric import get_ids_numeric_field + +__all__ = ["GridData", "load_core_grid", "get_psi_norm"] + + +@dataclass +class GridData: + """Dataclass for storing grid properties of the core profiles.""" + + rho_tor_norm: NDArray[np.float64] | None = None + """Normalized toroidal flux coordinate.""" + psi: NDArray[np.float64] | None = None + """Toroidal flux [Wb].""" + volume: NDArray[np.float64] | None = None + """Volume enclosed by the flux surface [m^3].""" + area: NDArray[np.float64] | None = None + """Area of the flux surface [m^2].""" + surface: NDArray[np.float64] | None = None + """Surface-averaged value of the profile on the flux surface.""" + + +def load_core_grid(grid_struct: IDSStructure) -> GridData: + """Load grid properties of the core profiles. + + The returned dictionary values for missing data are None. + + Parameters + ---------- + grid_struct + The IDS structure containing the grid data for 1D profiles. + + Returns + ------- + `.GridData` + Instance of the `.GridData` dataclass containing the grid properties for the core profiles. + """ + grid = GridData() + for field in fields(grid): + setattr(grid, field.name, get_ids_numeric_field(grid_struct, field.name)) + + return grid + + +def get_psi_norm( + psi: NDArray[np.float64] | None, + psi_axis: float, + psi_lcfs: float, + rho_tor_norm: NDArray[np.float64] | None, + psi_interpolator: Callable[[float], float] | None, +) -> NDArray[np.float64]: + """Calculate normalized poloidal flux. + + Parameters + ---------- + psi + Poloidal flux values from the core grid. + psi_axis + Poloidal flux at the magnetic axis. + psi_lcfs + Poloidal flux at the last closed flux surface. + rho_tor_norm + Normalized toroidal flux values. + psi_interpolator + Interpolator function to map `rho_tor_norm` to `psi_norm`. + Used only if ``psi`` is None. + + Returns + ------- + `NDArray[np.float64]` + Normalized poloidal flux values. + + Raises + ------ + RuntimeError + If both ``psi`` and ``rho_tor_norm`` are None, or if ``psi_interpolator`` is None when + ``psi`` is None. + """ + if psi is None: + if psi_interpolator is None: + raise RuntimeError( + "Unable to map rho_tor_norm to psi_norm grid: psi_interpolator is not provided." + ) + + if rho_tor_norm is None: + raise RuntimeError( + "No rho_tor_norm values are available in the core grid: unable to interpolate to psi_norm." + ) + + return np.array([psi_interpolator(rho) for rho in rho_tor_norm]) + + return (-psi / (2 * np.pi) - psi_axis) / (psi_lcfs - psi_axis) diff --git a/src/cherab/imas/ids/common/species.py b/src/cherab/imas/ids/common/species.py index 3499c87..097c584 100644 --- a/src/cherab/imas/ids/common/species.py +++ b/src/cherab/imas/ids/common/species.py @@ -200,14 +200,25 @@ def get_ion_state( break else: z_average = [] - if len(z_average): # probably, a bundle + if len(z_average): + warning_msg = ( + f"Warning: z_min or z_max is EMPTY_FLOAT for state index {state_index}. " + f"Using z_average to determine z_min and z_max." + ) z_min = ( int(state.z_min) if state.z_min != EMPTY_FLOAT else int(np.floor(min(z_average))) ) z_max = int(state.z_max) if state.z_max != EMPTY_FLOAT else int(np.ceil(max(z_average))) - else: # probably, a single ion + else: + warning_msg = ( + f"Warning: z_min or z_max is EMPTY_FLOAT for state index {state_index}. " + f"z_average is also empty. Using state_index + 1 as z_min and z_max." + ) z_min = int(state.z_min) if state.z_min != EMPTY_FLOAT else state_index + 1 z_max = int(state.z_max) if state.z_max != EMPTY_FLOAT else z_min + + print(warning_msg) + else: z_min = int(state.z_min) z_max = int(state.z_max) @@ -367,14 +378,14 @@ def get_elements(elements_aos: IDSStructArray) -> tuple[Element | Isotope, ...]: mass_number = int(round(element.a)) zn = int(round(element.z_n)) isotope = lookup_isotope(zn, number=mass_number) - if int(round(isotope.element.atomic_weight)) == mass_number: + if round(isotope.element.atomic_weight) == mass_number: # Prefer element over isotope isotope = isotope.element if getattr(element, "atoms_n", EMPTY_INT) == EMPTY_INT: atoms_n = 1 else: - atoms_n = int(round(getattr(element, "atoms_n", EMPTY_INT))) + atoms_n = round(getattr(element, "atoms_n", EMPTY_INT), ndigits=None) for _ in range(atoms_n): elements.append(isotope) diff --git a/src/cherab/imas/ids/core_profiles/__init__.py b/src/cherab/imas/ids/core_profiles/__init__.py index 92e024a..aaa0371 100644 --- a/src/cherab/imas/ids/core_profiles/__init__.py +++ b/src/cherab/imas/ids/core_profiles/__init__.py @@ -17,6 +17,6 @@ # under the Licence. """Subpackage for loading core profiles from IMAS IDS structures.""" -from .load_profiles import GridData, load_core_grid, load_core_profiles, load_core_species +from .load_profiles import load_core_profiles, load_core_species -__all__ = ["GridData", "load_core_grid", "load_core_profiles", "load_core_species"] +__all__ = ["load_core_profiles", "load_core_species"] diff --git a/src/cherab/imas/ids/core_profiles/load_profiles.py b/src/cherab/imas/ids/core_profiles/load_profiles.py index d107700..aef2dff 100644 --- a/src/cherab/imas/ids/core_profiles/load_profiles.py +++ b/src/cherab/imas/ids/core_profiles/load_profiles.py @@ -17,10 +17,9 @@ # under the Licence. """Module for loading core-profile-related data from IMAS IDS structures.""" -from dataclasses import astuple, dataclass, fields +from dataclasses import astuple, fields import numpy as np -from numpy.typing import NDArray from cherab.core.atomic import AtomicData from imas.ids_primitive import IDSNumericArray @@ -41,29 +40,11 @@ ) __all__ = [ - "GridData", - "load_core_grid", "load_core_profiles", "load_core_species", ] -@dataclass -class GridData: - """Dataclass for storing grid properties of the core profiles.""" - - rho_tor_norm: NDArray[np.float64] | None = None - """Normalized toroidal flux coordinate.""" - psi: NDArray[np.float64] | None = None - """Toroidal flux [Wb].""" - volume: NDArray[np.float64] | None = None - """Volume enclosed by the flux surface [m^3].""" - area: NDArray[np.float64] | None = None - """Area of the flux surface [m^2].""" - surface: NDArray[np.float64] | None = None - """Surface-averaged value of the profile on the flux surface.""" - - def _get_profile(ids_struct: IDSStructure, name: str, name2: str | None = None): data = getattr(ids_struct, name, None) if isinstance(data, IDSNumericArray): @@ -78,28 +59,6 @@ def _get_profile(ids_struct: IDSStructure, name: str, name2: str | None = None): return None -def load_core_grid(grid_struct: IDSStructure) -> GridData: - """Load grid properties of the core profiles. - - The returned dictionary values for missing data are None. - - Parameters - ---------- - grid_struct - The IDS structure containing the grid data for 1D profiles. - - Returns - ------- - `.GridData` - Instance of the `.GridData` dataclass containing the grid properties for the core profiles. - """ - grid = GridData() - for field in fields(grid): - setattr(grid, field.name, _get_profile(grid_struct, field.name)) - - return grid - - def load_core_profiles( species_struct: IDSStructure, species: SpeciesData, diff --git a/src/cherab/imas/ids/radiation/__init__.py b/src/cherab/imas/ids/radiation/__init__.py index 96d8ffc..bb7cddf 100644 --- a/src/cherab/imas/ids/radiation/__init__.py +++ b/src/cherab/imas/ids/radiation/__init__.py @@ -17,6 +17,14 @@ # under the Licence. """Subpackage for loading radiation data from IMAS IDS structures.""" -from .load_radiation import load_radiation_coefficients, load_radiation_emissivity +from .load_radiation import ( + EmissivityData, + load_core_emissivity, + load_ggd_emissivity, +) -__all__ = ["load_radiation_emissivity", "load_radiation_coefficients"] +__all__ = [ + "EmissivityData", + "load_core_emissivity", + "load_ggd_emissivity", +] diff --git a/src/cherab/imas/ids/radiation/load_radiation.py b/src/cherab/imas/ids/radiation/load_radiation.py index 1489802..2bb17ee 100644 --- a/src/cherab/imas/ids/radiation/load_radiation.py +++ b/src/cherab/imas/ids/radiation/load_radiation.py @@ -17,196 +17,197 @@ # under the Licence. """Module for loading radiation emissivity from the radiation IDS.""" +from dataclasses import dataclass +from typing import Literal + import numpy as np from numpy.typing import NDArray -from imas.ids_defs import EMPTY_INT -from imas.ids_primitive import IDSNumericArray from imas.ids_structure import IDSStructArray, IDSStructure -__all__ = ["load_radiation_emissivity", "load_radiation_coefficients"] +from ..common import get_ids_numeric_field +from ..common.ggd.load_data import get_ggd_subset_data +__all__ = [ + "EmissivityData", + "load_core_emissivity", + "load_ggd_emissivity", +] -def _get_emissivity( - ggd_struct: IDSStructure, - grid_subset_index: int, -) -> NDArray[np.float64] | None: - """Extract emissivity values for a given grid subset index from a ggd time-slice structure. - Parameters - ---------- - ggd_struct - A single element of the ``ggd`` (or ``process[i].ggd``) array-of-structures. - grid_subset_index - The ``grid_subset_index`` to match against. +@dataclass +class EmissivityData: + """Emissivity data.""" + + electron: NDArray[np.float64] | None = None + ion: NDArray[np.float64] | None = None + neutral: NDArray[np.float64] | None = None + + def sum(self) -> NDArray[np.float64] | None: + """Return the sum of all available emissivity arrays. + + Returns + ------- + `NDArray[numpy.float64]` or None + The sum of all available emissivity arrays, or None if no arrays are available. + """ + arrays = [arr for arr in (self.electron, self.ion, self.neutral) if arr is not None] + if not arrays: + return None + return np.sum(arrays, axis=0) + + +def _sum_profile_species_emissivity( + profile_1d: IDSStructure, + species_name: str, +) -> NDArray[np.float64] | None: + """Sum ``(:)/emissivity(:)`` for 1D profiles. Returns ------- `NDArray[numpy.float64]` or None - Emissivity values [W/m³], or ``None`` if the requested subset is not found. + Summed emissivity across all species entries with available data. """ - emissivity_arr = getattr(ggd_struct, "emissivity", None) - if not isinstance(emissivity_arr, IDSStructArray) or not len(emissivity_arr): + species_arr = getattr(profile_1d, species_name, None) + if not isinstance(species_arr, IDSStructArray) or not len(species_arr): return None - for item in emissivity_arr: - idx = getattr(item, "grid_subset_index", EMPTY_INT) - if idx == grid_subset_index: - values = getattr(item, "values", None) - if isinstance(values, IDSNumericArray) and len(values): - return np.asarray(values, dtype=np.float64) - - return None + total = None + for species in species_arr: + values = get_ids_numeric_field(species, "emissivity") + if values is None: + continue + total = values.copy() if total is None else total + values + return total -def load_radiation_emissivity( - radiation_ids, - process_index: int | None = None, - grid_subset_index: int = 5, -) -> NDArray[np.float64]: - """Load emissivity values from a radiation IDS time slice. - Parameters - ---------- - radiation_ids - Radiation IDS object (top-level or time slice) obtained from `~imas.db_entry.DBEntry`. - process_index - Index of the radiation process whose emissivity to load. - If ``None`` (default), the total emissivity is read from the top-level ``ggd`` array of the - IDS. - grid_subset_index - ``grid_subset_index`` identifier of the grid subset to read, by default 5 (``"Cells"``). +def _sum_ggd_species_emissivity( + ggd_struct: IDSStructure, + species_name: str, + grid_subset_index: int, + field: Literal["values", "coefficients"] = "values", +) -> NDArray[np.float64] | None: + """Sum ``(:)/emissivity(:)/(:)`` for one GGD structure. Returns ------- - `NDArray[numpy.float64]` - Emissivity values [W/m³] for each cell of the requested grid subset. - - Raises - ------ - RuntimeError - If the required AOS is empty or the requested subset cannot be found. + `NDArray[numpy.float64]` or None + Summed emissivity across all species entries on the requested subset. """ - if process_index is None: - # Total emissivity from the top-level ggd AOS - ggd = getattr(radiation_ids, "ggd", None) - if not isinstance(ggd, IDSStructArray) or not len(ggd): - raise RuntimeError("The 'ggd' AOS of the radiation IDS is empty.") + species_arr = getattr(ggd_struct, species_name, None) + if not isinstance(species_arr, IDSStructArray) or not len(species_arr): + return None - values = _get_emissivity(ggd[0], grid_subset_index) + total = None + for species in species_arr: + values = get_ggd_subset_data(species, "emissivity", grid_subset_index, field=field) if values is None: - raise RuntimeError( - f"Emissivity with grid_subset_index={grid_subset_index} not found " - "in radiation.ggd[0]." - ) - else: - # Per-process emissivity - processes = getattr(radiation_ids, "process", None) - if not isinstance(processes, IDSStructArray) or not len(processes): - raise RuntimeError("The 'process' AOS of the radiation IDS is empty.") - if process_index < 0 or process_index >= len(processes): - raise RuntimeError( - f"process_index={process_index} is out of range [0, {len(processes) - 1}]." - ) - - process = processes[process_index] - ggd = getattr(process, "ggd", None) - if not isinstance(ggd, IDSStructArray) or not len(ggd): - raise RuntimeError(f"The 'ggd' AOS of radiation.process[{process_index}] is empty.") - - values = _get_emissivity(ggd[0], grid_subset_index) - if values is None: - raise RuntimeError( - f"Emissivity with grid_subset_index={grid_subset_index} not found " - f"in radiation.process[{process_index}].ggd[0]." - ) + continue + total = values.copy() if total is None else total + values - return values + return total -def load_radiation_coefficients( - radiation_ids, - process_index: int = 0, - ion_index: int = 0, - emissivity_index: int = 0, - grid_subset_index: int | None = None, -) -> NDArray[np.float64]: - """Load JOREK-style emissivity coefficients from a radiation IDS time slice. +def load_core_emissivity(process: IDSStructure) -> EmissivityData: + """Load emissivity arrays from ``process.profiles_1d[0]``. - This accessor targets the coefficient layout typically used by JOREK: - ``radiation.process[i].ggd[0].ion[j].emissivity[k].coefficients``. + All species emissivity arrays are summed to produce a single array for each species type + (electron, ion, neutral). Parameters ---------- - radiation_ids - Radiation IDS object (top-level or time slice). - process_index - Index of the process in the IDS ``process`` AOS. - ion_index - Index of the ion entry in ``process[...].ggd[0].ion``. - emissivity_index - Index of the emissivity entry in ``ion[...].emissivity``. - grid_subset_index - Optional subset index constraint. If provided, coefficient data is accepted - only when ``grid_subset_index`` matches the emissivity entry. + process + The IDS structure containing the radiation process data. Returns ------- - `NDArray[numpy.float64]` - Coefficient array, typically shaped ``(num_vertices * num_modes, 4)``. + `.EmissivityData` + Core-region emissivity arrays. + + Examples + -------- + >>> load_core_emissivity(ids.process[0]) + EmissivityData( + electron=array([...]), + ion=array([...]), + neutral=array([...]) + ) + """ + profiles = getattr(process, "profiles_1d", None) + if not isinstance(profiles, IDSStructArray) or not len(profiles): + return EmissivityData() + + profile_1d = profiles[0] + electrons = getattr(profile_1d, "electrons", None) + + return EmissivityData( + electron=( + get_ids_numeric_field(electrons, "emissivity") + if isinstance(electrons, IDSStructure) + else None + ), + ion=( + get_ids_numeric_field(profile_1d, "emissivity_ion_total") + or _sum_profile_species_emissivity(profile_1d, "ion") + ), + neutral=( + get_ids_numeric_field(profile_1d, "emissivity_neutral_total") + or _sum_profile_species_emissivity(profile_1d, "neutral") + ), + ) + + +def load_ggd_emissivity( + process: IDSStructure, + grid_subset_index: int, + field: Literal["values", "coefficients"] = "values", +) -> EmissivityData: + """Load emissivity data from ``process.ggd[0]``. + + Parameters + ---------- + process + The IDS structure containing the radiation process data. + grid_subset_index + The index of the GGD subset to retrieve. + field + The field to retrieve from the GGD subset, by default "values". - Raises - ------ - RuntimeError - If the requested process/ion/emissivity structure is missing or empty. + Returns + ------- + `.EmissivityData` + GGD-region emissivity arrays for the requested grid subset. + + Examples + -------- + >>> load_ggd_emissivity(ids.process[0], grid_subset_index=5) + EmissivityData( + electron=array([...]), + ion=array([...]), + neutral=array([...]) + ) + + >>> load_ggd_emissivity(ids.process[0], grid_subset_index=1, field="coefficients") + EmissivityData( + electron=array([[...], [...]]), + ion=array([[...], [...]]), + neutral=array([[...], [...]]) + ) """ - processes = getattr(radiation_ids, "process", None) - if not isinstance(processes, IDSStructArray) or not len(processes): - raise RuntimeError("The 'process' AOS of the radiation IDS is empty.") - if process_index < 0 or process_index >= len(processes): - raise RuntimeError( - f"process_index={process_index} is out of range [0, {len(processes) - 1}]." - ) - - process = processes[process_index] ggd_arr = getattr(process, "ggd", None) if not isinstance(ggd_arr, IDSStructArray) or not len(ggd_arr): - raise RuntimeError(f"The 'ggd' AOS of radiation.process[{process_index}] is empty.") + return EmissivityData() ggd = ggd_arr[0] - ions = getattr(ggd, "ion", None) - if not isinstance(ions, IDSStructArray) or not len(ions): - raise RuntimeError( - f"No ion emissivity data found in radiation.process[{process_index}].ggd[0]." - ) - if ion_index < 0 or ion_index >= len(ions): - raise RuntimeError(f"ion_index={ion_index} is out of range [0, {len(ions) - 1}].") - - emissivities = getattr(ions[ion_index], "emissivity", None) - if not isinstance(emissivities, IDSStructArray) or not len(emissivities): - raise RuntimeError( - "No emissivity coefficients found in " - f"radiation.process[{process_index}].ggd[0].ion[{ion_index}]." - ) - if emissivity_index < 0 or emissivity_index >= len(emissivities): - raise RuntimeError( - f"emissivity_index={emissivity_index} is out of range [0, {len(emissivities) - 1}]." - ) - - emissivity = emissivities[emissivity_index] - if grid_subset_index is not None: - idx = getattr(emissivity, "grid_subset_index", EMPTY_INT) - if idx != grid_subset_index: - raise RuntimeError( - "Requested emissivity coefficients do not match the requested " - f"grid_subset_index={grid_subset_index}." - ) - - coefficients = getattr(emissivity, "coefficients", None) - if not isinstance(coefficients, IDSNumericArray) or not len(coefficients): - raise RuntimeError( - "Emissivity coefficients are missing in " - f"radiation.process[{process_index}].ggd[0].ion[{ion_index}]." - ) - - return np.asarray(coefficients, dtype=np.float64) + electrons = getattr(ggd, "electrons", None) + + return EmissivityData( + electron=( + get_ggd_subset_data(electrons, "emissivity", grid_subset_index, field=field) + if isinstance(electrons, IDSStructure) + else None + ), + ion=_sum_ggd_species_emissivity(ggd, "ion", grid_subset_index, field=field), + neutral=_sum_ggd_species_emissivity(ggd, "neutral", grid_subset_index, field=field), + ) diff --git a/src/cherab/imas/math/__init__.py b/src/cherab/imas/math/__init__.py index 48c5756..f89b267 100644 --- a/src/cherab/imas/math/__init__.py +++ b/src/cherab/imas/math/__init__.py @@ -17,5 +17,6 @@ # under the Licence. """Subpackage for mathematical utilities.""" +from .blend import blend_core_edge_functions as blend_core_edge_functions from .functions import * # noqa: F403 from .interpolators import * # noqa: F403 diff --git a/src/cherab/imas/math/blend.py b/src/cherab/imas/math/blend.py new file mode 100644 index 0000000..cd5016f --- /dev/null +++ b/src/cherab/imas/math/blend.py @@ -0,0 +1,151 @@ +# Copyright 2023 Euratom +# Copyright 2023 United Kingdom Atomic Energy Authority +# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnologicas +# +# Licensed under the EUPL, Version 1.1 or - as soon they will be approved by the +# European Commission - subsequent versions of the EUPL (the "Licence"); +# You may not use this work except in compliance with the Licence. +# You may obtain a copy of the Licence at: +# +# https://joinup.ec.europa.eu/software/page/eupl5 +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. +# +# See the Licence for the specific language governing permissions and limitations +# under the Licence. +"""Utilities for blending core and edge profile functions.""" + +from __future__ import annotations + +from typing import Literal, overload + +from raysect.core.math.function.float import Blend2D, Blend3D, Function2D, Function3D +from raysect.core.math.function.vector3d import Blend2D as VectorBlend2D +from raysect.core.math.function.vector3d import Blend3D as VectorBlend3D +from raysect.core.math.function.vector3d import Function2D as VectorFunction2D +from raysect.core.math.function.vector3d import Function3D as VectorFunction3D + +from cherab.core.math import AxisymmetricMapper, VectorAxisymmetricMapper + +__all__ = ["blend_core_edge_functions"] + + +@overload +def blend_core_edge_functions( + core_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + edge_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + mask: Function2D | Function3D, + return3d: Literal[True] = True, +) -> Function3D | VectorFunction3D | None: ... + + +@overload +def blend_core_edge_functions( + core_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + edge_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + mask: Function2D | Function3D, + return3d: Literal[False], +) -> Function2D | VectorFunction2D | None: ... + + +def blend_core_edge_functions( + core_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + edge_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, + mask: Function2D | Function3D, + return3d: bool = True, +) -> Function2D | Function3D | VectorFunction2D | VectorFunction3D | None: + """Blend core and edge functions using ``(1 - mask) * edge + mask * core``. + + Parameters + ---------- + core_func + A 2D or 3D core scalar/vector function. + edge_func + A 2D or 3D edge scalar/vector function. + mask + A 2D or 3D scalar mask function. + return3d + If True, map 2D outputs to 3D assuming axisymmetry, by default True. + + Returns + ------- + Function2D | Function3D | VectorFunction2D | VectorFunction3D | None + The blended function, or None if both inputs are None. + + Raises + ------ + TypeError + If function/mask types are unsupported. + RuntimeError + If scalar and vector functions are mixed. + """ + if core_func is None and edge_func is None: + return None + + if core_func is not None and not isinstance( + core_func, Function2D | Function3D | VectorFunction2D | VectorFunction3D + ): + raise TypeError("The core_func must be a 2D or 3D function.") + + if edge_func is not None and not isinstance( + edge_func, Function2D | Function3D | VectorFunction2D | VectorFunction3D + ): + raise TypeError("The edge_func must be a 2D or 3D function.") + + if not isinstance(mask, Function2D | Function3D): + raise TypeError("The mask must be a 2D or 3D function.") + + if core_func is None: + if isinstance(edge_func, Function2D) and return3d: + return AxisymmetricMapper(edge_func) + if isinstance(edge_func, VectorFunction2D) and return3d: + return VectorAxisymmetricMapper(edge_func) + return edge_func + + if edge_func is None: + if isinstance(core_func, Function2D) and return3d: + return AxisymmetricMapper(core_func) + if isinstance(core_func, VectorFunction2D) and return3d: + return VectorAxisymmetricMapper(core_func) + return core_func + + if ( + isinstance(core_func, Function2D) + and isinstance(edge_func, Function2D) + and isinstance(mask, Function2D) + ): + blended_func = Blend2D(edge_func, core_func, mask) + return AxisymmetricMapper(blended_func) if return3d else blended_func + + if ( + isinstance(core_func, VectorFunction2D) + and isinstance(edge_func, VectorFunction2D) + and isinstance(mask, Function2D) + ): + blended_func = VectorBlend2D(edge_func, core_func, mask) + return VectorAxisymmetricMapper(blended_func) if return3d else blended_func + + if isinstance(core_func, Function2D): + core_func = AxisymmetricMapper(core_func) + + if isinstance(core_func, VectorFunction2D): + core_func = VectorAxisymmetricMapper(core_func) + + if isinstance(edge_func, Function2D): + edge_func = AxisymmetricMapper(edge_func) + + if isinstance(edge_func, VectorFunction2D): + edge_func = VectorAxisymmetricMapper(edge_func) + + if isinstance(mask, Function2D): + mask = AxisymmetricMapper(mask) + + if isinstance(core_func, Function3D) and isinstance(edge_func, Function3D): + return Blend3D(edge_func, core_func, mask) + + if isinstance(core_func, VectorFunction3D) and isinstance(edge_func, VectorFunction3D): + return VectorBlend3D(edge_func, core_func, mask) + + raise RuntimeError("Cannot blend scalar and vector functions.") diff --git a/src/cherab/imas/plasma/blend.py b/src/cherab/imas/plasma/blend.py index 730c05a..5f131d2 100644 --- a/src/cherab/imas/plasma/blend.py +++ b/src/cherab/imas/plasma/blend.py @@ -22,26 +22,25 @@ import numpy as np from raysect.core.math import translate -from raysect.core.math.function.float import Blend2D, Blend3D, Function2D, Function3D -from raysect.core.math.function.vector3d import Blend2D as VectorBlend2D -from raysect.core.math.function.vector3d import Blend3D as VectorBlend3D +from raysect.core.math.function.float import Function2D, Function3D from raysect.core.math.function.vector3d import Function2D as VectorFunction2D -from raysect.core.math.function.vector3d import Function3D as VectorFunction3D from raysect.core.scenegraph._nodebase import _NodeBase from raysect.primitive import Cylinder, Subtract from scipy.constants import atomic_mass, electron_mass from cherab.core import AtomicData, Maxwellian, Plasma, Species -from cherab.core.math import AxisymmetricMapper, VectorAxisymmetricMapper +from cherab.core.math import VectorAxisymmetricMapper from cherab.tools.equilibrium import EFITEquilibrium from imas import DBEntry from imas.ids_structure import IDSStructure from ..ids.common import get_ids_time_slice from ..ids.common.ggd import load_grid -from ..ids.core_profiles import load_core_grid, load_core_species +from ..ids.common.grid_radial import get_psi_norm, load_core_grid +from ..ids.core_profiles import load_core_species from ..ids.edge_profiles import load_edge_species -from .core import get_core_interpolators, get_psi_norm, load_core_plasma +from ..math.blend import blend_core_edge_functions +from .core import get_core_interpolators, load_core_plasma from .edge import get_edge_interpolators, load_edge_plasma from .equilibrium import load_equilibrium, load_magnetic_field from .utility import ( @@ -112,6 +111,8 @@ def load_plasma( Alternative ``grid_ggd`` structure describing the grid. By default None. grid_subset_id Identifier of the grid subset (index or name). By default 5 (``"Cells"``). + The list of grid subset identifiers can be seen at + https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/ggd_subset_identifier.html equilibrium Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` used to map core profiles. By default None: the equilibrium is read from the same IMAS query as the @@ -483,116 +484,7 @@ def blend_core_edge_interpolators( setattr( interpolators, field.name, - _blend_core_edge_functions(core_func, edge_func, mask, return3d), + blend_core_edge_functions(core_func, edge_func, mask, return3d), ) return interpolators - - -def _blend_core_edge_functions( - core_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, - edge_func: Function2D | Function3D | VectorFunction2D | VectorFunction3D | None, - mask: Function2D | Function3D, - return3d: bool, -) -> Function2D | Function3D | VectorFunction2D | VectorFunction3D | None: - """Blend together the core and edge interpolating functions using the modulating mask function. - - Parameters - ---------- - core_func - A 2D or 3D core interpolator. - edge_func - A 2D or 3D edge interpolator. - mask - The 2D or 3D mask function used for blending: ``(1 - mask) * f_edge + mask * f_core``. - return3d - If True, return the 3D functions for 2D interpolators assuming - rotational symmetry, by default False. - - Returns - ------- - Function2D | Function3D | VectorFunction2D | VectorFunction3D | None - The blended function, or None if both input functions are None. - - Raises - ------ - TypeError - If the input functions are not 2D or 3D scalar/vector functions. - RuntimeError - If the core and edge functions have incompatible dimensions or types for blending. - """ - if core_func is None and edge_func is None: - return None - - # === Validation === - if core_func is not None and not isinstance( - core_func, Function2D | Function3D | VectorFunction2D | VectorFunction3D - ): - raise TypeError("The core_func must be a 2D or 3D function.") - - if edge_func is not None and not isinstance( - edge_func, Function2D | Function3D | VectorFunction2D | VectorFunction3D - ): - raise TypeError("The edge_func must be a 2D or 3D function.") - - if not isinstance(mask, Function2D | Function3D): - raise TypeError("The mask must be a 2D or 3D function.") - - # === Only one of the two functions is available === - if core_func is None: - if isinstance(edge_func, Function2D) and return3d: - return AxisymmetricMapper(edge_func) - elif isinstance(edge_func, VectorFunction2D) and return3d: - return VectorAxisymmetricMapper(edge_func) - else: - return edge_func - - if edge_func is None: - if isinstance(core_func, Function2D) and return3d: - return AxisymmetricMapper(core_func) - elif isinstance(core_func, VectorFunction2D) and return3d: - return VectorAxisymmetricMapper(core_func) - else: - return core_func - - # === Both functions are available: blend them together === - if ( - isinstance(core_func, Function2D) - and isinstance(edge_func, Function2D) - and isinstance(mask, Function2D) - ): - blended_func = Blend2D(edge_func, core_func, mask) - return AxisymmetricMapper(blended_func) if return3d else blended_func - - if ( - isinstance(core_func, VectorFunction2D) - and isinstance(edge_func, VectorFunction2D) - and isinstance(mask, Function2D) - ): - blended_func = VectorBlend2D(edge_func, core_func, mask) - return VectorAxisymmetricMapper(blended_func) if return3d else blended_func - - # unable to return 2D, convert to 3D - - if isinstance(core_func, Function2D): - core_func = AxisymmetricMapper(core_func) - - if isinstance(core_func, VectorFunction2D): - core_func = VectorAxisymmetricMapper(core_func) - - if isinstance(edge_func, Function2D): - edge_func = AxisymmetricMapper(edge_func) - - if isinstance(edge_func, VectorFunction2D): - edge_func = VectorAxisymmetricMapper(edge_func) - - if isinstance(mask, Function2D): - mask = AxisymmetricMapper(mask) - - if isinstance(core_func, Function3D) and isinstance(edge_func, Function3D): - return Blend3D(edge_func, core_func, mask) - - if isinstance(core_func, VectorFunction3D) and isinstance(edge_func, VectorFunction3D): - return VectorBlend3D(edge_func, core_func, mask) - - raise RuntimeError("Cannot blend scalar and vector functions.") diff --git a/src/cherab/imas/plasma/core.py b/src/cherab/imas/plasma/core.py index ef9c627..00273a3 100644 --- a/src/cherab/imas/plasma/core.py +++ b/src/cherab/imas/plasma/core.py @@ -36,7 +36,8 @@ from imas import DBEntry from ..ids.common import get_ids_time_slice -from ..ids.core_profiles import load_core_grid, load_core_species +from ..ids.common.grid_radial import get_psi_norm, load_core_grid +from ..ids.core_profiles import load_core_species from .equilibrium import load_equilibrium, load_magnetic_field from .utility import ZERO_VELOCITY, ProfileInterpolator, warn_unsupported_species @@ -301,7 +302,7 @@ def get_core_interpolators( continue data_1d = getattr(profile, field.name, None) if isinstance(data_1d, np.ndarray) and data_1d.size > 0: - extrapolation_range = max(0, psi_norm[0], 1.0 - psi_norm[-1]) + extrapolation_range = max(0.0, psi_norm[0], 1.0 - psi_norm[-1]) func = Interpolator1DArray( psi_norm, data_1d[index], "cubic", "nearest", extrapolation_range ) @@ -315,53 +316,3 @@ def get_core_interpolators( pass # TODO: handle velocity profile return interpolators - - -def get_psi_norm( - psi: NDArray[np.float64] | None, - psi_axis: float, - psi_lcfs: float, - rho_tor_norm: NDArray[np.float64] | None, - psi_interpolator: Callable[[float], float] | None, -) -> NDArray[np.float64]: - """Calculate normalized poloidal flux. - - Parameters - ---------- - psi - Poloidal flux values from the core grid. - psi_axis - Poloidal flux at the magnetic axis. - psi_lcfs - Poloidal flux at the last closed flux surface. - rho_tor_norm - Normalized toroidal flux values. - psi_interpolator - Interpolator function to map `rho_tor_norm` to `psi_norm`. - Used only if ``psi`` is None. - - Returns - ------- - `NDArray[np.float64]` - Normalized poloidal flux values. - - Raises - ------ - RuntimeError - If both ``psi`` and ``rho_tor_norm`` are None, or if ``psi_interpolator`` is None when - ``psi`` is None. - """ - if psi is None: - if psi_interpolator is None: - raise RuntimeError( - "Unable to map rho_tor_norm to psi_norm grid: psi_interpolator is not provided." - ) - - if rho_tor_norm is None: - raise RuntimeError( - "No rho_tor_norm values are available in the core grid: unable to interpolate to psi_norm." - ) - - return np.array([psi_interpolator(rho) for rho in rho_tor_norm]) - - return (-psi / (2 * np.pi) - psi_axis) / (psi_lcfs - psi_axis) diff --git a/src/cherab/imas/plasma/equilibrium.py b/src/cherab/imas/plasma/equilibrium.py index b1ea1ff..36fe916 100644 --- a/src/cherab/imas/plasma/equilibrium.py +++ b/src/cherab/imas/plasma/equilibrium.py @@ -53,7 +53,7 @@ def load_equilibrium( time_threshold: float = np.inf, with_psi_interpolator: Literal[True], **kwargs, -) -> tuple[EFITEquilibrium, Interpolator1DArray | None]: ... +) -> tuple[EFITEquilibrium, Interpolator1DArray]: ... def load_equilibrium( @@ -63,7 +63,7 @@ def load_equilibrium( time_threshold: float = np.inf, with_psi_interpolator: bool = False, **kwargs, -) -> tuple[EFITEquilibrium, Interpolator1DArray | None] | EFITEquilibrium: +) -> tuple[EFITEquilibrium, Interpolator1DArray] | EFITEquilibrium: """Load plasma equilibrium from the equilibrium IDS and create an `EFITEquilibrium` object. Parameters @@ -87,11 +87,16 @@ def load_equilibrium( ------- equilibrium : `~cherab.tools.equilibrium.efit.EFITEquilibrium` The plasma equilibrium object. - psi_interpolator : `~raysect.core.math.function.float.function1d.interpolate.Interpolator1DArray` | None + psi_interpolator : `~raysect.core.math.function.float.function1d.interpolate.Interpolator1DArray` If ``with_psi_interpolator`` is True and ``rho_tor_norm`` is available, returns the ``psi_norm(rho_tor_norm)`` interpolator. - If rho_tor_norm is not available, returns None. - Otherwise, returns only the equilibrium object. + If rho_tor_norm is not available, raise a `RuntimeError`. + + Raises + ------ + RuntimeError + If the equilibrium IDS does not have a time slice or if ``rho_tor_norm`` is not available + when ``with_psi_interpolator`` is True. """ with DBEntry(*args, **kwargs) as entry: equilibrium_ids = get_ids_time_slice( @@ -127,19 +132,19 @@ def load_equilibrium( if not with_psi_interpolator: return equilibrium + else: + if eq_data.rho_tor_norm is None: + raise RuntimeError("rho_tor_norm is not available in the equilibrium IDS.") + + psi_interpolator = Interpolator1DArray( + eq_data.rho_tor_norm, + eq_data.psi_norm, + "cubic", + "none", + 0, + ) - if eq_data.rho_tor_norm is None: - return equilibrium, None - - psi_interpolator = Interpolator1DArray( - eq_data.rho_tor_norm, - eq_data.psi_norm, - "cubic", - "none", - 0, - ) - - return equilibrium, psi_interpolator + return equilibrium, psi_interpolator def load_magnetic_field( diff --git a/tests/conftest.py b/tests/conftest.py index 6f03df7..d531288 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,14 +1,76 @@ import shutil +from contextlib import suppress +from functools import lru_cache from pathlib import Path +from uuid import uuid4 import pytest +from imas import DBEntry +from imas.ids_defs import MEMORY_BACKEND from cherab.core.atomic.elements import neon -from cherab.imas.datasets import iter_jintrac, iter_jorek, iter_solps +from cherab.imas.datasets import ( + iter_jintrac, + iter_jintrac_radiation_values, + iter_jorek, + iter_solps, +) from cherab.openadas import OpenADAS from cherab.openadas.repository import populate +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--force-imas-memory-backend-tests", + action="store_true", + default=False, + help="Run tests marked 'requires_imas_memory_backend' even if the IMAS memory backend probe fails.", + ) + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "markers", + "requires_imas_memory_backend: test requires a working IMAS memory backend.", + ) + + +@lru_cache(maxsize=1) +def _probe_imas_memory_backend() -> tuple[bool, str]: + token = uuid4().int + entry = None + try: + entry = DBEntry( + MEMORY_BACKEND, + f"cherab_pytest_{token & 0xFFFF:04x}", + 1 + token % 2_000_000_000, + 1 + (token >> 31) % 2_000_000_000, + ) + entry.create() + except Exception as exc: + return False, f"IMAS memory backend is unavailable on this machine: {exc}" + finally: + if entry is not None: + with suppress(Exception): + entry.close() + + return True, "" + + +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + if config.getoption("--force-imas-memory-backend-tests"): + return + + available, reason = _probe_imas_memory_backend() + if available: + return + + skip_marker = pytest.mark.skip(reason=reason) + for item in items: + if item.get_closest_marker("requires_imas_memory_backend") is not None: + item.add_marker(skip_marker) + + @pytest.fixture(scope="session", autouse=True) def populate_openadas_repository(): """Fixture to populate the OpenADAS repository before running tests.""" @@ -50,3 +112,10 @@ def path_iter_jorek(tmp_path_factory: pytest.TempPathFactory) -> str: """Fixture to provide the path to a sample JOREK IMAS dataset.""" path = Path(iter_jorek()) return _copy_dataset_to_tmp(path, tmp_path_factory) + + +@pytest.fixture(scope="session") +def path_iter_jintrac_radiation_values(tmp_path_factory: pytest.TempPathFactory) -> str: + """Fixture to provide the path to a synthetic JINTRAC radiation values dataset.""" + path = Path(iter_jintrac_radiation_values()) + return _copy_dataset_to_tmp(path, tmp_path_factory) diff --git a/tests/emitter/test_radiation.py b/tests/emitter/test_radiation.py index d6831e2..29da939 100644 --- a/tests/emitter/test_radiation.py +++ b/tests/emitter/test_radiation.py @@ -1,13 +1,17 @@ import os from pathlib import Path from typing import Literal, TypedDict +from uuid import uuid4 import numpy as np import pytest +from imas import DBEntry, IDSFactory +from imas.ids_defs import MEMORY_BACKEND from raysect.primitive import Cylinder, Subtract import cherab.imas.emitter.radiation as radiation_module from cherab.imas.emitter import load_radiation_emitter +from cherab.imas.plasma.equilibrium import load_equilibrium class _EmitterCacheKwargs(TypedDict, total=False): @@ -44,6 +48,84 @@ def _cache_kwargs( return kwargs +class _OpenEntryContext: + def __init__(self, entry: DBEntry): + self.entry = entry + + def __enter__(self) -> DBEntry: + return self.entry + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + +def _patch_dbentry_for_open_entries( + monkeypatch: pytest.MonkeyPatch, + entries: dict[tuple, DBEntry], +) -> None: + original_dbentry = radiation_module.DBEntry + + def _dbentry_router(*args, **kwargs): + if not kwargs and args in entries: + return _OpenEntryContext(entries[args]) + return original_dbentry(*args, **kwargs) + + monkeypatch.setattr(radiation_module, "DBEntry", _dbentry_router) + + +def _write_split_radiation_to_memory( + path_values_dataset: str, + *, + include_core: bool, + include_ggd: bool, + name: str, + pulse: int, + run: int, +) -> tuple[tuple, DBEntry]: + with DBEntry(path_values_dataset, "r") as entry: + equilibrium = entry.get("equilibrium", autoconvert=False) + radiation = entry.get("radiation", autoconvert=False) + + split_radiation = IDSFactory(equilibrium._version).new("radiation") + split_radiation.ids_properties.homogeneous_time = equilibrium.ids_properties.homogeneous_time + split_radiation.ids_properties.comment = radiation.ids_properties.comment + split_radiation.ids_properties.creation_date = radiation.ids_properties.creation_date + split_radiation.time = np.asarray(radiation.time, dtype=np.float64) + + split_radiation.grid_ggd.resize(1) + split_radiation.grid_ggd[0] = radiation.grid_ggd[0] + + split_radiation.process.resize(1) + proc_src = radiation.process[0] + proc_dst = split_radiation.process[0] + proc_dst.identifier.index = int(np.asarray(proc_src.identifier.index).item()) + proc_dst.identifier.name = proc_src.identifier.name + + if include_core: + proc_dst.profiles_1d.resize(1) + src = proc_src.profiles_1d[0] + dst = proc_dst.profiles_1d[0] + dst.grid.rho_tor_norm = np.asarray(src.grid.rho_tor_norm, dtype=np.float64) + dst.grid.psi = np.asarray(src.grid.psi, dtype=np.float64) + dst.electrons.emissivity = np.asarray(src.electrons.emissivity, dtype=np.float64) + + if include_ggd: + proc_dst.ggd.resize(1) + src = proc_src.ggd[0].electrons.emissivity[0] + dst = proc_dst.ggd[0].electrons.emissivity + dst.resize(1) + dst[0].grid_subset_index = int(np.asarray(src.grid_subset_index).item()) + dst[0].values = np.asarray(src.values, dtype=np.float64) + + args = (MEMORY_BACKEND, name, pulse, run) + entry = DBEntry(*args, data_version=equilibrium._version) + entry.create() + entry.put(equilibrium) + entry.put(split_radiation) + + return args, entry + + def test_load_radiation_emitter_coefficients_uses_default_phis( path_iter_jorek: str, monkeypatch: pytest.MonkeyPatch, @@ -132,7 +214,10 @@ def test_load_radiation_emitter_values_raises_for_jorek( path_iter_jorek: str, radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], ): - with pytest.raises(RuntimeError, match="The 'ggd' AOS of the radiation IDS is empty"): + with pytest.raises( + RuntimeError, + match="No emissivity values are available in either core or GGD radiation data.", + ): load_radiation_emitter( path_iter_jorek, "r", @@ -141,14 +226,218 @@ def test_load_radiation_emitter_values_raises_for_jorek( ) -def test_load_radiation_emitter_invalid_source_raises( - path_iter_jorek: str, +def test_load_radiation_emitter_values_with_core_and_ggd( + path_iter_jintrac_radiation_values: str, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + primitive = load_radiation_emitter( + path_iter_jintrac_radiation_values, + "r", + source="values", + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + + +@pytest.mark.requires_imas_memory_backend +def test_load_radiation_emitter_values_completes_from_second_ids_core_then_ggd( + path_iter_jintrac_radiation_values: str, + monkeypatch: pytest.MonkeyPatch, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + seed = uuid4().int + name = f"cherab_radiation_values_{seed & 0xFFFF:04x}" + pulse = 500000 + seed % 100000 + args1, entry1 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=True, + include_ggd=False, + name=name, + pulse=pulse, + run=1, + ) + args2, entry2 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=False, + include_ggd=True, + name=name, + pulse=pulse, + run=2, + ) + + try: + _patch_dbentry_for_open_entries(monkeypatch, {args1: entry1, args2: entry2}) + equilibrium = load_equilibrium(path_iter_jintrac_radiation_values, "r") + + primitive = load_radiation_emitter( + *args1, + args2=args2, + source="values", + equilibrium=equilibrium, + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + finally: + entry1.close() + entry2.close() + + +@pytest.mark.requires_imas_memory_backend +def test_load_radiation_emitter_values_completes_from_second_ids_ggd_then_core( + path_iter_jintrac_radiation_values: str, + monkeypatch: pytest.MonkeyPatch, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + seed = uuid4().int + name = f"cherab_radiation_values_{seed & 0xFFFF:04x}" + pulse = 600000 + seed % 100000 + args1, entry1 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=False, + include_ggd=True, + name=name, + pulse=pulse, + run=3, + ) + args2, entry2 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=True, + include_ggd=False, + name=name, + pulse=pulse, + run=4, + ) + + try: + _patch_dbentry_for_open_entries(monkeypatch, {args1: entry1, args2: entry2}) + equilibrium = load_equilibrium(path_iter_jintrac_radiation_values, "r") + + primitive = load_radiation_emitter( + *args1, + args2=args2, + source="values", + equilibrium=equilibrium, + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + finally: + entry1.close() + entry2.close() + + +@pytest.mark.requires_imas_memory_backend +def test_load_radiation_emitter_values_raises_for_duplicate_core_in_args2( + path_iter_jintrac_radiation_values: str, + monkeypatch: pytest.MonkeyPatch, radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], ): - with pytest.raises(RuntimeError, match="Unable to load emissivity"): + seed = uuid4().int + name = f"cherab_radiation_values_{seed & 0xFFFF:04x}" + pulse = 700000 + seed % 100000 + args1, entry1 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=True, + include_ggd=False, + name=name, + pulse=pulse, + run=5, + ) + args2, entry2 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=True, + include_ggd=False, + name=name, + pulse=pulse, + run=6, + ) + + try: + _patch_dbentry_for_open_entries(monkeypatch, {args1: entry1, args2: entry2}) + equilibrium = load_equilibrium(path_iter_jintrac_radiation_values, "r") + + with pytest.raises( + RuntimeError, + match="Duplicate core emissivity values are available in both radiation IDSs.", + ): + load_radiation_emitter( + *args1, + args2=args2, + source="values", + equilibrium=equilibrium, + **_cache_kwargs(radiation_interpolator_cache), + ) + finally: + entry1.close() + entry2.close() + + +@pytest.mark.requires_imas_memory_backend +def test_load_radiation_emitter_values_raises_for_duplicate_ggd_in_args2( + path_iter_jintrac_radiation_values: str, + monkeypatch: pytest.MonkeyPatch, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + seed = uuid4().int + name = f"cherab_radiation_values_{seed & 0xFFFF:04x}" + pulse = 800000 + seed % 100000 + args1, entry1 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=False, + include_ggd=True, + name=name, + pulse=pulse, + run=7, + ) + args2, entry2 = _write_split_radiation_to_memory( + path_iter_jintrac_radiation_values, + include_core=False, + include_ggd=True, + name=name, + pulse=pulse, + run=8, + ) + + try: + _patch_dbentry_for_open_entries(monkeypatch, {args1: entry1, args2: entry2}) + equilibrium = load_equilibrium(path_iter_jintrac_radiation_values, "r") + + with pytest.raises( + RuntimeError, + match="Duplicate GGD emissivity values are available in both radiation IDSs.", + ): + load_radiation_emitter( + *args1, + args2=args2, + source="values", + equilibrium=equilibrium, + **_cache_kwargs(radiation_interpolator_cache), + ) + finally: + entry1.close() + entry2.close() + + +def test_load_radiation_emitter_invalid_source_raises_early( + monkeypatch: pytest.MonkeyPatch, +): + class _UnexpectedDBEntry: + def __init__(self, *args, **kwargs): + raise AssertionError("DBEntry must not be instantiated for invalid source.") + + monkeypatch.setattr(radiation_module, "DBEntry", _UnexpectedDBEntry) + + with pytest.raises( + ValueError, + match="Invalid source 'unsupported'. Expected one of: 'auto', 'values', 'coefficients'.", + ): load_radiation_emitter( - path_iter_jorek, + "dummy_path", "r", source="unsupported", # type: ignore[arg-type] - **_cache_kwargs(radiation_interpolator_cache), )