creation functions according to NEP 35 and updated *_like behavior

docs/user/numpy.ipynb

@@ -611,13 +611,17 @@
    "source": [
     "## Additional Comments\n",
     "\n",
-    "What follows is a short discussion about how NumPy support is implemented in Pint's `Quantity` Object.\n",
+    "What follows is a short discussion about how NumPy support is implemented in Pint's `Quantity` object.\n",
     "\n",
     "For the supported functions, Pint expects certain units and attempts to convert the input (or inputs). For example, the argument of the exponential function (`numpy.exp`) must be dimensionless. Units will be simplified (converting the magnitude appropriately) and `numpy.exp` will be applied to the resulting magnitude. If the input is not dimensionless, a `DimensionalityError` exception will be raised.\n",
     "\n",
     "In some functions that take 2 or more arguments (e.g. `arctan2`), the second argument is converted to the units of the first. Again, a `DimensionalityError` exception will be raised if this is not possible. ndarray or downcast type arguments are generally treated as if they were dimensionless quantities, whereas Pint defers to its declared upcast types by always returning `NotImplemented` when they are encountered (see above).\n",
     "\n",
-    "To achive these function and ufunc overrides, Pint uses the ``__array_function__`` and ``__array_ufunc__`` protocols respectively, as recommened by NumPy. This means that functions and ufuncs that Pint does not explicitly handle will error, rather than return a value with units stripped (in contrast to Pint's behavior prior to v0.10). For more\n",
+    "Array creation functions (including those that support the NEP35 `like` keyword argument) will return a quantity with the same unit / underlying array type as the input array. The only the exceptions to this pattern are:\n",
+    "- `full` and `full_like`, which return the same units as the `fill_value` keyword argument (or a non-quantity if the `fill_value` is not a quantity)\n",
+    "- `arange`, which returns the same units as `start`, `stop`, and `step`. More specifically, `np.arange(Q_(1, \"m\"), Q_(2, \"m\"), Q_(1, \"mm\"), like=Quantity(1, \"s\"))` is valid and returns an array with shape `(1000,)`.\n",
+    "\n",
+    "To achieve these function and ufunc overrides, Pint uses the ``__array_function__`` and ``__array_ufunc__`` protocols respectively, as recommened by NumPy. This means that functions and ufuncs that Pint does not explicitly handle will error, rather than return a value with units stripped (in contrast to Pint's behavior prior to v0.10). For more\n",
     "information on these protocols, see <https://docs.scipy.org/doc/numpy-1.17.0/user/basics.dispatch.html>.\n",
     "\n",
     "This behaviour introduces some performance penalties and increased memory usage. Quantities that must be converted to other units require additional memory and CPU cycles. Therefore, for numerically intensive code, you might want to convert the objects first and then use directly the magnitude, such as by using Pint's `wraps` utility (see [wrapping](wrapping.rst)).\n",

pint/facets/numpy/numpy_func.py

@@ -527,22 +527,19 @@ def _meshgrid(*xi, **kwargs):
 
 
 @implements("full_like", "function")
-def _full_like(a, fill_value, dtype=None, order="K", subok=True, shape=None):
-    # Make full_like by multiplying with array from ones_like in a
-    # non-multiplicative-unit-safe way
+def _full_like(a, fill_value, **kwargs):
     if hasattr(fill_value, "_REGISTRY"):
-        return fill_value._REGISTRY.Quantity(
-            (
-                np.ones_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
-                * fill_value.m
-            ),
-            fill_value.units,
-        )
+        units = fill_value.units
+        fill_value_ = fill_value.m
     else:
-        return (
-            np.ones_like(a, dtype=dtype, order=order, subok=subok, shape=shape)
-            * fill_value
-        )
+        units = None
+        fill_value_ = fill_value
+
+    magnitude = np.full_like(a.m, fill_value=fill_value_, **kwargs)
+    if units is not None:
+        return fill_value._REGISTRY.Quantity(magnitude, units)
+    else:
+        return magnitude
 
 
 @implements("interp", "function")
@@ -904,9 +901,6 @@ def implementation(a, *args, **kwargs):
     "isreal",
     "iscomplex",
     "shape",
-    "ones_like",
-    "zeros_like",
-    "empty_like",
     "argsort",
     "argmin",
     "argmax",
@@ -935,6 +929,93 @@ def implementation(a, *args, **kwargs):
     implement_func("function", func_str, input_units=None, output_unit="variance")
 
 
+for func_str in ["ones_like", "zeros_like", "empty_like"]:
+    implement_func("function", func_str, input_units=None, output_unit="match_input")
+
+
+nep35_function_names = set()
+
+
+def register_nep35_function(func_str):
+    nep35_function_names.add(func_str)
+
+    def wrapper(f):
+        return f
+
+    return wrapper
+
+
+def implement_nep35_func(func_str):
+    # If NumPy is not available, do not attempt implement that which does not exist
+    if np is None:
+        return
+
+    func = getattr(np, func_str)
+
+    @register_nep35_function(func_str)
+    @implements(func_str, "function")
+    def implementation(*args, like, **kwargs):
+        args, kwargs = convert_to_consistent_units(*args, **kwargs)
+        result = func(*args, like=like.magnitude, **kwargs)
+        return like._REGISTRY.Quantity(result, like.units)
+
+
+# generic implementations
+for func_str in {
+    "array",
+    "asarray",
+    "asanyarray",
+    "arange",
+    "ones",
+    "zeros",
+    "empty",
+    "identity",
+    "eye",
+}:
+    implement_nep35_func(func_str)
+
+
+@register_nep35_function("full")
+@implements("full", "function")
+def _full(shape, fill_value, dtype=None, order="C", *, like):
+    if hasattr(fill_value, "_REGISTRY"):
+        units = fill_value.units
+        fill_value_ = fill_value.m
+    else:
+        units = None
+        fill_value_ = fill_value
+
+    magnitude = np.full(
+        shape=shape,
+        fill_value=fill_value_,
+        dtype=dtype,
+        order=order,
+        like=like.magnitude,
+    )
+    if units is not None:
+        return fill_value._REGISTRY.Quantity(magnitude, units)
+    else:
+        return like._REGISTRY.Quantity(magnitude, units)
+
+
+@register_nep35_function("arange")
+@implements("arange", "function")
+def _arange(start, stop=None, step=None, dtype=None, *, like):
+    args = [start, stop, step]
+    if any(_is_quantity(arg) for arg in args):
+        args, kwargs = convert_to_consistent_units(
+            start,
+            stop,
+            step,
+            pre_calc_units=like.units,
+            like=like,
+        )
+    else:
+        kwargs = {"like": like.magnitude}
+
+    return like._REGISTRY.Quantity(np.arange(*args, dtype=dtype, **kwargs), like.units)
+
+
 def numpy_wrap(func_type, func, args, kwargs, types):
     """Return the result from a NumPy function/ufunc as wrapped by Pint."""
 

pint/facets/numpy/quantity.py

@@ -22,6 +22,7 @@
     get_op_output_unit,
     matching_input_copy_units_output_ufuncs,
     matching_input_set_units_output_ufuncs,
+    nep35_function_names,
     numpy_wrap,
     op_units_output_ufuncs,
     set_units_ufuncs,
@@ -61,6 +62,10 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         return numpy_wrap("ufunc", ufunc, inputs, kwargs, types)
 
     def __array_function__(self, func, types, args, kwargs):
+        nep35_functions = {getattr(np, name) for name in nep35_function_names}
+        if func in nep35_functions:
+            kwargs["like"] = self
+
         return numpy_wrap("function", func, args, kwargs, types)
 
     _wrapped_numpy_methods = ["flatten", "astype", "item"]

pint/testsuite/helpers.py

@@ -109,6 +109,13 @@ def requires_not_array_function_protocol():
     )
 
 
+def requires_numpy_nep35():
+    return pytest.mark.skipif(
+        not version_parse(NUMPY_VER) >= version_parse("1.20.0"),
+        reason="Needs NEP 35, which is supported from numpy=1.20.0",
+    )
+
+
 def requires_numpy_previous_than(version):
     if not HAS_NUMPY:
         return pytest.mark.skip("Requires NumPy")

pint/testsuite/test_numpy.py

@@ -57,17 +57,22 @@ class TestNumpyArrayCreation(TestNumpyMethods):
 
     @helpers.requires_array_function_protocol()
     def test_ones_like(self):
-        self.assertNDArrayEqual(np.ones_like(self.q), np.array([[1, 1], [1, 1]]))
+        helpers.assert_quantity_equal(
+            np.ones_like(self.q), self.Q_([[1, 1], [1, 1]], self.q.units)
+        )
 
     @helpers.requires_array_function_protocol()
     def test_zeros_like(self):
-        self.assertNDArrayEqual(np.zeros_like(self.q), np.array([[0, 0], [0, 0]]))
+        helpers.assert_quantity_equal(
+            np.zeros_like(self.q), self.Q_([[0, 0], [0, 0]], self.q.units)
+        )
 
     @helpers.requires_array_function_protocol()
     def test_empty_like(self):
         ret = np.empty_like(self.q)
-        assert ret.shape == (2, 2)
-        assert isinstance(ret, np.ndarray)
+        expected = self.Q_(np.empty_like(self.q.magnitude), self.q.units)
+
+        helpers.assert_quantity_equal(ret, expected)
 
     @helpers.requires_array_function_protocol()
     def test_full_like(self):
@@ -77,6 +82,103 @@ def test_full_like(self):
         )
         self.assertNDArrayEqual(np.full_like(self.q, 2), np.array([[2, 2], [2, 2]]))
 
+    @helpers.requires_numpy_nep35()
+    def test_array(self):
+        x = [0, 1, 2, 3]
+        actual = np.array(x, like=self.q)
+        expected = self.Q_(x, self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_asarray(self):
+        x = [0, 1, 2, 3]
+        actual = np.asarray(x, like=self.q)
+        expected = self.Q_(x, self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_asanyarray(self):
+        x = [0, 1, 2, 3]
+        actual = np.asanyarray(x, like=self.q)
+        expected = self.Q_(x, self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_arange(self):
+        actual = np.arange(10, like=self.q)
+        expected = self.Q_(np.arange(10), self.q.units)
+        helpers.assert_quantity_equal(actual, expected)
+
+        actual = np.arange(
+            self.Q_(1, "kg"),
+            self.Q_(5, "kg"),
+            self.Q_(100, "g"),
+            like=self.Q_([0], "kg"),
+        )
+        expected = self.Q_(np.arange(1, 5, 0.1), "kg")
+        helpers.assert_quantity_equal(actual, expected)
+
+    # before 1.23.0, ones seems to be a pure python function with changing address
+    @helpers.requires_numpy_at_least("1.23.0")
+    @helpers.requires_numpy_nep35()
+    def test_ones(self):
+        shape = (2, 3)
+        actual = np.ones(shape=shape, like=self.q)
+        expected = self.Q_(np.ones(shape=shape), self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_zeros(self):
+        shape = (2, 3)
+        actual = np.zeros(shape=shape, like=self.q)
+        expected = self.Q_(np.zeros(shape=shape), self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_empty(self):
+        shape = (2, 3)
+        actual = np.empty(shape=shape, like=self.q)
+        expected = self.Q_(np.empty(shape=shape), self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    @helpers.requires_numpy_nep35()
+    def test_full(self):
+        shape = (2, 2)
+
+        actual = np.full(
+            shape=shape, fill_value=self.Q_(0, self.ureg.degC), like=self.q
+        )
+        expected = self.Q_([[0, 0], [0, 0]], self.ureg.degC)
+        helpers.assert_quantity_equal(actual, expected)
+
+        actual = np.full(shape=shape, fill_value=2, like=self.q)
+        expected = self.Q_([[2, 2], [2, 2]], "dimensionless")
+        helpers.assert_quantity_equal(actual, expected)
+
+    # before 1.23.0, identity seems to be a pure python function with changing address
+    @helpers.requires_numpy_at_least("1.23.0")
+    @helpers.requires_numpy_nep35()
+    def test_identity(self):
+        actual = np.identity(10, like=self.q)
+        expected = self.Q_(np.identity(10), self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
+    # before 1.23.0, eye seems to be a pure python function with changing address
+    @helpers.requires_numpy_at_least("1.23.0")
+    @helpers.requires_numpy_nep35()
+    def test_eye(self):
+        actual = np.eye(10, like=self.q)
+        expected = self.Q_(np.eye(10), self.q.units)
+
+        helpers.assert_quantity_equal(actual, expected)
+
 
 class TestNumpyArrayManipulation(TestNumpyMethods):
     # TODO