From e9ec4c6dc2edba640cf426c668edb78fe56469c1 Mon Sep 17 00:00:00 2001 From: Frank Milthaler Date: Wed, 2 Aug 2023 18:50:01 +0200 Subject: [PATCH] Vectorized computation of cumulative_returns (#127) rather than row-by-row operation with `apply` --- finquant/returns.py | 3 ++- tests/test_returns.py | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/finquant/returns.py b/finquant/returns.py index 6401cd2d..2a2e26c8 100644 --- 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 index 730a30f5..3ac33999 100644 --- 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)