csiro-coasts / emsarray

xarray extension that supports EMS model formats
BSD 3-Clause "New" or "Revised" License
13 stars 1 forks source link

Support plotting more kinds of grids #121

Open mx-moth opened 8 months ago

mx-moth commented 8 months ago

SCHISM datasets follow the unstructured grid conventions, but define some important data on the nodes. Currently it is impossible to plot these data because emsarray only supports plotting data that are defined on faces, not nodes. It is possible to plot this data in matplotlib using tricontourf.

Unstructured grids can also define values on edges - these are often flux values indicating transport between adjacent cells. Arakawa C grids also support this, although the data will be defined on both the back and the left edge grids. These fluxes could be plotted as arrows centred on the edge, or by averaging all the vectors around a cell and plotting that in the centre of a cell.

For cell based data, the currently supported plot types are scalar (coloured polygons from one variable) and vector (arrows composed of two variables).

A unified approach to plotting is proposed. Add a new Convention.make_artist() method. This method takes a DataArray, or a tuple of DataArrays, and returns one matplotlib Artist that represents the values passed in. The ability to pass in a tuple of DataArrays is required for plotting vectors where the x/y components are defined in separate variables.

emsarray will provide a base implementation. If the passed in value is a single data array defined on the default grid - which is a polygon grid by default - the data array is passed to Convention.make_poly_collection(). If the value is a tuple of two data arrays both defined on the default grid, the values are assumed to be vector components and passed to Convention.make_quiver().

Each Convention type can define further plots that make sense for their plot types. For example the UGrid and ArakawaC conventions can provide the ability to plot values on nodes. A series of mixin classes can be provided to assist. The following shows how a NodeMixin could be define that supports plotting on any arbitrary node grid:

class NodeMixin(Convention):

    @property
    @abc.abstractmethod
    def node_grid_kinds(self) -> set[GridKind]:
        """The grid kinds in this convention that are 'nodes'."""
        pass

    @abc.abstractmethod
    def get_points_for_grid_kind(self, grid_kind: GridKind):
        """
        Get the (x, y) coordinates of all nodes for the given grid kind.
        Each subclass must implement this.
        """
        pass

    @abc.abstractmethod
    def triangulate_grid_kind(self, grid_kind: GridKind) -> matplotlib.tri.Triangulation:
        # Override this method if your convention has a better way of deriving a triangulation
        return matplotlib.tri.delaunay(*self.get_points_for_grid_kind(grid_kind))

    def make_tricontourp(self, figure, data_array: xarray.DataArray, **kwargs):
        x, y = self.get_points_for_grid_kind(data_array)
        return matplotlib.tri.tricontourp(x, y, self.ravel(data_array), **kwargs)

    def make_artist(self, figure, value, **kwargs):
        if isinstance(value, xarray.DataArray):
            grid_kind = self.get_grid_kind(value)
            if grid_kind in self.node_grid_kinds:
                return self.make_tricontourp(figure, data_array)
        return super().make_artist(figure, value, **kwargs)

UGrid using this might look like:


class UGrid(NodeGridKind, Convention):
    node_grid_kinds = {UGridKind.node}

    def get_points_for_grid_kind(self, grid_kind):
        return (self.topology.node_x.values, self.topology.node_y.values)