RFR.py

from sklearn.ensemble import RandomForestRegressor
import networkx as nx
from networkx.drawing.nx_agraph import to_agraph
import pandas as pd
import matplotlib
from matplotlib import pyplot as plt
import shap
import numpy as np
import random
import copy
import itertools

def avgENSedges(ENS_dfs,members):
    IDS = []
    SOURCES = []
    TARGETS = []
    LAGS = []
    NON_ZEROS = []
    AVG_WEIGHT = []
    STD_WEIGHT = []

    for i in ENS_dfs[list(ENS_dfs.keys())[0]]['id']:
        s = None
        t = None
        l = None
        weights = []
        for member in members:            
            row = (ENS_dfs[member][ENS_dfs[member]['id'] == i]).iloc[0]
            if s == None:
                # s = row['source'].split(member)[0] + row['source'].split(member)[1]
                # t = row['Target'].split(member)[0] + row['Target'].split(member)[1]
                s = row['source'].replace(member,'ensAVG')
                t = row['Target'].replace(member,'ensAVG')
                l = row['lag']
            if row['Weight'] != 0:
                weights.append(row['Weight'])

        if len(weights) > 0:
            IDS.append(i)
            SOURCES.append(s)
            TARGETS.append(t)
            LAGS.append(l)
            NON_ZEROS.append(len(weights))
            AVG_WEIGHT.append(np.mean(weights))
            STD_WEIGHT.append(np.std(weights))


    df = pd.DataFrame({
                'id': IDS,
                'source': SOURCES,
                'lag': LAGS,
                'Target': TARGETS,
                'Weight': AVG_WEIGHT,
                'Std_Weight': STD_WEIGHT,
                'Non_zeros': NON_ZEROS

            })
    return df

def readAllFiles(FileZoneTypes,FileFeatTypes):
    r_data = {}
    for Z in FileZoneTypes:
        for F in FileFeatTypes:
            r_data[F+'_'+Z] = readRFRfeatures("Data_E3SM_limvar/"+Z+"_limvar_"+F+"_dailymean.csv")
    return r_data

def readHSWfile(FileZoneTypes,FileFeatTypes):
    data = readRFRfeatures('Data_HSW/ZonalHSW_LV_AOD_T050_T1000.csv')
    
    r_data = {}
    for Z in FileZoneTypes:
        for F in FileFeatTypes:
            keep_keys = []
            for k in data.keys():
                Zone,Feat = k.split('_')[0],k.split('_')[-1]
                if Zone == Z and Feat == F:
                    keep_keys.append(k)
                
            r_data[F+'_'+Z] = {key: data[key] for key in keep_keys}
    
    return r_data
    
def readRFRfeatures(filename):
    RAW_DATA = {}

    f = open(filename,'r')
    first = True

    for line in f.readlines():
        l = line.strip().split(',')
        if first:
            if l[-1][0] == '"':
                labels = [x[1:-1] for x in l]
            else:
                labels = [x for x in l]

            for k in labels:
                if k != 't':
                    RAW_DATA[k] = []
            first = False
        else:
            for i,v in enumerate(l):
                if i == 0:
                    continue
                RAW_DATA[labels[i]].append(float(v))

    for key in RAW_DATA.keys():
        RAW_DATA[key] = list( ([RAW_DATA[key]])[0] )

    return RAW_DATA

def normVal(val,range_v,norm_range=(0,1)):
    if norm_range == (0,1):
        return (val - range_v[0]) / (range_v[1] - range_v[0])
    elif norm_range == (-1,1):
        return val  / max(abs(range_v[1]),abs(range_v[0]))
    else:
        return None

def normByFeat(DATA,feat_types=['_AOD','_T1000','_T050'],norm_range=(0,1)):
    normed_data = {}
    feat_scalars = {}

    for ft in feat_types:
        feat_scalars[ft] = [None,None]
        temp_d = list(itertools.chain.from_iterable( [DATA[x] for x in DATA if ft in x] ))
        if temp_d == []:
            continue

        feat_scalars[ft][0] = min(temp_d)
        feat_scalars[ft][1] = max(temp_d)

        for key in DATA:
            if ft in key:
                normed_data[key] = [ normVal(val,feat_scalars[ft],norm_range) for val in DATA[key]]

    return normed_data,feat_scalars

def sampleIndices(indices,perc=0.75,replace=False):
    last = int(len(indices)*perc)
    if replace:
        ind_set = []
        while len(ind_set) < last:
            rand_ind = int(random.random()*len(indices))
            ind_set.append(indices[rand_ind])
    else:
        ind_set = list(copy.deepcopy(indices))
        random.shuffle(ind_set)
        ind_set = ind_set[:last]

    return ind_set

def pruneTopEdges(df,n=20,top_perc=None,minVal=0.00000001):
    #columns: source,Target,Weight
    df_new = df.loc[(df['Weight'] > minVal)]
    if top_perc == None and n == None:
        return df_new.sort_values(by = ['Weight'], ascending= False)
    # print(int(len(df_new)))
    if n == None:
        top_rows = int(len(df_new) * top_perc)
    else:
        top_rows = n
    return df_new.sort_values(by = ['Weight'], ascending= False).head(top_rows)

def bestLag(df,minVal=0.0):
    #columns: source,lag,Target,Weight
    df_new = df.loc[(df['Weight'] > minVal)]
    dfs = []
    for T in set(df['Target']):
        for S in set(df['source']):
            df_temp = df_new.loc[(df_new['source'] == S) & (df_new['Target'] == T)]
            df_sorted = df_temp.sort_values(by=['Weight','lag'],ascending=[False,True])
            dfs.append(df_sorted.head(1))

    return (pd.concat(dfs)).sort_values(by = ['Weight'], ascending= False,inplace=False)

def pruneEdges(df,percentile=0.90,skip_self_loops=True):
    #columns: source,Target,lag,Weight
    if skip_self_loops == True:
        TOTAL_W = (df.loc[df['Target'] != df['source']])['Weight'].sum()
    else:
        TOTAL_W = df['Weight'].sum()
    curr_w = 0.0
    count = 0

    for index, row in df.iterrows():
        count += 1
        if row['Target'] == row['source']:
            continue
        curr_w += row['Weight']
        print(count,curr_w / TOTAL_W)
        if curr_w / TOTAL_W >= percentile:
            break

    return df[:count]

    # TOP N edges per node
def pruneIncomingEdges(df,top=3,minVal=0.0):
    num_ens = max(df['Non_zeros'])
    df_new = df.loc[(df['Weight'] > minVal) & (df['Weight'] > df['Std_Weight']) & (df['Non_zeros'] == num_ens)]
#     df_new = df.loc[(df['Weight'] > minVal) & (df['Weight'] > df['Std_Weight'])]
    dfs = []
    for T in set(df['Target']):
        df_temp = df_new.loc[(df_new['Target'] == T)]
        df_sorted = df_temp.sort_values(by='Weight',ascending=False)
        dfs.append(df_sorted.head(top))
    df_return = pd.concat(dfs)
    df_return.sort_values(by = ['Weight'], ascending= False,inplace=True)
    return df_return

def genCompleteMap(FEATS,LAGS):
    map = {}
    for feat in FEATS:
        map[feat] = []
        for f in FEATS:
            for l in LAGS:
                map[feat].append( (f,l) )
    return map

def dfToMap(df):
    MAP = {}
    for Tar in set(df['Target']):
        MAP[Tar] = []

    for index, row in df.iterrows():
        MAP[ row['Target'] ].append( (row['source'],int(row['lag'])) )

    return MAP

def run(DATA,MAP,randVar=False,models=False,tree_depth=4,random_seed=None,show_plot=True):
    figure, axis = plt.subplots(len(MAP.keys()), 1,figsize=(15, 5),sharex=True) #,sharey=True

    id_val = 0
    results = {}
    r2_results = {}
    count = 0

    for TAR_KEY in sorted(MAP.keys()):
        VALUES = MAP[TAR_KEY]
        ID = []
        SOURCE = []
        LAG = []
        TARGET = []
        WEIGHT = []
        LABELS = []

        c_data = []
        max_l = max([x[1] for x in VALUES])

        for v in VALUES:
            LABELS.append(str(v[0])+'_-'+str(v[1]))
            diff_l = max_l-int(v[1])
            c_data.append( [x for x in DATA[v[0]][diff_l:-1*int(v[1])]]  )

        ###ADD RV####
        if randVar:
            temp_data = []
            for key in DATA:
                size = len(DATA[key])
                temp_data += DATA[key]

            mean = np.mean(temp_data)
            stddev = np.std(temp_data)

            rv_data = np.random.normal(loc=mean, scale=stddev,size=len(data))

            c_data.append(rv_data[:-1*max_l])

            LABELS.append('RV')
        ######

        data_INPUT = np.column_stack(c_data)
        data_TARGET = [ x for x in DATA[TAR_KEY][max_l:] ]

        weights = {}

        importances,rfr_model = getFeatureImportances(data_INPUT,
                                            data_TARGET,
                                            multiout=False,
                                            model=True,
                                            tree_depth=tree_depth,
                                            random_seed=random_seed)

        for i,d in enumerate(LABELS):
            if d == 'RV':
                continue
            id_val += 1
            ID.append(id_val)
            SOURCE.append(d.split('_-')[0])
            LAG.append( int(d.split('_-')[1]) )
            TARGET.append(TAR_KEY)
            WEIGHT.append(round(importances[i],5))


        df = pd.DataFrame({
            'id': ID,
            'source': SOURCE,
            'lag': LAG,
            'Target': TARGET,
            'Weight': WEIGHT
        })
        df.sort_values(by = ['Weight'], ascending= False,inplace=True)

        if models:
            results[TAR_KEY] = (df,rfr_model,(data_INPUT,data_TARGET))
        else:
            results[TAR_KEY] = df

        r2_results.update(plotPredictions(LABELS,rfr_model,data_INPUT,data_TARGET,ax=axis[count],name=TAR_KEY,padding=max_l,show_plot=show_plot))
        axis[count].axhline(y=0,color='k',linestyle='--')
        count += 1
    if show_plot:
        plt.show()

    df_merged = (pd.concat([res[0] for tar,res in results.items()])).sort_values(by = ['Weight','lag'], ascending= [False,True],inplace=False)
    return df_merged,results,r2_results

def getFeatureImportances(data_X,data_y,multiout=False,model=False,tree_depth=None,random_seed=0):
    forest = RandomForestRegressor(max_depth=tree_depth, random_state=random_seed)
    forest.fit(data_X, data_y)


    input_data = pd.DataFrame(data_X) # columns=feature_names)

    explainer = shap.TreeExplainer(forest)
    shap_values = explainer.shap_values(data_X) #,check_additivity=False)

#     print(len(shap_values),len(shap_values[0]),len(shap_values[0][0]))

#     shap.summary_plot(shap_values, input_data, plot_type="violin")
#     shap.summary_plot(shap_values, input_data, feature_names=["W","X","Y","Z","RV"], plot_type="bar")

#     shap.summary_plot(shap_values, input_data)

    if multiout:
        score = [abs(sv).mean(axis=0) for sv in shap_values]
    else:
        score = abs(shap_values).mean(axis=0)

    if model:
        return score,forest
    else:
        return score


def generatePredictions(INPUT,TARGET,ax,name,padding=0,show_plot=True):

    if show_plot:
        ax.plot(([None]*padding)+TARGET,label=name)
        ax.legend(loc=1)
        ax.plot(([None]*padding)+INPUT,label=str(name)+"*",color='orange')
        ax.legend(loc=1)

    from sklearn.metrics import r2_score
    from sklearn.metrics import explained_variance_score
    from sklearn.metrics import mean_absolute_error

    SCORE_r2 = r2_score(TARGET, INPUT)
    # SCORE_expVar = explained_variance_score(TARGET, INPUT)
    # SCORE_meanErr = mean_absolute_error(TARGET, INPUT)
    # print("EV",SCORE_expVar)
    return round(SCORE_r2,3) #, SCORE_expVar, SCORE_meanErr

def plotPredictions(LABELS,model,data_INPUT,data_TARGET,ax=None,name=None,padding=0,show_plot=True):

    predictions = model.predict(data_INPUT)

    predict_size = len(predictions.shape)
    if predict_size > 1:
        figure, axis = plt.subplots(len(LABELS), 1,figsize=(15, 5),sharex=True) #,sharey=True

        for i,k in enumerate(sorted(LABELS)):
                print(k,generatePredictions(predictions[:,i],data_TARGET[:,i],axis[i],k),show_plot)
        if show_plot:
            plt.show()
    else:
        # print(name,generatePredictions(list(predictions),list(data_TARGET),ax,name,padding=padding,show_plot=True))
        return {name:generatePredictions(list(predictions),list(data_TARGET),ax,name,padding=padding,show_plot=True)}


def plotGraphNetwork(data_frame,graphname="graph.png",inch_size=12,penwidth=None,
                    edge_cmap=plt.cm.Blues,
                    node_color={'AEROD_v':'blue','AOD':'blue',
                                'T050':'red','T1000':'orange','TREFHT':'orange',
                                'FLNT':'darkgreen','FSDSC':'darkgreen','FLA':'darkgreen'}):

    feat_names = set(data_frame['Target'])

    edges = {}
    for index, row in data_frame.iterrows():
        #source, Target, Weight
        src = row['source']
        lag = row['lag']
        key = (src,row['Target'])
        if key not in edges:
            edges[key] = {'lags':str(lag),'weight':float(row['Weight'])}
        else:
            edges[key]['lags'] += ',' + str(lag)
            edges[key]['weight'] += float(row['Weight'])

    edge_values = [value['weight'] for key,value in edges.items()]
    color_norm = matplotlib.colors.Normalize(vmin=0.0,
                                             vmax=max(edge_values)*1.02
                                            )

    G = nx.DiGraph()

    for key, value in edges.items():
        G.add_edge(key[0],key[1],
                    weight=value['weight'],
                    lags=value['lags'])

    edge_labels = dict([((n1, n2), d['lags'])
                        for n1, n2, d in G.edges(data=True)])

    nx.set_edge_attributes(G,
                            values=edge_labels,
                            name= "label")
    node_colormap = {}
    for node in G:
        c = 'black'
        for key,value in node_color.items():
            if key in node:
                c = value
        node_colormap[node] = c
    nx.set_node_attributes(G, node_colormap, name="color")

    if penwidth != None:
        edge_weights = dict([((n1, n2), d['weight'])
                            for n1, n2, d in G.edges(data=True)])

        nx.set_edge_attributes(G,
                                values=edge_weights,
                                name= "weight")

    edge_colors = dict([((n1, n2),matplotlib.colors.rgb2hex(edge_cmap(color_norm(d['weight']))))
                        for n1, n2, d in G.edges(data=True)])

    nx.set_edge_attributes(G,
                            values=edge_colors,
                            name= "color")

    A = to_agraph(G)
    if penwidth != None:
        for key,value in edge_weights.items():
            edge = A.get_edge(key[0],key[1])
            edge.attr['penwidth'] = value * penwidth

    for node_name in G:
        r = "same"
        if 'AEROD_v' in node_name or 'AOD' in node_name:
            r = "min"
        elif 'TREFHT' in node_name or 'T050' in node_name or 'T1000' in node_name:
            r = "max"
        B = A.add_subgraph([node_name],rank=r)
        B.graph_attr["rank"] = r
    A.layout('dot')


    # A.graph_attr.update(dpi=1000,edge_cmap=edge_cmap,fontsize=12)#,size=(inch_size,inch_size))
    A.graph_attr['dpi'] = 500
    A.graph_attr['size'] = "20,20!"
    A.graph_attr['fontsize'] = "8pt"


    A.draw(graphname)