ENH: Allow 2D input in cum_returns_final.

Scott Sanderson · Scott Sanderson · commit 7d39c4ab8cb6 · 2018-06-13T10:51:27.000-04:00
diff --git a/empyrical/stats.py b/empyrical/stats.py
@@ -250,32 +250,34 @@ def cum_returns_final(returns, starting_value=0):
 
     Parameters
     ----------
-    returns : pd.Series or np.ndarray
-       Returns of the strategy as a percentage, noncumulative.
-         - Time series with decimal returns.
-         - Example:
-            2015-07-16    -0.012143
-            2015-07-17    0.045350
-            2015-07-20    0.030957
-            2015-07-21    0.004902.
+    returns : pd.DataFrame, pd.Series, or np.ndarray
+       Noncumulative simple returns of one or more timeseries.
     starting_value : float, optional
        The starting returns.
 
     Returns
     -------
-    total_returns : float
-    """
+    total_returns : pd.Series, np.ndarray, or float
+        If input is 1-dimensional (a Series or 1D numpy array), the result is a
+        scalar.
 
+        If input is 2-dimensional (a DataFrame or 2D numpy array), the result
+        is a 1D array containing cumulative returns for each column of input.
+    """
     if len(returns) == 0:
         return np.nan
 
-    returns = np.nanprod(returns + 1)
+    if isinstance(returns, pd.DataFrame):
+        result = (returns + 1).prod()
+    else:
+        result = np.nanprod(returns + 1, axis=0)
+
     if starting_value == 0:
-        returns -= 1
+        result -= 1
     else:
-        returns *= starting_value
+        result *= starting_value
 
-    return returns
+    return result
 
 
 def aggregate_returns(returns, convert_to):
diff --git a/empyrical/tests/test_stats.py b/empyrical/tests/test_stats.py
@@ -1495,6 +1495,19 @@ def test_cum_returns_df(self, returns, starting_value, expected):
 
         self.assert_indexes_match(cum_returns, returns)
 
+    @parameterized.expand([
+        (df_input, 0, df_0_expected.iloc[-1]),
+        (df_input, 100, df_100_expected.iloc[-1]),
+    ])
+    def test_cum_returns_final_df(self, returns, starting_value, expected):
+        return_types = (pd.Series, np.ndarray)
+        result = self.empyrical(return_types=return_types).cum_returns_final(
+            returns,
+            starting_value=starting_value,
+        )
+        assert_almost_equal(np.array(result), expected, 5)
+        self.assert_indexes_match(result, expected)
+
     @property
     def empyrical(self):
         """
