Merge branch 'master' into fix-gridfile-loading

Conflicts:
	xbout/load.py

Author: bendudson

.github/workflows/master.yml

@@ -20,7 +20,7 @@ jobs:
     if: always()
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.x"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
       fail-fast: false
 
     steps:
@@ -61,7 +61,7 @@ jobs:
       run: |
         sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev
         python -m pip install --upgrade pip
-        pip install xarray~=0.18.0 pandas~=1.4.0
+        pip install xarray~=2023.1.0 pandas~=1.4.0
     - name: Install package
       run: |
         pip install -e .[tests]

.github/workflows/pythonpackage.yml

@@ -31,7 +31,6 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev
         python -m pip install --upgrade pip
     - name: Install package
       run: |
@@ -59,9 +58,8 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        sudo apt-get update && sudo apt-get install libhdf5-dev libnetcdf-dev
         python -m pip install --upgrade pip
-        pip install xarray~=0.18.0 pandas~=1.4.0
+        pip install xarray~=2023.1.0 pandas~=1.4.0
     - name: Install package
       run: |
         pip install -e .[tests]

.github/workflows/ruff.yaml

@@ -5,11 +5,11 @@ jobs:
   dependency-review:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: '3.x'
       - name: Install ruff
         run: pip install ruff
       - name: Run ruff
-        run: ruff xbout
+        run: ruff check xbout

pyproject.toml

@@ -33,7 +33,7 @@ classifiers = [
 ]
 requires-python = ">=3.8"
 dependencies = [
-    "xarray>=0.18.0,<2022.9.0",
+    "xarray>=2023.01.0",
     "boutdata>=0.1.4",
     "dask[array]>=2.10.0",
     "gelidum>=0.5.3",
@@ -53,6 +53,9 @@ calc = [
     "xrft",
     "xhistogram",
 ]
+cherab = [
+    "cherab",
+]
 docs = [
     "sphinx >= 5.3",
     "sphinx-book-theme >= 0.4.0rc1",
@@ -79,5 +82,5 @@ write_to = "xbout/_version.py"
 [tool.setuptools]
 packages = ["xbout"]
 
-[tool.ruff]
+[tool.ruff.lint]
 ignore = ["E501"]

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

@@ -597,12 +597,12 @@ def interpolate_to_cartesian(
         n_toroidal = ds.sizes[zdim]
 
         # Create Cartesian grid to interpolate to
-        Xmin = ds["X_cartesian"].min()
-        Xmax = ds["X_cartesian"].max()
-        Ymin = ds["Y_cartesian"].min()
-        Ymax = ds["Y_cartesian"].max()
-        Zmin = ds["Z_cartesian"].min()
-        Zmax = ds["Z_cartesian"].max()
+        Xmin = ds["X_cartesian"].min().data[()]
+        Xmax = ds["X_cartesian"].max().data[()]
+        Ymin = ds["Y_cartesian"].min().data[()]
+        Ymax = ds["Y_cartesian"].max().data[()]
+        Zmin = ds["Z_cartesian"].min().data[()]
+        Zmax = ds["Z_cartesian"].max().data[()]
         newX_1d = xr.DataArray(np.linspace(Xmin, Xmax, nX), dims="X")
         newX = newX_1d.expand_dims({"Y": nY, "Z": nZ}, axis=[1, 2])
         newY_1d = xr.DataArray(np.linspace(Ymin, Ymax, nY), dims="Y")
@@ -614,9 +614,7 @@ def interpolate_to_cartesian(
         # 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