Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closes 3428 putmask optimization #3749

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 46 additions & 29 deletions arkouda/numpy/_numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -2291,59 +2291,76 @@ def array_equal(pda_a: pdarray, pda_b: pdarray, equal_nan: bool = False):
return ak_all(pda_a == pda_b)


def putmask(pda: pdarray, mask: Union[bool, pdarray], values: pdarray):
def putmask(
A : pdarray ,
mask : pdarray,
Values : pdarray
) : # doesn't return anything, as A is overwritten in place
"""
Overwrite elements of a pdarray at indices where mask is True
Overwrites elements of A with elements from B based upon a mask array.
Similar to numpy.putmask, where mask = False, A retains its original value,
but where mask = True, A is overwritten with the corresponding entry from Values.

This is similar to ak.where, except that (1) no new pdarray is created, and
(2) Values does not have to be the same size as A and mask.

Parameters
----------
pda : pdarray, source data, also output data
pda = input where mask is False, = values where mask is True
mask : a scalar boolean, or a pdarray of booleans
values : pdarray, replacement data

Returns
-------
None - pda is modified in-place

Notes
-----
If values.size != a.size, values is repeated and/or pruned as needed to
make sizes match, because ak.where requires matching sizes.
A : pdarray
Value(s) used when mask is False (see Notes for allowed dtypes)
mask : pdarray
Used to choose values from A or B, must be same size as A, and of type ak.bool_
Values : pdarray
Value(s) used when mask is False (see Notes for allowed dtypes)

Examples
-------
>>> a = ak.array(np.arange(10))
>>> ak.putmask (a,a>2,a**2)
>>> a
array ([0,1,2,9,16,25,36,49,64,81])

>>> a = ak.array(np.arange(10))
>>> values = ak.array([3,2])
>>> ak.putmask (a,a>2,values)
>>> a
array ([0,1,2,2,3,2,3,2,3,2])

Raises
------
TypeError
Raised if a and values are not the same type
RuntimeError
Raised if mask is not same size as A, or if A.dtype and Values.dtype are not
an allowed pair (see Notes for details).

"""
Notes
-----
A and mask must be the same size. Values can be any size.

from arkouda.pdarraysetops import concatenate
Allowed dtypes for A and Values conform to types accepted by numpy putmask.

If A is ak.float64, Values can be ak.float64, ak.int64, ak.uint64, ak.bool_.

# check for matching types
If A is ak.int64, Values can be ak.int64 or ak.bool_.

if values.dtype != pda.dtype:
raise TypeError("ak.putmask requires arrays of matching type")
If A is ak.uint64, Values can be ak.int64, ak.uint64, or ak.bool_.

# if values is not the same size as pda, repeat it and/or prune it as needed
If A is ak.bool_, Values must be ak.bool_.

growth = pda.size // values.size + (0 if pda.size % values.size == 0 else 1)
result = concatenate(growth * [values])
if result.size > pda.size:
reduction = result.size % pda.size
result = result[:-(reduction)]
Only one conditional clause is supported e.g., n < 5, n > 1, which is supported in numpy
is not currently supported in Arkouda

pda[:] = where(mask, result, pda) # pda[:] = allows us to return modified value
Only 1D pdarrays are implemented for now.
"""
generic_msg(
cmd=f"efunc3vv{mask.ndim}D",
args={
"func": "putmask",
"condition": mask,
"a": A,
"b": Values,
},
)
return


def eye(rows: int_scalars, cols: int_scalars, diag: int_scalars = 0, dt: type = akint64):
Expand Down
Loading
Loading