<jupyter_start><jupyter_text>Load libraries<jupyter_code>%load_ext autoreload
%autoreload 2
import random
import time
import json
import numpy as np
import pandas as pd
import xarray as xr
import xskillscore as xs

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

from scipy.stats import norm
from xskillscore import crps_quadrature, crps_ensemble

import sys
sys.path.append('/'.join(sys.path[0].split('/')[:-1]))
from modules.architectures import *

from modules.swag import data, models, utils, losses
from modules.swag.posteriors import SWAG

import modules.architectures as modelArchitectures
from modules.full_pipeline import compute_errors, _inner
from modules.plotting import plot_intervalmap, plot_signal, plot_rmses
from modules.data import hp_to_equiangular
from modules.test import compute_rmse_healpix, compute_R2, compute_rmse_equiangular, compute_weighted_rmse
config_file = 'config_s01_two_steps_small.json'
ep = 11
lead_time = 6

# load config
with open("../configs/" + config_file) as json_data_file:
    cfg = json.load(json_data_file)

# define paths
datadir = cfg['directories']['datadir']
input_dir = datadir + cfg['directories']['input_dir']
model_save_path = datadir + cfg['directories']['model_save_path']
pred_save_path = datadir + cfg['directories']['pred_save_path']
metrics_path = datadir + cfg['directories']['metrics_path']

# model parameters
max_lead_time = cfg['training_constants']['max_lead_time']
len_sqce = cfg['model_parameters']['len_sqce']
delta_t = cfg['model_parameters']['delta_t']
architecture_name = cfg['model_parameters']['architecture_name']
resolution = cfg['model_parameters']['resolution']

lead_times = np.arange(lead_time, max_lead_time+lead_time, lead_time)
def get_spread(description_member, description_member_median, nb_models=10, swag=False):
    ensemble = []
    description = description_member
    for i in range(nb_models):
        if swag:
            description = description.replace(f"realization{i}", f"realization{i+1}")
        else:
            description = description.replace(f"model{i}", f"model{i+1}")
        pred = xr.open_dataset(pred_save_path + description + '.nc').drop('lat').drop('lon').chunk('auto')
        ensemble.append(pred)
    
    ensemble_median = xr.open_dataset(pred_save_path + description_member_median + '.nc').drop('lat').drop('lon').chunk('auto')
    
    diff = [abs(member - ensemble_median) for member in ensemble]
    diff = xr.concat(diff, dim='member')
    
    dims = list(diff.dims).copy()
    dims.remove('lead_time')
    
    diff_mean = [diff['z'].mean(dims), diff['t'].mean(dims)]
    diff_mean = xr.merge(diff_mean).load()
    
    del ensemble
    del ensemble_median
    
    return diff_mean<jupyter_output><empty_output><jupyter_text>Deep Ensemble<jupyter_code># description_epoch + f'_s02_swag_freq10_epoch_11_swag_scale001_median',
# description_epoch + f'_s02_swag_freq10_epoch_11_swag_scale01_mean',
# description_epoch + f'_s02_swag_freq10_epoch_11_swag_scale03_median',
# description_epoch + f'_s02_multiswag_ensemble_epoch_11_swag_scale00_median',
# description_epoch + f'_s02_multiswag_ensemble_epoch_11_swag_scale01_aggregate_median',
compute = True
obs = xr.open_mfdataset(pred_save_path + 'observations_nearest.nc', combine='by_coords')

description_epoch = "all_const_len{}_delta_{}_architecture_".format(len_sqce, delta_t) + architecture_name

description_members = [
    description_epoch + f'_s02_ensemble_model0_epoch11',
    description_epoch + f'_s02_fixed_data_ensemble_ensemble_model0_epoch11'
]

descriptions_median = [
    description_epoch + f'_s02_ensemble_epoch_{ep}_median',
    description_epoch + f'_s02_fixed_data_ensemble_ensemble_epoch_{ep}_median'
]

if compute:
    spread_rand = get_spread(description_members[0], descriptions_median[0])
    spread_fixed = get_spread(description_members[1], descriptions_median[1])
if compute:
    spread_rand.to_netcdf(datadir + 'metrics/spread_' + descriptions_median[0].replace('_median', '') + '.nc')
    spread_fixed.to_netcdf(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
else:
    spread_rand = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[0].replace('_median', '') + '.nc')
    spred_fixed = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
rmse_rand_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[0] + '.nc')
rmse_fixed_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[1] + '.nc')
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), constrained_layout=True)

ax1.plot(lead_times, rmse_rand_median.z.values, label='RMSE Random split Median')
ax2.plot(lead_times, rmse_rand_median.t.values, label='RMSE Random split Median')

ax1.plot(lead_times, spread_rand.z.values, linestyle='dashed', label='Random split')
ax2.plot(lead_times, spread_rand.t.values, linestyle='dashed', label='Random split')

ax1.plot(lead_times, rmse_fixed_median.z.values, label='RMSE Fixed split Median')
ax2.plot(lead_times, rmse_fixed_median.t.values, label='RMSE Fixed split Median')

ax1.plot(lead_times, spread_fixed.z.values, linestyle='dashed', label='Fixed split')
ax2.plot(lead_times, spread_fixed.t.values, linestyle='dashed', label='Fixed split')

ax1.set_xlabel('Lead time (h)')
ax1.set_xticks(lead_times)
ax2.set_xticks(lead_times)
ax2.set_xlabel('Lead time (h)')
ax1.set_ylabel('Spread')
ax2.set_ylabel('Spread')
ax1.set_title('Z500')
ax2.set_title('T850')

f.suptitle('Comparison of ensemble spread for Deep Ensemble experiments', 
           fontsize=20)
ax1.legend()
ax2.legend()
plt.show()<jupyter_output><empty_output><jupyter_text>MultiSWAG<jupyter_code>compute = True
obs = xr.open_mfdataset(pred_save_path + 'observations_nearest.nc', combine='by_coords')

description_epoch = "all_const_len{}_delta_{}_architecture_".format(len_sqce, delta_t) + architecture_name

description_members_multiswag = [
    description_epoch + f'_s02_multiswag_ensemble_model1_epoch11_swag_scale00',
    description_epoch + f'_s02_multiswag_ensemble_model1_epoch11_swag_scale01_aggregate'
]

descriptions_median_multiswag = [
    description_epoch + f'_s02_multiswag_ensemble_epoch_11_swag_scale00_median',
    description_epoch + f'_s02_multiswag_ensemble_epoch_11_swag_scale01_aggregate_median'
]

if compute:
    spread_multiswa = get_spread(description_members_multiswag[0], descriptions_median_multiswag[0])
    spread_multiswa.to_netcdf(datadir + 'metrics/spread_' + descriptions_median_multiswag[0].replace('_median', '') + '.nc')
    
    spread_multiswag = get_spread(description_members_multiswag[1], descriptions_median_multiswag[1])
    spread_multiswag.to_netcdf(datadir + 'metrics/spread_' + descriptions_median_multiswag[1].replace('_median', '') + '.nc')
else:
    spread_multiswa = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median_multiswag[0].replace('_median', '') + '.nc')
    spread_multiswag = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median_multiswag[1].replace('_median', '') + '.nc')
rmse_multiswa_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median_multiswag[0] + '.nc')
rmse_multiswag_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median_multiswag[1] + '.nc')
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), constrained_layout=True)

ax1.plot(lead_times, rmse_multiswa_median.z.values, label='RMSE MultiSWA Median')
ax2.plot(lead_times, rmse_multiswa_median.t.values, label='RMSE MultiSWA Median')

ax1.plot(lead_times, spread_multiswa.z.values, linestyle='dashed', label='Spread MultiSWA')
ax2.plot(lead_times, spread_multiswa.t.values, linestyle='dashed', label='Spread MultiSWA')

ax1.plot(lead_times, rmse_multiswag_median.z.values, label='RMSE MultiSWAG Scale 0.1 Median')
ax2.plot(lead_times, rmse_multiswag_median.t.values, label='RMSE MultiSWAG Scale 0.1 Median')

ax1.plot(lead_times, spread_multiswag.z.values, linestyle='dashed', label='Spread MultiSWAG Scale 0.1')
ax2.plot(lead_times, spread_multiswag.t.values, linestyle='dashed', label='Spread MultiSWAG Scale 0.1')

ax1.set_xlabel('Lead time (h)')
ax1.set_xticks(lead_times)
ax2.set_xticks(lead_times)
ax2.set_xlabel('Lead time (h)')
ax1.set_ylabel('Spread')
ax2.set_ylabel('Spread')
ax1.set_title('Z500')
ax2.set_title('T850')

f.suptitle('Comparison of ensemble spread for MultiSWAG experiments', 
           fontsize=20)
ax1.legend()
ax2.legend()
plt.show()<jupyter_output><empty_output><jupyter_text>SWAG  Scale 0.1<jupyter_code>compute = True
obs = xr.open_mfdataset(pred_save_path + 'observations_nearest.nc', combine='by_coords')

description_epoch = "all_const_len{}_delta_{}_architecture_".format(len_sqce, delta_t) + architecture_name

description_members = [
    description_epoch + f'_s02_ensemble_model0_epoch11',
    description_epoch + f'_s02_swag_freq10_epoch11_swag_scale01_realization0'
]

descriptions_median = [
    description_epoch + f'_s02_ensemble_epoch_11_median',
    description_epoch + f'_s02_swag_freq10_epoch_11_swag_scale01_median'
]

if compute:
    spread_deep_ensemble = get_spread(description_members[0], descriptions_median[0])
    spread_swag_scale01 = get_spread(description_members[1], descriptions_median[1], swag=True)
if compute:
    spread_deep_ensemble.to_netcdf(datadir + 'metrics/spread_' + descriptions_median[0].replace('_median', '') + '.nc')
    spread_swag_scale01.to_netcdf(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
else:
    spread_deep_ensemble = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[0].replace('_median', '') + '.nc')
    spread_swag_scale01 = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
rmse_deep_ensemble_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[0] + '.nc')
rmse_swag_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[1] + '.nc')
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), constrained_layout=True)

ax1.plot(lead_times, rmse_deep_ensemble_median.z.values, label='RMSE Deep Ensemble Median')
ax2.plot(lead_times, rmse_deep_ensemble_median.t.values, label='RMSE Deep Ensemble Median')

ax1.plot(lead_times, spread_deep_ensemble.z.values, linestyle='dashed', label='Deep Ensemble')
ax2.plot(lead_times, spread_deep_ensemble.t.values, linestyle='dashed', label='Deep Ensemble')

ax1.plot(lead_times, rmse_swag_median.z.values, label='RMSE SWAG Scale 0.1 Median')
ax2.plot(lead_times, rmse_swag_median.t.values, label='RMSE SWAG Scale 0.1 Median')

ax1.plot(lead_times, spread_swag_scale01.z.values, linestyle='dashed', label='SWAG Scale 0.1')
ax2.plot(lead_times, spread_swag_scale01.t.values, linestyle='dashed', label='SWAG Scale 0.1')

ax1.set_xlabel('Lead time (h)')
ax1.set_xticks(lead_times)
ax2.set_xticks(lead_times)
ax2.set_xlabel('Lead time (h)')
ax1.set_ylabel('Spread')
ax2.set_ylabel('Spread')
ax1.set_title('Z500')
ax2.set_title('T850')

f.suptitle('Comparison of ensemble spread for SWAG Scale 0.1 and Deep Ensemble experiments', 
           fontsize=20)
ax1.legend()
ax2.legend()
plt.show()
diff_swag = rmse_swag_median - spread_swag_scale01
diff_deep_ensemble = rmse_deep_ensemble_median - spread_deep_ensemble 
dfs = [diff_swag['z'].to_dataframe(), diff_deep_ensemble['z'].to_dataframe(),
       diff_swag['t'].to_dataframe(), diff_deep_ensemble['t'].to_dataframe()]
df_diff = pd.concat(dfs, axis=1)
df_diff.columns = [ 'Z500 SWAG Scale 0.1', 'Z500 DeepEnsemble', 'T850 SWAG Scale 0.1', 'T850 DeepEnsemble']
df_diff<jupyter_output><empty_output><jupyter_text>Scale 0.3<jupyter_code>compute = True
obs = xr.open_mfdataset(pred_save_path + 'observations_nearest.nc', combine='by_coords')

description_epoch = "all_const_len{}_delta_{}_architecture_".format(len_sqce, delta_t) + architecture_name

description_members = [
    description_epoch + f'_s02_ensemble_model0_epoch11',
    description_epoch + f'_s02_swag_freq10_epoch11_swag_scale03_realization0'
]

descriptions_median = [
    description_epoch + f'_s02_ensemble_epoch_11_median',
    description_epoch + f'_s02_swag_freq10_epoch_11_swag_scale03_median'
]

spread_deep_ensemble = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[0].replace('_median', '') + '.nc')

if compute:
    spread_swag_scale03 = get_spread(description_members[1], descriptions_median[1], swag=True)
if compute:
    spread_swag_scale03.to_netcdf(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
else:
    spread_swag_scale03 = xr.open_dataset(datadir + 'metrics/spread_' + descriptions_median[1].replace('_median', '') + '.nc')
rmse_deep_ensemble_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[0] + '.nc')
rmse_swag_03_median = xr.open_dataset(datadir + 'metrics/rmse_' + descriptions_median[1] + '.nc')
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), constrained_layout=True)

ax1.plot(lead_times, rmse_deep_ensemble_median.z.values, label='RMSE Deep Ensemble Median')
ax2.plot(lead_times, rmse_deep_ensemble_median.t.values, label='RMSE Deep Ensemble Median')

ax1.plot(lead_times, spread_deep_ensemble.z.values, linestyle='dashed', label='Deep Ensemble')
ax2.plot(lead_times, spread_deep_ensemble.t.values, linestyle='dashed', label='Deep Ensemble')

ax1.plot(lead_times, rmse_swag_03_median.z.values, label='RMSE SWAG Scale 0.3 Median')
ax2.plot(lead_times, rmse_swag_03_median.t.values, label='RMSE SWAG Scale 0.3 Median')

ax1.plot(lead_times, spread_swag_scale03.z.values, linestyle='dashed', label='SWAG Scale 0.3')
ax2.plot(lead_times, spread_swag_scale03.t.values, linestyle='dashed', label='SWAG Scale 0.3')

ax1.set_xlabel('Lead time (h)')
ax1.set_xticks(lead_times)
ax2.set_xticks(lead_times)
ax2.set_xlabel('Lead time (h)')
ax1.set_ylabel('Spread')
ax2.set_ylabel('Spread')
ax1.set_title('Z500')
ax2.set_title('T850')

f.suptitle('Comparison of ensemble spread for SWAG Scale 0.3 and Deep Ensemble experiments', 
           fontsize=20)
ax1.legend()
ax2.legend()
plt.show()
diff_swag = rmse_swag_median - spread_swag_scale03
diff_deep_ensemble = rmse_deep_ensemble_median - spread_deep_ensemble 
dfs = [diff_swag['z'].to_dataframe(), diff_deep_ensemble['z'].to_dataframe(),
       diff_swag['t'].to_dataframe(), diff_deep_ensemble['t'].to_dataframe()]
df_diff = pd.concat(dfs, axis=1)
df_diff.columns = [ 'Z500 SWAG Scale 0.3', 'Z500 DeepEnsemble', 'T850 SWAG Scale 0.3', 'T850 DeepEnsemble']
df_diff<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13108228673392355, 'fraction_numerical': 0.02793219408413973, 'mean_word_length': 7.149888143176733, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '253489', 'license': 'MIT', 'n_tokens_mistral': 5763, 'n_tokens_neox': 5423, 'n_words': 811, 'path': 'notebooks/exploration_ensemble_spread.ipynb', 'repo_name': 'yasserhaddad/weather_prediction', 'size': '14571'}