Commit b9b19221 authored by Aaron Spring's avatar Aaron Spring 🚼
Browse files

groupby new week

parent b33110f1
Pipeline #232147 passed with stages
in 34 minutes and 19 seconds
This diff is collapsed.
......@@ -158,15 +158,25 @@ def add_year_week_coords(ds):
return ds
def add_year_week_coords(ds):
if 'week' not in ds.coords and 'year' not in ds.coords:
year = ds.forecast_time.dt.year.to_index().unique()
week = sorted(list(ds.forecast_time.dt.week.to_index().unique()))
weeks = week * len(year)
import numpy as np
years = np.repeat(year,len(week))
ds.coords["week"] = ("forecast_time", weeks)
ds.coords['week'].attrs['description'] = "This week represents the number of forecast_time starting from 1 to 53. Note: This week is different from the ISO week from groupby('forecast_time.weekofyear'), see https://en.wikipedia.org/wiki/ISO_week_date and https://renkulab.io/gitlab/aaron.spring/s2s-ai-challenge/-/issues/29"
ds.coords["year"] = ("forecast_time", years)
ds.coords['year'].attrs['long_name'] = "calendar year"
return ds
def make_probabilistic(ds, tercile_edges, member_dim='realization', mask=None, groupby_coord='week'):
"""Compute probabilities from ds (observations or forecasts) based on tercile_edges."""
# broadcast
if 'forecast_time' not in tercile_edges.dims and groupby_coord in tercile_edges.dims:
if groupby_coord == 'week':
ds = split_forecast_time_into_year_week(ds)
tercile_edges = tercile_edges.sel({groupby_coord:ds.coords[groupby_coord]})
else:
tercile_edges = tercile_edges.sel({groupby_coord:ds.forecast_time.dt[groupby_coord]})
ds = add_year_week_coords(ds)
tercile_edges = tercile_edges.sel({groupby_coord: ds.coords[groupby_coord]})
bn = ds < tercile_edges.isel(category_edge=0, drop=True) # below normal
n = (ds >= tercile_edges.isel(category_edge=0, drop=True)) & (ds < tercile_edges.isel(category_edge=1, drop=True)) # normal
an = ds >= tercile_edges.isel(category_edge=1, drop=True) # above normal
......@@ -192,8 +202,8 @@ def make_probabilistic(ds, tercile_edges, member_dim='realization', mask=None, g
'comment': 'All three tercile category probabilities must add up to 1.',
'variable_before_categorization': 'https://confluence.ecmwf.int/display/S2S/S2S+Surface+Air+Temperature'
}
if 'weekofyear' in ds_p.coords:
ds_p = ds_p.drop('weekofyear')
if groupby_coord in ds_p.coords:
ds_p = ds_p.drop(groupby_coord)
return ds_p
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment