diff --git a/pyproject.toml b/pyproject.toml index 39e926e8..0caa3826 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,9 @@ calc = [ "xrft", "xhistogram", ] +cherab = [ + "cherab", +] docs = [ "sphinx >= 5.3", "sphinx-book-theme >= 0.4.0rc1", diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 2a2c5ed6..3e54f402 100644 --- a/xbout/boutdataarray.py +++ b/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, + ) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index cc308fcb..e8b8a0a6 100644 --- a/xbout/boutdataset.py +++ b/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): """ diff --git a/xbout/cherab/__init__.py b/xbout/cherab/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xbout/cherab/grid.py b/xbout/cherab/grid.py new file mode 100644 index 00000000..94e48cf8 --- /dev/null +++ b/xbout/cherab/grid.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 diff --git a/xbout/cherab/triangulate.py b/xbout/cherab/triangulate.py new file mode 100644 index 00000000..5299ea0d --- /dev/null +++ b/xbout/cherab/triangulate.py @@ -0,0 +1,231 @@ +""" +Interface to Cherab + +This module performs triangulation of BOUT++ grids, making them +suitable for input to Cherab ray-tracing analysis. + +""" + +import numpy as np +import xarray as xr + + +class TriangularData: + """ + Represents a set of triangles with data constant on them. + Creates a Cherab Discrete2DMesh, and can then convert + that to a 3D (axisymmetric) emitting material. + """ + + def __init__(self, vertices, triangles, data): + self.vertices = vertices + self.triangles = triangles + self.data = data + + from raysect.core.math.function.float import Discrete2DMesh + + self.mesh = Discrete2DMesh( + self.vertices, self.triangles, self.data, limit=False, default_value=0.0 + ) + + def to_emitter( + self, + parent=None, + cylinder_zmin=None, + cylinder_zmax=None, + cylinder_rmin=None, + cylinder_rmax=None, + step: float = 0.01, + ): + """ + Returns a 3D Cherab emitter, by rotating the 2D mesh about the Z axis + + step: Volume integration step length [m] + + """ + from raysect.core import translate + from raysect.primitive import Cylinder, Subtract + from raysect.optical.material import VolumeTransform + from cherab.core.math import AxisymmetricMapper + from cherab.tools.emitters import RadiationFunction + + if cylinder_zmin is None: + cylinder_zmin = np.amin(self.vertices[:, 1]) + if cylinder_zmax is None: + cylinder_zmax = np.amax(self.vertices[:, 1]) + if cylinder_rmin is None: + cylinder_rmin = np.amin(self.vertices[:, 0]) + if cylinder_rmax is None: + cylinder_rmax = np.amax(self.vertices[:, 0]) + + rad_function_3d = AxisymmetricMapper(self.mesh) + + shift = translate(0, 0, cylinder_zmin) + emitting_material = VolumeTransform( + RadiationFunction(rad_function_3d, step=step), shift.inverse() + ) + + # Create an annulus by removing the middle from the cylinder. + return Subtract( + Cylinder(cylinder_rmax, cylinder_zmax - cylinder_zmin), + Cylinder(cylinder_rmin, cylinder_zmax - cylinder_zmin), + transform=shift, + parent=parent, + material=emitting_material, + ) + + def plot_2d(self, ax=None, nr: int = 150, nz: int = 150): + """ + Make a 2D plot of the data + + nr, nz - Number of samples in R and Z + """ + import matplotlib.pyplot as plt + + if ax is None: + fig, ax = plt.subplots() + + Rmin, Zmin = np.amin(self.vertices, axis=0) + Rmax, Zmax = np.amax(self.vertices, axis=0) + + from cherab.core.math import sample2d + + r, z, emiss_sampled = sample2d(self.mesh, (Rmin, Rmax, nr), (Zmin, Zmax, nz)) + + image = ax.imshow( + emiss_sampled.T, origin="lower", extent=(r.min(), r.max(), z.min(), z.max()) + ) + fig.colorbar(image) + ax.set_xlabel("r") + ax.set_ylabel("z") + + return ax + + +class Triangulate: + """ + Represents a set of triangles for a 2D mesh in R-Z + + """ + + def __init__(self, rm, zm): + """ + rm and zm define quadrilateral cell corners in 2D (R, Z) + + rm : [nx, ny, 4] + zm : [nx, ny, 4] + """ + assert zm.shape == rm.shape + assert len(rm.shape) == 3 + nx, ny, n = rm.shape + assert n == 4 + + # Build a list of vertices and a list of triangles + vertices = [] + triangles = [] + + def vertex_index(R, Z): + """ + Return the vertex index at given (R,Z) location. + Note: This is inefficient linear search + """ + # Check if there is already a vertex at this location + for i, v in enumerate(vertices): + vr, vz = v + d2 = (vr - R) ** 2 + (vz - Z) ** 2 + if d2 < 1e-10: + return i + # Not found so add a new vertex + vertices.append((R, Z)) + return len(vertices) - 1 + + for ix in range(nx): + for jy in range(ny): + # Adding cell (ix, jy) + # Get the vertex indices of the 4 corners + vertex_inds = [ + vertex_index(rm[ix, jy, n], zm[ix, jy, n]) for n in range(4) + ] + # Choose corners so triangles have the same sign + triangles.append(vertex_inds[0:3]) # Corners 1,2,3 + triangles.append(vertex_inds[:0:-1]) # Corners 4,3,2 + + self.vertices = np.array(vertices) + self.triangles = np.array(triangles) + self.cell_number = np.arange(nx * ny).reshape((nx, ny)) + + def __repr__(self): + return "" + + def plot_triangles(self, ax=None): + import matplotlib.pyplot as plt + + if ax is None: + fig, ax = plt.subplots() + + rs = self.vertices[self.triangles, 0] + zs = self.vertices[self.triangles, 1] + + # Close the triangles + rs = np.concatenate((rs, rs[:, 0:1]), axis=1) + zs = np.concatenate((zs, zs[:, 0:1]), axis=1) + + ax.plot(rs.T, zs.T, "k") + + return ax + + def with_data(self, da): + """ + Returns a new object containing vertices, triangles, and data + + Parameters + ---------- + + da : xarray.DataArray + Expected to have 'cherab_grid' attribute + and 'cell_number' coordinate. + Should only have 'x' and 'theta' dimensions. + + Returns + ------- + + A TriangularData object + """ + + if "cherab_grid" not in da.attrs: + raise ValueError("DataArray missing cherab_grid attribute") + + if "cell_number" not in da.coords: + raise ValueError("DataArray missing cell_number coordinate") + + da = da.squeeze() # Drop dimensions of size 1 + + # Check that extra dimensions (e.g time) have been dropped + # so that the data has the same dimensions as cell_number + if da.sizes != da.coords["cell_number"].sizes: + raise ValueError( + f"Data and cell_number coordinate have " + f"different sizes ({da.sizes} and " + f"{da.coords['cell_number'].sizes})" + ) + + if 2 * da.size == self.triangles.shape[0]: + # Data has not been sliced, so the size matches + # the number of triangles + + # Note: Two triangles per quad, so repeat data twice + return TriangularData( + self.vertices, self.triangles, da.data.flatten().repeat(2) + ) + + # Data size and number of triangles don't match. + # Use cell_number to work out which triangles to keep + + cells = da.coords["cell_number"].data.flatten() + triangles = np.concatenate( + (self.triangles[cells * 2, :], self.triangles[cells * 2 + 1, :]) + ) + + data = np.tile(da.data.flatten(), 2) + + return TriangularData(self.vertices, triangles, data)