Fix iteration on rolling windows to handle no padding

Also changes the object returned during iteration, so that it is the same
(modulo transposition) as the view created during rolling_window.construct()
along the rolling axis.

Author: kmsquire

xarray/core/rolling.py

@@ -263,11 +263,47 @@ def __init__(
     def __iter__(self):
         if len(self.dim) > 1:
             raise ValueError("__iter__ is only supported for 1d-rolling")
-        stops = np.arange(1, len(self.window_labels) + 1)
-        starts = stops - int(self.window[0])
-        starts[: int(self.window[0])] = 0
-        for (label, start, stop) in zip(self.window_labels, starts, stops):
-            window = self.obj.isel(**{self.dim[0]: slice(start, stop)})
+        dim = self.dim[0]
+        center = self.center[0]
+        pad = self.pad[0]
+        window = self.window[0]
+        center_offset = window // 2 if center else 0
+
+        pads = utils.get_pads(self.dim, self.window, self.center, self.pad)
+        start_pad, end_pad = pads[dim]
+
+        # Select the proper subset of labels, based on whether or not to center and/or pad
+        first_label_idx = 0 if pad else center_offset if center else window - 1
+        last_label_idx = (
+            len(self.obj[dim])
+            if pad or not center
+            else len(self.obj[dim]) - center_offset
+        )
+
+        labels = (
+            self.obj[dim][slice(first_label_idx, last_label_idx)]
+            if self.obj[dim].coords
+            else np.arange(last_label_idx - first_label_idx)
+        )
+
+        padded_obj = self.obj.pad(pads, mode="constant", constant_values=dtypes.NA)
+
+        if pad and not center:
+            first_stop = 1
+            last_stop = len(self.obj[dim])
+        elif pad and center:
+            first_stop = end_pad + 1
+            last_stop = len(self.obj[dim]) + end_pad
+        elif not pad:
+            first_stop = window
+            last_stop = len(self.obj[dim])
+
+        # These are indicies into the padded array, so we need to add start_pad
+        stops = np.arange(first_stop, last_stop + 1) + start_pad
+        starts = stops - window
+
+        for (label, start, stop) in zip(labels, starts, stops):
+            window = padded_obj.isel({self.dim[0]: slice(start, stop)})
 
             counts = window.count(dim=self.dim[0])
             window = window.where(counts >= self.min_periods)
@@ -486,15 +522,17 @@ def _counts(self, keep_attrs):
         # array is faster to be reduced than object array.
         # The use of skipna==False is also faster since it does not need to
         # copy the strided array.
+        output_dim_coords = self._get_rolling_dim_coords()
         counts = (
             self.obj.notnull(keep_attrs=keep_attrs)
             .rolling(
                 center={d: self.center[i] for i, d in enumerate(self.dim)},
+                pad={p: self.pad[i] for i, p in enumerate(self.pad)},
                 **{d: w for d, w in zip(self.dim, self.window)},
             )
             .construct(rolling_dim, fill_value=False, keep_attrs=keep_attrs)
             .sum(dim=list(rolling_dim.values()), skipna=False, keep_attrs=keep_attrs)
-        )
+        ).sel(output_dim_coords)
         return counts
 
     def _bottleneck_reduce(self, func, keep_attrs, **kwargs):

xarray/tests/test_dataarray.py

@@ -6506,23 +6506,16 @@ def test_isin(da):
 @pytest.mark.parametrize("da", (1, 2), indirect=True)
 @pytest.mark.parametrize("center", (True, False, None))
 @pytest.mark.parametrize("pad", (True, False, None))
-def test_rolling_iter(da, center, pad):
-    rolling_obj = da.rolling(time=7, center=center, pad=pad)
+@pytest.mark.parametrize("min_periods", (1, 6, None))
+@pytest.mark.parametrize("window", (6, 7))
+def test_rolling_iter(da, center, pad, min_periods, window):
+    rolling_obj = da.rolling(
+        time=window, center=center, pad=pad, min_periods=min_periods
+    )
     rolling_obj_mean = rolling_obj.mean()
 
-    if pad:
-        expected_times = da["time"]
-    else:
-        if center:
-            expected_times = da["time"][slice(3, -3)]
-        else:
-            expected_times = da["time"][slice(6, None)]
-
-    assert len(rolling_obj.window_labels) == len(expected_times)
-    assert_identical(rolling_obj.window_labels, expected_times)
-
     for i, (label, window_da) in enumerate(rolling_obj):
-        assert label == da["time"].isel(time=i)
+        assert label == rolling_obj_mean["time"].isel(time=i)
 
         actual = rolling_obj_mean.isel(time=i)
         expected = window_da.mean("time")
@@ -6531,10 +6524,14 @@ def test_rolling_iter(da, center, pad):
         # as well as the closeness of the values.
         assert_array_equal(actual.isnull(), expected.isnull())
         if (~actual.isnull()).sum() > 0:
-            np.allclose(
-                actual.values[actual.values.nonzero()],
-                expected.values[expected.values.nonzero()],
-            )
+            if actual.ndim == 0:
+                actual_values = actual.values
+                expected_values = expected.values
+            else:
+                actual_values = actual.values[actual.values.nonzero()]
+                expected_values = expected.values[expected.values.nonzero()]
+
+            assert np.allclose(actual_values, expected_values)
 
 
 @pytest.mark.parametrize("da", (1,), indirect=True)