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.

Sign up for GitHub

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

Jump to bottom

Add CuPy support #63

Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

Add CuPy support #63

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
934a248
Add CuPy support
dcherian Aug 9, 2022
d148074
More like arguments
dcherian May 12, 2023
a37c559
_full
dcherian Jun 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
43 changes: 30 additions & 13 deletions numpy_groupies/aggregate_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import numpy as np
from packaging.version import Version

from .utils import (
aggregate_common_doc,
Expand All @@ -13,13 +14,27 @@
check_fill_value,
input_validation,
iscomplexobj,
maxval,
minimum_dtype,
minimum_dtype_scalar,
minval,
maxval,
)


def _full(size, fill_value, *, dtype=None, like=None):
"""Backcompat for numpy < 1.20.0 which does not support the `like` kwarg"""
if (
like is not None # numpy bug?
and not np.isscalar(like) # scalars don't work
and Version(np.__version__) >= Version("1.20.0")
):
kwargs = {"like": like}
else:
kwargs = {}

return np.full(size, fill_value=fill_value, dtype=dtype, **kwargs)


def _sum(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype_scalar(fill_value, dtype, a)

Expand All @@ -44,7 +59,7 @@ def _sum(group_idx, a, size, fill_value, dtype=None):

def _prod(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype_scalar(fill_value, dtype, a)
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
if fill_value != 1:
ret[group_idx] = 1 # product starts from 1
np.multiply.at(ret, group_idx, a)
Expand All @@ -57,7 +72,7 @@ def _len(group_idx, a, size, fill_value, dtype=None):

def _last(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype(fill_value, dtype or a.dtype)
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
# repeated indexing gives last value, see:
# the phrase "leaving behind the last value" on this page:
# http://wiki.scipy.org/Tentative_NumPy_Tutorial
Expand All @@ -67,14 +82,14 @@ def _last(group_idx, a, size, fill_value, dtype=None):

def _first(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype(fill_value, dtype or a.dtype)
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
ret[group_idx[::-1]] = a[::-1] # same trick as _last, but in reverse
return ret


def _all(group_idx, a, size, fill_value, dtype=None):
check_boolean(fill_value)
ret = np.full(size, fill_value, dtype=bool)
ret = _full(size, fill_value, dtype=bool, like=a)
if not fill_value:
ret[group_idx] = True
ret[group_idx.compress(np.logical_not(a))] = False
Expand All @@ -83,7 +98,7 @@ def _all(group_idx, a, size, fill_value, dtype=None):

def _any(group_idx, a, size, fill_value, dtype=None):
check_boolean(fill_value)
ret = np.full(size, fill_value, dtype=bool)
ret = _full(size, fill_value, dtype=bool, like=a)
if fill_value:
ret[group_idx] = False
ret[group_idx.compress(a)] = True
Expand All @@ -93,7 +108,7 @@ def _any(group_idx, a, size, fill_value, dtype=None):
def _min(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype(fill_value, dtype or a.dtype)
dmax = maxval(fill_value, dtype)
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
if fill_value != dmax:
ret[group_idx] = dmax # min starts from maximum
np.minimum.at(ret, group_idx, a)
Expand All @@ -103,7 +118,7 @@ def _min(group_idx, a, size, fill_value, dtype=None):
def _max(group_idx, a, size, fill_value, dtype=None):
dtype = minimum_dtype(fill_value, dtype or a.dtype)
dmin = minval(fill_value, dtype)
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
if fill_value != dmin:
ret[group_idx] = dmin # max starts from minimum
np.maximum.at(ret, group_idx, a)
Expand All @@ -115,7 +130,7 @@ def _argmax(group_idx, a, size, fill_value, dtype=int, _nansqueeze=False):
group_max = _max(group_idx, a_, size, np.nan)
# nan should never be maximum, so use a and not a_
is_max = a == group_max[group_idx]
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
group_idx_max = group_idx[is_max]
(argmax,) = is_max.nonzero()
ret[group_idx_max[::-1]] = argmax[
Expand All @@ -129,7 +144,7 @@ def _argmin(group_idx, a, size, fill_value, dtype=int, _nansqueeze=False):
group_min = _min(group_idx, a_, size, np.nan)
# nan should never be minimum, so use a and not a_
is_min = a == group_min[group_idx]
ret = np.full(size, fill_value, dtype=dtype)
ret = _full(size, fill_value, dtype=dtype, like=a)
group_idx_min = group_idx[is_min]
(argmin,) = is_min.nonzero()
ret[group_idx_min[::-1]] = argmin[
Expand All @@ -144,11 +159,13 @@ def _mean(group_idx, a, size, fill_value, dtype=np.dtype(np.float64)):
counts = np.bincount(group_idx, minlength=size)
if iscomplexobj(a):
dtype = a.dtype # TODO: this is a bit clumsy
sums = np.empty(size, dtype=dtype)
sums = np.empty(size, dtype=dtype, like=a)
sums.real = np.bincount(group_idx, weights=a.real, minlength=size)
sums.imag = np.bincount(group_idx, weights=a.imag, minlength=size)
else:
sums = np.bincount(group_idx, weights=a, minlength=size).astype(dtype, copy=False)
sums = np.bincount(group_idx, weights=a, minlength=size).astype(
dtype, copy=False
)

with np.errstate(divide="ignore", invalid="ignore"):
ret = sums.astype(dtype, copy=False) / counts
Expand Down Expand Up @@ -223,7 +240,7 @@ def _generic_callable(
"""groups a by inds, and then applies foo to each group in turn, placing
the results in an array."""
groups = _array(group_idx, a, size, ())
ret = np.full(size, fill_value, dtype=dtype or np.float64)
ret = _full(size, fill_value, dtype=dtype or np.float64)

for i, grp in enumerate(groups):
if np.ndim(grp) == 1 and len(grp) > 0:
Expand Down
9 changes: 4 additions & 5 deletions numpy_groupies/utils_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ def _ravel_group_idx(group_idx, a, axis, size, order, method="ravel"):
size = []
for ii, s in enumerate(a.shape):
if method == "ravel":
ii_idx = group_idx_in if ii == axis else np.arange(s)
ii_idx = group_idx_in if ii == axis else np.arange(s, like=group_idx_in)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the only change; but requires numpy >= 1.20.0. Shall we add that as the min version?

I could also add a version check to switch between the old and new version.

Copy link
Owner

@ml31415 ml31415 Aug 10, 2022 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The last available numpy for py2 is 1.16.6 as far as I know. This would basically mean to completely drop the weave implementation, which so far is still the best benchmark to compete with for the numba implementation. Dropping py2 completely is probably unavoidable at some point, but if there is a way around it, remaining downwards compatible would be nice. So +1 for the switch!

ii_shape = [1] * ndim_a
ii_shape[ii] = s
group_idx.append(ii_idx.reshape(ii_shape))
Expand Down Expand Up @@ -259,10 +259,9 @@ def offset_labels(group_idx, inshape, axis, order, size):
group_idx = np.moveaxis(group_idx, axis, -1)
newshape = group_idx.shape[:-1] + (-1,)

group_idx = (
group_idx
+ np.arange(np.prod(newshape[:-1]), dtype=int).reshape(newshape) * size
)
offset_ = np.arange(np.prod(newshape[:-1]), dtype=int, like=group_idx).reshape(newshape)
dcherian marked this conversation as resolved.
Show resolved Hide resolved
group_idx = group_idx + offset_ * size

if axis not in (-1, len(inshape) - 1):
return np.moveaxis(group_idx, -1, axis)
else:
Expand Down