From f4a657c3c340d0c957a35ea7c1c559723a6468f9 Mon Sep 17 00:00:00 2001 From: Amanda Potts Date: Fri, 6 Sep 2024 15:35:34 -0400 Subject: [PATCH] Part of #3708: array_api to call functions from arkouda.pdarray_creation --- arkouda/array_api/creation_functions.py | 79 +++++----------------- arkouda/array_api/elementwise_functions.py | 24 ++----- arkouda/array_api/linalg.py | 65 +++--------------- arkouda/pdarraycreation.py | 8 ++- pytest.ini | 2 + tests/array_api/array_creation.py | 58 ++++++++++++++++ tests/array_api/elementwise_functions.py | 18 +++++ tests/array_api/linalg.py | 65 ++++++++++++++++++ 8 files changed, 179 insertions(+), 140 deletions(-) create mode 100644 tests/array_api/elementwise_functions.py create mode 100644 tests/array_api/linalg.py diff --git a/arkouda/array_api/creation_functions.py b/arkouda/array_api/creation_functions.py index 27063716cd..f4f87ce0c1 100644 --- a/arkouda/array_api/creation_functions.py +++ b/arkouda/array_api/creation_functions.py @@ -2,12 +2,11 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union, cast -from arkouda.client import generic_msg import numpy as np -from arkouda.pdarrayclass import create_pdarray, pdarray, _to_pdarray -from arkouda.pdarraycreation import scalar_array + from arkouda.numpy.dtypes import dtype as akdtype from arkouda.numpy.dtypes import resolve_scalar_dtype +from arkouda.pdarrayclass import _to_pdarray, pdarray if TYPE_CHECKING: from ._typing import ( @@ -17,6 +16,7 @@ NestedSequence, SupportsBufferProtocol, ) + import arkouda as ak @@ -83,9 +83,7 @@ def asarray( elif isinstance(obj, np.ndarray): return Array._new(_to_pdarray(obj, dt=dtype)) else: - raise ValueError( - "asarray not implemented for 'NestedSequence' or 'SupportsBufferProtocol'" - ) + raise ValueError("asarray not implemented for 'NestedSequence' or 'SupportsBufferProtocol'") def arange( @@ -155,9 +153,7 @@ def empty( ) -def empty_like( - x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> Array: +def empty_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array: """ Return a new array whose shape and dtype match the input array, without initializing entries. """ @@ -217,17 +213,8 @@ def eye( if n_cols is not None: cols = n_cols - repMsg = generic_msg( - cmd="eye", - args={ - "dtype": np.dtype(dtype).name, - "rows": n_rows, - "cols": cols, - "diag": k, - }, - ) - - return Array._new(create_pdarray(repMsg)) + from arkouda import dtype as akdtype + return Array._new(ak.eye(rows=n_rows, cols=cols, diag=k, dt=akdtype(dtype))) def from_dlpack(x: object, /) -> Array: @@ -312,9 +299,7 @@ def ones( return a -def ones_like( - x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> Array: +def ones_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array: """ Return a new array whose shape and dtype match the input array, filled with ones. """ @@ -328,15 +313,7 @@ def tril(x: Array, /, *, k: int = 0) -> Array: """ from .array_object import Array - repMsg = generic_msg( - cmd=f"tril{x._array.ndim}D", - args={ - "array": x._array.name, - "diag": k, - }, - ) - - return Array._new(create_pdarray(repMsg)) + return Array._new(ak.tril(x._array, diag=k)) def triu(x: Array, /, *, k: int = 0) -> Array: @@ -344,17 +321,10 @@ def triu(x: Array, /, *, k: int = 0) -> Array: Create a new array with the values from `x` above the `k`-th diagonal, and all other elements zero. """ - from .array_object import Array - repMsg = generic_msg( - cmd=f"triu{x._array.ndim}D", - args={ - "array": x._array.name, - "diag": k, - }, - ) + from .array_object import Array - return Array._new(create_pdarray(repMsg)) + return Array._new(ak.triu(x._array, k)) def zeros( @@ -372,31 +342,14 @@ def zeros( if device not in ["cpu", None]: raise ValueError(f"Unsupported device {device!r}") - if isinstance(shape, tuple): - if shape == (): - return Array._new(scalar_array(0, dtype=dtype)) - else: - ndim = len(shape) - else: - if shape == 0: - return Array._new(scalar_array(0, dtype=dtype)) - else: - ndim = 1 - - dtype = akdtype(dtype) # normalize dtype - dtype_name = cast(np.dtype, dtype).name + return_dtype = akdtype(dtype) + if dtype is None: + return_dtype = akdtype(ak.float64) - repMsg = generic_msg( - cmd=f"create<{dtype_name},{ndim}>", - args={"shape": shape}, - ) + return Array._new(ak.zeros(shape, return_dtype)) - return Array._new(create_pdarray(repMsg)) - -def zeros_like( - x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> Array: +def zeros_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array: """ Return a new array whose shape and dtype match the input array, filled with zeros. """ diff --git a/arkouda/array_api/elementwise_functions.py b/arkouda/array_api/elementwise_functions.py index 0f269c2899..9793d7fe0d 100644 --- a/arkouda/array_api/elementwise_functions.py +++ b/arkouda/array_api/elementwise_functions.py @@ -106,10 +106,7 @@ def bitwise_and(x1: Array, x2: Array, /) -> Array: """ Compute the element-wise bitwise AND of two arrays. """ - if ( - x1.dtype not in _integer_or_boolean_dtypes - or x2.dtype not in _integer_or_boolean_dtypes - ): + if x1.dtype not in _integer_or_boolean_dtypes or x2.dtype not in _integer_or_boolean_dtypes: raise TypeError("Only integer or boolean dtypes are allowed in bitwise_and") # Call result type here just to raise on disallowed type combinations _result_type(x1.dtype, x2.dtype) @@ -141,10 +138,7 @@ def bitwise_or(x1: Array, x2: Array, /) -> Array: """ Compute the element-wise bitwise OR of two arrays. """ - if ( - x1.dtype not in _integer_or_boolean_dtypes - or x2.dtype not in _integer_or_boolean_dtypes - ): + if x1.dtype not in _integer_or_boolean_dtypes or x2.dtype not in _integer_or_boolean_dtypes: raise TypeError("Only integer or boolean dtypes are allowed in bitwise_or") # Call result type here just to raise on disallowed type combinations _result_type(x1.dtype, x2.dtype) @@ -169,10 +163,7 @@ def bitwise_xor(x1: Array, x2: Array, /) -> Array: """ Compute the element-wise bitwise XOR of two arrays. """ - if ( - x1.dtype not in _integer_or_boolean_dtypes - or x2.dtype not in _integer_or_boolean_dtypes - ): + if x1.dtype not in _integer_or_boolean_dtypes or x2.dtype not in _integer_or_boolean_dtypes: raise TypeError("Only integer or boolean dtypes are allowed in bitwise_xor") # Call result type here just to raise on disallowed type combinations _result_type(x1.dtype, x2.dtype) @@ -410,14 +401,7 @@ def logical_not(x: Array, /) -> Array: """ Compute the element-wise logical NOT of a boolean array. """ - repMsg = ak.generic_msg( - cmd=f"efunc{x._array.ndim}D", - args={ - "func": "not", - "array": x._array, - }, - ) - return Array._new(ak.create_pdarray(repMsg)) + return ~x def logical_or(x1: Array, x2: Array, /) -> Array: diff --git a/arkouda/array_api/linalg.py b/arkouda/array_api/linalg.py index e297d8ce6f..6b9c8bdddc 100644 --- a/arkouda/array_api/linalg.py +++ b/arkouda/array_api/linalg.py @@ -1,36 +1,15 @@ from .array_object import Array -from arkouda.client import generic_msg -from arkouda.pdarrayclass import create_pdarray, broadcast_if_needed - def matmul(x1: Array, x2: Array, /) -> Array: """ Matrix product of two arrays. """ - from .array_object import Array - - if x1._array.ndim < 2 and x2._array.ndim < 2: - raise ValueError( - "matmul requires at least one array argument to have more than two dimensions" - ) - - x1b, x2b, tmp_x1, tmp_x2 = broadcast_if_needed(x1._array, x2._array) + from arkouda import matmul as ak_matmul - repMsg = generic_msg( - cmd=f"matMul{len(x1b.shape)}D", - args={ - "x1": x1b.name, - "x2": x2b.name, - }, - ) - - if tmp_x1: - del x1b - if tmp_x2: - del x2b + from .array_object import Array - return Array._new(create_pdarray(repMsg)) + return Array._new(ak_matmul(x1._array, x2._array)) def tensordot(): @@ -44,42 +23,16 @@ def matrix_transpose(x: Array) -> Array: """ Matrix product of two arrays. """ - from .array_object import Array + from arkouda import transpose as ak_transpose - if x._array.ndim < 2: - raise ValueError( - "matrix_transpose requires the array to have more than two dimensions" - ) - - repMsg = generic_msg( - cmd=f"transpose{x._array.ndim}D", - args={ - "array": x._array.name, - }, - ) - - return Array._new(create_pdarray(repMsg)) - - -def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: from .array_object import Array - x1b, x2b, tmp_x1, tmp_x2 = broadcast_if_needed(x1._array, x2._array) + return Array._new(ak_transpose(x._array)) - repMsg = generic_msg( - cmd=f"vecdot{len(x1b.shape)}D", - args={ - "x1": x1b.name, - "x2": x2b.name, - "bcShape": x1b.shape, - "axis": axis, - }, - ) - if tmp_x1: - del x1b +def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: + from arkouda import vecdot as ak_vecdot - if tmp_x2: - del x2b + from .array_object import Array - return Array._new(create_pdarray(repMsg)) + return Array._new(ak_vecdot(x1._array, x2._array)) diff --git a/arkouda/pdarraycreation.py b/arkouda/pdarraycreation.py index 81e063f867..e747a4c6e3 100644 --- a/arkouda/pdarraycreation.py +++ b/arkouda/pdarraycreation.py @@ -27,7 +27,6 @@ from arkouda.pdarrayclass import create_pdarray, pdarray from arkouda.strings import Strings - __all__ = [ "array", "zeros", @@ -284,6 +283,7 @@ def array( raise RuntimeError(f"Unhandled dtype {a.dtype}") else: from arkouda.util import _infer_shape_from_size + shape, ndim, full_size = _infer_shape_from_size(a.shape) # Do not allow arrays that are too large @@ -478,11 +478,15 @@ def zeros( if dtype_name not in NumericDTypes: raise TypeError(f"unsupported dtype {dtype}") from arkouda.util import _infer_shape_from_size + shape, ndim, full_size = _infer_shape_from_size(size) if ndim > get_max_array_rank(): raise ValueError(f"array rank {ndim} exceeds maximum of {get_max_array_rank()}") + if shape == (): + return scalar_array(0, dtype=dtype) + repMsg = generic_msg(cmd=f"create<{dtype_name},{ndim}>", args={"shape": shape}) return create_pdarray(repMsg, max_bits=max_bits) @@ -538,6 +542,7 @@ def ones( if dtype_name not in NumericDTypes: raise TypeError(f"unsupported dtype {dtype}") from arkouda.util import _infer_shape_from_size + shape, ndim, full_size = _infer_shape_from_size(size) if ndim > get_max_array_rank(): @@ -607,6 +612,7 @@ def full( if dtype_name not in NumericDTypes: raise TypeError(f"unsupported dtype {dtype}") from arkouda.util import _infer_shape_from_size + shape, ndim, full_size = _infer_shape_from_size(size) if ndim > get_max_array_rank(): diff --git a/pytest.ini b/pytest.ini index 551bf33042..2890f92272 100644 --- a/pytest.ini +++ b/pytest.ini @@ -5,7 +5,9 @@ testpaths = tests/array_api/array_creation.py tests/array_api/array_manipulation.py tests/array_api/binary_ops.py + tests/array_api/elementwise_functions.py tests/array_api/indexing.py + tests/array_api/linalg.py tests/array_api/searching_functions.py tests/array_api/set_functions.py tests/array_api/sorting.py diff --git a/tests/array_api/array_creation.py b/tests/array_api/array_creation.py index 047a412f4a..f12f9094d1 100644 --- a/tests/array_api/array_creation.py +++ b/tests/array_api/array_creation.py @@ -1,8 +1,11 @@ +from math import sqrt + import numpy as np import pytest import arkouda as ak import arkouda.array_api as xp +from arkouda.testing import assert_almost_equivalent # requires the server to be built with 2D array support SHAPES = [(), (0,), (0, 0), (1,), (5,), (2, 2), (5, 10)] @@ -44,3 +47,58 @@ def test_from_numpy(self): assert b.ndim == a.ndim assert b.shape == a.shape assert b.tolist() == a.tolist() + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_triu(self, data_type, prob_size): + from arkouda.array_api.creation_functions import triu as array_triu + + size = int(sqrt(prob_size)) + + # test on one square and two non-square matrices + + for rows, cols in [(size, size), (size + 1, size - 1), (size - 1, size + 1)]: + pda = xp.asarray(ak.randint(1, 10, (rows, cols))) + nda = pda.to_ndarray() + sweep = range(-(rows - 1), cols - 1) # sweeps the diagonal from LL to UR + for diag in sweep: + np_triu = np.triu(nda, diag) + ak_triu = array_triu(pda, k=diag)._array + assert_almost_equivalent(ak_triu, np_triu) + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_tril(self, data_type, prob_size): + from arkouda.array_api.creation_functions import tril as array_tril + + size = int(sqrt(prob_size)) + + # test on one square and two non-square matrices + + for rows, cols in [(size, size), (size + 1, size - 1), (size - 1, size + 1)]: + pda = xp.asarray(ak.randint(1, 10, (rows, cols))) + nda = pda.to_ndarray() + sweep = range(-(rows - 2), cols) # sweeps the diagonal from LL to UR + for diag in sweep: + np_tril = np.tril(nda, diag) + ak_tril = array_tril(pda, k=diag)._array + assert_almost_equivalent(np_tril, ak_tril) + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_eye(self, data_type, prob_size): + from arkouda.array_api.creation_functions import eye as array_eye + + size = int(sqrt(prob_size)) + + # test on one square and two non-square matrices + + for rows, cols in [(size, size), (size + 1, size - 1), (size - 1, size + 1)]: + sweep = range(-(cols - 1), rows) # sweeps the diagonal from LL to UR + for diag in sweep: + np_eye = np.eye(rows, cols, diag, dtype=data_type) + ak_eye = array_eye(rows, cols, k=diag, dtype=data_type)._array + assert_almost_equivalent(np_eye, ak_eye) diff --git a/tests/array_api/elementwise_functions.py b/tests/array_api/elementwise_functions.py new file mode 100644 index 0000000000..58fea13dd5 --- /dev/null +++ b/tests/array_api/elementwise_functions.py @@ -0,0 +1,18 @@ +import pytest + +import arkouda as ak +import arkouda.array_api as xp +from arkouda.testing import assert_equal + +# requires the server to be built with 2D array support +SHAPES = [(), (0,), (0, 0), (1,), (5,), (2, 2), (5, 10)] +SIZES = [1, 0, 0, 1, 5, 4, 50] +DIMS = [0, 1, 2, 1, 1, 2, 2] + + +class TestElemenwiseFunctions: + @pytest.mark.skip_if_max_rank_less_than(2) + def test_logical_not(self): + a = xp.asarray(ak.array([True, False, True, False])) + not_a = xp.asarray(ak.array([False, True, False, True])) + assert_equal(xp.logical_not(a)._array, not_a._array) diff --git a/tests/array_api/linalg.py b/tests/array_api/linalg.py new file mode 100644 index 0000000000..71d1fea024 --- /dev/null +++ b/tests/array_api/linalg.py @@ -0,0 +1,65 @@ +import pytest +from math import sqrt +import arkouda as ak +import arkouda.array_api as xp +from arkouda.testing import assert_almost_equivalent +import numpy as np + + +class TestLinalg: + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type1", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("data_type2", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_matmul(self, data_type1, data_type2, prob_size): + + size = int(sqrt(prob_size)) + + # test on one square and two non-square products + + for rows, cols in [(size, size), (size + 1, size - 1), (size - 1, size + 1)]: + arrayLeft = xp.asarray(ak.randint(0, 10, (rows, size), dtype=data_type1)) + ndaLeft = arrayLeft.to_ndarray() + arrayRight = xp.asarray(ak.randint(0, 10, (size, cols), dtype=data_type2)) + ndaRight = arrayRight.to_ndarray() + akProduct = xp.matmul(arrayLeft, arrayRight) + npProduct = np.matmul(ndaLeft, ndaRight) + assert_almost_equivalent(akProduct._array, npProduct) + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_transpose(self, data_type, prob_size): + + size = int(sqrt(prob_size)) + + # test on one square and two non-square matrices + + for rows, cols in [(size, size), (size + 1, size - 1), (size - 1, size + 1)]: + array = xp.asarray(ak.randint(1, 10, (rows, cols))) + nda = array.to_ndarray() + npa = np.transpose(nda) + ppa = xp.matrix_transpose(array)._array + assert np.allclose(ppa.to_ndarray(), npa) + + @pytest.mark.skip_if_max_rank_less_than(2) + @pytest.mark.parametrize("data_type1", [ak.int64, ak.float64, ak.bool_]) + @pytest.mark.parametrize("data_type2", [ak.int64, ak.float64]) + @pytest.mark.parametrize("prob_size", pytest.prob_size) + def test_vecdot(self, data_type1, data_type2, prob_size): + + depth = np.random.randint(2, 10) + width = prob_size // depth + + pda_a = xp.asarray(ak.randint(0, 10, (depth, width), dtype=data_type1)) + nda_a = pda_a.to_ndarray() + pda_b = xp.asarray(ak.randint(0, 10, (depth, width), dtype=data_type2)) + nda_b = pda_b.to_ndarray() + akProduct = xp.vecdot(pda_a, pda_b) + + # there is no vecdot in numpy (and vdot doesn't do the same thing). + # np.add.reduce does. + + npProduct = np.add.reduce(nda_a * nda_b) + assert_almost_equivalent(npProduct, akProduct._array)