csiro-coasts / emsarray

xarray extension that supports EMS model formats
BSD 3-Clause "New" or "Revised" License
13 stars 2 forks source link

Add `Convention.selector_for_indices()` #106

Closed mx-moth closed 2 months ago

mx-moth commented 1 year ago

This will take a list of indices (native or linear?) and return a xarray Dataset that can be passed to .isel(). Naive implementation:

def selectors_for_indices(
    self,
    indices: List[Index],
    *,
    include_linear_index: bool = False,
) -> xarray.Dataset:
    selectors = map(self.selector_for_index, indices)
    selector = pandas.DataFrame(selectors).to_xarray().drop_vars('index')
    if include_linear_index:
        selector = selector.assign_coords({'index': list(map(self.ravel_index, indices))})
    return selector

Usage:

points = shapely.points([...])
indices = [ds.ems.get_index_for_point(point) for point in points]
points_selector = ds.ems.selectors_for_indices(indices)
points_ds = ds.isel(points_selector)