From 510790641344cd4e0c864b153dc45cbcb49ed80f Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Fri, 12 Aug 2016 14:03:49 -0400 Subject: [PATCH] ENH: Add as_column to numpy_utils. --- zipline/utils/numpy_utils.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/zipline/utils/numpy_utils.py b/zipline/utils/numpy_utils.py index 59a32d20..84fbfd7e 100644 --- a/zipline/utils/numpy_utils.py +++ b/zipline/utils/numpy_utils.py @@ -364,3 +364,36 @@ def vectorized_is_element(array, choices): Array indicating whether each element of ``array`` was in ``choices``. """ return vectorize(choices.__contains__, otypes=[bool])(array) + + +def as_column(a): + """ + Convert an array of shape (N,) into an array of shape (N, 1). + + This is equivalent to `a[:, np.newaxis]`. + + Parameters + ---------- + a : np.ndarray + + Example + ------- + >>> import numpy as np + >>> a = np.arange(5) + >>> a + array([0, 1, 2, 3, 4]) + >>> as_column(a) + array([[0], + [1], + [2], + [3], + [4]]) + >>> as_column(a).shape + (5, 1) + """ + if a.ndim != 1: + raise ValueError( + "as_column expected an 1-dimensional array, " + "but got an array of shape %s" % a.shape + ) + return a[:, None]