Cherab interface for xBOUT

Raytracing library Cherab (https://www.cherab.info/)

For now only works for axisymmetric grids (x, theta).
Creates a triangulation of the BOUT++ grid, then can use that
to create a volumetric emitter from DataArrays.

Author: bendudson

pyproject.toml

@@ -53,6 +53,9 @@ calc = [
     "xrft",
     "xhistogram",
 ]
+cherab = [
+    "cherab",
+]
 docs = [
     "sphinx >= 5.3",
     "sphinx-book-theme >= 0.4.0rc1",

xbout/boutdataarray.py

@@ -1101,3 +1101,71 @@ def plot3d(self, ax=None, **kwargs):
         See plotfuncs.plot3d()
         """
         return plotfuncs.plot3d(self.data, **kwargs)
+
+    def with_cherab_grid(self):
+        """
+        Returns a new DataArray with a 'cherab_grid' attribute.
+
+        If called then the `cherab` package must be available.
+        """
+        # Import here so Cherab is required only if this method is called
+        from .cherab import grid
+
+        return grid.da_with_cherab_grid(self.data)
+
+    def as_cherab_data(self):
+        """
+        Returns a new cherab.TriangularData object.
+
+        If a Cherab grid has not been calculated then it will be created.
+        It is more efficient to first compute a Cherab grid for a whole
+        DataSet (using `with_cherab_grid`) and then call this function
+        on individual DataArrays.
+        """
+        if "cherab_grid" not in self.data.attrs:
+            # Calculate the Cherab triangulation
+            da = self.with_cherab_grid()
+        else:
+            da = self.data
+
+        return da.attrs["cherab_grid"].with_data(da)
+
+    def as_cherab_emitter(
+        self,
+        parent=None,
+        cylinder_zmin=None,
+        cylinder_zmax=None,
+        cylinder_rmin=None,
+        cylinder_rmax=None,
+        step: float = 0.01,
+    ):
+        """
+        Make a Cherab emitter (RadiationFunction), rotating a 2D mesh about the Z axis
+
+        Cherab (https://www.cherab.info/) is a python library for forward
+        modelling diagnostics based on spectroscopic plasma emission.
+        It is based on the Raysect (http://www.raysect.org/) scientific
+        ray-tracing framework.
+
+        Parameters
+        ----------
+        parent : Cherab scene (default None)
+            The Cherab scene to attach the emitter to
+        step : float (default 0.01 meters)
+            Volume integration step length [m]
+
+        Returns
+        -------
+
+        A cherab.tools.emitters.RadiationFunction
+
+        """
+
+        return self.as_cherab_data().to_emitter(
+            parent=parent,
+            cylinder_zmin=cylinder_zmin,
+            cylinder_zmax=cylinder_zmax,
+            cylinder_rmin=cyliner_rmin,
+            cylinder_rmax=cyliner_rmax,
+            step=step,
+        )

xbout/boutdataset.py

@@ -609,14 +609,12 @@ def interpolate_to_cartesian(
         newY = newY_1d.expand_dims({"X": nX, "Z": nZ}, axis=[0, 2])
         newZ_1d = xr.DataArray(np.linspace(Zmin, Zmax, nZ), dims="Z")
         newZ = newZ_1d.expand_dims({"X": nX, "Y": nY}, axis=[0, 1])
-        newR = np.sqrt(newX**2 + newY**2)
+        newR = np.sqrt(newX ** 2 + newY ** 2)
         newzeta = np.arctan2(newY, newX)
         # Define newzeta in range 0->2*pi
         newzeta = np.where(newzeta < 0.0, newzeta + 2.0 * np.pi, newzeta)
 
-        from scipy.interpolate import (
-            RegularGridInterpolator,
-        )
+        from scipy.interpolate import RegularGridInterpolator
 
         # Create Cylindrical coordinates for intermediate grid
         Rcyl_min = float_type(ds["R"].min())
@@ -664,10 +662,7 @@ def interp_single_time(da):
             )
 
             print("    do 3d interpolation")
-            return interp(
-                (newR, newZ, newzeta),
-                method="linear",
-            )
+            return interp((newR, newZ, newzeta), method="linear")
 
         for name, da in ds.data_vars.items():
             print(f"\ninterpolating {name}")
@@ -993,14 +988,7 @@ def to_restart(
         # Is this even possible without saving the guard cells?
         # Can they be recreated?
         restart_datasets, paths = _split_into_restarts(
-            self.data,
-            variables,
-            savepath,
-            nxpe,
-            nype,
-            tind,
-            prefix,
-            overwrite,
+            self.data, variables, savepath, nxpe, nype, tind, prefix, overwrite
         )
 
         with ProgressBar():
@@ -1357,6 +1345,17 @@ def is_list(variable):
 
         return anim
 
+    def with_cherab_grid(self):
+        """
+        Returns a new DataSet with a 'cherab_grid' attribute.
+
+        If called then the `cherab` package must be available.
+        """
+        # Import here so Cherab is required only if this method is called
+        from .cherab import grid
+
+        return grid.ds_with_cherab_grid(self.data)
+
 
 def _find_major_vars(data):
     """

xbout/cherab/__init__.py

@@ -0,0 +1,93 @@
+import numpy as np
+import xarray as xr
+
+from .triangulate import Triangulate
+
+
+def da_with_cherab_grid(da):
+    """
+    Convert an BOUT++ DataArray to a format that Cherab can use:
+    - A 'cell_number' coordinate
+    - A 'cherab_grid` attribute 
+
+    The cell_number coordinate enables the DataArray to be sliced
+    before input to Cherab.
+
+    Parameters
+    ----------
+    ds : xarray.DataArray
+
+    Returns
+    -------
+    updated_da
+    """
+    if "cherab_grid" in da.attrs:
+        # Already has required attribute
+        return da
+
+    if da.attrs["geometry"] != "toroidal":
+        raise ValueError("xhermes.plotting.cherab: Geometry must be toroidal")
+
+    if da.sizes["zeta"] != 1:
+        raise ValueError("xhermes.plotting.cherab: Zeta index must have size 1")
+
+    nx = da.sizes["x"]
+    ny = da.sizes["theta"]
+
+    # Cell corners
+    rm = np.stack(
+        (
+            da.coords["Rxy_upper_right_corners"],
+            da.coords["Rxy_upper_left_corners"],
+            da.coords["Rxy_lower_right_corners"],
+            da.coords["Rxy_lower_left_corners"],
+        ),
+        axis=-1,
+    )
+    zm = np.stack(
+        (
+            da.coords["Zxy_upper_right_corners"],
+            da.coords["Zxy_upper_left_corners"],
+            da.coords["Zxy_lower_right_corners"],
+            da.coords["Zxy_lower_left_corners"],
+        ),
+        axis=-1,
+    )
+
+    grid = Triangulate(rm, zm)
+
+    # Store the cell number as a coordinate.
+    # This allows slicing of arrays before passing to Cherab
+
+    # Create a DataArray for the vertices and triangles
+    cell_number = xr.DataArray(
+        grid.cell_number, dims=["x", "theta"], name="cell_number"
+    )
+
+    result = da.assign_coords(cell_number=cell_number)
+    result.attrs.update(cherab_grid=grid)
+    return result
+
+
+def ds_with_cherab_grid(ds):
+    """
+    Create an xarray DataSet with a Cherab grid
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+
+    Returns
+    -------
+    updated_ds
+    """
+
+    # The same operation works on a DataSet
+    ds = da_with_cherab_grid(ds)
+
+    # Add the Cherab grid as an attribute to all variables
+    grid = ds.attrs["cherab_grid"]
+    for var in ds.data_vars:
+        ds[var].attrs.update(cherab_grid=grid)
+
+    return ds