visualizers.py 文件源码

def heatmap(data, half=True, scale=1, vmin=-0.8, vmax=0.8, cmap='RdBu_r', **kwargs):
        """

        :param dataframe:
        :param half:
        :param scale:
        :param vmin:
        :param vmax:
        :param cmap:
        :param kwargs:
        :return:
        """
        figsize = (6 * scale, 4 * scale)
        for arg in kwargs.keys():
            if arg is 'figsize':
                figsize = kwargs[arg]
        if half:
            mask = np.zeros_like(data)
            mask[np.triu_indices_from(mask)] = True
        else:
            mask = None
        fig = plt.figure(figsize=figsize, dpi=300)
        fig.set_facecolor('white')
        axes = fig.add_subplot(111)

        with sns.plotting_context("notebook", font_scale=1):
            ax = sns.heatmap(data, mask=mask, vmin=vmin, vmax=vmax,
                             cmap=cmap, square=True, ax=axes)
            # ax.set_xticklabels(ax.xaxis.get_majorticklabels(), rotation=45)
            ax.tick_params(labelsize=3.5, length=0)
            # ax.set_yticklabels(ax.yaxis.get_majorticklabels(), rotation=45)
            cbar = ax.collections[0].colorbar
            cbar.set_ticks([vmin, 0, vmax])
            # cbar.set_ticklabels(['low', '20%', '75%', '100%'])