Vectorized computation of cumulative_returns (#127)

rather than row-by-row operation with `apply`
fmilthaler · Aug 2, 2023 · e9ec4c6 · e9ec4c6
1 parent d116bee
commit e9ec4c6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/finquant/returns.py b/finquant/returns.py
@@ -18,7 +18,8 @@ def cumulative_returns(data, dividend=0):
     :Output:
      :ret: a ``pandas.DataFrame`` of cumulative Returns of given stock prices.
     """
-    return data.dropna(axis=0, how="any").apply(lambda x: (x - x[0] + dividend) / x[0])
+    data = data.dropna(axis=0, how="any")
+    return ((data - data.iloc[0] + dividend) / data.iloc[0]).astype(np.float64)
 
 
 def daily_returns(data):

diff --git a/tests/test_returns.py b/tests/test_returns.py
@@ -20,6 +20,7 @@ def test_cumulative_returns():
     d = {"1": l1, "2": l2}
     df = pd.DataFrame(d)
     ret = cumulative_returns(df)
+    assert isinstance(ret, pd.DataFrame) and not ret.empty
     assert all(abs(ret["1"].values - orig[0]) <= 1e-15)
     assert all(abs(ret["2"].values - orig[1]) <= 1e-15)
     # with dividend of 0.2
@@ -28,6 +29,7 @@ def test_cumulative_returns():
         [0.005, -0.02, -0.045, -0.07, -0.095, -0.12, -0.145, -0.17, -0.195, -0.22],
     ]
     ret = cumulative_returns(df, 0.2)
+    assert isinstance(ret, pd.DataFrame) and not ret.empty
     assert all(abs(ret["1"].values - orig[0]) <= 1e-15)
     assert all(abs(ret["2"].values - orig[1]) <= 1e-15)