mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-11 04:24:13 +08:00
dmichalowicz
246 lines
8.1 KiB
Python
246 lines
8.1 KiB
Python
"""
|
|
Dependency-Graph representation of Pipeline API terms.
|
|
"""
|
|
from networkx import (
|
|
DiGraph,
|
|
topological_sort,
|
|
)
|
|
from six import itervalues, iteritems
|
|
from zipline.utils.memoize import lazyval
|
|
from zipline.pipeline.visualize import display_graph
|
|
|
|
from .term import LoadableTerm
|
|
|
|
|
|
class CyclicDependency(Exception):
|
|
pass
|
|
|
|
|
|
class TermGraph(DiGraph):
|
|
"""
|
|
Graph represention of Pipeline 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 final output terms.
|
|
|
|
Attributes
|
|
----------
|
|
outputs
|
|
offset
|
|
extra_rows
|
|
|
|
Methods
|
|
-------
|
|
ordered()
|
|
Return a topologically-sorted iterator over the terms in self.
|
|
"""
|
|
def __init__(self, terms):
|
|
super(TermGraph, self).__init__(self)
|
|
|
|
self._frozen = False
|
|
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._ordered = topological_sort(self)
|
|
|
|
# Mark that no more terms should be added to the graph.
|
|
self._frozen = True
|
|
|
|
@lazyval
|
|
def offset(self):
|
|
"""
|
|
For all pairs (term, input) such that `input` is an input to `term`,
|
|
compute a mapping::
|
|
|
|
(term, input) -> offset(term, input)
|
|
|
|
where ``offset(term, input)`` is the number of rows that ``term``
|
|
should truncate off the raw array produced for ``input`` before using
|
|
it. We compute this value as follows::
|
|
|
|
offset(term, input) = (extra_rows_computed(input)
|
|
- extra_rows_computed(term)
|
|
- requested_extra_rows(term, input))
|
|
Examples
|
|
--------
|
|
|
|
Case 1
|
|
~~~~~~
|
|
|
|
Factor A needs 5 extra rows of USEquityPricing.close, and Factor B
|
|
needs 3 extra rows of the same. Factor A also requires 5 extra rows of
|
|
USEquityPricing.high, which no other Factor uses. We don't require any
|
|
extra rows of Factor A or Factor B
|
|
|
|
We load 5 extra rows of both `price` and `high` to ensure we can
|
|
service Factor A, and the following offsets get computed::
|
|
|
|
offset[Factor A, USEquityPricing.close] == (5 - 0) - 5 == 0
|
|
offset[Factor A, USEquityPricing.high] == (5 - 0) - 5 == 0
|
|
offset[Factor B, USEquityPricing.close] == (5 - 0) - 3 == 2
|
|
offset[Factor B, USEquityPricing.high] raises KeyError.
|
|
|
|
Case 2
|
|
~~~~~~
|
|
|
|
Factor A needs 5 extra rows of USEquityPricing.close, and Factor B
|
|
needs 3 extra rows of Factor A, and Factor B needs 2 extra rows of
|
|
USEquityPricing.close.
|
|
|
|
We load 8 extra rows of USEquityPricing.close (enough to load 5 extra
|
|
rows of Factor A), and the following offsets get computed::
|
|
|
|
offset[Factor A, USEquityPricing.close] == (8 - 3) - 5 == 0
|
|
offset[Factor B, USEquityPricing.close] == (8 - 0) - 2 == 6
|
|
offset[Factor B, Factor A] == (3 - 0) - 3 == 0
|
|
|
|
Notes
|
|
-----
|
|
`offset(term, input) >= 0` for all valid pairs, since `input` must be
|
|
an input to `term` if the pair appears in the mapping.
|
|
|
|
This value is useful because we load enough rows of each input to serve
|
|
all possible dependencies. However, for any given dependency, we only
|
|
want to compute using the actual number of required extra rows for that
|
|
dependency. We can do so by truncating off the first `offset` rows of
|
|
the loaded data for `input`.
|
|
|
|
See Also
|
|
--------
|
|
zipline.pipeline.graph.TermGraph.offset
|
|
zipline.pipeline.engine.SimplePipelineEngine._inputs_for_term
|
|
zipline.pipeline.engine.SimplePipelineEngine._mask_and_dates_for_term
|
|
"""
|
|
extra = self.extra_rows
|
|
return {
|
|
# Another way of thinking about this is:
|
|
# How much bigger is the array for ``dep`` compared to ``term``?
|
|
# How much of that difference did I ask for.
|
|
(term, dep): (extra[dep] - extra[term]) - requested_extra_rows
|
|
for term in self
|
|
for dep, requested_extra_rows in term.dependencies.items()
|
|
}
|
|
|
|
@lazyval
|
|
def extra_rows(self):
|
|
"""
|
|
A dict mapping `term` -> `# of extra rows to load/compute of `term`.
|
|
|
|
This is always the maximum number of extra **input** rows required by
|
|
any Filter/Factor for which `term` is an input.
|
|
|
|
Notes
|
|
----
|
|
This value depends on the other terms in the graph that require `term`
|
|
**as an input**. This is not to be confused with `term.dependencies`,
|
|
which describes how many additional rows of `term`'s inputs we need to
|
|
load, and which is determined entirely by `Term` itself.
|
|
|
|
Example
|
|
-------
|
|
Our graph contains the following terms:
|
|
|
|
A = SimpleMovingAverage([USEquityPricing.high], window_length=5)
|
|
B = SimpleMovingAverage([USEquityPricing.high], window_length=10)
|
|
C = SimpleMovingAverage([USEquityPricing.low], window_length=8)
|
|
|
|
To compute N rows of A, we need N + 4 extra rows of `high`.
|
|
To compute N rows of B, we need N + 9 extra rows of `high`.
|
|
To compute N rows of C, we need N + 7 extra rows of `low`.
|
|
|
|
We store the following extra_row requirements:
|
|
|
|
self.extra_rows[high] = 9 # Ensures that we can service B.
|
|
self.extra_rows[low] = 7
|
|
|
|
See Also
|
|
--------
|
|
zipline.pipeline.graph.TermGraph.offset
|
|
zipline.pipeline.term.Term.dependencies
|
|
"""
|
|
return {
|
|
term: attrs['extra_rows']
|
|
for term, attrs in iteritems(self.node)
|
|
}
|
|
|
|
@property
|
|
def outputs(self):
|
|
"""
|
|
Dict mapping names to designated output terms.
|
|
"""
|
|
return self._outputs
|
|
|
|
def ordered(self):
|
|
"""
|
|
Return a topologically-sorted iterator over the terms in `self`.
|
|
"""
|
|
return iter(self._ordered)
|
|
|
|
@lazyval
|
|
def loadable_terms(self):
|
|
return tuple(term for term in self if isinstance(term, LoadableTerm))
|
|
|
|
def _add_to_graph(self, term, parents, extra_rows):
|
|
"""
|
|
Add `term` and all its inputs to the graph.
|
|
"""
|
|
if self._frozen:
|
|
raise ValueError("Can't mutate `TermGraph` after construction.")
|
|
# 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)
|
|
|
|
# Idempotent if term is already in the graph.
|
|
self.add_node(term)
|
|
|
|
# Make sure we're going to compute at least `extra_rows` of `term`.
|
|
self._ensure_extra_rows(term, extra_rows)
|
|
|
|
# Recursively add dependencies.
|
|
for dependency, additional_extra_rows in term.dependencies.items():
|
|
self._add_to_graph(
|
|
dependency,
|
|
parents,
|
|
extra_rows=extra_rows + additional_extra_rows,
|
|
)
|
|
self.add_edge(dependency, term)
|
|
|
|
parents.remove(term)
|
|
|
|
def _ensure_extra_rows(self, term, N):
|
|
"""
|
|
Ensure that we're going to compute at least N extra rows of `term`.
|
|
"""
|
|
attrs = self.node[term]
|
|
attrs['extra_rows'] = max(N, attrs.get('extra_rows', 0))
|
|
|
|
@lazyval
|
|
def jpeg(self):
|
|
return display_graph(self, 'jpeg')
|
|
|
|
@lazyval
|
|
def png(self):
|
|
return display_graph(self, 'png')
|
|
|
|
@lazyval
|
|
def svg(self):
|
|
return display_graph(self, 'svg')
|
|
|
|
def _repr_png_(self):
|
|
return self.png.data
|