diff --git a/tinysunset/tools/temp_means.py b/tinysunset/tools/temp_means.py index 2e651edf7c2db9984486cfa541b0a548e2baf14e..fd0e0582e81866ebbbd4104623d3d8f43b8bd21a 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