<reponame>tjiagoM/spatio-temporal-brain
import argparse

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.cluster.hierarchy import fcluster

from utils_datasets import STRUCT_COLUMNS


# python diffpool_plot_assignment.py --sweep_name 100_n_diffpool
# python diffpool_plot_assignment.py --sweep_name 100_n_e_diffpool

def plot_and_save_interp(arr, name, sweep_name):
    s_df = pd.DataFrame(arr, index=STRUCT_COLUMNS, columns=STRUCT_COLUMNS)

    # First create a dummy clustermap to know how the dendrogram is created and find the right mask next
    g_obj = sns.clustermap(s_df, yticklabels=1, xticklabels=1, dendrogram_ratio=(0.1, 0.2),
                           cbar_pos=(0, 0.85, 0.02, 0.15), cmap="viridis")
    mask_array = np.full(arr.shape, False)
    mask_array[np.tril_indices(mask_array.shape[0])] = True

    mask_after = mask_array[np.argsort(g_obj.dendrogram_row.reordered_ind), :]
    mask_after = mask_after[:, np.argsort(g_obj.dendrogram_col.reordered_ind)]

    g_obj = sns.clustermap(s_df, yticklabels=1, xticklabels=1, dendrogram_ratio=(0.1, 0.2),
                           cbar_pos = (0, 0.85, 0.02,0.15), cmap = "viridis", mask = mask_after,
                           linewidths = 0.5, linecolor = (0.7, 0.7, 0.7, 0.2))

    g_obj.ax_heatmap.set_xticklabels(g_obj.ax_heatmap.get_xmajorticklabels(), fontsize=7)
    g_obj.ax_heatmap.set_yticklabels(g_obj.ax_heatmap.get_ymajorticklabels(), fontsize=7)

    if sweep_name == '100_n_e_diffpool':
        granularities = [4, 8, 12]
    elif sweep_name == '100_n_diffpool':
        granularities = [4]

    for granularity_id in granularities:
        ord_ind = fcluster(g_obj.dendrogram_col.linkage, granularity_id, criterion='maxclust')

        tmp_df = pd.DataFrame(ord_ind, index=s_df.index, columns=['cluster'])

        for hemi_char in ['l_', 'r_']:
            t2_df = tmp_df[tmp_df.index.str.startswith(hemi_char)]
            t2_df.index = t2_df.index.map(lambda x: x.replace(hemi_char, ''))
            t2_df.to_csv(f'results/dp_clust_{granularity_id}_{sweep_name}_{hemi_char}{name}.csv', index_label='label')


    g_obj.savefig(f'figures/dp_interp_{sweep_name}_{name}.pdf')
    plt.close()


# num_males_test1 = 3305
# num_females_test1 = 3727

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--sweep_name')
    args = parser.parse_args()

    sweep_name = args.sweep_name

    s_male = np.load(f'results/dp_interp_{sweep_name}_male.npy')
    s_female = np.load(f'results/dp_interp_{sweep_name}_female.npy')

    s_total = s_male + s_female

    for s_arr, s_name in [(s_male, 'male'), (s_female, 'female'), (s_total, 'total')]:
        plot_and_save_interp(s_arr + s_arr.T, s_name, sweep_name)

import argparse

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.cluster.hierarchy import fcluster

from utils_datasets import STRUCT_COLUMNS


# python diffpool_plot_assignment.py --sweep_name 100_n_diffpool
# python diffpool_plot_assignment.py --sweep_name 100_n_e_diffpool

def plot_and_save_interp(arr, name, sweep_name):
    s_df = pd.DataFrame(arr, index=STRUCT_COLUMNS, columns=STRUCT_COLUMNS)

    # First create a dummy clustermap to know how the dendrogram is created and find the right mask next
    g_obj = sns.clustermap(s_df, yticklabels=1, xticklabels=1, dendrogram_ratio=(0.1, 0.2),
                           cbar_pos=(0, 0.85, 0.02, 0.15), cmap="viridis")
    mask_array = np.full(arr.shape, False)
    mask_array[np.tril_indices(mask_array.shape[0])] = True

    mask_after = mask_array[np.argsort(g_obj.dendrogram_row.reordered_ind), :]
    mask_after = mask_after[:, np.argsort(g_obj.dendrogram_col.reordered_ind)]

    g_obj = sns.clustermap(s_df, yticklabels=1, xticklabels=1, dendrogram_ratio=(0.1, 0.2),
                           cbar_pos = (0, 0.85, 0.02,0.15), cmap = "viridis", mask = mask_after,
                           linewidths = 0.5, linecolor = (0.7, 0.7, 0.7, 0.2))

    g_obj.ax_heatmap.set_xticklabels(g_obj.ax_heatmap.get_xmajorticklabels(), fontsize=7)
    g_obj.ax_heatmap.set_yticklabels(g_obj.ax_heatmap.get_ymajorticklabels(), fontsize=7)

    if sweep_name == '100_n_e_diffpool':
        granularities = [4, 8, 12]
    elif sweep_name == '100_n_diffpool':
        granularities = [4]

    for granularity_id in granularities:
        ord_ind = fcluster(g_obj.dendrogram_col.linkage, granularity_id, criterion='maxclust')

        tmp_df = pd.DataFrame(ord_ind, index=s_df.index, columns=['cluster'])

        for hemi_char in ['l_', 'r_']:
            t2_df = tmp_df[tmp_df.index.str.startswith(hemi_char)]
            t2_df.index = t2_df.index.map(lambda x: x.replace(hemi_char, ''))
            t2_df.to_csv(f'results/dp_clust_{granularity_id}_{sweep_name}_{hemi_char}{name}.csv', index_label='label')


    g_obj.savefig(f'figures/dp_interp_{sweep_name}_{name}.pdf')
    plt.close()


# num_males_test1 = 3305
# num_females_test1 = 3727

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--sweep_name')
    args = parser.parse_args()

    sweep_name = args.sweep_name

    s_male = np.load(f'results/dp_interp_{sweep_name}_male.npy')
    s_female = np.load(f'results/dp_interp_{sweep_name}_female.npy')

    s_total = s_male + s_female

    for s_arr, s_name in [(s_male, 'male'), (s_female, 'female'), (s_total, 'total')]:
        plot_and_save_interp(s_arr + s_arr.T, s_name, sweep_name)
# python diffpool_plot_assignment.py --sweep_name 100_n_diffpool # python diffpool_plot_assignment.py --sweep_name 100_n_e_diffpool # First create a dummy clustermap to know how the dendrogram is created and find the right mask next # num_males_test1 = 3305 # num_females_test1 = 3727