pyjanitor icon indicating copy to clipboard operation
pyjanitor copied to clipboard

[ENH] Expose sparse library functions

Open zbarry opened this issue 5 years ago • 0 comments

XArray supports the Sparse package but doesn't expose the functions to convert to/from sparse objects. These functions could be nicely packaged in pyjanitor to do so:

import sparse

@register_xarray_dataarray_method
def to_scipy_sparse(
    da: xr.DataArray,
) -> xr.DataArray:
    if isinstance(da.data, sparse.COO):
        return xr.apply_ufunc(sparse.COO.to_scipy_sparse, da)
    return da


@register_xarray_dataarray_method
def todense(
    da: xr.DataArray,
) -> xr.DataArray:
    if isinstance(da.data, sparse.COO):
        return xr.apply_ufunc(sparse.COO.todense, da)
    return da


@register_xarray_dataarray_method
def tocsc(
    da: xr.DataArray,
) -> xr.DataArray:
    if isinstance(da.data, sparse.COO):
        return xr.apply_ufunc(sparse.COO.tocsc, da)
    return da


@register_xarray_dataarray_method
def tocsr(
    da: xr.DataArray,
) -> xr.DataArray:
    if isinstance(da.data, sparse.COO):
        return xr.apply_ufunc(sparse.COO.tocsr, da)
    return da


@register_xarray_dataarray_method
def tosparse(
    da: xr.DataArray,
) -> xr.DataArray:
    if isinstance(da.data, sparse.COO):
        return da
    return xr.apply_ufunc(sparse.COO, da)

zbarry avatar Feb 22 '20 16:02 zbarry