diff --git a/setup.py b/setup.py index 1bcca359..ad3966f6 100644 --- a/setup.py +++ b/setup.py @@ -62,6 +62,7 @@ ext_modules = LazyCythonizingList([ ('zipline.assets._assets', ['zipline/assets/_assets.pyx']), ('zipline.lib.adjusted_array', ['zipline/lib/adjusted_array.pyx']), ('zipline.lib.adjustment', ['zipline/lib/adjustment.pyx']), + ('zipline.lib.rank', ['zipline/lib/rank.pyx']), ( 'zipline.data.ffc.loaders._us_equity_pricing', ['zipline/data/ffc/loaders/_us_equity_pricing.pyx'] diff --git a/tests/modelling/base.py b/tests/modelling/base.py new file mode 100644 index 00000000..adc757bc --- /dev/null +++ b/tests/modelling/base.py @@ -0,0 +1,114 @@ +""" +Base class for FFC unit tests. +""" +from functools import wraps +from unittest import TestCase + +from numpy import arange, prod +from numpy.random import randn, seed as random_seed +from pandas import date_range, Int64Index, DataFrame +from six import iteritems + +from zipline.assets import AssetFinder +from zipline.modelling.engine import SimpleFFCEngine +from zipline.modelling.graph import TermGraph +from zipline.utils.test_utils import make_simple_asset_info, ExplodingObject +from zipline.utils.tradingcalendar import trading_day + + +def with_defaults(**default_funcs): + """ + Decorator for providing dynamic default values for a method. + + Usages: + + @with_defaults(foo=lambda self: self.x + self.y) + def func(self, foo): + ... + + If a value is passed for `foo`, it will be used. Otherwise the function + supplied to `with_defaults` will be called with `self` as an argument. + """ + def decorator(f): + @wraps(f) + def method(self, *args, **kwargs): + for name, func in iteritems(default_funcs): + if name not in kwargs: + kwargs[name] = func(self) + return f(self, *args, **kwargs) + return method + return decorator + + +with_default_shape = with_defaults(shape=lambda self: self.default_shape) + + +class BaseFFCTestCase(TestCase): + + def setUp(self): + self.__calendar = date_range('2014', '2015', freq=trading_day) + self.__assets = assets = Int64Index(arange(1, 20)) + self.__finder = AssetFinder( + make_simple_asset_info( + assets, + self.__calendar[0], + self.__calendar[-1], + ), + db_path=':memory:', + create_table=True, + ) + self.__mask = self.__finder.lifetimes(self.__calendar[-10:]) + + @property + def default_shape(self): + """Default shape for methods that build test data.""" + return self.__mask.shape + + def run_terms(self, terms, initial_workspace, mask=None): + """ + Compute the given terms, seeding the workspace of our FFCEngine with + `initial_workspace`. + + Parameters + ---------- + terms : dict + Mapping from termname -> term object. + + Returns + ------- + results : dict + Mapping from termname -> computed result. + """ + engine = SimpleFFCEngine( + ExplodingObject(), + self.__calendar, + self.__finder, + ) + mask = mask if mask is not None else self.__mask + return engine.compute_chunk(TermGraph(terms), mask, initial_workspace) + + def build_mask(self, array): + ndates, nassets = array.shape + return DataFrame( + array, + # Use the **last** N dates rather than the first N so that we have + # space for lookbacks. + index=self.__calendar[-ndates:], + columns=self.__assets[:nassets], + dtype=bool, + ) + + @with_default_shape + def arange_data(self, shape, dtype=float): + """ + Build a block of testing data from numpy.arange. + """ + return arange(prod(shape), dtype=dtype).reshape(shape) + + @with_default_shape + def randn_data(self, seed, shape): + """ + Build a block of testing data from numpy.random.randn. + """ + random_seed(seed) + return randn(*shape) diff --git a/tests/modelling/test_factor.py b/tests/modelling/test_factor.py index 5286d5b2..8b0f14ad 100644 --- a/tests/modelling/test_factor.py +++ b/tests/modelling/test_factor.py @@ -1,21 +1,16 @@ """ Tests for Factor terms. """ -from unittest import TestCase - from numpy import ( array, + ones, ) -from numpy.testing import assert_array_equal -from pandas import ( - DataFrame, - date_range, - Int64Index, -) -from six import iteritems from zipline.errors import UnknownRankMethod from zipline.modelling.factor import TestingFactor +from zipline.utils.test_utils import check_arrays + +from .base import BaseFFCTestCase class F(TestingFactor): @@ -23,23 +18,17 @@ class F(TestingFactor): window_length = 0 -class FactorTestCase(TestCase): +class FactorTestCase(BaseFFCTestCase): def setUp(self): + super(FactorTestCase, self).setUp() self.f = F() - self.dates = date_range('2014-01-01', periods=5, freq='D') - self.assets = Int64Index(range(5)) - self.mask = DataFrame(True, index=self.dates, columns=self.assets) - - def tearDown(self): - pass def test_bad_input(self): - with self.assertRaises(UnknownRankMethod): self.f.rank("not a real rank method") - def test_rank(self): + def test_rank_ascending(self): # Generated with: # data = arange(25).reshape(5, 5).transpose() % 4 @@ -47,7 +36,7 @@ class FactorTestCase(TestCase): [1, 2, 3, 0, 1], [2, 3, 0, 1, 2], [3, 0, 1, 2, 3], - [0, 1, 2, 3, 0]]) + [0, 1, 2, 3, 0]], dtype=float) expected_ranks = { 'ordinal': array([[1., 3., 4., 5., 2.], [2., 4., 5., 1., 3.], @@ -76,14 +65,73 @@ class FactorTestCase(TestCase): [1., 2., 3., 4., 1.]]), } - # Test with the default, which should be 'ordinal'. - default_result = self.f.rank().compute_from_arrays([data], self.mask) - assert_array_equal(default_result, expected_ranks['ordinal']) - - # Test with each method passed explicitly. - for method, expected_result in iteritems(expected_ranks): - result = self.f.rank(method=method).compute_from_arrays( - [data], - self.mask, + def check(terms): + results = self.run_terms( + terms, + initial_workspace={self.f: data}, + mask=self.build_mask(ones((5, 5))), ) - assert_array_equal(result, expected_ranks[method]) + for method in terms: + check_arrays(results[method], expected_ranks[method]) + + check({meth: self.f.rank(method=meth) for meth in expected_ranks}) + check({ + meth: self.f.rank(method=meth, ascending=True) + for meth in expected_ranks + }) + # Not passing a method should default to ordinal. + check({'ordinal': self.f.rank()}) + check({'ordinal': self.f.rank(ascending=True)}) + + def test_rank_descending(self): + + # Generated with: + # data = arange(25).reshape(5, 5).transpose() % 4 + data = array([[0, 1, 2, 3, 0], + [1, 2, 3, 0, 1], + [2, 3, 0, 1, 2], + [3, 0, 1, 2, 3], + [0, 1, 2, 3, 0]], dtype=float) + expected_ranks = { + 'ordinal': array([[4., 3., 2., 1., 5.], + [3., 2., 1., 5., 4.], + [2., 1., 5., 4., 3.], + [1., 5., 4., 3., 2.], + [4., 3., 2., 1., 5.]]), + 'average': array([[4.5, 3., 2., 1., 4.5], + [3.5, 2., 1., 5., 3.5], + [2.5, 1., 5., 4., 2.5], + [1.5, 5., 4., 3., 1.5], + [4.5, 3., 2., 1., 4.5]]), + 'min': array([[4., 3., 2., 1., 4.], + [3., 2., 1., 5., 3.], + [2., 1., 5., 4., 2.], + [1., 5., 4., 3., 1.], + [4., 3., 2., 1., 4.]]), + 'max': array([[5., 3., 2., 1., 5.], + [4., 2., 1., 5., 4.], + [3., 1., 5., 4., 3.], + [2., 5., 4., 3., 2.], + [5., 3., 2., 1., 5.]]), + 'dense': array([[4., 3., 2., 1., 4.], + [3., 2., 1., 4., 3.], + [2., 1., 4., 3., 2.], + [1., 4., 3., 2., 1.], + [4., 3., 2., 1., 4.]]), + } + + def check(terms): + results = self.run_terms( + terms, + initial_workspace={self.f: data}, + mask=self.build_mask(ones((5, 5))), + ) + for method in terms: + check_arrays(results[method], expected_ranks[method]) + + check({ + meth: self.f.rank(method=meth, ascending=False) + for meth in expected_ranks + }) + # Not passing a method should default to ordinal. + check({'ordinal': self.f.rank(ascending=False)}) diff --git a/tests/modelling/test_filter.py b/tests/modelling/test_filter.py index 780e3ce0..00085e74 100644 --- a/tests/modelling/test_filter.py +++ b/tests/modelling/test_filter.py @@ -1,28 +1,54 @@ """ Tests for filter terms. """ -from unittest import TestCase +from operator import and_ from numpy import ( arange, + argsort, array, eye, float64, + full_like, nan, nanpercentile, - ones_like, + ones, putmask, ) -from numpy.testing import assert_array_equal - -from pandas import ( - DataFrame, - date_range, - Int64Index, -) from zipline.errors import BadPercentileBounds from zipline.modelling.factor import TestingFactor +from zipline.utils.test_utils import check_arrays + +from .base import BaseFFCTestCase + + +def rowwise_rank(array): + """ + Take a 2D array and return the 0-indexed sorted position of each element in + the array for each row. + + Example + ------- + In [5]: data + Out[5]: + array([[-0.141, -1.103, -1.0171, 0.7812, 0.07 ], + [ 0.926, 0.235, -0.7698, 1.4552, 0.2061], + [ 1.579, 0.929, -0.557 , 0.7896, -1.6279], + [-1.362, -2.411, -1.4604, 1.4468, -0.1885], + [ 1.272, 1.199, -3.2312, -0.5511, -1.9794]]) + + In [7]: argsort(argsort(data)) + Out[7]: + array([[2, 0, 1, 4, 3], + [3, 2, 0, 4, 1], + [4, 3, 1, 2, 0], + [2, 0, 1, 4, 3], + [4, 3, 0, 2, 1]]) + """ + # note that unlike scipy.stats.rankdata, the output here is 0-indexed, not + # 1-indexed. + return argsort(argsort(array)) class SomeFactor(TestingFactor): @@ -30,23 +56,17 @@ class SomeFactor(TestingFactor): window_length = 0 -class FilterTestCase(TestCase): +class SomeOtherFactor(TestingFactor): + inputs = () + window_length = 0 + + +class FilterTestCase(BaseFFCTestCase): def setUp(self): + super(FilterTestCase, self).setUp() self.f = SomeFactor() - self.dates = date_range('2014-01-01', periods=5, freq='D') - self.assets = Int64Index(range(5)) - self.mask = DataFrame(True, index=self.dates, columns=self.assets) - - def tearDown(self): - pass - - def maskframe(self, array): - return DataFrame( - array, - index=date_range('2014-01-01', periods=array.shape[0], freq='D'), - columns=arange(array.shape[1]), - ) + self.g = SomeOtherFactor() def test_bad_percentiles(self): f = self.f @@ -61,156 +81,214 @@ class FilterTestCase(TestCase): with self.assertRaises(BadPercentileBounds): f.percentile_between(min_, max_) - def test_rank_percentile_nice_partitions(self): - # Test case with nicely-defined partitions. + def test_top(self): + counts = 2, 3, 10 + data = self.randn_data(seed=5) # Arbitrary seed choice. + results = self.run_terms( + terms={'top_' + str(c): self.f.top(c) for c in counts}, + initial_workspace={self.f: data}, + ) + for c in counts: + result = results['top_' + str(c)] + + # Check that `min(c, num_assets)` assets passed each day. + passed_per_day = result.sum(axis=1) + check_arrays( + passed_per_day, + full_like(passed_per_day, min(c, data.shape[1])), + ) + + # Check that the top `c` assets passed. + expected = rowwise_rank(-data) < c + check_arrays(result, expected) + + def test_bottom(self): + counts = 2, 3, 10 + data = self.randn_data(seed=5) # Arbitrary seed choice. + results = self.run_terms( + terms={'bottom_' + str(c): self.f.bottom(c) for c in counts}, + initial_workspace={self.f: data}, + ) + for c in counts: + result = results['bottom_' + str(c)] + + # Check that `min(c, num_assets)` assets passed each day. + passed_per_day = result.sum(axis=1) + check_arrays( + passed_per_day, + full_like(passed_per_day, min(c, data.shape[1])), + ) + + # Check that the bottom `c` assets passed. + expected = rowwise_rank(data) < c + check_arrays(result, expected) + + def test_percentile_between(self): + + quintiles = range(5) + filter_names = ['pct_' + str(q) for q in quintiles] + iter_quintiles = zip(filter_names, quintiles) + + terms = { + name: self.f.percentile_between(q * 20.0, (q + 1) * 20.0) + for name, q in zip(filter_names, quintiles) + } + + # Test with 5 columns and no NaNs. eye5 = eye(5, dtype=float64) - eye6 = eye(6, dtype=float64) - nanmask = array([[0, 0, 0, 0, 0, 1], - [1, 0, 0, 0, 0, 0], - [0, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0], - [0, 0, 0, 1, 0, 0], - [0, 0, 0, 0, 1, 0]], dtype=bool) - nandata = eye6.copy() - putmask(nandata, nanmask, nan) - - for quintile in range(5): - factor = self.f.percentile_between( - quintile * 20.0, - (quintile + 1) * 20.0, - ) - # Test w/o any NaNs - result = factor.compute_from_arrays( - [eye5], - self.maskframe(ones_like(eye5, dtype=bool)), - ) - # Test with NaNs in the data. - nandata_result = factor.compute_from_arrays( - [nandata], - self.maskframe(ones_like(nandata, dtype=bool)), - ) - # Test with Falses in the mask. - nanmask_result = factor.compute_from_arrays( - [eye6], - self.maskframe(~nanmask), - ) - - assert_array_equal(nandata_result, nanmask_result) - + results = self.run_terms( + terms, + initial_workspace={self.f: eye5}, + mask=self.build_mask(ones((5, 5))), + ) + for name, quintile in iter_quintiles: + result = results[name] if quintile < 4: - # There are 4 0s and one 1 in each row, so the first 4 + # There are four 0s and one 1 in each row, so the first 4 # quintiles should be all the locations with zeros in the input # array. - assert_array_equal(result, ~eye5.astype(bool)) - # Should reject all the ones, plus the nans. - assert_array_equal( - nandata_result, - ~(nanmask | eye6.astype(bool)) - ) - + check_arrays(result, ~eye5.astype(bool)) else: - # The last quintile should contain all the 1s. - assert_array_equal(result, eye(5, dtype=bool)) - # Should accept all the 1s. - assert_array_equal(nandata_result, eye(6, dtype=bool)) + # The top quintile should match the sole 1 in each row. + check_arrays(result, eye5.astype(bool)) - def test_rank_percentile_nasty_partitions(self): - # Test case with nasty partitions: divide up 5 assets into quartiles. - data = arange(25, dtype=float).reshape(5, 5) % 4 - nandata = data.copy() - nandata[eye(5, dtype=bool)] = nan - for quartile in range(4): - lower_bound = quartile * 25.0 - upper_bound = (quartile + 1) * 25.0 - factor = self.f.percentile_between(lower_bound, upper_bound) + # Test with 6 columns, no NaNs, and one masked entry per day. + eye6 = eye(6, dtype=float64) + mask = array([[1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1], + [1, 0, 1, 1, 1, 1], + [1, 1, 0, 1, 1, 1], + [1, 1, 1, 0, 1, 1], + [1, 1, 1, 1, 0, 1]], dtype=bool) - # There isn't a nice definition of correct behavior here, so for - # now we guarantee the behavior of numpy.nanpercentile. - - result = factor.compute_from_arrays([data], self.mask) - min_value = nanpercentile(data, lower_bound, axis=1, keepdims=True) - max_value = nanpercentile(data, upper_bound, axis=1, keepdims=True) - assert_array_equal( - result, - (min_value <= data) & (data <= max_value), - ) - - nanresult = factor.compute_from_arrays([nandata], self.mask) - min_value = nanpercentile( - nandata, - lower_bound, - axis=1, - keepdims=True, - ) - max_value = nanpercentile( - nandata, - upper_bound, - axis=1, - keepdims=True, - ) - assert_array_equal( - nanresult, - (min_value <= nandata) & (nandata <= max_value), - ) - - def test_sequenced_filter(self): - first = SomeFactor() < 1 - first_input = eye(5) - first_result = first.compute_from_arrays([first_input], self.mask) - assert_array_equal(first_result, ~eye(5, dtype=bool)) - - # Second should pick out the fourth column. - second = SomeFactor().eq(3.0) - second_input = arange(25, dtype=float).reshape(5, 5) % 5 - - sequenced = first.then(second) - - result = sequenced.compute_from_arrays( - [first_result, second_input], - self.mask, + results = self.run_terms( + terms, + initial_workspace={self.f: eye6}, + mask=self.build_mask(mask) ) - expected_result = (first_result & (second_input == 3.0)) - assert_array_equal(result, expected_result) + for name, quintile in iter_quintiles: + result = results[name] + if quintile < 4: + # Should keep all values that were 0 in the base data and were + # 1 in the mask. + check_arrays(result, mask & (~eye6.astype(bool))), + else: + # Should keep all the 1s in the base data. + check_arrays(result, eye6.astype(bool)) + + # Test with 6 columns, no mask, and one NaN per day. Should have the + # same outcome as if we had masked the NaNs. + # In particular, the NaNs should never pass any filters. + eye6_withnans = eye6.copy() + putmask(eye6_withnans, ~mask, nan) + results = self.run_terms( + terms, + initial_workspace={self.f: eye6}, + mask=self.build_mask(mask) + ) + for name, quintile in iter_quintiles: + result = results[name] + if quintile < 4: + # Should keep all values that were 0 in the base data and were + # 1 in the mask. + check_arrays(result, mask & (~eye6.astype(bool))), + else: + # Should keep all the 1s in the base data. + check_arrays(result, eye6.astype(bool)) + + def test_percentile_nasty_partitions(self): + # Test percentile with nasty partitions: divide up 5 assets into + # quartiles. + # There isn't a nice mathematical definition of correct behavior here, + # so for now we guarantee the behavior of numpy.nanpercentile. This is + # mostly for regression testing in case we write our own specialized + # percentile calculation at some point in the future. + + data = arange(25, dtype=float).reshape(5, 5) % 4 + quartiles = range(4) + filter_names = ['pct_' + str(q) for q in quartiles] + + terms = { + name: self.f.percentile_between(q * 25.0, (q + 1) * 25.0) + for name, q in zip(filter_names, quartiles) + } + results = self.run_terms( + terms, + initial_workspace={self.f: data}, + mask=self.build_mask(ones((5, 5))), + ) + + for name, quartile in zip(filter_names, quartiles): + result = results[name] + lower = quartile * 25.0 + upper = (quartile + 1) * 25.0 + expected = and_( + nanpercentile(data, lower, axis=1, keepdims=True) <= data, + data <= nanpercentile(data, upper, axis=1, keepdims=True), + ) + check_arrays(result, expected) + + def test_sequenced_filter_order_independent(self): + data = self.arange_data() % 5 + results = self.run_terms( + { + # Sequencing is equivalent to &ing for commutative filters. + 'sequenced': (1.5 < self.f).then(self.f < 3.5), + 'anded': (1.5 < self.f) & (self.f < 3.5), + }, + initial_workspace={self.f: data}, + ) + expected = (1.5 < data) & (data < 3.5) + + check_arrays(results['sequenced'], expected) + check_arrays(results['anded'], expected) def test_sequenced_filter_order_dependent(self): - f = SomeFactor() < 1 + + first = self.f < 1 f_input = eye(5) - f_result = f.compute_from_arrays([f_input], self.mask) - assert_array_equal(f_result, ~eye(5, dtype=bool)) - g = SomeFactor().percentile_between(80, 100) - g_input = arange(25, dtype=float).reshape(5, 5) % 5 - g_result = g.compute_from_arrays([g_input], self.mask) - assert_array_equal(g_result, g_input == 4) + second = self.g.percentile_between(80, 100) + g_input = arange(25, dtype=float).reshape(5, 5) - result = f.then(g).compute_from_arrays( - [f_result, g_input], - self.mask, + initial_mask = self.build_mask(ones((5, 5))) + + terms = { + 'first': first, + 'second': second, + 'sequenced': first.then(second), + } + + results = self.run_terms( + terms, + initial_workspace={self.f: f_input, self.g: g_input}, + mask=initial_mask, ) - # Input data is strictly increasing, so the result should be the top - # value not filtered by first. - expected_result = array( + + # First should pass everything but the diagonal. + check_arrays(results['first'], ~eye(5, dtype=bool)) + + # Second should pass the largest value each day. Each row is strictly + # increasing, so we always select the last value. + expected_second = array( [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], - [0, 0, 0, 1, 0]], + [0, 0, 0, 0, 1]], dtype=bool, ) - assert_array_equal(result, expected_result) + check_arrays(results['second'], expected_second) - result = g.then(f).compute_from_arrays( - [g_result, f_input], - self.mask, - ) - - # Percentile calculated first, then diagonal is removed. - expected_result = array( + # When sequencing, we should remove the diagonal as an option before + # computing percentiles. On the last day, we should get the + # second-largest value, rather than the largest. + expected_sequenced = array( [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], - [0, 0, 0, 0, 0]], + [0, 0, 0, 1, 0]], # Different from previous! dtype=bool, ) - assert_array_equal(result, expected_result) + check_arrays(results['sequenced'], expected_sequenced) diff --git a/zipline/lib/rank.pyx b/zipline/lib/rank.pyx new file mode 100644 index 00000000..911e695b --- /dev/null +++ b/zipline/lib/rank.pyx @@ -0,0 +1,53 @@ +""" +Functions for ranking and sorting. +""" +cimport cython +from numpy cimport ( + float64_t, + import_array, + intp_t, + ndarray, + NPY_DOUBLE, + NPY_MERGESORT, + PyArray_ArgSort, + PyArray_DIMS, + PyArray_EMPTY, +) +from numpy import nan + +import_array() + + +cdef double NAN = nan + + +@cython.boundscheck(False) +@cython.wraparound(False) +@cython.embedsignature(True) +def 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 diff --git a/zipline/modelling/engine.py b/zipline/modelling/engine.py index eb57fe54..82cbba55 100644 --- a/zipline/modelling/engine.py +++ b/zipline/modelling/engine.py @@ -325,6 +325,7 @@ class SimpleFFCEngine(object): self._inputs_for_term(term, workspace, extra_rows), base_mask_for_term, ) + assert(workspace[term].shape == base_mask_for_term.shape) out = {} for name, term in iteritems(graph.outputs): diff --git a/zipline/modelling/factor/factor.py b/zipline/modelling/factor/factor.py index a6679f8f..105925bc 100644 --- a/zipline/modelling/factor/factor.py +++ b/zipline/modelling/factor/factor.py @@ -13,6 +13,7 @@ from zipline.errors import ( UnknownRankMethod, UnsupportedDataType, ) +from zipline.lib.rank import rankdata_2d_ordinal from zipline.modelling.term import ( CustomTermMixin, RequiredWindowLengthMixin, @@ -219,22 +220,27 @@ class Factor(Term): eq = binary_operator('==') - def rank(self, method='ordinal'): + def rank(self, method='ordinal', ascending=True): """ Construct a new Factor representing the sorted rank of each column within each row. - Returns - ------- + Parameters + ---------- ranks : zipline.modelling.factor.Rank A new factor that will compute the sorted indices of the data produced by `self`. method : str, {'ordinal', 'min', 'max', 'dense', 'average'} - The method used to assign ranks to tied elements. Default is - 'ordinal'. See `scipy.stats.rankdata` for a full description of - the semantics for each ranking method. + The method used to assign ranks to tied elements. See + `scipy.stats.rankdata` for a full description of the semantics for + each ranking method. Default is 'ordinal'. + ascending : bool, optional + Whether to return sorted rank in ascending or descending order. + Default is True. - The default is 'ordinal'. + Returns + ------- + ranks : zipline.modelling.factor.Rank Notes ----- @@ -247,10 +253,41 @@ class Factor(Term): See Also -------- - scipy.stats.rankdata : Underlying ranking algorithm. - zipline.modelling.factor.Rank : Class implementing core functionality. + scipy.stats.rankdata + zipline.lib.rank + zipline.modelling.factor.Rank """ - return Rank(self, method=method) + return Rank(self if ascending else -self, method=method) + + def top(self, N): + """ + Construct a Filter matching the top N asset values of self each day. + + Parameters + ---------- + N : int + Number of assets passing the returned filter each day. + + Returns + ------- + filter : zipline.modelling.filter.Filter + """ + return self.rank(ascending=False) <= N + + def bottom(self, N): + """ + Construct a Filter matching the bottom N asset values of self each day. + + Parameters + ---------- + N : int + Number of assets passing the returned filter each day. + + Returns + ------- + filter : zipline.modelling.filter.Filter + """ + return self.rank(ascending=True) <= N def percentile_between(self, min_percentile, max_percentile): """ @@ -361,22 +398,23 @@ class Rank(SingleInputMixin, Factor): For each row in the input, compute a like-shaped array of per-row ranks. """ - # FUTURE OPTIMIZATION: - # Write a less general `apply_to_rows` method in - # Cython that doesn't do all the extra work that apply_over_axis does. - - # FUTURE OPTIMIZATION: - # Look at bottleneck.nanrankdata, which is ~30% faster than numpy here, - # and does what we want with NaNs, but doesn't support `method`. - result = apply_along_axis( - rankdata, - 1, - arrays[0], - method=self._method, - ) - # rankdata will sort nan values into last place, but we want our nans - # to propagate, so explicitly re-apply - result[~mask.values] = nan + # OPTIMIZATION: Fast path the default value with our own specialized + # implementation. + if self._method == 'ordinal': + result = rankdata_2d_ordinal(arrays[0]) + 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, + arrays[0], + method=self._method, + ) + # rankdata will sort nan values into last place, but we want our + # nans to propagate, so explicitly re-apply + result[~mask.values] = nan return result def __repr__(self): diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py index 202c62b5..abf2d0c0 100644 --- a/zipline/utils/test_utils.py +++ b/zipline/utils/test_utils.py @@ -322,3 +322,18 @@ def check_arrays(left, right, err_msg='', verbose=True): if type(left) != type(right): raise AssertionError("%s != %s" % (type(left), type(right))) return assert_array_equal(left, right, err_msg=err_msg, verbose=True) + + +class UnexpectedAttributeAccess(Exception): + pass + + +class ExplodingObject(object): + """ + Object that will raise an exception on any attribute access. + + Useful for verifying that an object is never touched during a + function/method call. + """ + def __getattribute__(self, name): + raise UnexpectedAttributeAccess(name)