Replace cdo lsm interpolation with python one.

[JanStreffing]

Jan Gärtner found a nice method to interpolate OpenIFS lsm in python from FESOM mesh description file.
This should replace the cumbersome cdo and nco based script:

mesh_path = '/proj/awi/input/fesom2/dart/dart_griddes_nodes.nc'

fesom_mesh = xr.open_dataset(mesh_path)
model_lon = fesom_mesh.lon.values
model_lat = fesom_mesh.lat.values


# create interpolation grid
box = (-180, 180, 0, 90)
res = (3600, 900)
proj = 'pc'
int_x_steps, int_y_steps, int_lons, int_lats = region_cartopy(box, res, proj)

# calculate distance to and index of the fesom grid point closest to
# the grid point in the interpolation grid
distances, inds = create_indexes_and_distances(model_lon, model_lat,
                                               int_lons, int_lats,
                                               k=1, workers=10)

mask = np.array(fesom_mesh.coast)[inds]
mask.shape = int_lons.shape




def region_cartopy(box, res, projection="pc"):
    """ Computes coordinates for the region 
    Parameters
    ----------
    box : list
        List of left, right, bottom, top boundaries of the region in -180 180 degree format
    res: list
        List of two variables, defining number of points along x and y
    projection : str
        Options are:
            "pc" : cartopy PlateCarree
            "mer": cartopy Mercator
            "np" : cartopy NorthPolarStereo
            "sp" : cartopy SouthPolarStereo
    Returns
    -------
    x : numpy.array
        1 d array of coordinate values along x
    y : numpy.array
        1 d array of coordinate values along y
    lon : numpy.array
        2 d array of longitudes
    lat : numpy array
        2 d array of latitudes
    """
    if projection == "pc":
        projection_ccrs = ccrs.PlateCarree()
    elif projection == "mer":
        projection_ccrs = ccrs.Mercator()
    elif projection == "np":
        projection_ccrs = ccrs.NorthPolarStereo()
    elif projection == "sp":
        projection_ccrs = ccrs.SouthPolarStereo()

    if not res is None:
        lonNumber, latNumber = res
    else:
        lonNumber, latNumber = 500, 500
    left, right, down, up = box
    logging.info('Box %s, %s, %s, %s', left, right, down, up)
    fig, ax = plt.subplots(
        1,
        1,
        subplot_kw=dict(projection=projection_ccrs),
        constrained_layout=True,
        figsize=(10, 10),
    )
    ax.set_extent([left, right, down, up], crs=ccrs.PlateCarree())
    xmin, xmax = ax.get_xbound()
    ymin, ymax = ax.get_ybound()

    # res = scl_fac * 300. # last number is the grid resolution in meters (NEEDS TO BE CHANGED)
    # nx = int((xmax-xmin)/res)+1; ny = int((ymax-ymin)/res)+1
    x = np.linspace(xmin, xmax, lonNumber)
    y = np.linspace(ymin, ymax, latNumber)
    x2d, y2d = np.meshgrid(x, y)

    npstere = ccrs.PlateCarree()
    transformed2 = npstere.transform_points(projection_ccrs, x2d, y2d)
    lon = transformed2[:, :, 0]  # .ravel()
    lat = transformed2[:, :, 1]  # .ravel()
    fig.clear()
    plt.close(fig)
   
    return x, y, lon, lat



def create_indexes_and_distances(model_lon, model_lat, lons, lats, k=1, workers=2):
    """
    Creates KDTree object and query it for indexes of points in FESOM mesh that are close to the
    points of the target grid. Also return distances of the original points to target points.
    Parameters
    ----------
    mesh : fesom_mesh object
        pyfesom mesh representation
    lons/lats : array
        2d arrays with target grid values.
    k : int
        k-th nearest neighbors to return.
    n_jobs : int, optional
        Number of jobs to schedule for parallel processing. If -1 is given
        all processors are used. Default: 1.
    Returns
    -------
    distances : array of floats
        The distances to the nearest neighbors.
    inds : ndarray of ints
        The locations of the neighbors in data.
    """
    xs, ys, zs = lon_lat_to_cartesian(model_lon, model_lat)
    xt, yt, zt = lon_lat_to_cartesian(lons.flatten(), lats.flatten())

    tree = cKDTree(list(zip(xs, ys, zs)))
    distances, inds = tree.query(list(zip(xt, yt, zt)), k=k, workers=workers)

    return distances, inds