diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index cadf302..7aa6335 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -15,7 +15,7 @@ jobs: uses: actions/checkout@v5 - name: 🟨 Set up Pixi - uses: prefix-dev/setup-pixi@v0.9.2 + uses: prefix-dev/setup-pixi@v0.9.4 with: environments: lint @@ -35,8 +35,9 @@ jobs: uses: actions/checkout@v5 - name: 🟨 Set up Pixi - uses: prefix-dev/setup-pixi@v0.9.2 + uses: prefix-dev/setup-pixi@v0.9.4 with: + pixi-version: v0.62.2 # NOTE: pin until https://github.com/prefix-dev/pixi/issues/5424 is resolved environments: ${{ matrix.environment }} - name: 🧪 Execute pytest diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index c693324..03f9d95 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -23,8 +23,9 @@ jobs: uses: actions/checkout@v5 - name: 🟨 Set up Pixi - uses: prefix-dev/setup-pixi@v0.9.2 + uses: prefix-dev/setup-pixi@v0.9.4 with: + pixi-version: v0.62.2 # NOTE: pin until https://github.com/prefix-dev/pixi/issues/5424 is resolved environments: docs - name: 📝 Build docs diff --git a/.lefthook.yaml b/.lefthook.yaml index 2aeb41b..e37e0ae 100644 --- a/.lefthook.yaml +++ b/.lefthook.yaml @@ -31,18 +31,19 @@ pre-commit: - name: ruff glob: "*.{py,pyi}" - stage_fixed: true group: piped: true jobs: - name: ruff check run: pixi {run} ruff check --fix {staged_files} + stage_fixed: true - name: ruff format run: pixi {run} ruff format {staged_files} + stage_fixed: true - # - name: pyright + # - name: pyrefly # glob: "*.{py,pyi}" - # run: pixi {run} pyright {staged_files} + # run: pixi {run} pyrefly-check {staged_files} # - name: mypy # glob: "*.{py,pyi}" diff --git a/CHANGELOG.md b/CHANGELOG.md index 05b7371..d1775c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,27 @@ 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.3.0] - 2026-01-30 + +### Added + +- Implement Bolometry Observer functionality +- Add new notebook for creating emission profiles +- Add `*.pyi` files for cython sources +- Add `overload` decorator for better type hinting +- Add `ultraplot` package into `pixi` default environment + +### Changed + +- Add `pyrefly` package for type checking (still experimental) +- Refactor `load_equilibrium` and `load_magnetic_field_data` to use dataclasses +- Use `ultraplot` for plotting in examples and notebooks + +### Fixed + +- Fix `num_toroidal` parameter handling in `load_unstruct_grid_2d` function +- Fix typo in warning message in `load_profiles.py` + ## [0.2.1] - 2025-11-18 ### Added diff --git a/docs/notebooks/plasma/emission.ipynb b/docs/notebooks/plasma/emission.ipynb new file mode 100644 index 0000000..59037b4 --- /dev/null +++ b/docs/notebooks/plasma/emission.ipynb @@ -0,0 +1,730 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "326b42e5", + "metadata": {}, + "source": [ + "# Emission Models\n", + "\n", + "This notebook demonstrates how to construct emission models with `cherab.imas`-loaded plasmas.\n", + "\n", + "Prerequisites: [Pooch](https://www.fatiando.org/pooch/) must be installed to download the example data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a5b6f13", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import ultraplot as uplt\n", + "from raysect.core.math import Point3D, Vector3D\n", + "from raysect.core.workflow import MulticoreEngine\n", + "from raysect.optical import Ray, Spectrum, World\n", + "from rich import print as rprint\n", + "from rich.table import Table\n", + "\n", + "from cherab.core import Line, elements\n", + "from cherab.core.math import sample3d, sample3d_grid\n", + "from cherab.core.model import Bremsstrahlung, ExcitationLine, RecombinationLine\n", + "from cherab.imas.datasets import iter_jintrac\n", + "from cherab.imas.plasma import load_plasma\n", + "\n", + "# Set dark background for plots\n", + "uplt.rc.style = \"dark_background\"" + ] + }, + { + "cell_type": "markdown", + "id": "c1528131", + "metadata": {}, + "source": [ + "## Load full plasma from IMAS\n", + "\n", + "The details of loading a full plasma are covered in the [](./full_plasma.ipynb) notebook.\n", + "Here we simply load a plasma from an IMAS jintrac dataset." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4e29448", + "metadata": {}, + "outputs": [], + "source": [ + "world = World()\n", + "path = iter_jintrac()\n", + "plasma = load_plasma(path, \"r\", parent=world)" + ] + }, + { + "cell_type": "markdown", + "id": "05a5e0c8", + "metadata": {}, + "source": [ + "Show which species the plasma is composed of." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59aa884a", + "metadata": {}, + "outputs": [], + "source": [ + "species = sorted(\n", + " [s for s in plasma.composition],\n", + " key=lambda x: (x.element.symbol, x.charge),\n", + ")\n", + "\n", + "table = Table(title=\"Plasma Species\")\n", + "table.add_column(\"Element\", style=\"cyan\", justify=\"right\")\n", + "table.add_column(\"Charge\")\n", + "symbol = \"\"\n", + "for s in species:\n", + " _symbol = s.element.symbol\n", + " if _symbol != symbol:\n", + " table.add_row(s.element.symbol, f\"{s.charge:>2d}+\")\n", + " symbol = _symbol\n", + " else:\n", + " table.add_row(\"\", f\"{s.charge:>2d}+\")\n", + "rprint(table)" + ] + }, + { + "cell_type": "markdown", + "id": "cdd871cf", + "metadata": {}, + "source": [ + "## Construct emission models" + ] + }, + { + "cell_type": "markdown", + "id": "b0e94ab0", + "metadata": {}, + "source": [ + "### Retrieve atomic data\n", + "\n", + "[OpenADAS](https://open.adas.ac.uk) provides data for a variety of emission processes in plasmas.\n", + "Here we demonstrate how to include some of these processes in a Cherab simulation and generate synthetic spectra." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20e7af34", + "metadata": {}, + "outputs": [], + "source": [ + "from cherab.openadas import OpenADAS\n", + "from cherab.openadas.repository import populate\n", + "\n", + "# Download any missing ADAS data files from OpenADAS\n", + "populate()\n", + "\n", + "# Set up atomic data source\n", + "plasma.atomic_data = OpenADAS(permit_extrapolation=True)" + ] + }, + { + "cell_type": "markdown", + "id": "b85f1493", + "metadata": {}, + "source": [ + "Currently, neon and tangsten atomic data files (adf15) are to be downloaded manually as a workaround." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2689c2cb", + "metadata": {}, + "outputs": [], + "source": [ + "from cherab.openadas.install import install_files\n", + "\n", + "install_files(\n", + " {\n", + " \"adf15\": (\n", + " (elements.neon, 0, \"adf15/pec96#ne/pec96#ne_pjr#ne0.dat\"), # pjr: metastable resolved\n", + " (elements.neon, 1, \"adf15/pec96#ne/pec96#ne_pjr#ne1.dat\"),\n", + " # TODO: fix the parser to accept these files\n", + " # (elements.tungsten, 0, \"adf15/pec40#w/pec40#w_ic#w0.dat\"),\n", + " # (elements.tungsten, 1, \"adf15/pec40#w/pec40#w_ic#w1.dat\"),\n", + " )\n", + " },\n", + " download=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b9edcafd", + "metadata": {}, + "source": [ + "### Apply emission models to plasma species\n", + "\n", + "Here we create emission models for Bremsstrahlung, excitation lines, and recombination lines for the main plasma species (D, T, He, Ne, W).\n", + "Each line emission model is manually configured with `cherab`'s `Line` class to specify the desired transition." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db9a6923", + "metadata": {}, + "outputs": [], + "source": [ + "hydrogen_like_transitions = [\n", + " (2, 1),\n", + " (3, 1),\n", + " (3, 2),\n", + " (4, 1),\n", + " (4, 2),\n", + " (4, 3),\n", + " (5, 1),\n", + " (5, 2),\n", + " (5, 3),\n", + " (5, 4),\n", + " (6, 1),\n", + " (6, 2),\n", + " (6, 3),\n", + " (6, 4),\n", + " (6, 5),\n", + " (7, 1),\n", + " (7, 2),\n", + " (7, 3),\n", + " (7, 4),\n", + " (7, 5),\n", + " (7, 6),\n", + " (8, 1),\n", + " (8, 2),\n", + " (8, 3),\n", + " (8, 4),\n", + " (8, 5),\n", + " (8, 6),\n", + " (8, 7),\n", + " (9, 1),\n", + " (9, 2),\n", + " (9, 3),\n", + " (9, 4),\n", + " (9, 5),\n", + " (9, 6),\n", + " (9, 7),\n", + " (9, 8),\n", + " (10, 1),\n", + " (10, 2),\n", + " (10, 3),\n", + " (10, 4),\n", + " (10, 5),\n", + " (10, 6),\n", + " (10, 7),\n", + " (10, 8),\n", + " (10, 9),\n", + " (11, 1),\n", + " (11, 2),\n", + " (11, 3),\n", + " (11, 4),\n", + " (11, 5),\n", + " (11, 6),\n", + " (11, 7),\n", + " (11, 8),\n", + " (11, 9),\n", + " (11, 10),\n", + " (12, 1),\n", + " (12, 2),\n", + " (12, 3),\n", + " (12, 4),\n", + " (12, 5),\n", + " (12, 6),\n", + " (12, 7),\n", + " (12, 8),\n", + " (12, 9),\n", + " (12, 10),\n", + " (12, 11),\n", + "]\n", + "\n", + "lines = {\n", + " \"d0+\": [ # D 0+\n", + " *[Line(elements.deuterium, 0, t) for t in hydrogen_like_transitions],\n", + " ],\n", + " \"t0+\": [ # T 0+\n", + " *[Line(elements.tritium, 0, t) for t in hydrogen_like_transitions],\n", + " ],\n", + " \"he0+\": [ # He 0+\n", + " Line(elements.helium, 0, (\"1s1 4p1 1P1.0\", \"1s2 1S.0\")), # 52.22 nm\n", + " Line(elements.helium, 0, (\"1s1 3p1 1P1.0\", \"1s2 1S.0\")), # 53.70 nm\n", + " Line(elements.helium, 0, (\"1s1 2p1 1P1.0\", \"1s2 1S.0\")), # 58.44 nm\n", + " Line(elements.helium, 0, (\"1s1 4p1 3P4.0\", \"1s1 2s1 3S1.0\")), # 318.87 nm\n", + " Line(elements.helium, 0, (\"1s1 3p1 3P4.0\", \"1s1 2s1 3S1.0\")), # 388.97 nm\n", + " Line(elements.helium, 0, (\"1s1 4p1 1P1.0\", \"1s1 2s1 1S.0\")), # 396.57 nm\n", + " Line(elements.helium, 0, (\"1s1 4d1 3D7.0\", \"1s1 2p1 3P4.0\")), # 447.29 nm\n", + " Line(elements.helium, 0, (\"1s1 4s1 3S1.0\", \"1s1 2p1 3P4.0\")), # 471.48 nm\n", + " Line(elements.helium, 0, (\"1s1 4d1 1D2.0\", \"1s1 2p1 1P1.0\")), # 492.32 nm\n", + " Line(elements.helium, 0, (\"1s1 3p1 1P1.0\", \"1s1 2s1 1S.0\")), # 501.71 nm\n", + " Line(elements.helium, 0, (\"1s1 4s1 1S.0\", \"1s1 2p1 1P1.0\")), # 504.90 nm\n", + " Line(elements.helium, 0, (\"1s1 3d1 3D7.0\", \"1s1 2p1 3P4.0\")), # 587.75 nm\n", + " Line(elements.helium, 0, (\"1s1 3d1 1D2.0\", \"1s1 2p1 1P1.0\")), # 668.00 nm\n", + " Line(elements.helium, 0, (\"1s1 3s1 3S1.0\", \"1s1 2p1 3P4.0\")), # 706.76 nm\n", + " Line(elements.helium, 0, (\"1s1 3s1 1S.0\", \"1s1 2p1 1P1.0\")), # 728.33 nm\n", + " ],\n", + " \"he1+\": [ # He 1+\n", + " Line(elements.helium, 1, (2, 1)),\n", + " Line(elements.helium, 1, (3, 1)),\n", + " Line(elements.helium, 1, (3, 2)),\n", + " Line(elements.helium, 1, (4, 1)),\n", + " Line(elements.helium, 1, (4, 2)),\n", + " Line(elements.helium, 1, (4, 3)),\n", + " Line(elements.helium, 1, (5, 1)),\n", + " Line(elements.helium, 1, (5, 2)),\n", + " Line(elements.helium, 1, (5, 3)),\n", + " ],\n", + " \"ne0+\": [ # Ne 0+\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 3D7.0\", \"2s2 2p6 1S0.0\")), # 61.69 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 1P1.0\", \"2s2 2p6 1S0.0\")), # 61.87 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 3P4.0\", \"2s2 2p6 1S0.0\")), # 61.89 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3s1 1P1.0\", \"2s2 2p6 1S0.0\")), # 73.59 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1S0.0\", \"2s2 2p5 3s1 3P4.0\")), # 600.59 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1D2.0\", \"2s2 2p5 3s1 3P4.0\")), # 602.73 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1P1.0\", \"2s2 2p5 3s1 3P4.0\")), # 605.88 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 3D7.0\", \"2s2 2p5 3s1 3P4.0\")), # 643.56 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1S0.0\", \"2s2 2p5 3s1 1P1.0\")), # 665.39 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1D2.0\", \"2s2 2p5 3s1 1P1.0\")), # 668.01 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 3P4.0\", \"2s2 2p5 3s1 1P1.0\")), # 668.32 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 1P1.0\", \"2s2 2p5 3s1 1P1.0\")), # 671.89 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 3S1.0\", \"2s2 2p5 3s1 3P4.0\")), # 714.77 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3p1 3D7.0\", \"2s2 2p5 3s1 1P1.0\")), # 718.55 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 3D7.0\", \"2s2 2p5 3p1 3S1.0\")), # 723.07 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 1P1.0\", \"2s2 2p5 3p1 3S1.0\")), # 747.45 nm\n", + " Line(elements.neon, 0, (\"2s2 2p5 3d1 3P4.0\", \"2s2 2p5 3p1 3S1.0\")), # 751.26 nm\n", + " ],\n", + " \"ne1+\": [ # Ne 1+\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p5 2P2.5\")), # 30.52 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2F6.5\", \"2s2 2p5 2P2.5\")), # 32.63 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p5 2P2.5\")), # 32.68 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2S0.5\", \"2s2 2p5 2P2.5\")), # 32.71 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2P2.5\", \"2s2 2p5 2P2.5\")), # 32.75 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4P5.5\", \"2s2 2p5 2P2.5\")), # 35.61 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2F6.5\", \"2s2 2p5 2P2.5\")), # 35.64 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4F13.5\", \"2s2 2p5 2P2.5\")), # 35.69 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p5 2P2.5\")), # 35.70 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4D9.5\", \"2s2 2p5 2P2.5\")), # 35.85 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3s1 2S0.5\", \"2s2 2p5 2P2.5\")), # 36.18 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3s1 2D4.5\", \"2s2 2p5 2P2.5\")), # 40.63 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3s1 2P2.5\", \"2s2 2p5 2P2.5\")), # 44.64 nm\n", + " Line(elements.neon, 1, (\"2s1 2p6 2S0.5\", \"2s2 2p5 2P2.5\")), # 46.13 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 2P2.5\")), # 123.29 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 2D4.5\")), # 169.47 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 4P5.5\")), # 175.70 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2D4.5\", \"2s2 2p4 3s1 2P2.5\")), # 188.52 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 2P2.5\")), # 192.07 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2P2.5\", \"2s2 2p4 3p1 4P5.5\")), # 283.63 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 4P5.5\")), # 287.62 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4P5.5\", \"2s2 2p4 3p1 4P5.5\")), # 288.11 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2F6.5\", \"2s2 2p4 3p1 4P5.5\")), # 289.98 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4F13.5\", \"2s2 2p4 3p1 4P5.5\")), # 293.18 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p4 3p1 4P5.5\")), # 293.97 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 4S1.5\", \"2s2 2p4 3s1 4P5.5\")), # 298.38 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4D9.5\", \"2s2 2p4 3p1 4P5.5\")), # 304.01 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2P2.5\", \"2s2 2p4 3p1 4D9.5\")), # 310.65 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2D4.5\", \"2s2 2p4 3s1 4P5.5\")), # 314.70 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4P5.5\", \"2s2 2p4 3p1 4D9.5\")), # 316.02 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2F6.5\", \"2s2 2p4 3p1 4D9.5\")), # 318.07 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4F13.5\", \"2s2 2p4 3p1 4D9.5\")), # 322.14 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p4 3p1 4D9.5\")), # 323.10 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2D4.5\", \"2s2 2p4 3s1 2D4.5\")), # 323.19 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2P2.5\", \"2s2 2p4 3p1 2D4.5\")), # 329.21 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3s1 2S0.5\", \"2s2 2p4 3p1 4P5.5\")), # 329.57 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 4D9.5\", \"2s2 2p4 3s1 4P5.5\")), # 333.77 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 2D4.5\")), # 333.78 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2P2.5\", \"2s2 2p4 3s1 2P2.5\")), # 334.27 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4D9.5\", \"2s2 2p4 3p1 4D9.5\")), # 335.26 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4P5.5\", \"2s2 2p4 3p1 2D4.5\")), # 335.26 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2F6.5\", \"2s2 2p4 3p1 2D4.5\")), # 337.56 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4F13.5\", \"2s2 2p4 3p1 2D4.5\")), # 342.15 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 2D4.5\", \"2s2 2p4 3p1 2D4.5\")), # 343.23 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2S0.5\", \"2s2 2p4 3s1 2P2.5\")), # 350.79 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3d1 4D9.5\", \"2s2 2p4 3p1 2D4.5\")), # 356.99 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2F6.5\", \"2s2 2p4 3s1 2D4.5\")), # 357.21 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3s1 2S0.5\", \"2s2 2p4 3p1 4D9.5\")), # 366.61 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 2D4.5\", \"2s2 2p4 3s1 2P2.5\")), # 371.41 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 4P5.5\", \"2s2 2p4 3s1 4P5.5\")), # 371.82 nm\n", + " Line(elements.neon, 1, (\"2s2 2p4 3p1 4D9.5\", \"2s2 2p4 3s1 2P2.5\")), # 398.26 nm\n", + " ],\n", + "}" + ] + }, + { + "cell_type": "markdown", + "id": "970329c3", + "metadata": {}, + "source": [ + "Apply the Bremsstrahlung and line emission models to the plasma species." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28b5333b", + "metadata": {}, + "outputs": [], + "source": [ + "plasma.models = [\n", + " Bremsstrahlung(),\n", + " *[ExcitationLine(line) for line_list in lines.values() for line in line_list],\n", + " *[\n", + " RecombinationLine(line)\n", + " for line in lines[\"d0+\"] + lines[\"t0+\"] + lines[\"he0+\"] + lines[\"he1+\"] + lines[\"ne0+\"]\n", + " ],\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "ce058ee2", + "metadata": {}, + "source": [ + "### Visualize emission profiles\n", + "\n", + "Let us visualize the total emission from all processes on a 2D grid in the R-Z plane.\n", + "Due to the heavy computational load, we limit the grid resolution to 50 mm and wavelength range to\n", + "0.0125 nm - 1240 nm with 100 bins evenly spaced in log-scale." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "082d62b1", + "metadata": {}, + "outputs": [], + "source": [ + "# Grid specification\n", + "R_MIN, R_MAX = 4.0, 8.5\n", + "Z_MIN, Z_MAX = -4.5, 4.6\n", + "RES = 50.0e-3 # 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", + "extent = (R_MIN, R_MAX, Z_MIN, Z_MAX)\n", + "\n", + "# Wavelength sampling\n", + "wavelengths = np.logspace(\n", + " np.log10(0.012),\n", + " np.log10(1240),\n", + " 100 + 1,\n", + ")\n", + "# Define wavelength bands as list of (min, max) tuples\n", + "# NOTE: Because cherab can only handle linearly spaced wavelength bins, we define wavelength bands\n", + "# as adjacent pairs of the logarithmically spaced wavelengths so that we calculate spectrum at one\n", + "# point per band.\n", + "wavelength_bands = [\n", + " (wavelengths[i].item(), wavelengths[i + 1].item()) for i in range(len(wavelengths) - 1)\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "72502443", + "metadata": {}, + "source": [ + "Sample emission at each grid point integrating over the defined wavelength bands." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b5df6f6", + "metadata": {}, + "outputs": [], + "source": [ + "class Emission2DSample:\n", + " def __init__(self, plasma, r_pts, z_pts, wavelength_bands) -> None:\n", + " self.plasma = plasma\n", + " self.r_pts = r_pts\n", + " self.z_pts = z_pts\n", + " self.indices = list(np.ndindex(len(r_pts), len(z_pts)))\n", + " self.wavelength_bands = wavelength_bands\n", + " self.engine = MulticoreEngine()\n", + " self.emission_2d = np.zeros((len(r_pts), len(z_pts)))\n", + "\n", + " def run(self) -> None:\n", + " self.engine.run(\n", + " self.indices,\n", + " self._render,\n", + " self._update,\n", + " )\n", + "\n", + " def _render(self, idx) -> float:\n", + " i_r, i_z = idx\n", + " x = self.r_pts[i_r]\n", + " z = self.z_pts[i_z]\n", + "\n", + " n_e = self.plasma.electron_distribution.density(x, 0, z)\n", + " if n_e <= 0:\n", + " return idx, 0.0\n", + "\n", + " total = 0.0\n", + " for min_wl, max_wl in self.wavelength_bands:\n", + " spectrum = Spectrum(\n", + " min_wavelength=min_wl,\n", + " max_wavelength=max_wl,\n", + " bins=1,\n", + " )\n", + " total += self._emission(x, 0, z, spectrum)\n", + " return idx, total\n", + "\n", + " def _update(self, packed_result) -> None:\n", + " (i_r, i_z), value = packed_result\n", + " self.emission_2d[i_r, i_z] = value\n", + "\n", + " def _emission(self, x, y, z, spectrum: Spectrum) -> float:\n", + " \"\"\"Calculate total emission at a point for a given spectrum.\"\"\"\n", + " v = 0.0\n", + " for model in self.plasma.models:\n", + " v += model.emission(Point3D(x, y, z), Vector3D(0, 1, 0), spectrum).total()\n", + " return v\n", + "\n", + "\n", + "# Get R, Z sampling points\n", + "r_pts, _, z_pts, _ = sample3d(\n", + " lambda x, y, z: 0,\n", + " (R_MIN, R_MAX, n_r),\n", + " (0, 0, 1),\n", + " (Z_MIN, Z_MAX, n_z),\n", + ")\n", + "\n", + "# Create emission sampler and run\n", + "emission_sampler = Emission2DSample(plasma, r_pts, z_pts, wavelength_bands)\n", + "emission_sampler.run()" + ] + }, + { + "cell_type": "markdown", + "id": "fa2562dd", + "metadata": {}, + "source": [ + "Plot the 2-D emission map" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7bfccd5", + "metadata": {}, + "outputs": [], + "source": [ + "emission_2d = emission_sampler.emission_2d\n", + "emission_2d[emission_2d <= 0] = np.nan # for log plotting\n", + "\n", + "fig, ax = uplt.subplots()\n", + "ax.imshow(\n", + " emission_2d.T,\n", + " discrete=False,\n", + " extent=extent,\n", + " origin=\"lower\",\n", + " norm=\"log\",\n", + " cmap=\"rocket\",\n", + " colorbar=\"r\",\n", + " colorbar_kw=dict(\n", + " tickdir=\"out\",\n", + " formatter=\"log\",\n", + " label=\"[W/m³/sr]\",\n", + " ),\n", + ")\n", + "ax.format(\n", + " aspect=1,\n", + " title=f\"Emissivity integrated over {wavelengths[0]:.3f} – {wavelengths[-1]:.0f} nm\",\n", + " xlabel=\"$R$ [m]\",\n", + " ylabel=\"$Z$ [m]\",\n", + " xlocator=1,\n", + " ylocator=1,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "55a88f82", + "metadata": {}, + "source": [ + "## Measure line-of-sight Spectra\n", + "\n", + "Here we measure line-of-sight spectra through the above modeled emission." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "41d9851f", + "metadata": {}, + "outputs": [], + "source": [ + "class MeasureLoS:\n", + " def __init__(self, world) -> None:\n", + " self.engine = MulticoreEngine()\n", + " self.world = world\n", + " self.results: list[tuple[float, float]] = []\n", + "\n", + " def run(self, rays) -> np.ndarray:\n", + " self.engine.run(rays, self.render, self.update)\n", + " return np.asarray(\n", + " sorted(self.results, key=lambda x: x[0]),\n", + " )\n", + "\n", + " def render(self, ray: Ray) -> tuple[float, float]:\n", + " s = ray.trace(self.world)\n", + " return s.wavelengths.item(), s.samples.item()\n", + "\n", + " def update(self, packed_result) -> None:\n", + " x, y = packed_result\n", + " self.results.append((x, y))\n", + "\n", + "\n", + "# Generate rays with defined wavelength bands\n", + "origin = Point3D(8.4, 0, 1)\n", + "endpoint = Point3D(4.6, 0, -4)\n", + "direction = origin.vector_to(endpoint).normalise()\n", + "rays = [\n", + " Ray(\n", + " origin=origin,\n", + " direction=direction,\n", + " bins=1,\n", + " max_wavelength=band[1],\n", + " min_wavelength=band[0],\n", + " )\n", + " for band in wavelength_bands\n", + "]\n", + "\n", + "# Measure line of sight spectrum\n", + "measure = MeasureLoS(world)\n", + "spectra = measure.run(rays)" + ] + }, + { + "cell_type": "markdown", + "id": "e7694533", + "metadata": {}, + "source": [ + "### Visualize measured spectra" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e4f8c46f", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = uplt.subplots(\n", + " refaspect=2,\n", + " refwidth=6,\n", + ")\n", + "\n", + "ax.loglog(\n", + " spectra[:, 0],\n", + " spectra[:, 1],\n", + ")\n", + "ax.format(\n", + " xlabel=\"Wavelength [nm]\",\n", + " ylabel=\"Spectral Radiance [W/m²/sr/nm]\",\n", + " yformatter=\"log\",\n", + " grid=True,\n", + " gridalpha=0.3,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "4c0d714c", + "metadata": {}, + "source": [ + "Let us visualize the ray trajectory projected onto the R-Z plane along with the measured spectra." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2202b8cb", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axs = uplt.subplots([1, 2, 2, 2], share=False)\n", + "\n", + "# Plot 2-D emissivity map\n", + "image = axs[0].imshow(\n", + " emission_2d.T,\n", + " extent=extent,\n", + " origin=\"lower\",\n", + " norm=\"log\",\n", + " cmap=\"rocket\",\n", + ")\n", + "axs[0].colorbar(\n", + " image,\n", + " formatter=\"log\",\n", + " tickdir=\"out\",\n", + " loc=\"lr\",\n", + " orientation=\"vertical\",\n", + " ticklabelsize=\"small\",\n", + " length=5,\n", + " frameon=False,\n", + ")\n", + "\n", + "# Plot line of sight\n", + "axs[0].plot(\n", + " [origin.x, endpoint.x],\n", + " [origin.z, endpoint.z],\n", + " color=\"C0\",\n", + " label=\"Line of sight\",\n", + " legend=\"ur\",\n", + ")\n", + "\n", + "axs[0].format(\n", + " aspect=1,\n", + " xlim=(extent[0], extent[1]),\n", + " ylim=(extent[2], extent[3]),\n", + " xlabel=\"$R$ [m]\",\n", + " ylabel=\"$Z$ [m]\",\n", + " ylocator=1,\n", + " xlocator=1,\n", + " title=\"Emissivity [W/m³/sr]\",\n", + " gridalpha=0.1,\n", + ")\n", + "\n", + "# Plot 1-D measured spectrum\n", + "axs[1].loglog(spectra[:, 0], spectra[:, 1])\n", + "axs[1].format(\n", + " xlim=(wavelengths.min(), wavelengths.max()),\n", + " title=\"Ray-traced emission spectrum\",\n", + " xlabel=\"Wavelength [nm]\",\n", + " ylabel=\"Spectral Radiance [W/m²/sr/nm]\",\n", + " yformatter=\"log\",\n", + " grid=True,\n", + " gridalpha=0.3,\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.13.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/notebooks/plasma/full_plasma.ipynb b/docs/notebooks/plasma/full_plasma.ipynb index 1fa4cc1..d26d6cb 100644 --- a/docs/notebooks/plasma/full_plasma.ipynb +++ b/docs/notebooks/plasma/full_plasma.ipynb @@ -22,15 +22,19 @@ "outputs": [], "source": [ "import numpy as np\n", - "from matplotlib import pyplot as plt\n", + "import ultraplot as uplt\n", "from matplotlib.colors import SymLogNorm\n", + "from rich import print as rprint\n", + "from rich.table import Table\n", "\n", "from cherab.core.math import sample3d, sample3d_grid\n", "from cherab.imas.datasets import iter_jintrac\n", "from cherab.imas.plasma import load_plasma\n", "\n", "# Set dark background for plots\n", - "plt.style.use(\"dark_background\")" + "uplt.rc.style = \"dark_background\"\n", + "\n", + "uplt.rc[\"title.borderwidth\"] = 0" ] }, { @@ -49,20 +53,13 @@ "outputs": [], "source": [ "def plot_quantity(\n", - " quantity,\n", - " extent,\n", - " title=\"\",\n", - " clabel=\"\",\n", - " logscale=False,\n", - " symmetric=False,\n", - "):\n", + " ax: uplt.axes.Axes,\n", + " quantity: np.ndarray,\n", + " extent: tuple[float, float, float, float],\n", + " logscale: bool = False,\n", + " symmetric: bool = False,\n", + ") -> uplt.axes.Axes:\n", " \"\"\"Make a 2D plot of quantity, optionally on a log scale.\"\"\"\n", - "\n", - " fig, ax = plt.subplots(\n", - " dpi=150,\n", - " figsize=(4.0, 6.0),\n", - " layout=\"constrained\",\n", - " )\n", " if logscale:\n", " # Plot lowest values (mainly 0's) on linear map, as log(0) = -inf.\n", " linthresh = np.percentile(np.unique(quantity), 1)\n", @@ -79,6 +76,7 @@ " image = ax.imshow(\n", " quantity.T,\n", " extent=extent,\n", + " discrete=False,\n", " origin=\"lower\",\n", " vmin=-vmax,\n", " vmax=vmax,\n", @@ -88,19 +86,27 @@ " image = ax.imshow(\n", " quantity.T,\n", " extent=extent,\n", + " discrete=False,\n", " origin=\"lower\",\n", " norm=norm,\n", " cmap=\"gnuplot\",\n", " )\n", - " cbar = plt.colorbar(image, aspect=50)\n", - " cbar.set_label(clabel)\n", - " ax.set_xlim(extent[0], extent[1])\n", - " ax.set_ylim(extent[2], extent[3])\n", - " ax.set_xlabel(\"$R$ [m]\")\n", - " ax.set_ylabel(\"$Z$ [m]\")\n", - " ax.set_title(title)\n", "\n", - " return fig" + " ax.colorbar(\n", + " image,\n", + " formatter=\"log\",\n", + " tickminor=True,\n", + " )\n", + "\n", + " ax.format(\n", + " aspect=1,\n", + " xlabel=\"$R$ [m]\",\n", + " ylabel=\"$Z$ [m]\",\n", + " xlocator=1,\n", + " ylocator=1,\n", + " )\n", + "\n", + " return ax" ] }, { @@ -127,7 +133,8 @@ "metadata": {}, "source": [ "## Load the plasma object\n", - "The instance of the `~cherab.core.plasma.node.Plasma` class is created by loading the `core_profiles` and `edge_profiles` IDSs from the IMAS database.\n", + "\n", + "The instance of the [`Plasma`](https://www.cherab.info/plasmas/core_plasma_classes.html#cherab.core.Plasma) class is created by loading the `core_profiles` and `edge_profiles` IDSs from the IMAS database.\n", "\n", "The equilibrium information is automatically loaded from the `equilibrium` IDS if not already provided." ] @@ -157,8 +164,29 @@ "metadata": {}, "outputs": [], "source": [ + "# Organize plasma species by symbol and charge\n", + "species = {}\n", "for s in plasma.composition:\n", - " print(s)" + " symbol = s.element.symbol\n", + " if symbol not in species:\n", + " species[symbol] = []\n", + " species[symbol].append(s)\n", + "\n", + "# Sort species of each element by charge\n", + "for symbol in species:\n", + " species[symbol].sort(key=lambda x: x.charge)\n", + "\n", + "# Print table of plasma species\n", + "table = Table(title=\"Plasma Species\")\n", + "table.add_column(\"Element\", style=\"cyan\", justify=\"right\")\n", + "table.add_column(\"Charge\")\n", + "for symbol in species:\n", + " for i, s in enumerate(species[symbol]):\n", + " if i == 0:\n", + " table.add_row(symbol, f\"{s.charge:>2d}+\")\n", + " else:\n", + " table.add_row(\"\", f\"{s.charge:>2d}+\")\n", + "rprint(table)" ] }, { @@ -221,12 +249,16 @@ "n_e = n_e.squeeze()\n", "\n", "# Plot electron density\n", - "fig = plot_quantity(\n", + "fig, ax = uplt.subplots()\n", + "ax = plot_quantity(\n", + " ax,\n", " n_e,\n", " extent,\n", - " title=\"$n_\\\\mathrm{e}$ [/m³]\",\n", " logscale=True,\n", ")\n", + "ax.format(\n", + " title=\"$n_\\\\mathrm{e}$ [/m³]\",\n", + ")\n", "\n", "# Sample electron temperature\n", "te_plasma = sample3d_grid(\n", @@ -237,11 +269,15 @@ ").squeeze()\n", "\n", "# Plot electron temperature\n", - "fig = plot_quantity(\n", + "fig, ax = uplt.subplots()\n", + "ax = plot_quantity(\n", + " ax,\n", " te_plasma,\n", " extent,\n", - " title=\"$T_\\\\mathrm{e}$ [eV]\",\n", " logscale=True,\n", + ")\n", + "ax.format(\n", + " title=\"$T_\\\\mathrm{e}$ [eV]\",\n", ")" ] }, @@ -250,7 +286,7 @@ "id": "403fe1fc", "metadata": {}, "source": [ - "### Hydrogenic species (H, D, T) density and temperature" + "### Hydrogenic species (D, T) density and temperature" ] }, { @@ -260,94 +296,67 @@ "metadata": {}, "outputs": [], "source": [ - "for species in plasma.composition:\n", - " # Construct element label\n", + "def species_label(species) -> str:\n", + " \"\"\"Construct element symbol with charge label\"\"\"\n", " if species.charge == 0:\n", " label = f\"{species.element.symbol}$^{{0}}$\"\n", " elif species.charge == 1:\n", " label = f\"{species.element.symbol}$^{{+}}$\"\n", " else:\n", " label = f\"{species.element.symbol}$^{{{species.charge}+}}$\"\n", + " return label\n", + "\n", "\n", - " # Extract only hydrogenic species\n", - " if species.element.atomic_number == 1:\n", + "# === Density Plot ===\n", + "fig, axs = uplt.subplots(\n", + " nrows=2,\n", + " ncols=2,\n", + ")\n", + "axs.format(suptitle=\"Hydrogenic Species Density [/m³]\")\n", + "for i_row, s_list in enumerate([species[\"D\"], species[\"T\"]]):\n", + " for i_col, s in enumerate(s_list):\n", " # Sample density\n", " density = sample3d_grid(\n", - " species.distribution.density,\n", + " s.distribution.density,\n", " r_pts,\n", " [0],\n", " z_pts,\n", " ).squeeze()\n", "\n", " # Plot density\n", - " fig = plot_quantity(\n", + " ax = plot_quantity(\n", + " axs[i_row, i_col],\n", " density,\n", " extent,\n", - " title=\"Density of \" + label,\n", - " clabel=\"[1/m³]\",\n", " logscale=True,\n", " )\n", + " ax.format(urtitle=species_label(s))\n", + "\n", "\n", + "# === Temperature Plot ===\n", + "fig, axs = uplt.subplots(\n", + " nrows=2,\n", + " ncols=2,\n", + ")\n", + "axs.format(suptitle=\"Hydrogenic Species Temperature [eV]\")\n", + "for i_row, s_list in enumerate([species[\"D\"], species[\"T\"]]):\n", + " for i_col, s in enumerate(s_list):\n", " # Sample temperature\n", " temperature = sample3d_grid(\n", - " species.distribution.effective_temperature,\n", + " s.distribution.effective_temperature,\n", " r_pts,\n", " [0],\n", " z_pts,\n", " ).squeeze()\n", "\n", " # Plot deutrium temperature\n", - " fig = plot_quantity(\n", + " ax = plot_quantity(\n", + " axs[i_row, i_col],\n", " temperature,\n", " extent,\n", - " title=\"Temperature of \" + label,\n", - " clabel=\"[eV]\",\n", - " logscale=True,\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "df9fb179", - "metadata": {}, - "source": [ - "### Helium atom/ion density profiles" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4fb77b30", - "metadata": {}, - "outputs": [], - "source": [ - "for species in plasma.composition:\n", - " # Construct element label\n", - " if species.charge == 0:\n", - " label = f\"{species.element.symbol}$^{{0}}$\"\n", - " elif species.charge == 1:\n", - " label = f\"{species.element.symbol}$^{{+}}$\"\n", - " else:\n", - " label = f\"{species.element.symbol}$^{{{species.charge}+}}$\"\n", - "\n", - " # Extract only Helium species\n", - " if species.element.atomic_number == 2:\n", - " # Sample density\n", - " density = sample3d_grid(\n", - " species.distribution.density,\n", - " r_pts,\n", - " [0],\n", - " z_pts,\n", - " ).squeeze()\n", - "\n", - " # Plot density\n", - " fig = plot_quantity(\n", - " density,\n", - " extent,\n", - " title=\"Density of \" + label,\n", - " clabel=\"[1/m³]\",\n", " logscale=True,\n", - " )" + " )\n", + " ax.format(urtitle=species_label(s))" ] }, { @@ -365,35 +374,33 @@ "metadata": {}, "outputs": [], "source": [ - "for species in plasma.composition:\n", - " # Construct element label\n", - " if species.charge == 0:\n", - " label = f\"{species.element.symbol}$^{{0}}$\"\n", - " elif species.charge == 1:\n", - " label = f\"{species.element.symbol}$^{{+}}$\"\n", - " else:\n", - " label = f\"{species.element.symbol}$^{{{species.charge}+}}$\"\n", + "for s_list in [species[\"He\"], species[\"Ne\"], species[\"W\"]]:\n", + " # Each element gets its own figure\n", + " fig, axs = uplt.subplots(\n", + " ncols=len(s_list),\n", + " )\n", + " axs.format(suptitle=f\"{s_list[0].element.name.capitalize()} Density [/m³]\")\n", "\n", - " # Extract only non-hydrogenic and non-Helium species\n", - " if species.element.atomic_number != 1 and species.element.atomic_number != 2:\n", + " for i_col, s in enumerate(s_list):\n", " # Sample density\n", " density = sample3d_grid(\n", - " species.distribution.density,\n", + " s.distribution.density,\n", " r_pts,\n", " [0],\n", " z_pts,\n", " ).squeeze()\n", "\n", - " density[density < 0] = 0.0\n", + " # Make non-physical negative densities zero\n", + " density[density <= 0] = 0.0\n", "\n", " # Plot density\n", - " fig = plot_quantity(\n", + " ax = plot_quantity(\n", + " axs[i_col],\n", " density,\n", " extent,\n", - " title=\"Density of \" + label,\n", - " clabel=\"[1/m³]\",\n", " logscale=True,\n", - " )" + " )\n", + " ax.format(urtitle=species_label(s))" ] } ], @@ -413,7 +420,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.9" + "version": "3.13.11" } }, "nbformat": 4, diff --git a/docs/source/_static/custom.css b/docs/source/_static/custom.css new file mode 100644 index 0000000..d097255 --- /dev/null +++ b/docs/source/_static/custom.css @@ -0,0 +1,6 @@ +/* Optional: Ensure image respects container boundaries */ +.output_area.docutils.container img { + height: auto !important; + display: block !important; + object-fit: contain; +} \ No newline at end of file diff --git a/docs/source/conf.py b/docs/source/conf.py index d560367..9106803 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -151,6 +151,8 @@ }, ], } +html_static_path = ["_static"] +html_css_files = ["custom.css"] # Shorten Table Of Contents in API documentation object_description_options = [ @@ -168,6 +170,7 @@ "imas-python": ("https://imas-python.readthedocs.io/en/stable/", None), "rich": ("https://rich.readthedocs.io/en/stable/", None), "pooch": ("https://www.fatiando.org/pooch/latest/", None), + "ultraplot": ("https://ultraplot.readthedocs.io/en/stable/", None), } intersphinx_timeout = 10 diff --git a/pixi.toml b/pixi.toml index 4c0b016..0a59c7f 100644 --- a/pixi.toml +++ b/pixi.toml @@ -21,9 +21,8 @@ python = ["3.10.*", "3.11.*", "3.12.*", "3.13.*"] [package.host-dependencies] uv = "*" python = "*" -hatchling = "*" +hatch-cython = ">=0.6.0" hatch-vcs = "*" -hatch-cython = ">=0.6.0rc0" cython = ">=3.1,<4" llvm-openmp = "*" cherab = "1.5.*" @@ -41,6 +40,9 @@ ipython = "*" pooch = "*" rich = "*" +# Publication-quality plot +ultraplot = ">=1.72.0" + [tasks] ipython = { cmd = "ipython", description = "🐍 Start an IPython shell" } @@ -105,7 +107,7 @@ lefthook = "*" ruff = "*" typos = "*" mypy = "*" -basedpyright = "*" +pyrefly = "*" actionlint = "*" shellcheck = "*" validate-pyproject = "*" @@ -118,7 +120,7 @@ lefthook = { cmd = "lefthook", description = "🔗 Run lefthook" } hooks = { cmd = "lefthook install", description = "🔗 Install pre-commit hooks" } pre-commit = { cmd = "lefthook run pre-commit", description = "🔗 Run pre-commit checks" } mypy = { cmd = "mypy", description = "Type check with mypy" } -pyright = { cmd = "basedpyright", description = "Type check with basedpyright" } +pyrefly-check = { cmd = "pyrefly check", description = "Type check with pyrefly" } ruff-check = { cmd = "ruff check --fix", description = "Lint with ruff" } ruff-format = { cmd = "ruff format", description = "Format with ruff" } dprint = { cmd = "dprint fmt", description = "Format with dprint" } diff --git a/pyproject.toml b/pyproject.toml index ccde31c..9d798ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["hatchling", "hatch-vcs", "hatch-cython>=0.6.0rc0", "cherab==1.5.*"] +requires = ["hatch-vcs", "hatch-cython>=0.6.0", "cherab==1.5.*"] build-backend = "hatchling.build" [project] @@ -164,23 +164,12 @@ max-line-length = 100 # === Type checking settings === # ------------------------------ [tool.mypy] -files = ["src", "tests"] +files = ["src"] warn_unused_configs = true strict = true enable_error_code = ["ignore-without-code", "truthy-bool"] disable_error_code = ["no-any-return"] plugins = ["numpy.typing.mypy_plugin"] -[tool.basedpyright] -include = ["src", "tests"] -pythonPlatform = "All" -typeCheckingMode = "all" -reportMissingTypeStubs = true - -# no raysect/cherab type stubs; pytest fixtures -reportUnknownMemberType = false -reportUnknownVariableType = false -reportUnknownParameterType = false - -# ruff handles this -reportUnusedParameter = false +[tool.pyrefly] +project-includes = ["src"] diff --git a/src/cherab/imas/__init__.py b/src/cherab/imas/__init__.py index 37746bf..5368303 100644 --- a/src/cherab/imas/__init__.py +++ b/src/cherab/imas/__init__.py @@ -17,8 +17,6 @@ # under the Licence. """The Cherab subpackage for the IMAS.""" -from __future__ import annotations - from importlib.metadata import version as _version __version__ = _version("cherab-imas") diff --git a/src/cherab/imas/ggd/base_mesh.py b/src/cherab/imas/ggd/base_mesh.py index 8ba9925..3869110 100755 --- a/src/cherab/imas/ggd/base_mesh.py +++ b/src/cherab/imas/ggd/base_mesh.py @@ -19,6 +19,7 @@ from __future__ import annotations +from abc import abstractmethod from typing import Literal import matplotlib.axes @@ -62,15 +63,16 @@ def __init__( self._name: str = str(name) self._coordinate_system: str = str(coordinate_system) - self._interpolator: object | None = None - self._cell_centre: NDArray[np.float64] | None = None - self._cell_area: NDArray[np.float64] | None = None - self._cell_volume: NDArray[np.float64] | None = None - self._mesh_extent: dict[str, float] | None = None - self._num_cell: int = 0 + self._interpolator: object + self._cell_centre: NDArray[np.float64] + self._cell_area: NDArray[np.float64] + self._cell_volume: NDArray[np.float64] + self._mesh_extent: dict[str, float] + self._num_cell: int self._initial_setup() + @abstractmethod def _initial_setup(self) -> None: raise NotImplementedError("To be defined in subclass.") @@ -99,28 +101,29 @@ def coordinate_system(self) -> str: return self._coordinate_system @property - def cell_centre(self) -> NDArray[np.float64] | None: + def cell_centre(self) -> NDArray[np.float64]: """Coordinate of cell centres as ``(num_cell, dimension)`` array.""" return self._cell_centre @property - def cell_area(self) -> NDArray[np.float64] | None: + def cell_area(self) -> NDArray[np.float64]: """Cell areas as ``(num_cell,)`` array.""" return self._cell_area @property - def cell_volume(self) -> NDArray[np.float64] | None: + def cell_volume(self) -> NDArray[np.float64]: """Cell volume as ``(num_cell,)`` array.""" return self._cell_volume @property - def mesh_extent(self) -> dict[str, float] | None: + def mesh_extent(self) -> dict[str, float]: """Extent of the mesh. A dictionary with xmin, xmax, ymin and ymax, ... keys. """ return self._mesh_extent + @abstractmethod def subset(self, indices: ArrayLike, name: str | None = None) -> GGDGrid: """Create a subset grid from this instance. @@ -138,8 +141,9 @@ def subset(self, indices: ArrayLike, name: str | None = None) -> GGDGrid: """ raise NotImplementedError("To be defined in subclass.") + @abstractmethod def interpolator( - self, grid_data: ArrayLike, fill_value: float = 0.0 + self, grid_data: NDArray[np.float64], fill_value: float = 0.0 ) -> Function2D | Function3D: """Return an Function interpolator instance for the data defined on this grid. @@ -160,8 +164,9 @@ def interpolator( """ raise NotImplementedError("To be defined in subclass.") + @abstractmethod def vector_interpolator( - self, grid_vectors: ArrayLike, fill_vector: Vector3D = ZEROVECTOR + self, grid_vectors: NDArray[np.float64], fill_vector: Vector3D = ZEROVECTOR ) -> VectorFunction2D | VectorFunction3D: """Return a VectorFunction interpolator instance for the vector data defined on this grid. @@ -182,7 +187,9 @@ def vector_interpolator( """ raise NotImplementedError("To be defined in subclass.") - def plot_mesh(self, data: ArrayLike | None = None, ax: matplotlib.axes.Axes | None = None): + def plot_mesh( + self, data: NDArray[np.float64] | None = None, ax: matplotlib.axes.Axes | None = None + ): """Plot the grid geometry to a matplotlib figure. Parameters diff --git a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py index 73cb72a..679dcb6 100644 --- a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py +++ b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py @@ -255,7 +255,7 @@ def subset_faces(self, indices: ArrayLike, name: str | None = None) -> UnstructG grid._num_faces = len(indices) grid._num_toroidal = self._num_toroidal - index_flags = np.zeros(self._num_faces, dtype=bool) + index_flags: NDArray[np.bool_] = np.zeros(self._num_faces, dtype=bool) index_flags[indices] = True index_flags = np.tile(index_flags, self._num_toroidal) index_flags.setflags(write=False) @@ -278,7 +278,7 @@ def subset_faces(self, indices: ArrayLike, name: str | None = None) -> UnstructG for cell in cells_original: cells.append(inv_index[ist : ist + len(cell)]) ist += len(cell) - grid._cells = np.array(cells, dtype=np.int32) + grid._cells = np.vstack(cells, dtype=np.int32) grid._cells.setflags(write=False) grid._num_cell = len(grid._cells) grid._num_poloidal = grid._vertices.shape[0] // self._num_toroidal @@ -363,7 +363,7 @@ def subset(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2DE for cell in cells_original: cells.append(inv_index[ist : ist + len(cell)]) ist += len(cell) - grid._cells = np.array(cells, dtype=np.int32) + grid._cells = np.vstack(cells, dtype=np.int32) grid._cells.setflags(write=False) grid._num_cell = len(grid._cells) @@ -395,7 +395,9 @@ def subset(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2DE return grid @override - def interpolator(self, grid_data: ArrayLike, fill_value: float = 0) -> UnstructGridFunction3D: + def interpolator( + self, grid_data: NDArray[np.float64], fill_value: float = 0 + ) -> UnstructGridFunction3D: """Return an UnstructGridFunction3D interpolator instance for the data defined on this grid. On the second and subsequent calls, the interpolator is created as an instance diff --git a/src/cherab/imas/ggd/unstruct_2d_mesh.py b/src/cherab/imas/ggd/unstruct_2d_mesh.py index 08e9b92..7502045 100755 --- a/src/cherab/imas/ggd/unstruct_2d_mesh.py +++ b/src/cherab/imas/ggd/unstruct_2d_mesh.py @@ -294,7 +294,9 @@ def subset(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2D: return grid @override - def interpolator(self, grid_data: ArrayLike, fill_value: float = 0) -> UnstructGridFunction2D: + def interpolator( + self, grid_data: NDArray[np.float64], fill_value: float = 0 + ) -> UnstructGridFunction2D: """Return an `UnstructGridFunction2D` interpolator instance for the data defined on this grid. On the second and subsequent calls, the interpolator is created as an instance of the @@ -322,7 +324,7 @@ def interpolator(self, grid_data: ArrayLike, fill_value: float = 0) -> UnstructG @override def vector_interpolator( - self, grid_vectors: ArrayLike, fill_vector: Vector3D = ZERO_VECTOR + self, grid_vectors: NDArray[np.float64], fill_vector: Vector3D = ZERO_VECTOR ) -> UnstructGridVectorFunction2D: """Return an `UnstructGridVectorFunction2D` interpolator instance for the vector data defined on this grid. diff --git a/src/cherab/imas/ids/bolometer/__init__.py b/src/cherab/imas/ids/bolometer/__init__.py new file mode 100644 index 0000000..e00f925 --- /dev/null +++ b/src/cherab/imas/ids/bolometer/__init__.py @@ -0,0 +1,23 @@ +# 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 bolometer data from IMAS IDS structures.""" + +from . import utility +from ._camera import BoloCamera, BoloChannel, Geometry, load_cameras + +__all__ = ["load_cameras", "utility", "BoloCamera", "BoloChannel", "Geometry"] diff --git a/src/cherab/imas/ids/bolometer/_camera.py b/src/cherab/imas/ids/bolometer/_camera.py new file mode 100644 index 0000000..926db5c --- /dev/null +++ b/src/cherab/imas/ids/bolometer/_camera.py @@ -0,0 +1,322 @@ +# 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 bolometer cameras from IMAS bolometer IDS.""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass, field + +import numpy as np +from numpy.typing import NDArray +from raysect.core.math import Point3D, Vector3D, from_cylindrical + +from imas.ids_structure import IDSStructure +from imas.ids_toplevel import IDSToplevel + +from .utility import CameraType, GeometryType + + +@dataclass +class Geometry: + """Represent the geometric specification of a bolometer sensor head or slit aperture. + + The Geometry describes both simple rectangular / circular sensor faces and + arbitrary polygonal / polyline definitions via explicit coordinates. + It encapsulates a local orthonormal (or user supplied) basis, extents, and derived surface + properties. + """ + + centre: Point3D = field(default_factory=Point3D) + """Geometric centre (reference point) of the sensor/slit in global 3D coordinates.""" + type: GeometryType = GeometryType.RECTANGLE + """Enumerated shape type. + + Common values include `RECTANGLE`, `CIRCLE`, `OUTLINE`, etc. + Defaults to `.GeometryType.RECTANGLE`. + """ + basis_x: Vector3D = field(default_factory=lambda: Vector3D(1, 0, 0)) + """Local x-axis direction vector lying in the sensor/slit plane.""" + basis_y: Vector3D = field(default_factory=lambda: Vector3D(0, 1, 0)) + """Local y-axis direction vector lying in the sensor/slit plane.""" + basis_z: Vector3D = field(default_factory=lambda: Vector3D(0, 0, 1)) + """Local outward-facing normal vector perpendicular to the sensor/slit plane. + + This vector must be directed toward the radiation sources. + When None, it can be derived as the cross product of `basis_x` and `basis_y`. + """ + dx: float = 0.0 + """Width along `basis_x` for rectangular geometry. + + None when not applicable (e.g., circular or polygonal types). + """ + dy: float = 0.0 + """Width along `basis_y` for rectangular geometry. + + None when not applicable (e.g., circular or polygonal types). + """ + surface: float = 0.0 + """Precomputed surface area of the aperture face. + + If None, may be derived from `dx`/`dy` (rectangles), `radius` (circles), or `coords` (polygons) + during validation or runtime. + """ + radius: float = 0.0 + """Radius for circular geometry types. + + None if the geometry is not circular. + """ + coords: NDArray[np.float64] | None = None + """Coordinate array defining the outline in the `basis_x` and `basis_y` plane. + + This array is used to represent a complex geometric outline. + The array shape is ``(2, N)`` where ``N`` is the number of points. + """ + + +@dataclass +class BoloChannel: + """Represent a bolometer camera channel. + + Each channel contains one foil and associated several slits. + """ + + foil: Geometry + """Geometry of the foil used in the bolometer channel.""" + slits: list[Geometry] + """List of geometries representing the slits associated with the bolometer channel.""" + num_subcollimator: int + """Number of sub-collimators in the bolometer channel.""" + thickness_subcollimator: float + """Thickness of each sub-collimator in the bolometer channel [m].""" + + +@dataclass +class BoloCamera: + """Represent a bolometer camera and its associated data. + + This class encapsulates all the properties and data channels associated with + a bolometer camera used for plasma diagnostics. + """ + + name: str + """Unique identifier or label for the bolometer camera.""" + description: str + """Detailed description of the bolometer camera, including its purpose, location, or other relevant information.""" + type: CameraType + """Type of Bolometer camera: Pinhole/Collimator""" + channels: list[BoloChannel] + """List of individual bolometer channels belonging to this camera.""" + + def __repre__(self) -> str: + return f"" + + def __len__(self) -> int: + """Return the number of channels in the bolometer camera. + + Returns + ------- + int + Number of channels. + """ + return len(self.channels) + + def __iter__(self) -> Iterator[BoloChannel]: + """Return an iterator over the bolometer camera channels. + + Returns + ------- + `Iterator[BoloChannel]` + An iterator over the bolometer camera channels. + """ + return iter(self.channels) + + def __getitem__(self, index: int) -> BoloChannel: + """Return the bolometer channel at the specified index. + + Parameters + ---------- + index + Index of the channel to retrieve. + + Returns + ------- + `.BoloChannel` + The bolometer channel at the specified index. + """ + return self.channels[index] + + +def load_cameras(ids: IDSToplevel) -> list[BoloCamera]: + """Load bolometer cameras from the bolometer IDS. + + This function retrieves the camera information from the bolometer IDS and organizes it into a + structured format. + + Parameters + ---------- + ids + The bolometer IDS. + + Returns + ------- + `.BoloCamera` + A list of bolometer camera data structures. + + Raises + ------ + ValueError + If the provided IDS is not a bolometer IDS. + RuntimeError + If no cameras are found in the IDS. + """ + if not str(ids.metadata.name) == "bolometer": + raise ValueError(f"Invalid bolometer IDS ({ids.metadata.name}).") + + cameras = getattr(ids, "camera", []) + if not len(cameras): + raise RuntimeError("No camera found in IDS.") + + bolo_data: list[BoloCamera] = [] + + for camera in cameras: + name = str(camera.name) + description = str(camera.description) + camera_type = CameraType.from_value(camera.type.index.value) + + channels: list[BoloChannel] = [] + for channel in camera.channel: + # Sub-collimator + n_subcollimators = int(channel.subcollimators_n.value) + thickness_subcollimator = float(channel.subcollimators_separation.value) + + channels.append( + BoloChannel( + foil=load_geometry(channel.detector), + slits=[load_geometry(aperture) for aperture in channel.aperture], + num_subcollimator=n_subcollimators if n_subcollimators > 1 else 1, + thickness_subcollimator=thickness_subcollimator + if thickness_subcollimator > 0 + else 0.0, + ) + ) + + bolo_data.append( + BoloCamera( + name=name, + description=description, + type=camera_type, + channels=channels, + ) + ) + + return bolo_data + + +def load_geometry(sensor: IDSStructure) -> Geometry: + """Load the geometry of a sensor or aperture from the bolometer IDS. + + Parameters + ---------- + sensor + Detector or aperture structure object. + + Returns + ------- + `.Geometry` + Object with the following attributes: + + - ``'centre'``: Point3D with the coordinates of the sensor centre, + - ``'type'``: Geometry type of the sensor, + - ``'basis_x'``: Vector3D with the x-basis vector of the sensor, + - ``'basis_y'``: Vector3D with the y-basis vector of the sensor, + - ``'basis_z'``: Vector3D with the z-basis vector of the sensor, + - ``'dx'``: Width of the sensor in the x-direction [m], + - ``'dy'``: Width of the sensor in the y-direction [m], + - ``'surface'``: Surface area of the sensor [m²], + - ``'radius'``: Radius of the sensor [m] if the geometry type is circular, + - ``'coords'``: Outline coordinates for basis_x and basis_y if the geometry type is outline. + + Some keys may not be present depending on the geometry type. + + Raises + ------ + ValueError + If the geometry type is invalid. + """ + geometry = Geometry() + + geometry.centre = from_cylindrical( + sensor.centre.r, sensor.centre.z, np.rad2deg(sensor.centre.phi) + ) + + geometry_type = GeometryType.from_value(int(sensor.geometry_type.value)) + geometry.type = geometry_type + + match geometry_type: + case GeometryType.RECTANGLE: + # Unit vectors + geometry.basis_x = Vector3D( + sensor.x2_unit_vector.x, sensor.x2_unit_vector.y, sensor.x2_unit_vector.z + ) + geometry.basis_y = Vector3D( + sensor.x1_unit_vector.x, sensor.x1_unit_vector.y, sensor.x1_unit_vector.z + ) + geometry.basis_z = Vector3D( + sensor.x3_unit_vector.x, sensor.x3_unit_vector.y, sensor.x3_unit_vector.z + ) + + # Dimensions + geometry.dx = float(sensor.x2_width.value) + geometry.dy = float(sensor.x1_width.value) + + case GeometryType.CIRCULAR: + # Radius + radius = getattr(sensor, "radius", None) + if radius is None or radius <= 0: + raise ValueError(f"Invalid radius ({radius}).") + + geometry.radius = float(radius.value) + + # Unit vectors + geometry.basis_z = Vector3D( + sensor.x3_unit_vector.x, sensor.x3_unit_vector.y, sensor.x3_unit_vector.z + ) + geometry.basis_x = geometry.basis_z.orthogonal().normalise() + geometry.basis_y = geometry.basis_z.cross(geometry.basis_x).normalise() + + case GeometryType.OUTLINE: + # Outline coordinates for basis_x and basis_y + geometry.coords = np.vstack((sensor.outline.x2, sensor.outline.x1)) + + # Unit vectors + geometry.basis_x = Vector3D( + sensor.x2_unit_vector.x, sensor.x2_unit_vector.y, sensor.x2_unit_vector.z + ) + geometry.basis_y = Vector3D( + sensor.x1_unit_vector.x, sensor.x1_unit_vector.y, sensor.x1_unit_vector.z + ) + geometry.basis_z = Vector3D( + sensor.x3_unit_vector.x, sensor.x3_unit_vector.y, sensor.x3_unit_vector.z + ) + + # Surface area + surface = getattr(sensor, "surface", None) + geometry.surface = float(surface.value) if surface is not None else 0.0 + + return geometry diff --git a/src/cherab/imas/ids/bolometer/utility.py b/src/cherab/imas/ids/bolometer/utility.py new file mode 100644 index 0000000..3c4b915 --- /dev/null +++ b/src/cherab/imas/ids/bolometer/utility.py @@ -0,0 +1,101 @@ +# 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 bolometer utility functions.""" + +from enum import Enum +from typing import Self + +__all__ = ["CameraType", "GeometryType"] + + +class CameraType(Enum): + """Enum for camera type. + + The type of the bolometer camera. + + Attributes + ---------- + PINHOLE + The camera is a pinhole camera. + COLLIMATOR + The camera is a collimator camera. + OTHER + The camera is of other type. + """ + + PINHOLE = 1 + COLLIMATOR = 2 + OTHER = 0 + + @classmethod + def from_value(cls, value: int) -> Self: + """Get the camera type from a value. + + Parameters + ---------- + value + The integer value to convert to a camera type. + If the value is not a valid camera type, the default is `OTHER`. + + Returns + ------- + `.CameraType` + The corresponding `.CameraType` enum member. + """ + if value in cls._value2member_map_: + return cls(value) + return cls.OTHER # pyright: ignore[reportReturnType] + + +class GeometryType(Enum): + """Enum for geometry type. + + The geometry type of the bolometer foil or slit. + + Attributes + ---------- + OUTLINE + The geometry is defined by an outline. + CIRCULAR + The geometry is circular. + RECTANGLE + The geometry is rectangular. + """ + + OUTLINE = 1 + CIRCULAR = 2 + RECTANGLE = 3 + + @classmethod + def from_value(cls, value: int) -> Self: + """Get the geometry type from a value. + + Parameters + ---------- + value + The integer value to convert to a geometry type. + If the value is not a valid geometry type, the default is `RECTANGLE`. + + Returns + ------- + `.GeometryType` + The corresponding `.GeometryType` enum member. + """ + if value in cls._value2member_map_: + return cls(value) + return cls.RECTANGLE # pyright: ignore[reportReturnType] diff --git a/src/cherab/imas/ids/common/ggd/load_grid.py b/src/cherab/imas/ids/common/ggd/load_grid.py index 1a35ca1..ab29726 100644 --- a/src/cherab/imas/ids/common/ggd/load_grid.py +++ b/src/cherab/imas/ids/common/ggd/load_grid.py @@ -17,6 +17,8 @@ # under the Licence. """Module for loading GGD grids from IMAS grid_ggd IDS structures.""" +from typing import Literal, overload + from numpy import int32 from numpy.typing import NDArray @@ -30,6 +32,18 @@ __all__ = ["load_grid"] +@overload +def load_grid( + grid_ggd: IDSStructure, with_subsets: Literal[False] = False, num_toroidal: int | None = None +) -> UnstructGrid2D | UnstructGrid2DExtended: ... + + +@overload +def load_grid( + grid_ggd: IDSStructure, with_subsets: Literal[True], num_toroidal: int | None = None +) -> tuple[UnstructGrid2D, dict[str, NDArray[int32]], dict[str, int]]: ... + + def load_grid( grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int | None = 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 7c16821..43e0964 100644 --- a/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py +++ b/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py @@ -17,6 +17,9 @@ # under the Licence. """Module for loading unstructured 2D grids from IMAS grid_ggd IDS structure.""" +from enum import IntEnum +from typing import Literal, overload + import numpy as np from numpy.typing import NDArray @@ -27,9 +30,28 @@ __all__ = ["load_unstruct_grid_2d"] -VERTEX_DIMENSION = 0 -EDGE_DIMENSION = 1 -FACE_DIMENSION = 2 + +class DIMENSION(IntEnum): + """Enumeration for grid dimensions.""" + + VERTEX = 0 + EDGE = 1 + FACE = 2 + + +@overload +def load_unstruct_grid_2d( + grid_ggd: IDSStructure, space_index: int = 0, with_subsets: Literal[False] = False +) -> UnstructGrid2D: ... + + +@overload +def load_unstruct_grid_2d( + grid_ggd: IDSStructure, + space_index: int = 0, + *, + with_subsets: Literal[True], +) -> tuple[UnstructGrid2D, dict[str, NDArray[np.int32]], dict[str, int]]: ... def load_unstruct_grid_2d( @@ -70,15 +92,15 @@ def load_unstruct_grid_2d( grid_name = str(grid_ggd.identifier.name) # Reading vertices - num_vert = len(space.objects_per_dimension[VERTEX_DIMENSION].object) + num_vert = len(space.objects_per_dimension[DIMENSION.VERTEX].object) vertices = np.empty((num_vert, 2), dtype=np.float64) for i in range(num_vert): - vertices[i] = space.objects_per_dimension[VERTEX_DIMENSION].object[i].geometry[:2] + vertices[i] = space.objects_per_dimension[DIMENSION.VERTEX].object[i].geometry[:2] # Reading polygonal cells cells = [] winding_ok = True - for object in space.objects_per_dimension[FACE_DIMENSION].object: + for object in space.objects_per_dimension[DIMENSION.FACE].object: # getting cell from nodes cell = np.asarray_chkfinite(object.nodes, dtype=np.int32) - 1 # Fortran to C indexing if cell.size > 3: @@ -86,7 +108,7 @@ def load_unstruct_grid_2d( edge_dict = {} for boundary in object.boundary: n1, n2 = ( - space.objects_per_dimension[EDGE_DIMENSION].object[boundary.index - 1].nodes - 1 + space.objects_per_dimension[DIMENSION.EDGE].object[boundary.index - 1].nodes - 1 ) # Fortran to C indexing if n1 not in cell or n2 not in cell: # fail, error in the data edge_dict = {} @@ -120,7 +142,7 @@ def load_unstruct_grid_2d( cells.append(cell) if not winding_ok: - print("Warning! Unable to verify that the cell nodes are in the winging order.") + print("Warning! Unable to verify that the cell nodes are in the winding order.") grid = UnstructGrid2D(vertices, cells, name=grid_name) @@ -132,7 +154,7 @@ def load_unstruct_grid_2d( subsets = {} subset_id = {} for subset in grid_ggd.grid_subset: - dimension_is_2d = subset.dimension == FACE_DIMENSION + 1 # C to Fortran indexing + 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 ) diff --git a/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py b/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py index c963670..a6c60ee 100644 --- a/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py +++ b/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py @@ -17,6 +17,8 @@ # under the Licence. """Module for loading unstructured 3D grids from IMAS grid_ggd IDS structure.""" +from enum import IntEnum + import numpy as np from imas.ids_struct_array import IDSStructArray @@ -26,27 +28,44 @@ __all__ = ["load_unstruct_grid_2d_extended", "load_unstruct_grid_3d"] -# Constants for space indices and dimensions -SPACE_RZ = 0 -SPACE_FOURIER = 1 -VERTEX_DIMENSION = 0 -EDGE_DIMENSION = 1 -FACE_DIMENSION = 2 + +class SPACE(IntEnum): + """Enumeration for grid spaces.""" + + RZ = 0 + FOURIER = 1 + + +class DIMENSION(IntEnum): + """Enumeration for grid dimensions.""" + + VERTEX = 0 + EDGE = 1 + FACE = 2 + + +# Constants for tetrahedral cells and default number of toroidal points TETRA_IN_CELL = 5 NUM_TOROIDAL = 64 def load_unstruct_grid_2d_extended( - grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int = NUM_TOROIDAL + grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int | None = NUM_TOROIDAL ) -> UnstructGrid2DExtended: """Load unstructured 2D grid extended in 3D from the ``grid_ggd`` structure. + The grid is created by revolving the 2D cross-section around the torus with evenly spaced + toroidal angles. + + .. todo:: + `with_subsets` is not yet implemented. + Parameters ---------- grid_ggd The ``grid_ggd`` structure. with_subsets - Read grid subset data if True. + Read grid subset data if True, by default False. num_toroidal Number of toroidal points. If specifying more than 1, the grid will be extended in 3D by repeating the 2D @@ -64,11 +83,12 @@ def load_unstruct_grid_2d_extended( If the grid is not an unstructured extended 2D grid. """ # Validate num_toroidal + num_toroidal = num_toroidal or NUM_TOROIDAL if num_toroidal < 1: raise ValueError("The number of toroidal points must be greater than 0.") # Get the R-Z space - space: IDSStructArray = grid_ggd.space[SPACE_RZ] + space: IDSStructArray = grid_ggd.space[SPACE.RZ] # Check if the grid is 2D if len(space.objects_per_dimension) != 3: @@ -81,16 +101,16 @@ def load_unstruct_grid_2d_extended( # ========================================= # Reading vertices (poloidal and toroidal) # ========================================= - num_poloidal = len(space.objects_per_dimension[VERTEX_DIMENSION].object) + num_poloidal = len(space.objects_per_dimension[DIMENSION.VERTEX].object) num_vert = num_poloidal * num_toroidal vertices_rpz = np.empty((num_vert, 3), dtype=float) for i_phi, i_pol in np.ndindex(num_toroidal, num_poloidal): phi = 2.0 * np.pi * i_phi / num_toroidal vertices_rpz[i_pol + i_phi * num_poloidal] = ( - space.objects_per_dimension[VERTEX_DIMENSION].object[i_pol].geometry[0], + space.objects_per_dimension[DIMENSION.VERTEX].object[i_pol].geometry[0], phi, - space.objects_per_dimension[VERTEX_DIMENSION].object[i_pol].geometry[1], + space.objects_per_dimension[DIMENSION.VERTEX].object[i_pol].geometry[1], ) # Convert to cartesian coordinates vertices = np.empty((num_vert, 3), dtype=float) @@ -102,7 +122,7 @@ def load_unstruct_grid_2d_extended( # ========================================= # Reading cells indices # ========================================= - faces: IDSStructArray = space.objects_per_dimension[FACE_DIMENSION].object + faces: IDSStructArray = space.objects_per_dimension[DIMENSION.FACE].object num_faces = len(faces) cells = np.zeros((num_faces * num_toroidal, 8), dtype=np.int32) i_cell = 0 diff --git a/src/cherab/imas/ids/common/slice.py b/src/cherab/imas/ids/common/slice.py index 1711941..93ecd89 100644 --- a/src/cherab/imas/ids/common/slice.py +++ b/src/cherab/imas/ids/common/slice.py @@ -21,7 +21,7 @@ from numpy import inf -from imas import DBEntry +from imas.db_entry import DBEntry from imas.ids_defs import CLOSEST_INTERP from imas.ids_toplevel import IDSToplevel diff --git a/src/cherab/imas/ids/edge_profiles/load_profiles.py b/src/cherab/imas/ids/edge_profiles/load_profiles.py index 233595c..2e4f57b 100644 --- a/src/cherab/imas/ids/edge_profiles/load_profiles.py +++ b/src/cherab/imas/ids/edge_profiles/load_profiles.py @@ -237,7 +237,7 @@ def load_edge_species( else: species_type, species_id = get_neutral(neutral, elements) if species_id in composition[species_type]: - print("Warning! Skipping duplicated neutral: ") + print(f"Warning! Skipping duplicated neutral: {neutral.name.strip()}") else: composition[species_type][species_id] = load_edge_profiles( neutral, grid_subset_index diff --git a/src/cherab/imas/ids/equilibrium/load_equilibrium.py b/src/cherab/imas/ids/equilibrium/load_equilibrium.py index ac56b3d..2a24464 100644 --- a/src/cherab/imas/ids/equilibrium/load_equilibrium.py +++ b/src/cherab/imas/ids/equilibrium/load_equilibrium.py @@ -19,6 +19,8 @@ from __future__ import annotations +from dataclasses import dataclass + import numpy as np from numpy.typing import NDArray from raysect.core.math import Point2D @@ -31,9 +33,32 @@ __all__ = ["load_equilibrium_data"] +@dataclass +class Equilibrium2DData: + """Dataclass to hold 2D plasma equilibrium data.""" + + r: NDArray[np.float64] + z: NDArray[np.float64] + psi_grid: NDArray[np.float64] + psi_axis: float + psi_lcfs: float + magnetic_axis: Point2D + x_points: list[Point2D] + strike_points: list[Point2D] + psi_norm: NDArray[np.float64] + f: NDArray[np.float64] + q: NDArray[np.float64] + phi: NDArray[np.float64] | None + rho_tor_norm: NDArray[np.float64] | None + b_vacuum_radius: float + b_vacuum_magnitude: float + lcfs_polygon: NDArray[np.float64] + time: float + + def load_equilibrium_data( equilibrium_ids: IDSToplevel, -) -> dict[str, NDArray[np.float64] | float | Point2D | list[Point2D] | None]: +) -> Equilibrium2DData: """Load 2D plasma equilibrium data from the equilibrium IDS. Parameters @@ -43,8 +68,8 @@ def load_equilibrium_data( Returns ------- - dict[str, Any] - Dictionary with the following keys and values: + Equilibrium2DData + Dataclass instance containing the following attributes: :r: ``(N, )`` ndarray with R coordinates of rectangular grid, :z: ``(M, )`` ndarray with Z coordinates of rectangular grid, @@ -183,22 +208,22 @@ def load_equilibrium_data( time = float(equilibrium_ids.time[0]) - return { - "r": r, - "z": z, - "psi_grid": psi_grid, - "psi_axis": psi_axis, - "psi_lcfs": psi_lcfs, - "magnetic_axis": magnetic_axis, - "x_points": x_points, - "strike_points": strike_points, - "psi_norm": psi_norm, - "f": f, - "q": q, - "phi": phi, - "rho_tor_norm": rho_tor_norm, - "b_vacuum_radius": b_vacuum_radius, - "b_vacuum_magnitude": b_vacuum_magnitude, - "lcfs_polygon": lcfs_polygon, - "time": time, - } + return Equilibrium2DData( + r=r, + z=z, + psi_grid=psi_grid, + psi_axis=psi_axis, + psi_lcfs=psi_lcfs, + magnetic_axis=magnetic_axis, + x_points=x_points, + strike_points=strike_points, + psi_norm=psi_norm, + f=f, + q=q, + phi=phi, + rho_tor_norm=rho_tor_norm, + b_vacuum_radius=b_vacuum_radius, + b_vacuum_magnitude=b_vacuum_magnitude, + lcfs_polygon=lcfs_polygon, + time=time, + ) diff --git a/src/cherab/imas/ids/equilibrium/load_field.py b/src/cherab/imas/ids/equilibrium/load_field.py index 3f4ec58..d7e075e 100644 --- a/src/cherab/imas/ids/equilibrium/load_field.py +++ b/src/cherab/imas/ids/equilibrium/load_field.py @@ -17,6 +17,8 @@ # under the Licence. """Module for loading magnetic field data from the equilibrium IDS.""" +from dataclasses import dataclass + import numpy as np from imas.ids_defs import EMPTY_INT @@ -27,7 +29,18 @@ __all__ = ["load_magnetic_field_data"] -def load_magnetic_field_data(profiles_2d: IDSStructArray) -> dict: +@dataclass +class MagneticField2DData: + """Dataclass to hold 2D magnetic field data.""" + + r: np.ndarray + z: np.ndarray + b_field_r: np.ndarray + b_field_z: np.ndarray + b_field_phi: np.ndarray + + +def load_magnetic_field_data(profiles_2d: IDSStructArray) -> MagneticField2DData: """Load 2D profiles of the magnetic field components from equilibrium IDS. The magnetic field components are extracted from the ``profiles_2d`` IDS structure, @@ -40,7 +53,8 @@ def load_magnetic_field_data(profiles_2d: IDSStructArray) -> dict: Returns ------- - Dictionary with the following keys: + MagneticField2DData + Dataclass containing the magnetic field data: :r: ``(N,)`` ndarray with R coordinates of rectangular grid. :z: ``(M,)`` ndarray with Z coordinates of rectangular grid. @@ -66,24 +80,24 @@ def load_magnetic_field_data(profiles_2d: IDSStructArray) -> dict: + "rectangular grid for 2D profiles is not found and other grid types are not supported." ) - b_dict = {} - - b_dict["r"] = np.asarray_chkfinite(prof2d.grid.dim1) - b_dict["z"] = np.asarray_chkfinite(prof2d.grid.dim2) - shape = (b_dict["r"].size, b_dict["z"].size) + r = np.asarray_chkfinite(prof2d.grid.dim1) + z = np.asarray_chkfinite(prof2d.grid.dim2) + shape = (r.size, z.size) - b_dict["b_field_r"] = np.asarray_chkfinite(prof2d.b_field_r) - b_dict["b_field_phi"] = np.asarray_chkfinite(prof2d.b_field_phi) - b_dict["b_field_z"] = np.asarray_chkfinite(prof2d.b_field_z) + b_field_r = np.asarray_chkfinite(prof2d.b_field_r) + b_field_phi = np.asarray_chkfinite(prof2d.b_field_phi) + b_field_z = np.asarray_chkfinite(prof2d.b_field_z) - if ( - b_dict["b_field_r"].shape != shape - or b_dict["b_field_phi"].shape != shape - or b_dict["b_field_z"].shape != shape - ): + if b_field_r.shape != shape or b_field_phi.shape != shape or b_field_z.shape != shape: raise RuntimeError( "Unable to read magnetic field: " + "the shape of the magnetic field components does not match the grid shape." ) - return b_dict + return MagneticField2DData( + r=r, + z=z, + b_field_r=b_field_r, + b_field_z=b_field_z, + b_field_phi=b_field_phi, + ) diff --git a/src/cherab/imas/ids/wall/load3d.py b/src/cherab/imas/ids/wall/load3d.py index a28ee2f..29131ee 100644 --- a/src/cherab/imas/ids/wall/load3d.py +++ b/src/cherab/imas/ids/wall/load3d.py @@ -46,8 +46,17 @@ def load_wall_3d( dict[str, dict[str, np.ndarray]] Dictionary of wall components defined by vertices and triangles. The dictionary keys for components are assigns as follows: - ``"{grid_name}.{subset_name}.{material_name}"`` - E.g.: ``"FullTokamak.full_main_chamber_wall.Be"``. + ``"{grid_name}.{subset_name}.{material_name}"``. + So, the returning dictionary looks like: + ```python + { + "FullTokamak.full_main_chamber_wall.Be": { + "vertices": np.ndarray of shape (N, 3), + "triangles": np.ndarray of shape (M, 3), + }, + ... + } + ``` Raises ------ diff --git a/src/cherab/imas/math/functions/vector_functions.pyi b/src/cherab/imas/math/functions/vector_functions.pyi new file mode 100755 index 0000000..f5331e5 --- /dev/null +++ b/src/cherab/imas/math/functions/vector_functions.pyi @@ -0,0 +1,38 @@ +"""Module defining unit vector functions.""" + +from collections.abc import Callable + +from raysect.core.math import Vector3D +from raysect.core.math.function.vector3d import Function1D as VectorFunction1D +from raysect.core.math.function.vector3d import Function2D as VectorFunction2D +from raysect.core.math.function.vector3d import Function3D as VectorFunction3D + +__all__ = ["UnitVector1D", "UnitVector2D", "UnitVector3D"] + +class UnitVector1D(VectorFunction1D): + """Evaluates a unit vector for the given VectorFunction1D instance.""" + + _vector: VectorFunction1D + def __init__( + self, vector: VectorFunction1D | tuple[float, float, float] | Callable[[float], Vector3D] + ) -> None: ... + +class UnitVector2D(VectorFunction2D): + """Evaluates a unit vector for the given VectorFunction2D instance.""" + + _vector: VectorFunction2D + def __init__( + self, + vector: VectorFunction2D | tuple[float, float, float] | Callable[[float, float], Vector3D], + ) -> None: ... + +class UnitVector3D(VectorFunction3D): + """Evaluates a unit vector for the given VectorFunction3D instance.""" + + _vector: VectorFunction3D + def __init__( + self, + vector: VectorFunction3D + | tuple[float, float, float] + | Callable[[float, float, float], Vector3D], + ) -> None: ... diff --git a/src/cherab/imas/math/interpolators/struct_grid_2d_functions.pyi b/src/cherab/imas/math/interpolators/struct_grid_2d_functions.pyi new file mode 100755 index 0000000..fc0e96f --- /dev/null +++ b/src/cherab/imas/math/interpolators/struct_grid_2d_functions.pyi @@ -0,0 +1,71 @@ +"""Module defining simple interpolators for data defined on a 2D structured grid.""" + +from numpy import float64 +from numpy.typing import ArrayLike, NDArray +from raysect.core.math import Vector3D +from raysect.core.math.function.float import Function2D +from raysect.core.math.function.vector3d import Function2D as VectorFunction2D + +ZERO_VECTOR = Vector3D(0, 0, 0) + +class StructGridFunction2D(Function2D): + """Simple interpolator for the data defined on the 2D structured grid. + + Find the cell containing the point (x, y). + Return the data value for this cell or the `fill_value` if the points lies outside the grid. + + Parameters + ---------- + x : (L,) array_like + The corners of the quadrilateral cells along x axis. + y : (M,) array_like + The corners of the quadrilateral cells along y axis. + grid_data : (L-1, M-1) ndarray + Array containing data in the grid cells. + fill_value : float, optional + Value returned outside the gird, by default 0. + """ + + _x: NDArray[float64] + _y: NDArray[float64] + _grid_data: NDArray[float64] + _fill_value: float + + def __init__( + self, + x: ArrayLike, + y: ArrayLike, + grid_data: NDArray[float64], + fill_value: float = 0, + ) -> None: ... + +class StructGridVectorFunction2D(VectorFunction2D): + """Simple vector interpolator for the data defined on the 2D structured grid. + + Find the cell containing the point (x, y). + Return the 3D vector value this cell or the `fill_vector` if the points lies outside the grid. + + Parameters + ---------- + x : (L,) array_like + The corners of the quadrilateral cells along x axis. + y : (M,) array_like + The corners of the quadrilateral cells along y axis. + grid_vectors : (3, L-1, M-1) ndarray + Array containing 3D vectors in the grid cells. + fill_vector : Vector3D + 3D vector returned outside the gird, by default `Vector3D(0, 0, 0)`. + """ + + _x: NDArray[float64] + _y: NDArray[float64] + _grid_vectors: NDArray[float64] + _fill_vector: Vector3D + + def __init__( + self, + x: ArrayLike, + y: ArrayLike, + grid_vectors: NDArray[float64], + fill_vector: Vector3D = ZERO_VECTOR, + ) -> None: ... diff --git a/src/cherab/imas/math/interpolators/struct_grid_3d_functions.pyi b/src/cherab/imas/math/interpolators/struct_grid_3d_functions.pyi new file mode 100644 index 0000000..b58e066 --- /dev/null +++ b/src/cherab/imas/math/interpolators/struct_grid_3d_functions.pyi @@ -0,0 +1,77 @@ +"""Module defining simple interpolators for data defined on a 3D structured grid.""" + +from numpy import float64 +from numpy.typing import ArrayLike, NDArray +from raysect.core.math import Vector3D +from raysect.core.math.function.float import Function3D +from raysect.core.math.function.vector3d import Function3D as VectorFunction3D + +ZERO_VECTOR = Vector3D(0, 0, 0) + +class StructGridFunction3D(Function3D): + """Simple interpolator for the data defined on the 3D structured grid. + + Find the cell containing the point (x, y, z). + Return the data value for this cell or the `fill_value` if the points lies outside the grid. + + Parameters + ---------- + x : (L,) array_like + The corners of the quadrilateral cells along x axis. + y : (M,) array_like + The corners of the quadrilateral cells along y axis. + z : (N,) array_like + The corners of the quadrilateral cells along z axis. + grid_data : (L-1, M-1, N-1) ndarray + Array containing data in the grid cells. + fill_value : float, optional + A value returned outside the grid, by default 0. + """ + + _x: NDArray[float64] + _y: NDArray[float64] + _z: NDArray[float64] + _grid_data: NDArray[float64] + _fill_value: float + def __init__( + self, + x: ArrayLike, + y: ArrayLike, + z: ArrayLike, + grid_data: NDArray[float64], + fill_value: float = 0.0, + ) -> None: ... + +class StructGridVectorFunction3D(VectorFunction3D): + """Simple vector interpolator for the data defined on the 3D structured grid. + + Find the cell containing the point (x, y, z). + Return the 3D vector value this cell or the `fill_vector` if the points lies outside the grid. + + Parameters + ---------- + x : (L,) array_like + The corners of the quadrilateral cells along x axis. + y : (M,) array_like + The corners of the quadrilateral cells along y axis. + z : (N,) array_like + The corners of the quadrilateral cells along z axis. + grid_vectors : (3, L-1, M-1, N-1) ndarray + Array containing 3D vectors in the grid cells. + fill_vector : Vector3D + 3D vector returned outside the gird, by default `Vector3D(0, 0, 0)`. + """ + + _x: NDArray[float64] + _y: NDArray[float64] + _z: NDArray[float64] + _grid_vectors: NDArray[float64] + _fill_vector: Vector3D + def __init__( + self, + x: ArrayLike, + y: ArrayLike, + z: ArrayLike, + grid_vectors: NDArray[float64], + fill_vector: Vector3D = ZERO_VECTOR, + ) -> None: ... diff --git a/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyi b/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyi new file mode 100644 index 0000000..555835a --- /dev/null +++ b/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyi @@ -0,0 +1,149 @@ +"""Module defining simple interpolators for data defined on a 2D unstructured grid.""" + +from numpy import float64, int32 +from numpy.typing import ArrayLike, NDArray +from raysect.core.math import Vector3D +from raysect.core.math.function.float import Function2D +from raysect.core.math.function.float.function2d.interpolate.common import MeshKDTree2D +from raysect.core.math.function.vector3d import Function3D as VectorFunction3D + +ZERO_VECTOR = Vector3D(0, 0, 0) + +class UnstructGridFunction2D(Function2D): + """Simple interpolator for the data defined on the 2D unstructured grid. + + Find the cell containing the point (x, y) using the KDtree algorithm. + Return the data value for this cell or the `fill_value` if the grid does not contain the point. + + Parameters + ---------- + vertex_coords : (N,3) array_like + 2D array-like with the vertex coordinates of triangles. + triangles : (M,3) array_like + 2D integer array-like with the vertex indices forming the triangles. + triangle_to_cell_map : (M,) array_like + 1D integer array-like with the indices of the grid cells (polygons) containing the + triangles. + grid_data : (L,) ndarray + Array containing data in the grid cells. + fill_value : float, optional + Value returned outside the gird, by default 0. + """ + + _kdtree: MeshKDTree2D + _triangle_to_cell_map: NDArray[int32] + _grid_data: NDArray[float64] + _fill_value: float + + def __init__( + self, + vertex_coords: ArrayLike, + triangles: ArrayLike, + triangle_to_cell_map: ArrayLike, + grid_data: NDArray[float64], + fill_value: float = 0.0, + ) -> None: ... + @classmethod + def instance( + cls, + instance: UnstructGridFunction2D | UnstructGridVectorFunction2D, + grid_data: NDArray[float64] | None = None, + fill_value: float | None = None, + ) -> UnstructGridFunction2D: + """Create a new interpolator instance from an existing `UnstructGridFunction2D` or `UnstructGridVectorFunction2D` instance. + + The new interpolator instance will share the same internal acceleration data as the original + interpolator. The grid_data of the new instance can be redefined. + This method should be used if the user has multiple datasets that lie on the same mesh + geometry. Using this methods avoids the repeated rebuilding of the mesh acceleration + structures by sharing the geometry data between multiple interpolator objects. + + If created from the UnstructGridVectorFunction2D instance, the grid_data and the fill_value + must not be None. + + Parameters + ---------- + instance : UnstructGridFunction2D | UnstructGridVectorFunction2D + The instance from which to create the new interpolator. + grid_data : (L,) ndarray, optional + Array containing data in the grid cells. + fill_value : float, optional + Value returned outside the grid, by default None. + If None, inherited from the original instance. + + Returns + ------- + UnstructGridFunction2D + New interpolator instance. + """ + +class UnstructGridVectorFunction2D(VectorFunction3D): + """Simple vector interpolator for the data defined on the 2D unstructured grid. + + Find the cell containing the point (x, y) using the KDtree algorithm. + Return the 3D vector value for this cell or the `fill_vector` if the grid does not contain the + point. + + Parameters + ---------- + vertex_coords : (N,3) array_like + 2D array-like with the vertex coordinates of triangles. + triangles : (M,3) array_like + 2D integer array-like with the vertex indices forming the triangles. + triangle_to_cell_map : (M,1) array_like + 1D integer array-like with the indices of the grid cells (polygons) containing the + triangles. + grid_vectors : (3,K) ndarray + Array containing 3D vectors in the grid cells. + fill_vector : Vector3D + 3D vector returned outside the gird, by default `Vector3D(0, 0, 0)`. + """ + + _kdtree: MeshKDTree2D + _grid_vectors: NDArray[float64] + _triangle_to_cell_map: NDArray[int32] + _fill_vector: Vector3D + + def __init__( + self, + vertex_coords: ArrayLike, + triangles: ArrayLike, + triangle_to_cell_map: ArrayLike, + grid_vectors: NDArray[float64], + fill_vector: Vector3D = ZERO_VECTOR, + ) -> None: ... + @classmethod + def instance( + cls, + instance: UnstructGridVectorFunction2D | UnstructGridFunction2D, + grid_vectors: NDArray[float64] | None = None, + fill_vector: Vector3D | None = None, + ) -> UnstructGridVectorFunction2D: + """Create a new interpolator instance from an existing `UnstructGridVectorFunction2D` or `UnstructGridFunction2D` instance. + + The new interpolator instance will share the same internal acceleration + data as the original interpolator. The grid_vectors of the new instance can + be redefined. + This method should be used if the user has multiple datasets + that lie on the same mesh geometry. Using this methods avoids the + repeated rebuilding of the mesh acceleration structures by sharing the + geometry data between multiple interpolator objects. + + If created from the UnstructGridFunction2D instance, + the grid_vectors and the fill_vector must not be None. + + Parameters + ---------- + instance : UnstructGridVectorFunction2D | UnstructGridFunction2D + The instance from which to create the new interpolator. + grid_vectors : (3,L) ndarray, optional + Array containing vector grid data. + fill_vector : Vector3D, optional + 3D vector returned outside the grid, by default None. + If None, inherited from the original instance. + + Returns + ------- + UnstructGridVectorFunction2D + New interpolator instance. + """ diff --git a/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyx b/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyx index fca864b..f39ea9f 100755 --- a/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyx +++ b/src/cherab/imas/math/interpolators/unstruct_grid_2d_functions.pyx @@ -92,7 +92,7 @@ cdef class UnstructGridFunction2D(Function2D): @classmethod def instance(cls, object instance not None, np.ndarray grid_data=None, object fill_value=None): - """Creates a new interpolator instance from an existing `UnstructGridFunction2D` or + """Create a new interpolator instance from an existing `UnstructGridFunction2D` or `UnstructGridVectorFunction2D` instance. The new interpolator instance will share the same internal acceleration data as the original @@ -116,7 +116,7 @@ cdef class UnstructGridFunction2D(Function2D): Returns ------- - UnstructGridFunction2D | UnstructGridVectorFunction2D + UnstructGridFunction2D New interpolator instance. """ @@ -248,7 +248,7 @@ cdef class UnstructGridVectorFunction2D(VectorFunction2D): np.ndarray grid_vectors=None, Vector3D fill_vector=None ): - """Creates a new interpolator instance from an existing `UnstructGridVectorFunction2D` or + """Create a new interpolator instance from an existing `UnstructGridVectorFunction2D` or `UnstructGridFunction2D` instance. The new interpolator instance will share the same internal acceleration data as the original @@ -271,7 +271,7 @@ cdef class UnstructGridVectorFunction2D(VectorFunction2D): Returns ------- - UnstructGridFunction2D | UnstructGridVectorFunction2D + UnstructGridVectorFunction2D The new interpolator instance. """ diff --git a/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyi b/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyi new file mode 100644 index 0000000..a7a2727 --- /dev/null +++ b/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyi @@ -0,0 +1,149 @@ +"""Module defining simple interpolators for data defined on a 3D unstructured grid.""" + +from numpy import float64, int32 +from numpy.typing import ArrayLike, NDArray +from raysect.core.math import Vector3D +from raysect.core.math.function.float import Function3D +from raysect.core.math.function.float.function3d.interpolate.common import MeshKDTree3D +from raysect.core.math.function.vector3d import Function3D as VectorFunction3D + +ZERO_VECTOR = Vector3D(0, 0, 0) + +class UnstructGridFunction3D(Function3D): + """Simple interpolator for the data defined on the 3D unstructured grid. + + Find the cell containing the point (x, y, z) using the KDtree algorithm. + Return the data value for this cell or the `fill_value` if the grid does not contain the point. + + Parameters + ---------- + vertex_coords : (N,3) array_like + 3D array-like with the vertex coordinates of tetrahedra. + tetrahedra : (M,4) array_like + 3D integer array-like with the vertex indices forming the tetrahedra. + tetra_to_cell_map : (M,) array_like + 1D integer array-like with the indices of the grid cells (cube) containing the tetrahedra. + grid_data : (L,) ndarray + Array containing data in the grid cells. + fill_value : float, optional + Value returned outside the grid, by default 0. + """ + + _kdtree: MeshKDTree3D + _tetra_to_cell_map: NDArray[int32] + _grid_data: NDArray[float64] + _fill_value: float + + def __init__( + self, + vertex_coords: ArrayLike, + tetrahedra: ArrayLike, + tetra_to_cell_map: ArrayLike, + grid_data: NDArray[float64], + fill_value: float = 0.0, + ) -> None: ... + @classmethod + def instance( + cls, + instance: UnstructGridFunction3D | UnstructGridVectorFunction3D, + grid_data: NDArray[float64] | None = None, + fill_value: float | None = None, + ) -> UnstructGridFunction3D: + """Create a new interpolator instance from an existing `UnstructGridFunction3D` or `UnstructGridVectorFunction3D` instance. + + The new interpolator instance will share the same internal acceleration + data as the original interpolator. The grid_data of the new instance can + be redefined. + This method should be used if the user has multiple datasets + that lie on the same mesh geometry. Using this methods avoids the + repeated rebuilding of the mesh acceleration structures by sharing the + geometry data between multiple interpolator objects. + + If created from the UnstructGridVectorFunction3D instance, + the grid_data and the fill_value must not be None. + + Parameters + ---------- + instance : UnstructGridFunction3D | UnstructGridVectorFunction3D + The instance from which to create the new interpolator. + grid_data : (L,) ndarray, optional + Array containing data in the grid cells. + fill_value : float, optional + Value returned outside the grid, by default None. + If None, inherited from the original instance. + + Returns + ------- + UnstructGridFunction3D + New interpolator instance. + """ + +class UnstructGridVectorFunction3D(VectorFunction3D): + """Simple vector interpolator for the data defined on the 3D unstructured grid. + + Find the cell containing the point (x, y, z) using the KDtree algorithm. + Return the 3D vector value for this cell or the `fill_vector` if the grid does not contain the + point. + + Parameters + ---------- + vertex_coords : (N,3) array_like + 3D array-like with the vertex coordinates of tetrahedra. + tetrahedra : (M,4) array_like + 3D integer array-like with the vertex indices forming the tetrahedra. + tetra_to_cell_map : (M,) array_like + 1D integer array-like with the indices of the grid cells (cube) containing the tetrahedra. + grid_vectors : (3,L) ndarray + Array containing 3D vectors in the grid cells. + fill_vector : Vector3D, optional + 3D vector returned outside the grid, by default `Vector3D(0, 0, 0)`. + """ + + _kdtree: MeshKDTree3D + _grid_vectors: NDArray[float64] + _tetra_to_cell_map: NDArray[int32] + _fill_vector: Vector3D + + def __init__( + self, + vertex_coords: ArrayLike, + tetrahedra: ArrayLike, + tetra_to_cell_map: ArrayLike, + grid_vectors: NDArray[float64], + fill_vector: Vector3D = ZERO_VECTOR, + ) -> None: ... + @classmethod + def instance( + cls, + instance: UnstructGridVectorFunction3D | UnstructGridFunction3D, + grid_vectors: NDArray[float64] | None = None, + fill_vector: Vector3D | None = None, + ) -> UnstructGridVectorFunction3D: + """Create a new interpolator instance from an existing `UnstructGridVectorFunction3D` or `UnstructGridFunction3D` instance. + + The new interpolator instance will share the same internal acceleration + data as the original interpolator. The grid_vectors of the new instance can + be redefined. + This method should be used if the user has multiple datasets + that lie on the same mesh geometry. Using this methods avoids the + repeated rebuilding of the mesh acceleration structures by sharing the + geometry data between multiple interpolator objects. + + If created from the UnstructGridFunction3D instance, + the grid_vectors and the fill_vector must not be None. + + Parameters + ---------- + instance : UnstructGridVectorFunction3D | UnstructGridFunction3D + The instance from which to create the new interpolator. + grid_vectors : (3,L) ndarray, optional + Array containing vector grid data. + fill_vector : Vector3D, optional + 3D vector returned outside the grid, by default None. + If None, inherited from the original instance. + + Returns + ------- + UnstructGridVectorFunction3D + New interpolator instance. + """ diff --git a/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyx b/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyx index ea5e30b..17c843a 100644 --- a/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyx +++ b/src/cherab/imas/math/interpolators/unstruct_grid_3d_functions.pyx @@ -116,7 +116,7 @@ cdef class UnstructGridFunction3D(Function3D): Returns ------- - UnstructGridFunction3D | UnstructGridVectorFunction3D + UnstructGridFunction3D New interpolator instance. """ @@ -274,7 +274,7 @@ cdef class UnstructGridVectorFunction3D(VectorFunction3D): Returns ------- - UnstructGridVectorFunction3D | UnstructGridFunction3D + UnstructGridVectorFunction3D New interpolator instance. """ diff --git a/src/cherab/imas/math/tetrahedralize.pyi b/src/cherab/imas/math/tetrahedralize.pyi new file mode 100644 index 0000000..eedab05 --- /dev/null +++ b/src/cherab/imas/math/tetrahedralize.pyi @@ -0,0 +1,111 @@ +from numpy import float64, int32 +from numpy.typing import NDArray + +def cell_to_5tetra(cells: NDArray[int32]) -> NDArray[int32]: + """Generate tetrahedral indices by dividing one cell into 5 tetrahedra. + + One cubic-like cell having 8 vertices can be divided into a minimum of five tetrahedra. + Reference: https://www.mathworks.com/matlabcentral/mlc-downloads/downloads/submissions/48509/versions/3/previews/COMP_GEOM_TLBX/html/Divide_hypercube_5_simplices_3D.html + + .. note:: + If one side face of the cell is twisted, Adjacent cells do not share a plane with each + other. + In this case, it is required to alternate combinations of tetrahedra, or simply use + :func:`~cherab.imas.math.tetrahedralize.cell_to_6tetra`. + + Parameters + ---------- + cells : (N,8) ndarray [numpy.int32] + cell indices 2D array, the shape of which is :math:`(N, 8)`, where :math:`N` is the number + of cells. + + Returns + ------- + `(5N,4) ndarray` + tetrahedra indices array, the shape of which is :math:`(5N, 4)`. + + Examples + -------- + >>> import numpy as np + >>> from cherab.imas.math.tetrahedralize import cell_to_5tetra + >>> + >>> array = np.arange(16, dtype=np.int32).reshape((2, -1)) + >>> array + array([[ 0, 1, 2, 3, 4, 5, 6, 7], + [ 8, 9, 10, 11, 12, 13, 14, 15]], dtype=int32) + >>> cell_to_5tetra(array) + array([[ 0, 1, 3, 4], + [ 1, 3, 4, 6], + [ 3, 6, 7, 4], + [ 1, 2, 3, 6], + [ 1, 6, 4, 5], + [ 8, 9, 11, 12], + [ 9, 11, 12, 14], + [11, 14, 15, 12], + [ 9, 10, 11, 14], + [ 9, 14, 12, 13]], dtype=int32) + """ + +def cell_to_6tetra(cells: NDArray[int32]) -> NDArray[int32]: + """Generate tetrahedral indices by dividing one cell into 6 tetrahedra. + + One cubic-like cell having 8 vertices can be divided into six tetrahedra. + This manner is useful when the cell is twisted. + + Parameters + ---------- + cells : (N,8) ndarray [numpy.int32] + Cell indices 2D array, the shape of which is :math:`(N, 8)`, where :math:`N` is the number + of cells. + + Returns + ------- + `(6N,4) ndarray` + Tetrahedra indices array, the shape of which is :math:`(6N, 4)`. + + Examples + -------- + >>> import numpy as np + >>> from cherab.imas.math.tetrahedralize import cell_to_6tetra + >>> + >>> array = np.arange(16, dtype=np.int32).reshape((2, -1)) + >>> array + array([[ 0, 1, 2, 3, 4, 5, 6, 7], + [ 8, 9, 10, 11, 12, 13, 14, 15]], dtype=int32) + >>> cell_to_6tetra(array) + array([[ 6, 2, 1, 0], + [ 7, 3, 2, 0], + [ 0, 7, 6, 2], + [ 1, 5, 6, 4], + [ 0, 4, 6, 7], + [ 6, 4, 0, 1], + [14, 10, 9, 8], + [15, 11, 10, 8], + [ 8, 15, 14, 10], + [ 9, 13, 14, 12], + [ 8, 12, 14, 15], + [14, 12, 8, 9]], dtype=int32) + """ + +def calculate_tetra_volume(vertices: NDArray[float64], tetrahedra: NDArray[int32]) -> float: + """Calculate the volume of tetrahedra. + + Parameters + ---------- + vertices : ndarray [numpy.float64] + Vertices of tetrahedra. + tets : ndarray [numpy.int32] + Tetrahedra indices. + + Returns + ------- + float + Volume of tetrahedra. + + Examples + -------- + >>> vertices = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) + >>> tetras = np.array([[0, 1, 2, 3]], dtype=np.int32) + >>> calculate_tetra_volume(vertices, tetras) + 0.16666666666666666 + """ diff --git a/src/cherab/imas/math/tetrahedralize.pyx b/src/cherab/imas/math/tetrahedralize.pyx index 508fdce..203a76e 100644 --- a/src/cherab/imas/math/tetrahedralize.pyx +++ b/src/cherab/imas/math/tetrahedralize.pyx @@ -35,7 +35,7 @@ cpdef ndarray[int32_t, ndim=2] cell_to_5tetra(const int32_t[:, ::1] cells): Returns ------- - (5N,4) ndarray + `(5N,4) ndarray` tetrahedra indices array, the shape of which is :math:`(5N, 4)`. Examples @@ -43,7 +43,7 @@ cpdef ndarray[int32_t, ndim=2] cell_to_5tetra(const int32_t[:, ::1] cells): >>> import numpy as np >>> from cherab.imas.math.tetrahedralize import cell_to_5tetra >>> - >>> array = np.arrange(16, dtype=np.int32).reshape((2, -1)) + >>> array = np.arange(16, dtype=np.int32).reshape((2, -1)) >>> array array([[ 0, 1, 2, 3, 4, 5, 6, 7], [ 8, 9, 10, 11, 12, 13, 14, 15]], dtype=int32) @@ -109,7 +109,7 @@ cpdef ndarray[int32_t, ndim=2] cell_to_6tetra(const int32_t[:, ::1] cells): Returns ------- - (6N,4) ndarray + `(6N,4) ndarray` Tetrahedra indices array, the shape of which is :math:`(6N, 4)`. Examples @@ -117,7 +117,7 @@ cpdef ndarray[int32_t, ndim=2] cell_to_6tetra(const int32_t[:, ::1] cells): >>> import numpy as np >>> from cherab.imas.math.tetrahedralize import cell_to_6tetra >>> - >>> array = np.arrange(16, dtype=np.int32).reshape((2, -1)) + >>> array = np.arange(16, dtype=np.int32).reshape((2, -1)) >>> array array([[ 0, 1, 2, 3, 4, 5, 6, 7], [ 8, 9, 10, 11, 12, 13, 14, 15]], dtype=int32) diff --git a/src/cherab/imas/observer/__init__.py b/src/cherab/imas/observer/__init__.py new file mode 100644 index 0000000..a10d71c --- /dev/null +++ b/src/cherab/imas/observer/__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 creating observer objects from IMAS.""" + +from .bolometer import load_bolometers + +__all__ = ["load_bolometers"] diff --git a/src/cherab/imas/observer/bolometer.py b/src/cherab/imas/observer/bolometer.py new file mode 100644 index 0000000..eea49b5 --- /dev/null +++ b/src/cherab/imas/observer/bolometer.py @@ -0,0 +1,443 @@ +# 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 bolometer cameras from IMAS bolometer IDS.""" + +from __future__ import annotations + +from typing import overload + +from raysect.core.constants import ORIGIN +from raysect.core.math import Point3D, Vector3D, rotate_basis, translate +from raysect.core.scenegraph._nodebase import _NodeBase +from raysect.optical.material import AbsorbingSurface +from raysect.primitive import Box, Subtract, Union +from raysect.primitive.csg import CSGPrimitive + +from cherab.tools.observers.bolometry import BolometerCamera, BolometerFoil, BolometerSlit +from imas.db_entry import DBEntry + +from ..ids.bolometer import load_cameras +from ..ids.bolometer._camera import BoloCamera, Geometry +from ..ids.bolometer.utility import CameraType, GeometryType +from ..ids.common import get_ids_time_slice + +__all__ = ["load_bolometers"] + +X_AXIS = Vector3D(1, 0, 0) +Y_AXIS = Vector3D(0, 1, 0) +Z_AXIS = Vector3D(0, 0, 1) +THICKNESS = 2.0e-3 # Thickness of camera box walls +THICKNESS_INNER_LAYER = 0.5e-3 # Thickness of inner aperture layers +EPS = 1e-4 # Small epsilon value to avoid numerical issues + + +@overload +def load_bolometers( + backend_id: int, + db_name: str, + pulse: int, + run: int, + user_name: str | None = None, + data_version: str | None = None, + *, + parent: _NodeBase | None = None, + shot: int | None = None, + dd_version: str | None = None, + xml_path: str | None = None, +) -> list[BolometerCamera]: ... + + +@overload +def load_bolometers( + uri: str, + mode: str, + *, + parent: _NodeBase | None = None, + dd_version: str | None = None, + xml_path: str | None = None, +) -> list[BolometerCamera]: ... + + +def load_bolometers( + *args, + parent: _NodeBase | None = None, + **kwargs, +) -> list[BolometerCamera]: + """Load bolometer cameras from IMAS bolometer IDS. + + .. note:: + This function requires the Data dictionary v4.1.0 or later. + + Parameters + ---------- + *args + Arguments passed to `~imas.db_entry.DBEntry`. + parent + The parent node of `~cherab.tools.observers.bolometry.BolometerCamera` in the Raysect + scene-graph, by default None. + **kwargs + Keyword arguments passed to `~imas.db_entry.DBEntry` constructor. + + Returns + ------- + `list[BolometerCamera]` + List of `~cherab.tools.observers.bolometry.BolometerCamera` objects. + + Examples + -------- + >>> from raysect.optical import World + >>> world = World() + + If you have a local IMAS database and store the "bolometer.h5" file there: + + >>> bolometers = load_bolometers("imas:hdf5?path=path/to/db/", "r", parent=world) + + If you want to load netCDF files directly: + + >>> bolometers = load_bolometers("path/to/bolometer_file.nc", "r", parent=world) + """ + # Load bolometer IDS + with DBEntry(*args, **kwargs) as entry: + # Get available time slices + ids = get_ids_time_slice(entry, "bolometer") + + # Extract bolometer data + bolo_data = load_cameras(ids) + + bolometers: list[BolometerCamera] = [] + + for data in bolo_data: + # Skip empty cameras + if len(data.channels) == 0: + continue + + # ------------------ + # === Camera Box === + # ------------------ + camera_box = _create_camera_box(data) + camera = BolometerCamera(camera_geometry=camera_box, name=data.name, parent=parent) + + match data.type: + case CameraType.PINHOLE: + # ---------------------- + # === Slit (Pinhole) === + # ---------------------- + # Pick up only first aperture and use it for all channels + slit_data = data.channels[0].slits[0] + slit = BolometerSlit( + f"slit-{data.name}", + slit_data.centre, + slit_data.basis_x, + slit_data.dx, + slit_data.basis_y, + slit_data.dy, + curvature_radius=(slit_data.radius or 0.0) + if slit_data.type == GeometryType.CIRCULAR + else 0.0, + parent=camera, + ) + case CameraType.COLLIMATOR: + slit = None # Defined per channel below + case _: + raise NotImplementedError(f"Camera type {data.type} not supported yet.") + + for i_channel, channel in enumerate(data.channels): + # ------------------------- + # === Slit (Collimator) === + # ------------------------- + # Concatenate top plate slits into one slit if multiple sub-collimators + # NOTE: Top plate slits are accumulated in the beginning of the slits list + # NOTE: All slit geometry types are assumed to be the RECTANGLE. + # NOTE: Sub-collimators are assumed to be aligned in a row along the basis_y direction. + if (num_subcol := channel.num_subcollimator) > 1: + _slit_data = channel.slits[:num_subcol] + centre = ( + ORIGIN + + sum([ORIGIN.vector_to(s.centre) for s in _slit_data], Vector3D(0, 0, 0)) + / num_subcol + ) + basis_x = _slit_data[0].basis_x + basis_y = _slit_data[0].basis_y + dx = _slit_data[0].dx + dy = ( + ORIGIN + + ORIGIN.vector_to(_slit_data[0].centre) + - _slit_data[0].dy * 0.5 * basis_y + ).distance_to( + ORIGIN + + ORIGIN.vector_to(_slit_data[-1].centre) + + _slit_data[-1].dy * 0.5 * basis_y + ) + + slit_data = Geometry( + type=GeometryType.RECTANGLE, + centre=centre, + basis_x=basis_x, + basis_y=basis_y, + dx=dx, + dy=dy, + ) + else: + slit_data = channel.slits[0] + + if data.type == CameraType.COLLIMATOR: + # Create slit object + match slit_data.type: + case GeometryType.CIRCULAR | GeometryType.RECTANGLE: + slit = BolometerSlit( + f"slit-{data.name}-ch{i_channel}", + slit_data.centre, + slit_data.basis_x, + slit_data.dx, + slit_data.basis_y, + slit_data.dy, + curvature_radius=(slit_data.radius or 0.0) + if slit_data.type == GeometryType.CIRCULAR + else 0.0, + parent=camera, + ) + case _: + raise NotImplementedError("Outline geometry not supported yet.") + + # ------------ + # === Foil === + # ------------ + match channel.foil.type: + case GeometryType.CIRCULAR | GeometryType.RECTANGLE: + foil = BolometerFoil( + f"foil-{data.name}-ch{i_channel}", + channel.foil.centre, + channel.foil.basis_x, + channel.foil.dx, + channel.foil.basis_y, + channel.foil.dy, + slit, + curvature_radius=(channel.foil.radius or 0.0) + if channel.foil.type == GeometryType.CIRCULAR + else 0.0, + parent=camera, + ) + case _: + raise NotImplementedError("Outline geometry not supported yet.") + + # Add foil to the camera + camera.add_foil_detector(foil) + + bolometers.append(camera) + + return bolometers + + +def _create_camera_box(bolo_data: BoloCamera) -> CSGPrimitive: + """Create a camera housing box geometry. + + This box is represented as a Subtract primitive, where a smaller inner box is subtracted + from a larger outer box to create a hollow structure. + Additionally, top plate apertures and other apertures (inner layers) are included as Unions. + + Returns + ------- + `CSGPrimitive` + The camera housing box geometry. + """ + # Extract slits data + slits_top: list[Geometry] + slits_inner: list[list[Geometry]] = [] + if bolo_data.type == CameraType.PINHOLE: + # Pick up only first aperture and use it for all channels + slits_top = [bolo_data[0].slits[0]] + else: + # Consider sub-collimators for collimator cameras + num_subcol = bolo_data[0].num_subcollimator + num_channel = len(bolo_data) + _s: list[Geometry] = [] + slits_top = sum( + [bolo_data[i_ch].slits[:num_subcol] for i_ch in range(num_channel)], + _s, + ) + slits_inner = [ + sum( + [ + bolo_data[i_ch].slits[i_slit : i_slit + num_subcol] + for i_ch in range(num_channel) + ], + _s, + ) + for i_slit in range(num_subcol, len(bolo_data[0].slits), num_subcol) + ] + + # ------------------------------- + # === Local coordinate system === + # ------------------------------- + # The local origin is placed at the average position of top plate slits of all foil detectors, + # and the basis vectors are defined based on the average normal and basis_x vectors of the slits. + origin = ORIGIN + sum( + [ORIGIN.vector_to(slit.centre) for slit in slits_top], + Vector3D(0, 0, 0), + ) / len(slits_top) + + basis_z: Vector3D = sum( + [slit.basis_z for slit in slits_top], + Vector3D(0, 0, 0), + ).normalise() + + basis_x: Vector3D = sum( + [slit.basis_x for slit in slits_top], + Vector3D(0, 0, 0), + ).normalise() + + basis_y = basis_z.cross(basis_x).normalise() + + # Transformation matrix from local to global coordinate system + to_root = translate(*origin) * rotate_basis(basis_z, basis_y) + + # ------------------------------ + # === Determine box geometry === + # ------------------------------ + EPS_WIDTH = 1e-3 + EPS_HEIGHT = 1e-3 + EPS_DEPTH = 1e-3 + + # Collect all corner points of foils and slits + pts: list[Point3D] = [] + for geometry in [bolo_data[i].foil for i in range(len(bolo_data))] + slits_top: + pts += _get_verts(geometry) + + # Camera box dimensions (in local coordinate system) + box_width_u = max([origin.vector_to(p).dot(basis_x) for p in pts]) + EPS_WIDTH + box_width_l = min([origin.vector_to(p).dot(basis_x) for p in pts]) - EPS_WIDTH + box_height_u = max([origin.vector_to(p).dot(basis_y) for p in pts]) + EPS_HEIGHT + box_height_l = min([origin.vector_to(p).dot(basis_y) for p in pts]) - EPS_HEIGHT + box_depth_u = max([origin.vector_to(p).dot(basis_z) for p in pts]) # slit top plate + box_depth_l = min([origin.vector_to(p).dot(basis_z) for p in pts]) - EPS_DEPTH + + # ------------------------- + # === Create Hollow Box === + # ------------------------- + camera_box = Box( + lower=ORIGIN + box_width_l * X_AXIS + box_height_l * Y_AXIS + box_depth_l * Z_AXIS, + upper=ORIGIN + box_width_u * X_AXIS + box_height_u * Y_AXIS + box_depth_u * Z_AXIS, + name=f"inner-box-{bolo_data.name}", + ) + outside_box = Box( + lower=camera_box.lower - Vector3D(THICKNESS, THICKNESS, THICKNESS), + upper=camera_box.upper + Vector3D(THICKNESS, THICKNESS, THICKNESS), + name=f"outside-box-{bolo_data.name}", + ) + camera_box = Subtract(outside_box, camera_box, name=f"hollow-box-{bolo_data.name}") + + # ------------------------------------------ + # === Clip out Top Plate Apertures Holes === + # ------------------------------------------ + for slit in slits_top: + slit_verts = _get_verts(slit) + aperture_box = Box( + lower=ORIGIN + + min([origin.vector_to(p).dot(basis_x) for p in slit_verts]) * X_AXIS + + min([origin.vector_to(p).dot(basis_y) for p in slit_verts]) * Y_AXIS + + (box_depth_u - EPS) * Z_AXIS, + upper=ORIGIN + + max([origin.vector_to(p).dot(basis_x) for p in slit_verts]) * X_AXIS + + max([origin.vector_to(p).dot(basis_y) for p in slit_verts]) * Y_AXIS + + (box_depth_u + THICKNESS + EPS) * Z_AXIS, + name=f"aperture-box-top-{bolo_data.name}", + ) + camera_box = Subtract(camera_box, aperture_box, name=f"aperture-top-{bolo_data.name}") + + # ---------------------------------- + # === Add Inner Apertures Plates === + # ---------------------------------- + # NOTE: Assume all inner apertures use the same local coordinate system as the top plate. + for slits in slits_inner: + # Create Inner Aperture Plate layer + layer_depth_z = basis_z.dot(origin.vector_to(slits[0].centre)) + layer = Box( + lower=Point3D( + outside_box.lower.x, + outside_box.lower.y, + layer_depth_z - THICKNESS_INNER_LAYER * 0.5, + ), + upper=Point3D( + outside_box.upper.x, + outside_box.upper.y, + layer_depth_z + THICKNESS_INNER_LAYER * 0.5, + ), + name=f"inner-plate-{bolo_data.name}", + ) + + for slit in slits: + slit_verts = _get_verts(slit) + aperture_box = Box( + lower=ORIGIN + + min([origin.vector_to(p).dot(basis_x) for p in slit_verts]) * X_AXIS + + min([origin.vector_to(p).dot(basis_y) for p in slit_verts]) * Y_AXIS + + (layer_depth_z - THICKNESS_INNER_LAYER * 0.5 - EPS) * Z_AXIS, + upper=ORIGIN + + max([origin.vector_to(p).dot(basis_x) for p in slit_verts]) * X_AXIS + + max([origin.vector_to(p).dot(basis_y) for p in slit_verts]) * Y_AXIS + + (layer_depth_z + THICKNESS_INNER_LAYER * 0.5 + EPS) * Z_AXIS, + name=f"inner-aperture-box-{bolo_data.name}", + ) + layer = Subtract(layer, aperture_box, name=f"inner-aperture-plate-{bolo_data.name}") + + camera_box = Union(camera_box, layer) + + # Transform to global coordinate system + camera_box.transform = to_root + + # Apply absorbing material + camera_box.material = AbsorbingSurface() + + # Name the camera box + camera_box.name = f"camera-box-{bolo_data.name}" + + return camera_box + + +def _get_verts(geometry: Geometry) -> list[Point3D]: + """Get the geometry vertices. + + For example, if the geometry is rectangular, return 4 corner vertices. + + Parameters + ---------- + geometry + Geometry structure object. + + Returns + ------- + `list[Point3D]` + List of geometry vertices. + """ + dx = geometry.dx + dy = geometry.dy + basis_x = geometry.basis_x + basis_y = geometry.basis_y + center = geometry.centre + match geometry.type: + case GeometryType.RECTANGLE: + verts = [ + center + 0.5 * dx * basis_x + 0.5 * dy * basis_y, + center - 0.5 * dx * basis_x + 0.5 * dy * basis_y, + center - 0.5 * dx * basis_x - 0.5 * dy * basis_y, + center + 0.5 * dx * basis_x - 0.5 * dy * basis_y, + ] + case _: + raise NotImplementedError(f"Geometry type {geometry.type} not implemented yet.") + + return verts diff --git a/src/cherab/imas/plasma/edge.py b/src/cherab/imas/plasma/edge.py index 22773e6..0cd8db1 100644 --- a/src/cherab/imas/plasma/edge.py +++ b/src/cherab/imas/plasma/edge.py @@ -237,7 +237,7 @@ def get_edge_interpolators( return interpolators.freeze() -def _get_velocity_interpolators(grid, profiles, b_field=None): +def _get_velocity_interpolators(grid: GGDGrid, profiles, b_field=None): # Note: np.all(None == 0) returns False vrad = None if np.all(profiles["velocity_radial"] == 0) else profiles["velocity_radial"] vpol = None if np.all(profiles["velocity_poloidal"] == 0) else profiles["velocity_poloidal"] @@ -260,7 +260,7 @@ def _get_velocity_interpolators(grid, profiles, b_field=None): return _get_parallel_velocity_interpolators(grid, vpar, vrad, b_field) -def _get_cylindrical_velocity_interpolators(grid, vr, vz, vtor): +def _get_cylindrical_velocity_interpolators(grid: GGDGrid, vr, vz, vtor): if vr is None and vz is None and vtor is None: if grid.dimension == 2: return ConstantVector2D( @@ -279,7 +279,7 @@ def _get_cylindrical_velocity_interpolators(grid, vr, vz, vtor): return grid.vector_interpolator(np.array([vr, vtor, vz])) -def _get_parallel_velocity_interpolators(grid, vpar, vrad, b_field): +def _get_parallel_velocity_interpolators(grid: GGDGrid, vpar, vrad, b_field): if vpar is None and vrad is None: if grid.dimension == 2: # 2D case return ConstantVector2D( @@ -303,7 +303,7 @@ def _get_parallel_velocity_interpolators(grid, vpar, vrad, b_field): return vpar_i * parallel_vector + vrad_i * surface_normal -def _get_poloidal_velocity_interpolators(grid, vpol, vrad, vtor, b_field): +def _get_poloidal_velocity_interpolators(grid: GGDGrid, vpol, vrad, vtor, b_field): if vpol is None and vrad is None and vtor is None: if grid.dimension == 2: # 2D case return ConstantVector2D( @@ -330,7 +330,7 @@ def _get_poloidal_velocity_interpolators(grid, vpol, vrad, vtor, b_field): return vpol_i * poloidal_vector + vrad_i * surface_normal + vtor_i * toroidal_vector -def _get_components_from_vpar(grid, vpar, b_field): +def _get_components_from_vpar(grid: GGDGrid, vpar, b_field): vpol = np.zeros(grid.num_cell, dtype=np.float64) vtor = np.zeros(grid.num_cell, dtype=np.float64) diff --git a/src/cherab/imas/plasma/equilibrium.py b/src/cherab/imas/plasma/equilibrium.py index 4e725c9..2163988 100644 --- a/src/cherab/imas/plasma/equilibrium.py +++ b/src/cherab/imas/plasma/equilibrium.py @@ -17,10 +17,13 @@ # under the Licence. """Module for loading plasma equilibrium and magnetic field from the equilibrium IDS.""" +from typing import Literal, overload + import numpy as np from raysect.core.math.function.float import Interpolator1DArray, Interpolator2DArray from raysect.core.math.function.vector3d import FloatToVector3DFunction2D +from cherab.imas.ids.equilibrium.load_equilibrium import Equilibrium2DData from cherab.tools.equilibrium import EFITEquilibrium from imas import DBEntry @@ -30,6 +33,28 @@ __all__ = ["load_equilibrium", "load_magnetic_field"] +@overload +def load_equilibrium( + *args, + time: float = 0, + occurrence: int = 0, + time_threshold: float = np.inf, + with_psi_interpolator: Literal[False] = False, + **kwargs, +) -> EFITEquilibrium: ... + + +@overload +def load_equilibrium( + *args, + time: float = 0, + occurrence: int = 0, + time_threshold: float = np.inf, + with_psi_interpolator: Literal[True], + **kwargs, +) -> tuple[EFITEquilibrium, Interpolator1DArray | None]: ... + + def load_equilibrium( *args, time: float = 0, @@ -72,42 +97,45 @@ def load_equilibrium( entry, "equilibrium", time=time, occurrence=occurrence, time_threshold=time_threshold ) - equilibrium_dict = load_equilibrium_data(equilibrium_ids) + eq_data = load_equilibrium_data(equilibrium_ids) - cocos_11to3(equilibrium_dict) + cocos_11to3(eq_data) - equilibrium_dict["psi_norm"][0] = min(0, equilibrium_dict["psi_norm"][0]) - equilibrium_dict["psi_norm"][-1] = max(1.0, equilibrium_dict["psi_norm"][-1]) - - f_profile = np.array([equilibrium_dict["psi_norm"], equilibrium_dict["f"]]) - q_profile = np.array([equilibrium_dict["psi_norm"], equilibrium_dict["q"]]) + eq_data.psi_norm[0] = min(0, eq_data.psi_norm[0]) + eq_data.psi_norm[-1] = max(1.0, eq_data.psi_norm[-1]) + f_profile = np.array([eq_data.psi_norm, eq_data.f]) + q_profile = np.array([eq_data.psi_norm, eq_data.q]) equilibrium = EFITEquilibrium( - equilibrium_dict["r"], - equilibrium_dict["z"], - equilibrium_dict["psi_grid"], - equilibrium_dict["psi_axis"], - equilibrium_dict["psi_lcfs"], - equilibrium_dict["magnetic_axis"], - equilibrium_dict["x_points"], - equilibrium_dict["strike_points"], + eq_data.r, + eq_data.z, + eq_data.psi_grid, + eq_data.psi_axis, + eq_data.psi_lcfs, + eq_data.magnetic_axis, + eq_data.x_points, + eq_data.strike_points, f_profile, q_profile, - equilibrium_dict["b_vacuum_radius"], - equilibrium_dict["b_vacuum_magnitude"], - equilibrium_dict["lcfs_polygon"], - None, - equilibrium_dict["time"], + eq_data.b_vacuum_radius, + eq_data.b_vacuum_magnitude, + eq_data.lcfs_polygon, + None, # Limiter Polygon + eq_data.time, ) if not with_psi_interpolator: return equilibrium - if equilibrium_dict["rho_tor_norm"] is None: + if eq_data.rho_tor_norm is None: return equilibrium, None psi_interpolator = Interpolator1DArray( - equilibrium_dict["rho_tor_norm"], equilibrium_dict["psi_norm"], "cubic", "none", 0 + eq_data.rho_tor_norm, + eq_data.psi_norm, + "cubic", + "none", + 0, ) return equilibrium, psi_interpolator @@ -154,31 +182,55 @@ def load_magnetic_field( if not len(equilibrium_ids.time_slice): raise RuntimeError("Equilibrium IDS does not have a time slice.") - b_dict = load_magnetic_field_data(equilibrium_ids.time_slice[0].profiles_2d) + b_data = load_magnetic_field_data(equilibrium_ids.time_slice[0].profiles_2d) + + extra_range_r = b_data.r[-1] - b_data.r[0] + extra_range_z = b_data.z[-1] - b_data.z[0] - br = Interpolator2DArray(b_dict["r"], b_dict["z"], b_dict["b_field_r"], "cubic", "none", 0, 0) + br = Interpolator2DArray( + b_data.r, + b_data.z, + b_data.b_field_r, + "cubic", + "linear", + extra_range_r, + extra_range_z, + ) btor = Interpolator2DArray( - b_dict["r"], b_dict["z"], b_dict["b_field_phi"], "cubic", "none", 0, 0 + b_data.r, + b_data.z, + b_data.b_field_phi, + "cubic", + "linear", + extra_range_r, + extra_range_z, + ) + bz = Interpolator2DArray( + b_data.r, + b_data.z, + b_data.b_field_z, + "cubic", + "linear", + extra_range_r, + extra_range_z, ) - bz = Interpolator2DArray(b_dict["r"], b_dict["z"], b_dict["b_field_z"], "cubic", "none", 0, 0) - return FloatToVector3DFunction2D(br, btor, bz) -def cocos_11to3(equilibrium_dict: dict) -> None: +def cocos_11to3(eq_data: Equilibrium2DData) -> None: """Convert from COCOS 11 convention used in IMAS to COCOS 3 convention used in EFIT. - This function modifies the equilibrium_dict in place to convert the coordinates + This function modifies the eq_data in place to convert the coordinates and other relevant data from the COCOS 11 convention to the COCOS 3 convention. Parameters ---------- - equilibrium_dict - The equilibrium data dictionary to modify in place. + eq_data + `Equilibrium2DData` dataclass instance containing equilibrium data to be converted. """ - equilibrium_dict["psi_grid"] = -equilibrium_dict["psi_grid"] / (2.0 * np.pi) - equilibrium_dict["psi_axis"] = -equilibrium_dict["psi_axis"] / (2.0 * np.pi) - equilibrium_dict["psi_lcfs"] = -equilibrium_dict["psi_lcfs"] / (2.0 * np.pi) - equilibrium_dict["q"] = -equilibrium_dict["q"] - if equilibrium_dict["phi"] is not None: - equilibrium_dict["phi"] = -equilibrium_dict["phi"] + eq_data.psi_grid = -eq_data.psi_grid / (2.0 * np.pi) + eq_data.psi_axis = -eq_data.psi_axis / (2.0 * np.pi) + eq_data.psi_lcfs = -eq_data.psi_lcfs / (2.0 * np.pi) + eq_data.q = -eq_data.q + if eq_data.phi is not None: + eq_data.phi = -eq_data.phi diff --git a/tests/plasma/test_blend.py b/tests/plasma/test_blend.py index 9820e32..77e88c2 100644 --- a/tests/plasma/test_blend.py +++ b/tests/plasma/test_blend.py @@ -108,7 +108,7 @@ def test_load_plasma_error_handling(path_iter_jintrac: str): """Test error handling for invalid inputs.""" # Test with non-existent file path with pytest.raises((FileNotFoundError, RuntimeError, OSError)): - load_plasma("non_existent_path", "r") + load_plasma("non_existent_path.nc", "r") # Test with invalid mode with pytest.raises((ValueError, RuntimeError, OSError)): diff --git a/tests/plasma/test_core.py b/tests/plasma/test_core.py index 212564a..81d258f 100644 --- a/tests/plasma/test_core.py +++ b/tests/plasma/test_core.py @@ -101,7 +101,7 @@ def test_load_core_plasma_error_handling(path_iter_jintrac: str): """Test error handling for invalid inputs.""" # Test with non-existent file path with pytest.raises((FileNotFoundError, RuntimeError, OSError)): - load_core_plasma("non_existent_path", "r") + load_core_plasma("non_existent_path.nc", "r") # Test with invalid mode with pytest.raises((ValueError, RuntimeError, OSError)): diff --git a/tests/plasma/test_edge.py b/tests/plasma/test_edge.py index c2915bb..d1d47f1 100644 --- a/tests/plasma/test_edge.py +++ b/tests/plasma/test_edge.py @@ -109,7 +109,7 @@ def test_load_edge_plasma_error_handling(path_iter_jintrac: str): """Test error handling for invalid inputs.""" # Test with non-existent file path with pytest.raises((FileNotFoundError, RuntimeError, OSError)): - load_edge_plasma("non_existent_path", "r") + load_edge_plasma("non_existent_path.nc", "r") # Test with invalid mode with pytest.raises((ValueError, RuntimeError, OSError)): diff --git a/typos.toml b/typos.toml index fdfb275..dac3a3f 100644 --- a/typos.toml +++ b/typos.toml @@ -1,2 +1,3 @@ [default.extend-words] ist = "ist" +arange = "arange"