1325 invalid way of calling first element of pdseries (#1433)

* add test cases with index not starting at zero

* make the change and add docs

pymc_marketing/mmm/tvp.py

@@ -234,10 +234,27 @@ def create_time_varying_gp_multiplier(
 
 
 def infer_time_index(
-    date_series_new: pd.Series, date_series: pd.Series, time_resolution: int
+    date_series_new: pd.Series,
+    date_series: pd.Series,
+    time_resolution: int,
 ) -> npt.NDArray[np.int_]:
     """Infer the time-index given a new dataset.
 
     Infers the time-indices by calculating the number of days since the first date in the dataset.
+
+    Parameters
+    ----------
+    date_series_new : pd.Series
+        New date series.
+    date_series : pd.Series
+        Original date series.
+    time_resolution : int
+        Resolution of time points in days.
+
+    Returns
+    -------
+    np.ndarray
+        Time index.
+
     """
-    return (date_series_new - date_series[0]).dt.days.values // time_resolution
+    return (date_series_new - date_series.iloc[0]).dt.days.values // time_resolution

tests/mmm/test_tvp.py

@@ -149,28 +149,26 @@ def test_create_time_varying_intercept(coords, model_config):
         assert isinstance(result, pt.TensorVariable)
 
 
-@pytest.mark.parametrize("freq, time_resolution", [("D", 1), ("W", 7)])
-def test_infer_time_index_in_sample(freq, time_resolution):
-    date_series = pd.Series(pd.date_range(start="1/1/2022", periods=5, freq=freq))
-    date_series_new = date_series
-    expected = np.arange(0, 5)
-    result = infer_time_index(date_series_new, date_series, time_resolution)
-    np.testing.assert_array_equal(result, expected)
-
-
-@pytest.mark.parametrize("freq, time_resolution", [("D", 1), ("W", 7)])
-def test_infer_time_index_oos_forward(freq, time_resolution):
-    date_series = pd.Series(pd.date_range(start="1/1/2022", periods=5, freq=freq))
-    date_series_new = date_series + pd.Timedelta(5, unit=freq)
-    expected = np.arange(5, 10)
-    result = infer_time_index(date_series_new, date_series, time_resolution)
-    np.testing.assert_array_equal(result, expected)
-
-
-@pytest.mark.parametrize("freq, time_resolution", [("D", 1), ("W", 7)])
-def test_infer_time_index_oos_backward(freq, time_resolution):
-    date_series = pd.Series(pd.date_range(start="1/1/2022", periods=5, freq=freq))
-    date_series_new = date_series - pd.Timedelta(5, unit=freq)
-    expected = np.arange(-5, 0)
+@pytest.mark.parametrize(
+    "freq, time_resolution",
+    [
+        pytest.param("D", 1, id="daily"),
+        pytest.param("W", 7, id="weekly"),
+    ],
+)
+@pytest.mark.parametrize(
+    "index",
+    [np.arange(5), np.arange(5) + 10],
+    ids=["zero-start", "non-zero-start"],
+)
+@pytest.mark.parametrize(
+    "offset, expected",
+    [(0, np.arange(0, 5)), (5, np.arange(5, 10)), (-5, np.arange(-5, 0))],
+    ids=["in-sample", "oos_forward", "oos_backward"],
+)
+def test_infer_time_index(freq, time_resolution, index, offset, expected):
+    dates = pd.date_range(start="1/1/2022", periods=5, freq=freq)
+    date_series = pd.Series(dates, index=index)
+    date_series_new = date_series + pd.Timedelta(offset, unit=freq)
     result = infer_time_index(date_series_new, date_series, time_resolution)
     np.testing.assert_array_equal(result, expected)