BUG: fix unexpected return of np.pad with mode=wrap (#22575)

np.pad with mode="wrap" returns unexpected result that original data is not strictly looped in padding. This may happen in some occassions when padding widths in the same dimension are unbalanced (see added testcase in test_arraypad.py and the related issue). The reason is the function pad makes iterative calls of _set_wrap_both() in the above situation, yet period for padding is not correctly computed in each iteration.

The bug is fixed by guaranteeing that period is always a multiple of original data size, and also be the possible maximum for computation efficiency.

Closes #22464

Co-authored-by: Lars Grüter <lagru+github@mailbox.org>

1 files changed, 18 insertions, 13 deletions

diff --git a/numpy/lib/arraypad.py b/numpy/lib/arraypad.py
index 28123246d..b06a645d8 100644
--- a/numpy/lib/arraypad.py
+++ b/numpy/lib/arraypad.py
@@ -378,7 +378,7 @@ def _set_reflect_both(padded, axis, width_pair, method, include_edge=False):
     return left_pad, right_pad
 
 
-def _set_wrap_both(padded, axis, width_pair):
+def _set_wrap_both(padded, axis, width_pair, original_period):
     """
     Pad `axis` of `arr` with wrapped values.
 
@@ -391,6 +391,8 @@ def _set_wrap_both(padded, axis, width_pair):
     width_pair : (int, int)
         Pair of widths that mark the pad area on both sides in the given
         dimension.
+    original_period : int
+        Original length of data on `axis` of `arr`.
 
     Returns
     -------
@@ -400,6 +402,9 @@ def _set_wrap_both(padded, axis, width_pair):
     """
     left_pad, right_pad = width_pair
     period = padded.shape[axis] - right_pad - left_pad
+    # Avoid wrapping with only a subset of the original area by ensuring period
+    # can only be a multiple of the original area's length.
+    period = period // original_period * original_period
 
     # If the current dimension of `arr` doesn't contain enough valid values
     # (not part of the undefined pad area) we need to pad multiple times.
@@ -410,14 +415,12 @@ def _set_wrap_both(padded, axis, width_pair):
 
     if left_pad > 0:
         # Pad with wrapped values on left side
-        # First slice chunk from right side of the non-pad area.
+        # First slice chunk from left side of the non-pad area.
         # Use min(period, left_pad) to ensure that chunk is not larger than
-        # pad area
-        right_slice = _slice_at_axis(
-            slice(-right_pad - min(period, left_pad),
-                  -right_pad if right_pad != 0 else None),
-            axis
-        )
+        # pad area.
+        slice_end = left_pad + period
+        slice_start = slice_end - min(period, left_pad)
+        right_slice = _slice_at_axis(slice(slice_start, slice_end), axis)
         right_chunk = padded[right_slice]
 
         if left_pad > period:
@@ -431,11 +434,12 @@ def _set_wrap_both(padded, axis, width_pair):
 
     if right_pad > 0:
         # Pad with wrapped values on right side
-        # First slice chunk from left side of the non-pad area.
+        # First slice chunk from right side of the non-pad area.
         # Use min(period, right_pad) to ensure that chunk is not larger than
-        # pad area
-        left_slice = _slice_at_axis(
-            slice(left_pad, left_pad + min(period, right_pad),), axis)
+        # pad area.
+        slice_start = -right_pad - period
+        slice_end = slice_start + min(period, right_pad)
+        left_slice = _slice_at_axis(slice(slice_start, slice_end), axis)
         left_chunk = padded[left_slice]
 
         if right_pad > period:
@@ -867,11 +871,12 @@ def pad(array, pad_width, mode='constant', **kwargs):
     elif mode == "wrap":
         for axis, (left_index, right_index) in zip(axes, pad_width):
             roi = _view_roi(padded, original_area_slice, axis)
+            original_period = padded.shape[axis] - right_index - left_index
             while left_index > 0 or right_index > 0:
                 # Iteratively pad until dimension is filled with wrapped
                 # values. This is necessary if the pad area is larger than
                 # the length of the original values in the current dimension.
                 left_index, right_index = _set_wrap_both(
-                    roi, axis, (left_index, right_index))
+                    roi, axis, (left_index, right_index), original_period)
 
     return padded