how to query lat and lon in a DataArray with dims x and y?

[raybellwaves]

Take the non-uniform example in the docs (http://xarray.pydata.org/en/stable/plotting.html#multidimensional-coordinates)

import xarray as xr
import numpy as np

lon, lat = np.meshgrid(np.linspace(-20, 20, 5), np.linspace(0, 30, 4))
lon += lat / 10
lat += lon / 10
da = xr.DataArray(
    np.arange(20).reshape(4, 5),
    dims=["y", "x"],
    coords={"lat": (("y", "x"), lat), "lon": (("y", "x"), lon)},
)

How would I query a lon-lat point? e.g. lat = 10.1 and lon = 10.

[ahuang11]

Something like this maybe.

import xarray as xr
import numpy as np
from scipy.spatial import cKDTree

lon, lat = np.meshgrid(np.linspace(-20, 20, 5), np.linspace(0, 30, 4))
lon += lat / 10
lat += lon / 10
da = xr.DataArray(
    np.arange(20).reshape(4, 5),
    dims=["y", "x"],
    coords={"lat": (("y", "x"), lat), "lon": (("y", "x"), lon)},
)

query_lats = [10.1, -2]
query_lons = [10, 19]
query_coords = np.stack((query_lats, query_lons)).T

lats_lons = np.vstack([da['lat'].values.ravel(), da['lon'].values.ravel()]).T
kdtree = cKDTree(lats_lons)
dist, index = kdtree.query(query_coords)
da_sel = da.stack({'index': ['y', 'x']}).isel(index=index).unstack()
da_sel

Or use balltree/haversine for more lat/lon accuracy: https://stackoverflow.com/questions/10549402/kdtree-for-longitude-latitude

Or another way,

import xarray as xr
import numpy as np
from scipy.spatial import cKDTree

lon, lat = np.meshgrid(np.linspace(-20, 20, 5), np.linspace(0, 30, 4))
lon += lat / 10
lat += lon / 10
da = xr.DataArray(
    np.arange(20).reshape(4, 5),
    dims=["y", "x"],
    coords={"lat": (("y", "x"), lat), "lon": (("y", "x"), lon)},
)

query_lats = [10.1]
query_lons = [10]

tol = 2
da.where(
    (da['lat'] > (np.min(query_lats) - tol)) &
    (da['lat'] < (np.max(query_lats) + tol)) &
    (da['lon'] > (np.min(query_lons) - tol)) &
    (da['lon'] < (np.max(query_lons) + tol)),
    drop=True
)

[keewis]

using where or falling back to numpy are currently the only options available in xarray itself. We're in the process of changing this (see #4979), so if you have any feedback feel free to comment there.

[djhoese]

A similar question was asked on stackoverflow: https://stackoverflow.com/questions/66433948/get-nearest-pixel-value-from-satellite-image-using-latitude-longitude-coordinate/66439989#66439989

[dcherian]

Also see xoak: https://xoak.readthedocs.io/en/latest/

[raybellwaves]

Nice work @benbovy. xoak looks great.

points = xr.Dataset(
    {
        "lat": 10.1,
        "lon": 10,
    }
)
da.xoak.set_index(["lat", "lon"], "sklearn_geo_balltree")
da.xoak.sel(lat=points.lat, lon=points.lon)

[benbovy]

Also @willirath!

For latitude / longitude points, you might want to try pys2index (registered as "s2point" in xoak). For large sets of points it may be dozen times faster than the "sklearn_geo_balltree" index.