Does UXarray support regridding unstructured grids to rectilinear grids?

[tomvothecoder]

Hi again! For my upcoming E3SM demo with UXarray and xCDAT, we want to regrid E3SM data to a rectilinear grid in order to use xCDAT's spatial averager. Is there a way to do this with UXarray? There was a discussion about regridding to lat-lon here.

What we found

We found uxarray.UxDataset.nearest_neighbor_remap, but we aren't sure if it supports rectilinear grids.

Our Current Workarounds

1. ncremap

• https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/754286611/Regridding+E3SM+Data+with+ncremap
• Example usage in this xCDAT notebook

2. The simple_regrid function shown below

import xcdat as xc

 def simple_regrid(
    ds: xr.Dataset,
    target_var: str,
    nlat: xr.DataArray,
    nlon: xr.DataArray,
    infill: bool = True,
) -> xr.Dataset:
    """
    Nearest neighbor mapping of 2D fields from E3SM column format to a
    user-defined rectilinear grid.

    Parameters:
    -----------
    ds : xr.Dataset
        Source dataset to remap
    target_var : str
        Name of variable to remap
    nlat : xr.DataArray
        Target latitude values
    nlon : xr.DataArray
        Target longitude values
    infill : bool
        Flag to infill (with extrapolation) missing values (default True)

    Returns:
    --------
    xr.Dataset

    Notes:
    ------
    This regridding tool is intended as a simple regridding tool and is not fit
    for most scientific applications, but may be useful for data quick-looks.
    """
    dset_out = []

    # Loop over time steps and remap one map at a time.
    for i in range(len(ds.time)):
        # Get data
        lat = ds.lat
        lon = ds.lon
        data = ds[target_var].isel(time=i)

        # Target grid
        LON, LAT = np.meshgrid(nlon, nlat)
        shp = LAT.shape

        # Create a nearest-neighbor tree for the grid
        tree = spatial.cKDTree(list(zip(LAT.flat, LON.flat)))
        _, ind = tree.query(np.array([lat, lon]).T)
        n = tree.n

        # Sum of data in each grid box
        X = np.bincount(ind, weights=data, minlength=n)
        # Total number of matches
        cnt = np.bincount(ind, weights=np.ones_like(data), minlength=n)

        # Going to get divide by zero here for grid boxes with no data
        with np.errstate(divide="ignore", invalid="ignore"):
            grid = X / cnt

        # Reshape to regular grid
        grid = grid.reshape(shp)
        grid = xr.DataArray(
            data=grid,
            dims=["lat", "lon"],
            coords={"lat": nlat, "lon": nlon},
            name=target_var,
        )
        dset_out.append(grid)

    # Concatenate time steps and create xr.Dataset
    ds_final = xr.concat(dset_out, dim=ds.time).to_dataset()

    # Incorporate bounds from original dataset
    if "time_bnds" in ds.data_vars:
        ds_final["time_bnds"] = ds.time_bnds

    # Add missing bounds
    ds_final = ds_final.bounds.add_missing_bounds(["X", "Y", "T"])

    # Infill (if desired)
    if infill:
        ds_final[target_var] = ds_final[target_var].interpolate_na(
            dim="lon", method="nearest", fill_value="extrapolate"
        )

    return ds_final

# Define regrid targets
target_var = "TREFHT"
nlat, _ = xc.create_axis(
    "lat",
    np.arange(-88.75, 90, 2.5),
    attrs={"axis": "Y", "units": "degrees_north"},
    generate_bounds=False,
)
nlon, _ = xc.create_axis(
    "lon",
    np.arange(1.25, 360, 2.5),
    attrs={"axis": "X", "units": "degrees_east"},
    generate_bounds=False,
)

# call simple regridder
uxds_r = simple_regrid(uxds, target_var, nlat, nlon)

[philipc2]

Hi @tomvothecoder

We currently only support unstructured to unstructured regridding, however support for structured and rectilinear grids is something we plan to support.

Will create an issue out of this so our intern @aaronzedwick can start working on this when he returns

[tomvothecoder]

Thanks for the clarification @philipc2! This feature would be awesome to keep code in Python workflows!

[rajeeja]

Hey @tomvothecoder,

On chrysalis, how do I run the workarounds you mention?:

ncremap - > samples mesh files and exact options to pass to this tool
The simple_regrid function shown below - uxds? - uxds_r = simple_regrid(uxds, target_var, nlat, nlon)

[tomvothecoder]

Hey @rajeeja, I updated the post description with a link to an xCDAT notebook using ncremap. I also fixed code example for the simple_regrid function to show the run script.

@pochedls might be able to help you more here (he wrote the simple_regrid function).

[pochedls]

The simple_regrid function isn't a robust regridder – it was just meant as a simple, informal nearest neighbor remapper for E3SM column data (using scipy.spatial.kdtree) - mainly for demonstration purposes (but we ended up using ncremap instead). As a result, it hardcodes assumptions (like lat exists in the dataset).

With that said, if you source the function above, it should work with a basic 2D E3SM dataarray. It looks like it is missing from scipy import spatial at the top. I also notice that the first few lines can be removed from the loop (# Get data and # Target grid stuff can be computed once before the loop).

This nearest neighbor approach doesn't take into account cell edges – so there is missing data (and the infilling option in the function is extremely crude).

[rljacob]

Is everyone sure they want to do regridding in python? Some past experience showed me that a compiled tool like TempestRemap will be much much faster then anything in python.

[paullric]

Python has a certain flexibility that is more difficult to duplicate in a compiled tool. Plus there's a larger community of developers for Python. It would certainly be worthwhile to investigate how much slower the regridding would be in Python, and leverage the capabilities we've been developing.

[rljacob]

Can our interface to YAC do this? @rajeeja ?

[rajeeja]

Can our interface to YAC do this? @rajeeja ?

No yet, our interface in #1440 isn't ready yet; to support it, we likely need structured-target-specific handling in the YAC adapter; #1440 PR moved uxarray over to YAC’s low-level yac.core API, and that path is now working for unstructured-to-unstructured remapping (nnn and conservative). this is an uxarray adapter gap, not necessarily a YAC limitation. YAC itself appears to support regular 2D grids, but our current backend still treats the target too much like the unstructured case. We likely need structured-target-specific handling plus dedicated tests before answering “yes”.

[rajeeja]

Can our interface to YAC do this? @rajeeja ?

No yet, our interface in #1440 isn't ready yet; to support it, we likely need structured-target-specific handling in the YAC adapter; #1440 PR moved uxarray over to YAC’s low-level yac.core API, and that path is now working for unstructured-to-unstructured remapping (nnn and conservative). this is an uxarray adapter gap, not necessarily a YAC limitation. YAC itself appears to support regular 2D grids, but our current backend still treats the target too much like the unstructured case. We likely need structured-target-specific handling plus dedicated tests before answering “yes”.

I'll take this on a followup PR