diff --git a/.dprint.jsonc b/.dprint.jsonc
index 7c5bf4a..83bc718 100644
--- a/.dprint.jsonc
+++ b/.dprint.jsonc
@@ -12,7 +12,7 @@
     "formatComments": true,
     "braceSpacing": false,
   },
-  "excludes": ["**/*.lock", "dist/"],
+  "excludes": ["**/*-lock.json", "**/*.lock", "**/node_modules/", "dist/"],
   "plugins": [
     "https://plugins.dprint.dev/json-0.20.0.wasm",
     "https://plugins.dprint.dev/markdown-0.19.0.wasm",
diff --git a/.gitignore b/.gitignore
index cc5d10b..3778b9b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -220,5 +220,8 @@ __marimo__/
 # Streamlit
 .streamlit/secrets.toml
 
+# Type Stubs
+typings/
+
 # Plots
 demos/ITER/plots/
\ No newline at end of file
diff --git a/.lefthook.yaml b/.lefthook.yaml
index 88cec02..2aeb41b 100644
--- a/.lefthook.yaml
+++ b/.lefthook.yaml
@@ -1,7 +1,7 @@
-lefthook: pixi run --environment=lint lefthook
+lefthook: pixi run --no-progress --environment=lint --frozen lefthook
 
 templates:
-  run: run --environment=lint
+  run: run --quiet --no-progress --environment=lint --frozen
 
 colors: true
 
@@ -65,18 +65,6 @@ pre-commit:
       glob: "pyproject.toml"
       run: pixi {run} validate-pyproject {staged_files}
 
-    - name: docstring-lint
-      stage_fixed: true
-      group:
-        piped: true
-        jobs:
-          - name: docformatter
-            glob: "*.{py,pyx,pyi}"
-            run: pixi {run} docformatter --in-place {staged_files}
-          - name: numpydoc
-            glob: "*.{py,pyi}"
-            run: pixi {run} numpydoc {staged_files}
-
     - name: cython-lint
       glob: "*.{pyx,pxd}"
       run: pixi {run} cython-lint {staged_files}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c80fb6f..05b7371 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,8 +2,30 @@
 
 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).
+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.2.1] - 2025-11-18
+
+### Added
+
+- Add some missing type hints.
+
+### Changed
+
+- Migrate docstring linting/formatting from `numpydoc` and `docformatter` to `ruff`
+- Migrate `numpydoc` Python API reference to `napoleon` extension for Sphinx (to enjoy type hinting support)
+- Update docstrings to be compatible with the `napoleon` style
+
+### Fixed
+
+- Fix values assignment in `load_equilibrium_data` function (convert to python types)
+- Bug fix in `load_unstruct_grid_2d` function (incorrect `cells.append(cell)` line)
+
+### Removed
+
+- Remove `numpydoc` dependency (migrated to `ruff` for docstring linting as well)
+- Remove `docformatter` dependency (migrated to `ruff` for docstring formatting as well)
 
 ## [0.2.0] - 2025-11-04
 
diff --git a/demos/ITER/camera_geometry_matrix_gif.py b/demos/ITER/camera_geometry_matrix_gif.py
index 305692a..4fc43d4 100644
--- a/demos/ITER/camera_geometry_matrix_gif.py
+++ b/demos/ITER/camera_geometry_matrix_gif.py
@@ -62,7 +62,7 @@
 unsaturated_fraction = 0.95
 
 
-def read_geometry_matrix_without_reflections(
+def read_geometry_matrix_without_reflections(  # noqa: D103
     shot, run, user, database="ITER_MD", backend=HDF5_BACKEND
 ):
     geometry_matrix = {}
@@ -96,7 +96,7 @@ def read_geometry_matrix_without_reflections(
     return geometry_matrix
 
 
-def read_geometry_matrix_interpolated(shot, run, user, database="ITER_MD", backend=HDF5_BACKEND):
+def read_geometry_matrix_interpolated(shot, run, user, database="ITER_MD", backend=HDF5_BACKEND):  # noqa: D103
     geometry_matrix = {}
 
     with DBEntry(backend, database, shot, run, user) as entry:
@@ -235,7 +235,7 @@ def read_geometry_matrix_interpolated(shot, run, user, database="ITER_MD", backe
         power = volume[index] * line_emission[il, it][index]
         points = np.array([r2d[index], z2d[index]]).T
 
-        def f(task, node_points, interp_data, points, power):
+        def _f(task, node_points, interp_data, points, power):
             ibegin = task[0]
             iend = task[1]
             interp_matrix = griddata(
@@ -252,7 +252,7 @@ def f(task, node_points, interp_data, points, power):
 
         with ctx.Pool(processes=NPROC) as pool:
             results = pool.map(
-                f,
+                _f,
                 tasks,
                 node_points=node_points,
                 interp_data=interp_data,
diff --git a/demos/ITER/core_plasma_plot_profiles.py b/demos/ITER/core_plasma_plot_profiles.py
index 1489dda..e4a05ea 100644
--- a/demos/ITER/core_plasma_plot_profiles.py
+++ b/demos/ITER/core_plasma_plot_profiles.py
@@ -10,6 +10,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.colors import SymLogNorm
+from matplotlib.figure import Figure
 
 from cherab.core.math import sample3d_grid, samplevector2d
 from cherab.imas.datasets import iter_jintrac
@@ -19,9 +20,26 @@
 plt.ion()
 
 
-def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False):
-    """Make a 2D plot of quantity, with a title, optionally on a log scale."""
-
+def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False) -> Figure:
+    """Make a 2D plot of quantity, with a title, optionally on a log scale.
+
+    Parameters
+    ----------
+    quantity
+        2D array of the quantity to plot.
+    extent
+        The extent of the plot in the form [xmin, xmax, ymin, ymax].
+    title
+        The title of the plot.
+    logscale
+        Whether to use a logarithmic scale for the color map.
+    symmetric
+        Whether to use a symmetric color map around zero.
+
+    Returns
+    -------
+    The matplotlib Figure object containing the plot.
+    """
     fig = plt.figure(figsize=(4.0, 6.0), layout="constrained")
     ax = fig.add_subplot(111)
     if logscale:
diff --git a/demos/ITER/edge_plasma_mesh_and_values.py b/demos/ITER/edge_plasma_mesh_and_values.py
index 0cdccbe..7dc2b43 100644
--- a/demos/ITER/edge_plasma_mesh_and_values.py
+++ b/demos/ITER/edge_plasma_mesh_and_values.py
@@ -10,6 +10,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.colors import SymLogNorm
+from matplotlib.figure import Figure
 
 from cherab.imas.datasets import iter_jintrac
 from cherab.imas.ids.common import get_ids_time_slice
@@ -21,7 +22,26 @@
 grid_subset_name = "Cells"
 
 
-def plot_grid_quantity(grid, quantity, title="", logscale=False, symmetric=False):
+def plot_grid_quantity(grid, quantity, title="", logscale=False, symmetric=False) -> Figure:
+    """Make a 2D plot of a grid quantity, with a title, optionally on a log scale.
+
+    Parameters
+    ----------
+    grid
+        The grid object.
+    quantity
+        1D array of the quantity to plot on the grid.
+    title
+        The title of the plot.
+    logscale
+        Whether to use a logarithmic scale for the color map.
+    symmetric
+        Whether to use a symmetric color map around zero.
+
+    Returns
+    -------
+    The matplotlib Figure object containing the plot.
+    """
     ax = grid.plot_mesh(data=quantity)
 
     if logscale:
diff --git a/demos/ITER/edge_plasma_plot_profiles.py b/demos/ITER/edge_plasma_plot_profiles.py
index d648ff1..1dda2ab 100644
--- a/demos/ITER/edge_plasma_plot_profiles.py
+++ b/demos/ITER/edge_plasma_plot_profiles.py
@@ -9,6 +9,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.colors import SymLogNorm
+from matplotlib.figure import Figure
 
 from cherab.core.math import sample3d, sample3d_grid, samplevector2d, samplevector3d_grid
 from cherab.imas.datasets import iter_jintrac
@@ -18,9 +19,26 @@
 plt.ion()
 
 
-def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False):
-    """Make a 2D plot of quantity, with a title, optionally on a log scale."""
-
+def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False) -> Figure:
+    """Make a 2D plot of quantity, with a title, optionally on a log scale.
+
+    Parameters
+    ----------
+    quantity
+        2D array of the quantity to plot.
+    extent
+        The extent of the plot in the form [xmin, xmax, ymin, ymax].
+    title
+        The title of the plot.
+    logscale
+        Whether to use a logarithmic scale for the color map.
+    symmetric
+        Whether to use a symmetric color map around zero.
+
+    Returns
+    -------
+    The matplotlib Figure object containing the plot.
+    """
     fig = plt.figure(figsize=(4.0, 6.0), layout="constrained")
     ax = fig.add_subplot(111)
     if logscale:
diff --git a/demos/ITER/full_plasma_plot_profiles.py b/demos/ITER/full_plasma_plot_profiles.py
index c017028..6e802ce 100644
--- a/demos/ITER/full_plasma_plot_profiles.py
+++ b/demos/ITER/full_plasma_plot_profiles.py
@@ -10,6 +10,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.colors import SymLogNorm
+from matplotlib.figure import Figure
 
 from cherab.core.math import sample3d, sample3d_grid, samplevector2d, samplevector3d_grid
 from cherab.imas.datasets import iter_jintrac
@@ -19,9 +20,26 @@
 plt.ion()
 
 
-def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False):
-    """Make a 2D plot of quantity, with a title, optionally on a log scale."""
-
+def plot_quantity(quantity, extent, title="", logscale=False, symmetric=False) -> Figure:
+    """Make a 2D plot of quantity, with a title, optionally on a log scale.
+
+    Parameters
+    ----------
+    quantity
+        2D array of the quantity to plot.
+    extent
+        The extent of the plot in the form [xmin, xmax, ymin, ymax].
+    title
+        The title of the plot.
+    logscale
+        Whether to use a logarithmic scale for the color map.
+    symmetric
+        Whether to use a symmetric color map around zero.
+
+    Returns
+    -------
+    The matplotlib figure object.
+    """
     fig = plt.figure(figsize=(4.0, 6.0), layout="constrained")
     ax = fig.add_subplot(111)
     if logscale:
diff --git a/docs/notebooks/plasma/full_plasma.ipynb b/docs/notebooks/plasma/full_plasma.ipynb
index 9afc426..1fa4cc1 100644
--- a/docs/notebooks/plasma/full_plasma.ipynb
+++ b/docs/notebooks/plasma/full_plasma.ipynb
@@ -393,7 +393,7 @@
     "            title=\"Density of \" + label,\n",
     "            clabel=\"[1/m³]\",\n",
     "            logscale=True,\n",
-    "        )\n"
+    "        )"
    ]
   }
  ],
diff --git a/docs/source/_templates/class.rst b/docs/source/_templates/class.rst
index 24e9d54..b1bafd7 100644
--- a/docs/source/_templates/class.rst
+++ b/docs/source/_templates/class.rst
@@ -5,31 +5,29 @@
 .. autoclass:: {{ objname }}
    :show-inheritance:
    :members:
+   :special-members: __call__, __getitem__
    :inherited-members:
-
-   {% block methods %}
-   {% if all_methods %}
+{% block methods %}
+{% if all_methods %}
    .. rubric:: {{ _('Methods') }}
 
    .. autosummary::
       :template: method.rst
-
-   {% for item in all_methods %}
-   {%- if not item.startswith('_') or item in ['__call__', '__getitem__'] %}
+{% for item in all_methods %}
+{%- if not item.startswith('_') or item in ['__call__', '__getitem__'] %}
       ~{{ name }}.{{ item }}
-   {% endif %}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
-
-   {% block attributes %}
-   {% if attributes %}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{% endblock %}
+{% block attributes %}
+{% if attributes %}
    .. rubric:: {{ _('Attributes') }}
 
    .. autosummary::
       :template: attribute.rst
-   {% for item in attributes %}
+{% for item in attributes %}
       ~{{ name }}.{{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
+{%- endfor %}
+{%- endif %}
+{% endblock %}
diff --git a/docs/source/_templates/module.rst b/docs/source/_templates/module.rst
index ee4ea73..9a10035 100644
--- a/docs/source/_templates/module.rst
+++ b/docs/source/_templates/module.rst
@@ -1,55 +1,54 @@
 {{ fullname | escape | underline}}
 
 .. automodule:: {{ fullname }}
-
-   {% block attributes %}
-   {% if attributes %}
+{% block attributes %}
+{% if attributes %}
    .. rubric:: {{ _('Module Attributes') }}
 
    .. autosummary::
       :toctree:
-   {% for item in attributes %}
+{% for item in attributes %}
       {{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
+{%- endfor %}
+{%- endif %}
+{%- endblock %}
 
-   {% block functions %}
-   {% if functions %}
+{%- block functions %}
+{%- if functions %}
    .. rubric:: {{ _('Functions') }}
 
    .. autosummary::
       :toctree:
-   {% for item in functions %}
+{% for item in functions %}
       {{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
+{%- endfor %}
+{%- endif %}
+{%- endblock %}
 
-   {% block classes %}
-   {% if classes %}
+{%- block classes %}
+{%- if classes %}
    .. rubric:: {{ _('Classes') }}
 
    .. autosummary::
       :toctree:
       :template: class.rst
-   {% for item in classes %}
+{% for item in classes %}
       {{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
+{%- endfor %}
+{%- endif %}
+{%- endblock %}
 
-   {% block exceptions %}
-   {% if exceptions %}
+{%- block exceptions %}
+{%- if exceptions %}
    .. rubric:: {{ _('Exceptions') }}
 
    .. autosummary::
       :toctree:
-   {% for item in exceptions %}
+{% for item in exceptions %}
       {{ item }}
-   {%- endfor %}
-   {% endif %}
-   {% endblock %}
+{%- endfor %}
+{%- endif %}
+{%- endblock %}
 
 {% block modules %}
 {% if modules %}
@@ -63,5 +62,5 @@
 {% for item in modules %}
    {{ item }}
 {%- endfor %}
-{% endif %}
+{%- endif %}
 {% endblock %}
diff --git a/docs/source/conf.py b/docs/source/conf.py
index a406b96..d560367 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -20,6 +20,8 @@
     "sphinx.ext.intersphinx",
     "sphinx.ext.mathjax",
     "sphinx.ext.todo",
+    "sphinx.ext.napoleon",
+    "sphinx_api_relink",
     "sphinx_copybutton",
     "sphinx_codeautolink",
     "sphinx_design",
@@ -27,7 +29,6 @@
     "sphinx_immaterial",
     "IPython.sphinxext.ipython_console_highlighting",
     "myst_parser",
-    "numpydoc",
     "nbsphinx",
 ]
 
@@ -35,6 +36,10 @@
 
 # autodoc config
 autodoc_member_order = "bysource"
+autodoc_typehints = "signature"
+autodoc_type_aliases = {
+    "ArrayLike": "ArrayLike",
+}
 
 # autosummary config
 autosummary_generate = True
@@ -42,9 +47,10 @@
 autosummary_imported_members = True
 autosummary_ignore_module_all = False
 
-# numpydoc config
-numpydoc_show_class_members = False
-numpydoc_xref_param_type = True
+# napoleon config
+napoleon_google_docstring = False
+napoleon_numpy_docstring = True
+napoleon_attr_annotations = True
 
 # todo config
 todo_include_todos = True
@@ -117,8 +123,8 @@
         "toc.sticky",
         "content.tabs.link",
         # "content.code.copy",
-        "content.action.edit",
-        "content.action.view",
+        # "content.action.edit",
+        # "content.action.view",
         "content.tooltips",
         "announce.dismiss",
     ],
@@ -129,7 +135,6 @@
             "media": "(prefers-color-scheme: light)",
             "scheme": "default",
             "primary": "indigo",
-            "accent": "green",
             "toggle": {
                 "icon": "material/lightbulb-outline",
                 "name": "Switch to dark mode",
@@ -138,8 +143,7 @@
         {
             "media": "(prefers-color-scheme: dark)",
             "scheme": "slate",
-            "primary": "light-blue",
-            "accent": "lime",
+            "primary": "indigo",
             "toggle": {
                 "icon": "material/lightbulb",
                 "name": "Switch to light mode",
@@ -162,6 +166,8 @@
     "raysect": ("http://www.raysect.org", None),
     "cherab": ("https://www.cherab.info", None),
     "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),
 }
 
 intersphinx_timeout = 10
@@ -172,6 +178,7 @@
 linkcode_link_text = "Source"
 
 # -- NBSphinx configuration ---------------------------------------------------
+# nbsphinx_execute = "never"
 nbsphinx_prolog = r"""
 {% set docname = 'docs/' + env.doc2path(env.docname, base=None)|string %}
 
diff --git a/pixi.toml b/pixi.toml
index 0299f32..4c0b016 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -30,6 +30,7 @@ cherab = "1.5.*"
 
 [package.run-dependencies]
 imas-python = "*"
+typing-extensions = ">=4.5"
 
 # ---------------------------------
 # === Development Configuration ===
@@ -60,7 +61,6 @@ sphinx-design = "*"
 sphinx-immaterial = "*"
 linkify-it-py = "*"
 myst-parser = "*"
-numpydoc = "*"
 # For notebook support
 ipykernel = "*"
 ipywidgets = "*"
@@ -69,6 +69,7 @@ pandoc = "*"
 
 [feature.docs.pypi-dependencies]
 sphinx-github-style = "*"
+sphinx-api-relink = "*"
 
 [feature.docs.tasks]
 doc-build = { cmd = [
@@ -105,12 +106,10 @@ ruff = "*"
 typos = "*"
 mypy = "*"
 basedpyright = "*"
-numpydoc = "*"
 actionlint = "*"
 shellcheck = "*"
 validate-pyproject = "*"
 cython-lint = "*"
-docformatter = "*"
 blacken-docs = "*"
 taplo = "*"
 
@@ -128,8 +127,6 @@ taplo = { cmd = "taplo fmt", description = "Format toml files with taplo" }
 actionlint = { cmd = "actionlint", description = "Lint actions with actionlint" }
 blacken-docs = { cmd = "blacken-docs", description = "Format Python markdown blocks with Black" }
 validate-pyproject = { cmd = "validate-pyproject pyproject.toml", description = "Validate pyproject.toml" }
-numpydoc = { cmd = "numpydoc lint", description = "Validate docstrings with numpydoc" }
-docformatter = { cmd = "docformatter --in-place", description = "Format docstrings with docformatter" }
 cython-lint = { cmd = "cython-lint", description = "Lint Cython files" }
 lint = { cmd = "lefthook run pre-commit --all-files --force", description = "🧹 Run all linters" }
 
diff --git a/pyproject.toml b/pyproject.toml
index ff9e8c4..ccde31c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,11 @@ classifiers = [
   "Topic :: Software Development :: Libraries",
   "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ["cherab==1.5.*", "imas-python[netcdf]"]
+dependencies = [
+  "cherab==1.5.*",
+  "imas-python[netcdf]",
+  "typing-extensions; python_version < '3.12'",
+]
 dynamic = ["version"]
 
 [project.optional-dependencies]
@@ -124,19 +128,31 @@ select = [
   "W",   # Warning
   "UP",  # pyupgrade
   "NPY", # numpy specific rules
+  "D",   # pydocstyle
+  "DOC", # pydoclint
 ]
+preview = true
 ignore = [
   # Recommended ignores by ruff when using formatter
   "E501", # line too long
   "N803", # argument name should be lowercase
   "N806", # variable in function should be lowercase
+  "D107", # missing docstring in __init__
 ]
 
 [tool.ruff.lint.isort]
 known-first-party = ["cherab"]
 
 [tool.ruff.lint.per-file-ignores]
-"**/*.ipynb" = ["W292"] # no newline at end of file
+"**/*.ipynb" = [
+  "W292",
+  "D",
+  "DOC",
+] # no newline at end of file and skip docstring checks
+"tests/**" = ["D", "DOC"] # skip docstring checks in tests
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
 
 [tool.ruff.format]
 docstring-code-format = true
@@ -144,30 +160,6 @@ docstring-code-format = true
 [tool.cython-lint]
 max-line-length = 100
 
-[tool.docformatter]
-recursive = true
-wrap-summaries = 100
-wrap-descriptions = 100
-blank = false
-
-[tool.numpydoc_validation]
-checks = [
-  "all",  # report on all checks, except the below
-  "GL01", # Docstring text (summary) should start in the line immediately
-  "EX01", # No examples section found
-  "SA01", # See Also section not found
-  "ES01", # No extended summary found
-  "GL08", # The object does not have a docstring
-  "RT01", # No Returns section found
-  "PR01", # Parameters {missing_params} not documented
-  "SS06", # Summary should fit in a single line
-]
-# remember to use single quotes for regex in TOML
-exclude = [ # don't report on objects that match any of these regex
-  '\.__repr__$',
-]
-override_SS05 = ['^Process ', '^Assess ', '^Access ']
-
 # ------------------------------
 # === Type checking settings ===
 # ------------------------------
@@ -177,6 +169,7 @@ 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"]
@@ -187,6 +180,7 @@ reportMissingTypeStubs = true
 # no raysect/cherab type stubs; pytest fixtures
 reportUnknownMemberType = false
 reportUnknownVariableType = false
+reportUnknownParameterType = false
 
 # ruff handles this
 reportUnusedParameter = false
diff --git a/src/cherab/imas/__init__.py b/src/cherab/imas/__init__.py
index ba7c205..37746bf 100644
--- a/src/cherab/imas/__init__.py
+++ b/src/cherab/imas/__init__.py
@@ -19,6 +19,6 @@
 
 from __future__ import annotations
 
-from importlib.metadata import version
+from importlib.metadata import version as _version
 
-__version__ = version("cherab-imas")
+__version__ = _version("cherab-imas")
diff --git a/src/cherab/imas/datasets/__init__.py b/src/cherab/imas/datasets/__init__.py
index f330db5..3f7b9a2 100644
--- a/src/cherab/imas/datasets/__init__.py
+++ b/src/cherab/imas/datasets/__init__.py
@@ -1,4 +1,4 @@
-"""Sample Dataset utilities for fetching and processing data.
+r"""Sample Dataset utilities for fetching and processing data.
 
 Usage of Datasets
 =================
@@ -37,7 +37,7 @@
 
 For Windows::
 
-    "C:\\Users\\<user>\\AppData\\Local\\<AppAuthor>\\cherab\\imas\\Cache"
+    "C:\Users\<user>\AppData\Local\<AppAuthor>\cherab\imas\Cache"
 
 In environments with constrained network connectivity for various security
 reasons or on systems without continuous internet connections, one may manually
diff --git a/src/cherab/imas/datasets/_fetchers.py b/src/cherab/imas/datasets/_fetchers.py
index 51fbc29..92eeccd 100644
--- a/src/cherab/imas/datasets/_fetchers.py
+++ b/src/cherab/imas/datasets/_fetchers.py
@@ -1,4 +1,4 @@
-"""This module provides functionality to fetch IMAS sample datasets."""
+"""Provide functionality to fetch IMAS sample datasets."""
 
 from ...imas import __version__
 from ._registry import registry
@@ -17,7 +17,7 @@
     )
 
 
-def fetch_data(dataset_name: str, data_fetcher=data_fetcher, show_progress=True) -> str:
+def fetch_data(dataset_name: str, data_fetcher=data_fetcher, show_progress: bool = True) -> str:
     if data_fetcher is None:
         raise ImportError(
             "Missing optional dependency 'pooch' required for cherab.imas.datasets module. "
@@ -46,7 +46,7 @@ def fetch_data(dataset_name: str, data_fetcher=data_fetcher, show_progress=True)
 
 
 def iter_jintrac() -> str:
-    """Fetche and return the path to the ITER JINTRAC sample dataset.
+    """Fetch and return the path to the ITER JINTRAC sample dataset.
 
     Returns
     -------
@@ -64,7 +64,7 @@ def iter_jintrac() -> str:
 
 
 def iter_solps() -> str:
-    """Fetche and return the path to the ITER SOLPS sample dataset.
+    """Fetch and return the path to the ITER SOLPS sample dataset.
 
     Returns
     -------
@@ -82,7 +82,7 @@ def iter_solps() -> str:
 
 
 def iter_jorek() -> str:
-    """Fetche and return the path to the ITER JOREK sample dataset.
+    """Fetch and return the path to the ITER JOREK sample dataset.
 
     Returns
     -------
diff --git a/src/cherab/imas/datasets/_progress.py b/src/cherab/imas/datasets/_progress.py
index b711ddf..148b497 100644
--- a/src/cherab/imas/datasets/_progress.py
+++ b/src/cherab/imas/datasets/_progress.py
@@ -1,4 +1,4 @@
-"""This module provides a rich progress bar for dataset downloads."""
+"""Progress bar implementation for pooch using rich."""
 
 try:
     from rich.progress import (
diff --git a/src/cherab/imas/datasets/_registry.py b/src/cherab/imas/datasets/_registry.py
index c04458c..d1fb3df 100644
--- a/src/cherab/imas/datasets/_registry.py
+++ b/src/cherab/imas/datasets/_registry.py
@@ -1,4 +1,4 @@
-"""This file serves as the IMAS sample dataset registry."""
+"""Serve as the IMAS sample dataset registry."""
 
 registry = {
     "iter_disruption_113112_1.nc": "md5:f8747539b8d5a4b6974cc1bb33f2924a",  # JOREK 3D simulation
diff --git a/src/cherab/imas/datasets/_utils.py b/src/cherab/imas/datasets/_utils.py
index 6afb8e9..f3dcc82 100644
--- a/src/cherab/imas/datasets/_utils.py
+++ b/src/cherab/imas/datasets/_utils.py
@@ -82,7 +82,7 @@ def clear_cache(
 
     Parameters
     ----------
-    datasets : Callable[..., str] | Collection[Callable[..., str]], optional
+    datasets
         `cherab.imas.datasets` method or a list/tuple of the same whose cached data files are to be
         removed. If `None`, all the cached data files are removed.
 
diff --git a/src/cherab/imas/ggd/base_mesh.py b/src/cherab/imas/ggd/base_mesh.py
index 5bdcd4b..8ba9925 100755
--- a/src/cherab/imas/ggd/base_mesh.py
+++ b/src/cherab/imas/ggd/base_mesh.py
@@ -17,7 +17,17 @@
 # under the Licence.
 """Module defining the base class for general grids (GGD)."""
 
-from raysect.core.math import Vector3D
+from __future__ import annotations
+
+from typing import Literal
+
+import matplotlib.axes
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from raysect.core.math.function.float import Function2D, Function3D
+from raysect.core.math.function.vector3d.function2d import Function2D as VectorFunction2D
+from raysect.core.math.function.vector3d.function3d import Function3D as VectorFunction3D
+from raysect.core.math.vector import Vector3D
 
 ZEROVECTOR = Vector3D(0, 0, 0)
 
@@ -30,96 +40,107 @@ class GGDGrid:
 
     Parameters
     ----------
-    name : str
-        Name of the grid, by default ''.
-    dimension : int
+    name
+        Name of the grid.
+    dimension
         Grid dimensions, by default 1.
-    coordinate_system : str
-        Coordinate system, by default 'cartesian'.
+    coordinate_system
+        Coordinate system, by default ``"cartesian"``.
     """
 
-    def __init__(self, name="", dimension=1, coordinate_system="cartesian"):
+    def __init__(
+        self,
+        name: str = "",
+        dimension: int = 1,
+        coordinate_system: Literal["cylindrical", "cartesian"] = "cartesian",
+    ):
         dimension = int(dimension)
         if dimension < 1:
             raise ValueError("Attribute dimension must be >= 1.")
 
-        self._dimension = dimension
-        self._name = str(name)
-        self._coordinate_system = str(coordinate_system)
+        self._dimension: int = dimension
+        self._name: str = str(name)
+        self._coordinate_system: str = str(coordinate_system)
 
-        self._interpolator = None
-        self._cell_centre = None
-        self._cell_area = None
-        self._cell_volume = None
-        self._mesh_extent = None
-        self._num_cell = 0
+        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._initial_setup()
 
-    def _initial_setup(self):
+    def _initial_setup(self) -> None:
         raise NotImplementedError("To be defined in subclass.")
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Grid name."""
         return self._name
 
     @name.setter
-    def name(self, value):
+    def name(self, value: str):
         self._name = str(value)
 
     @property
-    def dimension(self):
+    def dimension(self) -> int:
         """Grid dimension."""
         return self._dimension
 
     @property
-    def num_cell(self):
+    def num_cell(self) -> int:
         """Number of grid cells."""
         return self._num_cell
 
     @property
-    def coordinate_system(self):
+    def coordinate_system(self) -> str:
         """Coordinate system."""
         return self._coordinate_system
 
     @property
-    def cell_centre(self):
-        """Coordinate of cell centres as (num_cell, dimension) array."""
+    def cell_centre(self) -> NDArray[np.float64] | None:
+        """Coordinate of cell centres as ``(num_cell, dimension)`` array."""
         return self._cell_centre
 
     @property
-    def cell_area(self):
-        """Cell areas as (num_cell,) array."""
+    def cell_area(self) -> NDArray[np.float64] | None:
+        """Cell areas as ``(num_cell,)`` array."""
         return self._cell_area
 
     @property
-    def cell_volume(self):
-        """Cell volume as (num_cell,) array."""
+    def cell_volume(self) -> NDArray[np.float64] | None:
+        """Cell volume as ``(num_cell,)`` array."""
         return self._cell_volume
 
     @property
-    def mesh_extent(self):
+    def mesh_extent(self) -> dict[str, float] | None:
         """Extent of the mesh.
 
         A dictionary with xmin, xmax, ymin and ymax, ... keys.
         """
         return self._mesh_extent
 
-    def subset(self, indices, name=None):
+    def subset(self, indices: ArrayLike, name: str | None = None) -> GGDGrid:
         """Create a subset grid from this instance.
 
         Parameters
         ----------
-        indices : array_like
+        indices
             Indices of the cells of the original grid in the subset.
-        name : str, optional
-            Name of the grid subset. Default is instance.name + ' subset'.
-        """
+        name
+            Name of the grid subset. Default is ``instance.name + " subset"``.
 
+        Returns
+        -------
+        GGDGrid
+            Subset grid instance.
+        """
         raise NotImplementedError("To be defined in subclass.")
 
-    def interpolator(self, grid_data, fill_value=0.0):
+    def interpolator(
+        self, grid_data: ArrayLike, fill_value: float = 0.0
+    ) -> Function2D | Function3D:
         """Return an Function interpolator instance for the data defined on this grid.
 
         On the second and subsequent calls, the interpolator is created as an instance of the
@@ -127,20 +148,21 @@ def interpolator(self, grid_data, fill_value=0.0):
 
         Parameters
         ----------
-        grid_data : array_like
+        grid_data
             Array containing data in the grid cells.
-        fill_value : float, optional
+        fill_value
             A value returned outside the grid, by default is 0.0.
 
         Returns
         -------
-        Function2D | Function3D
+        `Function2D` | `Function3D`
             Interpolator instance.
         """
-
         raise NotImplementedError("To be defined in subclass.")
 
-    def vector_interpolator(self, grid_vectors, fill_vector=ZEROVECTOR):
+    def vector_interpolator(
+        self, grid_vectors: ArrayLike, fill_vector: Vector3D = ZEROVECTOR
+    ) -> VectorFunction2D | VectorFunction3D:
         """Return a VectorFunction interpolator instance for the vector data defined on this grid.
 
         On the second and subsequent calls, the interpolator is created as an instance of the
@@ -148,28 +170,26 @@ def vector_interpolator(self, grid_vectors, fill_vector=ZEROVECTOR):
 
         Parameters
         ----------
-        grid_vectors : (3, num_cell) array_like
-            Array containing 3D vectors in the grid cells.
-        fill_vector : Vector3D, optional
-            3D vector returned outside the grid, by default Vector3D(0, 0, 0).
+        grid_vectors
+            ``(3, num_cell)`` Array containing 3D vectors in the grid cells.
+        fill_vector
+            3D vector returned outside the grid, by default ``Vector3D(0, 0, 0)``.
 
         Returns
         -------
-        VectorFunction2D | VectorFunction3D
+        `VectorFunction2D` | `VectorFunction3D`
             Interpolator instance.
         """
-
         raise NotImplementedError("To be defined in subclass.")
 
-    def plot_mesh(self, data=None, ax=None):
+    def plot_mesh(self, data: ArrayLike | None = None, ax: matplotlib.axes.Axes | None = None):
         """Plot the grid geometry to a matplotlib figure.
 
         Parameters
         ----------
-        data : array_like, optional
+        data
             Data array defined on the grid.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot on. If None, a new figure and axes are created.
         """
-
         raise NotImplementedError("To be defined in subclass.")
diff --git a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py
index fc651eb..73cb72a 100644
--- a/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py
+++ b/src/cherab/imas/ggd/unstruct_2d_extend_mesh.py
@@ -17,13 +17,23 @@
 # under the Licence.
 """Module defining unstructured 2D-extended mesh class and related methods."""
 
+from __future__ import annotations
+
+import sys
 from typing import Literal
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override  # pyright: ignore[reportUnreachable]
+
+import matplotlib.axes
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.collections import PolyCollection
 from matplotlib.tri import Triangulation
-from raysect.core.math import Vector3D
+from numpy.typing import ArrayLike, NDArray
+from raysect.core.math.vector import Vector3D
 
 from ..math import UnstructGridFunction3D, UnstructGridVectorFunction3D
 from ..math.tetrahedralize import calculate_tetra_volume, cell_to_5tetra
@@ -43,27 +53,29 @@ class UnstructGrid2DExtended(GGDGrid):
 
     Parameters
     ----------
-    vertices : (N,3) array_like
-        Array containing coordinates of the cell vertices in the (X, Y, Z) space.
-    cells : (M,8) array_like
-        Array containing the vertex indices in clockwise or counterclockwise order for each cubic
-        cell in the list (the starting vertex must not be included twice).
-    num_faces : int
+    vertices
+        Array-like of shape ``(N, 3)`` containing coordinates of the cell vertices in the (X, Y, Z)
+        space.
+    cells
+        Array-like of shape ``(M, 8)`` containing the vertex indices in clockwise or
+        counterclockwise order for each cubic cell in the list
+        (the starting vertex must not be included twice).
+    num_faces
         Number of faces at the poloidal plane.
-    num_poloidal : int
+    num_poloidal
         Number of poloidal points.
-    num_toroidal : int
+    num_toroidal
         Number of toroidal points.
-    name : str
-        Name of the grid, by default 'Cells'.
-    coordinate_system : {'cylindrical', 'cartesian'}, optional
-        Coordinate system of the grid, by default 'cylindrical'.
+    name
+        Name of the grid, by default ``"Cells"``.
+    coordinate_system
+        Coordinate system of the grid, by default ``"cylindrical"``.
     """
 
     def __init__(
         self,
-        vertices,
-        cells,
+        vertices: ArrayLike,
+        cells: ArrayLike,
         num_faces: int,
         num_poloidal: int,
         num_toroidal: int,
@@ -99,18 +111,19 @@ def __init__(
                 + f"The shape of 'cells' is {cells.shape}."
             )
 
-        self._vertices = vertices
-        self._cells = cells
-        self._num_faces = num_faces
-        self._num_poloidal = num_poloidal
-        self._num_toroidal = num_toroidal
+        self._vertices: NDArray[np.float64] = vertices
+        self._cells: NDArray[np.int32] = cells
+        self._num_faces: int = num_faces
+        self._num_poloidal: int = num_poloidal
+        self._num_toroidal: int = num_toroidal
 
         super().__init__(name, 3, coordinate_system)
 
+    @override
     def _initial_setup(self):
         self._interpolator = None
 
-        self._num_cell = self._cells.shape[0]
+        self._num_cell: int = self._cells.shape[0]
 
         # Extract grid points at the first poloidal plane
         r = self._vertices[: self._num_poloidal, 0]
@@ -155,43 +168,43 @@ def _initial_setup(self):
         self._cell_centre.setflags(write=False)
 
     @property
-    def num_poloidal(self):
+    def num_poloidal(self) -> int:
         """Number of poloidal grid points."""
         return self._num_poloidal
 
     @property
-    def num_toroidal(self):
+    def num_toroidal(self) -> int:
         """Number of toroidal grid points."""
         return self._num_toroidal
 
     @property
-    def num_faces(self):
+    def num_faces(self) -> int:
         """Number of faces at the poloidal plane."""
         return self._num_faces
 
     @property
-    def vertices(self):
-        """Mesh vertex coordinates as (N, 3) array in the (X, Y, Z) space."""
+    def vertices(self) -> NDArray[np.float64]:
+        """Mesh vertex coordinates as ``(N, 3)`` array in the (X, Y, Z) space."""
         return self._vertices
 
     @property
-    def cells(self):
-        """Mesh cells as (M, 8) array."""
+    def cells(self) -> NDArray[np.int32]:
+        """Mesh cells as ``(M, 8)`` array."""
         return self._cells
 
     @property
-    def tetrahedra(self):
-        """Mesh tetrahedra as (5M, 4) array."""
+    def tetrahedra(self) -> NDArray[np.int32]:
+        """Mesh tetrahedra as ``(5M, 4)`` array."""
         return self._tetrahedra
 
     @property
-    def tetra_to_cell_map(self):
-        """Array of shape (5M,) mapping every tetrahedral index to a grid cell ID."""
+    def tetra_to_cell_map(self) -> NDArray[np.int32]:
+        """Array of shape ``(5M,)`` mapping every tetrahedral index to a grid cell ID."""
         return self._tetra_to_cell_map
 
     @property
-    def cell_to_tetra_map(self):
-        """Array of shape (M, 2) mapping every grid cell index to tetrahedral IDs.
+    def cell_to_tetra_map(self) -> NDArray[np.int32]:
+        """Array of shape ``(M, 2)`` mapping every grid cell index to tetrahedral IDs.
 
         The first column is the index of the first tetrahedron forming the cell.
         The second column is the number of tetrahedra forming the cell.
@@ -201,7 +214,7 @@ def cell_to_tetra_map(self):
         """
         return self._cell_to_tetra_map
 
-    def subset_faces(self, indices, name: str | None = None):
+    def subset_faces(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2DExtended:
         """Create a subset UnstructGrid2DExtended from this instance.
 
         The subset is defined by the indices of the faces at the poloidal plane.
@@ -213,10 +226,20 @@ def subset_faces(self, indices, name: str | None = None):
 
         Parameters
         ----------
-        indices : array_like
+        indices
             Indices of the faces of the original grid in the subset.
-        name : str, optional
-            Name of the grid subset. Default is `instance.name` + `' subset'`.
+        name
+            Name of the grid subset. Default is ``instance.name + " subset"``.
+
+        Returns
+        -------
+        `.UnstructGrid2DExtended`
+            Subset instance.
+
+        Raises
+        ------
+        ValueError
+            If any of the indices of the faces is out of range.
         """
         grid = UnstructGrid2DExtended.__new__(UnstructGrid2DExtended)
 
@@ -295,7 +318,8 @@ def subset_faces(self, indices, name: str | None = None):
 
         return grid
 
-    def subset(self, indices, name: str | None = None):
+    @override
+    def subset(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2DExtended:
         """Create a subset UnstructGrid2DExtended from this instance.
 
         .. warning::
@@ -304,12 +328,16 @@ def subset(self, indices, name: str | None = None):
 
         Parameters
         ----------
-        indices : array_like
+        indices
             Indices of the cells of the original grid in the subset.
-        name : str, optional
-            Name of the grid subset. Default is `instance.name` + `' subset'`.
-        """
+        name
+            Name of the grid subset. Default is ``instance.name + " subset"``.
 
+        Returns
+        -------
+        `.UnstructGrid2DExtended`
+            Subset instance.
+        """
         grid = UnstructGrid2DExtended.__new__(UnstructGrid2DExtended)
 
         grid._name = name or self._name + " subset"
@@ -366,7 +394,8 @@ def subset(self, indices, name: str | None = None):
 
         return grid
 
-    def interpolator(self, grid_data, fill_value: float = 0):
+    @override
+    def interpolator(self, grid_data: ArrayLike, 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
@@ -374,14 +403,14 @@ def interpolator(self, grid_data, fill_value: float = 0):
 
         Parameters
         ----------
-        grid_data : array_like
+        grid_data
             Array containing data in the grid cells.
-        fill_value : float, optional
+        fill_value
             Value returned outside the grid, by default 0.
 
         Returns
         -------
-        UnstructGridFunction3D
+        `.UnstructGridFunction3D`
             Interpolator instance.
         """
         if self._interpolator is None:
@@ -392,23 +421,25 @@ def interpolator(self, grid_data, fill_value: float = 0):
 
         return UnstructGridFunction3D.instance(self._interpolator, grid_data, fill_value)
 
-    def vector_interpolator(self, grid_vectors, fill_vector: Vector3D = ZERO_VECTOR):
-        """Return an `UnstructGridVectorFunction3D` interpolator instance for the vector data
-        defined on this grid.
+    @override
+    def vector_interpolator(
+        self, grid_vectors: ArrayLike, fill_vector: Vector3D = ZERO_VECTOR
+    ) -> UnstructGridVectorFunction3D:
+        """Return an `UnstructGridVectorFunction3D` interpolator instance for the vector data defined on this grid.
 
         On the second and subsequent calls, the interpolator is created as an instance
         of the previously created interpolator sharing the same KDtree structure.
 
         Parameters
         ----------
-        grid_vectors : (3,K) array_like
-            Array containing 3D vectors in the grid cells.
-        fill_vector : Vector3D, optional
-            3D vector returned outside the grid, by default `Vector3D(0, 0, 0)`.
+        grid_vectors
+            ``(3, L)`` Array containing 3D vectors in the grid cells.
+        fill_vector
+            3D vector returned outside the grid.
 
         Returns
         -------
-        UnstructGridVectorFunction3D
+        `.UnstructGridVectorFunction3D`
             Interpolator instance.
         """
         if self._interpolator is None:
@@ -423,7 +454,14 @@ def vector_interpolator(self, grid_vectors, fill_vector: Vector3D = ZERO_VECTOR)
 
         return UnstructGridVectorFunction3D.instance(self._interpolator, grid_vectors, fill_vector)
 
+    @override
     def __getstate__(self):
+        """Get the state of the UnstructGrid2DExtended instance for serialization.
+
+        Returns
+        -------
+        Dictionary with the instance attributes.
+        """
         state = {
             "name": self._name,
             "dimension": self._dimension,
@@ -438,6 +476,7 @@ def __getstate__(self):
         return state
 
     def __setstate__(self, state):
+        """Restore the state of the UnstructGrid2DExtended instance."""
         self._name = state["name"]
         self._dimension = state["dimension"]
         self._coordinate_system = state["coordinate_system"]
@@ -454,39 +493,50 @@ def __setstate__(self, state):
 
         self._initial_setup()
 
-    def plot_tetra_mesh(self, data=None, ax=None):
+    def plot_tetra_mesh(
+        self, data: ArrayLike | None = None, ax: matplotlib.axes.Axes | None = None
+    ) -> None:
         """Plot the tetrahedral mesh grid geometry.
 
+        .. warning::
+            Plotting of tetrahedral mesh is not implemented yet.
+
         Parameters
         ----------
-        data : array_like
+        data
             Data array defined on the tetrahedral mesh.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot the mesh. If None, a new figure is created.
 
-        Returns
-        -------
-        matplotlib.axes.Axes
-            Matplotlib axes with the plotted mesh.
+        Raises
+        ------
+        NotImplementedError
+            Plotting of tetrahedral mesh is not implemented yet.
         """
         raise NotImplementedError("Plotting of tetrahedral mesh is not implemented yet.")
 
-    def plot_mesh(self, data=None, ax=None, **grid_styles: str | float):
+    @override
+    def plot_mesh(
+        self,
+        data: ArrayLike | None = None,
+        ax: matplotlib.axes.Axes | None = None,
+        **grid_styles: str | float,
+    ) -> matplotlib.axes.Axes:
         """Plot the polygonal mesh grid geometry at the first poloidal plane to a matplotlib figure.
 
         Parameters
         ----------
-        data : array_like
+        data
             Data array defined on the polygonal mesh at the poloidal plane.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot the mesh. If None, a new figure is created.
-        **grid_styles : str | float, optional
+        **grid_styles
             Styles for the grid lines and faces,
             by default ``{"facecolor": "none", "edgecolor": "b", "linewidth": 0.25}``.
 
         Returns
         -------
-        matplotlib.axes.Axes
+        `~matplotlib.axes.Axes`
             Matplotlib axes with the plotted mesh.
         """
         if ax is None:
@@ -519,25 +569,31 @@ def plot_mesh(self, data=None, ax=None, **grid_styles: str | float):
 
         return ax
 
-    def plot_tri_mesh(self, data, ax=None, cmap="viridis"):
-        """Plot the data defined on the triangular mesh at the poloidal plane to a matplotlib
-        figure.
+    def plot_tri_mesh(
+        self, data: ArrayLike, ax: matplotlib.axes.Axes | None = None, cmap: str = "viridis"
+    ) -> matplotlib.axes.Axes:
+        """Plot the data defined on the triangular mesh at the poloidal plane to a matplotlib figure.
 
         Parameters
         ----------
-        data : array_like
+        data
             Data array defined on the triangular mesh at the poloidal plane.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot the mesh. If None, a new figure is created.
-        cmap : str, optional
+        cmap
             Colormap to use for the data, by default 'viridis'.
 
         Returns
         -------
-        matplotlib.axes.Axes
+        `~matplotlib.axes.Axes`
             Matplotlib axes with the plotted mesh.
+
+        Raises
+        ------
+        ValueError
+            If the data array does not have the same number of faces as the grid.
         """
-        data = np.asarray(data)
+        data = np.asarray_chkfinite(data)
         if data.shape[0] != self._num_faces:
             raise ValueError("The data array must have the same number of faces as the grid.")
         data = np.repeat(data, 2)
diff --git a/src/cherab/imas/ggd/unstruct_2d_mesh.py b/src/cherab/imas/ggd/unstruct_2d_mesh.py
index ed90f34..08e9b92 100755
--- a/src/cherab/imas/ggd/unstruct_2d_mesh.py
+++ b/src/cherab/imas/ggd/unstruct_2d_mesh.py
@@ -17,14 +17,24 @@
 # under the Licence.
 """Module defining unstructured 2D mesh class and related methods."""
 
+from __future__ import annotations
+
+import sys
+from typing import Literal
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override  # pyright: ignore[reportUnreachable]
+import matplotlib.axes
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.collections import PolyCollection
-from raysect.core.math import Vector3D
+from numpy.typing import ArrayLike, NDArray
 from raysect.core.math.polygon import triangulate2d
+from raysect.core.math.vector import Vector3D
 
-from cherab.imas.math import UnstructGridFunction2D, UnstructGridVectorFunction2D
-
+from ..math import UnstructGridFunction2D, UnstructGridVectorFunction2D
 from .base_mesh import GGDGrid
 
 __all__ = ["UnstructGrid2D"]
@@ -41,20 +51,26 @@ class UnstructGrid2D(GGDGrid):
 
     Parameters
     ----------
-    vertices : (N, 2) array_like
-        Array-like of shape (N, 2) containing coordinates of the polygon vertices.
-    cells : list of (N,) array_like
-        List of (N,)-shaped arrays containing the vertex indices in clockwise or
+    vertices
+        Array-like of shape ``(N, 2)`` containing coordinates of the polygon vertices.
+    cells
+        List of ``(N,)``-shaped arrays containing the vertex indices in clockwise or
         counterclockwise order for each polygonal cell in the list (the starting vertex must not be
         included twice).
-    name : str, optional
-        Name of the grid, by default 'Cells'.
-    coordinate_system : {'cylindrical', 'cartesian'}, optional
-        Coordinate system of the grid, by default 'cylindrical'.
+    name
+        Name of the grid, by default ``'Cells'``.
+    coordinate_system
+        Coordinate system of the grid, by default ``'cylindrical'``.
     """
 
-    def __init__(self, vertices, cells, name: str = "Cells", coordinate_system="cylindrical"):
-        vertices = np.array(vertices, dtype=np.float64)
+    def __init__(
+        self,
+        vertices: ArrayLike,
+        cells: list[ArrayLike],
+        name: str = "Cells",
+        coordinate_system: Literal["cylindrical", "cartesian"] = "cylindrical",
+    ):
+        vertices = np.asarray_chkfinite(vertices, dtype=np.float64)
         vertices.setflags(write=False)
 
         if vertices.ndim != 2:
@@ -76,16 +92,16 @@ def __init__(self, vertices, cells, name: str = "Cells", coordinate_system="cyli
             if len(cell) < 3:
                 raise ValueError(f"Cell {np.array2string(cell)} is not a polygon.")
 
-        self._vertices = vertices
-
-        self._cells = tuple(cells)
+        self._vertices: NDArray[np.float64] = vertices
+        self._cells: tuple[ArrayLike, ...] = tuple(cells)
 
         super().__init__(name, 2, coordinate_system)
 
+    @override
     def _initial_setup(self) -> None:
         self._interpolator = None
 
-        self._num_cell = len(self._cells)
+        self._num_cell: int = len(self._cells)
 
         x = self._vertices[:, 0]
         y = self._vertices[:, 1]
@@ -152,28 +168,28 @@ def _initial_setup(self) -> None:
             self._cell_volume.setflags(write=False)
 
     @property
-    def vertices(self):
-        """Mesh vertex coordinates as (N, 2) array."""
+    def vertices(self) -> NDArray[np.float64]:
+        """Mesh vertex coordinates as ``(N, 2)`` array."""
         return self._vertices
 
     @property
-    def cells(self):
-        """List of K polygonal cells."""
+    def cells(self) -> tuple[ArrayLike, ...]:
+        """List of ``K`` polygonal cells."""
         return self._cells
 
     @property
-    def triangles(self):
-        """Mesh triangles as (M, 3) array."""
+    def triangles(self) -> NDArray[np.int32]:
+        """Mesh triangles as ``(M, 3)`` array."""
         return self._triangles
 
     @property
-    def triangle_to_cell_map(self):
-        """Array of shape (M,) mapping every triangle index to a grid cell ID."""
+    def triangle_to_cell_map(self) -> NDArray[np.int32]:
+        """Array of shape ``(M,)`` mapping every triangle index to a grid cell ID."""
         return self._triangle_to_cell_map
 
     @property
-    def cell_to_triangle_map(self):
-        """Array of shape (K, 2) mapping every grid cell index to triangle IDs.
+    def cell_to_triangle_map(self) -> NDArray[np.int32]:
+        """Array of shape ``(K, 2)`` mapping every grid cell index to triangle IDs.
 
         The first column is the index of the first triangle forming the cell.
         The second column is the number of triangles forming the cell.
@@ -183,17 +199,22 @@ def cell_to_triangle_map(self):
         """
         return self._cell_to_triangle_map
 
-    def subset(self, indices, name=None):
+    @override
+    def subset(self, indices: ArrayLike, name: str | None = None) -> UnstructGrid2D:
         """Create a subset UnstructGrid2D from this instance.
 
         Parameters
         ----------
-        indices : array_like
+        indices
             Indices of the cells of the original grid in the subset.
-        name : str, optional
-            Name of the grid subset. Default is `instance.name` + `' subset'`.
-        """
+        name
+            Name of the grid subset. Default is ``instance.name + " subset"``.
 
+        Returns
+        -------
+        `.UnstructGrid2D`
+            Subset instance.
+        """
         grid = UnstructGrid2D.__new__(UnstructGrid2D)
 
         grid._name = name or self.name + " subset"
@@ -272,23 +293,23 @@ def subset(self, indices, name=None):
 
         return grid
 
-    def interpolator(self, grid_data, fill_value=0):
-        """Return an `UnstructGridFunction2D` interpolator instance for the data defined on this
-        grid.
+    @override
+    def interpolator(self, grid_data: ArrayLike, 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
         previously created interpolator sharing the same KDtree structure.
 
         Parameters
         ----------
-        grid_data : array_like
+        grid_data
             Array containing data in the grid cells.
-        fill_value : float, optional
+        fill_value
             Value returned outside the grid, by default is 0.
 
         Returns
         -------
-        UnstructGridFunction2D
+        `.UnstructGridFunction2D`
             Interpolator instance.
         """
         if self._interpolator is None:
@@ -299,23 +320,25 @@ def interpolator(self, grid_data, fill_value=0):
 
         return UnstructGridFunction2D.instance(self._interpolator, grid_data, fill_value)
 
-    def vector_interpolator(self, grid_vectors, fill_vector=ZERO_VECTOR):
-        """Return an `UnstructGridVectorFunction2D` interpolator instance for the vector data
-        defined on this grid.
+    @override
+    def vector_interpolator(
+        self, grid_vectors: ArrayLike, fill_vector: Vector3D = ZERO_VECTOR
+    ) -> UnstructGridVectorFunction2D:
+        """Return an `UnstructGridVectorFunction2D` interpolator instance for the vector data defined on this grid.
 
         On the second and subsequent calls, the interpolator is created as an instance of the
         previously created interpolator sharing the same KDtree structure.
 
         Parameters
         ----------
-        grid_vectors : (3, K) array_like
-            Array containing 3D vectors in the grid cells.
-        fill_vector : Vector3D, optional
-            3D vector returned outside the grid, by default Vector3D(0, 0, 0).
+        grid_vectors
+            ``(3, K)`` Array containing 3D vectors in the grid cells.
+        fill_vector
+            3D vector returned outside the grid.
 
         Returns
         -------
-        UnstructGridVectorFunction2D
+        `.UnstructGridVectorFunction2D`
             Interpolator instance.
         """
         if self._interpolator is None:
@@ -330,7 +353,14 @@ def vector_interpolator(self, grid_vectors, fill_vector=ZERO_VECTOR):
 
         return UnstructGridVectorFunction2D.instance(self._interpolator, grid_vectors, fill_vector)
 
+    @override
     def __getstate__(self):
+        """Serialize the state of the UnstructGrid2D instance for pickling.
+
+        Returns
+        -------
+        Dictionary with the instance attributes.
+        """
         state = {
             "name": self._name,
             "dimension": self._dimension,
@@ -341,6 +371,7 @@ def __getstate__(self):
         return state
 
     def __setstate__(self, state):
+        """Restore the state of the UnstructGrid2D instance from the serialized state."""
         self._name = state["name"]
         self._dimension = state["dimension"]
         self._coordinate_system = state["coordinate_system"]
@@ -350,24 +381,33 @@ def __setstate__(self, state):
 
         self._initial_setup()
 
-    def plot_triangle_mesh(self, data=None, ax=None):
+    def plot_triangle_mesh(
+        self, data: ArrayLike | None = None, ax: matplotlib.axes.Axes | None = None
+    ) -> matplotlib.axes.Axes:
         """Plot the triangle mesh grid geometry to a matplotlib figure.
 
         Parameters
         ----------
-        data : array_like
+        data
             Data array defined on the polygonal mesh.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot on. If None, a new figure and axes are created.
+
+        Returns
+        -------
+        `~matplotlib.axes.Axes`
+            The matplotlib axes with the plotted mesh.
         """
         if ax is None:
             _, ax = plt.subplots(constrained_layout=True)
 
         verts = self._vertices[self._triangles]
         if data is None:
-            collection_mesh = PolyCollection(verts, facecolor="none", edgecolor="b", linewidth=0.25)
+            collection_mesh = PolyCollection(
+                [verts], facecolor="none", edgecolor="b", linewidth=0.25
+            )
         else:
-            collection_mesh = PolyCollection(verts)
+            collection_mesh = PolyCollection([verts])
             collection_mesh.set_array(data[self._triangle_to_cell_map])
         ax.add_collection(collection_mesh)
         ax.set_aspect(1)
@@ -383,23 +423,38 @@ def plot_triangle_mesh(self, data=None, ax=None):
 
         return ax
 
-    def plot_mesh(self, data=None, ax=None):
+    @override
+    def plot_mesh(
+        self,
+        data: ArrayLike | None = None,
+        ax: matplotlib.axes.Axes | None = None,
+        **grid_styles: str | float,
+    ) -> matplotlib.axes.Axes:
         """Plot the polygonal mesh grid geometry to a matplotlib figure.
 
         Parameters
         ----------
-        data : array_like, optional
+        data
             Data array defined on the polygonal mesh.
-        ax : matplotlib.axes.Axes, optional
+        ax
             Matplotlib axes to plot on. If None, a new figure and axes are created.
-        """
 
+        Returns
+        -------
+        `~matplotlib.axes.Axes`
+            The matplotlib axes with the plotted mesh.
+        """
         if ax is None:
             _, ax = plt.subplots(constrained_layout=True)
 
+        # Set default grid line styles if not provided
+        grid_styles.setdefault("facecolor", "none")
+        grid_styles.setdefault("edgecolor", "b")
+        grid_styles.setdefault("linewidth", 0.25)
+
         verts = [self._vertices[cell] for cell in self._cells]
         if data is None:
-            collection_mesh = PolyCollection(verts, facecolor="none", edgecolor="b", linewidth=0.25)
+            collection_mesh = PolyCollection(verts, **grid_styles)
         else:
             collection_mesh = PolyCollection(verts)
             collection_mesh.set_array(data)
diff --git a/src/cherab/imas/ids/common/__init__.py b/src/cherab/imas/ids/common/__init__.py
index 411769e..41e2c65 100644
--- a/src/cherab/imas/ids/common/__init__.py
+++ b/src/cherab/imas/ids/common/__init__.py
@@ -17,7 +17,7 @@
 # under the Licence.
 """Subpackage for common utilities for loading data from IMAS IDS structures."""
 
-from . import species
+from . import ggd, species
 from .slice import get_ids_time_slice
 
-__all__ = ["species", "get_ids_time_slice"]
+__all__ = ["species", "get_ids_time_slice", "ggd"]
diff --git a/src/cherab/imas/ids/common/ggd/load_grid.py b/src/cherab/imas/ids/common/ggd/load_grid.py
index c52c307..1a35ca1 100644
--- a/src/cherab/imas/ids/common/ggd/load_grid.py
+++ b/src/cherab/imas/ids/common/ggd/load_grid.py
@@ -17,37 +17,75 @@
 # under the Licence.
 """Module for loading GGD grids from IMAS grid_ggd IDS structures."""
 
+from numpy import int32
+from numpy.typing import NDArray
+
 from imas.ids_structure import IDSStructure
 
+from ....ggd.unstruct_2d_extend_mesh import UnstructGrid2DExtended
+from ....ggd.unstruct_2d_mesh import UnstructGrid2D
 from .load_unstruct_2d import load_unstruct_grid_2d
 from .load_unstruct_3d import load_unstruct_grid_2d_extended
 
-__all__ = ["load_grid", "load_unstruct_grid_2d"]
+__all__ = ["load_grid"]
+
+
+def load_grid(
+    grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int | None = None
+) -> (
+    UnstructGrid2D
+    | tuple[UnstructGrid2D, dict[str, NDArray[int32]], dict[str, int]]
+    | UnstructGrid2DExtended
+):
+    """Load grid from the ``grid_ggd`` structure.
 
+    The ``grid_ggd`` structure is expected to follow the IMAS GGD grid definition.
+    Please see: https://imas-data-dictionary.readthedocs.io/en/latest/ggd_guide/doc.html#the-grid-ggd-aos
 
-def load_grid(grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int | None = None):
-    """Load grid from the grid_ggd structure.
+    .. warning::
+        This function currently supports only unstructured 2D grids and unstructured 2D grids
+        extended in 3D (mainly used in JOREK).
+        Loading of structured grids and unstructured 3D grids will be implemented in the future.
 
     Parameters
     ----------
-    grid_ggd : IDSStructure
-        The grid_ggd structure.
-    with_subsets : bool, optional
+    grid_ggd
+        The ``grid_ggd`` structure.
+    with_subsets
         Read grid subset data, by default is False.
-    num_toroidal : int, optional
+    num_toroidal
         Number of toroidal points, by default None.
 
     Returns
     -------
-    grid : UnstructGrid2D
-        Grid object that depends on the grid type (structured/unstructured, 2D/3D).
-    subsets : dict, optional
+    grid : `.UnstructGrid2D` | `.UnstructGrid2DExtended`
+        Grid object that depends on the grid type (2D/2D-extended/3D).
+    subsets : `dict[str, NDArray[int32]]`
         Dictionary with grid subsets for each subset name containing the indices of the cells from
         that subset. Note that 'Cells' subset is included only if cell indices are specified.
-    subset_id : dict, optional
+    subset_id : `dict[str, int]`
         Dictionary with grid subset indices.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the grid type is unsupported or if no spaces are found in the ``grid_ggd`` structure.
+    NotImplementedError
+        If the grid type is not yet implemented.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from imas import DBEntry
+        from cherab.imas.ids.common import get_ids_time_slice
+        from cherab.imas.ids.common.ggd import load_grid
+
+        with DBEntry("imas:hdf5?path=/work/imas/shared/imasdb/ITER/4/123356/1", "r") as entry:
+            ids = get_ids_time_slice(entry, "edge_profiles", 0)
+
+        grid, subsets, subset_id = load_grid(ids.grid_ggd[0], with_subsets=True)
+    """
     spaces = get_standard_spaces(grid_ggd)
 
     if not len(spaces):
@@ -95,15 +133,19 @@ def get_standard_spaces(grid_ggd: IDSStructure) -> list[IDSStructure]:
 
     Parameters
     ----------
-    grid_ggd : IDSStructure
+    grid_ggd
         The grid_ggd structure.
 
     Returns
     -------
     list[IDSStructure]
         List of standard spaces.
-    """
 
+    Raises
+    ------
+    ValueError
+        If no spaces are defined in the grid_ggd structure.
+    """
     if not len(grid_ggd.space):
         error_massage = "Unable to read the grid. Grid space is not defined."
         if len(grid_ggd.path):
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 aa1e759..7c16821 100644
--- a/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py
+++ b/src/cherab/imas/ids/common/ggd/load_unstruct_2d.py
@@ -18,11 +18,13 @@
 """Module for loading unstructured 2D grids from IMAS grid_ggd IDS structure."""
 
 import numpy as np
+from numpy.typing import NDArray
 
-from cherab.imas.ggd import UnstructGrid2D
 from imas.ids_defs import EMPTY_INT
 from imas.ids_structure import IDSStructure
 
+from ....ggd.unstruct_2d_mesh import UnstructGrid2D
+
 __all__ = ["load_unstruct_grid_2d"]
 
 VERTEX_DIMENSION = 0
@@ -30,29 +32,35 @@
 FACE_DIMENSION = 2
 
 
-def load_unstruct_grid_2d(grid_ggd: IDSStructure, space_index=0, with_subsets=False):
+def load_unstruct_grid_2d(
+    grid_ggd: IDSStructure, space_index: int = 0, with_subsets: bool = False
+) -> UnstructGrid2D | tuple[UnstructGrid2D, dict[str, NDArray[np.int32]], dict[str, int]]:
     """Load unstructured 2D grid from the grid_ggd structure.
 
     Parameters
     ----------
-    grid_ggd : IDSStructure
+    grid_ggd
         The grid_ggd structure.
-    space_index : int, optional
+    space_index
         The index of the grid space, by default 0.
-    with_subsets : bool, optional
+    with_subsets
         Read grid subset data, by default False.
 
     Returns
     -------
-    grid : UnstructGrid2D
+    grid : `.UnstructGrid2D`
         Unstructured 2D grid object.
-    subsets : dict, optional
+    subsets : `dict[str, NDArray[numpy.int32]]`
         Dictionary with grid subsets for each subset name containing the indices of the cells from
         that subset. Note that 'Cells' subset is included only if cell indices are specified.
-    subset_id : dict, optional
+    subset_id : `dict[str, int]`
         Dictionary with grid subset indices.
-    """
 
+    Raises
+    ------
+    ValueError
+        If the specified space does not contain a 2D grid.
+    """
     space = grid_ggd.space[space_index]
 
     # Check if the grid is 2D
@@ -109,7 +117,7 @@ def load_unstruct_grid_2d(grid_ggd: IDSStructure, space_index=0, with_subsets=Fa
             else:
                 winding_ok = False
 
-        cells.append(cell)
+            cells.append(cell)
 
     if not winding_ok:
         print("Warning! Unable to verify that the cell nodes are in the winging order.")
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 997d594..c963670 100644
--- a/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py
+++ b/src/cherab/imas/ids/common/ggd/load_unstruct_3d.py
@@ -19,6 +19,7 @@
 
 import numpy as np
 
+from imas.ids_struct_array import IDSStructArray
 from imas.ids_structure import IDSStructure
 
 from ....ggd import UnstructGrid2DExtended
@@ -36,34 +37,38 @@
 
 
 def load_unstruct_grid_2d_extended(
-    grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int | None = None
-):
-    """Load unstructured 2D grid extended in 3D from the grid_ggd structure.
+    grid_ggd: IDSStructure, with_subsets: bool = False, num_toroidal: int = NUM_TOROIDAL
+) -> UnstructGrid2DExtended:
+    """Load unstructured 2D grid extended in 3D from the ``grid_ggd`` structure.
 
     Parameters
     ----------
-    grid_ggd : IDSStructure
-        The grid_ggd structure.
-    with_subsets : bool, optional
-        Read grid subset data if True, by default False.
-    num_toroidal : int, optional
-        Number of toroidal points, by default None.
+    grid_ggd
+        The ``grid_ggd`` structure.
+    with_subsets
+        Read grid subset data if True.
+    num_toroidal
+        Number of toroidal points.
+        If specifying more than 1, the grid will be extended in 3D by repeating the 2D
+        cross-section around the torus with evenly spaced toroidal angles.
 
     Returns
     -------
-    UnstructGrid2DExtended
+    `.UnstructGrid2DExtended`
         The unstructured 2D grid extended in 3D.
+
+    Raises
+    ------
+    ValueError
+        If the number of toroidal points is less than 1.
+        If the grid is not an unstructured extended 2D grid.
     """
-    # Check if the number of toroidal points is specified
-    if isinstance(num_toroidal, int):
-        if num_toroidal < 1:
-            raise ValueError("The number of toroidal points must be greater than 0.")
-        num_toroidal = num_toroidal
-    else:
-        num_toroidal = NUM_TOROIDAL
+    # Validate num_toroidal
+    if num_toroidal < 1:
+        raise ValueError("The number of toroidal points must be greater than 0.")
 
     # Get the R-Z space
-    space = 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:
@@ -71,7 +76,7 @@ def load_unstruct_grid_2d_extended(
             "The load_unstruct_grid_2d_extended() supports only unstructured extended 2D grids."
         )
 
-    grid_name = grid_ggd.identifier.name
+    grid_name = str(grid_ggd.identifier.name)
 
     # =========================================
     # Reading vertices (poloidal and toroidal)
@@ -97,7 +102,7 @@ def load_unstruct_grid_2d_extended(
     # =========================================
     # Reading cells indices
     # =========================================
-    faces = space.objects_per_dimension[FACE_DIMENSION].object
+    faces: IDSStructArray = space.objects_per_dimension[FACE_DIMENSION].object
     num_faces = len(faces)
     cells = np.zeros((num_faces * num_toroidal, 8), dtype=np.int32)
     i_cell = 0
@@ -122,23 +127,23 @@ def load_unstruct_grid_2d_extended(
 
 
 def load_unstruct_grid_3d(grid_ggd: IDSStructure, space_index: int = 0, with_subsets: bool = False):
-    """Load unstructured 3D grid from the grid_ggd structure.
+    """Load unstructured 3D grid from the ``grid_ggd`` structure.
 
     .. warning::
         This function is a placeholder for future implementation.
 
     Parameters
     ----------
-    grid_ggd : IDSStructure
-        The grid_ggd structure.
-    space_index : int, optional
+    grid_ggd
+        The ``grid_ggd`` structure.
+    space_index
         The index of the space to read, by default 0.
-    with_subsets : bool, optional
+    with_subsets
         Read grid subset data if True, by default False.
 
     Returns
     -------
-    UnstructGrid3D
+    `.UnstructGrid3D`
         The unstructured 3D grid.
     """
     raise NotImplementedError("Loading unstructured 3D grids will be implemented in the future.")
diff --git a/src/cherab/imas/ids/common/slice.py b/src/cherab/imas/ids/common/slice.py
index 3ac8d69..1711941 100644
--- a/src/cherab/imas/ids/common/slice.py
+++ b/src/cherab/imas/ids/common/slice.py
@@ -44,15 +44,15 @@ def get_ids_time_slice(
 
     Parameters
     ----------
-    entry : `~imas.db_entry.DBEntry`
+    entry
         The IMAS entry. The entry must be opened in read mode.
-    ids_name : str
+    ids_name
         The name of the IDS.
-    time : float, optional
+    time
         The time in seconds of the requested time slice, by default is 0.
-    occurrence : int, optional
+    occurrence
         The occurrence of the IDS, by default is 0.
-    time_threshold : float, optional
+    time_threshold
         The maximum allowed time difference in seconds between the actual time of the nearest time
         slice and the given time, by default is infinity.
 
@@ -60,6 +60,29 @@ def get_ids_time_slice(
     -------
     `~imas.ids_toplevel.IDSToplevel`
         The requested IDS time slice.
+
+    Raises
+    ------
+    ValueError
+        If `.time` or `.time_threshold` is negative.
+    RuntimeError
+        If the requested IDS is empty.
+    RuntimeError
+        If the time difference between the actual time of the nearest time slice and the given time
+        exceeds the specified threshold.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        from imas import DBEntry
+        from cherab.imas.ids.common import get_ids_time_slice
+
+        with DBEntry(
+            "imas://uda.iter.org/uda?path=/work/imas/shared/imasdb/ITER/3/123072/3&backend=hdf5",
+            "r",
+        ) as entry:
+            ids = get_ids_time_slice(entry, "equilibrium", time=0.0)
     """
     if time < 0:
         raise ValueError(f"Argument 'time' must be >=0 ({time} s).")
@@ -72,6 +95,7 @@ def get_ids_time_slice(
             time,
             CLOSEST_INTERP,
             occurrence=occurrence,
+            autoconvert=False,
         )
     except NotImplementedError:
         # Fallback to `get` method to retrieve the entire IDS
@@ -81,7 +105,7 @@ def get_ids_time_slice(
             RuntimeWarning,
             stacklevel=2,
         )
-        ids = entry.get(ids_name, occurrence=occurrence)
+        ids = entry.get(ids_name, occurrence=occurrence, autoconvert=False)
 
     if not len(ids.time):
         raise RuntimeError(f"The '{ids_name}' IDS is empty.")
diff --git a/src/cherab/imas/ids/common/species.py b/src/cherab/imas/ids/common/species.py
index 5ef819b..c3bb5ad 100644
--- a/src/cherab/imas/ids/common/species.py
+++ b/src/cherab/imas/ids/common/species.py
@@ -41,20 +41,20 @@ def get_ion_state(
 
     Parameters
     ----------
-    state : IDSStructure
+    state
         The ion_state structure from IMAS.
-    state_index : int
+    state_index
         Index of the state in the list of states.
-    elements : list[`~cherab.core.atomic.elements.Element`]
+    elements
         List of elements that make up the ion state.
-    grid_subset_index : int, optional
+    grid_subset_index
         The grid subset index to use for 1D profiles, by default None.
 
     Returns
     -------
-    species_type : str
+    species_type : `str`
         The type of species: 'ion', 'ion_bundle', 'molecule', or 'molecular_bundle'.
-    species_id : frozenset
+    species_id : `frozenset`
         A frozenset of key-value pairs that uniquely identify the species.
     """
     if state.z_min == EMPTY_FLOAT or state.z_max == EMPTY_FLOAT:
@@ -116,16 +116,16 @@ def get_neutral_state(state: IDSStructure, elements: list[Element]) -> tuple[str
 
     Parameters
     ----------
-    state : IDSStructure
+    state
         The neutral_state structure from IMAS.
-    elements : list[`~cherab.core.atomic.elements.Element`]
+    elements
         List of elements that make up the neutral state.
 
     Returns
     -------
-    species_type : str
+    species_type : `str`
         The type of species: 'molecule' or 'ion'.
-    species_id : frozenset
+    species_id : `frozenset`
         A frozenset of key-value pairs that uniquely identify the species.
     """
     state_dict = {"name": state.name.strip()}
@@ -161,16 +161,16 @@ def get_ion(ion: IDSStructure, elements: list[Element]) -> tuple[str, frozenset]
 
     Parameters
     ----------
-    ion : IDSStructure
+    ion
         The ion structure from IMAS.
-    elements : list[`~cherab.core.atomic.elements.Element`]
+    elements
         List of elements that make up the ion.
 
     Returns
     -------
-    species_type : str
+    species_type : `str`
         The type of species: 'molecule' or 'ion'.
-    species_id : frozenset
+    species_id : `frozenset`
         A frozenset of key-value pairs that uniquely identify the species.
     """
     z_ion = int(ion.z_ion) if ion.z_ion != EMPTY_FLOAT else elements[0].atomic_number
@@ -198,21 +198,21 @@ def get_ion(ion: IDSStructure, elements: list[Element]) -> tuple[str, frozenset]
     return "ion", species_id
 
 
-def get_neutral(neutral, elements):
+def get_neutral(neutral: IDSStructure, elements: list[Element]) -> tuple[str, frozenset]:
     """Get a unique identifier for a neutral or molecule.
 
     Parameters
     ----------
-    neutral : IDSStructure
+    neutral
         The neutral structure from IMAS.
-    elements : list[`~cherab.core.atomic.elements.Element`]
+    elements
         List of elements that make up the neutral state.
 
     Returns
     -------
-    species_type : str
+    species_type : `str`
         The type of species: 'molecule' or 'ion'.
-    species_id : frozenset
+    species_id : `frozenset`
         A frozenset of key-value pairs that uniquely identify the species.
     """
     if len(elements) > 1:
@@ -244,7 +244,7 @@ def get_element_list(element_aos: IDSStructArray) -> list[Element]:
 
     Parameters
     ----------
-    element_aos : IDSStructArray
+    element_aos
         The element_aos structure from IMAS.
 
     Returns
diff --git a/src/cherab/imas/ids/core_profiles/load_profiles.py b/src/cherab/imas/ids/core_profiles/load_profiles.py
index f569872..4745904 100644
--- a/src/cherab/imas/ids/core_profiles/load_profiles.py
+++ b/src/cherab/imas/ids/core_profiles/load_profiles.py
@@ -41,9 +41,9 @@ def load_core_profiles(
 
     Parameters
     ----------
-    species_struct : IDSStructure
+    species_struct
         IDS structure containing the profiles for a single species.
-    backup_species_struct : IDSStructure | None, optional
+    backup_species_struct
         The backup ids structure that is used if the profile is missing in species_struct,
         by default None.
 
@@ -53,7 +53,6 @@ def load_core_profiles(
         Dictionary with the profiles: ``density``, ``density_thermal``, ``density_fast``,
         ``temperature`` and ``z_average_1d``.
     """
-
     profiles = {
         "density": None,
         "density_thermal": None,
@@ -77,16 +76,15 @@ def load_core_grid(grid_struct: IDSStructure) -> dict[str, np.ndarray | None]:
 
     Parameters
     ----------
-    grid_struct : IDSStructure
+    grid_struct
         The IDS structure containing the grid data for 1D profiles.
 
     Returns
     -------
     dict[str, np.ndarray | None]
-        A dictionary with the following keys: ``rho_tor_norm``, ``psi``, ``volume``,
+        Dictionary with the following keys: ``rho_tor_norm``, ``psi``, ``volume``,
         ``area``, ``surface``.
     """
-
     grid = {
         "rho_tor_norm": None,
         "volume": None,
@@ -162,15 +160,19 @@ def load_core_species(profiles_struct: IDSStructure) -> dict[str, dict[str, np.n
 
     Parameters
     ----------
-    profiles_struct : IDSStructure
+    profiles_struct
         The IDS structure containing the core profiles data.
 
     Returns
     -------
     dict[str, dict[str, ndarray | None]]
         Dictionary with the species and their profiles.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If unable to determine the species type or identifier.
+    """
     species_types = ("molecule", "molecular_bundle", "ion", "ion_bundle")
     composition = {species_type: {} for species_type in species_types}
 
diff --git a/src/cherab/imas/ids/edge_profiles/load_profiles.py b/src/cherab/imas/ids/edge_profiles/load_profiles.py
index b82f61c..233595c 100644
--- a/src/cherab/imas/ids/edge_profiles/load_profiles.py
+++ b/src/cherab/imas/ids/edge_profiles/load_profiles.py
@@ -46,11 +46,11 @@ def load_edge_profiles(
 
     Parameters
     ----------
-    species_struct : IDSStructure
+    species_struct
         The ids structure containing the profiles for a single species.
-    grid_subset_index : int, optional
-        Identifier index of the grid subset, by default 5 ("Cells").
-    backup_species_struct : IDSStructure, optional
+    grid_subset_index
+        Identifier index of the grid subset, by default 5 (``"Cells"``).
+    backup_species_struct
         The backup ids structure that is used if the profile is missing in
         species_struct, by default None.
 
@@ -61,7 +61,6 @@ def load_edge_profiles(
         ``velocity_radial``, ``velocity_parallel``, ``velocity_poloidal``, ``velocity_phi``,
         ``velocity_r``, ``velocity_z``, ``z_average``.
     """
-
     profiles = {
         "density": None,
         "density_fast": None,
@@ -107,8 +106,10 @@ def load_edge_profiles(
 def load_edge_species(
     ggd_struct: IDSStructure, grid_subset_index: int = 5
 ) -> dict[str, dict[str, np.ndarray | None]]:
-    """Load edge plasma species and their profiles from a given GGD structure for a given grid and
-    subset indices.
+    """Load edge plasma species and their profiles from a given GGD structure.
+
+    The profiles are taken from the arrays with the given grid and subset indices
+    (e.g. ``ggd[i1].electrons``, ``ggd[i1].ion[i2].states[i3]``).
 
     The returned dictionary has the following structure.
 
@@ -168,17 +169,21 @@ def load_edge_species(
 
     Parameters
     ----------
-    ggd_struct : IDSStructure
+    ggd_struct
         The ggd ids structure containing the profiles.
-    grid_subset_index : int, optional
-        Identifier index of the grid subset, by default 5 ("Cells").
+    grid_subset_index
+        Identifier index of the grid subset, by default 5 (``"Cells"``).
 
     Returns
     -------
     dict[str, dict[str, ndarray | None]]
         Dictionary with plasma profiles.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If unable to determine the species due to missing element information.
+    """
     species_types = ("molecule", "molecular_bundle", "ion", "ion_bundle")
     composition = {species_type: {} for species_type in species_types}
 
diff --git a/src/cherab/imas/ids/equilibrium/load_equilibrium.py b/src/cherab/imas/ids/equilibrium/load_equilibrium.py
index 8e60d69..ac56b3d 100644
--- a/src/cherab/imas/ids/equilibrium/load_equilibrium.py
+++ b/src/cherab/imas/ids/equilibrium/load_equilibrium.py
@@ -17,9 +17,10 @@
 # under the Licence.
 """Module for loading 2D plasma equilibrium data from the equilibrium IDS."""
 
-from typing import Any
+from __future__ import annotations
 
 import numpy as np
+from numpy.typing import NDArray
 from raysect.core.math import Point2D
 
 from imas.ids_defs import EMPTY_FLOAT, EMPTY_INT
@@ -30,12 +31,14 @@
 __all__ = ["load_equilibrium_data"]
 
 
-def load_equilibrium_data(equilibrium_ids: IDSToplevel) -> dict[str, Any]:
+def load_equilibrium_data(
+    equilibrium_ids: IDSToplevel,
+) -> dict[str, NDArray[np.float64] | float | Point2D | list[Point2D] | None]:
     """Load 2D plasma equilibrium data from the equilibrium IDS.
 
     Parameters
     ----------
-    equilibrium_ids : `~imas.ids_toplevel.IDSToplevel`
+    equilibrium_ids
         The time-slice of the equilibrium IDS.
 
     Returns
@@ -43,26 +46,30 @@ def load_equilibrium_data(equilibrium_ids: IDSToplevel) -> dict[str, Any]:
     dict[str, Any]
         Dictionary with the following keys and values:
 
-        :r: (N, ) ndarray with R coordinates of rectangular grid,
-        :z: (M, ) ndarray with Z coordinates of rectangular grid,
-        :psi_grid: (N, M) ndarray with psi grid values,
+        :r: ``(N, )`` ndarray with R coordinates of rectangular grid,
+        :z: ``(M, )`` ndarray with Z coordinates of rectangular grid,
+        :psi_grid: ``(N, M)`` ndarray with psi grid values,
         :psi_axis: The psi value at the magnetic axis,
         :psi_lcfs: The psi value at the LCFS,
         :magnetic_axis: Point2D containing the coordinates of the magnetic axis,
         :x_points: list or tuple of Point2D x-points,
         :strike_points: list or tuple of Point2D strike-points,
-        :psi_norm: (K, ) ndarray with the values of psi_norm,
-        :f: (K, ) ndarray with the current flux profile on psi_norm,
-        :q: (K, ) ndarray with the safety factor (q) profile on psi_norm,
-        :phi: (K, ) ndarray with the toroidal flux profile on psi_norm,
-        :rho_tor_norm: (K, ) ndarray with the normalised toroidal flux coordinate on psi_norm,
+        :psi_norm: ``(K, )`` ndarray with the values of psi_norm,
+        :f: ``(K, )`` ndarray with the current flux profile on psi_norm,
+        :q: ``(K, )`` ndarray with the safety factor (q) profile on psi_norm,
+        :phi: ``(K, )`` ndarray with the toroidal flux profile on psi_norm,
+        :rho_tor_norm: ``(K, )`` ndarray with the normalised toroidal flux coordinate on psi_norm,
         :b_vacuum_radius: Vacuum B-field reference radius (in meters),
         :b_vacuum_magnitude: Vacuum B-Field magnitude at the reference radius,
-        :lcfs_polygon: (2, L) ndarray of ``[[x0, ...], [y0, ...]]`` vertices specifying the LCFS
+        :lcfs_polygon: ``(2, L)`` ndarray of ``[[x0, ...], [y0, ...]]`` vertices specifying the LCFS
             boundary
         :time: The time stamp of the time-slice (in seconds).
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the required data is not available in the IDS.
+    """
     if not len(equilibrium_ids.time_slice):
         raise RuntimeError("Equilibrium IDS does not have a time slice.")
 
@@ -105,28 +112,33 @@ def load_equilibrium_data(equilibrium_ids: IDSToplevel) -> dict[str, Any]:
     psi_axis = global_quantities.psi_axis
     if psi_axis == EMPTY_FLOAT:
         raise RuntimeError("Unable to read equilibrium: 'global_quantities.psi_axis' is not set.")
+    else:
+        psi_axis = float(psi_axis)
 
     psi_lcfs = global_quantities.psi_boundary
     if psi_lcfs == EMPTY_FLOAT:
         raise RuntimeError(
             "Unable to read equilibrium: 'global_quantities.psi_boundary' is not set."
         )
+    else:
+        psi_lcfs = float(psi_lcfs)
 
     psi_norm = (psi1d - psi_axis) / (psi_lcfs - psi_axis)
     psi_norm, index = np.unique(psi_norm, return_index=True)
-    f = profiles_1d.f[index]
-    q = profiles_1d.q[index]
+    f: NDArray[np.float64] = profiles_1d.f[index]
+    q: NDArray[np.float64] = profiles_1d.q[index]
 
     # additional 1D profiles
-    phi = profiles_1d.phi[index] if len(profiles_1d.phi) else None
-    rho_tor_norm = profiles_1d.rho_tor_norm[index] if len(profiles_1d.rho_tor_norm) else None
-
+    phi: NDArray[np.float64] | None = profiles_1d.phi[index] if len(profiles_1d.phi) else None
+    rho_tor_norm: NDArray[np.float64] | None = (
+        profiles_1d.rho_tor_norm[index] if len(profiles_1d.rho_tor_norm) else None
+    )
     r_axis = global_quantities.magnetic_axis.r
     z_axis = global_quantities.magnetic_axis.z
     if r_axis == EMPTY_FLOAT or z_axis == EMPTY_FLOAT:
         raise RuntimeError("Unable to read equilibrium: magnetic axis is not set.")
 
-    magnetic_axis = Point2D(r_axis.item(), z_axis.item())
+    magnetic_axis = Point2D(r_axis.value, z_axis.value)
 
     boundary = equilibrium_ids.time_slice[0].boundary
 
@@ -134,13 +146,13 @@ def load_equilibrium_data(equilibrium_ids: IDSToplevel) -> dict[str, Any]:
     if hasattr(boundary, "x_point"):  # DD4 no longer has x_point
         for x_point in boundary.x_point:
             if x_point.r != EMPTY_FLOAT and x_point.z != EMPTY_FLOAT:
-                x_points.append(Point2D(x_point.r, x_point.z))
+                x_points.append(Point2D(x_point.r.value, x_point.z.value))
 
     strike_points = []
     if hasattr(boundary, "strike_point"):  # DD4 no longer has strike_point
         for strike_point in boundary.strike_point:
             if strike_point.r != EMPTY_FLOAT and strike_point.z != EMPTY_FLOAT:
-                strike_points.append(Point2D(strike_point.r, strike_point.z))
+                strike_points.append(Point2D(strike_point.r.value, strike_point.z.value))
 
     r_lcfs = boundary.outline.r
     z_lcfs = boundary.outline.z
@@ -160,12 +172,16 @@ def load_equilibrium_data(equilibrium_ids: IDSToplevel) -> dict[str, Any]:
     b_vacuum_radius = equilibrium_ids.vacuum_toroidal_field.r0
     if b_vacuum_radius == EMPTY_FLOAT:
         raise RuntimeError("Unable to read equilibrium: vacuum_toroidal_field.r0 is not set.")
+    else:
+        b_vacuum_radius = float(b_vacuum_radius)
 
     b_vacuum_magnitude = equilibrium_ids.vacuum_toroidal_field.b0[0]
     if b_vacuum_magnitude == EMPTY_FLOAT:
         raise RuntimeError("Unable to read equilibrium: vacuum_toroidal_field.b0 is not set.")
+    else:
+        b_vacuum_magnitude = float(b_vacuum_magnitude)
 
-    time = equilibrium_ids.time[0]
+    time = float(equilibrium_ids.time[0])
 
     return {
         "r": r,
diff --git a/src/cherab/imas/ids/equilibrium/load_field.py b/src/cherab/imas/ids/equilibrium/load_field.py
index fcfc4ae..3f4ec58 100644
--- a/src/cherab/imas/ids/equilibrium/load_field.py
+++ b/src/cherab/imas/ids/equilibrium/load_field.py
@@ -28,26 +28,32 @@
 
 
 def load_magnetic_field_data(profiles_2d: IDSStructArray) -> dict:
-    """Load 2D profiles of the magnetic field components from the profiles_2d structure of the
-    equilibrium IDS.
+    """Load 2D profiles of the magnetic field components from equilibrium IDS.
+
+    The magnetic field components are extracted from the ``profiles_2d`` IDS structure,
+    assuming that the profiles are defined on a rectangular grid.
 
     Parameters
     ----------
-    profiles_2d : IDSStructArray
-        The profiles_2d structure of the equilibrium IDS.
+    profiles_2d
+        The ``profiles_2d`` structure of the equilibrium IDS.
 
     Returns
     -------
-    dict[str, ndarray]
-        Dictionary with the following keys:
-
-        :r: (N,) ndarray with R coordinates of rectangular grid.
-        :z: (M,) ndarray with Z coordinates of rectangular grid.
-        :b_field_r: (N, M) ndarray with R component of the magnetic field.
-        :b_field_z: (N, M) ndarray with Z component of the magnetic field.
-        :b_field_phi: (N, M) ndarray with toroidal component of the magnetic field.
+    Dictionary with the following keys:
+
+        :r: ``(N,)`` ndarray with R coordinates of rectangular grid.
+        :z: ``(M,)`` ndarray with Z coordinates of rectangular grid.
+        :b_field_r: ``(N, M)`` ndarray with R component of the magnetic field.
+        :b_field_z: ``(N, M)`` ndarray with Z component of the magnetic field.
+        :b_field_phi: ``(N, M)`` ndarray with toroidal component of the magnetic field.
+
+    Raises
+    ------
+    RuntimeError
+        If unable to read the magnetic field due to unsupported grid type or
+        mismatched array shapes.
     """
-
     rectangular_grid = False
     for prof2d in profiles_2d:
         if prof2d.grid_type.index == RECTANGULAR_GRID or prof2d.grid_type.index == EMPTY_INT:
diff --git a/src/cherab/imas/ids/wall/load2d.py b/src/cherab/imas/ids/wall/load2d.py
index 8acbf66..daa819a 100644
--- a/src/cherab/imas/ids/wall/load2d.py
+++ b/src/cherab/imas/ids/wall/load2d.py
@@ -29,7 +29,7 @@ def load_wall_2d(description_2d: IDSStructure) -> dict[str, np.ndarray]:
 
     Parameters
     ----------
-    description_2d : IDSStructure
+    description_2d
         IDS structure with 2D description of the wall.
 
     Returns
@@ -37,7 +37,6 @@ def load_wall_2d(description_2d: IDSStructure) -> dict[str, np.ndarray]:
     dict[str, np.ndarray]
         Dictionary of wall unit outlines given in RZ coordinates.
     """
-
     wall_outline = {}
 
     for unit in description_2d.limiter.unit:
diff --git a/src/cherab/imas/ids/wall/load3d.py b/src/cherab/imas/ids/wall/load3d.py
index 6a9514d..a28ee2f 100644
--- a/src/cherab/imas/ids/wall/load3d.py
+++ b/src/cherab/imas/ids/wall/load3d.py
@@ -36,9 +36,9 @@ def load_wall_3d(
 
     Parameters
     ----------
-    description_ggd : IDSStructure
+    description_ggd
         A description_ggd structure from the 'wall' IDS.
-    subsets : list[str], optional
+    subsets
         List of names of specific ggd subsets to load, by default None (loads all subsets).
 
     Returns
@@ -48,8 +48,12 @@ def load_wall_3d(
         The dictionary keys for components are assigns as follows:
         ``"{grid_name}.{subset_name}.{material_name}"``
         E.g.: ``"FullTokamak.full_main_chamber_wall.Be"``.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the grid_ggd AOS is empty in the given description_ggd.
+    """
     if not len(description_ggd.grid_ggd):
         raise RuntimeError("The grid_ggd AOS is empty in the given description_ggd.")
 
diff --git a/src/cherab/imas/math/functions/__init__.py b/src/cherab/imas/math/functions/__init__.py
index 2faf912..5829fa6 100644
--- a/src/cherab/imas/math/functions/__init__.py
+++ b/src/cherab/imas/math/functions/__init__.py
@@ -15,6 +15,11 @@
 #
 # See the Licence for the specific language governing permissions and limitations
 # under the Licence.
+"""Mathematical function utilities.
+
+This module provides vector function implementations including unit vector classes for 1D, 2D, and
+3D spaces.
+"""
 
 from .vector_functions import UnitVector1D, UnitVector2D, UnitVector3D
 
diff --git a/src/cherab/imas/math/interpolators/__init__.py b/src/cherab/imas/math/interpolators/__init__.py
index ed8a0a3..4acb3bd 100644
--- a/src/cherab/imas/math/interpolators/__init__.py
+++ b/src/cherab/imas/math/interpolators/__init__.py
@@ -15,6 +15,13 @@
 #
 # See the Licence for the specific language governing permissions and limitations
 # under the Licence.
+"""Interpolators for structured and unstructured grid functions.
+
+This module provides interpolation classes for 2D and 3D grid-based functions, supporting both
+structured and unstructured grids. These interpolators are designed to work with IMAS data
+structures and provide efficient evaluation of scalar and vector functions defined on various grid
+types.
+"""
 
 from .struct_grid_2d_functions import StructGridFunction2D, StructGridVectorFunction2D
 from .struct_grid_3d_functions import StructGridFunction3D, StructGridVectorFunction3D
diff --git a/src/cherab/imas/plasma/blend.py b/src/cherab/imas/plasma/blend.py
index 86fdfa9..38c95fa 100644
--- a/src/cherab/imas/plasma/blend.py
+++ b/src/cherab/imas/plasma/blend.py
@@ -67,62 +67,69 @@ def load_plasma(
     parent: _NodeBase | None = None,
     **kwargs,
 ) -> Plasma:
-    """Load core and edge profiles and Create a `~cherab.core.plasma.node.Plasma` object.
+    """Load core and edge profiles and create a `~cherab.core.plasma.node.Plasma` object.
 
-    If ``edge_profiles`` IDS is empty, this returns core plasma only.
-    If ``core_profiles`` IDS is empty, this returns edge plasma only.
+    If the ``edge_profiles`` IDS is empty, returns only the core plasma.
+    If the ``core_profiles`` IDS is empty, returns only the edge plasma.
+
+    To load the edge plasma from a different IMAS entry, use `edge_args` and `edge_kwargs`
+    to pass different arguments to the `~imas.db_entry.DBEntry` constructor.
 
     Parameters
     ----------
     *args
         Arguments passed to the `~imas.db_entry.DBEntry` constructor.
-    time : float, optional
-        Time moment for the core plasma, by default 0.
-    occurrence_core : int, optional
-        Instance index of the 'core_profiles' IDS, by default 0.
-    edge_args : tuple, optional
-        Arguments passed to the `~imas.db_entry.DBEntry` constructor for the edge plasma if different
-        from the core plasma, by default None: use the same as for the core plasma.
-    edge_kwargs : dict, optional
-        Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor for the edge plasma if different
-        from the core plasma, by default None: use the same as for the core plasma.
-    time_edge : float, optional
-        Time moment for the edge plasma if different from the 'time', by default None.
-    occurrence_edge : int, optional
-        Instance index of the 'edge_profiles' IDS, by default 0.
-    grid_ggd : IDSStructure, optional
-        Alternative grid_ggd structure with the grid description, by default None.
-    grid_subset_id : int | str, optional
-        Identifier of the grid subset. Either index or name, by default 5 ("Cells").
-    equilibrium : EFITEquilibrium, optional
-        Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` object used to map core
-        profiles, by default None. equilibrium is read from the same IMAS query as the core profiles.
-        This parameter is ignored if core plasma is not available.
-    b_field : VectorFunction2D, optional
-        Alternative 2D interpolator of the magnetic field vector (Br, Btor, Bz).
-        Default is None. The magnetic field will be loaded from the 'equilibrium' IDS.
-    psi_interpolator : Callable[[float], float], optional
-        Alternative `psi_norm(rho_tor_norm)` interpolator.
-        Used only if 'psi' is missing in the core grid, by default None.
-        Obtained from the 'equilibrium' IDS.
-    mask : Function2D | Function3D, optional
+    time
+        Time for the core plasma, by default 0.
+    occurrence_core
+        Occurrence index of the ``core_profiles`` IDS, by default 0.
+    edge_args
+        Arguments passed to the `~imas.db_entry.DBEntry` constructor for the edge plasma
+        if different from the core plasma. By default None: uses the same as `*args`.
+    edge_kwargs
+        Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor for the edge plasma
+        if different from the core plasma. By default None: uses the same as `**kwargs`.
+    time_edge
+        Time for the edge plasma. If None, uses `time`. By default None.
+    occurrence_edge
+        Occurrence index of the ``edge_profiles`` IDS, by default 0.
+    grid_ggd
+        Alternative ``grid_ggd`` structure describing the grid. By default None.
+    grid_subset_id
+        Identifier of the grid subset (index or name). By default 5 (``"Cells"``).
+    equilibrium
+        Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` used to map core
+        profiles. By default None: the equilibrium is read from the same IMAS query as the
+        core profiles. Ignored if the core plasma is not available.
+    b_field
+        Alternative 2D interpolator of the magnetic field vector (Br, Bphi, Bz).
+        By default None: the magnetic field is loaded from the ``equilibrium`` IDS.
+    psi_interpolator
+        Alternative ``psi_norm(rho_tor_norm)`` interpolator.
+        Used only if ``psi`` is missing in the core grid. By default None.
+        Obtained from the ``equilibrium`` IDS.
+    mask
         Mask function used for blending: ``(1 - mask) * f_edge + mask * f_core``.
-        Default is None. Use `EFITEquilibrium.inside_lcfs` as a mask function.
-    time_threshold : float, optional
-        Sets the maximum allowable difference between the specified time and the nearest
+        By default, uses `~cherab.tools.equilibrium.efit.EFITEquilibrium`'s `inside_lcfs`.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest
         available time, by default `numpy.inf`.
-    parent : _NodeBase, optional
-        Parent node in the Raysect scene-graph, by default None.
-        Normally, `~raysect.optical.scenegraph.world.World` instance.
+    parent
+        Parent node in the Raysect scene graph, by default None.
+        Typically a `~raysect.optical.scenegraph.world.World` instance.
     **kwargs
         Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor.
 
     Returns
     -------
     `~cherab.core.plasma.node.Plasma`
-        Plasma object with core and edge profiles.
-    """
+        Plasma object with core and/or edge profiles.
 
+    Raises
+    ------
+    RuntimeError
+        If neither core nor edge profiles are available.
+    """
     edge_args = edge_args or args
     edge_kwargs = edge_kwargs or kwargs
     if time_edge is None:
@@ -355,13 +362,13 @@ def blend_core_edge_interpolators(
 
     Parameters
     ----------
-    core_interpolators : dict
+    core_interpolators
         Dictionary with 2D or 3D core profiles interpolators.
-    edge_interpolators : dict
+    edge_interpolators
         Dictionary with 2D or 3D edge profiles interpolators.
-    mask : Function2D | Function3D
+    mask
         Mask function used for blending: ``(1 - mask) * f_edge + mask * f_core``.
-    return3d : bool, optional
+    return3d
         If True, return the 3D functions for 2D interpolators assuming
         rotational symmetry, by default False.
 
@@ -370,7 +377,6 @@ def blend_core_edge_interpolators(
     dict
         Dictionary with blended interpolators.
     """
-
     interpolators = {}
 
     for core_key, core_func in core_interpolators.items():
@@ -394,13 +400,13 @@ def blend_core_edge_functions(
 
     Parameters
     ----------
-    core_func : Function2D | Function3D | VectorFunction2D | VectorFunction3D | None
+    core_func
         A 2D or 3D core interpolator.
-    edge_func : Function2D | Function3D | VectorFunction2D | VectorFunction3D | None
+    edge_func
         A 2D or 3D edge interpolator.
-    mask : Function2D | Function3D
-        The 2D or 3D mask function used for blending: (1 - mask) * f_edge + mask * f_core.
-    return3d : bool, optional
+    mask
+        The 2D or 3D mask function used for blending: ``(1 - mask) * f_edge + mask * f_core``.
+    return3d
         If True, return the 3D functions for 2D interpolators assuming
         rotational symmetry, by default False.
 
@@ -408,8 +414,12 @@ def blend_core_edge_functions(
     -------
     Function3D | VectorFunction3D | None
         Blended interpolator.
-    """
 
+    Raises
+    ------
+    ValueError
+        If both core and edge functions are None.
+    """
     if core_func is None and edge_func is None:
         return None
 
diff --git a/src/cherab/imas/plasma/core.py b/src/cherab/imas/plasma/core.py
index 79c2c96..7f9c37e 100644
--- a/src/cherab/imas/plasma/core.py
+++ b/src/cherab/imas/plasma/core.py
@@ -24,6 +24,7 @@
 from raysect.core.math.function.float import Function2D, Function3D, Interpolator1DArray
 from raysect.core.math.function.vector3d import Constant3D as ConstantVector3D
 from raysect.core.math.function.vector3d import Function2D as VectorFunction2D
+from raysect.core.scenegraph._nodebase import _NodeBase
 from raysect.primitive import Cylinder, Subtract
 from scipy.constants import atomic_mass, electron_mass
 
@@ -49,38 +50,41 @@ def load_core_plasma(
     b_field: VectorFunction2D | None = None,
     psi_interpolator: Callable[[float], float] | None = None,
     time_threshold: float = np.inf,
-    parent=None,
+    parent: _NodeBase | None = None,
     **kwargs,
 ) -> Plasma:
     """Load core profiles and Create a `~cherab.core.plasma.node.Plasma` object.
 
-    Prefer 'density_thermal' over 'density' profile.
+    Prefer ``density_thermal`` over ``density`` profile.
 
     Parameters
     ----------
     *args
         Arguments passed to the `~imas.db_entry.DBEntry` constructor.
-    time : float, optional
-        Time moment, by default 0.
-    occurrence : int, optional
-        Instance index of the 'core_profiles' IDS, by default 0.
-    equilibrium : EFITEquilibrium, optional
-        Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` object used to map core
-        profiles, by default None. equilibrium is read from the same IMAS query as the core profiles.
-        This parameter is ignored if core plasma is not available.
-    b_field : VectorFunction2D, optional
-        An alternative 2D interpolator of the magnetic field vector (Br, Btor, Bz).
-        Default is None. The magnetic field will be loaded from the 'equilibrium' IDS.
-    psi_interpolator : Callable[[float], float], optional
-        Alternative `psi_norm(rho_tor_norm)` interpolator.
-        Used only if 'psi' is missing in the core grid, by default None.
-        Obtained from the 'equilibrium' IDS.
-    time_threshold : float, optional
-        Sets the maximum allowable difference between the specified time and the nearest
+    time
+        Time for the core plasma, by default 0.
+    occurrence
+        Occurrence index of the ``core_profiles`` IDS, by default 0.
+    equilibrium
+        Alternative `~cherab.tools.equilibrium.efit.EFITEquilibrium` used to map core
+        profiles. By default None: the equilibrium is read from the same IMAS query as the
+        core profiles. Ignored if the core plasma is not available.
+    b_field
+        Alternative 2D interpolator of the magnetic field vector (Br, Bphi, Bz).
+        By default None: the magnetic field is loaded from the ``equilibrium`` IDS.
+    psi_interpolator
+        Alternative ``psi_norm(rho_tor_norm)`` interpolator.
+        Used only if ``psi`` is missing in the core grid. By default None.
+        Obtained from the ``equilibrium`` IDS.
+    mask
+        Mask function used for blending: ``(1 - mask) * f_edge + mask * f_core``.
+        By default, uses `~cherab.tools.equilibrium.efit.EFITEquilibrium`'s `inside_lcfs`.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest
         available time, by default `numpy.inf`.
-    parent : _NodeBase, optional
-        Parent node in the Raysect scene-graph, by default None.
-        Normally, `~raysect.optical.scenegraph.world.World` instance.
+    parent
+        Parent node in the Raysect scene graph, by default None.
+        Typically a `~raysect.optical.scenegraph.world.World` instance.
     **kwargs
         Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor.
 
@@ -88,8 +92,15 @@ def load_core_plasma(
     -------
     `~cherab.core.plasma.node.Plasma`
         Plasma object with core profiles.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the ``profiles_1d`` AOS in core_profiles IDS is empty.
+    ValueError
+        If the ``equilibrium`` argument is not an `~cherab.tools.equilibrium.efit.EFITEquilibrium`
+        instance when provided.
+    """
     with DBEntry(*args, **kwargs) as entry:
         core_profiles_ids = get_ids_time_slice(
             entry, "core_profiles", time=time, occurrence=occurrence, time_threshold=time_threshold
@@ -204,23 +215,28 @@ def get_core_interpolators(
 
     Parameters
     ----------
-    psi_norm : ndarray
+    psi_norm
         Normalized poloidal flux values.
-    profiles : dict
+    profiles
        Dictionary with core plasma profiles.
-    equilibrium : EFITEquilibrium
+    equilibrium
         `EFITEquilibrium` object used to map core profiles.
-    return3d : bool, optional
+    return3d
         If True, return the 3D interpolators assuming rotational symmetry, by default False.
 
     Returns
     -------
     dict[str, Function3D | Function2D | None]
        Dictionary with core interpolators.
-    """
 
+    Raises
+    ------
+    ValueError
+        If the ``equilibrium`` argument is not an `~cherab.tools.equilibrium.efit.EFITEquilibrium`
+        instance.
+    """
     if not isinstance(equilibrium, EFITEquilibrium):
-        raise ValueError("Argiment equilibrium must be a EFITEquilibrium instance.")
+        raise ValueError("Argument equilibrium must be a EFITEquilibrium instance.")
 
     psi_norm, index = np.unique(psi_norm, return_index=True)
 
@@ -252,22 +268,28 @@ def get_psi_norm(
 
     Parameters
     ----------
-    psi : ndarray | None
+    psi
         Poloidal flux values from the core grid.
-    psi_axis : float
+    psi_axis
         Poloidal flux at the magnetic axis.
-    psi_lcfs : float
+    psi_lcfs
         Poloidal flux at the last closed flux surface.
-    rho_tor_norm : ndarray | None
+    rho_tor_norm
         Normalized toroidal flux values.
-    psi_interpolator : Callable[[float], float] | None
-        Interpolator function to map rho_tor_norm to psi_norm.
-        Used only if 'psi' is None.
+    psi_interpolator
+        Interpolator function to map `rho_tor_norm` to `psi_norm`.
+        Used only if ``psi`` is None.
 
     Returns
     -------
     ndarray
         Normalized poloidal flux values.
+
+    Raises
+    ------
+    RuntimeError
+        If both ``psi`` and ``rho_tor_norm`` are None, or if ``psi_interpolator`` is None when
+        ``psi`` is None.
     """
     if psi is None:
         if psi_interpolator is None:
diff --git a/src/cherab/imas/plasma/edge.py b/src/cherab/imas/plasma/edge.py
index 9a21ca4..22773e6 100644
--- a/src/cherab/imas/plasma/edge.py
+++ b/src/cherab/imas/plasma/edge.py
@@ -62,23 +62,23 @@ def load_edge_plasma(
     ----------
     *args
         Arguments passed to the `~imas.db_entry.DBEntry` constructor.
-    time : float, optional
-        Time moment for the edge plasma, by default 0.
-    occurrence : int, optional
-        Instance index of the 'edge_profiles' IDS, by default 0.
-    grid_ggd : IDSStructure, optional
-        Alternative grid_ggd structure with the grid description, by default None.
-    grid_subset_id : int | str, optional
-        Identifier of the grid subset. Either index or name, by default 5 ("Cells").
-    b_field : VectorFunction2D, optional
-        Alternative 2D interpolator of the magnetic field vector (Br, Btor, Bz).
-        Default is None. The magnetic field will be loaded from the 'equilibrium' IDS.
-    time_threshold : float, optional
-        Sets the maximum allowable difference between the specified time and the nearest
+    time
+        Time for the edge plasma, by default 0.
+    occurrence
+        Occurrence index of the ``edge_profiles`` IDS, by default 0.
+    grid_ggd
+        Alternative ``grid_ggd`` structure with the grid description, by default None.
+    grid_subset_id
+        Identifier of the grid subset. Either index or name, by default 5 (``"Cells"``).
+    b_field
+        Alternative 2D interpolator of the magnetic field vector (Br, Bphi, Bz).
+        Default is None. The magnetic field will be loaded from the ``equilibrium`` IDS.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest
         available time, by default `numpy.inf`.
-    parent : _NodeBase, optional
-        Parent node in the Raysect scene-graph, by default None.
-        Normally, `~raysect.optical.scenegraph.world.World` instance.
+    parent
+        Parent node in the Raysect scene graph, by default None.
+        Typically a `~raysect.optical.scenegraph.world.World` instance.
     **kwargs
         Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor.
 
@@ -86,8 +86,12 @@ def load_edge_plasma(
     -------
     `~cherab.core.plasma.node.Plasma`
         Plasma object with edge profiles.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the ``grid_ggd`` or ``ggd`` AOS of the edge_profiles IDS is empty.
+    """
     with DBEntry(*args, **kwargs) as entry:
         edge_profiles_ids = get_ids_time_slice(
             entry, "edge_profiles", time=time, occurrence=occurrence, time_threshold=time_threshold
@@ -197,13 +201,13 @@ def get_edge_interpolators(
 
     Parameters
     ----------
-    grid : GGDGrid
+    grid
         GGD-compatible grid object.
-    profiles : dict[str, np.ndarray | None]
+    profiles
         Dictionary with edge plasma profiles.
-    b_field : VectorFunction2D, optional
+    b_field
         2D interpolator of the magnetic field vector (Br, Btor, Bz), by default None.
-    return3d : bool, optional
+    return3d
         If True, return the 3D functions for 2D interpolators assuming
         rotational symmetry, by default False.
 
@@ -212,7 +216,6 @@ def get_edge_interpolators(
     dict[str, Function3D | Function2D | None]
         Dictionary with edge interpolators.
     """
-
     interpolators = RecursiveDict()
 
     for prof_key, profile in profiles.items():
diff --git a/src/cherab/imas/plasma/equilibrium.py b/src/cherab/imas/plasma/equilibrium.py
index 9c16004..4e725c9 100644
--- a/src/cherab/imas/plasma/equilibrium.py
+++ b/src/cherab/imas/plasma/equilibrium.py
@@ -44,14 +44,14 @@ def load_equilibrium(
     ----------
     *args
         Arguments passed to `~imas.db_entry.DBEntry`.
-    time : float, optional
-        Time moment, by default 0.
-    occurrence : int, optional
-        Instance index of the 'equilibrium' IDS, by default 0.
-    time_threshold : float, optional
-        Maximum allowable difference between the specified time and the nearest available
+    time
+        Time for the equilibrium, by default 0.
+    occurrence
+        Occurrence index of the ``equilibrium`` IDS, by default 0.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest available
         time, by default `numpy.inf`.
-    with_psi_interpolator : bool, optional
+    with_psi_interpolator
         If True, returns the ``psi_norm(rho_tor_norm)`` interpolator; otherwise, returns only the
         equilibrium object.
     **kwargs
@@ -59,9 +59,9 @@ def load_equilibrium(
 
     Returns
     -------
-    equilibrium : EFITEquilibrium
-        The `EFITEquilibrium` object.
-    psi_interpolator : Function1D, optional
+    equilibrium : `~cherab.tools.equilibrium.efit.EFITEquilibrium`
+        The plasma equilibrium object.
+    psi_interpolator : `~raysect.core.math.function.float.function1d.interpolate.Interpolator1DArray` | None
         If ``with_psi_interpolator`` is True and ``rho_tor_norm`` is available, returns the
         ``psi_norm(rho_tor_norm)`` interpolator.
         If rho_tor_norm is not available, returns None.
@@ -126,22 +126,26 @@ def load_magnetic_field(
     ----------
     *args
         Arguments passed to `~imas.db_entry.DBEntry`.
-    time : float, optional
-        Time moment, by default 0.
-    occurrence : int, optional
-        Instance index of the 'equilibrium' IDS, by default 0.
-    time_threshold : float, optional
-        Sets the maximum allowable difference between the specified time and the nearest available
+    time
+        Time for the equilibrium, by default 0.
+    occurrence
+        Occurrence index of the ``equilibrium`` IDS, by default 0.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest available
         time, by default `numpy.inf`.
     **kwargs
         Keyword arguments passed to `~imas.db_entry.DBEntry`.
 
     Returns
     -------
-    VectorFunction2D
+    `~raysect.core.math.function.vector3d.function2d.base.Function2D`
         The magnetic field interpolator.
-    """
 
+    Raises
+    ------
+    RuntimeError
+        If the equilibrium IDS does not have a time slice.
+    """
     with DBEntry(*args, **kwargs) as entry:
         equilibrium_ids = get_ids_time_slice(
             entry, "equilibrium", time=time, occurrence=occurrence, time_threshold=time_threshold
@@ -169,10 +173,9 @@ def cocos_11to3(equilibrium_dict: dict) -> None:
 
     Parameters
     ----------
-    equilibrium_dict : dict
+    equilibrium_dict
         The equilibrium data dictionary to modify in place.
     """
-
     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)
diff --git a/src/cherab/imas/plasma/utility.py b/src/cherab/imas/plasma/utility.py
index 6bed7af..e28a3f3 100644
--- a/src/cherab/imas/plasma/utility.py
+++ b/src/cherab/imas/plasma/utility.py
@@ -20,14 +20,14 @@
 __all__ = ["warn_unsupported_species", "get_subset_name_index"]
 
 
-def warn_unsupported_species(composition: dict[str, dict], species_type: str):
+def warn_unsupported_species(composition: dict[str, dict], species_type: str) -> None:
     """Warn if species of a given type are present in the composition dictionary.
 
     Parameters
     ----------
-    composition : dict[str, dict]
+    composition
         Dictionary with species composition.
-    species_type : str
+    species_type
         Type of species to check for (e.g., 'ion_bundle', 'molecular_bundle').
     """
     if species_type in composition and len(composition[species_type]):
@@ -47,17 +47,22 @@ def get_subset_name_index(subset_id_dict: dict, grid_subset_id: int | str) -> tu
 
     Parameters
     ----------
-    subset_id_dict : dict
+    subset_id_dict
         Dictionary with grid subset indices.
-    grid_subset_id : int | str
+    grid_subset_id
         Identifier of the grid subset. Either index or name.
 
     Returns
     -------
-    grid_subset_name : str
+    grid_subset_name
         Name of the grid subset.
-    grid_subset_index : int
+    grid_subset_index
         Index of the grid subset.
+
+    Raises
+    ------
+    ValueError
+        If the grid subset with the given identifier is not found.
     """
     subset_id = subset_id_dict.copy()
     subset_id.update({value: key for key, value in subset_id.items()})
diff --git a/src/cherab/imas/wall/wall.py b/src/cherab/imas/wall/wall.py
index 20575ef..5cb6ba5 100644
--- a/src/cherab/imas/wall/wall.py
+++ b/src/cherab/imas/wall/wall.py
@@ -41,33 +41,32 @@ def load_wall_mesh(
     parent: _NodeBase | None = None,
     **kwargs,
 ) -> dict[str, Mesh]:
-    """Load machine wall components from IMAS wall IDS and Create a dictionary of Raysect mesh
-    primitives.
+    """Load machine wall components from IMAS wall IDS and Create Raysect mesh primitives.
 
     Parameters
     ----------
     *args
         Arguments passed to the `~imas.db_entry.DBEntry` constructor.
-    time : float, optional
-        Time moment for the edge plasma, by default 0.
-    occurrence : int, optional
-        Instance index of the 'wall' IDS, by default 0.
-    desc_ggd_index : int
-        Index of 'description_ggd', by default 0.
-    subsets : list[str], optional
+    time
+        Time for the wall, by default 0.
+    occurrence
+        Occurrence index of the ``wall`` IDS, by default 0.
+    desc_ggd_index
+        Index of ``description_ggd``, by default 0.
+    subsets
         List of names of specific ggd subsets to load, by default None (loads all subsets).
-    materials : dict[str, Material], optional
+    materials
         Optional dictionary with Raysect materials for each wall component, by default None.
         Use component names as keys. The components are split by their grid subsets and for
         each grid subset by materials. All elements of the grid subset that share the same material
         are combined into a single component. The component names are assigns as follows:
         ``"{grid_name}.{subset_name}.{material_name}"``
         E.g.: ``"TokamakWall.full_main_chamber_wall.Be"``.
-    time_threshold : float, optional
-        Sets the maximum allowable difference between the specified time and the nearest
-        available time, by default `numpy.inf`.
-    parent : _NodeBase, optional
-        Parent node in the Raysect scene-graph, by default None.
+    time_threshold
+        Maximum allowed difference between the requested time and the nearest available
+        time, by default `numpy.inf`.
+    parent
+        Parent node in the Raysect scene graph, by default None.
         Normally, `~raysect.optical.scenegraph.world.World` instance.
     **kwargs
         Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor.
@@ -82,12 +81,11 @@ def load_wall_mesh(
     >>> from raysect.optical import World
     >>> world = World()
     >>> meshes = load_wall_mesh(
-    ...     "imas:uda?path=/work/imas/shared/imasdb/ITER_MD/3/116100/1001/", "r", parent=world
+    ...     "imas:hdf5?path=/work/imas/shared/imasdb/ITER_MD/3/116100/1001/", "r", parent=world
     ... )
     >>> meshes
     {'FullTokamak.none.none': <raysect.primitive.mesh.mesh.Mesh at 0x1766322a0>}
     """
-
     with DBEntry(*args, **kwargs) as entry:
         wall_ids = get_ids_time_slice(
             entry, "wall", time=time, occurrence=occurrence, time_threshold=time_threshold
@@ -117,21 +115,21 @@ def load_wall_outline(
     ----------
     *args
         Arguments passed to the `~imas.db_entry.DBEntry` constructor.
-    occurrence : int, optional
-        Instance index of the 'wall' IDS, by default 0.
-    desc_index : int, optional
-        Index of 'description_2d', by default 0.
+    occurrence
+        Occurrence index of the ``wall`` IDS, by default 0.
+    desc_index
+        Index of ``description_2d``, by default 0.
     **kwargs
         Keyword arguments passed to the `~imas.db_entry.DBEntry` constructor.
 
     Returns
     -------
     dict[str, (N, 2) ndarray]
-       Dictionary of wall unit outlines (N, 2) array given in RZ coordinates.
+       Dictionary of wall unit outlines ``(N, 2)`` array given in RZ coordinates.
 
     Examples
     --------
-    >>> load_wall_outline("imas:uda?path=/work/imas/shared/imasdb/ITER_MD/3/116000/5/", "r")
+    >>> load_wall_outline("imas:hdf5?path=/work/imas/shared/imasdb/ITER_MD/3/116000/5/", "r")
     {'First Wall': array([[ 4.11129713, -2.49559808],
                           [ 4.11129713, -1.48329401],
                           ...
@@ -140,7 +138,6 @@ def load_wall_outline(
                         ...
                         [ 6.36320019, -3.24460006]])}
     """
-
     with DBEntry(*args, **kwargs) as entry:
         description2d = entry.get("wall", occurrence=occurrence, autoconvert=False).description_2d[
             desc_index
diff --git a/tests/plasma/test_blend.py b/tests/plasma/test_blend.py
index 302c1f7..9820e32 100644
--- a/tests/plasma/test_blend.py
+++ b/tests/plasma/test_blend.py
@@ -38,7 +38,6 @@ def test_load_plasma_electron_distribution(path_iter_jintrac: str):
 
 def test_load_plasma_with_time(path_iter_jintrac: str):
     """Test loading plasma with specific time parameter."""
-
     # Test with default time (should not raise error)
     plasma1 = load_plasma(path_iter_jintrac, "r", time=0.0)
     assert isinstance(plasma1, Plasma)
@@ -50,7 +49,6 @@ def test_load_plasma_with_time(path_iter_jintrac: str):
 
 def test_load_plasma_with_occurrence(path_iter_jintrac: str):
     """Test loading plasma with different occurrence values."""
-
     # Test with default occurrence core
     plasma1 = load_plasma(path_iter_jintrac, "r", occurrence_core=0)
     assert isinstance(plasma1, Plasma)
@@ -75,7 +73,6 @@ def test_load_plasma_with_parent(path_iter_jintrac: str):
 
 def test_load_plasma_with_magnetic_field(path_iter_jintrac: str):
     """Test loading plasma with external magnetic field."""
-
     # Load magnetic field separately
     try:
         b_field = load_magnetic_field(path_iter_jintrac, "r")
@@ -90,7 +87,6 @@ def test_load_plasma_with_magnetic_field(path_iter_jintrac: str):
 
 def test_load_plasma_time_threshold(path_iter_jintrac: str):
     """Test loading plasma with time threshold parameter."""
-
     # Test with large time threshold (should work)
     plasma = load_plasma(path_iter_jintrac, "r", time_threshold=500)
     assert isinstance(plasma, Plasma)