From cd9f69a727346765a70dbdff22d316a9d37e7c73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= Date: Tue, 14 Nov 2023 12:21:32 +0100 Subject: [PATCH] API: Add cross to numpy.linalg --- .../upcoming_changes/25145.new_feature.rst | 6 ++ doc/source/reference/array_api.rst | 1 - doc/source/reference/routines.linalg.rst | 1 + numpy/_core/numeric.py | 2 + numpy/_core/numeric.pyi | 28 ++++----- numpy/linalg/__init__.py | 1 + numpy/linalg/__init__.pyi | 1 + numpy/linalg/_linalg.py | 61 +++++++++++++++++-- numpy/linalg/_linalg.pyi | 28 +++++++++ numpy/linalg/tests/test_linalg.py | 20 ++++++ numpy/typing/tests/data/reveal/linalg.pyi | 4 ++ tools/ci/array-api-skips.txt | 3 - 12 files changed, 134 insertions(+), 22 deletions(-) create mode 100644 doc/release/upcoming_changes/25145.new_feature.rst diff --git a/doc/release/upcoming_changes/25145.new_feature.rst b/doc/release/upcoming_changes/25145.new_feature.rst new file mode 100644 index 000000000000..6498fb7ec435 --- /dev/null +++ b/doc/release/upcoming_changes/25145.new_feature.rst @@ -0,0 +1,6 @@ +``cross`` for `numpy.linalg` +---------------------------- + +`numpy.linalg.cross` has been added. It computes the cross product of two +(arrays of) 3-dimensional vectors. It differs from `numpy.cross` by accepting +three-dimensional vectors only. This function is compatible with Array API. diff --git a/doc/source/reference/array_api.rst b/doc/source/reference/array_api.rst index 09615c2e793b..ca362947afee 100644 --- a/doc/source/reference/array_api.rst +++ b/doc/source/reference/array_api.rst @@ -138,7 +138,6 @@ Function instead of method These functions are in the ``linalg`` sub-namespace in the array API, but are only in the top-level namespace in NumPy: -- ``cross`` - ``diagonal`` - ``matmul`` (*) - ``outer`` diff --git a/doc/source/reference/routines.linalg.rst b/doc/source/reference/routines.linalg.rst index 02a001f6d203..715056f80eb6 100644 --- a/doc/source/reference/routines.linalg.rst +++ b/doc/source/reference/routines.linalg.rst @@ -64,6 +64,7 @@ Matrix and vector products einsum_path linalg.matrix_power kron + linalg.cross Decompositions -------------- diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index 96ba76df3c18..85f2ace81514 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -1510,6 +1510,8 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): -------- inner : Inner product outer : Outer product. + linalg.cross : An Array API compatible variation of ``np.cross``, + which accepts (arrays of) 3-element vectors only. ix_ : Construct index arrays. Notes diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi index 72b595acbbfb..26163c945c4f 100644 --- a/numpy/_core/numeric.pyi +++ b/numpy/_core/numeric.pyi @@ -488,8 +488,8 @@ def moveaxis( @overload def cross( - a: _ArrayLikeUnknown, - b: _ArrayLikeUnknown, + x1: _ArrayLikeUnknown, + x2: _ArrayLikeUnknown, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -497,8 +497,8 @@ def cross( ) -> NDArray[Any]: ... @overload def cross( - a: _ArrayLikeBool_co, - b: _ArrayLikeBool_co, + x1: _ArrayLikeBool_co, + x2: _ArrayLikeBool_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -506,8 +506,8 @@ def cross( ) -> NoReturn: ... @overload def cross( - a: _ArrayLikeUInt_co, - b: _ArrayLikeUInt_co, + x1: _ArrayLikeUInt_co, + x2: _ArrayLikeUInt_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -515,8 +515,8 @@ def cross( ) -> NDArray[unsignedinteger[Any]]: ... @overload def cross( - a: _ArrayLikeInt_co, - b: _ArrayLikeInt_co, + x1: _ArrayLikeInt_co, + x2: _ArrayLikeInt_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -524,8 +524,8 @@ def cross( ) -> NDArray[signedinteger[Any]]: ... @overload def cross( - a: _ArrayLikeFloat_co, - b: _ArrayLikeFloat_co, + x1: _ArrayLikeFloat_co, + x2: _ArrayLikeFloat_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -533,8 +533,8 @@ def cross( ) -> NDArray[floating[Any]]: ... @overload def cross( - a: _ArrayLikeComplex_co, - b: _ArrayLikeComplex_co, + x1: _ArrayLikeComplex_co, + x2: _ArrayLikeComplex_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -542,8 +542,8 @@ def cross( ) -> NDArray[complexfloating[Any, Any]]: ... @overload def cross( - a: _ArrayLikeObject_co, - b: _ArrayLikeObject_co, + x1: _ArrayLikeObject_co, + x2: _ArrayLikeObject_co, axisa: int = ..., axisb: int = ..., axisc: int = ..., diff --git a/numpy/linalg/__init__.py b/numpy/linalg/__init__.py index 8c48ce60cd30..e7c27b36b431 100644 --- a/numpy/linalg/__init__.py +++ b/numpy/linalg/__init__.py @@ -26,6 +26,7 @@ Matrix and vector products -------------------------- + cross multi_dot matrix_power diff --git a/numpy/linalg/__init__.pyi b/numpy/linalg/__init__.pyi index 3075aced376b..2f889b423dec 100644 --- a/numpy/linalg/__init__.pyi +++ b/numpy/linalg/__init__.pyi @@ -21,6 +21,7 @@ from numpy.linalg._linalg import ( multi_dot as multi_dot, trace as trace, diagonal as diagonal, + cross as cross, ) from numpy._pytesttester import PytestTester diff --git a/numpy/linalg/_linalg.py b/numpy/linalg/_linalg.py index 48a801ac6584..7ae99112b8c0 100644 --- a/numpy/linalg/_linalg.py +++ b/numpy/linalg/_linalg.py @@ -12,7 +12,7 @@ __all__ = ['matrix_power', 'solve', 'tensorsolve', 'tensorinv', 'inv', 'cholesky', 'eigvals', 'eigvalsh', 'pinv', 'slogdet', 'det', 'svd', 'eig', 'eigh', 'lstsq', 'norm', 'qr', 'cond', 'matrix_rank', - 'LinAlgError', 'multi_dot', 'trace', 'diagonal'] + 'LinAlgError', 'multi_dot', 'trace', 'diagonal', 'cross'] import functools import operator @@ -26,7 +26,8 @@ add, multiply, sqrt, sum, isfinite, finfo, errstate, moveaxis, amin, amax, prod, abs, atleast_2d, intp, asanyarray, object_, matmul, swapaxes, divide, count_nonzero, isnan, sign, argsort, sort, - reciprocal, overrides, diagonal as _core_diagonal, trace as _core_trace + reciprocal, overrides, diagonal as _core_diagonal, trace as _core_trace, + cross as _core_cross, ) from numpy.lib._twodim_base_impl import triu, eye from numpy.lib.array_utils import normalize_axis_index @@ -2937,14 +2938,14 @@ def diagonal(x, /, *, offset=0): See Also -------- numpy.diagonal + """ return _core_diagonal(x, offset, axis1=-2, axis2=-1) # trace -def _trace_dispatcher( - x, /, *, offset=None, dtype=None): +def _trace_dispatcher(x, /, *, offset=None, dtype=None): return (x,) @@ -2990,5 +2991,57 @@ def trace(x, /, *, offset=0, dtype=None): See Also -------- numpy.trace + """ return _core_trace(x, offset, axis1=-2, axis2=-1, dtype=dtype) + + +# cross + +def _cross_dispatcher(x1, x2, /, *, axis=None): + return (x1, x2,) + + +@array_function_dispatch(_cross_dispatcher) +def cross(x1, x2, /, *, axis=-1): + """ + Returns the cross product of 3-element vectors. + + If ``x1`` and/or ``x2`` are multi-dimensional arrays, then + the cross-product of each pair of corresponding 3-element vectors + is independently computed. + + This function is Array API compatible, contrary to + :func:`numpy.cross`. + + Parameters + ---------- + x1 : array_like + The first input array. + x2 : array_like + The second input array. Must be compatible with ``x1`` for all + non-compute axes. The size of the axis over which to compute + the cross-product must be the same size as the respective axis + in ``x1``. + axis : int, optional + The axis (dimension) of ``x1`` and ``x2`` containing the vectors for + which to compute the cross-product. Default: ``-1``. + + Returns + ------- + out : ndarray + An array containing the cross products. + + See Also + -------- + numpy.cross + + """ + if x1.shape[axis] != 3 or x2.shape[axis] != 3: + raise ValueError( + "Both input arrays must be (arrays of) 3-dimensional vectors, " + f"but they are {x1.shape[axis]} and {x2.shape[axis]} " + "dimensional instead." + ) + + return _core_cross(x1, x2, axis=axis) diff --git a/numpy/linalg/_linalg.pyi b/numpy/linalg/_linalg.pyi index a00f250aacce..14410ade24c7 100644 --- a/numpy/linalg/_linalg.pyi +++ b/numpy/linalg/_linalg.pyi @@ -14,6 +14,8 @@ from numpy import ( generic, floating, complexfloating, + signedinteger, + unsignedinteger, int32, float64, complex128, @@ -25,6 +27,7 @@ from numpy._typing import ( NDArray, ArrayLike, _ArrayLikeInt_co, + _ArrayLikeUInt_co, _ArrayLikeFloat_co, _ArrayLikeComplex_co, _ArrayLikeTD64_co, @@ -307,3 +310,28 @@ def trace( offset: SupportsIndex = ..., dtype: DTypeLike = ..., ) -> Any: ... + +@overload +def cross( + a: _ArrayLikeUInt_co, + b: _ArrayLikeUInt_co, + axis: int = ..., +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def cross( + a: _ArrayLikeInt_co, + b: _ArrayLikeInt_co, + axis: int = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def cross( + a: _ArrayLikeFloat_co, + b: _ArrayLikeFloat_co, + axis: int = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cross( + a: _ArrayLikeComplex_co, + b: _ArrayLikeComplex_co, + axis: int = ..., +) -> NDArray[complexfloating[Any, Any]]: ... diff --git a/numpy/linalg/tests/test_linalg.py b/numpy/linalg/tests/test_linalg.py index d174643d1ecb..c749714c7e2d 100644 --- a/numpy/linalg/tests/test_linalg.py +++ b/numpy/linalg/tests/test_linalg.py @@ -2228,3 +2228,23 @@ def test_trace(): expected = np.array([36, 116, 196]) assert_equal(actual, expected) + + +def test_cross(): + + x = np.arange(9).reshape((3, 3)) + actual = np.linalg.cross(x, x + 1) + expected = np.array([ + [-1, 2, -1], + [-1, 2, -1], + [-1, 2, -1], + ]) + + assert_equal(actual, expected) + + with assert_raises_regex( + ValueError, + r"input arrays must be \(arrays of\) 3-dimensional vectors" + ): + x_2dim = x[:, 1:] + np.linalg.cross(x_2dim, x_2dim) diff --git a/numpy/typing/tests/data/reveal/linalg.pyi b/numpy/typing/tests/data/reveal/linalg.pyi index a3c08b41c6d3..d8264d9311a5 100644 --- a/numpy/typing/tests/data/reveal/linalg.pyi +++ b/numpy/typing/tests/data/reveal/linalg.pyi @@ -106,3 +106,7 @@ assert_type(np.linalg.multi_dot([AR_i8, AR_f8]), Any) assert_type(np.linalg.multi_dot([AR_f8, AR_c16]), Any) assert_type(np.linalg.multi_dot([AR_O, AR_O]), Any) assert_type(np.linalg.multi_dot([AR_m, AR_m]), Any) + +assert_type(np.linalg.cross(AR_i8, AR_i8), npt.NDArray[np.signedinteger[Any]]) +assert_type(np.linalg.cross(AR_f8, AR_f8), npt.NDArray[np.floating[Any]]) +assert_type(np.linalg.cross(AR_c16, AR_c16), npt.NDArray[np.complexfloating[Any, Any]]) diff --git a/tools/ci/array-api-skips.txt b/tools/ci/array-api-skips.txt index 0ee346f472dc..e2446dec2e17 100644 --- a/tools/ci/array-api-skips.txt +++ b/tools/ci/array-api-skips.txt @@ -42,7 +42,6 @@ array_api_tests/test_data_type_functions.py::test_isdtype array_api_tests/test_data_type_functions.py::test_astype # missing names -array_api_tests/test_has_names.py::test_has_names[linalg-cross] array_api_tests/test_has_names.py::test_has_names[linalg-matmul] array_api_tests/test_has_names.py::test_has_names[linalg-matrix_norm] array_api_tests/test_has_names.py::test_has_names[linalg-matrix_transpose] @@ -76,7 +75,6 @@ array_api_tests/test_has_names.py::test_has_names[array_method-to_device] array_api_tests/test_has_names.py::test_has_names[array_attribute-device] # missing linalg names -array_api_tests/test_linalg.py::test_cross array_api_tests/test_linalg.py::test_matrix_norm array_api_tests/test_linalg.py::test_matrix_transpose array_api_tests/test_linalg.py::test_outer @@ -125,7 +123,6 @@ array_api_tests/test_signatures.py::test_func_signature[bitwise_right_shift] array_api_tests/test_signatures.py::test_func_signature[pow] array_api_tests/test_signatures.py::test_func_signature[matrix_transpose] array_api_tests/test_signatures.py::test_func_signature[vecdot] -array_api_tests/test_signatures.py::test_extension_func_signature[linalg.cross] array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matmul] array_api_tests/test_signatures.py::test_extension_func_signature[linalg.cholesky] array_api_tests/test_signatures.py::test_extension_func_signature[linalg.matrix_norm]