diff --git a/doc/whats-new.rst b/doc/whats-new.rst index dbd3ebbfe7e..fcb14a7c29a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -75,6 +75,11 @@ Enhancements - In :py:meth:`~xarray.Dataset.to_zarr`, passing ``mode`` is not mandatory if ``append_dim`` is set, as it will automatically be set to ``'a'`` internally. By `David Brochart `_. + +- Added the ability to initialize an empty or full DataArray + with a single value. (:issue:`277`) + By `Gerardo Rivera `_. + - :py:func:`~xarray.Dataset.to_netcdf()` now supports the ``invalid_netcdf`` kwarg when used with ``engine="h5netcdf"``. It is passed to :py:func:`h5netcdf.File`. By `Ulrich Herter `_. diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index e3e543b6621..4b8d8acb513 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -158,6 +158,24 @@ def _infer_coords_and_dims( return new_coords, dims +def _check_data_shape(data, coords, dims): + if data is dtypes.NA: + data = np.nan + if coords is not None and utils.is_scalar(data, include_0d=False): + if utils.is_dict_like(coords): + if dims is None: + return data + else: + data_shape = tuple( + as_variable(coords[k], k).size if k in coords.keys() else 1 + for k in dims + ) + else: + data_shape = tuple(as_variable(coord, "foo").size for coord in coords) + data = np.full(data_shape, data) + return data + + class _LocIndexer: def __init__(self, data_array: "DataArray"): self.data_array = data_array @@ -234,7 +252,7 @@ class DataArray(AbstractArray, DataWithCoords): def __init__( self, - data: Any, + data: Any = dtypes.NA, coords: Union[Sequence[Tuple], Mapping[Hashable, Any], None] = None, dims: Union[Hashable, Sequence[Hashable], None] = None, name: Hashable = None, @@ -323,6 +341,7 @@ def __init__( if encoding is None: encoding = getattr(data, "encoding", None) + data = _check_data_shape(data, coords, dims) data = as_compatible_data(data) coords, dims = _infer_coords_and_dims(data.shape, coords, dims) variable = Variable(dims, data, attrs, encoding, fastpath=True) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index a8bb5532f1f..9e0037b4da0 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -29,8 +29,6 @@ import numpy as np import pandas as pd -from .pycompat import dask_array_type - K = TypeVar("K") V = TypeVar("V") T = TypeVar("T") @@ -269,16 +267,20 @@ def either_dict_or_kwargs( return cast(Mapping[Hashable, T], kw_kwargs) -def is_scalar(value: Any) -> bool: +def is_scalar(value: Any, include_0d: bool = True) -> bool: """Whether to treat a value as a scalar. Any non-iterable, string, or 0-D array """ + from .variable import NON_NUMPY_SUPPORTED_ARRAY_TYPES + + if include_0d: + include_0d = getattr(value, "ndim", None) == 0 return ( - getattr(value, "ndim", None) == 0 + include_0d or isinstance(value, (str, bytes)) or not ( - isinstance(value, (Iterable,) + dask_array_type) + isinstance(value, (Iterable,) + NON_NUMPY_SUPPORTED_ARRAY_TYPES) or hasattr(value, "__array_function__") ) ) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 798ef66c82e..f623ec9976f 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -1446,6 +1446,32 @@ def test_rename(self): renamed_kwargs = self.dv.x.rename(x="z").rename("z") assert_identical(renamed, renamed_kwargs) + def test_init_value(self): + expected = DataArray( + np.full((3, 4), 3), dims=["x", "y"], coords=[range(3), range(4)] + ) + actual = DataArray(3, dims=["x", "y"], coords=[range(3), range(4)]) + assert_identical(expected, actual) + + expected = DataArray( + np.full((1, 10, 2), 0), + dims=["w", "x", "y"], + coords={"x": np.arange(10), "y": ["north", "south"]}, + ) + actual = DataArray(0, dims=expected.dims, coords=expected.coords) + assert_identical(expected, actual) + + expected = DataArray( + np.full((10, 2), np.nan), coords=[("x", np.arange(10)), ("y", ["a", "b"])] + ) + actual = DataArray(coords=[("x", np.arange(10)), ("y", ["a", "b"])]) + assert_identical(expected, actual) + + with pytest.raises(KeyError): + DataArray(np.array(1), coords={"x": np.arange(10)}, dims=["x"]) + with raises_regex(ValueError, "does not match the 0 dim"): + DataArray(np.array(1), coords=[("x", np.arange(10))]) + def test_swap_dims(self): array = DataArray(np.random.randn(3), {"y": ("x", list("abc"))}, "x") expected = DataArray(array.values, {"y": list("abc")}, dims="y")