python matplotlib_routines.py

Posted 2021-05-09

def corr_plot(df, ylabel2=None, title=None):
    if isinstance(ylabel2, pd.Series):
        ylabel2 = ylabel2.values
    corrmat = df.corr().round(2)
    cols = df.columns.values
    ncols = len(cols)
    
    plt.figure(figsize = (len(cols)/2 , len(cols)/2))
    if title is not None:
        plt.title(title)
        
    mask = np.zeros_like(corrmat)
    mask[np.triu_indices_from(mask,k=1)] = True
    ax = sns.heatmap(corrmat, mask=mask, cmap='coolwarm', linewidths=0.5, annot=True, cbar=False, vmin=0, vmax=1)
    
    if ylabel2 is not None:
        assert len(ylabel2) == len(cols)
        for kth in range(ncols):
            ax.text(kth + 1.1, ncols - 0.6 - kth, ylabel2[kth], fontsize=11) 
            
# trick from https://stackoverflow.com/questions/4325733/save-a-subplot-in-matplotlib
def full_extent(self, ax, pad=0.2):
  """Get the full extent of an axes """
  # For text objects, we need to draw the figure first, otherwise the extents
  # are undefined.
  ax.figure.canvas.draw()
  items = ax.get_xticklabels() + ax.get_yticklabels()
  items += ax.texts
  items += [ax, ax.title, ax.xaxis.label, ax.yaxis.label]
  bbox = Bbox.union([item.get_window_extent() for item in items])
  return bbox.expanded(1.0 + pad, 1.0 + pad)