Part of #3708: array_api to call functions from arkouda.pdarray_creation

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,33 +313,18 @@ 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:
     """
     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.
     """

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:

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))

arkouda/pdarraycreation.py

@@ -484,6 +484,9 @@ def zeros(
     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)

arkouda/pdarraysetops.py

@@ -450,7 +450,7 @@ def multiarray_setop_validation(
 def union1d(
     pda1: groupable,
     pda2: groupable,
-) -> Union[pdarray, groupable]:
+) -> groupable:
     """
     Find the union of two arrays/List of Arrays.
 
@@ -532,7 +532,7 @@ def union1d(
         c = [concatenate(x, ordered=False) for x in zip(ua, ub)]
         g = GroupBy(c)
         k, ct = g.size()
-        return k
+        return list(k)
     else:
         raise TypeError(
             f"Both pda1 and pda2 must be pdarray, List, or Tuple. Received {type(pda1)} and {type(pda2)}"

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

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)