cluster_map_order.py

import pandas as pd
import seaborn as sns


def order_by_correlation_cluster(data: pd.DataFrame, x: str, y: str, values: str) -> pd.DataFrame:
    """Takes a tidy dataframe, pivots it, and creates a clustermap based on 
    correlations between rows. Returns the original data with a new 'order' 
    column that contains the cardinal ordering of the y axis from the clustermap.
    Made with plotting in altair in mind.
    """
    
    ordered_corr = (
        sns.clustermap(
            data
            .pivot(index=y, columns=x, values=values)
            .corr()
        )
        .data2d
        .index
        .values
    )
    order_map = dict(zip(ordered_corr, range(ordered_corr.shape[0])))
    data["order"] = data[y].map(order_map)
    return data