Add to_scipy_sparse/from_scipy_sparse xarray accessor methods

janitor/utils.py

@@ -144,7 +144,8 @@ def import_message(
 
     Generic message for indicating to the user when a function relies on an
     optional module / package that is not currently installed. Includes
-    installation instructions. Used in `chemistry.py` and `biology.py`.
+    installation instructions. Used in `chemistry.py`, `biology.py`, and
+    `xarray/functions.py`.
 
     Args:
         submodule: `pyjanitor` submodule that needs an external dependency.

janitor/xarray/functions.py

@@ -9,6 +9,18 @@
 import pandas_flavor as pf
 import xarray as xr
 
+from ..utils import import_message
+
+try:
+    import sparse
+except ImportError:
+    import_message(
+        submodule="xarray",
+        package="sparse",
+        conda_channel="conda-forge",
+        pip_install=True,
+    )
+
 
 @pf.register_xarray_dataarray_method
 def clone_using(
@@ -148,3 +160,117 @@ def convert_datetime_to_number(
     times = da_or_ds.coords[dim].data / np.timedelta64(1, time_units)
 
     return da_or_ds.assign_coords({dim: times})
+
+
+@pf.register_xarray_dataarray_method
+def to_scipy_sparse(da: xr.DataArray):
+    """
+    Convert a 2-dimensional `DataArray` to a `scipy.sparse` matrix.
+
+    Uses the [sparse](https://github.com/pydata/sparse) package as an
+    intermediate step, so only non-zero entries are kept. `scipy.sparse`
+    only supports 2-dimensional matrices, so `da` must be 2-dimensional.
+
+    Examples:
+        Converting a mostly-zero `DataArray` to a `scipy.sparse` matrix:
+
+        >>> import numpy as np
+        >>> import xarray as xr
+        >>> import janitor.xarray
+        >>> da = xr.DataArray(
+        ...     np.array([[0, 0, 3], [4, 0, 0]]),
+        ...     dims=["row", "col"],
+        ... )
+        >>> mat = da.to_scipy_sparse()  # doctest: +SKIP
+        >>> mat.toarray()  # doctest: +SKIP
+        array([[0, 0, 3],
+               [4, 0, 0]])
+
+    Args:
+        da: The `DataArray` supplied by the method itself. Must have
+            exactly 2 dimensions.
+
+    Raises:
+        ValueError: If `da` does not have exactly 2 dimensions.
+
+    Returns:
+        A `scipy.sparse` CSR matrix containing the non-zero values of `da`.
+    """
+    if da.ndim != 2:
+        raise ValueError(
+            "`to_scipy_sparse` only supports 2-dimensional DataArrays, "
+            f"but the supplied DataArray has {da.ndim} dimension(s)."
+        )
+
+    if isinstance(da.data, sparse.COO):
+        return da.data.tocsr()
+
+    return sparse.COO.from_numpy(np.asarray(da.data)).tocsr()
+
+
+@pf.register_xarray_dataarray_method
+def from_scipy_sparse(
+    da: xr.DataArray,
+    scipy_sparse_mat,
+    use_coords: bool = False,
+    use_attrs: bool = False,
+    new_name: str = None,
+) -> xr.DataArray:
+    """
+    Given a 2-dimensional `scipy.sparse` matrix, return a `DataArray`
+    that mirrors the dimension names (and, optionally, coordinates and
+    attrs) of the supplied `DataArray`.
+
+    This is the inverse of `to_scipy_sparse`, and is implemented as a
+    thin wrapper around `clone_using`, converting `scipy_sparse_mat` to
+    a `sparse.COO`-backed array first. Since `scipy_sparse_mat` will
+    generally not have the same shape as `da`, `use_coords` defaults to
+    `False` here (unlike in `clone_using`), so shapes are not required
+    to match unless you explicitly ask for coordinates to be copied over.
+
+    Examples:
+        Rebuilding a `DataArray` from a `scipy.sparse` matrix, keeping
+        only the dimension names of the original:
+
+        >>> import numpy as np
+        >>> import scipy.sparse
+        >>> import xarray as xr
+        >>> import janitor.xarray
+        >>> da = xr.DataArray(
+        ...     np.zeros((2, 3)),
+        ...     dims=["row", "col"],
+        ...     name="original",
+        ... )
+        >>> mat = scipy.sparse.csr_matrix([[0, 0, 3], [4, 0, 0]])
+        >>> new_da = da.from_scipy_sparse(mat)  # doctest: +SKIP
+        >>> new_da.dims  # doctest: +SKIP
+        ('row', 'col')
+        >>> new_da.data.todense()  # doctest: +SKIP
+        array([[0, 0, 3],
+               [4, 0, 0]])
+
+    Args:
+        da: The `DataArray` supplied by the method itself, used as the
+            template for dimension names and (optionally) coordinates
+            and attrs.
+        scipy_sparse_mat: A 2-dimensional `scipy.sparse` matrix.
+        use_coords: If `True`, use the coordinates of `da` for the
+            coordinates of the newly-generated array. `da` and
+            `scipy_sparse_mat` must have the same shape in that case.
+            Defaults to `False`.
+        use_attrs: If `True`, copy over the `attrs` from `da`.
+        new_name: If set, use as the name of the returned `DataArray`.
+            Otherwise, use the name of `da`.
+
+    Returns:
+        A new `DataArray`, with `sparse.COO`-backed data restored from
+            `scipy_sparse_mat`.
+    """
+    coo = sparse.COO.from_scipy_sparse(scipy_sparse_mat)
+
+    return da.clone_using(
+        coo,
+        use_coords=use_coords,
+        use_attrs=use_attrs,
+        new_name=new_name,
+    )