ttclust.py 文件源码

def plot_dendro(linkage, logname, cutoff, color_list,clusters_list):
    """
    DESCRIPTION
    This function will create the dendrogram graph with the corresponding
    cluster color.
    Args:
        linkage (numpy array) : linkage matrix
        output (str) : output logfile name
        cutoff (float) : cutoff used for clustering
        color_list (list) : HEX code color for each cluster
        cluster_list (list) : list of all cluster (Cluster object)
    Returns:
        None
    """
    if mpl.__version__[0] == "2":
        STYLE = "classic"
        if STYLE in plt.style.available:
            plt.style.use(STYLE)
    fig = plt.figure()
    #Convert RGB color to HEX color
    color_hex = [mpl.colors.rgb2hex(x) for x in color_list]
    sch.set_link_color_palette(color_hex)
    #clusters_list
    color_member = {}
    for cl in clusters_list:
        for frm in cl.frames:
            color_member[frm] = mpl.colors.rgb2hex(color_list[cl.id-1])

    #Attribute the correct color for each branch.
    #adapte from Ulrich Stern code in StackOverflow http://stackoverflow.com/a/38208611
    link_cols = {}
    for i, i12 in enumerate(linkage[:,:2].astype(int)):
        c1, c2 = (link_cols[x] if x > len(linkage) else color_member[x] for x in i12)
        link_cols[i+1+len(linkage)] = c1 if c1 == c2 else "#808080"

    #Dendrogram creation
    # Override the default linewidth.
    den = sch.dendrogram(linkage, color_threshold=float(cutoff), above_threshold_color="#808080", link_color_func=lambda x: link_cols[x])

    #Graph parameters
    plt.title("Clustering Dendrogram")
    ax = plt.axes()
    ax.set_xticklabels([])
    plt.axhline(y=float(cutoff), color = "grey") # cutoff value vertical line
    ax.set_ylabel("Distance (AU)")
    ax.set_xlabel("Frames")

    plt.savefig("{0}/{0}-den.png".format(logname), format="png", dpi=DPI, transparent=True)
    plt.close()