diff --git a/python/src/robyn/modeling/ridge/ridge_data_builder.py b/python/src/robyn/modeling/ridge/ridge_data_builder.py index ece54ecfa..b4988487e 100644 --- a/python/src/robyn/modeling/ridge/ridge_data_builder.py +++ b/python/src/robyn/modeling/ridge/ridge_data_builder.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Optional, Tuple import pandas as pd import numpy as np +from scipy.signal import lfilter class RidgeDataBuilder: @@ -166,12 +167,11 @@ def _hyper_collector( return hyper_collect def _geometric_adstock(self, x: pd.Series, theta: float) -> pd.Series: - # print(f"Before adstock: {x.head()}") - y = x.copy() - for i in range(1, len(x)): - y.iloc[i] += theta * y.iloc[i - 1] - # print(f"After adstock: {y.head()}") - return y + + x_array = x.values + # Use lfilter to efficiently compute the geometric transformation + y = lfilter([1], [1, -theta], x_array) + return pd.Series(y, index=x.index) def _hill_transformation( self, x: pd.Series, alpha: float, gamma: float diff --git a/python/tests/unit/modeling/ridge/test_adstock.py b/python/tests/unit/modeling/ridge/test_adstock.py new file mode 100644 index 000000000..cf9740c9e --- /dev/null +++ b/python/tests/unit/modeling/ridge/test_adstock.py @@ -0,0 +1,31 @@ +import pytest +import numpy as np +import pandas as pd +from scipy.signal import lfilter +from robyn.modeling.ridge.ridge_data_builder import RidgeDataBuilder + + +def original_geometric_adstock(x: pd.Series, theta: float) -> pd.Series: + y = x.copy() + for i in range(1, len(x)): + y.iloc[i] += theta * y.iloc[i - 1] + return y + + +@pytest.mark.parametrize("theta", [0, 0.5, 0.8, 1]) +def test_geometric_adstock(theta): + x = pd.Series(np.random.rand(10_000)) # Random test data + + # Instantiate the RidgeDataBuilder object (without requiring real data) + dummy_data = None + ridge_builder = RidgeDataBuilder(dummy_data, dummy_data) + + # Call the actual function from the RidgeDataBuilder instance + optimized = ridge_builder._geometric_adstock(x, theta) + + # Compute the expected output using the original function + original = original_geometric_adstock(x, theta) + + assert np.allclose( + original, optimized, atol=1e-6 + ), f"Mismatch found for theta={theta}"