From 4beb590fa4df5625157a11efa97b7797e74dbdab Mon Sep 17 00:00:00 2001
From: lpalma <lluis.palma@bsc.es>
Date: Mon, 1 Jul 2024 23:08:35 +0200
Subject: [PATCH] added proper seasonal means comp

---
 tinysunset/tools/temp_means.py | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/tinysunset/tools/temp_means.py b/tinysunset/tools/temp_means.py
index 2e651ed..fd0e058 100644
--- a/tinysunset/tools/temp_means.py
+++ b/tinysunset/tools/temp_means.py
@@ -1,16 +1,37 @@
 import xarray as xr
 import pandas as pd
+import numpy as np
+
+_MONTH_INT_2_ACRONYM = {
+    1: "J", 2: "F", 3: "M", 4: "A", 5: "M", 6: "J",
+    7: "J", 8: "A", 9: "S", 10: "O", 11: "N", 12: "D"
+}
+
+def get_season_str(month, nmonths):
+    return ''.join(_MONTH_INT_2_ACRONYM[(month + i - 1) % 12 + 1] for i in range(nmonths))
 
 def monthly2seasonal(ds: xr.Dataset, nmonths: int = 3) -> xr.Dataset:
+    
     # Sort by time and create a complete monthly time series
+    
     ds = ds.sortby("time")
+
     start, end = ds.time.to_index()[[0, -1]]
+    
     new_time_index = pd.date_range(start=f"{start.year}-01-01", 
                                     end=f"{end.year}-12-31", 
                                     freq="MS")
-    
-    # Reindex, compute seasonal means, and trim to original time range
-    return (ds.reindex(time=new_time_index)
-            .rolling(time=nmonths, center=True).mean()
+
+    ds = (ds.reindex(time=new_time_index)
+            .rolling(time=nmonths).mean().shift(time=-nmonths+1)
             .sel(time=slice(start, end))
-            .dropna(dim="time"))
\ No newline at end of file
+            .dropna(dim="time"))
+    
+    # create coord with the aggregated months
+    ds = ds.assign_coords(
+        season=("time", np.vectorize(get_season_str)(ds.time.dt.month.values, nmonths))
+        )
+
+    ds.attrs["frequency"] = f"{nmonths}M"
+    
+    return ds 
\ No newline at end of file
-- 
GitLab