bradyrx / esmtools

:octopus: a toolbox for Earth System Model analysis :octopus:
https://esmtools.readthedocs.io/en/latest/
MIT License
27 stars 4 forks source link

sel based on non-dim coordinate #74

Closed aaronspring closed 3 years ago

aaronspring commented 4 years ago

Ocean output is often on curvilinear grids. Therefore the lon, lat info is not stored in x,y. I often just want to have a look at a certain lon, lat box.

def cslice(ds, cx='lon', xmin=0, xmax=360, cy='lat', ymin=0, ymax=90):
    """Subselect from non-dimensional coordinates.

    Args:
        ds (xro): large xarray object with dim non-dim coords `cx` and `cy`.
        cx (str): coordinate name in x direction. Defaults to 'lon'.
        xmin (int, float): minimum x coordinate. Defaults to 0.
        xmax (int, float): maximum x coordinate. Defaults to 360.
        cy (str): coordinate name in y direction. Defaults to 'lat'.
        ymin (int, float): minimum y coordinate. Defaults to 0.
        ymax (int, float): maximum y coordinate. Defaults to 90.

    Returns:
        xro: subset of input ds.

    """
    assert xmin < xmax
    assert ymin < ymax
    ds_sub = ds.where((ds[cx] >= xmin) & (ds[cx] <= xmax))
    ds_sub = ds_sub.where((ds[cy] >= ymin) & (ds[cy] <= ymax))
    return ds_sub

This is a long standing xr issue: https://github.com/pydata/xarray/issues/2028 https://github.com/pydata/xarray/issues/1603 a few idea how to accomplish this but apparently some performance issues. The implemention with to_index() doesnt suit us because lon(x,y) and lat(x,y). I also tried to implement this with slice but failed.

Would you also need something like this @bradyrx ? Do you have any other ideas? csel could be a nice accessor for grid

bradyrx commented 4 years ago

I've always used this approach: https://github.com/bradyrx/esmtools/blob/master/esmtools/spatial.py (see extract_region). I think it's fairly quick since it uses numpy to minimize the xy coordinate and find the closest index on the curvilinear grid to that coordinate pair. Then I use the indices for the corners of the box with .isel() to extract the region.

aaronspring commented 4 years ago

do you think it would make sense adding this to the accessors? couldnt xgrid and ygrid be determined automatically xygrid something with x or y by default? use grid.extract_region(ds, coords)

bradyrx commented 4 years ago

I'll reopen this. I think it would be nice to have it as accessors.

aaronspring commented 3 years ago

xarray will soon have this