From a483455183d750876610f57bb5243e897eb67b7c Mon Sep 17 00:00:00 2001
From: Scott Sanderson <scoutoss@gmail.com>
Date: Tue, 29 Sep 2015 16:30:03 -0400
Subject: [PATCH] ENH: Add zipline.utils.numpy_utils.

Currently provides two functions for doing fancy things with array
strides: `repeat_first_axis` and `repeat_last_axis`.
---
 tests/test_doctests.py       |  4 ++
 zipline/utils/numpy_utils.py | 90 ++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 zipline/utils/numpy_utils.py

diff --git a/tests/test_doctests.py b/tests/test_doctests.py
index b71648c9..f5406fbc 100644
--- a/tests/test_doctests.py
+++ b/tests/test_doctests.py
@@ -13,6 +13,7 @@ from zipline.utils import (
     memoize,
     preprocess,
     test_utils,
+    numpy_utils,
 )
 
 
@@ -70,3 +71,6 @@ class DoctestTestCase(TestCase):
 
     def test_cache_docs(self):
         self._check_docs(cache)
+
+    def test_numpy_utils_docs(self):
+        self._check_docs(numpy_utils)
diff --git a/zipline/utils/numpy_utils.py b/zipline/utils/numpy_utils.py
new file mode 100644
index 00000000..b9a12c17
--- /dev/null
+++ b/zipline/utils/numpy_utils.py
@@ -0,0 +1,90 @@
+"""
+Utilities for working with numpy arrays.
+"""
+from numpy.lib.stride_tricks import as_strided
+
+
+def repeat_first_axis(array, count):
+    """
+    Restride `array` to repeat `count` times along the first axis.
+
+    Parameters
+    ----------
+    array : np.array
+        The array to restride.
+    count : int
+        Number of times to repeat `array`.
+
+    Returns
+    -------
+    result : array
+        Array of shape (count,) + array.shape, composed of `array` repeated
+        `count` times along the first axis.
+
+    Example
+    -------
+    >>> from numpy import arange
+    >>> a = arange(3); a
+    array([0, 1, 2])
+    >>> repeat_first_axis(a, 2)
+    array([[0, 1, 2],
+           [0, 1, 2]])
+    >>> repeat_first_axis(a, 4)
+    array([[0, 1, 2],
+           [0, 1, 2],
+           [0, 1, 2],
+           [0, 1, 2]])
+
+    Notes
+    ----
+    The resulting array will share memory with `array`.  If you need to assign
+    to the input or output, you should probably make a copy first.
+
+    See Also
+    --------
+    repeat_last_axis
+    """
+    return as_strided(array, (count,) + array.shape, (0,) + array.strides)
+
+
+def repeat_last_axis(array, count):
+    """
+    Restride `array` to repeat `count` times along the last axis.
+
+    Parameters
+    ----------
+    array : np.array
+        The array to restride.
+    count : int
+        Number of times to repeat `array`.
+
+    Returns
+    -------
+    result : array
+        Array of shape array.shape + (count,) composed of `array` repeated
+        `count` times along the last axis.
+
+    Example
+    -------
+    >>> from numpy import arange
+    >>> a = arange(3); a
+    array([0, 1, 2])
+    >>> repeat_last_axis(a, 2)
+    array([[0, 0],
+           [1, 1],
+           [2, 2]])
+    >>> repeat_last_axis(a, 4)
+    array([[0, 0, 0, 0],
+           [1, 1, 1, 1],
+           [2, 2, 2, 2]])
+
+    Notes
+    ----
+    The resulting array will share memory with `array`.  If you need to assign
+    to the input or output, you should probably make a copy first.
+
+    See Also
+    --------
+    repeat_last_axis
+    """
+    return as_strided(array, array.shape + (count,), array.strides + (0,))