alternative.py

import os
os.chdir('C:/Users/morenodu/OneDrive - Stichting Deltares/Documents/PhD/Paper_drought/data')
import xarray as xr 
import matplotlib.pyplot as plt
import pandas as pd
from bias_correction_masked import *
from return_period_storyline import return_period_storyline
from extrapolation_test import *
import pickle

#%% 1st - load data generated by bias correction

DS_ec_earth_PD_us = xr.open_dataset("ds_ec_earth_PD_us_lr.nc",decode_times=True, use_cftime=True)
DS_ec_earth_2C_us = xr.open_dataset("ds_ec_earth_2C_us_lr.nc",decode_times=True, use_cftime=True)
DS_ec_earth_3C_us = xr.open_dataset("ds_ec_earth_3C_us_lr.nc",decode_times=True, use_cftime=True)

# load and treat bias-adjusted ec earth projections
DS_ec_earth_PD_brs = xr.open_dataset("ds_ec_earth_PD_brs.nc",decode_times=True, use_cftime=True)
DS_ec_earth_2C_brs = xr.open_dataset("ds_ec_earth_2C_brs.nc",decode_times=True, use_cftime=True)
DS_ec_earth_3C_brs = xr.open_dataset("ds_ec_earth_3C_brs.nc",decode_times=True, use_cftime=True)

# load and treat bias-adjusted ec earth projections
DS_ec_earth_PD_ar = xr.open_dataset("ds_ec_earth_PD_ar.nc",decode_times=True, use_cftime=True)
DS_ec_earth_2C_ar = xr.open_dataset("ds_ec_earth_2C_ar.nc",decode_times=True, use_cftime=True)
DS_ec_earth_3C_ar = xr.open_dataset("ds_ec_earth_3C_ar.nc",decode_times=True, use_cftime=True)

# load and treat bias-adjusted ec earth projections
DS_ec_earth_PD_brc = xr.open_dataset("ds_ec_earth_PD_brc.nc",decode_times=True, use_cftime=True)
DS_ec_earth_2C_brc = xr.open_dataset("ds_ec_earth_2C_brc.nc",decode_times=True, use_cftime=True)


#%% 2nd - Convert datasets to dataframe adjusted for ML
df_features_ec_season_us, df_features_ec_season_2C_us, df_features_ec_season_3C_us = function_conversion(DS_ec_earth_PD_us, DS_ec_earth_2C_us,DS_cli_ec_3C=DS_ec_earth_3C_us, months_to_be_used=[7,8])

df_features_ec_season_brs, df_features_ec_season_2C_brs, df_features_ec_season_3C_brs = function_conversion(DS_ec_earth_PD_brs, DS_ec_earth_2C_brs,DS_cli_ec_3C=DS_ec_earth_3C_brs, months_to_be_used=[1,2], water_year= True)

df_features_ec_season_ar, df_features_ec_season_2C_ar, df_features_ec_season_3C_ar = function_conversion(DS_ec_earth_PD_ar, DS_ec_earth_2C_ar,DS_cli_ec_3C=DS_ec_earth_3C_ar, months_to_be_used=[1,2,3], water_year= True)

# Careful with this dataset as it needs better handling of water year and low performance - TRY AGAIN
df_features_ec_season_brc, df_features_ec_season_2C_brc = function_conversion(DS_ec_earth_PD_brc, DS_ec_earth_2C_brc, months_to_be_used=[1,2], water_year= True)

#%% PDFs for all regions
list_feautures_names = ['temperature Celisus Degrees','Diurnal Temperature Range (C)', 'Precipitation mm/month']
for scenario in ['present', '2C']:
    if scenario == 'present':
        df_feature_scen_us, df_feature_scen_brs, df_feature_scen_ar, df_feature_scen_brc = df_features_ec_season_us, df_features_ec_season_brs, df_features_ec_season_ar, df_features_ec_season_brc
    
    elif scenario == '2C':
        df_feature_scen_us, df_feature_scen_brs, df_feature_scen_ar, df_feature_scen_brc = df_features_ec_season_2C_us, df_features_ec_season_2C_brs, df_features_ec_season_2C_ar, df_features_ec_season_2C_brc

    for feature in range(len(df_features_ec_season_us.columns)):
        df_con_hist = pd.concat( [df_feature_scen_us.iloc[:,feature],df_feature_scen_brs.iloc[:,feature],df_feature_scen_ar.iloc[:,feature], df_feature_scen_brc.iloc[:,feature] ],axis=1)
        df_con_hist.columns = ['us','south brazil','argentina','central brazil']
        
        plt.figure(figsize = (6,6), dpi=144)
        fig = sns.displot(df_con_hist,kind="kde", aspect=1, linewidth=3,fill=True,alpha=.2)
        fig.set(xlabel=df_feature_scen_us.columns[feature])
    
        fig._legend.set_bbox_to_anchor((.6, 0.9))
        # plt.tight_layout()
        plt.show()
        
        
#%% Comparison between EC-Earth bias adjusted and CRU for a single region
# PDFs
for feature in range(len(df_clim_agg_chosen.columns)):
    df_cru_hist = pd.DataFrame( df_clim_agg_chosen.iloc[:,feature])
    df_cru_hist['Scenario'] = 'CRU'
    df_ec_hist = pd.DataFrame( df_features_ec_season_us.iloc[:,feature])
    df_ec_hist['Scenario'] = 'EC-earth'
    df_hist_us = pd.concat( [df_cru_hist,df_ec_hist],axis=0)
    df_hist_us.index = range(len(df_hist_us))
    
    plt.figure(figsize = (6,6), dpi=144)
    fig = sns.kdeplot( data = df_hist_us, x=df_clim_agg_chosen.columns[feature], hue="Scenario", fill=True, alpha=.2, common_norm = False)
    fig.set(xlabel=df_clim_agg_chosen.columns[feature])
    plt.show()
    
    df_runs = pd.DataFrame(np.repeat(range(0, 400), 5), index = df_ec_hist.index, columns = ['run'])  
    df_feat = pd.concat( [df_ec_hist.iloc[:,0], df_runs], axis = 1)
    
    df_feat2 = pd.DataFrame(df_ec_hist.iloc[:,0])
    df_feat2.index = np.tile([2011,2012,2013,2014,2015], 400) 
    df_feat2.sort_index(inplace=True)
    
    df_ensemble = pd.DataFrame( np.reshape(df_feat.iloc[:,0].values, (5,400)))
    df_ensemble.index.name = 'index' 
    
    # sns.lineplot(x=df_feat2.index, y=df_feat2.columns[0], data=df_feat2, ci=95, estimator='mean')
    
    plt.figure(figsize = (6,6), dpi=144)
    fig = sns.lineplot( data = df_cru_hist[42:], x=df_cru_hist[42:].index, y=df_clim_agg_chosen.columns[feature])
    fig = sns.lineplot( data = df_feat2, x=df_feat2.index, y=df_feat2.columns[0], legend= False)
    fig.set(xlabel=df_clim_agg_chosen.columns[feature])
    plt.show()
    
# Quantil quantile mapping
import scipy.stats as stats
for i in [0,1,2]:   
    stats.probplot(df_clim_agg_chosen.iloc[62:,i], dist=stats.beta, sparams=(2,3),plot=plt,fit=False)
    stats.probplot(df_features_ec_season_us.iloc[:,i], dist=stats.beta, sparams=(2,3),plot=plt,fit=False)
    plt.show()
       
#%% ####3RD part - Random Forest predictions 
with open('rf_model_us.pickle', 'rb') as f:
    brf_model_us = pickle.load(f)
table_scores_us, table_events_prob2012_us = predictions_permutation(brf_model_us, df_clim_agg_chosen, df_features_ec_season_us, df_features_ec_season_2C_us, df_features_ec_season_3C_us, df_clim_2012_us  )

table_scores_brs, table_events_prob2012_brs = predictions_permutation(brf_model_brs, input_features_brs, df_features_ec_season_brs, df_features_ec_season_2C_brs, df_features_ec_season_3C_brs, df_clim_2012 = df_clim_2012_brs )

table_scores_ar, table_events_prob2012_ar = predictions_permutation(brf_model_ar, input_features_ar,df_features_ec_season_ar, df_features_ec_season_2C_ar, df_features_ec_season_3C_ar,  df_clim_2012 = df_clim_2012_ar  )

#really weird results - poor ML performance
table_scores_brc, table_events_prob2012_brc = predictions_permutation(brf_model_brc, df_features_ec_season_brc, df_features_ec_season_2C_brc )

#%% ############### 4TH PART - Compound analysis

#US
df_joint_or_rf_us, table_JO_prob2012_us = compound_exploration(brf_model_us, df_features_ec_season_us, df_features_ec_season_2C_us, df_features_ec_season_3C_us, df_clim_2012_us )

#BRS
df_joint_or_rf_brs, table_JO_prob2012_brs = compound_exploration(brf_model_brs, df_features_ec_season_brs, df_features_ec_season_2C_brs, df_features_ec_season_3C_brs, df_clim_2012 = df_clim_2012_brs )

#AR
df_joint_or_rf_ar, table_JO_prob2012_ar = compound_exploration(brf_model_ar, df_features_ec_season_ar, df_features_ec_season_2C_ar, df_features_ec_season_3C_ar, df_clim_2012 = df_clim_2012_ar )

#BRC CAUTION - results unstable
df_joint_or_rf_brc, table_JO_prob2012_brc = compound_exploration(brf_model_brc, df_features_ec_season_brc, df_features_ec_season_2C_brc )

#%%% Return period figure with storyline
# US
mean_conditions_similar_2012_2C_us, mean_conditions_similar_2012_3C_us = return_period_storyline(
    df_features_ec_season_us, df_features_ec_season_2C_us, df_clim_agg_chosen,
    table_JO_prob2012_us, table_events_prob2012_us, brf_model_us,
    df_clim_2012_us, df_joint_or_rf_us, proof_total_us, df_features_ec_season_3C_us)

# BRS
mean_conditions_similar_2012_2C_brs, mean_conditions_similar_2012_3C_brs = return_period_storyline(
    df_features_ec_season_brs, df_features_ec_season_2C_brs, input_features_brs,
    table_JO_prob2012_brs, table_events_prob2012_brs, brf_model_brs,
    df_clim_2012_brs, df_joint_or_rf_brs, compoundness_obs_brs, df_features_ec_season_3C_brs)

# ARG
mean_conditions_similar_2012_2C_ar, mean_conditions_similar_2012_3C_ar = return_period_storyline(
    df_features_ec_season_ar, df_features_ec_season_2C_ar, input_features_ar,
    table_JO_prob2012_ar, table_events_prob2012_ar, brf_model_ar,
    df_clim_2012_ar, df_joint_or_rf_ar, compoundness_obs_ar, df_features_ec_season_3C_ar)


#%%

max_PD = df_features_ec_season_us.max(axis=0) -df_clim_agg_chosen.max(axis=0)
min_PD = df_features_ec_season_us.min(axis=0) - df_clim_agg_chosen.min(axis=0)

max_2C = df_features_ec_season_2C_us.max(axis=0) -df_clim_agg_chosen.max(axis=0)
min_2C = df_features_ec_season_2C_us.min(axis=0) - df_clim_agg_chosen.min(axis=0)

max_3C = df_features_ec_season_3C_us.max(axis=0) - df_clim_agg_chosen.max(axis=0)
min_3C = df_features_ec_season_3C_us.min(axis=0) - df_clim_agg_chosen.min(axis=0)

max_min_def = pd.DataFrame([max_2C, max_3C, min_2C, min_3C], index = ['Max 2C','Max 3C', 'Min 2C', 'Min 3C'])


# EXTRAPOLATION TEST
df_features_ec_season_us_new, df_features_ec_season_2C_us_new, df_features_ec_season_3C_us_new = extrapolation_test(df_clim_agg_chosen, df_features_ec_season_us, df_features_ec_season_2C_us, df_features_ec_season_3C_us,
                       brf_model_us, df_clim_2012_us, table_JO_prob2012_us, table_events_prob2012_us,
                       df_joint_or_rf_us, proof_total_us)




def out_in_range(out_data, in_data, scenario):
        print("MAX",scenario, out_data[out_data > in_data.max(axis=0)].count())
        print("MIN",scenario, out_data[out_data < in_data.min(axis=0)].count())
        
        max_count = out_data[out_data > in_data.max(axis=0)].count()
        min_count = out_data[out_data < in_data.min(axis=0)].count()
        max_min = pd.concat([max_count, min_count], axis = 1)    
        
        for feature in in_data:
            out_column = out_data.loc[:,feature]
            in_column = in_data.loc[:,feature]
            
            out_column[out_column > in_column.max()] = in_column.max()
            out_column[out_column < in_column.min()] = in_column.min()
        
        print("MAX CORRECTED",scenario, out_data[out_data > in_data.max(axis=0)].count())
        print("MIN CORRECTED",scenario, out_data[out_data < in_data.min(axis=0)].count())
        
        return max_min
        
    
max_min_PD = out_in_range(df_features_ec_season_us, df_clim_agg_chosen, scenario = 'PD')
max_min_2C = out_in_range(df_features_ec_season_2C_us, df_clim_agg_chosen, scenario = '2C')
max_min_3C = out_in_range(df_features_ec_season_3C_us, df_clim_agg_chosen, scenario = '3C')

table_exceedance = pd.DataFrame( pd.concat([max_min_PD, max_min_2C, max_min_3C], axis = 1) )
table_exceedance.columns = ['PD Above', 'PD Below', '2C Above','2C Below', '3C Above','3C Below']



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

df_features_ec_season_2C_us.mean(axis=0) - df_clim_agg_chosen.max(axis=0)

dev_2012 = (df_clim_2012_us - df_clim_agg_chosen.mean(axis=0)) / df_clim_agg_chosen.std(axis=0, ddof=0)

from sklearn import preprocessing
scaler = preprocessing.StandardScaler().fit(df_clim_agg_chosen)
ar_scaled = scaler.transform(df_clim_agg_chosen)
ar_2C_scaled = scaler.transform(df_features_ec_season_us)
ar_3C_scaled = scaler.transform(df_features_ec_season_3C_us)

df_scaled = pd.DataFrame(ar_scaled, index=df_clim_agg_chosen.index, columns = df_clim_agg_chosen.columns)
df_2C_scaled = pd.DataFrame(ar_2C_scaled, index=df_features_ec_season_us.index, columns = df_clim_agg_chosen.columns)
df_3C_scaled = pd.DataFrame(ar_3C_scaled, index=df_features_ec_season_us.index, columns = df_clim_agg_chosen.columns)

plt.figure(figsize=(9,5), dpi=500)
fig = sns.boxplot(x='variable', y='value', data=pd.melt(df_scaled))
fig = sns.boxplot(x='variable', y='value', data=pd.melt(df_3C_scaled))


fig = sns.scatterplot(x = df_scaled.columns, y = df_scaled.loc[2012,:], data =df_scaled.loc[2012,:], 
                color = 'red', alpha = 1, s = 80, zorder=100, label = '2012 season')
plt.xlabel('')
plt.legend(bbox_to_anchor=(0.95, 0.5), loc=2, borderaxespad=0.)
fig.set_ylabel('STD units')
plt.tight_layout()