From 2cd616222b9212069dd1bcfa4f4ad5dafb00b2f2 Mon Sep 17 00:00:00 2001 From: jeremiah-corrado <62707311+jeremiah-corrado@users.noreply.github.com> Date: Wed, 10 Apr 2024 10:33:55 -0600 Subject: [PATCH] Array API Manipulation Function Improvements (#3056) * refactor some array manipulation commands to reduce uneeded comm with locale 0 Signed-off-by: Jeremiah Corrado * add implementation of tile, unstack, and repeat to maniputation-functions module. Implement rank-reducing indexing for ND arrays. Bug fixes and performance improvements in manipultation functions. Signed-off-by: Jeremiah Corrado * fix flake8, mypy and python compat errors Signed-off-by: Jeremiah Corrado * add error handling to expandDims and stack for cases where max array rank would be exceeded Signed-off-by: Jeremiah Corrado * add special cases for flatten and unflatten to reshapeMsg Signed-off-by: Jeremiah Corrado * replace rank reducing slices in flatten/unflatten with explicit put/get operations Signed-off-by: Jeremiah Corrado * use AryUtil flatten/unflatten in SetMsg Signed-off-by: Jeremiah Corrado --------- Signed-off-by: Jeremiah Corrado --- arkouda/array_api/__init__.py | 8 + arkouda/array_api/_manipulation_functions.py | 124 ++-- arkouda/pdarrayclass.py | 49 +- src/AryUtil.chpl | 148 +++++ src/ManipulationMsg.chpl | 626 +++++++++++++++---- src/ReductionMsg.chpl | 6 +- src/SetMsg.chpl | 65 -- src/compat/e-132/ArkoudaAryUtilCompat.chpl | 4 +- src/compat/eq-131/ArkoudaAryUtilCompat.chpl | 4 +- src/compat/eq-133/ArkoudaAryUtilCompat.chpl | 4 +- src/compat/eq-134/ArkoudaAryUtilCompat.chpl | 4 +- src/compat/ge-20/ArkoudaAryUtilCompat.chpl | 6 +- tests/array_api/array_manipulation.py | 118 +++- 13 files changed, 860 insertions(+), 306 deletions(-) diff --git a/arkouda/array_api/__init__.py b/arkouda/array_api/__init__.py index c353f6fda5..2eee5a4aea 100644 --- a/arkouda/array_api/__init__.py +++ b/arkouda/array_api/__init__.py @@ -118,11 +118,15 @@ concat, expand_dims, flip, + moveaxis, permute_dims, + repeat, reshape, roll, squeeze, stack, + tile, + unstack, ) from ._searching_functions import argmax, argmin, nonzero, where @@ -255,11 +259,15 @@ "concat", "expand_dims", "flip", + "moveaxis", "permute_dims", + "repeat", "reshape", "roll", "squeeze", "stack", + "tile", + "unstack", ] __all__ += ["argmax", "argmin", "nonzero", "where"] diff --git a/arkouda/array_api/_manipulation_functions.py b/arkouda/array_api/_manipulation_functions.py index dc5248e813..1880b41d19 100644 --- a/arkouda/array_api/_manipulation_functions.py +++ b/arkouda/array_api/_manipulation_functions.py @@ -5,18 +5,13 @@ from typing import List, Optional, Tuple, Union, cast from arkouda.client import generic_msg from arkouda.pdarrayclass import create_pdarray +from arkouda.pdarraycreation import scalar_array from arkouda.util import broadcast_dims import numpy as np def broadcast_arrays(*arrays: Array) -> List[Array]: - """ - Array API compatible wrapper for :py:func:`np.broadcast_arrays `. - - See its docstring for more information. - """ - shapes = [a.shape for a in arrays] bcShape = shapes[0] for shape in shapes[1:]: @@ -50,15 +45,9 @@ def broadcast_to(x: Array, /, shape: Tuple[int, ...]) -> Array: raise ValueError(f"Failed to broadcast array: {e}") -# Note: the function name is different here def concat( arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: Optional[int] = 0 ) -> Array: - """ - Array API compatible wrapper for :py:func:`np.concatenate `. - - See its docstring for more information. - """ # TODO: type promotion across input arrays return Array._new( @@ -81,11 +70,6 @@ def concat( def expand_dims(x: Array, /, *, axis: int) -> Array: - """ - Array API compatible wrapper for :py:func:`np.expand_dims `. - - See its docstring for more information. - """ try: return Array._new( create_pdarray( @@ -106,11 +90,6 @@ def expand_dims(x: Array, /, *, axis: int) -> Array: def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array: - """ - Array API compatible wrapper for :py:func:`np.flip `. - - See its docstring for more information. - """ axisList = [] if axis is not None: axisList = list(axis) if isinstance(axis, tuple) else [axis] @@ -137,15 +116,22 @@ def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> def moveaxis( x: Array, source: Union[int, Tuple[int, ...]], destination: Union[int, Tuple[int, ...]], / ) -> Array: - raise NotImplementedError("moveaxis is not yet implemented") + perm = list(range(x.ndim)) + if isinstance(source, tuple): + if isinstance(destination, tuple): + for s, d in zip(source, destination): + perm[s] = d + else: + raise ValueError("source and destination must both be tuples if source is a tuple") + elif isinstance(destination, int): + perm[source] = destination + else: + raise ValueError("source and destination must both be integers if source is a tuple") + return permute_dims(x, axes=tuple(perm)) -def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array: - """ - Array API compatible wrapper for :py:func:`np.transpose `. - See its docstring for more information. - """ +def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array: try: return Array._new( create_pdarray( @@ -166,17 +152,33 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array: def repeat(x: Array, repeats: Union[int, Array], /, *, axis: Optional[int] = None) -> Array: - raise NotImplementedError("repeat is not yet implemented") + if isinstance(repeats, int): + reps = Array._new(scalar_array(repeats)) + else: + reps = repeats + + if axis is None: + return Array._new( + create_pdarray( + cast( + str, + generic_msg( + cmd=f"repeatFlat{x.ndim}D", + args={ + "name": x._array, + "repeats": reps._array, + }, + ), + ) + ) + ) + else: + raise NotImplementedError("repeat with 'axis' argument is not yet implemented") def reshape( x: Array, /, shape: Tuple[int, ...], *, copy: Optional[bool] = None ) -> Array: - """ - Array API compatible wrapper for :py:func:`np.reshape `. - - See its docstring for more information. - """ # TODO: figure out copying semantics (currently always creates a copy) try: @@ -205,11 +207,6 @@ def roll( *, axis: Optional[Union[int, Tuple[int, ...]]] = None, ) -> Array: - """ - Array API compatible wrapper for :py:func:`np.roll `. - - See its docstring for more information. - """ axisList = [] if axis is not None: axisList = list(axis) if isinstance(axis, tuple) else [axis] @@ -240,11 +237,6 @@ def roll( def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array: - """ - Array API compatible wrapper for :py:func:`np.squeeze `. - - See its docstring for more information. - """ nAxes = len(axis) if isinstance(axis, tuple) else 1 try: return Array._new( @@ -267,11 +259,6 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array: def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array: - """ - Array API compatible wrapper for :py:func:`np.stack `. - - See its docstring for more information. - """ # TODO: type promotion across input arrays return Array._new( create_pdarray( @@ -291,8 +278,43 @@ def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array: - raise NotImplementedError("tile is not yet implemented") + if len(repetitions) > x.ndim: + xr = reshape(x, (1,) * (len(repetitions) - x.ndim) + x.shape) + reps = repetitions + elif len(repetitions) < x.ndim: + xr = x + reps = (1,) * (x.ndim - len(repetitions)) + repetitions + else: + xr = x + reps = repetitions + + return Array._new( + create_pdarray( + cast( + str, + generic_msg( + cmd=f"tile{xr.ndim}D", + args={ + "name": xr._array, + "reps": reps, + }, + ), + ) + ) + ) def unstack(x: Array, /, *, axis: int = 0) -> Tuple[Array, ...]: - raise NotImplementedError("unstack is not yet implemented") + resp = cast( + str, + generic_msg( + cmd=f"unstack{x.ndim}D", + args={ + "name": x._array, + "axis": axis, + "numReturnArrays": x.shape[axis], + }, + ), + ) + + return tuple([Array._new(create_pdarray(a)) for a in resp.split("+")]) diff --git a/arkouda/pdarrayclass.py b/arkouda/pdarrayclass.py index 72bee0a1d6..4ac5f51f34 100755 --- a/arkouda/pdarrayclass.py +++ b/arkouda/pdarrayclass.py @@ -657,18 +657,41 @@ def __getitem__(self, key): return create_pdarray(repMsg) if isinstance(key, tuple): - allScalar = True + if len(key) > self.ndim: + raise IndexError(f"too many indices ({len(key)}) for array with {self.ndim} dimensions") + + # replace '...' with the appropriate number of ':' + elipsis_axis_idx = -1 + for dim, k in enumerate(key): + if isinstance(k, type(Ellipsis)): + if elipsis_axis_idx != -1: + raise IndexError("array index can only have one ellipsis") + else: + elipsis_axis_idx = dim + + if elipsis_axis_idx != -1: + key = tuple( + key[:elipsis_axis_idx] + + (slice(None),) * (self.ndim - len(key) + 1) + + key[(elipsis_axis_idx+1):] + ) + + # parse the key tuple + num_scalar = 0 + scalar_axes = [] starts = [] stops = [] strides = [] for dim, k in enumerate(key): if isinstance(k, slice): - allScalar = False (start, stop, stride) = k.indices(self.shape[dim]) starts.append(start) stops.append(stop) strides.append(stride) elif np.isscalar(k) and (resolve_scalar_dtype(k) in ["int64", "uint64"]): + num_scalar += 1 + scalar_axes.append(dim) + if k < 0: # Interpret negative key as offset from end of array k += int(self.shape[dim]) @@ -678,15 +701,14 @@ def __getitem__(self, key): ) else: # treat this as a single element slice - # TODO: implement rank-reducing slices starts.append(k) stops.append(k + 1) strides.append(1) else: raise IndexError(f"Unhandled key type: {k} ({type(k)})") - if allScalar: - # use simpler indexing (and return a scalar) if we got a tuple of only scalars + if num_scalar == len(key): + # all scalars: use simpler indexing (and return a scalar) repMsg = generic_msg( cmd=f"[int]{self.ndim}D", args={ @@ -706,7 +728,22 @@ def __getitem__(self, key): "strides": tuple(strides), }, ) - return create_pdarray(repMsg) + maybe_degen_arr = create_pdarray(repMsg) + + if num_scalar > 0: + # reduce the array rank if there are any scalar indices + # note: squeeze requires the non-default ManipulationMsg server module + repMsg = generic_msg( + cmd=f"squeeze{maybe_degen_arr.ndim}Dx{maybe_degen_arr.ndim - num_scalar}D", + args={ + "name": maybe_degen_arr, + "nAxes": num_scalar, + "axes": scalar_axes, + }, + ) + return create_pdarray(repMsg) + else: + return maybe_degen_arr if isinstance(key, pdarray) and self.ndim == 1: kind, _ = translate_np_dtype(key.dtype) diff --git a/src/AryUtil.chpl b/src/AryUtil.chpl index db2eca7f26..7e10cb4fe6 100644 --- a/src/AryUtil.chpl +++ b/src/AryUtil.chpl @@ -10,6 +10,7 @@ module AryUtil use BitOps; use GenSymIO; use PrivateDist; + use Communication; use ArkoudaPOSIXCompat; use ArkoudaCTypesCompat; @@ -731,4 +732,151 @@ module AryUtil } return s; } + + /* + unflatten a 1D array into a multi-dimensional array of the given shape + */ + proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws { + var unflat = makeDistArray((...shape), t); + + if N == 1 { + unflat = a; + return unflat; + } + + // ranges of flat indices owned by each locale + const flatLocRanges = [loc in Locales] d.localSubdomain(loc).dim(0); + + coforall loc in Locales do on loc { + const lduf = unflat.domain.localSubdomain(), + lastRank = lduf.dim(N-1); + + // iterate over each slice of contiguous memory in the local subdomain + forall idx in domOffAxis(lduf, N-1) with ( + const ord = new orderer(shape), + const dufc = unflat.domain, + in flatLocRanges + ) { + var idxTup: (N-1)*int; + for i in 0..<(N-1) do idxTup[i] = idx[i]; + + const low = ((...idxTup), lastRank.low), + high = ((...idxTup), lastRank.high), + flatSlice = ord.indexToOrder(low)..ord.indexToOrder(high); + + // compute which locales in the input array this slice corresponds to + var locInStart, locInStop = 0; + for (flr, locID) in zip(flatLocRanges, 0.. 1 + { + var flat = makeDistArray(d.size, t); + + // ranges of flat indices owned by each locale + const flatLocRanges = [loc in Locales] flat.domain.localSubdomain(loc).dim(0); + + coforall loc in Locales do on loc { + const ld = d.localSubdomain(), + lastRank = ld.dim(d.rank-1); + + // iterate over each slice of contiguous memory in the local subdomain + forall idx in domOffAxis(ld, d.rank-1) with ( + const ord = new orderer(d.shape), + const dc = d, + in flatLocRanges + ) { + var idxTup: (d.rank-1)*int; + for i in 0..<(d.rank-1) do idxTup[i] = idx[i]; + + const low = ((...idxTup), lastRank.low), + high = ((...idxTup), lastRank.high), + flatSlice = ord.indexToOrder(low)..ord.indexToOrder(high); + + // compute which locales in the output array this slice corresponds to + var locOutStart, locOutStop = 0; + for (flr, locID) in zip(flatLocRanges, 0.. order for the input array's indices + // e.g., order = k + (nz * j) + (nz * ny * i) + inline proc indexToOrder(idx: rank*int): int { + var order = 0; + for param i in 0.. {0..0, 0..<4, 0..0} + // - for 'nonBCIndex' = (0, 0, 0), outSliceIdx = (0..<5, 0..0, 0..<3) + // - for 'nonBCIndex' = (0, 1, 0), outSliceIdx = (0..<5, 1..1, 0..<3) + // - etc. + // */ + // forall nonBCIndex in domOffAxis(eOut.a.domain, bcDimsList.toArray()) { + // const nbcT = if ndOut == 1 then (nonBCIndex,) else nonBCIndex; + + + // var outSliceIdx: ndOut*range; + // for i in 0.. 1 - { - var ret = idx; - ret[axis] += startOffsets[arrIdx]; - return ret; - } - - inline proc imap(arrIdx: int, idx: int): int - where nd == 1 - do return idx + startOffsets[arrIdx]; - // copy the data from the input arrays to the output array - for (arrIdx, arr) in zip(eIns.domain, eIns) do - forall idx in arr.a.domain with (var agg = newDstAggregator(t)) do - agg.copy(eOut.a[imap(arrIdx, idx)], arr.a[idx]); + forall (arrIdx, arr) in zip(eIns.domain, eIns) with (in startOffsets) { + forall idx in arr.a.domain with (var agg = newDstAggregator(t)) { + var outIdx = if nd == 1 then (idx,) else idx; + outIdx[axis] += startOffsets[arrIdx]; + agg.copy(eOut.a[outIdx], arr.a[idx]); + } + } const repMsg = "created " + st.attrib(rname); mLogger.info(getModuleName(),pn,getLineNumber(),repMsg); @@ -198,7 +212,6 @@ module ManipulationMsg { when DType.UInt64 do return doConcat(uint); when DType.Float64 do return doConcat(real); when DType.Bool do return doConcat(bool); - when DType.BigInt do return doConcat(bigint); otherwise { var errorMsg = notImplementedError(pn,dtype2str(dt)); mLogger.error(getModuleName(),pn,getLineNumber(),errorMsg); @@ -257,10 +270,9 @@ module ManipulationMsg { var eOut = st.addEntry(rname, + reduce sizes, t); // copy the data from the input arrays to the output array - for arrIdx in 0.. nd then return (false, ret); @@ -437,19 +466,14 @@ module ManipulationMsg { const eIn = toSymEntry(gEnt, t, nd); var eOut = st.addEntry(rname, (...eIn.tupShape), t); - // mapping between the input and output array indices with all axes flipped - inline proc imap(idx: nd*int): nd*int { - var ret: nd*int; - for param i in 0.. order for the output array's indices - // e.g., order = k + (nz * j) + (nz * ny * i) - inline proc indexToOrder(idx: ndOut*int): int { - var order = 0; - for param i in 0.. 1 + { + var shapeOut: (N-1)*int; + if numReturnArrays != shape[axis] { + return (false, shapeOut); + } else { + var i = 0; + for ii in 0..N { + if ii == axis { + continue; + } else { + shapeOut[i] = shape[ii]; + i += 1; + } + } + return (true, shapeOut); + } + } + + + // see: https://numpy.org/doc/stable/reference/generated/numpy.repeat.html#numpy.repeat + // flattens the input array and repeats each element 'repeats' times + // if 'repeats' is an array, it must have the same number of elements as the input array + @arkouda.registerND + proc repeatFlatMsg(cmd: string, msgArgs: borrowed MessageArgs, st: borrowed SymTab, param nd: int): MsgTuple throws { + param pn = Reflection.getRoutineName(); + const name = msgArgs.getValueOf("name"), + repeats = msgArgs.getValueOf("repeats"), + rname = st.nextName(); + + var gEnt: borrowed GenSymEntry = getGenericTypedArrayEntry(name, st), + gEntRepeats: borrowed GenSymEntry = getGenericTypedArrayEntry(repeats, st); + + proc doRepeatFlat(type t): MsgTuple throws { + const eIn = toSymEntry(gEnt, t, nd), + eRepeats = toSymEntry(gEntRepeats, int, 1), + aFlat = if nd == 1 then eIn.a else flatten(eIn.a); + + if eRepeats.a.size == 1 { + const rep = eRepeats.a[0], + eOut = st.addEntry(rname, aFlat.size * rep, t); + + // simple case: repeat each element of the input array 'rep' times + forall i in aFlat.domain do eOut.a[i*rep..#rep] = aFlat[i]; + + } else if eRepeats.a.size == aFlat.size { + // repeat each element of the input array by the corresponding element of 'repeats' + + // // serial algorithm: + // var start = 0; + // for idx in aFlat.domain { + // eOut.a[start..#eRepeats.a[idx]] = aFlat[idx]; + // start += eRepeats.a[idx]; + // } + + // compute the number of repeated elements in the output array owned by each task + const nTasksPerLoc = here.maxTaskPar; + var nRepsPerTask: [0.. 1 - { - var flat = makeDistArray({0.. order for the input array's indices - // e.g., order = k + (nz * j) + (nz * ny * i) - inline proc indexToOrder(idx: rank*int): int { - var order = 0; - for param i in 0..