test41

833be2ac-b34d-4412-a5ef-071fe3aa27d4.py

def cluster_order(d):
    pdist = hierarchy.distance.pdist(d.values)
    linkage = hierarchy.linkage(pdist, method="complete")
    idx = hierarchy.fcluster(linkage, 0.5 * pdist.max(), "distance")
    idx = np.argsort(idx)
    return d.iloc[idx, idx]


def draw(df, symmetric=True):

    if symmetric:
        df = df.corr()
    else:
        df = df.T.dot(df) / df.sum()

    df = cluster_order(df)
    size = df.shape[0]
    mask = np.triu(np.repeat(True, size))
    cmap = matplotlib.cm.get_cmap("vlag")
    if symmetric:
        max_limit = np.abs(df.values[~mask]).max()
        min_limit = -max_limit
    else:
        max_limit = df.values[~mask].max()
        min_limit = 0
        mid = len(cmap.colors) // 2
        cmap = ListedColormap(cmap.colors[mid:])

    fig, ax = plt.subplots(figsize=(5, 5))

    sns.heatmap(
        df,
        square=True,
        cmap=cmap,
        vmin=min_limit,
        vmax=max_limit,
        ax=ax,
        cbar_kws={"shrink": 0.8},
    )

    ax.set_title("Los Altos City Council - 2020")
    ax.set_xlabel("")
    ax.set_ylabel("")

    return df.index.values.astype(str)