Skip to content

Commit

Permalink
Fix sparse typing (#8387)
Browse files Browse the repository at this point in the history
* Fix sparse typing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update variable.py

* typos

* Update variable.py

* Update variable.py

* Update variable.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update variable.py

* Update variable.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
Illviljan and pre-commit-ci[bot] authored Oct 31, 2023
1 parent cfe4d71 commit 8a2f29b
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 12 deletions.
22 changes: 22 additions & 0 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -2595,6 +2595,28 @@ def argmax(
"""
return self._unravel_argminmax("argmax", dim, axis, keep_attrs, skipna)

def _as_sparse(self, sparse_format=_default, fill_value=_default) -> Variable:
"""
Use sparse-array as backend.
"""
from xarray.namedarray.utils import _default as _default_named

if sparse_format is _default:
sparse_format = _default_named

if fill_value is _default:
fill_value = _default_named

out = super()._as_sparse(sparse_format, fill_value)
return cast("Variable", out)

def _to_dense(self) -> Variable:
"""
Change backend from sparse to np.array.
"""
out = super()._to_dense()
return cast("Variable", out)


class IndexVariable(Variable):
"""Wrapper for accommodating a pandas.Index in an xarray.Variable.
Expand Down
9 changes: 7 additions & 2 deletions xarray/namedarray/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,7 @@ class _sparsearrayfunction(
Corresponds to np.ndarray.
"""

def todense(self) -> NDArray[_ScalarType_co]:
def todense(self) -> np.ndarray[Any, _DType_co]:
...


Expand All @@ -262,9 +262,14 @@ class _sparsearrayapi(
Corresponds to np.ndarray.
"""

def todense(self) -> NDArray[_ScalarType_co]:
def todense(self) -> np.ndarray[Any, _DType_co]:
...


# NamedArray can most likely use both __array_function__ and __array_namespace__:
_sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi)

sparseduckarray = Union[
_sparsearrayfunction[_ShapeType_co, _DType_co],
_sparsearrayapi[_ShapeType_co, _DType_co],
]
18 changes: 8 additions & 10 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
_DType_co,
_ScalarType_co,
_ShapeType_co,
_sparsearrayfunction_or_api,
_SupportsImag,
_SupportsReal,
)
Expand Down Expand Up @@ -810,9 +811,9 @@ def _as_sparse(
self,
sparse_format: Literal["coo"] | Default = _default,
fill_value: ArrayLike | Default = _default,
) -> Self:
) -> NamedArray[Any, _DType_co]:
"""
use sparse-array as backend.
Use sparse-array as backend.
"""
import sparse

Expand All @@ -832,18 +833,15 @@ def _as_sparse(
raise ValueError(f"{sparse_format} is not a valid sparse format") from exc

data = as_sparse(astype(self, dtype).data, fill_value=fill_value)
return self._replace(data=data)
return self._new(data=data)

def _to_dense(self) -> Self:
def _to_dense(self) -> NamedArray[Any, _DType_co]:
"""
Change backend from sparse to np.array
Change backend from sparse to np.array.
"""
from xarray.namedarray._typing import _sparsearrayfunction_or_api

if isinstance(self._data, _sparsearrayfunction_or_api):
# return self._replace(data=self._data.todense())
data_: np.ndarray[Any, Any] = self._data.todense()
return self._replace(data=data_)
data_dense: np.ndarray[Any, _DType_co] = self._data.todense()
return self._new(data=data_dense)
else:
raise TypeError("self.data is not a sparse array")

Expand Down

0 comments on commit 8a2f29b

Please sign in to comment.