interaction-heatmaps.py

def interactions_heatmap(shap_interaction_values, n, column_names, plot_heatmap=True):
    """
    Sums over first tensor of shap interaction values dimension (axis =0), returns a matrix dim nfeatures x nfeatures of interaction values
        
        Parameters
        ----------
        shap_interaction_values (np.array): tensor of interaction values produced by shap estimator.shap_interaction_values()
        n (int) - n top features to be kept
        column_names (List[str], np.array[]) - names of colums for feature names
        plot_heatmap (bool) - default True should the image be plotted
    
    """
    
    tmp = np.abs(shap_interaction_values).sum(0) #sum over samples - feature x feature matrix
    for i in range(tmp.shape[0]):
        tmp[i,i] = 0
    inds = np.argsort(-tmp.sum(0))[:n] #order by importance
    tmp2 = tmp[inds,:][:,inds]
    
    if plot_heatmap:
        pl.figure(figsize=(12,12))
        ax = pl.imshow(tmp2, cmap = colors.red_transparent_blue)
        pl.yticks(range(tmp2.shape[0]), column_names[inds], rotation=50.4, horizontalalignment="right")
        pl.xticks(range(tmp2.shape[0]), column_names[inds], rotation=50.4, horizontalalignment="left")
        pl.gca().xaxis.tick_top()
        pl.colorbar(ax)
        pl.show()
        
    return pd.DataFrame(tmp2, columns = column_names[inds], index = column_names[inds])