date_range: update freq string (MESMER-group#504)

* date_range: update freq string * add freq conditionals
veni-vidi-vici-dormivi · Aug 21, 2024 · ea5004c · ea5004c
1 parent efb77f3
commit ea5004c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/tests/unit/test_auto_regression.py b/tests/unit/test_auto_regression.py
@@ -4,6 +4,7 @@
 import pandas as pd
 import pytest
 import xarray as xr
+from packaging.version import Version
 
 import mesmer
 from mesmer.core.utils import LinAlgWarning, _check_dataarray_form, _check_dataset_form
@@ -638,14 +639,15 @@ def test_fit_autoregression_monthly_np_with_noise(slope, intercept, std):
 
 
 def test_fit_auto_regression_monthly():
+    freq = "ME" if Version(pd.__version__) >= Version("2.2") else "M"
     n_years = 20
     n_gridcells = 10
     np.random.seed(0)
     data = xr.DataArray(
         np.random.normal(size=(n_years * 12, n_gridcells)),
         dims=("time", "gridcell"),
         coords={
-            "time": pd.date_range("2000-01-01", periods=n_years * 12, freq="M"),
+            "time": pd.date_range("2000-01-01", periods=n_years * 12, freq=freq),
             "gridcell": np.arange(n_gridcells),
         },
     )
@@ -728,6 +730,8 @@ def test_draw_autoregression_monthly_np_rng():
 
 
 def test_draw_auto_regression_monthly():
+    freq = "ME" if Version(pd.__version__) >= Version("2.2") else "M"
+
     n_gridcells = 10
     n_realisations = 5
     n_years = 10
@@ -756,7 +760,7 @@ def test_draw_auto_regression_monthly():
         },
     )
 
-    time = pd.date_range("2000-01-01", periods=n_years * 12, freq="M")
+    time = pd.date_range("2000-01-01", periods=n_years * 12, freq=freq)
     time = xr.DataArray(time, dims="time", coords={"time": time})
 
     result = mesmer.stats.draw_auto_regression_monthly(

diff --git a/tests/unit/test_harmonic_model.py b/tests/unit/test_harmonic_model.py
@@ -204,8 +204,11 @@ def test_fit_harmonic_model_checks():
     with pytest.raises(TypeError):
         mesmer.stats.fit_harmonic_model(yearly_predictor, monthly_target.values)
 
-    yearly_predictor["time"] = pd.date_range("2000-01-01", periods=10, freq="Y")
-    monthly_target["time"] = pd.date_range("2000-02-01", periods=10 * 12, freq="M")
+    freq = "YE" if Version(pd.__version__) >= Version("2.2") else "Y"
+    yearly_predictor["time"] = pd.date_range("2000-01-01", periods=10, freq=freq)
+
+    freq = "ME" if Version(pd.__version__) >= Version("2.2") else "M"
+    monthly_target["time"] = pd.date_range("2000-02-01", periods=10 * 12, freq=freq)
     with pytest.raises(ValueError, match="Monthly target data must start with January"):
         mesmer.stats.fit_harmonic_model(yearly_predictor, monthly_target)
 
@@ -214,8 +217,12 @@ def test_fit_harmonic_model_time_dim():
     # test if the time dimension can be different from "time"
     yearly_predictor = trend_data_2D(n_timesteps=10, n_lat=3, n_lon=2)
     monthly_target = trend_data_2D(n_timesteps=10 * 12, n_lat=3, n_lon=2)
-    yearly_predictor["time"] = pd.date_range("2000-01-01", periods=10, freq="Y")
-    monthly_target["time"] = pd.date_range("2000-01-01", periods=10 * 12, freq="M")
+
+    freq = "YE" if Version(pd.__version__) >= Version("2.2") else "Y"
+    yearly_predictor["time"] = pd.date_range("2000-01-01", periods=10, freq=freq)
+
+    freq = "ME" if Version(pd.__version__) >= Version("2.2") else "M"
+    monthly_target["time"] = pd.date_range("2000-01-01", periods=10 * 12, freq=freq)
 
     time_dim = "dates"
     monthly_target = monthly_target.rename({"time": time_dim})