Repository: yt-project/unyt
Pull request #293 (2022-10-11 17:36:09)

ENH: implement 8 new NEP 18 functions
implements https://github.com/yt-project/unyt/issues/289#issuecomment-1274964636

TODO:
- [x] get CI green for existing tests
- [x] add new tests


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 683266c..9be14eb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,4 +36,5 @@ version_scheme = "post-release"
 addopts = "--ignore=benchmarks --ignore=paper --ignore=unyt/_mpl_array_converter"
 filterwarnings = [
     "error",
+    "ignore:Matplotlib is currently using agg, which is a non-GUI backend, so cannot show the figure.:UserWarning"
 ]
diff --git a/setup.py b/setup.py
index f27ccf7..336617f 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@ with open("README.rst") as readme_file:
 with open("HISTORY.rst") as history_file:
     history = history_file.read()
 
-requirements = ["numpy>=1.17.5", "sympy>=1.5"]
+requirements = ["numpy>=1.17.5", "sympy>=1.5", "packaging>=20.9"]
 
 test_requirements = ["pytest"]
 
diff --git a/tox.ini b/tox.ini
index 2f5b2e5..51bc2c0 100644
--- a/tox.ini
+++ b/tox.ini
@@ -24,7 +24,7 @@ deps =
     dask[array,diagnostics]
     setuptools_scm
 commands =
-    pytest --cov=unyt --cov-append --doctest-modules --doctest-plus --doctest-rst --basetemp={envtmpdir} -W once
+    pytest --cov=unyt --cov-append --doctest-modules --doctest-plus --doctest-rst --basetemp={envtmpdir}
     coverage report --omit='.tox/*'
 
 [testenv:py38-versions]
@@ -61,7 +61,7 @@ depends = begin
 commands =
     # don't do doctests in rst files due to lack of way to specify optional
     # test dependencies there
-    pytest --cov=unyt --cov-append --doctest-modules --doctest-plus --basetemp={envtmpdir} -W once
+    pytest --cov=unyt --cov-append --doctest-modules --doctest-plus --basetemp={envtmpdir}
     coverage report --omit='.tox/*'
 
 [testenv:py38-docs]
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index f3e9899..ae29bdd 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -1,7 +1,9 @@
 import numpy as np
+from packaging.version import Version
 
-from unyt.exceptions import UnitConversionError
+from unyt.exceptions import UnitConversionError, UnitInconsistencyError
 
+NUMPY_VERSION = Version(np.__version__)
 _HANDLED_FUNCTIONS = {}
 
 
@@ -186,3 +188,90 @@ def histogramdd(
         density=density,
     )
     return counts, tuple(_bin * u for _bin, u in zip(bins, units))
+
+
+def _validate_units_consistency(arrs):
+    # NOTE: we cannot validate that all arrays are unyt_arrays
+    # by using this as a guard clause in unyt_array.__array_function__
+    # because it's already a necessary condition for numpy to use our
+    # custom implementations
+    a1 = arrs[0]
+    if not all(a.units == a1.units for a in arrs[1:]):
+        raise UnitInconsistencyError(*(a.units for a in arrs))
+
+
+@implements(np.concatenate)
+def concatenate(arrs, /, axis=0, out=None, dtype=None, casting="same_kind"):
+    _validate_units_consistency(arrs)
+    if NUMPY_VERSION >= Version("1.20"):
+        v = np.concatenate._implementation(
+            [_.ndview for _ in arrs], axis=axis, out=out, dtype=dtype, casting=casting
+        )
+    else:
+        v = np.concatenate._implementation(
+            [_.ndview for _ in arrs],
+            axis=axis,
+            out=out,
+        )
+    return v * arrs[0].units
+
+
+@implements(np.cross)
+def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
+    prod_units = a.units * b.units
+    return (
+        np.cross._implementation(
+            a.ndview, b.ndview, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis
+        )
+        * prod_units
+    )
+
+
+@implements(np.intersect1d)
+def intersect1d(arr1, arr2, /, assume_unique=False, return_indices=False):
+    _validate_units_consistency((arr1, arr2))
+    retv = np.intersect1d._implementation(
+        arr1.ndview,
+        arr2.ndview,
+        assume_unique=assume_unique,
+        return_indices=return_indices,
+    )
+    if return_indices:
+        return retv
+    else:
+        return retv * arr1.units
+
+
+@implements(np.union1d)
+def union1d(arr1, arr2, /):
+    _validate_units_consistency((arr1, arr2))
+    return np.union1d._implementation(arr1.ndview, arr2.ndview) * arr1.units
+
+
+@implements(np.linalg.norm)
+def norm(x, /, ord=None, axis=None, keepdims=False):
+    return (
+        np.linalg.norm._implementation(x.ndview, ord=ord, axis=axis, keepdims=keepdims)
+        * x.units
+    )
+
+
+@implements(np.vstack)
+def vstack(tup, /):
+    _validate_units_consistency(tup)
+    return np.vstack._implementation([_.ndview for _ in tup]) * tup[0].units
+
+
+@implements(np.hstack)
+def hstack(tup, /):
+    _validate_units_consistency(tup)
+    return np.vstack._implementation([_.ndview for _ in tup]) * tup[0].units
+
+
+@implements(np.stack)
+def stack(arrays, /, axis=0, out=None):
+    _validate_units_consistency(arrays)
+    return (
+        np.stack._implementation([_.ndview for _ in arrays], axis=axis, out=out)
+        * arrays[0].units
+    )
diff --git a/unyt/array.py b/unyt/array.py
index 9c47d26..47c6acb 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -2214,15 +2214,16 @@ def uconcatenate(arrs, axis=0):
     >>> from unyt import cm
     >>> A = [1, 2, 3]*cm
     >>> B = [2, 3, 4]*cm
-    >>> uconcatenate((A, B))
+    >>> uconcatenate((A, B)) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     unyt_array([1, 2, 3, 2, 3, 4], 'cm')
 
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.uconcatenate", replacement="use numpy.concatenate", since_version="3.0"
-    # )
-    v = np.concatenate(arrs, axis=axis)
+    warn_deprecated(
+        "unyt.uconcatenate", replacement="use numpy.concatenate", since_version="3.0"
+    )
+    v = np.concatenate._implementation(arrs, axis=axis)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
 
@@ -2234,11 +2235,16 @@ def ucross(arr1, arr2, registry=None, axisa=-1, axisb=-1, axisc=-1, axis=None):
     See the documentation of numpy.cross for full
     details.
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.ucross", replacement="use numpy.cross", since_version="3.0"
-    # )
-    v = np.cross(arr1, arr2, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis)
+    warn_deprecated(
+        "unyt.ucross",
+        replacement=(
+            "use numpy.cross (note that the *registry* argument will not be ported)"
+        ),
+        since_version="3.0",
+    )
+    v = np.cross._implementation(
+        arr1, arr2, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis
+    )
     units = arr1.units * arr2.units
     arr = unyt_array(v, units, registry=registry)
     return arr
@@ -2256,17 +2262,16 @@ def uintersect1d(arr1, arr2, assume_unique=False):
     >>> from unyt import cm
     >>> A = [1, 2, 3]*cm
     >>> B = [2, 3, 4]*cm
-    >>> uintersect1d(A, B)
+    >>> uintersect1d(A, B) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     unyt_array([2, 3], 'cm')
 
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.uintersect1d", replacement="use numpy.intersect1d", since_version="3.0"
-    # )
-    v = np.intersect1d(arr1, arr2, assume_unique=assume_unique)
-    v = _validate_numpy_wrapper_units(v, [arr1, arr2])
-    return v
+    warn_deprecated(
+        "unyt.uintersect1d", replacement="use numpy.intersect1d", since_version="3.0"
+    )
+    return np.intersect1d(arr1, arr2, assume_unique=assume_unique)
 
 
 def uunion1d(arr1, arr2):
@@ -2281,17 +2286,16 @@ def uunion1d(arr1, arr2):
     >>> from unyt import cm
     >>> A = [1, 2, 3]*cm
     >>> B = [2, 3, 4]*cm
-    >>> uunion1d(A, B)
+    >>> uunion1d(A, B) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     unyt_array([1, 2, 3, 4], 'cm')
 
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.uunion1d", replacement="use numpy.union1d", since_version="3.0"
-    # )
-    v = np.union1d(arr1, arr2)
-    v = _validate_numpy_wrapper_units(v, [arr1, arr2])
-    return v
+    warn_deprecated(
+        "unyt.uunion1d", replacement="use numpy.union1d", since_version="3.0"
+    )
+    return np.union1d(arr1, arr2)
 
 
 def unorm(data, ord=None, axis=None, keepdims=False):
@@ -2305,14 +2309,15 @@ def unorm(data, ord=None, axis=None, keepdims=False):
     --------
     >>> from unyt import km
     >>> data = [1, 2, 3]*km
-    >>> print(unorm(data))
+    >>> print(unorm(data)) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     3.7416573867739413 km
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.unorm", replacement="use numpy.norm", since_version="3.0"
-    # )
-    norm = np.linalg.norm(data, ord=ord, axis=axis, keepdims=keepdims)
+    warn_deprecated(
+        "unyt.unorm", replacement="use numpy.linalg.norm", since_version="3.0"
+    )
+    norm = np.linalg.norm._implementation(data, ord=ord, axis=axis, keepdims=keepdims)
     if norm.shape == ():
         return unyt_quantity(norm, data.units)
     return unyt_array(norm, data.units)
@@ -2328,12 +2333,14 @@ def udot(op1, op2):
     >>> from unyt import km, s
     >>> a = np.eye(2)*km
     >>> b = (np.ones((2, 2)) * 2)*s
-    >>> print(udot(a, b))
+    >>> print(udot(a, b)) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     [[2. 2.]
      [2. 2.]] km*s
     """
     warn_deprecated("unyt.udot", replacement="use numpy.dot", since_version="3.0")
-    dot = np.dot(op1.d, op2.d)
+    dot = np.dot._implementation(op1.d, op2.d)
     units = op1.units * op2.units
     if dot.shape == ():
         return unyt_quantity(dot, units)
@@ -2350,15 +2357,14 @@ def uvstack(arrs):
     >>> from unyt import km
     >>> a = [1, 2, 3]*km
     >>> b = [2, 3, 4]*km
-    >>> print(uvstack([a, b]))
+    >>> print(uvstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     [[1 2 3]
      [2 3 4]] km
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.uvstack", replacement="use numpy.vstack", since_version="3.0"
-    # )
-    v = np.vstack(arrs)
+    warn_deprecated("unyt.uvstack", replacement="use numpy.vstack", since_version="3.0")
+    v = np.vstack._implementation(arrs)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
 
@@ -2373,20 +2379,21 @@ def uhstack(arrs):
     >>> from unyt import km
     >>> a = [1, 2, 3]*km
     >>> b = [2, 3, 4]*km
-    >>> print(uhstack([a, b]))
+    >>> print(uhstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     [1 2 3 2 3 4] km
     >>> a = [[1],[2],[3]]*km
     >>> b = [[2],[3],[4]]*km
-    >>> print(uhstack([a, b]))
+    >>> print(uhstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     [[1 2]
      [2 3]
      [3 4]] km
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.uhstack", replacement="use numpy.hstack", since_version="3.0"
-    # )
-    v = np.hstack(arrs)
+    warn_deprecated("unyt.uhstack", replacement="use numpy.hstack", since_version="3.0")
+    v = np.hstack._implementation(arrs)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
 
@@ -2406,15 +2413,14 @@ def ustack(arrs, axis=0):
     >>> from unyt import km
     >>> a = [1, 2, 3]*km
     >>> b = [2, 3, 4]*km
-    >>> print(ustack([a, b]))
+    >>> print(ustack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    DeprecationWarning: ...
     [[1 2 3]
      [2 3 4]] km
     """
-    # TODO: see https://github.com/yt-project/unyt/issues/289
-    # warn_deprecated(
-    #    "unyt.ustack", replacement="use numpy.stack", since_version="3.0"
-    # )
-    v = np.stack(arrs, axis=axis)
+    warn_deprecated("unyt.ustack", replacement="use numpy.stack", since_version="3.0")
+    v = np.stack._implementation(arrs, axis=axis)
     v = _validate_numpy_wrapper_units(v, arrs)
     return v
 
diff --git a/unyt/exceptions.py b/unyt/exceptions.py
index a6a66e0..8414a61 100644
--- a/unyt/exceptions.py
+++ b/unyt/exceptions.py
@@ -47,6 +47,32 @@ class UnitOperationError(ValueError):
         return err
 
 
+class UnitInconsistencyError(UnitOperationError):
+    """An error raised for non-allower numpy calls on arrays with different units.
+
+    Example
+    -------
+
+    >>> import unyt as u
+    >>> import numpy as np
+    >>> np.concatenate([[1, 2, 3] * u.cm, [4, 5, 6] * u.s])\
+ # doctest: +IGNORE_EXCEPTION_DETAIL +NORMALIZE_WHITESPACE
+    Traceback (most recent call last):
+    ...
+    unyt.exceptions.UnitInconsistencyError: Expected all unyt_array arguments to
+    have identical units. Received mixed units (cm, s)
+    """
+
+    def __init__(self, *units):
+        self.units = units
+
+    def __str__(self):
+        return (
+            "Expected all unyt_array arguments to have identical units. "
+            f"Received mixed units ({', '.join(str(u) for u in self.units)})"
+        )
+
+
 class UnitConversionError(Exception):
     """An error raised when converting to a unit with different dimensions.
 


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index 81c5dbf..2c6f782 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -6,6 +6,7 @@ import pytest
 
 from unyt import cm, g, km, s
 from unyt.array import unyt_array
+from unyt.exceptions import UnitInconsistencyError
 
 
 def test_array_repr():
@@ -199,3 +200,89 @@ def test_histogramdd():
     assert 1 * xbins.units == 1 * x.units
     assert 1 * ybins.units == 1 * y.units
     assert 1 * zbins.units == 1 * z.units
+
+
+def test_concatenate():
+    x1 = np.random.normal(size=100) * cm
+    x2 = np.random.normal(size=100) * cm
+    res = np.concatenate((x1, x2))
+    assert 1 * res.units == 1 * cm
+    assert res.shape == (200,)
+
+
+def test_concatenate_different_units():
+    x1 = np.random.normal(size=100) * cm
+    x2 = np.random.normal(size=100) * s
+    with pytest.raises(
+        UnitInconsistencyError,
+        match=(
+            r"Expected all unyt_array arguments to have identical units\. "
+            r"Received mixed units \(cm, s\)"
+        ),
+    ):
+        np.concatenate((x1, x2))
+
+
+def test_cross():
+    x1 = np.random.random_sample((2, 2)) * cm
+    x2 = np.eye(2) * s
+    res = np.cross(x1, x2)
+    assert res.units == cm * s
+
+
+def test_intersect1d():
+    x1 = [1, 2, 3, 4, 5, 6, 7, 8] * cm
+    x2 = [0, 2, 4, 6, 8] * cm
+    res = np.intersect1d(x1, x2)
+    assert 1 * res.units == 1 * cm
+    np.testing.assert_array_equal(res, [2, 4, 6, 8])
+
+
+def test_intersect1d_return_indices():
+    x1 = [1, 2, 3, 4, 5, 6, 7, 8] * cm
+    x2 = [0, 2, 4, 6, 8] * cm
+    ures = np.intersect1d(x1, x2, return_indices=True)
+    rres = np.intersect1d(x1.d, x2.d, return_indices=True)
+    np.testing.assert_array_equal(ures, rres)
+
+
+def test_union1d():
+    x1 = [-1, 0, 1] * cm
+    x2 = [-2, -1, -3] * cm
+    res = np.union1d(x1, x2)
+    assert 1 * res.units == 1 * cm
+    np.testing.assert_array_equal(res, [-3, -2, -1, 0, 1])
+
+
+def test_linalg_norm():
+    x = [1, 1, 1] * s
+    res = np.linalg.norm(x)
+    assert 1 * res.units == 1 * s
+    assert res == pytest.approx(np.sqrt(3))
+
+
+def test_vstack():
+    x1 = [0, 1, 2] * cm
+    x2 = [3, 4, 5] * cm
+    res = np.vstack((x1, x2))
+    assert 1 * res.units == 1 * cm
+    np.testing.assert_array_equal(res, [[0, 1, 2], [3, 4, 5]])
+
+
+def test_hstack():
+    x1 = np.array([[0, 1, 2], [3, 4, 5]]) * cm
+    x2 = np.array([[6, 7, 8], [9, 10, 11]]) * cm
+    res = np.hstack((x1, x2))
+    assert 1 * res.units == 1 * cm
+    np.testing.assert_array_equal(res, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]])
+
+
+@pytest.mark.parametrize(
+    "axis, expected", [(0, [[0, 1, 2], [3, 4, 5]]), (1, [[0, 3], [1, 4], [2, 5]])]
+)
+def test_stack(axis, expected):
+    x1 = [0, 1, 2] * cm
+    x2 = [3, 4, 5] * cm
+    res = np.stack((x1, x2), axis=axis)
+    assert 1 * res.units == 1 * cm
+    np.testing.assert_array_equal(res, expected)
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index 6f8558f..f7746e3 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -2045,41 +2045,56 @@ def test_numpy_wrappers():
     arr_norm_answer = [norm_answer, np.sqrt(4**2 + 5**2 + 6**2)]
     dot_answer = 14
 
-    assert_array_equal(unyt_array(catenate_answer, "cm"), uconcatenate((a1, a2)))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unyt_array(catenate_answer, "cm"), uconcatenate((a1, a2)))
     assert_array_equal(catenate_answer, np.concatenate((a1, a2)))
 
-    assert_array_equal(unyt_array(intersect_answer, "cm"), uintersect1d(a1, a2))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unyt_array(intersect_answer, "cm"), uintersect1d(a1, a2))
     assert_array_equal(intersect_answer, np.intersect1d(a1, a2))
 
-    assert_array_equal(unyt_array(union_answer, "cm"), uunion1d(a1, a2))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unyt_array(union_answer, "cm"), uunion1d(a1, a2))
     assert_array_equal(union_answer, np.union1d(a1, a2))
 
-    assert_array_equal(
-        unyt_array(cross_answer, "cm**2"), ucross(a1, a1 + (2 * a1.units))
-    )
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(
+            unyt_array(cross_answer, "cm**2"), ucross(a1, a1 + (2 * a1.units))
+        )
     assert_array_equal(cross_answer, np.cross(a1.v, a1.v + 2))
 
-    assert_array_equal(unorm(a1), unyt_quantity(norm_answer, "cm"))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unorm(a1), unyt_quantity(norm_answer, "cm"))
     assert_array_equal(np.linalg.norm(a1), norm_answer)
-    assert_array_equal(unorm(a3, axis=1), unyt_array(arr_norm_answer, "cm"))
+
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unorm(a3, axis=1), unyt_array(arr_norm_answer, "cm"))
     assert_array_equal(np.linalg.norm(a3, axis=1), arr_norm_answer)
 
     with pytest.warns(DeprecationWarning):
         assert_array_equal(udot(a1, a1), unyt_quantity(dot_answer, "cm**2"))
 
-    assert_array_equal(np.array(catenate_answer), uconcatenate((a1.v, a2.v)))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(np.array(catenate_answer), uconcatenate((a1.v, a2.v)))
     with pytest.raises(RuntimeError):
-        uconcatenate((a1, a2.v))
+        with pytest.warns(DeprecationWarning):
+            uconcatenate((a1, a2.v))
     with pytest.raises(RuntimeError):
-        uconcatenate((a1.to("m"), a2))
-    assert_array_equal(unyt_array(vstack_answer, "cm"), uvstack([a2, a4]))
+        with pytest.warns(DeprecationWarning):
+            uconcatenate((a1.to("m"), a2))
+
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unyt_array(vstack_answer, "cm"), uvstack([a2, a4]))
     assert_array_equal(vstack_answer, np.vstack([a2, a4]))
-    assert_array_equal(unyt_array(vstack_answer, "cm"), ustack([a2, a4]))
+    with pytest.warns(DeprecationWarning):
+        assert_array_equal(unyt_array(vstack_answer, "cm"), ustack([a2, a4]))
     assert_array_equal(vstack_answer, np.stack([a2, a4]))
 
-    assert_array_equal(
-        unyt_array(vstack_answer_last_axis, "cm"), ustack([a2, a4], axis=-1)
-    )
+    with pytest.warns(DeprecationWarning):
+
+        assert_array_equal(
+            unyt_array(vstack_answer_last_axis, "cm"), ustack([a2, a4], axis=-1)
+        )
     assert_array_equal(vstack_answer_last_axis, np.stack([a2, a4], axis=-1))