mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-16 11:18:11 +08:00
Scott Sanderson
49bb8264dc
- Added test coverage for grouped and masked top/bottom. - Added test coverage for grouped rank on datetime factors. - Fixed an issue where grouped rank would fail on datetime inputs because unary-negative isn't defined for datetimes. We now instead directly invoke a function from rank.pyx that does the normalizations as neeeded. - Fixed an issue where GroupedRowTransform assumed that it produced the same dtype as its input. This isn't true for rank() of a datetime-dtype factor. GroupedRowTransform now takes a required dtype parameter. - Similarly, fixed an issue where GroupedRowTransform assumed that its missing_value was the same as its parent's, which isn't true for rank() of a datetime-dtype factor. GroupedRowTransform now takes a required dtype parameter. - Fixed an issue where Factor.demean() and Factor.zscore() weren't properly cached because their static_identity included a closure that was dynamically generated on each invocation. They both now always use a function defined at module scope.
120 lines
3.3 KiB
Cython
120 lines
3.3 KiB
Cython
"""
|
|
Functions for ranking and sorting.
|
|
"""
|
|
cimport cython
|
|
from cpython cimport bool
|
|
from numpy cimport (
|
|
float64_t,
|
|
import_array,
|
|
intp_t,
|
|
ndarray,
|
|
NPY_DOUBLE,
|
|
NPY_MERGESORT,
|
|
PyArray_ArgSort,
|
|
PyArray_DIMS,
|
|
PyArray_EMPTY,
|
|
)
|
|
from numpy import apply_along_axis, float64, isnan, nan
|
|
from scipy.stats import rankdata
|
|
|
|
from zipline.utils.numpy_utils import (
|
|
is_float,
|
|
float64_dtype,
|
|
int64_dtype,
|
|
datetime64ns_dtype,
|
|
)
|
|
|
|
|
|
import_array()
|
|
|
|
|
|
cpdef is_missing(ndarray data, object missing_value):
|
|
"""
|
|
Generic is_missing function that handles quirks with NaN.
|
|
"""
|
|
if is_float(data) and isnan(missing_value):
|
|
return isnan(data)
|
|
return (data == missing_value)
|
|
|
|
|
|
def rankdata_1d_descending(ndarray data, str method):
|
|
"""
|
|
1D descending version of scipy.stats.rankdata.
|
|
"""
|
|
return rankdata(-(data.view(float64)), method=method)
|
|
|
|
|
|
def masked_rankdata_2d(ndarray data,
|
|
ndarray mask,
|
|
object missing_value,
|
|
str method,
|
|
bool ascending):
|
|
"""
|
|
Compute masked rankdata on data on float64, int64, or datetime64 data.
|
|
"""
|
|
cdef str dtype_name = data.dtype.name
|
|
if dtype_name not in ('float64', 'int64', 'datetime64[ns]'):
|
|
raise TypeError(
|
|
"Can't compute rankdata on array of dtype %r." % dtype_name
|
|
)
|
|
|
|
cdef ndarray missing_locations = (~mask | is_missing(data, missing_value))
|
|
|
|
# Interpret the bytes of integral data as floats for sorting.
|
|
data = data.copy().view(float64)
|
|
data[missing_locations] = nan
|
|
if not ascending:
|
|
data = -data
|
|
|
|
# OPTIMIZATION: Fast path the default case with our own specialized
|
|
# Cython implementation.
|
|
if method == 'ordinal':
|
|
result = rankdata_2d_ordinal(data)
|
|
else:
|
|
# FUTURE OPTIMIZATION:
|
|
# Write a less general "apply to rows" method that doesn't do all
|
|
# the extra work that apply_along_axis does.
|
|
result = apply_along_axis(rankdata, 1, data, method=method)
|
|
|
|
# On SciPy >= 0.17, rankdata returns integers for any method except
|
|
# average.
|
|
if result.dtype.name != 'float64':
|
|
result = result.astype('float64')
|
|
|
|
# rankdata will sort missing values into last place, but we want our nans
|
|
# to propagate, so explicitly re-apply.
|
|
result[missing_locations] = nan
|
|
return result
|
|
|
|
|
|
@cython.boundscheck(False)
|
|
@cython.wraparound(False)
|
|
@cython.embedsignature(True)
|
|
cpdef rankdata_2d_ordinal(ndarray[float64_t, ndim=2] array):
|
|
"""
|
|
Equivalent to:
|
|
|
|
numpy.apply_over_axis(scipy.stats.rankdata, 1, array, method='ordinal')
|
|
"""
|
|
cdef:
|
|
int nrows, ncols
|
|
ndarray[intp_t, ndim=2] sort_idxs
|
|
ndarray[float64_t, ndim=2] out
|
|
|
|
nrows = array.shape[0]
|
|
ncols = array.shape[1]
|
|
|
|
# scipy.stats.rankdata explicitly uses MERGESORT instead of QUICKSORT for
|
|
# the ordinal branch. c.f. commit ab21d2fee2d27daca0b2c161bbb7dba7e73e70ba
|
|
sort_idxs = PyArray_ArgSort(array, 1, NPY_MERGESORT)
|
|
|
|
# Roughly, "out = np.empty_like(array)"
|
|
out = PyArray_EMPTY(2, PyArray_DIMS(array), NPY_DOUBLE, False)
|
|
|
|
cdef intp_t i, j
|
|
for i in range(nrows):
|
|
for j in range(ncols):
|
|
out[i, sort_idxs[i, j]] = j + 1.0
|
|
|
|
return out
|