# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: NPL-3.7.9
#     language: python
#     name: npl-3.7.9
# ---

# # Figures using detrended and resampled data
#
# ### Author: <NAME>, [github](https://github.com/chrisrwp/synthetic-ensemble/SIA/Resampled_figures.ipynb) <br>
#
# #### Input:
# - Detrended models and observations (both relative to the ensemble/average trend and the individual trends)
# - Resampled detrended anomalies
#
# #### Output figures:
# - Example plot of resampling observations
# - Comparison of the distribution of anomalies when detrending is realtive to the ensemble verses individual 
# - Ratio plots of $\sigma_{mem}$ to $\sigma_{LE}$ and $\sigma_{mem}$ to $\sigma_{obs}$
# - **Figures 1, 5 (A1), 6, 10**

import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import matplotlib.patheffects as pe
import seaborn as sns
import xarray as xr
import numpy as np
import math as m
import scipy.stats as stats

# ## Load data

# +
data_path = '/glade/scratch/cwpowell/Synthetic_ensemble/'

month_names = ['January', 'February', 'March', 'April', 'May', 'June', 'July', 
               'August', 'September', 'October', 'November', 'December']

month_names_short = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']

model_names       = ['CanESM2', 'CESM1', 'CSIRO_MK36', 'GFDL_CM3', 'GFDL_ESM2M', 'MPI_ESM1'] #for loading areas
model_print_names = ['CanESM2', 'CESM1', 'CSIRO MK36', 'GFDL CM3', 'GFDL ESM2M', 'MPI ESM1'] #for printing on graphs

member_numbers = [50, 40, 30, 20, 30, 100]

upper_quart = np.array(member_numbers)*0.75
lower_quart = np.array(member_numbers)*0.25

        # CanESM2, CESM1, CSIRO MK36, GFDL CM3, GFDL ESM2M, MPI ESM2
colors = ['m',     'b',   'g',        'orange',  'k',        'tab:olive' ]
obs_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']

# +
#load SIA and SIE files
SIA_SIE_CDR = xr.open_dataset('/glade/scratch/cwpowell/Synthetic_ensemble/Raw_data/observations/NSIDC_CDR_v4/SIA_SIE_CDR_BT_NT_79-20_filled.nc')

#load detrended files 
CLIVAR_detrended_ensemble   = xr.open_dataset(data_path+'SIA/SIA_detrended/CLIVAR_SIA_detrended_ensemble_79-20.nc')
CLIVAR_detrended_individual = xr.open_dataset(data_path+'SIA/SIA_detrended/CLIVAR_SIA_detrended_individual_79-20.nc')

Obs_detrended_average    = xr.open_dataset(data_path+'SIA/SIA_detrended/Obs_SIA_detrended_average_79-20.nc')
Obs_detrended_individual = xr.open_dataset(data_path+'SIA/SIA_detrended/Obs_SIA_detrended_individual_79-20.nc')
# -

#resampled SIA
Obs_SIA_resampled_ind = xr.open_dataset('/glade/scratch/cwpowell/Synthetic_ensemble/SIA/SIA_resampled/Obs_SIA_resampled_individual_79-20.nc')

# +
#load sigma files
# sigma_LE_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_LE_individual_79-20.nc')
# sigma_LE_ensemble   = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_LE_ensemble_79-20.nc')

# sigma_mem_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_mem_individual_79-20.nc')
# sigma_mem_ensemble   = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_mem_ensemble_79-20.nc')

# sigma_obs_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_obs_individual_79-20.nc')
# sigma_obs_average    = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_obs_average_79-20.nc')
# sigma_obs_jackknife  = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_obs_jackknife_79-20.nc')

# +
#load sigma files resampled 10000 times
sigma_LE_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_LE_individual_79-20.nc')
sigma_mem_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_mem_individual_10000.nc')
sigma_obs_individual = xr.open_dataset(data_path+'SIA/SIA_resampled/Sigma_obs_individual_10000.nc')

mu_mem = xr.open_dataset(data_path+'SIA/SIA_resampled/Mu_mem_individual_10000.nc')
mu_obs = xr.open_dataset(data_path+'SIA/SIA_resampled/Mu_obs_individual_10000.nc')


# -

# # Example of resampling with 5 part plot

def resample_boot2(time_period, data, resamp_n):
    '''
    Resample a 1D time series using a 2 year block boostrap size with replacement
    
    Parameters
    ----------
    time_period : integer,
        For 1979-2020 use 42 as the total number of years in that time period
    data : 1 dimensional xarray dataarray,
        For 1979-2020 this is an array of shape [42] 
    resamp_n : int
        The number of times to resample the data
    
    Returns
    ----------
        2D xarray dataarray object of n resamplings of the input data, shape: (time_period, 1000)
    '''  
    #create an xarray dataarray of indexes for half the length of the time period, year_i coordinates 1,3,5...
    boot_2_first_ind = xr.DataArray(data   = np.random.randint(0,time_period-2, (resamp_n,int(time_period/2))), 
                                    coords = {'resampling':np.arange(1,resamp_n+1), 'year_i':np.arange(1,time_period+1,2)},
                                    dims   = ['resampling', 'year_i'])

    #create an identical dataarray but with each element incremented by 1, year_i coordinates 2,4,6....
    boot_2_second_ind = (boot_2_first_ind+1).copy()
    boot_2_second_ind['year_i'] = np.arange(2,time_period+2,2)

    #concatenate the two arrays with the coordinates in order, this allows a 2 year block boostrap size
    all_boot_2_ind = xr.concat((boot_2_first_ind, boot_2_second_ind), dim='year_i').sortby('year_i')
    
    #create an array with the starting element of the flattened array for each resampling 0, 42, 84...
    ind_base = np.repeat(np.arange(0,time_period*resamp_n,time_period),time_period)
    
    #add together the base indexes (0,42,84...) with the randomly chosen indexes within the original data
    ind_1_d = np.ravel(all_boot_2_ind) + ind_base
    
    #copy the original data 1000 times as a 1D array so it will have the same indexes as we just made for ind_1_d
    data_1000 = np.ravel(np.tile(data,(time_period,resamp_n)))
    
    #select the randomly generated indexes from the flattened copied original data, reshape and save to xarray dataarray

    resampled_boot_2 = xr.DataArray(data = np.reshape(data_1000[ind_1_d], (resamp_n, time_period)),
                                    coords = {'resampling':np.arange(1,resamp_n+1), 'year_i':np.arange(1,time_period+1,1)},
                                    dims   = ['resampling', 'year_i'])

    return(resampled_boot_2, all_boot_2_ind, ind_1_d, data_1000)


# +
#do caclulations for plot
sep_SIA_CDR = SIA_SIE_CDR['CDR_SIA'].sel(time=SIA_SIE_CDR['time.month']==9)
yrs = np.arange(1979,2021)
sep_SIA_CDR['time'] = yrs
sep_SIA_CDR = sep_SIA_CDR.rename({'time':'year'})

coef_true = np.polyfit(yrs, sep_SIA_CDR ,1)
lin_trend = xr.DataArray(data=coef_true[0]*yrs+coef_true[1], coords={'year':yrs}, dims=['year'])
anoms_trend = sep_SIA_CDR - lin_trend

resamp_SIA_anoms = resample_boot2(42, anoms_trend, 1000)

all_resamp = resamp_SIA_anoms[0].std('year_i')
# -

# ## Figure 1 as submitted

# +
#make plot for september
fig = plt.figure(constrained_layout=True, figsize=[12,7.5])
gs = fig.add_gridspec(3,6)

ax0 = fig.add_subplot(gs[0,:-3])
ax1 = fig.add_subplot(gs[0,-3:])
ax2 = fig.add_subplot(gs[1,:-3])
ax3 = fig.add_subplot(gs[2,:-3])
ax4 = fig.add_subplot(gs[1:,-3:])

######################## plot SIA and the linear trend on axes [0,0] #############################
ax0.plot(yrs, sep_SIA_CDR, c='r', linewidth=1.5, alpha=1) #linewidth=0.75, alpha=0.75)
ax0.plot(yrs, coef_true[0]*yrs+coef_true[1], c='0.5',linewidth=2, linestyle='--')

ax0.set_xlim(1979,2020)
ax0.set_xlabel('Year', fontsize=12)
ax0.set_xticklabels(np.arange(1975, 2021, 5), fontsize=12);
ax0.set_ylim(2.7,7.3)
ax0.set_ylabel(r'$September \ SIA \ [10^6 \ km^2]$', fontsize=12)
ax0.set_yticklabels(np.arange(1,8,1), fontsize=12);

############################ plot anomalies of SIA on axes [0,1] ##################################
ax1.bar(yrs, anoms_trend, color=plt.get_cmap("jet")(np.linspace(0,1,42)), edgecolor='k')
# for i in range(len(yrs)):
#     ax1.text(yrs[i], 1.35, np.char.mod('%s', np.arange(0,43))[i], fontsize=6)

ax1.text(1980, 1, r'$\sigma = {}$'.format(np.round(np.std(anoms_trend.values),3)), fontsize=14)
   
ax1.axhline(0, color='k', linewidth=0.5);

ax1.set_xlim(1978.3,2020.7)
ax1.set_xlabel('Year', fontsize=12)
ax1.set_xticklabels(np.arange(1975, 2021, 5), fontsize=12);
ax1.set_ylim(-1.37,1.3);
ax1.set_ylabel(r'$SIA \ Anomalies \ [10^6 \ km^2]$', fontsize=12);
# ax1.set_yticklabels(np.arange(-1.5, 1.5, 0.5), fontsize=12);

###
ax2.bar(yrs, resamp_SIA_anoms[0].sel(resampling=1), color=plt.get_cmap("jet")(resamp_SIA_anoms[1].sel(resampling=1)/42), edgecolor='k')
# for i in range(len(yrs)):
#     ax2.text(yrs[i], 1.35, resamp_SIA_anoms[1].sel(resampling=1)[i].values, fontsize=6)

ax2.text(1980, 1, r'$\sigma = {}$'.format(np.round(np.std(resamp_SIA_anoms[0].sel(resampling=1)).values,3)), fontsize=14)

ax2.axhline(0, color='k', linewidth=0.5);

ax2.set_xlim(1978.3,2020.7)
ax2.set_xlabel('Year', fontsize=12)
ax2.set_xticklabels(np.arange(1975, 2021, 5), fontsize=12);
ax2.set_ylim(-1.37,1.3);
ax2.set_ylabel(r'$SIA \ Anomalies \ [10^6 \ km^2]$', fontsize=12);
# ax2.set_yticklabels(np.arange(-1.5, 1.5, 0.5), fontsize=12);

###

ax3.bar(yrs, resamp_SIA_anoms[0].sel(resampling=2), color=plt.get_cmap("jet")(resamp_SIA_anoms[1].sel(resampling=2)/42), edgecolor='k')
# for i in range(len(yrs)):
#     ax3.text(yrs[i], 1.35, resamp_SIA_anoms[1].sel(resampling=2)[i].values, fontsize=6)

ax3.text(1980, 1, r'$\sigma = {}$'.format(np.round(np.std(resamp_SIA_anoms[0].sel(resampling=2)).values,3)), fontsize=14)

ax3.axhline(0, color='k', linewidth=0.5);

ax3.set_xlim(1978.3,2020.7)
ax3.set_xlabel('Year', fontsize=12)
ax3.set_xticklabels(np.arange(1975, 2021, 5), fontsize=12);
ax3.set_ylim(-1.37,1.3);
ax3.set_ylabel(r'$SIA \ Anomalies \ [10^6 \ km^2]$', fontsize=12);
# ax3.set_yticklabels(np.arange(-1.5, 1.5, 0.5), fontsize=12);

############ normal distribution of resampled standard deviations with time on axes [1,2] #############
#calculate data for PDF of anomalies
distrib = stats.norm(all_resamp.mean('resampling'), all_resamp.std('resampling'))
x_vals  = np.linspace(all_resamp.mean('resampling')-all_resamp.std('resampling')*5, all_resamp.mean('resampling')+all_resamp.std('resampling')*5,100)
probs   = [distrib.pdf(x) for x in x_vals]

ax4.hist(all_resamp, bins=np.arange(0.2, 0.6, 0.01), density=True, color='0.5', edgecolor='k', alpha=0.6);
ax4.plot(x_vals, probs, c='k', linewidth=2)#, path_effects=[pe.Stroke(linewidth=2.5, foreground='0'), pe.Normal()]);
ax4.axvline(np.std(anoms_trend), c='r', linewidth=2, linestyle='--');

#add annotations of SD, skew and kurtosis
ax4.annotate(r'$\sigma_\sigma={}$'.format(np.round(all_resamp.std('resampling').values,3)),
                              xy=(0.75, 0.87), xycoords='axes fraction', fontsize=14, color='0.5'); 
ax4.annotate(r'$\mu_\sigma={}$'.format(np.round(all_resamp.mean('resampling').values,3)),
                              xy=(0.75, 0.93), xycoords='axes fraction', fontsize=14, color='0.5');
ax4.annotate(r'$\sigma={}$'.format(np.round(np.std(anoms_trend).values,3)),
                              xy=(0.1, 0.93), xycoords='axes fraction', fontsize=14, color='r');

ax4.set_xlim(0.25,0.65);
ax4.set_xlabel(r'$Resampled \ Standard \ Deviations \ [10^6 \ km^2]$', fontsize=12);
ax4.set_xticklabels(np.round(np.arange(0.2,0.67,0.05),2), fontsize=12);
ax4.set_ylabel('Normalized Frequency', fontsize=12);
ax4.set_yticklabels(np.round(np.arange(0,15,2)), fontsize=12);


ax0.text(0.05, 1.15, 'a', transform=ax0.transAxes, fontsize=16, fontweight='bold', va='top', ha='right');
ax1.text(0.05, 1.15, 'b', transform=ax1.transAxes, fontsize=16, fontweight='bold', va='top', ha='right');
ax2.text(0.05, 1.15, 'c', transform=ax2.transAxes, fontsize=16, fontweight='bold', va='top', ha='right');
ax3.text(0.05, 1.15, 'd', transform=ax3.transAxes, fontsize=16, fontweight='bold', va='top', ha='right');
ax4.text(0.05, 1.05, 'e', transform=ax4.transAxes, fontsize=16, fontweight='bold', va='top', ha='right');


# #save the figures
fig.savefig(data_path+'SIA/figures/Resampled_figures/Sep_SIA_5_part_resampling_example.pdf')
fig.savefig(data_path+'SIA/figures/Resampled_figures/Sep_SIA_5_part_resampling_example.png', dpi=400)


# -

# # Plot individual and ensemble detrending model anomalies distribution differences

def x_vals_probs_PDF(data, mean_0=False):
    if mean_0 == True:
        mean_ = 0
    else:
        mean_ = data.mean()
        
    distrib = stats.norm(mean_, data.std())
    x_vals  = np.linspace(mean_-data.std()*5, mean_+data.std()*5,100)
    probs   = [distrib.pdf(x) for x in x_vals]
    
    return(x_vals, probs)


# +
#plot a single month, distributions of anomalies individual and ensemble detrended
bins_ = np.arange(-1.85,1.86,0.1)

for month_ in [9]:#np.arange(1,13):
    
    fig = plt.figure(figsize=[10,6])

    for model_i, model_name in enumerate(model_names):
        data = CLIVAR_detrended_ensemble[model_name].sel(member=slice(0,member_numbers[model_i]+1)).sel(time=CLIVAR_detrended_ensemble['time.month']==month_)
        plt.plot(x_vals_probs_PDF(data)[0], x_vals_probs_PDF(data)[1], c=colors[model_i], label=model_print_names[model_i]);

        data = CLIVAR_detrended_individual[model_name].sel(member=slice(0,member_numbers[model_i]+1)).sel(time=CLIVAR_detrended_ensemble['time.month']==month_)
        plt.plot(x_vals_probs_PDF(data)[0], x_vals_probs_PDF(data)[1], c=colors[model_i], linestyle='--');

    plt.ylim(bottom=0.01);
    plt.ylabel('Normalized frequency', fontsize=14)
    plt.xlabel(r'$SIA \ Anomaly \ [10^6 \ km^2]$', fontsize=14)
    plt.xlim(-1.75, 1.75);
    plt.legend();
    plt.title('Anomalies, {}, for individual and ensemble detrended'.format(month_names[month_-1]), fontsize=16);

#     fig.savefig(data_path+'SIA/figures/Anomalies_comparison_ind_ens_{}.png'.format(str(month_).zfill(2)), dpi=400)
    
#     plt.close('all')
# -

# # Plot $\sigma_{LE}$, $\sigma_{mem}$, $\sigma_{obs}$ differences individual and ensemble detrending

# +
# #plot individual member differences
# (sigma_mem_ensemble['CESM1'] - sigma_mem_individual['CESM1']).plot()
# plt.xlim(0.5,40.5);

# +
#plot absolute values of sigma and the differences ind verses ens
fig, axes = plt.subplots(2,3,figsize=[15.5,7])

################# absolute value comparisons ###############
#sigma_LE
for model_i, model_name in enumerate(model_names):
    axes[0][0].plot(month_names_short, sigma_LE_ensemble[model_name], label=model_name, c=colors[model_i])
    axes[0][0].plot(month_names_short, sigma_LE_individual[model_name], label=model_name, c=colors[model_i], linestyle='--')
axes[0][0].set_ylabel(r'$\sigma_{LE}$', fontsize=15)    
axes[0][0].set_xlabel('')
axes[0][0].set_xlim(0,11)
axes[0][0].set_ylim(0.02,0.1)
    
#sigma_mem
for model_i, model_name in enumerate(model_names):
    axes[0][1].plot(month_names_short, sigma_mem_ensemble[model_name].mean('member'), label=model_name, c=colors[model_i])
    axes[0][1].plot(month_names_short, sigma_mem_individual[model_name].mean('member'), label=model_name, c=colors[model_i], linestyle='--')
axes[0][1].set_ylabel(r'$\sigma_{mem}$', fontsize=15)
axes[0][1].set_xlabel('')
axes[0][1].set_xlim(0,11)
axes[0][1].set_ylim(0.02,0.1)

#sigma_obs
for i, dataset in enumerate(list(sigma_obs_average.keys())):
    axes[0][2].plot(month_names_short, sigma_obs_jackknife[dataset].mean('jackknife'), label=dataset, c=obs_colors[i])
    axes[0][2].plot(month_names_short, sigma_obs_individual[dataset], label=dataset, c=obs_colors[i], linestyle='--')
axes[0][2].set_ylabel(r'$\sigma_{obs}$', fontsize=15)
axes[0][2].set_xlabel('')
axes[0][2].set_xlim(0,11)
axes[0][2].set_ylim(0.02,0.1)


############### relative value comparisons #################
#sigma_LE
for model_i, model_name in enumerate(model_names):
    axes[1][0].plot(month_names_short, ((sigma_LE_ensemble[model_name] - sigma_LE_individual[model_name]) / sigma_LE_individual[model_name])*100, 
                    label=model_name, c=colors[model_i])
axes[1][0].set_ylabel('Ensemble - Individual [%]', fontsize=13)
axes[1][0].set_xlabel('Month', fontsize=13)
axes[1][0].set_xlim(0,11)
axes[1][0].axhline(0, c='0.3', linestyle='--')

#sigma_mem
for model_i, model_name in enumerate(model_names):
    axes[1][1].plot(month_names_short, ((sigma_mem_ensemble[model_name].mean('member') - sigma_mem_individual[model_name].mean('member')) / sigma_mem_individual[model_name].mean('member'))*100, 
                    label=model_name, c=colors[model_i])
axes[1][1].set_ylabel('Ensemble - Individual [%]', fontsize=13)
axes[1][1].set_xlabel('Month', fontsize=13)
axes[1][1].set_xlim(0,11)
axes[1][1].axhline(0, c='0.3', linestyle='--')

#sigma_obs
for i, dataset in enumerate(list(sigma_obs_average.keys())):
    axes[1][2].plot(month_names_short, ((sigma_obs_jackknife[dataset].mean('jackknife') - sigma_obs_individual[dataset]) / sigma_obs_individual[dataset])*100, label=dataset)
axes[1][2].set_ylabel('Jackknife - Individual [%]', fontsize=13)
axes[1][2].set_xlabel('Month', fontsize=13)
axes[1][2].set_xlim(0,11)
axes[1][2].axhline(0, c='0.3', linestyle='--')

################ set legend ##################
legend_elements = [Patch(facecolor=colors[0], label='CanESM2'),
                   Patch(facecolor=colors[1], label='CESM1'),
                   Patch(facecolor=colors[2], label='CSIRO MK36'),
                   Patch(facecolor=colors[3], label='GFDL CM3'),
                   Patch(facecolor=colors[4], label='GFDL ESM2M'),
                   Patch(facecolor=colors[5], label='MPI ESM1'),
                   Patch(facecolor=obs_colors[0], label=list(sigma_obs_average.keys())[0]),
                   Patch(facecolor=obs_colors[1], label=list(sigma_obs_average.keys())[1]),
                   Patch(facecolor=obs_colors[2], label=list(sigma_obs_average.keys())[2]),
                   Patch(facecolor=obs_colors[3], label=list(sigma_obs_average.keys())[3]),
                   Patch(facecolor=obs_colors[4], label=list(sigma_obs_average.keys())[4]),
                   Patch(facecolor=obs_colors[5], label=list(sigma_obs_average.keys())[5])]

extra_legend = plt.legend(handles=legend_elements, bbox_to_anchor=(-0.72, -0.25), loc='upper center', borderaxespad=-0.25, ncol=6, fontsize=12)
# plt.tight_layout()

# fig.savefig(data_path+'SIA/figures/Sigma_LE_mem_obs_comparison.png', dpi=400)
# -

# # $\sigma_{mem}/\sigma_{LE}$  and $\sigma_{mem}/\sigma_{obs}$ Ratios

# +
# print mem/LE ratios
# sig_LE  = sigma_LE_individual.copy()
# sig_mem = sigma_mem_individual.copy()

# for model_i, model_name in enumerate(model_names):
#     print(model_name)
#     upper = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.floor(upper_quart[model_i])).max('member')
#     lower = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.ceil(lower_quart[model_i])).max('member')
    
#     print('mean', '\t', *(np.round(sig_mem[model_name].mean('member')/sig_LE[model_name],3).values), sep='\t')
#     print('median', '\t', *(np.round(sig_mem[model_name].median('member')/sig_LE[model_name],3).values), sep='\t')
#     print('max', '\t', *(np.round(sig_mem[model_name].max('member')/sig_LE[model_name],3).values), sep='\t')
#     print('min', '\t', *(np.round(sig_mem[model_name].min('member')/sig_LE[model_name],3).values), sep='\t')
#     print('UQ', '\t', *(np.round(upper/sig_LE[model_name], 3).values), sep='\t')
#     print('LQ', '\t', *(np.round(lower/sig_LE[model_name], 3).values), sep='\t')

# print mem/obs ratios
# sig_obs  = sigma_obs_individual['HadISST1'].copy()
# sig_mem = sigma_mem_individual.copy()

# for model_i, model_name in enumerate(model_names):
#     print(model_name)
#     upper = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.floor(upper_quart[model_i])).max('member')
#     lower = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.ceil(lower_quart[model_i])).max('member')
    
#     print('mean', '\t', *(np.round(sig_mem[model_name].mean('member')/sig_obs,3).values), sep='\t')
#     print('median', '\t', *(np.round(sig_mem[model_name].median('member')/sig_obs,3).values), sep='\t')
#     print('max', '\t', *(np.round(sig_mem[model_name].max('member')/sig_obs,3).values), sep='\t')
#     print('min', '\t', *(np.round(sig_mem[model_name].min('member')/sig_obs,3).values), sep='\t')
#     print('UQ', '\t', *(np.round(upper/sig_obs, 3).values), sep='\t')
#     print('LQ', '\t', *(np.round(lower/sig_obs, 3).values), sep='\t')
# -

# ## Fig 10 as submitted

# +
#plot sigma ratios with median and interquartile ranges
LE_mem = 1
ind_ens = 'ind'

if ind_ens == 'ind':
    sig_LE  = sigma_LE_individual.copy()
    sig_mem = sigma_mem_individual.copy()
    sig_obs = sigma_obs_individual.copy()
else:
    sig_LE  = sigma_LE_ensemble.copy()
    sig_mem = sigma_mem_ensemble.copy()
#     sig_obs = sigma_obs_average.copy()
    sig_obs = sigma_obs_jackknife.mean('jackknife').copy()


for obs_dataset in ['foo']:#list(sigma_obs_individual.keys()):    

    fig = plt.figure(figsize=[11.5,7])
    plt.axhline(1, c='0.5', linewidth=2)

    for model_i, model_name in enumerate(model_names):

        #define the upper and lower quarties of sigma_mem members
        upper = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.floor(upper_quart[model_i])).max('member')
        lower = sig_mem[model_name].where(sig_mem[model_name].rank('member')==m.ceil(lower_quart[model_i])).max('member')

        if LE_mem == 1: #plot LE/mem
            plt.plot(month_names_short, sig_mem[model_name].median('member')/sig_LE[model_name], 
                     c=colors[model_i], linewidth=2.5, marker='o', markersize=6)
#             plt.plot(month_names_short, sig_mem[model_name].mean('member')/sig_LE[model_name], 
#                      c=colors[model_i], linewidth=2.5, marker='o', markersize=6)

            plt.fill_between(month_names_short, upper/sig_LE[model_name], 
                             lower/sig_LE[model_name], color=colors[model_i], alpha=0.1)
#             plt.fill_between(month_names_short, sig_mem[model_name].max('member')/sig_LE[model_name], 
#                              sig_mem[model_name].min('member')/sig_LE[model_name], color=colors[model_i], alpha=0.1)

        else: #plot mem/obs
            plt.plot(month_names_short, sig_mem[model_name].median('member')/sig_obs[obs_dataset].values, 
                     c=colors[model_i], linewidth=2, marker='o')

            plt.fill_between(month_names_short, lower/sig_obs[obs_dataset].values, 
                             upper/sig_obs[obs_dataset].values, color=colors[model_i], alpha=0.1)

    #########################################################################

    plt.xticks(np.arange(0,12,1))
    plt.grid()
    plt.xlim(0,11)
    if LE_mem == 1:
        plt.ylim(0.41,1.15) #0.4,1.2 for interquartile range
#         plt.ylim(0.2,2.1) #0.4,1.2 for interquartile range
    else:
        plt.ylim(0.3,2.4) #0.3, 2.6 for interquartile range
    plt.xlabel('Month', fontsize=18)
    if LE_mem == 1:
        plt.ylabel('Resampled member relative to'+'\n'+r'Large Ensemble ($\bar{\sigma}_{mem} \ / \ \sigma_{LE}$)', fontsize=18)
    else:
        plt.ylabel('Model relative to observations', fontsize=18)
    plt.xticks(fontsize=16)
    plt.yticks(fontsize=16)

    legend_elements = [Patch(facecolor=colors[0], label='CanESM2'),
                       Patch(facecolor=colors[1], label='CESM1'),
                       Patch(facecolor=colors[2], label='CSIRO MK36'),
                       Patch(facecolor=colors[3], label='GFDL CM3'),
                       Patch(facecolor=colors[4], label='GFDL ESM2M'),
                       Patch(facecolor=colors[5], label='MPI ESM1'),]

    extra_legend = plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.24), loc='lower center', borderaxespad=-0.25, ncol=3, fontsize=14)

    plt.tight_layout()
    if LE_mem == 1:
        fig.savefig(data_path+'SIA/figures/Resampled_figures/Ratio_LE_mem_medians_iq_range_{}.png'.format(ind_ens), 
                    bbox_inches='tight', dpi=400)
        fig.savefig(data_path+'SIA/figures/Resampled_figures/Ratio_LE_mem_medians_iq_range_{}.pdf'.format(ind_ens), bbox_inches='tight')
    else:
        fig.savefig(data_path+'SIA/figures/Resampled_figures/Ratio_mem_obs_{}_medians_iq_range_{}.png'.format(obs_dataset, ind_ens), dpi=400)
#     plt.close('all')
# -

# ## Fig 5, A1 as submitted

# +
sigma_mu = 'mu'

if sigma_mu == 'sigma':
    mem = sigma_mem_individual.copy()
    obs = sigma_obs_individual.to_array().sel(variable=['CDR','BT','NT','HadISST1'])
else:
    mem = mu_mem.copy()
    obs = mu_obs.to_array().sel(variable=['CDR','BT','NT','HadISST1'])
    
fig = plt.figure(figsize=[14,6])

for model_i, model_name in enumerate(model_names):
    boxplot_model = np.empty((member_numbers[model_i],12))
    
    for memb_i, memb in enumerate(np.arange(1,member_numbers[model_i]+1)):    
        boxplot_model[memb_i] = mem[model_name].sel(member=memb).values
    
    b_plot = plt.boxplot(boxplot_model, widths=0.07, positions=np.arange(0.735+model_i*0.1,12.735+model_i*0.1,1), 
                         whis=(0,100), labels=month_names_short, patch_artist=True);
    
    #set the colors
    for patch in b_plot['boxes']: 
        patch.set_facecolor(colors[model_i]) 
    
    for element in ['boxes', 'whiskers', 'fliers', 'means', 'medians', 'caps']:
        plt.setp(b_plot[element], color=colors[model_i], linewidth=1.5)
        
    plt.setp(b_plot['medians'], color='0.8', linewidth=2)
    
marker_list = ['o',8,9,'d','s']
for var_i in range(4):
    for month_ in np.arange(1,13):
        plt.plot([-0.35+month_, 0.35+month_], [obs.isel(variable=var_i).sel(month=month_), obs.isel(variable=var_i).sel(month=month_)], c='r', alpha=1, zorder=9)
#     plt.scatter(np.arange(0.6,12.6), obs.isel(variable=var_i), c='r', marker=marker_list[var_i])

legend_elements = [Patch(facecolor=colors[0], label='CanESM2'),
                   Patch(facecolor=colors[1], label='CESM1'),
                   Patch(facecolor=colors[2], label='CSIRO MK3.6'),
                   Patch(facecolor=colors[3], label='GFDL CM3'),
                   Patch(facecolor=colors[4], label='GFDL ESM2M'),
                   Patch(facecolor=colors[5], label='MPI ESM1'),
                   Patch(facecolor='r', label='Observations')]

extra_legend = plt.legend(handles=legend_elements, bbox_to_anchor=(0.01, 0.97), loc='upper left', borderaxespad=0, ncol=2, fontsize=14)
plt.gca().add_artist(extra_legend);
    
plt.xticks(np.arange(1,13), fontsize=16)
plt.xlabel('Month', fontsize=18)
plt.yticks(fontsize=16)

if sigma_mu == 'sigma':
    plt.ylim(0.01,0.142) #for sigma
    plt.ylabel(r'$\bar{\sigma}_{mem} \ or \ \sigma_{obs} \ \ [10^6 \ km^2]$', fontsize=18);
else:
    plt.ylim(0.1,0.85) #for mu
    plt.ylabel(r'$\bar{\mu}_{mem} \ or \ \mu_{obs} \ \ [10^6 \ km^2]$', fontsize=18);

plt.tight_layout()

fig.savefig(data_path+'SIA/figures/Resampled_figures/Mem_obs_bar_chart_{}.png'.format(sigma_mu), dpi=400)
fig.savefig(data_path+'SIA/figures/Resampled_figures/Mem_obs_bar_chart_{}.pdf'.format(sigma_mu))
# -

# ## Fig 6 as submitted 

# +
dataset_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#9467bd']
sig_obs_array = sigma_obs_individual.to_array().sel(variable=['CDR','BT','NT','HadISST1'])
mu_obs_array  = mu_obs.to_array().sel(variable=['CDR','BT','NT','HadISST1'])

fig, axes = plt.subplots(2,2,figsize=[12,9])

#first plot of values of sigma_obs
for var_i in range(4):
    axes[0][0].plot(month_names_short, sig_obs_array.isel(variable=var_i), color=dataset_colors[var_i],
                    label=sig_obs_array['variable'].isel(variable=var_i).values, linewidth=2)

axes[0][0].grid()

axes[0][0].set_xlabel('Month', fontsize=18)
axes[0][0].set_ylabel(r'$\sigma_{obs} \ [10^6 \ km^2]$', fontsize=18)

axes[0][0].set_xlim(0,11)
axes[0][0].set_xticklabels(month_names_short, fontsize=14)
axes[0][0].set_yticklabels(['','0.02','0.03','0.04','0.05','0.06'], fontsize=14)

axes[0][0].legend(fontsize=13)

#second plot of observational uncertainty
percent_uncertainty = (sig_obs_array.max('variable').values - sig_obs_array.min('variable').values) / sig_obs_array.mean('variable').values

axes[0][1].plot(month_names_short, percent_uncertainty*100, color='r', linewidth=2.5)
axes[0][1].grid()
axes[0][1].set_xlim(0,11)
axes[0][1].set_xticklabels(month_names_short, fontsize=14)
axes[0][1].set_yticklabels(np.arange(0,110,20), fontsize=14)
axes[0][1].set_xlabel('Month', fontsize=18)
axes[0][1].set_ylabel('Observational Uncertainty [%]', fontsize=18);


#first plot of values of sigma_obs
for var_i in range(4):
    axes[1][0].plot(month_names_short, mu_obs_array.isel(variable=var_i), color=dataset_colors[var_i],
                    label=sig_obs_array['variable'].isel(variable=var_i).values, linewidth=2)

axes[1][0].grid()

axes[1][0].set_xlabel('Month', fontsize=18)
axes[1][0].set_ylabel(r'$\mu_{obs} \ [10^6 \ km^2]$', fontsize=18)

axes[1][0].set_xlim(0,11)
axes[1][0].set_xticklabels(month_names_short, fontsize=14)
axes[1][0].set_yticklabels(['','0.20','0.25','0.30','0.35','0.40','0.45','0.50'], fontsize=14)

axes[1][0].legend(fontsize=13)

#second plot of observational uncertainty
percent_uncertainty = (mu_obs_array.max('variable').values - mu_obs_array.min('variable').values) / mu_obs_array.mean('variable').values

axes[1][1].plot(month_names_short, percent_uncertainty*100, color='r', linewidth=2.5)
axes[1][1].grid()
axes[1][1].set_xlim(0,11)
axes[1][1].set_ylim(10,34.9)
axes[1][1].set_xticklabels(month_names_short, fontsize=14)
axes[1][1].set_yticklabels(np.arange(10,41,5), fontsize=14)
axes[1][1].set_xlabel('Month', fontsize=18)
axes[1][1].set_ylabel('Observational Uncertainty [%]', fontsize=18);


axes[0][0].text(-0.03, 1.02, 'a', transform=axes[0][0].transAxes, fontsize=18, fontweight='bold', va='top', ha='right');
axes[0][1].text(-0.03, 1.02, 'b', transform=axes[0][1].transAxes, fontsize=18, fontweight='bold', va='top', ha='right');
axes[1][0].text(-0.03, 1.02, 'c', transform=axes[1][0].transAxes, fontsize=18, fontweight='bold', va='top', ha='right');
axes[1][1].text(-0.03, 1.02, 'd', transform=axes[1][1].transAxes, fontsize=18, fontweight='bold', va='top', ha='right');


plt.tight_layout()

fig.savefig(data_path+'SIA/figures/Resampled_figures/Sigma_mu_obs_5_datasets_with_percent.png', bbox_inches='tight', dpi=400)
fig.savefig(data_path+'SIA/figures/Resampled_figures/Sigma_mu_obs_5_datasets_with_percent.pdf', bbox_inches='tight')

StarcoderdataJupyterScriptsDedupFiltered