diff --git a/xarray/core/common.py b/xarray/core/common.py index 5beb5234d4c..b7f33c23ca8 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -773,10 +773,10 @@ def isin(self, test_elements): -------- numpy.isin """ - from .computation import apply_ufunc from .dataset import Dataset from .dataarray import DataArray from .variable import Variable + from .computation import apply_ufunc if isinstance(test_elements, Dataset): raise TypeError( @@ -794,6 +794,46 @@ def isin(self, test_elements): dask='allowed', ) + def astype(self, dtype, casting = 'unsafe', copy = True): + """ + Copy of the xarray object, with data cast to a specified type. + Leaves coordinate dtype unchanged. + + Parameters + ---------- + dtype : str or dtype + Typecode or data-type to which the array is cast. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'unsafe' + for backwards compatibility. + + * 'no' means the data types should not be cast at all. + * 'equiv' means only byte-order changes are allowed. + * 'safe' means only casts which can preserve values are allowed. + * 'same_kind' means only safe casts or casts within a kind, + like float64 to float32, are allowed. + * 'unsafe' means any data conversions may be done. + copy : bool, optional + By default, astype always returns a newly allocated array. If this + is set to False and the `dtype` requirement is satisfied, the input + array is returned instead of a copy. + + See also + -------- + np.ndarray.astype + dask.array.Array.astype + """ + from .computation import apply_ufunc + + return apply_ufunc( + duck_array_ops.astype, + self, + dtype, + keep_attrs=True, + kwargs={'casting':casting,'copy':copy}, + dask='allowed' + ) + def __enter__(self): return self diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index faea30cdd99..a43b6c67f0e 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -95,6 +95,9 @@ def isnull(data): tensordot = _dask_or_eager_func('tensordot', n_array_args=2) +def astype(data, dtype, **kwargs): + return data.astype(dtype, **kwargs) + def asarray(data): return data if isinstance(data, dask_array_type) else np.asarray(data) diff --git a/xarray/core/ops.py b/xarray/core/ops.py index d9e8ceb65d5..dc9f38a0d89 100644 --- a/xarray/core/ops.py +++ b/xarray/core/ops.py @@ -36,7 +36,7 @@ NUMPY_SAME_METHODS = ['item', 'searchsorted'] # methods which don't modify the data shape, so the result should still be # wrapped in an Variable/DataArray -NUMPY_UNARY_METHODS = ['astype', 'argsort', 'clip', 'conj', 'conjugate'] +NUMPY_UNARY_METHODS = ['argsort', 'clip', 'conj', 'conjugate'] PANDAS_UNARY_FUNCTIONS = ['isnull', 'notnull'] # methods which remove an axis REDUCE_METHODS = ['all', 'any'] diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 5ec85c159a2..f9081820e79 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -302,6 +302,40 @@ def data(self, data): "replacement data must match the Variable's shape") self._data = data + def astype(self, dtype, casting = 'unsafe', copy = True): + """ + Copy of the Variable object, with data cast to a specified type. + + Parameters + ---------- + dtype : str or dtype + Typecode or data-type to which the array is cast. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + Controls what kind of data casting may occur. Defaults to 'unsafe' + for backwards compatibility. + + * 'no' means the data types should not be cast at all. + * 'equiv' means only byte-order changes are allowed. + * 'safe' means only casts which can preserve values are allowed. + * 'same_kind' means only safe casts or casts within a kind, + like float64 to float32, are allowed. + * 'unsafe' means any data conversions may be done. + copy : bool, optional + By default, astype always returns a newly allocated array. If this + is set to False and the `dtype` requirement is satisfied, the input + array is returned instead of a copy. + + See also + -------- + np.ndarray.astype + dask.array.Array.astype + """ + self.data = duck_array_ops.astype(self.data, + dtype, + casting = casting, + copy = copy) + return self + def load(self, **kwargs): """Manually trigger loading of this variable's data from disk or a remote source into memory and return this variable. diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 32ab3a634cb..3877dc14a8f 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -1446,8 +1446,8 @@ def test_dataset_getitem(self): def test_array_interface(self): assert_array_equal(np.asarray(self.dv), self.x) - # test patched in methods assert_array_equal(self.dv.astype(float), self.v.astype(float)) + # test patched in methods assert_array_equal(self.dv.argsort(), self.v.argsort()) assert_array_equal(self.dv.clip(2, 3), self.v.clip(2, 3)) # test ufuncs @@ -1458,6 +1458,10 @@ def test_array_interface(self): bar = Variable(['x', 'y'], np.zeros((10, 20))) assert_equal(self.dv, np.maximum(self.dv, bar)) + def test_astype_attrs(self): + mda = self.mda.astype(bool) + assert list(mda.attrs.items()) == list(self.mda.attrs.items()) + def test_is_null(self): x = np.random.RandomState(42).randn(5, 6) x[x < 0] = np.nan diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 722d1af14f7..a826af162dc 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -343,7 +343,6 @@ def test_array_interface(self): x = np.arange(5) v = self.cls(['x'], x) assert_array_equal(np.asarray(v), x) - # test patched in methods assert_array_equal(v.astype(float), x.astype(float)) # think this is a break, that argsort changes the type assert_identical(v.argsort(), v.to_base_variable())