aesop.py 文件源码

def plotESD(esd, filename=None, cmap='hot'):
    """Summary
    Function to display an electrostatic similarity heatmap from a previously
    run ElecSimilarity class.

    Parameters
    ----------
    esd : ndarray
        ESD matrix from ElecSimilarity class (ElecSimilarity.esd).
    filename : str, optional
        If the resulting plot should be written to disk, specify a filename.
        Otherwise, the image will only be saved.
    cmap : str, optional
        Colormap from matplotlib to use.

    Returns
    -------
    None
        Writes image to disk, if desired.
    """
    # plt.style.use('seaborn-talk')
    fig, ax = plt.subplots(sharey=True)
    heatmap = ax.pcolor(esd.esd, cmap=cmap, vmin=0, vmax=2)
    ax.set_xlim(0, esd.esd.shape[0])
    ax.set_ylim(0, esd.esd.shape[1])
    ax.set_xticks(np.arange(esd.esd.shape[0]) + 0.5, minor=False)
    ax.set_yticks(np.arange(esd.esd.shape[1]) + 0.5, minor=False)
    ax.set_xticklabels(esd.ids, rotation=90)
    ax.set_yticklabels(esd.ids)
    fig.colorbar(heatmap)
    plt.tight_layout()
    if filename is not None:
        fig.savefig(filename)


##########################################################################
# Function to plot ESD dendrogram
##########################################################################