mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-29 20:33:46 +08:00
Scott Sanderson
6e8a4b8144
- Add an `ascending=True` keyword to `rank()`. - Add `top(N)` and `bottom(N)` methods to Factor. These return Filters that pass the top and bottom N elements each day. - Add a slightly faster path for rank(method='ordinal'). I had originally thought the fast path was 2-3x faster because I had my benchmark data axes flipped. The actual speedup is only 5-10%, which means it probably wasn't worth the effort to Cythonize...but we have a slightly faster version now so we might as well use it. - Refactor test_filter and test_factor to make it easier to implement and test transformations on factors. These tests now subclass BaseFFCTestCase, which provides facilities for passing a dict of terms and an "initial_workspace", the values for which are used by SimpleFFCEngine rather than needing to manually manage the inputs and outputs of each term.
138 lines
5.0 KiB
Python
138 lines
5.0 KiB
Python
"""
|
|
Tests for Factor terms.
|
|
"""
|
|
from numpy import (
|
|
array,
|
|
ones,
|
|
)
|
|
|
|
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):
|
|
inputs = ()
|
|
window_length = 0
|
|
|
|
|
|
class FactorTestCase(BaseFFCTestCase):
|
|
|
|
def setUp(self):
|
|
super(FactorTestCase, self).setUp()
|
|
self.f = F()
|
|
|
|
def test_bad_input(self):
|
|
with self.assertRaises(UnknownRankMethod):
|
|
self.f.rank("not a real rank method")
|
|
|
|
def test_rank_ascending(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([[1., 3., 4., 5., 2.],
|
|
[2., 4., 5., 1., 3.],
|
|
[3., 5., 1., 2., 4.],
|
|
[4., 1., 2., 3., 5.],
|
|
[1., 3., 4., 5., 2.]]),
|
|
'average': array([[1.5, 3., 4., 5., 1.5],
|
|
[2.5, 4., 5., 1., 2.5],
|
|
[3.5, 5., 1., 2., 3.5],
|
|
[4.5, 1., 2., 3., 4.5],
|
|
[1.5, 3., 4., 5., 1.5]]),
|
|
'min': array([[1., 3., 4., 5., 1.],
|
|
[2., 4., 5., 1., 2.],
|
|
[3., 5., 1., 2., 3.],
|
|
[4., 1., 2., 3., 4.],
|
|
[1., 3., 4., 5., 1.]]),
|
|
'max': array([[2., 3., 4., 5., 2.],
|
|
[3., 4., 5., 1., 3.],
|
|
[4., 5., 1., 2., 4.],
|
|
[5., 1., 2., 3., 5.],
|
|
[2., 3., 4., 5., 2.]]),
|
|
'dense': array([[1., 2., 3., 4., 1.],
|
|
[2., 3., 4., 1., 2.],
|
|
[3., 4., 1., 2., 3.],
|
|
[4., 1., 2., 3., 4.],
|
|
[1., 2., 3., 4., 1.]]),
|
|
}
|
|
|
|
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) 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)})
|