<jupyter_start><jupyter_text>Debias precipitation using a gamma distributionSo same thing as the other debiasing notebooks, but this time we use a gamma distribution instead of a gaussian.This should lead to a slightly better debiasing.<jupyter_code>import dask
import dask.array as da
import dask.distributed
import datetime
import matplotlib.pyplot as plt
import os
import numpy as np
import pandas as pd
import pathlib
import scipy.stats
import typing
import xarray as xr
import xskillscore as xs
from crims2s.dask import create_dask_cluster
from crims2s.util import fix_dataset_dims
INPUT_TRAIN = '***BASEDIR***training-input/0.3.0/netcdf'
OBSERVATIONS = '***BASEDIR***/processed/training-output-reference/'
BENCHNMARK = '***BASEDIR***training-output-benchmark/'<jupyter_output><empty_output><jupyter_text>Boost dask cluster<jupyter_code>cluster = create_dask_cluster()
cluster.scale(jobs=2)
client = dask.distributed.Client(cluster)
client<jupyter_output><empty_output><jupyter_text>Generic Functions<jupyter_code>def extract_train_validation_from_lead_time(xr_data) -> typing.Tuple:
    xr_data_sub_train = xr_data.sel(forecast_year=slice(None, 2018))
    xr_data_sub_val = xr_data.sel(forecast_year=slice(2019, None))
    
    return xr_data_sub_train, xr_data_sub_val
def compute_and_correct_bias(data_center_train, data_center_val, obs_train):
    
    bias = (obs_train - data_center_train).mean(dim=['lead_time', 'forecast_year'])
    corrected_bias = data_center_val + bias
    
    return bias, corrected_bias
def add_biweekly_dim(dataset):
    weeklys = []
    for s in [slice('0D', '13D'), slice('14D', '27D'), slice('28D', '41D')]:
        weekly_forecast = dataset.sel(lead_time=s)

        first_lead = pd.to_timedelta(weekly_forecast.lead_time[0].item())

        weekly_forecast = weekly_forecast.expand_dims(dim='biweekly_forecast').assign_coords(biweekly_forecast=[first_lead])
        weekly_forecast = weekly_forecast.assign_coords(lead_time=(weekly_forecast.lead_time - first_lead))
        weeklys.append(weekly_forecast)
        
    weeklys[2] -= weeklys[1].isel(lead_time=-1)
    weeklys[1] -= weeklys[0].isel(lead_time=-1)
        
    return xr.concat(weeklys, dim='biweekly_forecast').transpose('forecast_year', 'forecast_dayofyear', 'biweekly_forecast', ...)<jupyter_output><empty_output><jupyter_text>Read DATA<jupyter_code>CENTER = 'ncep'
FIELD = 'tp'
input_path = pathlib.Path(INPUT_TRAIN)
input_files_tp = sorted([f for f in input_path.iterdir() if CENTER in f.stem and FIELD in f.stem])
input_files_tp[:10]
ecmwf_tp_raw = xr.open_mfdataset(input_files_tp, preprocess=fix_dataset_dims)
ecmwf_tp_raw.isel(lead_time=0, realization=0, forecast_year=0, forecast_dayofyear=0).tp.plot()<jupyter_output><empty_output><jupyter_text>There are non-zero values on the first lead time, so it's the values accumulated after 24h<jupyter_code>ecmwf_tp = add_biweekly_dim(ecmwf_tp_raw)
ecmwf_tp
n_smaller = (ecmwf_tp.isel(lead_time=-1) < ecmwf_tp.isel(lead_time=0))
n_smaller
n_smaller.sum(dim=['realization', 'latitude', 'longitude', 'forecast_dayofyear', 'forecast_year']).compute()<jupyter_output><empty_output><jupyter_text>Observations<jupyter_code>obs_path = pathlib.Path(OBSERVATIONS)
obs_files = [f for f in obs_path.iterdir() if 'tp' in f.stem]
obs_files[:4]
obs_tp_raw = xr.open_mfdataset(obs_files)
obs_tp_raw = obs_tp_raw.assign_coords(lead_time=obs_tp_raw.lead_time - obs_tp_raw.lead_time[0])
obs_tp = add_biweekly_dim(obs_tp_raw)
obs_tp
obs_tp = obs_tp.isel(lead_time=-1) - obs_tp.isel(lead_time=0)<jupyter_output><empty_output><jupyter_text>Split in train test<jupyter_code>ecmwf_tp_train, ecmwf_tp_val = extract_train_validation_from_lead_time(ecmwf_tp)
obs_tp_train, obs_tp_val = extract_train_validation_from_lead_time(obs_tp)
ecmwf_tp_train
ecmwf_tp_train.isel(biweekly_forecast=1, forecast_dayofyear=10, latitude=30, longitude=30)
obs_tp_train<jupyter_output><empty_output><jupyter_text>Fit Gamma distribution<jupyter_code>one_slice = ecmwf_tp_train.isel(biweekly_forecast=1, forecast_dayofyear=0).compute()
(one_slice.tp < 0.0).sum(dim=['forecast_year', 'realization']).plot()
scipy.stats.gamma.fit(one_slice.tp.data)
xr.apply_ufunc(scipy.stats.gamma.fit)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11308256446652472, 'fraction_numerical': 0.0134847409510291, 'mean_word_length': 7.439121756487026, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '516082', 'license': 'MIT', 'n_tokens_mistral': 1680, 'n_tokens_neox': 1600, 'n_words': 248, 'path': 'notebooks/landryda/gamma-tp-debias.ipynb', 'repo_name': 'crim-ca/crims2s', 'size': '4227'}