Add unstructured mesh support with MeshInfo

[stephenworsley]

This adds back end support for UGRID style unstructured meshes.
It is built in anticipation of Iris adding UGRID support and therefore the code is separated into an experimental directory. Adding and testing against this code now ensures that the design of iris-esmf-regrid will maintain the flexibility to add UGRID support when Iris allows it.

This builds on top of #22 and is therefore a draft.

NB: (See plots example below) this PR has facilitated regridding of unstructured FESOM data to a rectilinear grid.

[codecov]

Codecov Report

Merging #24 (0415780) into main (77f2119) will increase coverage by 1.35%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main      #24      +/-   ##
==========================================
+ Coverage   94.71%   96.06%   +1.35%     
==========================================
  Files           7        9       +2     
  Lines         227      305      +78     
==========================================
+ Hits          215      293      +78     
  Misses         12       12              

Impacted Files	Coverage Δ
esmf_regrid/esmf_regridder.py	90.90% <ø> (ø)
esmf_regrid/experimental/unstructured_regrid.py	100.00% <100.00%> (ø)
...t/experimental/unstuctured_regrid/test_MeshInfo.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 77f2119...0415780. Read the comment docs.

[stephenworsley]

From this PR i was able to do the regrid the following file http://esgf3.dkrz.de/thredds/fileServer/cmip6/CMIP/AWI/AWI-CM-1-1-MR/historical/r1i1p1f1/Omon/tos/gn/v20181218/tos_Omon_AWI-CM-1-1-MR_historical_r1i1p1f1_gn_201101-201412.nc to a lat/lon grid which looked like the attached picture.

This was done with the following code:

import numpy as np
from numpy import ma
import scipy.sparse
import iris
import esmf_regrid.esmf_regridder
from esmf_regrid.experimental.unstructured_regrid import MeshInfo

fn = "/downloads/tos_Omon_AWI-CM-1-1-MR_historical_r1i1p1f1_gn_201101-201412.nc"
cube = iris.load_cube(fn)

lat_bounds = cube.coords("latitude")[0].bounds
lon_bounds = cube.coords("longitude")[0].bounds

# the size of the input mesh can be reduced for the sake of debugging.
# n = 35000
n = lat_bounds.shape[0]

lat_bounds = lat_bounds[:n]
lon_bounds = lon_bounds[:n]

shape = lat_bounds.shape

# many of the bounds are repeated, this causes ESMF to fail.
# we mask these points and work around them.
lat_mask = lat_bounds[:, 1:] == lat_bounds[:, :-1]
lon_mask = lon_bounds[:, 1:] == lon_bounds[:, :-1]
mask = np.logical_and(lat_mask, lon_mask)
mask = np.concatenate([np.zeros([shape[0], 1]), mask], axis=1)

# unmasked bounds are extracted.
masked_lats = ma.array(lat_bounds, mask=mask)
masked_lons = ma.array(lon_bounds, mask=mask)
masked_node_lats = masked_lats.flatten()
masked_node_lons = masked_lons.flatten()
unmasked_lats = ma.compressed(masked_node_lats)
unmasked_lons = ma.compressed(masked_node_lons)

# appropriate indices for face node connectivity are calculated.
num_masked = mask.sum(axis=1)
num_masked_shifted = np.concatenate([[0], num_masked[:-1]])
cumulative_masked = np.cumsum(num_masked_shifted)
cumulative_masked = np.broadcast_to(cumulative_masked, (shape[1], len(cumulative_masked))).T

# data is shaped into a form MeshInfo understands.
node_coords = np.stack([unmasked_lons, unmasked_lats], axis=1)
fnc = np.arange(np.product(shape)).reshape(shape)
fnc = fnc - cumulative_masked
fnc = ma.array(fnc, mask=mask)
nsi = 0

# create the source mesh.
mesh = MeshInfo(node_coords, fnc, nsi)

# define target grid
grid_lats = np.linspace(-90, 90, 60)
grid_lons = np.linspace(-180, 180, 80)
grid = esmf_regrid.esmf_regridder.GridInfo(grid_lons[:-1], grid_lats[:-1], grid_lons, grid_lats)

# regridding is done in sections in order to reduce memory usage.
matrices = []
nm = 7
for x in range(nm):
    lower = int(n*x/nm)
    upper = int(n*(x+1)/nm)
    f = fnc[lower:upper]
    ns = int(f.min())
    ne = int(f.max())+1
    nc = node_coords[ns:ne]
    submesh = MeshInfo(nc, f, ns)
    srg = esmf_regrid.esmf_regridder.Regridder(submesh, grid)
    matrices.append(srg.weight_matrix)
    print(f"{x+1}/{nm} complete")

wm = scipy.sparse.hstack(matrices)

# the regridder is made using precalculated weights.
rg = esmf_regrid.esmf_regridder.Regridder(mesh, grid, wm)

# regrid data from the cube.
out_data = rg.regrid(cube.data[0, :n])

# create a cube from the resulting data.
out_cube = iris.cube.Cube(out_data, standard_name=cube.standard_name)
grid_lon_bounds = np.stack([grid_lons[:-1], grid_lons[1:]]).T
grid_lat_bounds = np.stack([grid_lats[:-1], grid_lats[1:]]).T
lon_coord = iris.coords.DimCoord(grid_lons[:-1], bounds=grid_lon_bounds, standard_name="longitude")
lat_coord = iris.coords.DimCoord(grid_lats[:-1], bounds=grid_lat_bounds, standard_name="latitude")
out_cube.add_dim_coord(lon_coord, 0)
out_cube.add_dim_coord(lat_coord, 1)

# defer importing of plotting packages to save memory.
import matplotlib.pyplot as plt
import iris.quickplot as qplt
# import cartopy.crs as ccrs

# an alternate plot method without using cubes. Note that the data requires transposing.

# ax = plt.axes(projection=ccrs.PlateCarree())
#
# ax.coastlines()
# plt.pcolormesh(grid_lons, grid_lats, out_data.T)

# plot the data.
qplt.pcolormesh(out_cube)
plt.gca().coastlines()

plt.show()

[stephenworsley]

The above exercise revealed a couple of things about what ESMF can and can't handle. ESMF will fail to build a mesh with the following:

• A face with two (consecutive) nodes at the same location.
• A node which does not belong to any face.

On the other hand, the following does not cause ESMF to fail:

• Two nodes at the same location, but attached to different faces.

This may be worth bearing in mind if we wish to generalise functions for building MeshInfo objects from non-UGRID formats. There may also be an argument for checking the validity of MeshInfo objects before passing them to ESMF.

I also want to comment about the memory saving trick I did with splitting up the mesh and recombining the weights matrices. This was necessary because I was running out of memory on my machine without this trick. This is something I had been experimenting with before and may consider adding as a feature, but still feels somewhat ad hoc. It works here because the source grid has a significantly higher resolution than the target grid and therefore the memory required by the regridder is roughly proportional to the size of the source grid. There may be a generalisation of this technique which can improve performance for a broader class of source/target grids, but I suspect such a technique would end up being significantly more complex. With that said, the fact that such a technique is possible to be implemented ad hoc is a good sign that the design of the regridder is working well.

[stephenworsley]

For reference, the original data looks like this.

Since this is using a scatter plot (since it is unstructured) here is the regrid data using the same scatter plot for comparison.

[zklaus]

LGTM, good work! 🍻