diff --git a/ci/requirements/py36.yml b/ci/requirements/py36.yml index 53da54e3b90..985782cbd23 100644 --- a/ci/requirements/py36.yml +++ b/ci/requirements/py36.yml @@ -35,4 +35,4 @@ dependencies: - iris>=1.10 - pydap - lxml - + - quantities diff --git a/ci/requirements/py37.yml b/ci/requirements/py37.yml index 538d4679a79..72e993fa7a2 100644 --- a/ci/requirements/py37.yml +++ b/ci/requirements/py37.yml @@ -32,5 +32,6 @@ dependencies: - cfgrib>=0.9.2 - lxml - pydap + - quantities - pip: - numbagg diff --git a/setup.cfg b/setup.cfg index 114f71f4a9f..56a5e58191d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -82,6 +82,8 @@ ignore_missing_imports = True ignore_missing_imports = True [mypy-pytest.*] ignore_missing_imports = True +[mypy-quantities.*] +ignore_missing_imports = True [mypy-rasterio.*] ignore_missing_imports = True [mypy-scipy.*] diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index fcd0400566f..17a1e7d4764 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -162,10 +162,12 @@ def trapz(y, x, axis): def asarray(data): + from .npcompat import _asarray + return ( data if (isinstance(data, dask_array_type) or hasattr(data, "__array_function__")) - else np.asarray(data) + else _asarray(data) ) diff --git a/xarray/core/npcompat.py b/xarray/core/npcompat.py index 22c14d9ff40..6fd341bb237 100644 --- a/xarray/core/npcompat.py +++ b/xarray/core/npcompat.py @@ -35,6 +35,8 @@ import numpy as np +from .options import OPTIONS + try: from numpy import isin except ImportError: @@ -378,3 +380,12 @@ def __array_function__(self, *args, **kwargs): IS_NEP18_ACTIVE = _is_nep18_active() + + +def _asarray(data): + # options get set after import, so this needs to be done in a + # function + if OPTIONS["enable_experimental_ndarray_subclass_support"]: + return np.asanyarray(data) + else: + return np.asarray(data) diff --git a/xarray/core/options.py b/xarray/core/options.py index c5086268f48..1e61219bed5 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -8,6 +8,9 @@ CMAP_SEQUENTIAL = "cmap_sequential" CMAP_DIVERGENT = "cmap_divergent" KEEP_ATTRS = "keep_attrs" +ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT = ( + "enable_experimental_ndarray_subclass_support" +) OPTIONS = { @@ -19,6 +22,7 @@ CMAP_SEQUENTIAL: "viridis", CMAP_DIVERGENT: "RdBu_r", KEEP_ATTRS: "default", + ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT: False, } _JOIN_OPTIONS = frozenset(["inner", "outer", "left", "right", "exact"]) @@ -35,6 +39,7 @@ def _positive_integer(value): FILE_CACHE_MAXSIZE: _positive_integer, WARN_FOR_UNCLOSED_FILES: lambda value: isinstance(value, bool), KEEP_ATTRS: lambda choice: choice in [True, False, "default"], + ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT: lambda value: isinstance(value, bool), } @@ -98,6 +103,9 @@ class set_options: attrs, ``False`` to always discard them, or ``'default'`` to use original logic that attrs should only be kept in unambiguous circumstances. Default: ``'default'``. + - ``enable_experimental_ndarray_subclass_support``: whether or not + to enable the support for subclasses of numpy's ndarray. + Default: ``False``. You can use ``set_options`` either as a context manager: diff --git a/xarray/core/variable.py b/xarray/core/variable.py index bc8da10dd0c..4a141705bb8 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -18,7 +18,7 @@ VectorizedIndexer, as_indexable, ) -from .npcompat import IS_NEP18_ACTIVE +from .npcompat import IS_NEP18_ACTIVE, _asarray from .options import _get_keep_attrs from .pycompat import dask_array_type, integer_types from .utils import ( @@ -193,6 +193,8 @@ def as_compatible_data(data, fastpath=False): data[mask] = fill_value else: data = np.asarray(data) + elif isinstance(data, np.matrix): + data = np.asarray(data) if not isinstance(data, np.ndarray): if hasattr(data, "__array_function__"): @@ -208,8 +210,7 @@ def as_compatible_data(data, fastpath=False): ) # validate whether the data is valid data types - data = np.asarray(data) - + data = _asarray(data) if isinstance(data, np.ndarray): if data.dtype.kind == "O": data = _possibly_convert_objects(data) @@ -244,6 +245,21 @@ def _as_array_or_item(data): return data +def _as_any_array_or_item(data): + """ Return the given values as a numpy array subclass instance, or + individual item if it's a 0d datetime64 or timedelta64 array. + + The same caveats as for ``_as_array_or_item`` apply. + """ + data = _asarray(data) + if data.ndim == 0: + if data.dtype.kind == "M": + data = np.datetime64(data, "ns") + elif data.dtype.kind == "m": + data = np.timedelta64(data, "ns") + return data + + class Variable( common.AbstractArray, arithmetic.SupportsArithmetic, utils.NdimSizeLenMixin ): @@ -322,7 +338,7 @@ def data(self): ): return self._data else: - return self.values + return _as_any_array_or_item(self._data) @data.setter def data(self, data): diff --git a/xarray/tests/test_quantities.py b/xarray/tests/test_quantities.py new file mode 100644 index 00000000000..72423d48cda --- /dev/null +++ b/xarray/tests/test_quantities.py @@ -0,0 +1,140 @@ +import numpy as np +import pytest + +from xarray import DataArray, set_options + +try: + import quantities as pq + + has_quantities = True +except ImportError: + has_quantities = False + +pytestmark = pytest.mark.skipif(not has_quantities, reason="requires python-quantities") + + +set_options(enable_experimental_ndarray_subclass_support=True) + + +def assert_equal_with_units(a, b): + a = a if not isinstance(a, DataArray) else a.data + b = b if not isinstance(b, DataArray) else b.data + + assert (hasattr(a, "units") and hasattr(b, "units")) and a.units == b.units + + assert (hasattr(a, "magnitude") and hasattr(b, "magnitude")) and np.allclose( + a.magnitude, b.magnitude + ) + + +def test_without_subclass_support(): + with set_options(enable_experimental_ndarray_subclass_support=False): + data_array = DataArray(data=np.arange(10) * pq.m) + assert not hasattr(data_array.data, "units") + + +@pytest.mark.filterwarnings("ignore:the matrix subclass:PendingDeprecationWarning") +def test_matrix(): + matrix = np.matrix([[1, 2], [3, 4]]) + da = DataArray(matrix) + + assert not isinstance(da.data, np.matrix) + + +def test_masked_array(): + masked = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]]) + da = DataArray(masked) + assert not isinstance(da.data, np.ma.MaskedArray) + + +def test_units_in_data_and_coords(): + data = np.arange(10) * pq.m + x = np.arange(10) * pq.s + xp = x.rescale(pq.ms) + data_array = DataArray(data=data, coords={"x": x, "xp": ("x", xp)}, dims=["x"]) + + assert_equal_with_units(data, data_array) + assert_equal_with_units(xp, data_array.xp) + + +def test_arithmetics(): + x = np.arange(10) + y = np.arange(20) + + coords = {"x": x, "y": y} + + f = (np.arange(10 * 20).reshape(10, 20) + 1) * pq.V + g = np.arange(10 * 20).reshape(10, 20) * pq.A + + a = DataArray(data=f, coords=coords, dims=("x", "y")) + b = DataArray(data=g, coords=coords, dims=("x", "y")) + + assert_equal_with_units(a * b, f * g) + + # swapped dimension order + g = np.arange(20 * 10).reshape(20, 10) * pq.V + b = DataArray(data=g, coords=coords, dims=("y", "x")) + assert_equal_with_units(a + b, f + g.T) + + # broadcasting + g = (np.arange(10) + 1) * pq.m + b = DataArray(data=g, coords={"x": x}, dims=["x"]) + assert_equal_with_units(a / b, f / g[:, None]) + + +@pytest.mark.xfail(reason="units don't survive through combining yet") +def test_combine(): + from xarray import concat + + data = (np.arange(15) + 10) * pq.m + y = np.arange(len(data)) + + data_array = DataArray(data=data, coords={"y": y}, dims=["y"]) + a = data_array[:5] + b = data_array[5:] + + assert_equal_with_units(concat([a, b], dim="y"), data_array) + + +def test_unit_checking(): + coords = {"x": np.arange(10), "y": np.arange(20)} + + f = np.arange(10 * 20).reshape(10, 20) * pq.A + g = np.arange(10 * 20).reshape(10, 20) * pq.V + + a = DataArray(f, coords=coords, dims=("x", "y")) + b = DataArray(g, coords=coords, dims=("x", "y")) + with pytest.raises(ValueError, match="Unable to convert between units"): + a + b + + +@pytest.mark.xfail(reason="units in indexes not supported") +def test_units_in_indexes(): + """ Test if units survive through xarray indexes. + Indexes are borrowed from Pandas, and Pandas does not support + units. Therefore, we currently don't intend to support units on + indexes either. + """ + data = np.arange(15) * 10 + x = np.arange(len(data)) * pq.A + + data_array = DataArray(data=data, coords={"x": x}, dims=["x"]) + assert_equal_with_units(data_array.x, x) + + +def test_sel(): + data = np.arange(10 * 20).reshape(10, 20) * pq.m / pq.s + x = np.arange(data.shape[0]) * pq.m + y = np.arange(data.shape[1]) * pq.s + + data_array = DataArray(data=data, coords={"x": x, "y": y}, dims=("x", "y")) + assert_equal_with_units(data_array.sel(y=y[0]), data[:, 0]) + + +def test_mean(): + data = np.arange(10 * 20).reshape(10, 20) * pq.V + x = np.arange(data.shape[0]) + y = np.arange(data.shape[1]) + + data_array = DataArray(data=data, coords={"x": x, "y": y}, dims=("x", "y")) + assert_equal_with_units(data_array.mean("x"), data.mean(0))