From c5f4d00bf137aead51bba77ac4aec5d7906ba450 Mon Sep 17 00:00:00 2001 From: Wes McKinney Date: Thu, 21 Mar 2013 23:16:32 -0400 Subject: [PATCH] ENH: prototype data structure for managing a rolling datapanel Manage a rolling window collection of collection of panels for computation purposes. --- tests/test_data_util.py | 108 ++++++++++++++++++++++++++++++++++++++++ zipline/utils/data.py | 86 ++++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 tests/test_data_util.py create mode 100644 zipline/utils/data.py diff --git a/tests/test_data_util.py b/tests/test_data_util.py new file mode 100644 index 00000000..ded2c8ca --- /dev/null +++ b/tests/test_data_util.py @@ -0,0 +1,108 @@ +# +# Copyright 2013 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from collections import deque + +import numpy as np + +import pandas as pd +import pandas.util.testing as tm + +from zipline.utils.data import RollingPanel + + +class TestRollingPanel(unittest.TestCase): + + def test_basics(self): + items = ['foo', 'bar', 'baz'] + minor = ['A', 'B', 'C', 'D'] + + window = 10 + + rp = RollingPanel(window, items, minor, cap_multiple=2) + + dates = pd.date_range('2000-01-01', periods=30, tz='utc') + + major_deque = deque() + + frames = {} + + for i in range(30): + frame = pd.DataFrame(np.random.randn(3, 4), index=items, + columns=minor) + date = dates[i] + + rp.add_frame(date, frame) + + frames[date] = frame + major_deque.append(date) + + if i >= window: + major_deque.popleft() + + result = rp.get_current() + expected = pd.Panel(frames, items=list(major_deque), + major_axis=items, minor_axis=minor) + tm.assert_panel_equal(result, expected.swapaxes(0, 1)) + + +def f(option='clever', n=500, copy=False): + items = range(5) + minor = range(20) + window = 100 + periods = n + + dates = pd.date_range('2000-01-01', periods=periods, tz='utc') + frames = {} + + if option == 'clever': + rp = RollingPanel(window, items, minor, cap_multiple=2) + major_deque = deque() + dummy = pd.DataFrame(np.random.randn(len(items), len(minor)), + index=items, columns=minor) + + for i in range(periods): + frame = dummy * (1 + 0.001 * i) + date = dates[i] + + rp.add_frame(date, frame) + + frames[date] = frame + major_deque.append(date) + + if i >= window: + del frames[major_deque.popleft()] + + result = rp.get_current() + if copy: + result = result.copy() + else: + major_deque = deque() + dummy = pd.DataFrame(np.random.randn(len(items), len(minor)), + index=items, columns=minor) + + for i in range(periods): + frame = dummy * (1 + 0.001 * i) + date = dates[i] + frames[date] = frame + major_deque.append(date) + + if i >= window: + del frames[major_deque.popleft()] + + result = pd.Panel(frames, items=list(major_deque), + major_axis=items, minor_axis=minor) diff --git a/zipline/utils/data.py b/zipline/utils/data.py new file mode 100644 index 00000000..ca9457b6 --- /dev/null +++ b/zipline/utils/data.py @@ -0,0 +1,86 @@ +# +# Copyright 2013 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pandas as pd + + +def _ensure_index(x): + if not isinstance(x, pd.Index): + x = pd.Index(x) + + return x + + +class RollingPanel(object): + """ + Preallocation strategies for rolling window over expanding data set + + Restrictions: major_axis can only be a DatetimeIndex for now + """ + + def __init__(self, window, items, minor_axis, cap_multiple=2, + dtype=np.float64): + self.pos = 0 + self.window = window + + self.items = _ensure_index(items) + self.minor_axis = _ensure_index(minor_axis) + + self.cap_multiple = cap_multiple + self.cap = cap_multiple * window + + self.dtype = dtype + self.buffer = np.empty((len(items), self.cap, len(minor_axis)), + dtype=dtype) + self.index_buf = np.empty(self.cap, dtype='M8[ns]') + + def add_frame(self, tick, frame): + """ + + TODO: this assumes the DataFrame has the right shape + """ + if self.pos == self.cap: + self._roll_data() + + self.buffer[:, self.pos, :] = frame.values + self.index_buf[self.pos] = tick + + self.pos += 1 + + def get_current(self): + """ + Get a Panel that is the current data in view. It is not safe to persist + these objects because internal data might change + """ + where = slice(max(self.pos - self.window, 0), self.pos) + major_axis = pd.DatetimeIndex(self.index_buf[where], tz='utc') + return pd.Panel(self.buffer[:, where, :], self.items, major_axis, + self.minor_axis) + + def _roll_data(self): + """ + Roll window worth of data up to position zero. + Save the effort of having to expensively roll at each iteration + """ + self.buffer[:, :self.window, :] = self.buffer[:, -self.window:] + self.index_buf[:self.window] = self.index_buf[-self.window:] + self.pos = self.window + + +class NaiveRollingPanel(object): + + def __init__(self, window, items, minor_axis, cap_multiple=2): + pass