benamaurer
/
aoc-2022
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
aoc-2022/venv/Lib/site-packages/pandas/_libs/groupby.pyx

1783 lines
54 KiB
Cython
Raw Normal View History

Give it up for day 1 of smooth brain shenanigans! https://pbs.twimg.com/media/E1bxikWWEAEGYHU.png
2022-12-01 16:50:29 +00:00
cimport cython
from cython cimport (
Py_ssize_t,
floating,
)
from libc.stdlib cimport (
free,
malloc,
)
import numpy as np
cimport numpy as cnp
from numpy cimport (
complex64_t,
complex128_t,
float32_t,
float64_t,
int8_t,
int16_t,
int32_t,
int64_t,
intp_t,
ndarray,
uint8_t,
uint16_t,
uint32_t,
uint64_t,
)
from numpy.math cimport NAN
cnp.import_array()
from pandas._libs cimport util
from pandas._libs.algos cimport (
get_rank_nan_fill_val,
kth_smallest_c,
)
from pandas._libs.algos import (
ensure_platform_int,
groupsort_indexer,
rank_1d,
take_2d_axis1_float64_float64,
)
from pandas._libs.dtypes cimport (
numeric_object_t,
numeric_t,
)
from pandas._libs.missing cimport checknull
cdef int64_t NPY_NAT = util.get_nat()
_int64_max = np.iinfo(np.int64).max
cdef float64_t NaN = <float64_t>np.NaN
cdef enum InterpolationEnumType:
INTERPOLATION_LINEAR,
INTERPOLATION_LOWER,
INTERPOLATION_HIGHER,
INTERPOLATION_NEAREST,
INTERPOLATION_MIDPOINT
cdef inline float64_t median_linear(float64_t* a, int n) nogil:
cdef:
int i, j, na_count = 0
float64_t result
float64_t* tmp
if n == 0:
return NaN
# count NAs
for i in range(n):
if a[i] != a[i]:
na_count += 1
if na_count:
if na_count == n:
return NaN
tmp = <float64_t*>malloc((n - na_count) * sizeof(float64_t))
j = 0
for i in range(n):
if a[i] == a[i]:
tmp[j] = a[i]
j += 1
a = tmp
n -= na_count
if n % 2:
result = kth_smallest_c(a, n // 2, n)
else:
result = (kth_smallest_c(a, n // 2, n) +
kth_smallest_c(a, n // 2 - 1, n)) / 2
if na_count:
free(a)
return result
@cython.boundscheck(False)
@cython.wraparound(False)
def group_median_float64(
ndarray[float64_t, ndim=2] out,
ndarray[int64_t] counts,
ndarray[float64_t, ndim=2] values,
ndarray[intp_t] labels,
Py_ssize_t min_count=-1,
) -> None:
"""
Only aggregates on axis=0
"""
cdef:
Py_ssize_t i, j, N, K, ngroups, size
ndarray[intp_t] _counts
ndarray[float64_t, ndim=2] data
ndarray[intp_t] indexer
float64_t* ptr
assert min_count == -1, "'min_count' only used in sum and prod"
ngroups = len(counts)
N, K = (<object>values).shape
indexer, _counts = groupsort_indexer(labels, ngroups)
counts[:] = _counts[1:]
data = np.empty((K, N), dtype=np.float64)
ptr = <float64_t*>cnp.PyArray_DATA(data)
take_2d_axis1_float64_float64(values.T, indexer, out=data)
with nogil:
for i in range(K):
# exclude NA group
ptr += _counts[0]
for j in range(ngroups):
size = _counts[j + 1]
out[j, i] = median_linear(ptr, size)
ptr += size
@cython.boundscheck(False)
@cython.wraparound(False)
def group_cumprod_float64(
float64_t[:, ::1] out,
const float64_t[:, :] values,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
bint skipna=True,
) -> None:
"""
Cumulative product of columns of `values`, in row groups `labels`.
Parameters
----------
out : np.ndarray[np.float64, ndim=2]
Array to store cumprod in.
values : np.ndarray[np.float64, ndim=2]
Values to take cumprod of.
labels : np.ndarray[np.intp]
Labels to group by.
ngroups : int
Number of groups, larger than all entries of `labels`.
is_datetimelike : bool
Always false, `values` is never datetime-like.
skipna : bool
If true, ignore nans in `values`.
Notes
-----
This method modifies the `out` parameter, rather than returning an object.
"""
cdef:
Py_ssize_t i, j, N, K, size
float64_t val
float64_t[:, ::1] accum
intp_t lab
N, K = (<object>values).shape
accum = np.ones((ngroups, K), dtype=np.float64)
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
for j in range(K):
val = values[i, j]
if val == val:
accum[lab, j] *= val
out[i, j] = accum[lab, j]
else:
out[i, j] = NaN
if not skipna:
accum[lab, j] = NaN
ctypedef fused int64float_t:
int64_t
uint64_t
float32_t
float64_t
@cython.boundscheck(False)
@cython.wraparound(False)
def group_cumsum(
int64float_t[:, ::1] out,
ndarray[int64float_t, ndim=2] values,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
bint skipna=True,
const uint8_t[:, :] mask=None,
uint8_t[:, ::1] result_mask=None,
) -> None:
"""
Cumulative sum of columns of `values`, in row groups `labels`.
Parameters
----------
out : np.ndarray[ndim=2]
Array to store cumsum in.
values : np.ndarray[ndim=2]
Values to take cumsum of.
labels : np.ndarray[np.intp]
Labels to group by.
ngroups : int
Number of groups, larger than all entries of `labels`.
is_datetimelike : bool
True if `values` contains datetime-like entries.
skipna : bool
If true, ignore nans in `values`.
mask: np.ndarray[uint8], optional
Mask of values
result_mask: np.ndarray[int8], optional
Mask of out array
Notes
-----
This method modifies the `out` parameter, rather than returning an object.
"""
cdef:
Py_ssize_t i, j, N, K, size
int64float_t val, y, t, na_val
int64float_t[:, ::1] accum, compensation
uint8_t[:, ::1] accum_mask
intp_t lab
bint isna_entry, isna_prev = False
bint uses_mask = mask is not None
N, K = (<object>values).shape
if uses_mask:
accum_mask = np.zeros((ngroups, K), dtype="uint8")
accum = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
compensation = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
na_val = _get_na_val(<int64float_t>0, is_datetimelike)
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = _treat_as_na(val, is_datetimelike)
if not skipna:
if uses_mask:
isna_prev = accum_mask[lab, j]
else:
isna_prev = _treat_as_na(accum[lab, j], is_datetimelike)
if isna_prev:
if uses_mask:
result_mask[i, j] = True
# Be deterministic, out was initialized as empty
out[i, j] = 0
else:
out[i, j] = na_val
continue
if isna_entry:
if uses_mask:
result_mask[i, j] = True
# Be deterministic, out was initialized as empty
out[i, j] = 0
else:
out[i, j] = na_val
if not skipna:
if uses_mask:
accum_mask[lab, j] = True
else:
accum[lab, j] = na_val
else:
# For floats, use Kahan summation to reduce floating-point
# error (https://en.wikipedia.org/wiki/Kahan_summation_algorithm)
if int64float_t == float32_t or int64float_t == float64_t:
y = val - compensation[lab, j]
t = accum[lab, j] + y
compensation[lab, j] = t - accum[lab, j] - y
else:
t = val + accum[lab, j]
accum[lab, j] = t
out[i, j] = t
@cython.boundscheck(False)
@cython.wraparound(False)
def group_shift_indexer(
int64_t[::1] out,
const intp_t[::1] labels,
int ngroups,
int periods,
) -> None:
cdef:
Py_ssize_t N, i, j, ii, lab
int offset = 0, sign
int64_t idxer, idxer_slot
int64_t[::1] label_seen = np.zeros(ngroups, dtype=np.int64)
int64_t[:, ::1] label_indexer
N, = (<object>labels).shape
if periods < 0:
periods = -periods
offset = N - 1
sign = -1
elif periods > 0:
offset = 0
sign = 1
if periods == 0:
with nogil:
for i in range(N):
out[i] = i
else:
# array of each previous indexer seen
label_indexer = np.zeros((ngroups, periods), dtype=np.int64)
with nogil:
for i in range(N):
# reverse iterator if shifting backwards
ii = offset + sign * i
lab = labels[ii]
# Skip null keys
if lab == -1:
out[ii] = -1
continue
label_seen[lab] += 1
idxer_slot = label_seen[lab] % periods
idxer = label_indexer[lab, idxer_slot]
if label_seen[lab] > periods:
out[ii] = idxer
else:
out[ii] = -1
label_indexer[lab, idxer_slot] = ii
@cython.wraparound(False)
@cython.boundscheck(False)
def group_fillna_indexer(
ndarray[intp_t] out,
ndarray[intp_t] labels,
ndarray[intp_t] sorted_labels,
ndarray[uint8_t] mask,
str direction,
int64_t limit,
bint dropna,
) -> None:
"""
Indexes how to fill values forwards or backwards within a group.
Parameters
----------
out : np.ndarray[np.intp]
Values into which this method will write its results.
labels : np.ndarray[np.intp]
Array containing unique label for each group, with its ordering
matching up to the corresponding record in `values`.
sorted_labels : np.ndarray[np.intp]
obtained by `np.argsort(labels, kind="mergesort")`; reversed if
direction == "bfill"
values : np.ndarray[np.uint8]
Containing the truth value of each element.
mask : np.ndarray[np.uint8]
Indicating whether a value is na or not.
direction : {'ffill', 'bfill'}
Direction for fill to be applied (forwards or backwards, respectively)
limit : Consecutive values to fill before stopping, or -1 for no limit
dropna : Flag to indicate if NaN groups should return all NaN values
Notes
-----
This method modifies the `out` parameter rather than returning an object
"""
cdef:
Py_ssize_t i, N, idx
intp_t curr_fill_idx=-1
int64_t filled_vals = 0
N = len(out)
# Make sure all arrays are the same size
assert N == len(labels) == len(mask)
with nogil:
for i in range(N):
idx = sorted_labels[i]
if dropna and labels[idx] == -1: # nan-group gets nan-values
curr_fill_idx = -1
elif mask[idx] == 1: # is missing
# Stop filling once we've hit the limit
if filled_vals >= limit and limit != -1:
curr_fill_idx = -1
filled_vals += 1
else: # reset items when not missing
filled_vals = 0
curr_fill_idx = idx
out[idx] = curr_fill_idx
# If we move to the next group, reset
# the fill_idx and counter
if i == N - 1 or labels[idx] != labels[sorted_labels[i + 1]]:
curr_fill_idx = -1
filled_vals = 0
@cython.boundscheck(False)
@cython.wraparound(False)
def group_any_all(
int8_t[:, ::1] out,
const int8_t[:, :] values,
const intp_t[::1] labels,
const uint8_t[:, :] mask,
str val_test,
bint skipna,
bint nullable,
) -> None:
"""
Aggregated boolean values to show truthfulness of group elements. If the
input is a nullable type (nullable=True), the result will be computed
using Kleene logic.
Parameters
----------
out : np.ndarray[np.int8]
Values into which this method will write its results.
labels : np.ndarray[np.intp]
Array containing unique label for each group, with its
ordering matching up to the corresponding record in `values`
values : np.ndarray[np.int8]
Containing the truth value of each element.
mask : np.ndarray[np.uint8]
Indicating whether a value is na or not.
val_test : {'any', 'all'}
String object dictating whether to use any or all truth testing
skipna : bool
Flag to ignore nan values during truth testing
nullable : bool
Whether or not the input is a nullable type. If True, the
result will be computed using Kleene logic
Notes
-----
This method modifies the `out` parameter rather than returning an object.
The returned values will either be 0, 1 (False or True, respectively), or
-1 to signify a masked position in the case of a nullable input.
"""
cdef:
Py_ssize_t i, j, N = len(labels), K = out.shape[1]
intp_t lab
int8_t flag_val, val
if val_test == 'all':
# Because the 'all' value of an empty iterable in Python is True we can
# start with an array full of ones and set to zero when a False value
# is encountered
flag_val = 0
elif val_test == 'any':
# Because the 'any' value of an empty iterable in Python is False we
# can start with an array full of zeros and set to one only if any
# value encountered is True
flag_val = 1
else:
raise ValueError("'bool_func' must be either 'any' or 'all'!")
out[:] = 1 - flag_val
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
for j in range(K):
if skipna and mask[i, j]:
continue
if nullable and mask[i, j]:
# Set the position as masked if `out[lab] != flag_val`, which
# would indicate True/False has not yet been seen for any/all,
# so by Kleene logic the result is currently unknown
if out[lab, j] != flag_val:
out[lab, j] = -1
continue
val = values[i, j]
# If True and 'any' or False and 'all', the result is
# already determined
if val == flag_val:
out[lab, j] = flag_val
# ----------------------------------------------------------------------
# group_sum, group_prod, group_var, group_mean, group_ohlc
# ----------------------------------------------------------------------
ctypedef fused mean_t:
float64_t
float32_t
complex64_t
complex128_t
ctypedef fused sum_t:
mean_t
int64_t
uint64_t
object
@cython.wraparound(False)
@cython.boundscheck(False)
def group_sum(
sum_t[:, ::1] out,
int64_t[::1] counts,
ndarray[sum_t, ndim=2] values,
const intp_t[::1] labels,
const uint8_t[:, :] mask,
uint8_t[:, ::1] result_mask=None,
Py_ssize_t min_count=0,
bint is_datetimelike=False,
) -> None:
"""
Only aggregates on axis=0 using Kahan summation
"""
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
sum_t val, t, y
sum_t[:, ::1] sumx, compensation
int64_t[:, ::1] nobs
Py_ssize_t len_values = len(values), len_labels = len(labels)
bint uses_mask = mask is not None
bint isna_entry
if len_values != len_labels:
raise ValueError("len(index) != len(labels)")
nobs = np.zeros((<object>out).shape, dtype=np.int64)
# the below is equivalent to `np.zeros_like(out)` but faster
sumx = np.zeros((<object>out).shape, dtype=(<object>out).base.dtype)
compensation = np.zeros((<object>out).shape, dtype=(<object>out).base.dtype)
N, K = (<object>values).shape
if sum_t is object:
# NB: this does not use 'compensation' like the non-object track does.
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
# not nan
if not checknull(val):
nobs[lab, j] += 1
if nobs[lab, j] == 1:
# i.e. we haven't added anything yet; avoid TypeError
# if e.g. val is a str and sumx[lab, j] is 0
t = val
else:
t = sumx[lab, j] + val
sumx[lab, j] = t
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
out[i, j] = None
else:
out[i, j] = sumx[i, j]
else:
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
# not nan
# With dt64/td64 values, values have been cast to float64
# instead if int64 for group_sum, but the logic
# is otherwise the same as in _treat_as_na
if uses_mask:
isna_entry = mask[i, j]
elif (sum_t is float32_t or sum_t is float64_t
or sum_t is complex64_t or sum_t is complex64_t):
# avoid warnings because of equality comparison
isna_entry = not val == val
elif sum_t is int64_t and is_datetimelike and val == NPY_NAT:
isna_entry = True
else:
isna_entry = False
if not isna_entry:
nobs[lab, j] += 1
y = val - compensation[lab, j]
t = sumx[lab, j] + y
compensation[lab, j] = t - sumx[lab, j] - y
sumx[lab, j] = t
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
# if we are integer dtype, not is_datetimelike, and
# not uses_mask, then getting here implies that
# counts[i] < min_count, which means we will
# be cast to float64 and masked at the end
# of WrappedCythonOp._call_cython_op. So we can safely
# set a placeholder value in out[i, j].
if uses_mask:
result_mask[i, j] = True
elif (sum_t is float32_t or sum_t is float64_t
or sum_t is complex64_t or sum_t is complex64_t):
out[i, j] = NAN
elif sum_t is int64_t:
out[i, j] = NPY_NAT
else:
# placeholder, see above
out[i, j] = 0
else:
out[i, j] = sumx[i, j]
@cython.wraparound(False)
@cython.boundscheck(False)
def group_prod(
int64float_t[:, ::1] out,
int64_t[::1] counts,
ndarray[int64float_t, ndim=2] values,
const intp_t[::1] labels,
const uint8_t[:, ::1] mask,
uint8_t[:, ::1] result_mask=None,
Py_ssize_t min_count=0,
) -> None:
"""
Only aggregates on axis=0
"""
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
int64float_t val, count
int64float_t[:, ::1] prodx
int64_t[:, ::1] nobs
Py_ssize_t len_values = len(values), len_labels = len(labels)
bint isna_entry, uses_mask = mask is not None
if len_values != len_labels:
raise ValueError("len(index) != len(labels)")
nobs = np.zeros((<object>out).shape, dtype=np.int64)
prodx = np.ones((<object>out).shape, dtype=(<object>out).base.dtype)
N, K = (<object>values).shape
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
elif int64float_t is float32_t or int64float_t is float64_t:
isna_entry = not val == val
else:
isna_entry = False
if not isna_entry:
nobs[lab, j] += 1
prodx[lab, j] *= val
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
# else case is not possible
if uses_mask:
result_mask[i, j] = True
# Be deterministic, out was initialized as empty
out[i, j] = 0
elif int64float_t is float32_t or int64float_t is float64_t:
out[i, j] = NAN
else:
# we only get here when < mincount which gets handled later
pass
else:
out[i, j] = prodx[i, j]
@cython.wraparound(False)
@cython.boundscheck(False)
@cython.cdivision(True)
def group_var(
floating[:, ::1] out,
int64_t[::1] counts,
ndarray[floating, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
int64_t ddof=1,
) -> None:
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
floating val, ct, oldmean
floating[:, ::1] mean
int64_t[:, ::1] nobs
Py_ssize_t len_values = len(values), len_labels = len(labels)
assert min_count == -1, "'min_count' only used in sum and prod"
if len_values != len_labels:
raise ValueError("len(index) != len(labels)")
nobs = np.zeros((<object>out).shape, dtype=np.int64)
mean = np.zeros((<object>out).shape, dtype=(<object>out).base.dtype)
N, K = (<object>values).shape
out[:, :] = 0.0
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
# not nan
if val == val:
nobs[lab, j] += 1
oldmean = mean[lab, j]
mean[lab, j] += (val - oldmean) / nobs[lab, j]
out[lab, j] += (val - mean[lab, j]) * (val - oldmean)
for i in range(ncounts):
for j in range(K):
ct = nobs[i, j]
if ct <= ddof:
out[i, j] = NAN
else:
out[i, j] /= (ct - ddof)
@cython.wraparound(False)
@cython.boundscheck(False)
def group_mean(
mean_t[:, ::1] out,
int64_t[::1] counts,
ndarray[mean_t, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
bint is_datetimelike=False,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
) -> None:
"""
Compute the mean per label given a label assignment for each value.
NaN values are ignored.
Parameters
----------
out : np.ndarray[floating]
Values into which this method will write its results.
counts : np.ndarray[int64]
A zeroed array of the same shape as labels,
populated by group sizes during algorithm.
values : np.ndarray[floating]
2-d array of the values to find the mean of.
labels : np.ndarray[np.intp]
Array containing unique label for each group, with its
ordering matching up to the corresponding record in `values`.
min_count : Py_ssize_t
Only used in sum and prod. Always -1.
is_datetimelike : bool
True if `values` contains datetime-like entries.
mask : ndarray[bool, ndim=2], optional
Not used.
result_mask : ndarray[bool, ndim=2], optional
Not used.
Notes
-----
This method modifies the `out` parameter rather than returning an object.
`counts` is modified to hold group sizes
"""
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
mean_t val, count, y, t, nan_val
mean_t[:, ::1] sumx, compensation
int64_t[:, ::1] nobs
Py_ssize_t len_values = len(values), len_labels = len(labels)
assert min_count == -1, "'min_count' only used in sum and prod"
if len_values != len_labels:
raise ValueError("len(index) != len(labels)")
# the below is equivalent to `np.zeros_like(out)` but faster
nobs = np.zeros((<object>out).shape, dtype=np.int64)
sumx = np.zeros((<object>out).shape, dtype=(<object>out).base.dtype)
compensation = np.zeros((<object>out).shape, dtype=(<object>out).base.dtype)
N, K = (<object>values).shape
nan_val = NPY_NAT if is_datetimelike else NAN
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
# not nan
if val == val and not (is_datetimelike and val == NPY_NAT):
nobs[lab, j] += 1
y = val - compensation[lab, j]
t = sumx[lab, j] + y
compensation[lab, j] = t - sumx[lab, j] - y
sumx[lab, j] = t
for i in range(ncounts):
for j in range(K):
count = nobs[i, j]
if nobs[i, j] == 0:
out[i, j] = nan_val
else:
out[i, j] = sumx[i, j] / count
@cython.wraparound(False)
@cython.boundscheck(False)
def group_ohlc(
int64float_t[:, ::1] out,
int64_t[::1] counts,
ndarray[int64float_t, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
) -> None:
"""
Only aggregates on axis=0
"""
cdef:
Py_ssize_t i, j, N, K, lab
int64float_t val
uint8_t[::1] first_element_set
bint isna_entry, uses_mask = not mask is None
assert min_count == -1, "'min_count' only used in sum and prod"
if len(labels) == 0:
return
N, K = (<object>values).shape
if out.shape[1] != 4:
raise ValueError('Output array must have 4 columns')
if K > 1:
raise NotImplementedError("Argument 'values' must have only one dimension")
if int64float_t is float32_t or int64float_t is float64_t:
out[:] = np.nan
else:
out[:] = 0
first_element_set = np.zeros((<object>counts).shape, dtype=np.uint8)
if uses_mask:
result_mask[:] = True
with nogil:
for i in range(N):
lab = labels[i]
if lab == -1:
continue
counts[lab] += 1
val = values[i, 0]
if uses_mask:
isna_entry = mask[i, 0]
elif int64float_t is float32_t or int64float_t is float64_t:
isna_entry = val != val
else:
isna_entry = False
if isna_entry:
continue
if not first_element_set[lab]:
out[lab, 0] = out[lab, 1] = out[lab, 2] = out[lab, 3] = val
first_element_set[lab] = True
if uses_mask:
result_mask[lab] = False
else:
out[lab, 1] = max(out[lab, 1], val)
out[lab, 2] = min(out[lab, 2], val)
out[lab, 3] = val
@cython.boundscheck(False)
@cython.wraparound(False)
def group_quantile(
ndarray[float64_t, ndim=2] out,
ndarray[numeric_t, ndim=1] values,
ndarray[intp_t] labels,
ndarray[uint8_t] mask,
const intp_t[:] sort_indexer,
const float64_t[:] qs,
str interpolation,
) -> None:
"""
Calculate the quantile per group.
Parameters
----------
out : np.ndarray[np.float64, ndim=2]
Array of aggregated values that will be written to.
values : np.ndarray
Array containing the values to apply the function against.
labels : ndarray[np.intp]
Array containing the unique group labels.
sort_indexer : ndarray[np.intp]
Indices describing sort order by values and labels.
qs : ndarray[float64_t]
The quantile values to search for.
interpolation : {'linear', 'lower', 'highest', 'nearest', 'midpoint'}
Notes
-----
Rather than explicitly returning a value, this function modifies the
provided `out` parameter.
"""
cdef:
Py_ssize_t i, N=len(labels), ngroups, grp_sz, non_na_sz, k, nqs
Py_ssize_t grp_start=0, idx=0
intp_t lab
InterpolationEnumType interp
float64_t q_val, q_idx, frac, val, next_val
int64_t[::1] counts, non_na_counts
assert values.shape[0] == N
if any(not (0 <= q <= 1) for q in qs):
wrong = [x for x in qs if not (0 <= x <= 1)][0]
raise ValueError(
f"Each 'q' must be between 0 and 1. Got '{wrong}' instead"
)
inter_methods = {
'linear': INTERPOLATION_LINEAR,
'lower': INTERPOLATION_LOWER,
'higher': INTERPOLATION_HIGHER,
'nearest': INTERPOLATION_NEAREST,
'midpoint': INTERPOLATION_MIDPOINT,
}
interp = inter_methods[interpolation]
nqs = len(qs)
ngroups = len(out)
counts = np.zeros(ngroups, dtype=np.int64)
non_na_counts = np.zeros(ngroups, dtype=np.int64)
# First figure out the size of every group
with nogil:
for i in range(N):
lab = labels[i]
if lab == -1: # NA group label
continue
counts[lab] += 1
if not mask[i]:
non_na_counts[lab] += 1
with nogil:
for i in range(ngroups):
# Figure out how many group elements there are
grp_sz = counts[i]
non_na_sz = non_na_counts[i]
if non_na_sz == 0:
for k in range(nqs):
out[i, k] = NaN
else:
for k in range(nqs):
q_val = qs[k]
# Calculate where to retrieve the desired value
# Casting to int will intentionally truncate result
idx = grp_start + <int64_t>(q_val * <float64_t>(non_na_sz - 1))
val = values[sort_indexer[idx]]
# If requested quantile falls evenly on a particular index
# then write that index's value out. Otherwise interpolate
q_idx = q_val * (non_na_sz - 1)
frac = q_idx % 1
if frac == 0.0 or interp == INTERPOLATION_LOWER:
out[i, k] = val
else:
next_val = values[sort_indexer[idx + 1]]
if interp == INTERPOLATION_LINEAR:
out[i, k] = val + (next_val - val) * frac
elif interp == INTERPOLATION_HIGHER:
out[i, k] = next_val
elif interp == INTERPOLATION_MIDPOINT:
out[i, k] = (val + next_val) / 2.0
elif interp == INTERPOLATION_NEAREST:
if frac > .5 or (frac == .5 and q_val > .5): # Always OK?
out[i, k] = next_val
else:
out[i, k] = val
# Increment the index reference in sorted_arr for the next group
grp_start += grp_sz
# ----------------------------------------------------------------------
# group_nth, group_last, group_rank
# ----------------------------------------------------------------------
cdef inline bint _treat_as_na(numeric_object_t val, bint is_datetimelike) nogil:
if numeric_object_t is object:
# Should never be used, but we need to avoid the `val != val` below
# or else cython will raise about gil acquisition.
raise NotImplementedError
elif numeric_object_t is int64_t:
return is_datetimelike and val == NPY_NAT
elif numeric_object_t is float32_t or numeric_object_t is float64_t:
return val != val
else:
# non-datetimelike integer
return False
cdef numeric_object_t _get_min_or_max(numeric_object_t val, bint compute_max, bint is_datetimelike):
"""
Find either the min or the max supported by numeric_object_t; 'val' is a
placeholder to effectively make numeric_object_t an argument.
"""
return get_rank_nan_fill_val(
not compute_max,
val=val,
is_datetimelike=is_datetimelike,
)
cdef numeric_t _get_na_val(numeric_t val, bint is_datetimelike):
cdef:
numeric_t na_val
if numeric_t == float32_t or numeric_t == float64_t:
na_val = NaN
elif numeric_t is int64_t and is_datetimelike:
na_val = NPY_NAT
else:
# Used in case of masks
na_val = 0
return na_val
# TODO(cython3): GH#31710 use memorviews once cython 0.30 is released so we can
# use `const numeric_object_t[:, :] values`
@cython.wraparound(False)
@cython.boundscheck(False)
def group_last(
numeric_object_t[:, ::1] out,
int64_t[::1] counts,
ndarray[numeric_object_t, ndim=2] values,
const intp_t[::1] labels,
const uint8_t[:, :] mask,
uint8_t[:, ::1] result_mask=None,
Py_ssize_t min_count=-1,
bint is_datetimelike=False,
) -> None:
"""
Only aggregates on axis=0
"""
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
numeric_object_t val
ndarray[numeric_object_t, ndim=2] resx
ndarray[int64_t, ndim=2] nobs
bint uses_mask = mask is not None
bint isna_entry
# TODO(cython3):
# Instead of `labels.shape[0]` use `len(labels)`
if not len(values) == labels.shape[0]:
raise AssertionError("len(index) != len(labels)")
min_count = max(min_count, 1)
nobs = np.zeros((<object>out).shape, dtype=np.int64)
if numeric_object_t is object:
resx = np.empty((<object>out).shape, dtype=object)
else:
resx = np.empty_like(out)
N, K = (<object>values).shape
if numeric_object_t is object:
# TODO(cython3): De-duplicate once conditional-nogil is available
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = checknull(val)
if not isna_entry:
# NB: use _treat_as_na here once
# conditional-nogil is available.
nobs[lab, j] += 1
resx[lab, j] = val
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
out[i, j] = None
else:
out[i, j] = resx[i, j]
else:
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = _treat_as_na(val, is_datetimelike)
if not isna_entry:
nobs[lab, j] += 1
resx[lab, j] = val
for i in range(ncounts):
for j in range(K):
# TODO(cython3): the entire next block can be shared
# across 3 places once conditional-nogil is available
if nobs[i, j] < min_count:
# if we are integer dtype, not is_datetimelike, and
# not uses_mask, then getting here implies that
# counts[i] < min_count, which means we will
# be cast to float64 and masked at the end
# of WrappedCythonOp._call_cython_op. So we can safely
# set a placeholder value in out[i, j].
if uses_mask:
result_mask[i, j] = True
elif numeric_object_t is float32_t or numeric_object_t is float64_t:
out[i, j] = NAN
elif numeric_object_t is int64_t:
# Per above, this is a placeholder in
# non-is_datetimelike cases.
out[i, j] = NPY_NAT
else:
# placeholder, see above
out[i, j] = 0
else:
out[i, j] = resx[i, j]
# TODO(cython3): GH#31710 use memorviews once cython 0.30 is released so we can
# use `const numeric_object_t[:, :] values`
@cython.wraparound(False)
@cython.boundscheck(False)
def group_nth(
numeric_object_t[:, ::1] out,
int64_t[::1] counts,
ndarray[numeric_object_t, ndim=2] values,
const intp_t[::1] labels,
const uint8_t[:, :] mask,
uint8_t[:, ::1] result_mask=None,
int64_t min_count=-1,
int64_t rank=1,
bint is_datetimelike=False,
) -> None:
"""
Only aggregates on axis=0
"""
cdef:
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
numeric_object_t val
ndarray[numeric_object_t, ndim=2] resx
ndarray[int64_t, ndim=2] nobs
bint uses_mask = mask is not None
bint isna_entry
# TODO(cython3):
# Instead of `labels.shape[0]` use `len(labels)`
if not len(values) == labels.shape[0]:
raise AssertionError("len(index) != len(labels)")
min_count = max(min_count, 1)
nobs = np.zeros((<object>out).shape, dtype=np.int64)
if numeric_object_t is object:
resx = np.empty((<object>out).shape, dtype=object)
else:
resx = np.empty_like(out)
N, K = (<object>values).shape
if numeric_object_t is object:
# TODO(cython3): De-duplicate once conditional-nogil is available
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = checknull(val)
if not isna_entry:
# NB: use _treat_as_na here once
# conditional-nogil is available.
nobs[lab, j] += 1
if nobs[lab, j] == rank:
resx[lab, j] = val
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
out[i, j] = None
else:
out[i, j] = resx[i, j]
else:
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = _treat_as_na(val, is_datetimelike)
if not isna_entry:
nobs[lab, j] += 1
if nobs[lab, j] == rank:
resx[lab, j] = val
# TODO: de-dup this whole block with group_last?
for i in range(ncounts):
for j in range(K):
if nobs[i, j] < min_count:
# if we are integer dtype, not is_datetimelike, and
# not uses_mask, then getting here implies that
# counts[i] < min_count, which means we will
# be cast to float64 and masked at the end
# of WrappedCythonOp._call_cython_op. So we can safely
# set a placeholder value in out[i, j].
if uses_mask:
result_mask[i, j] = True
# set out[i, j] to 0 to be deterministic, as
# it was initialized with np.empty. Also ensures
# we can downcast out if appropriate.
out[i, j] = 0
elif numeric_object_t is float32_t or numeric_object_t is float64_t:
out[i, j] = NAN
elif numeric_object_t is int64_t:
# Per above, this is a placeholder in
# non-is_datetimelike cases.
out[i, j] = NPY_NAT
else:
# placeholder, see above
out[i, j] = 0
else:
out[i, j] = resx[i, j]
@cython.boundscheck(False)
@cython.wraparound(False)
def group_rank(
float64_t[:, ::1] out,
ndarray[numeric_object_t, ndim=2] values,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
str ties_method="average",
bint ascending=True,
bint pct=False,
str na_option="keep",
const uint8_t[:, :] mask=None,
) -> None:
"""
Provides the rank of values within each group.
Parameters
----------
out : np.ndarray[np.float64, ndim=2]
Values to which this method will write its results.
values : np.ndarray of numeric_object_t values to be ranked
labels : np.ndarray[np.intp]
Array containing unique label for each group, with its ordering
matching up to the corresponding record in `values`
ngroups : int
This parameter is not used, is needed to match signatures of other
groupby functions.
is_datetimelike : bool
True if `values` contains datetime-like entries.
ties_method : {'average', 'min', 'max', 'first', 'dense'}, default 'average'
* average: average rank of group
* min: lowest rank in group
* max: highest rank in group
* first: ranks assigned in order they appear in the array
* dense: like 'min', but rank always increases by 1 between groups
ascending : bool, default True
False for ranks by high (1) to low (N)
na_option : {'keep', 'top', 'bottom'}, default 'keep'
pct : bool, default False
Compute percentage rank of data within each group
na_option : {'keep', 'top', 'bottom'}, default 'keep'
* keep: leave NA values where they are
* top: smallest rank if ascending
* bottom: smallest rank if descending
mask : np.ndarray[bool] or None, default None
Notes
-----
This method modifies the `out` parameter rather than returning an object
"""
cdef:
Py_ssize_t i, k, N
ndarray[float64_t, ndim=1] result
const uint8_t[:] sub_mask
N = values.shape[1]
for k in range(N):
if mask is None:
sub_mask = None
else:
sub_mask = mask[:, k]
result = rank_1d(
values=values[:, k],
labels=labels,
is_datetimelike=is_datetimelike,
ties_method=ties_method,
ascending=ascending,
pct=pct,
na_option=na_option,
mask=sub_mask,
)
for i in range(len(result)):
if labels[i] >= 0:
out[i, k] = result[i]
# ----------------------------------------------------------------------
# group_min, group_max
# ----------------------------------------------------------------------
@cython.wraparound(False)
@cython.boundscheck(False)
cdef group_min_max(
numeric_t[:, ::1] out,
int64_t[::1] counts,
ndarray[numeric_t, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
bint is_datetimelike=False,
bint compute_max=True,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
):
"""
Compute minimum/maximum of columns of `values`, in row groups `labels`.
Parameters
----------
out : np.ndarray[numeric_t, ndim=2]
Array to store result in.
counts : np.ndarray[int64]
Input as a zeroed array, populated by group sizes during algorithm
values : array
Values to find column-wise min/max of.
labels : np.ndarray[np.intp]
Labels to group by.
min_count : Py_ssize_t, default -1
The minimum number of non-NA group elements, NA result if threshold
is not met
is_datetimelike : bool
True if `values` contains datetime-like entries.
compute_max : bint, default True
True to compute group-wise max, False to compute min
mask : ndarray[bool, ndim=2], optional
If not None, indices represent missing values,
otherwise the mask will not be used
result_mask : ndarray[bool, ndim=2], optional
If not None, these specify locations in the output that are NA.
Modified in-place.
Notes
-----
This method modifies the `out` parameter, rather than returning an object.
`counts` is modified to hold group sizes
"""
cdef:
Py_ssize_t i, j, N, K, lab, ngroups = len(counts)
numeric_t val
ndarray[numeric_t, ndim=2] group_min_or_max
int64_t[:, ::1] nobs
bint uses_mask = mask is not None
bint isna_entry
# TODO(cython3):
# Instead of `labels.shape[0]` use `len(labels)`
if not len(values) == labels.shape[0]:
raise AssertionError("len(index) != len(labels)")
min_count = max(min_count, 1)
nobs = np.zeros((<object>out).shape, dtype=np.int64)
group_min_or_max = np.empty_like(out)
group_min_or_max[:] = _get_min_or_max(<numeric_t>0, compute_max, is_datetimelike)
N, K = (<object>values).shape
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
counts[lab] += 1
for j in range(K):
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = _treat_as_na(val, is_datetimelike)
if not isna_entry:
nobs[lab, j] += 1
if compute_max:
if val > group_min_or_max[lab, j]:
group_min_or_max[lab, j] = val
else:
if val < group_min_or_max[lab, j]:
group_min_or_max[lab, j] = val
for i in range(ngroups):
for j in range(K):
if nobs[i, j] < min_count:
# if we are integer dtype, not is_datetimelike, and
# not uses_mask, then getting here implies that
# counts[i] < min_count, which means we will
# be cast to float64 and masked at the end
# of WrappedCythonOp._call_cython_op. So we can safely
# set a placeholder value in out[i, j].
if uses_mask:
result_mask[i, j] = True
# set out[i, j] to 0 to be deterministic, as
# it was initialized with np.empty. Also ensures
# we can downcast out if appropriate.
out[i, j] = 0
elif numeric_t is float32_t or numeric_t is float64_t:
out[i, j] = NAN
elif numeric_t is int64_t:
# Per above, this is a placeholder in
# non-is_datetimelike cases.
out[i, j] = NPY_NAT
else:
# placeholder, see above
out[i, j] = 0
else:
out[i, j] = group_min_or_max[i, j]
@cython.wraparound(False)
@cython.boundscheck(False)
def group_max(
numeric_t[:, ::1] out,
int64_t[::1] counts,
ndarray[numeric_t, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
bint is_datetimelike=False,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
) -> None:
"""See group_min_max.__doc__"""
group_min_max(
out,
counts,
values,
labels,
min_count=min_count,
is_datetimelike=is_datetimelike,
compute_max=True,
mask=mask,
result_mask=result_mask,
)
@cython.wraparound(False)
@cython.boundscheck(False)
def group_min(
numeric_t[:, ::1] out,
int64_t[::1] counts,
ndarray[numeric_t, ndim=2] values,
const intp_t[::1] labels,
Py_ssize_t min_count=-1,
bint is_datetimelike=False,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
) -> None:
"""See group_min_max.__doc__"""
group_min_max(
out,
counts,
values,
labels,
min_count=min_count,
is_datetimelike=is_datetimelike,
compute_max=False,
mask=mask,
result_mask=result_mask,
)
@cython.boundscheck(False)
@cython.wraparound(False)
cdef group_cummin_max(
numeric_t[:, ::1] out,
ndarray[numeric_t, ndim=2] values,
const uint8_t[:, ::1] mask,
uint8_t[:, ::1] result_mask,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
bint skipna,
bint compute_max,
):
"""
Cumulative minimum/maximum of columns of `values`, in row groups `labels`.
Parameters
----------
out : np.ndarray[numeric_t, ndim=2]
Array to store cummin/max in.
values : np.ndarray[numeric_t, ndim=2]
Values to take cummin/max of.
mask : np.ndarray[bool] or None
If not None, indices represent missing values,
otherwise the mask will not be used
result_mask : ndarray[bool, ndim=2], optional
If not None, these specify locations in the output that are NA.
Modified in-place.
labels : np.ndarray[np.intp]
Labels to group by.
ngroups : int
Number of groups, larger than all entries of `labels`.
is_datetimelike : bool
True if `values` contains datetime-like entries.
skipna : bool
If True, ignore nans in `values`.
compute_max : bool
True if cumulative maximum should be computed, False
if cumulative minimum should be computed
Notes
-----
This method modifies the `out` parameter, rather than returning an object.
"""
cdef:
numeric_t[:, ::1] accum
Py_ssize_t i, j, N, K
numeric_t val, mval, na_val
uint8_t[:, ::1] seen_na
intp_t lab
bint na_possible
bint uses_mask = mask is not None
bint isna_entry
accum = np.empty((ngroups, (<object>values).shape[1]), dtype=values.dtype)
accum[:] = _get_min_or_max(<numeric_t>0, compute_max, is_datetimelike)
na_val = _get_na_val(<numeric_t>0, is_datetimelike)
if uses_mask:
na_possible = True
# Will never be used, just to avoid uninitialized warning
na_val = 0
elif numeric_t is float64_t or numeric_t is float32_t:
na_possible = True
elif is_datetimelike:
na_possible = True
else:
# Will never be used, just to avoid uninitialized warning
na_possible = False
if na_possible:
seen_na = np.zeros((<object>accum).shape, dtype=np.uint8)
N, K = (<object>values).shape
with nogil:
for i in range(N):
lab = labels[i]
if lab < 0:
continue
for j in range(K):
if not skipna and na_possible and seen_na[lab, j]:
if uses_mask:
result_mask[i, j] = 1
# Set to 0 ensures that we are deterministic and can
# downcast if appropriate
out[i, j] = 0
else:
out[i, j] = na_val
else:
val = values[i, j]
if uses_mask:
isna_entry = mask[i, j]
else:
isna_entry = _treat_as_na(val, is_datetimelike)
if not isna_entry:
mval = accum[lab, j]
if compute_max:
if val > mval:
accum[lab, j] = mval = val
else:
if val < mval:
accum[lab, j] = mval = val
out[i, j] = mval
else:
seen_na[lab, j] = 1
out[i, j] = val
@cython.boundscheck(False)
@cython.wraparound(False)
def group_cummin(
numeric_t[:, ::1] out,
ndarray[numeric_t, ndim=2] values,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
bint skipna=True,
) -> None:
"""See group_cummin_max.__doc__"""
group_cummin_max(
out=out,
values=values,
mask=mask,
result_mask=result_mask,
labels=labels,
ngroups=ngroups,
is_datetimelike=is_datetimelike,
skipna=skipna,
compute_max=False,
)
@cython.boundscheck(False)
@cython.wraparound(False)
def group_cummax(
numeric_t[:, ::1] out,
ndarray[numeric_t, ndim=2] values,
const intp_t[::1] labels,
int ngroups,
bint is_datetimelike,
const uint8_t[:, ::1] mask=None,
uint8_t[:, ::1] result_mask=None,
bint skipna=True,
) -> None:
"""See group_cummin_max.__doc__"""
group_cummin_max(
out=out,
values=values,
mask=mask,
result_mask=result_mask,
labels=labels,
ngroups=ngroups,
is_datetimelike=is_datetimelike,
skipna=skipna,
compute_max=True,
)