Pyramid create for creating pyramids with custom funcs

[ahuang11]

Closes #94

Adapted pyramid_coarsen from ndpyramid for a project, but I would like to add a bit of functionality back so I can simply import it.

Basic usage:

def sel_coarsen(ds, factor, dims, **kwargs):
    return ds.sel(**{dim: slice(None, None, factor) for dim in dims})

factors = [4, 2, 1]
pyramid = pyramid_create(
    temperature,
    dims=('lat', 'lon'),
    factors=factors,
    boundary='trim',
    func=sel_coarsen,
    method_label=method_label,
    type_label='pick',
)

I think it could eventually support #10, but haven't looked at it in detail.

More advanced, applicable use case:

import h5py
import xarray as xr
import dask.array as da
import datatree as dt
from ndpyramid import pyramid_create
from tsdownsample import MinMaxLTTBDownsampler


def _help_downsample(data, time, n_out):
    indices = MinMaxLTTBDownsampler().downsample(time, data, n_out=n_out)
    return data[indices], indices


def apply_downsample(ts_ds, factor, dims):
    dim = dims[0]
    n_out = len(ts_ds["data"]) // factor
    ts_ds_downsampled, indices = xr.apply_ufunc(
        _help_downsample,
        ts_ds["data"],
        ts_ds[dim],
        kwargs=dict(n_out=n_out),
        input_core_dims=[[dim], [dim]],
        output_core_dims=[[dim], ["indices"]],
        exclude_dims=set((dim,)),
        vectorize=True,
        dask="parallelized",
        dask_gufunc_kwargs=dict(output_sizes={dim: n_out, "indices": n_out}),
    )
    ts_ds_downsampled[dim] = ts_ds[dim].isel(time=indices.values[0])
    return ts_ds_downsampled.rename("data")


def build_dataset(f, data_key, dims):
    coords = {f[dim] for dim in dims.values()}
    data = f[data_key]
    ds = xr.DataArray(
        da.from_array(data, name="data", chunks=(data.shape[0], 1)),
        dims=dims,
        coords=coords,
    ).to_dataset()
    return ds


f = h5py.File("allensdk_cache/session_715093703/probe_810755797_lfp.nwb", "r")
ts_ds = build_dataset(
    f,
    "acquisition/probe_810755797_lfp_data/data",
    {
        "time": "acquisition/probe_810755797_lfp_data/timestamps",
        "channel": "acquisition/probe_810755797_lfp_data/electrodes",
    },
).isel(channel=[0, 1, 2, 3, 4])
ts_dt = pyramid_create(
    ts_ds,
    factors=[1, 2],
    dims=["time"],
    func=apply_downsample,
    type_label="pick",
    method_label="pyramid_downsample",
)
ts_dt

[maxrjones]

Thanks @ahuang11, this is nice! I appreciate that you shared your use case to show how this would broaden the usefulness of the library! I made a couple small docs suggestions. Also, can you please add pyramid_create to docs/api.rst so that it is rendered in the API docs?

[maxrjones]

Thanks @ahuang11!

[ahuang11]

Thanks for merging! I was wondering when you envision the next release to be?

[maxrjones]

Thanks for merging! I was wondering when you envision the next release to be?

I can put out a release this week

[ahuang11]

I was wondering if you had a chance to put out a new release yet?

[maxrjones]

I was wondering if you had a chance to put out a new release yet?

Thanks for the reminder - just put out a new release