Merge #956

956: Add np.pad and np.resize implementations r=hgrecco a=jthielen

As requested in pydata/xarray#3643 (comment), this PR adds implementations for `np.pad` and `np.resize`, with accompanying tests.

- ~~Closes (insert issue number)~~ (#905 follow-up)
- [x] Executed ``black -t py36 . && isort -rc . && flake8`` with no errors
- [x] The change is fully covered by automated unit tests
- [x] Documented in docs/ as appropriate
- ~~Added an entry to the CHANGES file~~ (#905 follow-up)


Co-authored-by: Jon Thielen <github@jont.cc>

docs/numpy.rst

@@ -112,7 +112,7 @@ The following ufuncs_ can be applied to a Quantity object:
 
 And the following NumPy functions:
 
-- alen, amax, amin, append, argmax, argmin, argsort, around, atleast_1d, atleast_2d, atleast_3d, average, block, broadcast_to, clip, column_stack, compress, concatenate, copy, copyto, count_nonzero, cross, cumprod, cumproduct, cumsum, diagonal, diff, dot, dstack, ediff1d, einsum, empty_like, expand_dims, fix, flip, full_like, gradient, hstack, insert, interp, isclose, iscomplex, isin, isreal, linspace, mean, median, meshgrid, moveaxis, nan_to_num, nanargmax, nanargmin, nancumprod, nancumsum, nanmax, nanmean, nanmedian, nanmin, nanpercentile, nanstd, nanvar, ndim, nonzero, ones_like, percentile, ptp, ravel, result_type, rollaxis, rot90, round\_, searchsorted, shape, size, sort, squeeze, stack, std, sum, swapaxes, tile, transpose, trapz, trim_zeros, unwrap, var, vstack, where, zeros_like
+- alen, amax, amin, append, argmax, argmin, argsort, around, atleast_1d, atleast_2d, atleast_3d, average, block, broadcast_to, clip, column_stack, compress, concatenate, copy, copyto, count_nonzero, cross, cumprod, cumproduct, cumsum, diagonal, diff, dot, dstack, ediff1d, einsum, empty_like, expand_dims, fix, flip, full_like, gradient, hstack, insert, interp, isclose, iscomplex, isin, isreal, linspace, mean, median, meshgrid, moveaxis, nan_to_num, nanargmax, nanargmin, nancumprod, nancumsum, nanmax, nanmean, nanmedian, nanmin, nanpercentile, nanstd, nanvar, ndim, nonzero, ones_like, pad, percentile, ptp, ravel, resize, result_type, rollaxis, rot90, round\_, searchsorted, shape, size, sort, squeeze, stack, std, sum, swapaxes, tile, transpose, trapz, trim_zeros, unwrap, var, vstack, where, zeros_like
 
 And the following `NumPy ndarray methods`_:
 

pint/numpy_func.py

@@ -597,6 +597,31 @@ def _isin(element, test_elements, assume_unique=False, invert=False):
     return np.isin(element.m, test_elements, assume_unique=assume_unique, invert=invert)
 
 
+@implements("pad", "function")
+def _pad(array, pad_width, mode="constant", **kwargs):
+    def _recursive_convert(arg, unit):
+        if iterable(arg):
+            return tuple(_recursive_convert(a, unit=unit) for a in arg)
+        elif _is_quantity(arg):
+            return arg.m_as(unit)
+        else:
+            return arg
+
+    # pad only dispatches on array argument, so we know it is a Quantity
+    units = array.units
+
+    # Handle flexible constant_values and end_values, converting to units if Quantity
+    # and ignoring if not
+    if mode == "constant":
+        kwargs["constant_values"] = _recursive_convert(kwargs["constant_values"], units)
+    elif mode == "linear_ramp":
+        kwargs["end_values"] = _recursive_convert(kwargs["end_values"], units)
+
+    return units._REGISTRY.Quantity(
+        np.pad(array._magnitude, pad_width, mode=mode, **kwargs), units
+    )
+
+
 # Implement simple matching-unit or stripped-unit functions based on signature
 
 
@@ -675,6 +700,7 @@ def implementation(*args, **kwargs):
     ("tile", "A", True),
     ("rot90", "m", True),
     ("insert", ["arr", "values"], True),
+    ("resize", "a", True),
 ]:
     implement_consistent_units_by_argument(func_str, unit_arguments, wrap_output)
 

pint/testsuite/test_numpy.py

@@ -1025,6 +1025,102 @@ def test_array_protocol_unavailable(self):
         for attr in ("__array_struct__", "__array_interface__"):
             self.assertRaises(AttributeError, getattr, self.q, attr)
 
+    @helpers.requires_array_function_protocol()
+    def test_resize(self):
+        self.assertQuantityEqual(
+            np.resize(self.q, (2, 4)), [[1, 2, 3, 4], [1, 2, 3, 4]] * self.ureg.m
+        )
+
+    @helpers.requires_array_function_protocol()
+    def test_pad(self):
+        # Tests reproduced with modification from NumPy documentation
+        a = [1, 2, 3, 4, 5] * self.ureg.m
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "constant", constant_values=(4, 600 * self.ureg.cm)),
+            [4, 4, 1, 2, 3, 4, 5, 6, 6, 6] * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "edge"), [1, 1, 1, 2, 3, 4, 5, 5, 5, 5] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "linear_ramp", end_values=(5, -4) * self.ureg.m),
+            [5, 3, 1, 2, 3, 4, 5, 2, -1, -4] * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2,), "maximum"), [5, 5, 1, 2, 3, 4, 5, 5, 5] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2,), "mean"), [3, 3, 1, 2, 3, 4, 5, 3, 3] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2,), "median"), [3, 3, 1, 2, 3, 4, 5, 3, 3] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(self.q, ((3, 2), (2, 3)), "minimum"),
+            [
+                [1, 1, 1, 2, 1, 1, 1],
+                [1, 1, 1, 2, 1, 1, 1],
+                [1, 1, 1, 2, 1, 1, 1],
+                [1, 1, 1, 2, 1, 1, 1],
+                [3, 3, 3, 4, 3, 3, 3],
+                [1, 1, 1, 2, 1, 1, 1],
+                [1, 1, 1, 2, 1, 1, 1],
+            ]
+            * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "reflect"), [3, 2, 1, 2, 3, 4, 5, 4, 3, 2] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "reflect", reflect_type="odd"),
+            [-1, 0, 1, 2, 3, 4, 5, 6, 7, 8] * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "symmetric"), [2, 1, 1, 2, 3, 4, 5, 5, 4, 3] * self.ureg.m
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "symmetric", reflect_type="odd"),
+            [0, 1, 1, 2, 3, 4, 5, 5, 6, 7] * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.pad(a, (2, 3), "wrap"), [4, 5, 1, 2, 3, 4, 5, 1, 2, 3] * self.ureg.m
+        )
+
+        def pad_with(vector, pad_width, iaxis, kwargs):
+            pad_value = kwargs.get("padder", 10)
+            vector[: pad_width[0]] = pad_value
+            vector[-pad_width[1] :] = pad_value
+
+        b = self.Q_(np.arange(6).reshape((2, 3)), "degC")
+        self.assertQuantityEqual(
+            np.pad(b, 2, pad_with),
+            self.Q_(
+                [
+                    [10, 10, 10, 10, 10, 10, 10],
+                    [10, 10, 10, 10, 10, 10, 10],
+                    [10, 10, 0, 1, 2, 10, 10],
+                    [10, 10, 3, 4, 5, 10, 10],
+                    [10, 10, 10, 10, 10, 10, 10],
+                    [10, 10, 10, 10, 10, 10, 10],
+                ],
+                "degC",
+            ),
+        )
+        self.assertQuantityEqual(
+            np.pad(b, 2, pad_with, padder=100),
+            self.Q_(
+                [
+                    [100, 100, 100, 100, 100, 100, 100],
+                    [100, 100, 100, 100, 100, 100, 100],
+                    [100, 100, 0, 1, 2, 100, 100],
+                    [100, 100, 3, 4, 5, 100, 100],
+                    [100, 100, 100, 100, 100, 100, 100],
+                    [100, 100, 100, 100, 100, 100, 100],
+                ],
+                "degC",
+            ),
+        )  # Note: Does not support Quantity pad_with vectorized callable use
+
 
 @unittest.skip
 class TestBitTwiddlingUfuncs(TestUFuncs):