<jupyter_start><jupyter_text>Generate function to standardize data based on different criteria<jupyter_code>import sys
sys.path.append('/'.join(sys.path[0].split('/')[:-1]))

import xarray as xr
import matplotlib.pyplot as plt
import numpy as np
import time
import os
import healpy as hp
datadir = "../data/healpix/"
input_dir = datadir + "5.625deg_nearest/"
pred_save_path = datadir + "predictions/"

train_years = ('1979', '2012')
val_years = ('2013', '2016')
test_years = ('2017', '2018')

nodes = 12*16*16
max_lead_time = 5*24
lead_time = 6
out_features = 2
nb_timesteps = 2

obs = xr.open_mfdataset(pred_save_path + 'observations_nearest.nc', combine='by_coords')
chunk_size = 521
z500 = xr.open_mfdataset(f'{input_dir}geopotential_500/*.nc', combine='by_coords', chunks={'time':chunk_size}).rename({'z':'z500'})
t850 = xr.open_mfdataset(f'{input_dir}temperature_850/*.nc', combine='by_coords', chunks={'time':chunk_size}).rename({'t':'t850'})
rad = xr.open_mfdataset(f'{input_dir}toa_incident_solar_radiation/*.nc', combine='by_coords', chunks={'time':chunk_size})

z500 = z500.isel(time=slice(7, None))
t850 = t850.isel(time=slice(7, None))

constants = xr.open_dataset(f'{input_dir}constants/constants_5.625deg.nc').rename({'orography' :'orog'})
constants = constants.assign(cos_lon=lambda x: np.cos(np.deg2rad(x.lon)))
constants = constants.assign(sin_lon=lambda x: np.sin(np.deg2rad(x.lon)))

temp = xr.DataArray(np.zeros(z500.dims['time']), coords=[('time', z500.time.values)])
constants, _ = xr.broadcast(constants, temp)
ds = xr.merge([z500, t850, rad], compat='override')
#ds = xr.merge([z500, t850, orog, lats, lsm, slt, rad], compat='override')

ds_train = ds.sel(time=slice(*train_years))
ds_valid = ds.sel(time=slice(*val_years))
ds_test = ds.sel(time=slice(*test_years))<jupyter_output><empty_output><jupyter_text>Standardize based on monthly statistics<jupyter_code>aux = ds_train.groupby('time.month').mean().compute()
aux.to_netcdf(f'{input_dir}' + 'mean_train_features_dynamic_month.nc')
aux = ds_train.groupby('time.month').std().compute()
aux.to_netcdf(f'{input_dir}' + 'std_train_features_dynamic_month.nc')
ds_train.groupby('time.month').std().to_netcdf(f'{input_dir}' + 'std_train_features_dynamic_month.nc')
z500_train = z500.sel(time=slice(*train_years))
aux_z500 = z500.groupby('time.month') - z500_train.groupby('time.month').mean()
aux_z500 = aux_z500.groupby('time.month') / z500_train.groupby('time.month').std()
aux_z500.compute()
aux_z500.to_netcdf(input_dir +'geopotential_500/month/'+ 'monthly_standardized_data.nc')
t850_train = t850.sel(time=slice(*train_years))
aux_t850 = t850.groupby('time.month') - t850_train.groupby('time.month').mean()
aux_t850 = aux_t850.groupby('time.month') / t850_train.groupby('time.month').std()
aux_t850.compute()

aux_t850.to_netcdf(input_dir +'temperature_850/month/'+ 'monthly_standardized_data.nc')
rad_train = rad.sel(time=slice(*train_years))
aux_rad = rad.groupby('time.month') - rad_train.groupby('time.month').mean()
aux_rad = aux_rad.groupby('time.month') / rad_train.groupby('time.month').std()

aux_rad.compute()
aux_rad.to_netcdf(input_dir +'toa_incident_solar_radiation/month/'+ 'monthly_standardized_data.nc')
z500_train = z500.sel(time=slice(*train_years))
aux_z500 = z500.groupby('time.month') - z500_train.groupby('time.month').mean()
aux_z500 = aux_z500.groupby('time.month') / z500_train.groupby('time.month').std()
def get_standarization_params(ds, mean='mean', std='std', time='point'):
    """
    Compute parameters for standardizing data.
    
    Parameters
    ----------
    ds : xarray Dataset
        Dataset containing the input data
    mean: str
    std:  str
    time: str
    """
    if mean.lower() not in ['mean', 'median']:
        raise ValueError("mean parameter not valid. \
        It must be: \
        \n-'mean' or \
        \n-'median'")
    if std.lower() not in ['std','iqr']:
        raise ValueError("std parameter not valid. \
        It must be: \
        \n-'std' for Standard Deviation or \
        \n-'iqr' for Interquartile Range ")
    if time.lower() not in ['point','week', 'month']:
        raise ValueError("time parameter not valid. \
        It must be: \
        \n-'point' for statistic using complete dataset \
        \n-'week' for statistic using weekly moving window values\
        \n-'month' for statistic using monthly (30d) moving window values")
    
    if time == 'point':
        if mean == 'mean':
            ds_mean = ds.mean(('time','node')).compute()
        else:
            ds_mean = ds.median('time').mean('node').compute()
    
        if std == 'std':
            ds_std = ds.std(('time','node')).compute()
        else:
            q1, q3 = ds.quantile([0.25, 0.75], dim='time').compute().mean('node').to_array().values[0]
            ds_std = q3 - q1
    else:
        if time == 'week': t = 7  
        else: t=30
            
        if mean == 'mean':
            ds_mean = ds.rolling(time=24*t, center=True).mean().mean('node')
        else:
            ds_mean = ds.rolling(time=24*t, center=True).median().mean('node')
        if std == 'std':
            ds_std = ds.rolling(time=24*t, center=True).std().mean('node')
        else:
            quantiles = ds.rolling(time=24*t, center=True)\
            .construct('tmp').quantile([0.25, 0.75], dim='tmp').mean('node').to_array()
            ds_std = quantiles[:,1] - quantiles[:,0]
    
    return ds_mean, ds_std
z500_mean, z500_std = get_standarization_params(z500, 'mean', 'std', 'point')
z500_mean
t850_mean, t850_std = get_standarization_params(t850, 'mean', 'std', 'point')
t850_mean
toa_mean, toa_std = get_standarization_params(rad, 'mean', 'std', 'point')
z500_standardized = (z500 - z500_mean) / z500_std
z500_mean_week, z500_std_week = get_standarization_params(z500, 'mean', 'std', 'week')
z500_mean_month, z500_std_month = get_standarization_params(z500, 'mean', 'std', 'month')
z500_median, z500_iqr = get_standarization_params(z500, 'median', 'iqr', 'point')
z500_median_week, z500_iqr_week = get_standarization_params(z500, 'median', 'iqr', 'week')
z500_mean_std = (z500 - z500_mean)/z500_std
z500_mean_std_week = (z500 - z500_mean_week)/z500_std_week
z500_mean_std_week
z500_median_std_week = (z500 - z500_median_week)/z500_iqr_week<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16862496030485868, 'fraction_numerical': 0.05303270879644331, 'mean_word_length': 4.448961937716263, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '618549', 'license': 'MIT', 'n_tokens_mistral': 2556, 'n_tokens_neox': 2322, 'n_words': 398, 'path': 'notebooks/other_notebooks/standardize_data.ipynb', 'repo_name': 'natbolon/weather_prediction', 'size': '6298'}