diff --git a/.binder/environment.yml b/.binder/environment.yml index 34b11e8..1819bc5 100644 --- a/.binder/environment.yml +++ b/.binder/environment.yml @@ -4,3 +4,4 @@ channels: dependencies: - cherab-imas - ultraplot + - plotly diff --git a/CHANGELOG.md b/CHANGELOG.md index f3b6a35..b996afc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ 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.0] - Unreleased +## [0.5.0] - 2026-06-24 ### Added @@ -15,6 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add notebook.link environment configuration for interactive notebook demos for future use in documentation. - Add Fourier-Bezier reconstruction module and integrate it into math utilities - Add JINTRAC/JOREK dataset fixtures and interpolator cache tests +- Add radiation data loading modules and emitter initialization +- Add JOREK-based 3D radiation emitter visualization notebook +- Add 3D grid visualization notebook for ITER JOREK data +- Add unit tests for `get_ids_time_slice` function with fallback handling ### Changed @@ -25,11 +29,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Enhance interpolator caching behavior in grid classes and add subset validation in `load_unstruct_grid_2d_extended` - Update pixi channels and `doc-serve` task to support normal conda-forge settings and dynamic port arguments - Refactor dataset handling to make `pooch` a required dependency and simplify cache directory handling +- Implement re-slicing of IDS via IMAS memory backend with fallback handling +- Update `environment.yml` to include additional dependencies for binder +- Update docstrings for interpolator cache parameters in grid classes ### Fixed - Ensure bolometer dataset cache directories are created before use - Fix total power calculation and exception handling in the emission notebook workflow +- Fix charge type conversion in `load_core_plasma` function +- Fix grid subset handling in plasma loading functions +- Ensure coefficients array is C-contiguous in `FourierBezierConstructor` ### Removed diff --git a/docs/notebooks/radiation/3d_radiation.ipynb b/docs/notebooks/radiation/3d_radiation.ipynb new file mode 100644 index 0000000..9b42387 --- /dev/null +++ b/docs/notebooks/radiation/3d_radiation.ipynb @@ -0,0 +1,356 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 3D Radiation Emitter\n", + "\n", + "This notebook demonstrates how to create a 3D radiation emitter from IMAS radiation IDS.\n", + "\n", + "The example test data was calculated by JOREK for an ITER disruption scenario." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "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", + "# Set dark background for plots\n", + "uplt.rc.style = \"dark_background\"\n", + "\n", + "TOGW = 1e-9 # [W] -> [GW]" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Retrieve ITER JOREK sample data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "path = iter_jorek()" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "## Create 3D radiation emitter from IMAS IDS" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "world = World()\n", + "emitter = load_radiation_emitter(\n", + " path,\n", + " \"r\",\n", + " time=0.0042853,\n", + " parent=world,\n", + " interpolator_cache=\"disk\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Visualize the emitter in 2D\n", + "\n", + "Sample 2D visualization of the radiation function at all toroidal angles." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "R_MIN, R_MAX = 4.0, 8.5\n", + "Z_MIN, Z_MAX = -4.5, 4.4\n", + "RES = 0.01 # resolution of grid in [m]\n", + "n_r = round((R_MAX - R_MIN) / RES) + 1\n", + "n_z = round((Z_MAX - Z_MIN) / RES) + 1\n", + "dr, dz = (R_MAX - R_MIN) / (n_r - 1), (Z_MAX - Z_MIN) / (n_z - 1)\n", + "\n", + "# (r, z) coordinates at phi=0\n", + "r_pts = np.linspace(R_MIN, R_MAX, n_r, endpoint=True)\n", + "z_pts = np.linspace(Z_MIN, Z_MAX, n_z, endpoint=True)\n", + "\n", + "# Toroidal angles to sample the radiation function at\n", + "n_phi = 64\n", + "d_phi = 360 / n_phi\n", + "eps = 1e-6 # small offset to avoid sampling at exactly 0 and 360 degrees\n", + "phis = np.linspace(0 + d_phi + eps, 360 + d_phi - eps, n_phi + 1, endpoint=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "n_phi = len(phis)\n", + "rad = np.zeros((n_r, n_z, n_phi), dtype=float)\n", + "for i, j, k in np.ndindex(n_r, n_z, n_phi):\n", + " rad[i, j, k] = emitter.material.radiation_function(\n", + " r_pts[i] * np.cos(np.deg2rad(phis[k])),\n", + " r_pts[i] * np.sin(np.deg2rad(phis[k])),\n", + " z_pts[j],\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "### Cross-sectional view of the radiation\n", + "\n", + "2D poloidal cross-sections of the radiation function at two toroidal angles." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "phi_selected = [95, 330] # toroidal angles to visualize\n", + "\n", + "# Get the maximum value\n", + "vmax = -np.inf\n", + "for phi in phi_selected:\n", + " i_phi = np.argmin(np.abs(phis - phi))\n", + " vmax = max(vmax, np.max(rad[:, :, i_phi] * TOGW))\n", + "\n", + "\n", + "fig, axs = uplt.subplots(ncols=len(phi_selected), spanx=False)\n", + "for i, phi in enumerate(phi_selected):\n", + " i_phi = np.argmin(np.abs(phis - phi))\n", + "\n", + " im = axs[i].pcolormesh(\n", + " r_pts,\n", + " z_pts,\n", + " rad[:, :, i_phi].T * TOGW,\n", + " shading=\"auto\",\n", + " cmap=\"inferno\",\n", + " discrete=False,\n", + " vmax=vmax,\n", + " vmin=0,\n", + " )\n", + " axs[i].format(\n", + " axesedgecolor=f\"C{i}\",\n", + " urtitle=f\"{phis[i_phi]:.1f}°\",\n", + " titleborder=False,\n", + " aspect=\"equal\",\n", + " xlabel=\"$R$ [m]\",\n", + " ylabel=\"$Z$ [m]\",\n", + " )\n", + "\n", + "fig.colorbar(\n", + " im,\n", + " ax=axs,\n", + " loc=\"b\",\n", + " label=\"[GW/m³]\",\n", + " shrink=0.5,\n", + " tickminor=True,\n", + ");" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "### Z-integrated X-Y projection\n", + "\n", + "The radiation distribution is integrated along the z-axis and projected onto the X-Y plane.\n", + "\n", + "Each toroidal angle at which the cross-section was plotted above is also shown as a colored line in the X-Y projection." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": { + "tags": [ + "nbsphinx-thumbnail" + ] + }, + "outputs": [], + "source": [ + "fig, ax = uplt.subplots(\n", + " proj=\"polar\",\n", + " dpi=150,\n", + " refwidth=4,\n", + ")\n", + "\n", + "im = ax.pcolormesh(\n", + " np.deg2rad(phis),\n", + " r_pts,\n", + " rad.sum(axis=1) * TOGW * dz,\n", + " shading=\"auto\",\n", + " cmap=\"inferno\",\n", + " discrete=False,\n", + " vmin=0,\n", + ")\n", + "fig.colorbar(\n", + " im,\n", + " ax=ax,\n", + " label=\"[GW/m²]\",\n", + " loc=\"b\",\n", + " shrink=0.35,\n", + " ticks=2,\n", + " tickminor=True,\n", + " orientation=\"horizontal\",\n", + " pad=-15,\n", + ")\n", + "\n", + "for i, phi in enumerate(phi_selected):\n", + " ax.axvline(\n", + " np.deg2rad(phis[np.argmin(np.abs(phis - phi))]),\n", + " color=uplt.set_alpha(f\"C{i}\", 0.5),\n", + " lw=4,\n", + " )\n", + "ax.format(\n", + " rmin=r_pts.min(),\n", + " rmax=r_pts.max(),\n", + " rformatter=\"none\",\n", + " r0=r_pts.min() / r_pts.max(),\n", + " thetalocator=30,\n", + " rlabelpos=0,\n", + " rlines=1,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## Visualize the emitter in 3D\n", + "\n", + "Resample the radiation function between the toroidal angle 270° and 360° (quarter of the torus) in (X, Y, Z) coordinates and visualize the result in 3D.\n", + "\n", + "To plot 3D data, we use the [plotly](https://plotly.com/python/) library. The 3D radiation distribution is visualized as a volume rendering." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "import plotly.graph_objects as go\n", + "from plotly import io\n", + "\n", + "io.renderers.default = \"notebook\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "dr = dz = 10e-2\n", + "x_pts = np.arange(0, R_MAX, dr)\n", + "y_pts = np.arange(0, -R_MAX, -dr)\n", + "z_pts = np.arange(Z_MIN, Z_MAX, dz)\n", + "X, Y, Z = np.meshgrid(x_pts, y_pts, z_pts)\n", + "\n", + "rad_xyz = np.zeros((len(x_pts), len(y_pts), len(z_pts)), dtype=float)\n", + "for i, j, k in np.ndindex(len(x_pts), len(y_pts), len(z_pts)):\n", + " rad_xyz[i, j, k] = emitter.material.radiation_function(x_pts[i], y_pts[j], z_pts[k])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "fig = go.Figure(\n", + " data=go.Volume(\n", + " x=X.flatten(),\n", + " y=Y.flatten(),\n", + " z=Z.flatten(),\n", + " value=rad_xyz.flatten(),\n", + " isomin=1e8,\n", + " isomax=rad_xyz.max(),\n", + " opacity=0.1, # needs to be small to see through all surfaces\n", + " surface_count=21, # needs to be a large number for good volume rendering\n", + " )\n", + ")\n", + "fig.update_layout(\n", + " scene=dict(\n", + " xaxis_title=\"X [m]\",\n", + " yaxis_title=\"Y [m]\",\n", + " zaxis_title=\"Z [m]\",\n", + " aspectmode=\"data\",\n", + " ),\n", + " margin=dict(l=0, r=0, b=0, t=0),\n", + " template=\"plotly_dark\",\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "docs", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/notebooks/radiation/grid_3d.ipynb b/docs/notebooks/radiation/grid_3d.ipynb new file mode 100644 index 0000000..7e4f026 --- /dev/null +++ b/docs/notebooks/radiation/grid_3d.ipynb @@ -0,0 +1,349 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 3D Grid Visualization\n", + "\n", + "This notebook demonstrates how to load and visualize a 3-D grid from Generalized Grid Description (GGD) in IMAS.\n", + "\n", + "The example test data was calculated by JOREK for an ITER disruption scenario." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import ultraplot as uplt\n", + "from imas import DBEntry\n", + "from rich import print as rprint\n", + "from rich.table import Table\n", + "\n", + "from cherab.imas.datasets import iter_jorek\n", + "from cherab.imas.ids.common import get_ids_time_slice\n", + "from cherab.imas.ids.common.ggd import load_grid\n", + "\n", + "# Set dark background for plots\n", + "uplt.rc.style = \"dark_background\"" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Retrieve ITER JOREK sample data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "path = iter_jorek()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "with DBEntry(path, \"r\") as entry:\n", + " ids = get_ids_time_slice(entry, \"radiation\")\n", + " grid = load_grid(ids.grid_ggd[0])" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "Show the grid specification" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "table = Table(show_header=False, title=\"Grid specification\")\n", + "table.add_row(\"Grid name\", str(grid.name))\n", + "table.add_row(\"Number Faces\", str(grid.num_faces))\n", + "table.add_row(\"Number Toroidal\", str(grid.num_toroidal))\n", + "table.add_row(\"Number Cell\", str(grid.num_cell))\n", + "table.add_row(\"Shape of vertices array\", str(grid.vertices.shape))\n", + "table.add_row(\"Shape of cells array\", str(grid.cells.shape))\n", + "rprint(table)" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## Visualize the cross-section of the grid\n", + "\n", + "Show the grid lines and cell volume map in the poloidal cross-section of the grid." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axs = uplt.subplots(ncols=2)\n", + "# Plot the grid mesh in 2D cross-section\n", + "grid.plot_mesh(ax=axs[0], edgecolor=\"red\")\n", + "\n", + "# Plot the cell volume map in 2D cross-section\n", + "ax = grid.plot_mesh(ax=axs[1], data=grid.cell_volume[: grid.num_faces])\n", + "ax.collections[0].set_cmap(\"magma\")\n", + "\n", + "ax.colorbar(\n", + " ax.collections[0],\n", + " tickdir=\"out\",\n", + " loc=\"lr\",\n", + " orientation=\"vertical\",\n", + " ticklabelsize=\"small\",\n", + " length=5,\n", + " frame=False,\n", + ")\n", + "ax.format(\n", + " urtitle=\"Cell volume [m$^3$]\",\n", + " titleborder=False,\n", + ")\n", + "\n", + "# Format the axes\n", + "axs.format(xlocator=1, ylocator=1)" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "Plot center points of cells\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": { + "tags": [ + "nbsphinx-thumbnail" + ] + }, + "outputs": [], + "source": [ + "# Extract cell center points for one toroidal slice\n", + "cell_centers = grid.cell_centre[: grid.num_faces, :]\n", + "\n", + "# Calculate the (r, z) coordinates\n", + "r_coords = np.hypot(cell_centers[:, 0], cell_centers[:, 1])\n", + "z_coords = cell_centers[:, 2]\n", + "\n", + "fig, ax = uplt.subplots()\n", + "\n", + "# Plot grid and cell centers in the poloidal cross-section\n", + "grid.plot_mesh(\n", + " ax=ax,\n", + " edgecolor=\"red\",\n", + ")\n", + "ax.scatter(r_coords, z_coords, s=1, c=\"C0\")\n", + "ax.format(\n", + " xlocator=1,\n", + " ylocator=1,\n", + ")\n", + "\n", + "# Zoom in around divertor region\n", + "ix = ax.inset(\n", + " [7.0, -9, 6, 6],\n", + " transform=\"data\",\n", + " zoom_kw={\"ec\": \"grape3\", \"ls\": \"--\", \"lw\": 2},\n", + ")\n", + "grid.plot_mesh(ax=ix, edgecolor=\"red\")\n", + "ix.scatter(r_coords, z_coords, s=1, c=\"C0\")\n", + "ix.format(\n", + " xlim=(ax.get_xlim()[0], 6.2),\n", + " ylim=(ax.get_ylim()[0], -3.0),\n", + " aspect=\"equal\",\n", + " color=\"grape9\",\n", + " linewidth=1.5,\n", + " ticklabelweight=\"bold\",\n", + " xlocator=1,\n", + " ylocator=1,\n", + " xformatter=\"none\",\n", + " yformatter=\"none\",\n", + " xlabel=\"\",\n", + " ylabel=\"\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Cut out quarter of the torus" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "grid_cut = grid.subset(np.arange(grid.num_faces * grid.num_toroidal // 4))" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "Visualize the 3D grid in a quarter of the torus.\n", + "[plotly](https://plotly.com/python/) is used to visualize the tetrahedral mesh in the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "import plotly.graph_objects as go\n", + "from plotly import io\n", + "\n", + "io.renderers.default = \"notebook\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "tetra = grid_cut.tetrahedra\n", + "faces = np.vstack(\n", + " [\n", + " tetra[:, [0, 1, 2]],\n", + " tetra[:, [0, 1, 3]],\n", + " tetra[:, [0, 2, 3]],\n", + " tetra[:, [1, 2, 3]],\n", + " ]\n", + ")\n", + "\n", + "# Keep only boundary faces (faces appearing exactly once).\n", + "faces_sorted = np.sort(faces, axis=1)\n", + "_, inverse, counts = np.unique(faces_sorted, axis=0, return_inverse=True, return_counts=True)\n", + "surface_faces = faces[counts[inverse] == 1]\n", + "\n", + "# Build unique boundary edges from boundary triangles.\n", + "tri_edges = np.vstack(\n", + " [\n", + " surface_faces[:, [0, 1]],\n", + " surface_faces[:, [1, 2]],\n", + " surface_faces[:, [2, 0]],\n", + " ]\n", + ")\n", + "boundary_edges = np.unique(np.sort(tri_edges, axis=1), axis=0)\n", + "\n", + "verts = grid_cut.vertices\n", + "x, y, z = verts[:, 0], verts[:, 1], verts[:, 2]\n", + "\n", + "# Customize edge appearance here.\n", + "edge_color = \"black\"\n", + "edge_width = 2\n", + "\n", + "# Convert edge index pairs to line segments separated by NaN for Plotly.\n", + "edge_xyz = verts[boundary_edges]\n", + "xe = np.column_stack(\n", + " [edge_xyz[:, 0, 0], edge_xyz[:, 1, 0], np.full(len(boundary_edges), np.nan)]\n", + ").ravel()\n", + "ye = np.column_stack(\n", + " [edge_xyz[:, 0, 1], edge_xyz[:, 1, 1], np.full(len(boundary_edges), np.nan)]\n", + ").ravel()\n", + "ze = np.column_stack(\n", + " [edge_xyz[:, 0, 2], edge_xyz[:, 1, 2], np.full(len(boundary_edges), np.nan)]\n", + ").ravel()\n", + "\n", + "fig = go.Figure(\n", + " data=[\n", + " go.Mesh3d(\n", + " x=x,\n", + " y=y,\n", + " z=z,\n", + " i=surface_faces[:, 0],\n", + " j=surface_faces[:, 1],\n", + " k=surface_faces[:, 2],\n", + " color=\"deepskyblue\",\n", + " flatshading=True,\n", + " hoverinfo=\"skip\",\n", + " showscale=False,\n", + " name=\"Surface\",\n", + " ),\n", + " go.Scatter3d(\n", + " x=xe,\n", + " y=ye,\n", + " z=ze,\n", + " mode=\"lines\",\n", + " line=dict(color=edge_color, width=edge_width),\n", + " name=\"Surface boundary\",\n", + " hoverinfo=\"skip\",\n", + " ),\n", + " ]\n", + ")\n", + "\n", + "fig.update_layout(\n", + " title=\"Quarter-torus tetrahedral mesh\",\n", + " scene=dict(\n", + " xaxis_title=\"X [m]\",\n", + " yaxis_title=\"Y [m]\",\n", + " zaxis_title=\"Z [m]\",\n", + " aspectmode=\"data\",\n", + " ),\n", + " margin=dict(l=0, r=0, b=0, t=40),\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "docs", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/examples.md b/docs/source/examples.md index b8fc4b0..d64a3e5 100644 --- a/docs/source/examples.md +++ b/docs/source/examples.md @@ -14,6 +14,17 @@ Each notebook is designed to showcase specific functionalities and provide pract notebooks/plasma/* ``` +## Radiation + +```{eval-rst} +.. nbgallery:: + :name: radiation-gallery + :glob: + :reversed: + + notebooks/radiation/* +``` + ## Observer ```{eval-rst} diff --git a/pixi.toml b/pixi.toml index 5764ed5..4dbdded 100644 --- a/pixi.toml +++ b/pixi.toml @@ -45,6 +45,7 @@ ultraplot = ">=1.72.0" # Visualization for 3D-related stuff plotly = ">=6.5.2,<7" +python-kaleido = ">=1.2.0,<2" [tasks] ipython = { cmd = "ipython", description = "🐍 Start an IPython shell" } diff --git a/src/cherab/imas/emitter/__init__.py b/src/cherab/imas/emitter/__init__.py new file mode 100644 index 0000000..e1249e4 --- /dev/null +++ b/src/cherab/imas/emitter/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2023 Euratom +# Copyright 2023 United Kingdom Atomic Energy Authority +# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# 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. +"""Subpackage for creating emitter objects from IMAS.""" + +from .radiation import load_radiation_emitter + +__all__ = [ + "load_radiation_emitter", +] diff --git a/src/cherab/imas/emitter/radiation.py b/src/cherab/imas/emitter/radiation.py new file mode 100644 index 0000000..e3369b3 --- /dev/null +++ b/src/cherab/imas/emitter/radiation.py @@ -0,0 +1,231 @@ +# Copyright 2023 Euratom +# Copyright 2023 United Kingdom Atomic Energy Authority +# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# 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. +"""Module for loading radiation emissivity from IMAS-like IDS objects and creating emitter objects.""" + +from pathlib import Path +from typing import Literal + +import numpy as np +from raysect.core.math import translate +from raysect.core.math.function.float import Function2D +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 imas import DBEntry +from imas.ids_structure import IDSStructure + +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 ..math import FourierBezierConstructor +from ..plasma.utility import get_subset_name_index + +__all__ = ["load_radiation_emitter"] + + +def load_radiation_emitter( + *args, + time: float = 0, + occurrence: int = 0, + process_index: int | None = None, + ion_index: int = 0, + emissivity_index: int = 0, + grid_ggd: IDSStructure | None = None, + grid_subset_id: int | str = 5, + num_toroidal: int | None = None, + phis: np.ndarray | None = None, + source: Literal["auto", "values", "coefficients"] = "auto", + 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. + + Parameters + ---------- + *args + Positional arguments passed to `imas.DBEntry`. + time + Time slice to load from the IDS, by default 0. + occurrence + Occurrence of the radiation IDS to load, 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. + grid_ggd + Alternative grid GGD structure to use if the radiation IDS grid is empty, by default None. + grid_subset_id + ID or name of the grid subset to use, by default 5 (``"Cells"``). + num_toroidal + Number of toroidal subdivisions for 3D grid extension, by default None. + 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"``. + step + Step size for the radiation function interpolator, by default 0.01 m. + parent + Parent node in the Raysect scenegraph, by default None. + time_threshold + Maximum allowed time difference when loading from IDS, by default ``inf``. + interpolator_cache + Interpolator cache strategy, by default ``"memory"``. + Each strategy is described in the `.InterpolatorCacheMode` type alias. + interpolator_cache_dir + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). + **kwargs + Additional keyword arguments passed to `imas.DBEntry`. + + Returns + ------- + `~raysect.primitive.csg.Subtract` or `~raysect.primitive.Cylinder` + Cylindrical emitter primitive with + `~cherab.tools.emitters.radiation_function.RadiationFunction` material. + + Raises + ------ + RuntimeError + If the radiation IDS or its emissivity data cannot be loaded. + """ + with DBEntry(*args, **kwargs) as entry: + radiation_ids = get_ids_time_slice( + entry, + "radiation", + time=time, + occurrence=occurrence, + time_threshold=time_threshold, + ) + + 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." + ) + + grid_ggd_struct = grid_ggd or radiation_ids.grid_ggd[0] + + try: + grid, subsets, subset_id = load_grid( + grid_ggd_struct, + with_subsets=True, + num_toroidal=num_toroidal, + ) + 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 + + 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, + ) + 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, + ) + + constructor = FourierBezierConstructor(grid_ggd_struct, 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() + + 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) + + emitter = RadiationFunction(rad_func, step=step) + + 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"] + + if radius_inner > 0: + primitive = Subtract( + Cylinder(radius_outer, height), Cylinder(radius_inner, height), parent=parent + ) + else: + primitive = Cylinder(radius_outer, height, parent=parent) + + primitive.transform = translate(0, 0, zmin) + primitive.material = emitter + primitive.name = f"RadiationEmitter_{radiation_ids.time[0]}s, uri {entry.uri}" + + return primitive diff --git a/src/cherab/imas/ggd/base_mesh.py b/src/cherab/imas/ggd/base_mesh.py index 22dcbd4..8308ff2 100755 --- a/src/cherab/imas/ggd/base_mesh.py +++ b/src/cherab/imas/ggd/base_mesh.py @@ -41,6 +41,11 @@ CellSelection: TypeAlias = Sequence[SupportsIndex] | NDArray[np.integer[Any]] InterpolatorCacheMode: TypeAlias = Literal["none", "memory", "disk"] """Cache mode for interpolator templates. + +- ``"none"``: No caching; build a new interpolator template on each call. +- ``"memory"``: Cache the interpolator template in memory for the current process lifetime. +- ``"disk"``: Cache the interpolator template on disk for reuse across processes and sessions. + The cache directory can be specified via the `interpolator_cache_dir` argument. """ InterpolatorT = TypeVar("InterpolatorT") @@ -465,11 +470,13 @@ def interpolator( fill_value A value returned outside the grid, by default is 0.0. interpolator_cache - Cache mode for the interpolator. See `InterpolatorCacheMode`. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- @@ -500,11 +507,13 @@ def vector_interpolator( fill_vector 3D vector returned outside the grid, by default ``Vector3D(0, 0, 0)``. interpolator_cache - Cache mode for the interpolator. See `InterpolatorCacheMode`. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- diff --git a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py index a069580..8572ced 100644 --- a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py +++ b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py @@ -444,11 +444,13 @@ def interpolator( fill_value Value returned outside the grid, by default 0.0. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- @@ -494,13 +496,15 @@ def vector_interpolator( grid_vectors ``(3, L)`` Array containing 3D vectors in the grid cells. fill_vector - 3D vector returned outside the grid, by default is `Vector3D(0, 0, 0)`. + 3D vector returned outside the grid, by default `Vector3D(0, 0, 0)`. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- diff --git a/src/cherab/imas/ggd/unstruct_2d_mesh.py b/src/cherab/imas/ggd/unstruct_2d_mesh.py index 2bb148f..91abb86 100755 --- a/src/cherab/imas/ggd/unstruct_2d_mesh.py +++ b/src/cherab/imas/ggd/unstruct_2d_mesh.py @@ -355,11 +355,13 @@ def interpolator( fill_value Value returned outside the grid, by default 0.0. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- @@ -405,13 +407,15 @@ def vector_interpolator( grid_vectors ``(3, K)`` Array containing 3D vectors in the grid cells. fill_vector - 3D vector returned outside the grid, by default is `Vector3D(0, 0, 0)`. + 3D vector returned outside the grid, by default `Vector3D(0, 0, 0)`. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- diff --git a/src/cherab/imas/ggd/unstruct_3d_mesh.py b/src/cherab/imas/ggd/unstruct_3d_mesh.py index b7ec3fb..08cbf6b 100644 --- a/src/cherab/imas/ggd/unstruct_3d_mesh.py +++ b/src/cherab/imas/ggd/unstruct_3d_mesh.py @@ -271,11 +271,13 @@ def interpolator( fill_value Value returned outside the grid, by default 0.0. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- @@ -321,13 +323,15 @@ def vector_interpolator( grid_vectors ``(3, L)`` Array containing 3D vectors in the grid cells. fill_vector - 3D vector returned outside the grid, by default is `Vector3D(0, 0, 0)`. + 3D vector returned outside the grid, by default `Vector3D(0, 0, 0)`. interpolator_cache - Cache mode for the interpolator. + Cache mode for the interpolator, by default ``"memory"``. + The cache mode is described in the `.InterpolatorCacheMode` type alias. interpolator_cache_dir - Directory used for disk cache mode. + Directory used when ``interpolator_cache="disk"``, by default None + (uses the system cache directory, e.g., ``~/.cache/cherab/imas/interpolators``). interpolator_cache_namespace - Namespace prefix to avoid cache-key collisions. + Namespace prefix to avoid cache-key collisions, by default ``"ggd"``. Returns ------- diff --git a/src/cherab/imas/ids/common/slice.py b/src/cherab/imas/ids/common/slice.py index 93ecd89..f6fe2e9 100644 --- a/src/cherab/imas/ids/common/slice.py +++ b/src/cherab/imas/ids/common/slice.py @@ -18,16 +18,63 @@ """Module for common functions used to get IDS time slices.""" import warnings +from uuid import uuid4 from numpy import inf from imas.db_entry import DBEntry -from imas.ids_defs import CLOSEST_INTERP +from imas.ids_defs import CLOSEST_INTERP, MEMORY_BACKEND from imas.ids_toplevel import IDSToplevel __all__ = ["get_ids_time_slice"] +def _slice_via_memory_backend( + ids: IDSToplevel, + ids_name: str, + time: float, + occurrence: int, +) -> IDSToplevel: + """Re-slice an IDS by round-tripping through the IMAS memory backend. + + Parameters + ---------- + ids + The IDS to re-slice. + ids_name + The name of the IDS. + time + The time in seconds of the requested time slice. + occurrence + The occurrence of the IDS. + + Returns + ------- + `~imas.ids_toplevel.IDSToplevel` + The re-sliced IDS. + """ + token = uuid4().int + # Use per-call identifiers so repeated/concurrent calls do not collide. + temp_entry = DBEntry( + MEMORY_BACKEND, + f"cherab_tmp_{token & 0xFFFF:04x}", + 1 + token % 2_000_000_000, + 1 + (token >> 31) % 2_000_000_000, + ) + temp_entry.create() + try: + temp_entry.put(ids, occurrence=occurrence) + return temp_entry.get_slice( + ids_name, + time, + CLOSEST_INTERP, + occurrence=occurrence, + autoconvert=False, + ) + finally: + temp_entry.close() + + def get_ids_time_slice( entry: DBEntry, ids_name: str, @@ -39,8 +86,9 @@ def get_ids_time_slice( .. note:: If the `~imas.db_entry.DBEntry.get_slice` method is not implemented for the given IMAS entry - URI, this function will fall back to using the `~imas.db_entry.DBEntry.get` method and - return the entire IDS. + URI, this function falls back to `~imas.db_entry.DBEntry.get` and tries to re-slice the IDS + by round-tripping through the IMAS memory backend. If that second step fails, it returns the + full IDS with a warning. Parameters ---------- @@ -89,6 +137,8 @@ def get_ids_time_slice( if time_threshold < 0: raise ValueError(f"Argument 'time_threshold' must be >=0 ({time_threshold} s).") + is_time_sliced = True + try: ids = entry.get_slice( ids_name, @@ -98,23 +148,60 @@ def get_ids_time_slice( autoconvert=False, ) except NotImplementedError: - # Fallback to `get` method to retrieve the entire IDS - warnings.warn( - f"The 'get_slice' method is not implemented for the URI '{entry.uri}'. " - + "Falling back to 'get' method to retrieve the entire IDS.", - RuntimeWarning, - stacklevel=2, - ) ids = entry.get(ids_name, occurrence=occurrence, autoconvert=False) + is_time_sliced = len(ids.time) <= 1 + get_slice_unavailable_msg = ( + f"The 'get_slice' method is not implemented for the URI '{entry.uri}'." + ) + + if not is_time_sliced: + try: + ids = _slice_via_memory_backend(ids, ids_name, time, occurrence) + is_time_sliced = True + warnings.warn( + get_slice_unavailable_msg + + " " + + "Falling back to 'get' and re-slicing via the IMAS memory backend.", + RuntimeWarning, + stacklevel=2, + ) + except Exception as exc: + warnings.warn( + get_slice_unavailable_msg + + " " + + "Fallback re-slicing via the IMAS memory backend failed. " + + "Returning the full " + + f"'{ids_name}' IDS without reducing to a single time slice. Error: {exc}", + RuntimeWarning, + stacklevel=2, + ) + else: + warnings.warn( + get_slice_unavailable_msg + + " " + + "Falling back to 'get' method because the returned IDS contains a single time slice.", + RuntimeWarning, + stacklevel=2, + ) if not len(ids.time): raise RuntimeError(f"The '{ids_name}' IDS is empty.") - if abs(ids.time[0] - time) > time_threshold: + nearest_time = min(ids.time, key=lambda t: abs(float(t) - time)) + + if abs(float(nearest_time) - time) > time_threshold: raise RuntimeError( - f"The time difference between the actual time ({ids.time[0]} s) " + f"The time difference between the actual time ({nearest_time} s) " + f"of the nearest '{ids_name}' time slice and the given time ({time} s) " + f"exceeds the specified threshold ({time_threshold} s)." ) + if not is_time_sliced: + warnings.warn( + f"Returning '{ids_name}' IDS with {len(ids.time)} time slices because a single-time " + + "fallback could not be constructed.", + RuntimeWarning, + stacklevel=2, + ) + return ids diff --git a/src/cherab/imas/ids/radiation/__init__.py b/src/cherab/imas/ids/radiation/__init__.py new file mode 100644 index 0000000..96d8ffc --- /dev/null +++ b/src/cherab/imas/ids/radiation/__init__.py @@ -0,0 +1,22 @@ +# Copyright 2023 Euratom +# Copyright 2023 United Kingdom Atomic Energy Authority +# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# 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. +"""Subpackage for loading radiation data from IMAS IDS structures.""" + +from .load_radiation import load_radiation_coefficients, load_radiation_emissivity + +__all__ = ["load_radiation_emissivity", "load_radiation_coefficients"] diff --git a/src/cherab/imas/ids/radiation/load_radiation.py b/src/cherab/imas/ids/radiation/load_radiation.py new file mode 100644 index 0000000..1489802 --- /dev/null +++ b/src/cherab/imas/ids/radiation/load_radiation.py @@ -0,0 +1,212 @@ +# Copyright 2023 Euratom +# Copyright 2023 United Kingdom Atomic Energy Authority +# Copyright 2023 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# 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. +"""Module for loading radiation emissivity from the radiation IDS.""" + +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"] + + +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. + + Returns + ------- + `NDArray[numpy.float64]` or None + Emissivity values [W/m³], or ``None`` if the requested subset is not found. + """ + emissivity_arr = getattr(ggd_struct, "emissivity", None) + if not isinstance(emissivity_arr, IDSStructArray) or not len(emissivity_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 + + +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"``). + + 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. + """ + 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.") + + 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 " + "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]." + ) + + return values + + +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. + + This accessor targets the coefficient layout typically used by JOREK: + ``radiation.process[i].ggd[0].ion[j].emissivity[k].coefficients``. + + 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. + + Returns + ------- + `NDArray[numpy.float64]` + Coefficient array, typically shaped ``(num_vertices * num_modes, 4)``. + + Raises + ------ + RuntimeError + If the requested process/ion/emissivity structure is missing or empty. + """ + 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.") + + 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) diff --git a/src/cherab/imas/math/functions/fourier_bezier.pxd b/src/cherab/imas/math/functions/fourier_bezier.pxd index 0f9b59f..4620bd6 100644 --- a/src/cherab/imas/math/functions/fourier_bezier.pxd +++ b/src/cherab/imas/math/functions/fourier_bezier.pxd @@ -1,5 +1,4 @@ from libc.math cimport cos, sin -cimport numpy as np DEF TO_RAD = 3.14159265358979323846 / 180.0 diff --git a/src/cherab/imas/math/functions/fourier_bezier.pyx b/src/cherab/imas/math/functions/fourier_bezier.pyx index af80423..a9f8bcf 100644 --- a/src/cherab/imas/math/functions/fourier_bezier.pyx +++ b/src/cherab/imas/math/functions/fourier_bezier.pyx @@ -140,7 +140,7 @@ cdef class FourierBezierConstructor: self._num_faces = len(sp_rz.objects_per_dimension[DomainType.face].object) self._num_vertices = len(sp_rz.objects_per_dimension[DomainType.vertex].object) self._num_toroidal_modes = len(sp_fourier.objects_per_dimension[DomainType.vertex].object) - self._fourier_periodicity = sp_fourier.geometry_type.index + self._fourier_periodicity = int(np.asarray(sp_fourier.geometry_type.index).item()) self._vertex_indices = np.empty((self._num_faces, 4), dtype=np.int32) self._scale_factors = np.empty((self._num_faces, 4, 4), dtype=np.double) @@ -194,6 +194,10 @@ cdef class FourierBezierConstructor: ) if coefficients.shape[1] != 4: raise ValueError("Coefficients array must have 4 columns.") + + if not coefficients.flags["C_CONTIGUOUS"]: + coefficients = np.ascontiguousarray(coefficients) + self._coefficients = self._set_coefficients(coefficients) @cython.boundscheck(False) @@ -581,9 +585,9 @@ def py_fourier_mode(double phi, int index, int periodicity = 1) -> float: .. math:: Z_l(\varphi) = \begin{cases} 1 - & \text{if } l = 0, \ + & \text{if } l = 0, \\ \sin\left(\displaystyle\frac{l}{2}n_\mathrm{p} \varphi\right) - & \text{if } l \text{ is even}, \ + & \text{if } l \text{ is even}, \\ \cos\left(\displaystyle\frac{l + 1}{2}n_\mathrm{p} \varphi\right) & \text{if } l \text{ is odd}. \end{cases} diff --git a/src/cherab/imas/plasma/blend.py b/src/cherab/imas/plasma/blend.py index 294dc2c..730c05a 100644 --- a/src/cherab/imas/plasma/blend.py +++ b/src/cherab/imas/plasma/blend.py @@ -270,11 +270,17 @@ def load_plasma( grid_ggd = grid_ggd or edge_profiles_ids.grid_ggd[0] 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 np.all(subsets[grid_subset_name] != np.arange(grid.num_cell, dtype=int)): - # To reduce memory usage, create the sub-grid only if needed. - grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) + 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)): + # To reduce memory usage, create the sub-grid only if needed. + grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) + except ValueError: + print( + f"Warning! Grid subset with identifier '{grid_subset_id}' not found in {subset_id}.", + ) + grid_subset_index = 5 # Default to 'Cells' subset index 5 composition_edge = load_edge_species( edge_profiles_ids.ggd[0], @@ -413,7 +419,7 @@ def load_plasma( # Add the blended species to the plasma composition for (element, charge), interp in species.items(): try: - species = plasma.composition.get(element, int(charge)) + species = plasma.composition.get(element, charge) print(f"Warning! Skipping {species}: already defined") continue except ValueError: @@ -426,7 +432,7 @@ def load_plasma( element.atomic_weight * atomic_mass, ) - plasma.composition.add(Species(element, int(charge), distribution)) + plasma.composition.add(Species(element, charge, distribution)) # === Ion Bundle Species === # Ion bundles are split into their constituent charge states at the composition level. diff --git a/src/cherab/imas/plasma/core.py b/src/cherab/imas/plasma/core.py index 96a96ac..616ad23 100644 --- a/src/cherab/imas/plasma/core.py +++ b/src/cherab/imas/plasma/core.py @@ -232,7 +232,7 @@ def load_core_plasma( charge = profile.species.z_min try: - species = plasma.composition.get(element, int(charge)) + species = plasma.composition.get(element, charge) print(f"Warning! Skipping {species}: already defined") continue except ValueError: @@ -247,7 +247,7 @@ def load_core_plasma( element.atomic_weight * atomic_mass, ) - plasma.composition.add(Species(element, int(charge), distribution)) + plasma.composition.add(Species(element, charge, distribution)) # === Ion Bundles === # Ion bundles are split into their constituent charge states at the composition level. diff --git a/src/cherab/imas/plasma/edge.py b/src/cherab/imas/plasma/edge.py index 07f03cd..67c3b97 100644 --- a/src/cherab/imas/plasma/edge.py +++ b/src/cherab/imas/plasma/edge.py @@ -156,11 +156,16 @@ def load_edge_plasma( grid_ggd = grid_ggd or edge_profiles_ids.grid_ggd[0] 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 np.all(subsets[grid_subset_name] != np.arange(grid.num_cell, dtype=int)): - # To reduce memory usage, create the sub-grid only if needed. - grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) + 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)): + # To reduce memory usage, create the sub-grid only if needed. + grid = grid.subset(subsets[grid_subset_name], name=grid_subset_name) + except ValueError: + print( + f"Warning! Grid subset with identifier '{grid_subset_id}' not found in {subset_id}.", + ) + grid_subset_index = 5 # Default to 'Cells' subset index 5 # Load species composition composition = load_edge_species( @@ -239,7 +244,7 @@ def load_edge_plasma( charge = profile.species.z_min try: - species = plasma.composition.get(element, int(charge)) + species = plasma.composition.get(element, charge) print(f"Warning! Skipping {species}: already defined") continue except ValueError: @@ -254,7 +259,7 @@ def load_edge_plasma( element.atomic_weight * atomic_mass, ) - plasma.composition.add(Species(element, int(charge), distribution)) + plasma.composition.add(Species(element, charge, distribution)) # === Ion Bundles === # Ion bundles are split into their constituent charge states at the composition level. diff --git a/tests/conftest.py b/tests/conftest.py index c71b876..6f03df7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,10 @@ +import shutil +from pathlib import Path + import pytest from cherab.core.atomic.elements import neon +from cherab.imas.datasets import iter_jintrac, iter_jorek, iter_solps from cherab.openadas import OpenADAS from cherab.openadas.repository import populate @@ -14,3 +18,35 @@ def populate_openadas_repository(): except Exception: print("Populating OpenADAS repository...") populate() + + +def _copy_dataset_to_tmp(path: Path, tmp_path_factory: pytest.TempPathFactory) -> str: + """Copy a dataset into a temporary location, handling files and directories.""" + tmp_path = tmp_path_factory.mktemp("cherab-imas-data") + target = tmp_path / path.name + if path.is_dir(): + shutil.copytree(path, target) + else: + shutil.copy2(path, target) + return str(target) + + +@pytest.fixture(scope="session") +def path_iter_jintrac(tmp_path_factory: pytest.TempPathFactory) -> str: + """Fixture to provide the path to a sample JINTRAC IMAS dataset.""" + path = Path(iter_jintrac()) + return _copy_dataset_to_tmp(path, tmp_path_factory) + + +@pytest.fixture(scope="session") +def path_iter_solps(tmp_path_factory: pytest.TempPathFactory) -> str: + """Fixture to provide the path to a sample SOLPS IMAS dataset.""" + path = Path(iter_solps()) + return _copy_dataset_to_tmp(path, tmp_path_factory) + + +@pytest.fixture(scope="session") +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) diff --git a/tests/emitter/test_radiation.py b/tests/emitter/test_radiation.py new file mode 100644 index 0000000..d6831e2 --- /dev/null +++ b/tests/emitter/test_radiation.py @@ -0,0 +1,154 @@ +import os +from pathlib import Path +from typing import Literal, TypedDict + +import numpy as np +import pytest +from raysect.primitive import Cylinder, Subtract + +import cherab.imas.emitter.radiation as radiation_module +from cherab.imas.emitter import load_radiation_emitter + + +class _EmitterCacheKwargs(TypedDict, total=False): + interpolator_cache: Literal["memory", "disk"] + interpolator_cache_dir: str | Path + + +@pytest.fixture(scope="session") +def radiation_interpolator_cache( + tmp_path_factory: pytest.TempPathFactory, +) -> tuple[Literal["memory", "disk"], Path | None]: + """Choose interpolator cache mode for radiation tests. + + Default is "memory" for speed. Set CHERAB_IMAS_RADIATION_TEST_CACHE=disk to + reuse cache artifacts across test cases inside one session. + """ + mode = os.getenv("CHERAB_IMAS_RADIATION_TEST_CACHE", "memory") + if mode not in {"memory", "disk"}: + mode = "memory" + + if mode == "disk": + return mode, tmp_path_factory.mktemp("cherab-imas-radiation-cache") + + return mode, None + + +def _cache_kwargs( + cache_cfg: tuple[Literal["memory", "disk"], Path | None], +) -> _EmitterCacheKwargs: + mode, cache_dir = cache_cfg + kwargs: _EmitterCacheKwargs = {"interpolator_cache": mode} + if cache_dir is not None: + kwargs["interpolator_cache_dir"] = cache_dir + return kwargs + + +def test_load_radiation_emitter_coefficients_uses_default_phis( + path_iter_jorek: str, + monkeypatch: pytest.MonkeyPatch, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + captured: dict[str, np.ndarray] = {} + original_constructor = radiation_module.FourierBezierConstructor + + class _SpyFourierBezierConstructor: + def __init__(self, *args, **kwargs): + self._inner = original_constructor(*args, **kwargs) + + def average_gaussian_faces_per_toroidal(self, phis): + captured["phis"] = np.asarray(phis, dtype=np.float64).copy() + return self._inner.average_gaussian_faces_per_toroidal(phis) + + monkeypatch.setattr(radiation_module, "FourierBezierConstructor", _SpyFourierBezierConstructor) + + primitive = load_radiation_emitter( + path_iter_jorek, + "r", + source="coefficients", + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + assert primitive.name.startswith("RadiationEmitter_") + + phis_used = captured["phis"] + assert phis_used.ndim == 1 + assert phis_used.size > 0 + assert np.all(np.diff(phis_used) > 0) + assert 0.0 < phis_used[0] < 360.0 + assert 0.0 < phis_used[-1] < 360.0 + + +def test_load_radiation_emitter_auto_falls_back_to_coefficients( + path_iter_jorek: str, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + primitive = load_radiation_emitter( + path_iter_jorek, + "r", + source="auto", + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + + +def test_load_radiation_emitter_coefficients_uses_given_phis( + path_iter_jorek: str, + monkeypatch: pytest.MonkeyPatch, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + captured: dict[str, np.ndarray] = {} + original_constructor = radiation_module.FourierBezierConstructor + + class _SpyFourierBezierConstructor: + def __init__(self, *args, **kwargs): + self._inner = original_constructor(*args, **kwargs) + + def average_gaussian_faces_per_toroidal(self, phis): + captured["phis"] = np.asarray(phis, dtype=np.float64).copy() + return self._inner.average_gaussian_faces_per_toroidal(phis) + + monkeypatch.setattr(radiation_module, "FourierBezierConstructor", _SpyFourierBezierConstructor) + + expected_phis = np.array([15.0, 105.0, 195.0, 285.0], dtype=np.float64) + primitive = load_radiation_emitter( + path_iter_jorek, + "r", + source="coefficients", + phis=expected_phis, + **_cache_kwargs(radiation_interpolator_cache), + ) + + assert isinstance(primitive, (Subtract, Cylinder)) + assert primitive.material is not None + np.testing.assert_allclose(captured["phis"], expected_phis) + + +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"): + load_radiation_emitter( + path_iter_jorek, + "r", + source="values", + **_cache_kwargs(radiation_interpolator_cache), + ) + + +def test_load_radiation_emitter_invalid_source_raises( + path_iter_jorek: str, + radiation_interpolator_cache: tuple[Literal["memory", "disk"], Path | None], +): + with pytest.raises(RuntimeError, match="Unable to load emissivity"): + load_radiation_emitter( + path_iter_jorek, + "r", + source="unsupported", # type: ignore[arg-type] + **_cache_kwargs(radiation_interpolator_cache), + ) diff --git a/tests/ggd/conftest.py b/tests/ggd/conftest.py index 9a2348f..52e92d5 100644 --- a/tests/ggd/conftest.py +++ b/tests/ggd/conftest.py @@ -1,11 +1,7 @@ -import shutil -from pathlib import Path - import numpy as np import pytest from imas import DBEntry -from cherab.imas.datasets import iter_jintrac, iter_jorek from cherab.imas.ggd.unstruct_2d_extend_mesh import UnstructGrid2DExtended from cherab.imas.ggd.unstruct_2d_mesh import UnstructGrid2D from cherab.imas.ggd.unstruct_3d_mesh import UnstructGrid3D @@ -13,31 +9,6 @@ from cherab.imas.ids.common.ggd import load_grid -def _copy_dataset_to_tmp(path: Path, tmp_path_factory: pytest.TempPathFactory) -> str: - """Copy a dataset into a temporary location, handling files and directories.""" - tmp_path = tmp_path_factory.mktemp("cherab-imas-data") - target = tmp_path / path.name - if path.is_dir(): - shutil.copytree(path, target) - else: - shutil.copy2(path, target) - return str(target) - - -@pytest.fixture(scope="session") -def path_iter_jintrac(tmp_path_factory) -> str: - """Fixture to provide the path to a sample JINTRAC IMAS dataset.""" - path = Path(iter_jintrac()) - return _copy_dataset_to_tmp(path, tmp_path_factory) - - -@pytest.fixture(scope="session") -def path_iter_jorek(tmp_path_factory) -> 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="module") def jintrac_unstruct_2d_grid(path_iter_jintrac: str) -> UnstructGrid2D: """Fixture to provide a compact UnstructGrid2D loaded from JINTRAC.""" diff --git a/tests/plasma/conftest.py b/tests/plasma/conftest.py index 7e33065..dda407b 100644 --- a/tests/plasma/conftest.py +++ b/tests/plasma/conftest.py @@ -1,33 +1 @@ -import shutil -from pathlib import Path - -import pytest - -from cherab.imas.datasets import iter_jintrac, iter_jorek, iter_solps - - -@pytest.fixture(scope="session") -def path_iter_jintrac(tmp_path_factory) -> str: - """Fixture to provide the path to a sample JINTRAC IMAS dataset.""" - path = Path(iter_jintrac()) - tmp_path = tmp_path_factory.mktemp("cherab-imas-data") - shutil.move(path, tmp_path) - return str(tmp_path / path.name) - - -@pytest.fixture(scope="session") -def path_iter_solps(tmp_path_factory) -> str: - """Fixture to provide the path to a sample SOLPS IMAS dataset.""" - path = Path(iter_solps()) - tmp_path = tmp_path_factory.mktemp("cherab-imas-data") - shutil.move(path, tmp_path) - return str(tmp_path / path.name) - - -@pytest.fixture(scope="session") -def path_iter_jorek(tmp_path_factory) -> str: - """Fixture to provide the path to a sample JOREK IMAS dataset.""" - path = Path(iter_jorek()) - tmp_path = tmp_path_factory.mktemp("cherab-imas-data") - shutil.move(path, tmp_path) - return str(tmp_path / path.name) +"""Plasma test fixtures are provided by tests/conftest.py.""" diff --git a/tests/test_ids_time_slice.py b/tests/test_ids_time_slice.py new file mode 100644 index 0000000..2a6ba90 --- /dev/null +++ b/tests/test_ids_time_slice.py @@ -0,0 +1,60 @@ +import numpy as np +import pytest +from imas.db_entry import DBEntry + +from cherab.imas.ids.common import get_ids_time_slice + + +class _FakeIDS: + def __init__(self, time): + self.time = np.asarray(time) + + +class _FallbackOnlyEntry(DBEntry): + def __init__(self, ids, uri: str = "imas://example"): + self._ids = ids + self.uri = uri + + def get_slice(self, *args, **kwargs): + raise NotImplementedError + + def get(self, *args, **kwargs): + return self._ids + + +def test_get_ids_time_slice_warns_for_single_time_fallback() -> None: + entry = _FallbackOnlyEntry(_FakeIDS([0.25])) + + with pytest.warns(RuntimeWarning, match="Falling back to 'get' method"): + ids = get_ids_time_slice(entry, "radiation", time=0.25) + + assert ids.time[0] == pytest.approx(0.25) + + +def test_get_ids_time_slice_reslices_multi_time_fallback(monkeypatch) -> None: + entry = _FallbackOnlyEntry(_FakeIDS([0.1, 0.2, 0.3])) + + def _fake_reslice(*args, **kwargs): + return _FakeIDS([0.2]) + + monkeypatch.setattr("cherab.imas.ids.common.slice._slice_via_memory_backend", _fake_reslice) + + with pytest.warns(RuntimeWarning, match="re-slicing via the IMAS memory backend"): + ids = get_ids_time_slice(entry, "radiation", time=0.2) + + assert len(ids.time) == 1 + assert ids.time[0] == pytest.approx(0.2) + + +def test_get_ids_time_slice_warns_when_multi_time_reslice_fails(monkeypatch) -> None: + entry = _FallbackOnlyEntry(_FakeIDS([0.1, 0.2, 0.3])) + + def _failing_reslice(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr("cherab.imas.ids.common.slice._slice_via_memory_backend", _failing_reslice) + + with pytest.warns(RuntimeWarning, match="Returning 'radiation' IDS with 3 time slices"): + ids = get_ids_time_slice(entry, "radiation", time=0.2) + + assert len(ids.time) == 3