Use updated pandas frequency string

pymc-devs · Feb 14, 2024 · faf1c87 · faf1c87
1 parent d32b71a
commit faf1c87
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/pymc_experimental/tests/statespace/utilities/test_helpers.py b/pymc_experimental/tests/statespace/utilities/test_helpers.py
@@ -19,8 +19,15 @@
 
 
 def load_nile_test_data():
+    from importlib.metadata import version
+
     nile = pd.read_csv("pymc_experimental/tests/statespace/test_data/nile.csv", dtype={"x": floatX})
-    nile.index = pd.date_range(start="1871-01-01", end="1970-01-01", freq="AS-Jan")
+    major, minor, rev = map(int, version("pandas").split("."))
+    if major >= 2 and minor >= 2 and rev >= 0:
+        freq_str = "YS-JAN"
+    else:
+        freq_str = "AS-JAN"
+    nile.index = pd.date_range(start="1871-01-01", end="1970-01-01", freq=freq_str)
     nile.rename(columns={"x": "height"}, inplace=True)
     nile = (nile - nile.mean()) / nile.std()
     nile = nile.astype(floatX)