diff --git a/doc/whats-new.rst b/doc/whats-new.rst index f6851d82e62..bc3bd210a51 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -39,6 +39,9 @@ Bug Fixes By `Ian Hunt-Isaak `_ - Coerce masked dask arrays to filled (:issue:`9374` :pull:`11157`). By `Julia Signell `_ +- Expose ``coord_pad_mode`` and associated parameters in ``Dataset.pad`` and + ``DataArray.pad`` (:issue:`3868` :issue:`6425` :pull:`11213`). By `Ian Cooke + `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 840b7242d64..67586fb363c 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -5840,6 +5840,19 @@ def pad( end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, reflect_type: PadReflectOptions = None, keep_attrs: bool | None = None, + coord_pad_mode: PadModeOptions | None = None, + coord_end_values: int + | tuple[int, int] + | Mapping[Any, tuple[int, int]] + | None = None, + coord_constant_values: float + | tuple[float, float] + | Mapping[Any, tuple[float, float]] + | None = None, + coord_stat_length: ( + int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None + ) = None, + coord_reflect_type: PadReflectOptions = None, **pad_width_kwargs: Any, ) -> Self: """Pad this array along one or more dimensions. @@ -5921,6 +5934,11 @@ def pad( If True, the attributes (``attrs``) will be copied from the original object to the new one. If False, the new object will be returned without attributes. + coord_pad_mode : see ``mode``, but this will be used to extend the coordinates + coord_end_values : see ``end_values`` + coord_constant_values : see ``constant_values`` + coord_stat_length : see ``stat_length`` + coord_reflect_type : see ``reflect_type`` **pad_width_kwargs The keyword arguments form of ``pad_width``. One of ``pad_width`` or ``pad_width_kwargs`` must be provided. @@ -5980,6 +5998,23 @@ def pad( * x (x) float64 32B nan 0.0 1.0 nan z (x) float64 32B nan 100.0 200.0 nan * y (y) int64 32B 10 20 30 40 + + Note that the default behavior of coordinate padding uses NaNs (or NaTs), which + requires that integer types be promoted to floats. Providing values via + ``coord_constant_values`` then follows the coercion behavior mentioned above, + where array types are preserved: + + >>> da.pad(x=1, constant_values=1.23456789, coord_constant_values=-999.9) + Size: 128B + array([[ 1, 1, 1, 1], + [ 0, 1, 2, 3], + [10, 11, 12, 13], + [ 1, 1, 1, 1]]) + Coordinates: + * x (x) int64 32B -999 0 1 -999 + z (x) int64 32B -999 100 200 -999 + * y (y) int64 32B 10 20 30 40 + """ ds = self._to_temp_dataset().pad( pad_width=pad_width, @@ -5989,6 +6024,11 @@ def pad( end_values=end_values, reflect_type=reflect_type, keep_attrs=keep_attrs, + coord_pad_mode=coord_pad_mode, + coord_end_values=coord_end_values, + coord_constant_values=coord_constant_values, + coord_stat_length=coord_stat_length, + coord_reflect_type=coord_reflect_type, **pad_width_kwargs, ) return self._from_temp_dataset(ds) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 149415847f3..4866e14c45a 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -8975,6 +8975,16 @@ def pad( end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, reflect_type: PadReflectOptions = None, keep_attrs: bool | None = None, + coord_pad_mode: PadModeOptions | None = None, + coord_end_values: int + | tuple[int, int] + | Mapping[Any, tuple[int, int]] + | None = None, + coord_constant_values: T_DatasetPadConstantValues | None = None, + coord_stat_length: ( + int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None + ) = None, + coord_reflect_type: PadReflectOptions = None, **pad_width_kwargs: Any, ) -> Self: """Pad this dataset along one or more dimensions. @@ -9058,6 +9068,11 @@ def pad( If True, the attributes (``attrs``) will be copied from the original object to the new one. If False, the new object will be returned without attributes. + coord_pad_mode : see ``mode``, but this will be used to extend the coordinates + coord_end_values : see ``end_values`` + coord_constant_values : see ``constant_values`` + coord_stat_length : see ``stat_length`` + coord_reflect_type : see ``reflect_type`` **pad_width_kwargs The keyword arguments form of ``pad_width``. One of ``pad_width`` or ``pad_width_kwargs`` must be provided. @@ -9092,17 +9107,31 @@ def pad( """ pad_width = either_dict_or_kwargs(pad_width, pad_width_kwargs, "pad") - if mode in ("edge", "reflect", "symmetric", "wrap"): - coord_pad_mode = mode - coord_pad_options = { - "stat_length": stat_length, - "constant_values": constant_values, - "end_values": end_values, - "reflect_type": reflect_type, - } - else: - coord_pad_mode = "constant" - coord_pad_options = {} + if not coord_pad_mode: + # if not provided, use mode dependent default + if mode in ("edge", "reflect", "symmetric", "wrap"): + coord_pad_mode = mode + else: + coord_pad_mode = "constant" + + coord_pad_options = { + "stat_length": coord_stat_length, + "constant_values": coord_constant_values, + "end_values": coord_end_values, + "reflect_type": coord_reflect_type, + } + + if coord_pad_mode in ("edge", "reflect", "symmetric", "wrap"): + # This block is for backward compatibility, if we can break that, + # then this block would be unnecessary + if coord_pad_options["stat_length"] is None: + coord_pad_options["stat_length"] = stat_length + if coord_pad_options["constant_values"] is None: + coord_pad_options["constant_values"] = constant_values + if coord_pad_options["end_values"] is None: + coord_pad_options["end_values"] = end_values + if coord_pad_options["reflect_type"] is None: + coord_pad_options["reflect_type"] = reflect_type if keep_attrs is None: keep_attrs = _get_keep_attrs(default=True) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index f00c5e4aed7..e14787db948 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -7538,6 +7538,54 @@ def test_pad_keep_attrs(self, keep_attrs, attrs, expected) -> None: ) xr.testing.assert_identical(actual, expected) + @pytest.mark.parametrize( + ["coord_pad_mode", "coord_pad_kwargs", "coord0", "coord9"], + [ + pytest.param("constant", {}, np.nan, np.nan, id="constant(default)"), + pytest.param( + "constant", {"coord_constant_values": 5}, 5, 5, id="constant(5)" + ), + pytest.param("edge", {}, 10, 50, id="edge"), + pytest.param( + "linear_ramp", {"coord_end_values": (-1, 7)}, -1, 7, id="linear_ramp" + ), + pytest.param( + "maximum", {"coord_stat_length": (3, 4)}, 30, 50, id="maximum" + ), + pytest.param("mean", {"coord_stat_length": (3, 4)}, 20, 35, id="mean"), + pytest.param("median", {"coord_stat_length": (3, 4)}, 20, 35, id="median"), + pytest.param( + "minimum", {"coord_stat_length": (3, 4)}, 10, 20, id="minimum" + ), + pytest.param( + "reflect", {"coord_reflect_type": "odd"}, -20, 70, id="reflect odd" + ), + pytest.param( + "reflect", {"coord_reflect_type": "even"}, 40, 30, id="reflect even" + ), + pytest.param("symmetric", {}, 30, 40, id="symmetric"), + pytest.param("wrap", {}, 30, 20, id="wrap"), + ], + ) + def test_pad_coord_pad_mode( + self, coord_pad_mode, coord_pad_kwargs, coord0, coord9 + ) -> None: + ds = xr.Dataset( + {"a": ("x", [1, 2, 3, 4, 5])}, coords={"x": [10, 20, 30, 40, 50]} + ) + pad_width = {"x": (3, 2)} # pad 5 elem array out to 10 + padded = ds.pad(pad_width, coord_pad_mode=coord_pad_mode, **coord_pad_kwargs) + + assert padded.sizes["x"] == 10 + if np.isnan(coord0): + assert np.isnan(padded["x"].values[0]) + else: + assert padded["x"].values[0] == coord0 + if np.isnan(coord9): + assert np.isnan(padded["x"].values[-1]) + else: + assert padded["x"].values[-1] == coord9 + def test_astype_attrs(self) -> None: data = create_test_data(seed=123) data.attrs["foo"] = "bar"