diff --git a/zipline/pipeline/loaders/synthetic.py b/zipline/pipeline/loaders/synthetic.py index f1783596..8c50d869 100644 --- a/zipline/pipeline/loaders/synthetic.py +++ b/zipline/pipeline/loaders/synthetic.py @@ -6,6 +6,7 @@ from bcolz import ctable from numpy import ( arange, array, + eye, float64, full, iinfo, @@ -86,6 +87,29 @@ class PrecomputedLoader(PipelineLoader): return out +class EyeLoader(PrecomputedLoader): + """ + A PrecomputedLoader that emits arrays containing 1s on the diagonal and 0s + elsewhere. + + Parameters + ---------- + columns : list[BoundColumn] + Columns that this loader should know about. + dates : iterable[datetime-like] + Same as PrecomputedLoader. + sids : iterable[int-like] + Same as PrecomputedLoader + """ + def __init__(self, columns, dates, sids): + shape = (len(dates), len(sids)) + super(EyeLoader, self).__init__( + {column: eye(shape, dtype=column.dtype) for column in columns}, + dates, + sids, + ) + + class SyntheticDailyBarWriter(BcolzDailyBarWriter): """ Bcolz writer that creates synthetic data based on asset lifetime metadata.