diff --git a/tests/modelling/test_term.py b/tests/modelling/test_term.py index 5720eabd..5a2ceb34 100644 --- a/tests/modelling/test_term.py +++ b/tests/modelling/test_term.py @@ -4,7 +4,6 @@ Tests for Term. from itertools import product from unittest import TestCase -from networkx import topological_sort from numpy import ( float32, uint32, @@ -20,8 +19,8 @@ from zipline.errors import ( TermInputsNotSpecified, WindowLengthNotSpecified, ) -from zipline.modelling.engine import build_dependency_graph from zipline.modelling.factor import Factor +from zipline.modelling.graph import TermGraph from zipline.modelling.expression import NUMEXPR_MATH_FUNCS @@ -69,6 +68,18 @@ def gen_equivalent_factors(): yield SomeFactorAlias() +def to_dict(l): + """ + Convert a list to a dict with keys drawn from '0', '1', '2', ... + + Example + ------- + >>> to_dict([2, 3, 4]) + {'0': 2, '1': 3, '2': 4} + """ + return dict(zip(map(str, range(len(l))), l)) + + class DependencyResolutionTestCase(TestCase): def setup(self): @@ -81,12 +92,9 @@ class DependencyResolutionTestCase(TestCase): """ Test dependency resolution for a single factor. """ - - build_dependency_graph([SomeFactor()]) - def check_output(graph): - resolution_order = topological_sort(graph) + resolution_order = list(graph.ordered()) self.assertEqual(len(resolution_order), 3) self.assertEqual( @@ -98,29 +106,29 @@ class DependencyResolutionTestCase(TestCase): self.assertEqual(graph.node[SomeDataSet.bar]['extra_rows'], 4) for foobar in gen_equivalent_factors(): - check_output(build_dependency_graph([foobar])) + check_output(TermGraph(to_dict([foobar]))) def test_single_factor_instance_args(self): """ Test dependency resolution for a single factor with arguments passed to the constructor. """ - graph = build_dependency_graph( - [SomeFactor([SomeDataSet.bar, SomeDataSet.buzz], window_length=5)] - ) - resolution_order = topological_sort(graph) + bar, buzz = SomeDataSet.bar, SomeDataSet.buzz + graph = TermGraph(to_dict([SomeFactor([bar, buzz], window_length=5)])) + + resolution_order = list(graph.ordered()) self.assertEqual(len(resolution_order), 3) self.assertEqual( set([resolution_order[0], resolution_order[1]]), - set([SomeDataSet.bar, SomeDataSet.buzz]), + set([bar, buzz]), ) self.assertEqual( resolution_order[-1], - SomeFactor([SomeDataSet.bar, SomeDataSet.buzz], window_length=5), + SomeFactor([bar, buzz], window_length=5), ) - self.assertEqual(graph.node[SomeDataSet.bar]['extra_rows'], 4) - self.assertEqual(graph.node[SomeDataSet.buzz]['extra_rows'], 4) + self.assertEqual(graph.extra_rows[bar], 4) + self.assertEqual(graph.extra_rows[buzz], 4) def test_reuse_atomic_terms(self): """ @@ -129,8 +137,8 @@ class DependencyResolutionTestCase(TestCase): f1 = SomeFactor([SomeDataSet.foo, SomeDataSet.bar]) f2 = SomeOtherFactor([SomeDataSet.bar, SomeDataSet.buzz]) - graph = build_dependency_graph([f1, f2]) - resolution_order = topological_sort(graph) + graph = TermGraph(to_dict([f1, f2])) + resolution_order = list(graph.ordered()) # bar should only appear once. self.assertEqual(len(resolution_order), 5) diff --git a/zipline/modelling/engine.py b/zipline/modelling/engine.py index d10d2309..eb57fe54 100644 --- a/zipline/modelling/engine.py +++ b/zipline/modelling/engine.py @@ -8,19 +8,14 @@ from abc import ( from operator import and_ from six import ( iteritems, + itervalues, with_metaclass, ) from six.moves import ( reduce, - zip, zip_longest, ) -from networkx import ( - DiGraph, - get_node_attributes, - topological_sort, -) from numpy import ( add, empty_like, @@ -33,81 +28,10 @@ from pandas import ( from zipline.lib.adjusted_array import ensure_ndarray from zipline.errors import NoFurtherDataError +from zipline.modelling.classifier import Classifier from zipline.modelling.factor import Factor from zipline.modelling.filter import Filter - - -# TODO: Move this somewhere else. -class CyclicDependency(Exception): - pass - - -def build_dependency_graph(terms): - """ - Build a dependency graph containing the given terms and their dependencies. - - Parameters - ---------- - terms : iterable - An iterable of zipline.modelling.term.Term. - - Returns - ------- - dependencies : networkx.DiGraph - A directed graph representing the dependencies of the desired inputs. - - Each node in the graph has an `extra_rows` attribute, indicating how - many, if any, extra rows we should compute for the node. Extra rows - are most often needed when a term is an input to a rolling window - computation. For example, if we compute a 30 day moving average of - price from day X to day Y, we need to load price data for the range - from day (X - 29) to day Y. - """ - dependencies = DiGraph() - parents = set() - for term in terms: - _add_to_graph(term, dependencies, parents, extra_rows=0) - # No parents should be left between top-level terms. - assert not parents - return dependencies - - -def _add_to_graph(term, - dependencies, - parents, - extra_rows): - """ - Add the term and all its inputs to dependencies. - """ - # If we've seen this node already as a parent of the current traversal, - # it means we have an unsatisifiable dependency. This should only be - # possible if the term's inputs are mutated after construction. - if term in parents: - raise CyclicDependency(term) - parents.add(term) - - try: - existing = dependencies.node[term] - except KeyError: - # We're not yet in the graph: add the term with the specified number of - # extra rows. - dependencies.add_node(term, extra_rows=extra_rows) - else: - # We're already in the graph because we've been traversed by - # another parent. Ensure that we have enough extra rows to satisfy - # all of our parents. - existing['extra_rows'] = max(extra_rows, existing['extra_rows']) - - for subterm in term.inputs: - _add_to_graph( - subterm, - dependencies, - parents, - extra_rows=extra_rows + term.extra_input_rows, - ) - dependencies.add_edge(subterm, term) - - parents.remove(term) +from zipline.modelling.graph import TermGraph class FFCEngine(with_metaclass(ABCMeta)): @@ -218,7 +142,7 @@ class SimpleFFCEngine(object): 5. Stick the values computed in (4) into a DataFrame and return it. - Step 0 is performed in `build_dependency_graph`. + Step 0 is performed by `zipline.modelling.graph.TermGraph`. Step 1 is performed in `self.build_lifetimes_matrix`. Step 2 is performed in `self.compute_chunk`. Steps 3, 4, and 5 are performed in self._format_factor_matrix. @@ -233,47 +157,36 @@ class SimpleFFCEngine(object): "start_date=%s, end_date=%s" % (start_date, end_date) ) - graph = build_dependency_graph(terms.values()) - ordered_terms = topological_sort(graph) - extra_row_counts = get_node_attributes(graph, 'extra_rows') - max_extra_rows = max(extra_row_counts.values()) + graph = TermGraph(terms) + max_extra_rows = graph.max_extra_rows lifetimes = self.build_lifetimes_matrix( start_date, end_date, max_extra_rows, ) - lifetimes_between_dates = lifetimes[max_extra_rows:] + raw_outputs = self.compute_chunk(graph, lifetimes, {}) + lifetimes_between_dates = lifetimes[max_extra_rows:] dates = lifetimes_between_dates.index.values assets = lifetimes_between_dates.columns.values - raw_outputs = self.compute_chunk( - ordered_terms, - extra_row_counts, - lifetimes, - ) - # We only need filters and factors to compute the final output matrix. - raw_filters = [lifetimes_between_dates.values] - raw_factors = [] - factor_names = [] + filters, factors = {}, {} for name, term in iteritems(terms): - extra = extra_row_counts[term] - if isinstance(term, Factor): - factor_names.append(name) - raw_factors.append(raw_outputs[term][extra:]) + if isinstance(term, Filter): + filters[name] = raw_outputs[name] + elif isinstance(term, Factor): + factors[name] = raw_outputs[name] + elif isinstance(term, Classifier): + continue + else: + raise ValueError("Unknown term type: %s" % term) - elif isinstance(term, Filter): - raw_filters.append(raw_outputs[term][extra:]) - - return self._format_factor_matrix( - dates, - assets, - raw_filters, - raw_factors, - factor_names, - ) + # Treat base_mask as an implicit filter. + # TODO: Is there a clean way to make this actually just be a filter? + filters['base'] = lifetimes_between_dates.values + return self._format_factor_matrix(dates, assets, filters, factors) def build_lifetimes_matrix(self, start_date, end_date, extra_rows): """ @@ -313,8 +226,8 @@ class SimpleFFCEngine(object): ), ) - # Build lifetimes matrix reaching back as far start_date plus - # max_extra_rows. + # Build lifetimes matrix reaching back to `extra_rows` days before + # `start_date.` lifetimes = finder.lifetimes( calendar[start_idx - extra_rows:end_idx] ) @@ -330,7 +243,7 @@ class SimpleFFCEngine(object): existed = lifetimes.iloc[extra_rows:].any() return lifetimes.loc[:, existed] - def _inputs_for_term(self, term, workspace, extra_row_counts): + def _inputs_for_term(self, term, workspace, extra_rows): """ Compute inputs for the given term. @@ -345,7 +258,7 @@ class SimpleFFCEngine(object): return [ workspace[input_].traverse( term.window_length, - offset=extra_row_counts[input_] - term_extra_rows + offset=extra_rows[input_] - term_extra_rows ) for input_ in term.inputs ] @@ -353,26 +266,45 @@ class SimpleFFCEngine(object): return [ ensure_ndarray( workspace[input_][ - extra_row_counts[input_] - term_extra_rows: + extra_rows[input_] - term_extra_rows: ], ) for input_ in term.inputs ] - def compute_chunk(self, ordered_terms, extra_row_counts, base_mask): + def compute_chunk(self, graph, base_mask, initial_workspace): """ - Compute the FFC terms in the graph based on the assets and dates - defined by base_mask. + Compute the FFC terms in the graph for the requested start and end + dates. - Returns a dictionary mapping terms to computed arrays. + Parameters + ---------- + graph : zipline.modelling.graph.TermGraph + + Returns + ------- + results : dict + Dictionary mapping requested results to outputs. """ loader = self._loader - max_extra_rows = max(extra_row_counts.values()) - workspace = {term: None for term in ordered_terms} - for term in ordered_terms: + extra_rows = graph.extra_rows + max_extra_rows = graph.max_extra_rows + + workspace = {} + if initial_workspace is not None: + workspace.update(initial_workspace) + + for term in graph.ordered(): + # Subclasses are allowed to pre-populate computed values for terms, + # and in the future we may pre-compute atomic terms coming from the + # same dataset. In both cases, it's possible that we already have + # an entry for this term. + if term in workspace: + continue + base_mask_for_term = base_mask.iloc[ - max_extra_rows - extra_row_counts[term]: + max_extra_rows - extra_rows[term]: ] if term.atomic: # FUTURE OPTIMIZATION: Scan the resolution order for terms in @@ -390,41 +322,39 @@ class SimpleFFCEngine(object): else: compute = term.compute_from_arrays workspace[term] = compute( - self._inputs_for_term(term, workspace, extra_row_counts), + self._inputs_for_term(term, workspace, extra_rows), base_mask_for_term, ) - return workspace - def _format_factor_matrix(self, - dates, - assets, - filter_data, - factor_data, - factor_names): + out = {} + for name, term in iteritems(graph.outputs): + # Truncate off extra rows from outputs. + out[name] = workspace[term][extra_rows[term]:] + return out + + def _format_factor_matrix(self, dates, assets, filters, factors): """ Convert raw computed filters/factors into a DataFrame for public APIs. Parameters ---------- dates : np.array[datetime64] - Index for raw data in filter_data/factor_data. + Row index for arrays in `filters` and `factors.` assets : np.array[int64] - Column labels for raw data in filter_data/factor_data. - filter_data : list[ndarray[bool]] - Raw filters data. - factor_data : list[ndarray] - Raw factor data. - factor_names : list[str] - Names of factors to use as keys. + Column index for arrays in `filters` and `factors.` + filters : dict + Dict mapping filter names -> computed filters. + factors : dict + Dict mapping factor names -> computed factors. Returns ------- factor_matrix : pd.DataFrame - A DataFrame with the following indices: + The indices of `factor_matrix` are as follows: - index : two-tiered MultiIndex of (date, asset). For each date, we - return a row for each asset that passed all filters on that - date. + index : two-tiered MultiIndex of (date, asset). + For each date, we return a row for each asset that passed all + filters on that date. columns : keys from `factor_data` Each date/asset/factor triple contains the computed value of the given @@ -433,19 +363,28 @@ class SimpleFFCEngine(object): # FUTURE OPTIMIZATION: Cythonize all of this. # Boolean mask of values that passed all filters. - unioned = reduce(and_, filter_data) + unioned = reduce(and_, itervalues(filters)) - # Parallel arrays of (x,y) coords for all date/asset pairs that passed + # Parallel arrays of (x,y) coords for (date, asset) pairs that passed # all filters. Each entry here will correspond to a row in our output # frame. nonzero_xs, nonzero_ys = unioned.nonzero() + # Raw arrays storing (date, asset) pairs. + # These will form the index of our output frame. raw_dates_index = empty_like(nonzero_xs, dtype='datetime64[ns]') raw_assets_index = empty_like(nonzero_xs, dtype=int) - factor_outputs = [ - empty_like(nonzero_xs, dtype=factor.dtype) - for factor in factor_data - ] + + # Mapping from column_name -> array. + # This will be the `data` arg to our output frame. + columns = { + name: empty_like(nonzero_xs, dtype=factor.dtype) + for name, factor in iteritems(factors) + } + # We're going to iterate over `iteritems(columns)` a whole bunch of + # times down below. It's faster to construct iterate over a tuple of + # pairs. + columns_iter = tuple(iteritems(columns)) # This is tricky. @@ -457,25 +396,25 @@ class SimpleFFCEngine(object): # each date, the running total of the number of assets that passed our # filters on or before that date. - # This means that (bounds[i - 1], bounds[i]) gives us the slice bounds - # of rows in our output DataFrame corresponding to each date. - dt_start = 0 + # This means that (bounds[i - 1], bounds[i]) gives us the indices of + # the first and last rows in our output frame for each date in `dates`. bounds = add.accumulate(unioned.sum(axis=1)) - for dt_idx, dt_end in enumerate(bounds): + day_start = 0 + for day_idx, day_end in enumerate(bounds): - row_bounds = slice(dt_start, dt_end) - column_indices = nonzero_ys[row_bounds] + day_bounds = slice(day_start, day_end) + column_indices = nonzero_ys[day_bounds] - raw_dates_index[row_bounds] = dates[dt_idx] - raw_assets_index[row_bounds] = assets[column_indices] - for computed, output in zip(factor_data, factor_outputs): - output[row_bounds] = computed[dt_idx, column_indices] + raw_dates_index[day_bounds] = dates[day_idx] + raw_assets_index[day_bounds] = assets[column_indices] + for name, colarray in columns_iter: + colarray[day_bounds] = factors[name][day_idx, column_indices] # Upper bound of current row becomes lower bound for next row. - dt_start = dt_end + day_start = day_end return DataFrame( - dict(zip(factor_names, factor_outputs)), + data=columns, index=MultiIndex.from_arrays( [ raw_dates_index, diff --git a/zipline/modelling/graph.py b/zipline/modelling/graph.py new file mode 100644 index 00000000..5ed0150b --- /dev/null +++ b/zipline/modelling/graph.py @@ -0,0 +1,116 @@ +""" +FFC-specific extensions to networkx.DiGraph +""" +from networkx import ( + DiGraph, + topological_sort, +) +from six import itervalues, iteritems + + +class CyclicDependency(Exception): + pass + + +class TermGraph(DiGraph): + """ + Graph represention of FFC Term dependencies + + Each node in the graph has an `extra_rows` attribute, indicating how many, + if any, extra rows we should compute for the node. Extra rows are most + often needed when a term is an input to a rolling window computation. For + example, if we compute a 30 day moving average of price from day X to day + Y, we need to load price data for the range from day (X - 29) to day Y. + + Parameters + ---------- + terms : dict + A dict mapping names to terms. + + Attributes + ---------- + outputs + extra_rows + max_extra_rows + + Methods + ------- + ordered() + Return a topologically-sorted iterator over the terms in self. + """ + def __init__(self, terms): + super(TermGraph, self).__init__(self) + parents = set() + for term in itervalues(terms): + self._add_to_graph(term, parents, extra_rows=0) + # No parents should be left between top-level terms. + assert not parents + + self._outputs = terms + self._extra_rows = { + term: attrs['extra_rows'] + for term, attrs in iteritems(self.node) + } + self._max_extra_rows = max(itervalues(self._extra_rows)) + self._ordered = topological_sort(self) + + @property + def outputs(self): + """ + Dict mapping names to designated output terms. + """ + return self._outputs + + @property + def extra_rows(self): + """ + Dict mapping term -> number of extra rows to compute for term. + """ + return self._extra_rows + + @property + def max_extra_rows(self): + """ + Maximum number of extra rows required to compute any term in the graph. + """ + return self._max_extra_rows + + def ordered(self): + """ + Return a topologically-sorted iterator over the terms in `self`. + """ + return iter(self._ordered) + + def _add_to_graph(self, term, parents, extra_rows): + """ + Add `term` and all its inputs to the graph. + """ + # If we've seen this node already as a parent of the current traversal, + # it means we have an unsatisifiable dependency. This should only be + # possible if the term's inputs are mutated after construction. + if term in parents: + raise CyclicDependency(term) + parents.add(term) + + try: + existing = self.node[term] + except KeyError: + # `term` is not yet in the graph: add it with the specified number + # of extra rows. + self.add_node(term, extra_rows=extra_rows) + else: + # `term` is already in the graph because we've been traversed by + # another parent. Ensure that we have enough extra rows to satisfy + # all of our parents. + existing['extra_rows'] = max(extra_rows, existing['extra_rows']) + + extra_rows_for_subterms = extra_rows + term.extra_input_rows + for subterm in term.inputs: + self._add_to_graph( + subterm, + parents, + extra_rows=extra_rows_for_subterms + ) + self.add_edge(subterm, term) + + parents.remove(term)