<jupyter_start><jupyter_text>Distance Correlation<jupyter_code>import dcor
nmad_threshold = 5
caps_dim = caps_corr.shape[1]
num_features = cat_corr.shape[1]
dcorr_mat = np.zeros((num_features, caps_dim))
for i in range(num_features):
    x = caps_corr.T
    y = cat_corr[:,i]
    finite_mask = np.isfinite(y)
    y = y[finite_mask]
    x = x[:,finite_mask]
    median = np.median(y)
    nmad = np.abs(stats.median_abs_deviation(y, scale="normal"))
    mad_mask = (y>= (median - nmad_threshold*nmad)) & (y<= (median + nmad_threshold*nmad))
    y = y[mad_mask]
    x = x[:,mad_mask]
    y = np.repeat(y[np.newaxis,:], x.shape[0], 0)
    dcorr_mat[i] = dcor.rowwise(dcor.distance_correlation, x, y, compile_mode=dcor.CompileMode.COMPILE_PARALLEL)
    print(f"{cat.columns.to_list()[i]} percent rejected: {(~mad_mask).sum()*100/len(mad_mask)}")
y_labels_phot = [
    r"$u$",
    r"$g$",
    r"$r$",
    r"$i$",
    r"$z$",
    r"$u-g$",
    r"$u-r$",
    r"$u-i$",
    r"$u-z$",
    r"$g-r$",
    r"$g-i$",
    r"$g-z$",
    r"$r-i$",
    r"$r-z$",
    r"$i-z$",
    r"$n_{r}$",
    r"$R_{90, r}$",
]

y_labels_spec = [
    r"$z_{spec}$",
    r"$\mathrm{M}_{u}$",
    r"$\mathrm{M}_{g}$",
    r"$\mathrm{M}_{r}$",
    r"$\mathrm{M}_{i}$",
    r"$\mathrm{M}_{z}$",
    "log(M$_{\star}$)",
    "log(SFR)",
    "log(sSFR)",
    r"$\sigma_{v}$",
]
fig, ax = plt.subplots(2, 1, figsize=(25, 20), sharex=True, gridspec_kw={'height_ratios': [1.7, 1]})
_ = sns.heatmap(
    dcorr_mat[:17],
    linewidths=0.2,
    annot=True,
    cmap="rocket",
    cbar=False,
    xticklabels=dim_names,
    yticklabels=y_labels_phot,
    #     yticklabels=cat.columns.to_list(),
    ax=ax[0],
    robust=True,
    annot_kws={"fontsize": 20},
    vmin=0,
    vmax=1,
)

_ = sns.heatmap(
    dcorr_mat[17:],
    linewidths=0.2,
    annot=True,
    cmap="rocket",
    cbar=False,
    xticklabels=dim_names,
    yticklabels=y_labels_spec,
    #     yticklabels=cat.columns.to_list(),
    ax=ax[1],
    robust=True,
    annot_kws={"fontsize": 20},
    vmin=0,
    vmax=1,
)
fig.subplots_adjust(hspace=0.05)

cbar = fig.colorbar(ax[0].collections[0], ax=ax)
cbar.ax.tick_params(axis="both", which="major", labelsize=25)
cbar.ax.set_ylabel("Distance Correlation", fontsize=40, labelpad=30)


ax[0].tick_params(axis="both", which="major", labelsize=25, labeltop=True, bottom=False, top=True, left=True, right=False)
ax[0].tick_params(axis="both", which="minor", labelsize=25)

ax[1].tick_params(axis="both", which="major", labelsize=25, labeltop=False, bottom=True, top=False, left=True, right=False)
ax[1].tick_params(axis="both", which="minor", labelsize=25)


ax[1].set_xlabel("Capsule Dimension", size=40)
fig.text(0.05,0.4,"Galaxy Property", size=40, rotation=90 )

fig.savefig("./figs/correlations.pdf", dpi=300, bbox_inches="tight")
# plt.scatter(cat["sersicN_r"],cat["sersic_R50_r"],marker=".")
# plt.xlabel("n")
# plt.ylabel("R")<jupyter_output><empty_output><jupyter_text>Correlations among capsule dims<jupyter_code>caps_dim = caps_corr.shape[1]
dcorr_caps_mat = np.zeros((caps_dim, caps_dim))
for i in range(caps_dim):
    x = caps_corr.T
    y = caps_corr[:,i]
    y = np.repeat(y[np.newaxis,:], x.shape[0], 0)
    dcorr_caps_mat[i] = dcor.rowwise(dcor.distance_correlation, x, y, compile_mode=dcor.CompileMode.COMPILE_PARALLEL)
fig, ax = plt.subplots(1, 1, figsize=(25, 20))
ax = sns.heatmap(
    dcorr_caps_mat,
    linewidths=0.2,
    annot=True,
    cmap="rocket",
    xticklabels=dim_names,
    yticklabels=dim_names,
#     yticklabels=cat.columns.to_list(),
    ax=ax,
    robust=True,
    annot_kws={"fontsize": 20},
    vmin=0,
    vmax=1,
)
cbar = ax.collections[0].colorbar
cbar.ax.tick_params(axis="both", which="major", labelsize=25)
cbar.ax.set_ylabel("Distance Correlation", fontsize=40)
ax.tick_params(axis="both", which="major", labelsize=25, labeltop=True)
ax.tick_params(axis="both", which="minor", labelsize=25)
ax.set_xlabel("Capsule Dimension", size=40)
ax.set_ylabel("Capsule Dimension", size=40)
fig.savefig("./figs/correlations.pdf", dpi=300, bbox_inches="tight")
from scipy.cluster import hierarchy
from scipy.spatial.distance import squareform
corr_linkage = hierarchy.linkage(dcorr_caps_mat, method="ward")#, optimal_ordering=True)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(30,15))
dendro = hierarchy.dendrogram(corr_linkage, labels=dim_names, leaf_rotation=90, ax=ax1,)

dendro_idx = np.arange(0, len(dendro['ivl']))


ax2 = sns.heatmap(
    dcorr_caps_mat[dendro['leaves'], :][:, dendro['leaves']],
    linewidths=0.2,
    annot=True,
    cmap="rocket",
    xticklabels=dendro['ivl'],
    yticklabels=dendro['ivl'],
#     yticklabels=cat.columns.to_list(),
    ax=ax2,
    robust=True,
    annot_kws={"fontsize": 20},
    vmin=0,
    vmax=1,
)
cbar = ax2.collections[0].colorbar
cbar.ax.tick_params(axis="both", which="major", labelsize=25)
cbar.ax.set_ylabel("Distance Correlation", fontsize=40)
ax2.tick_params(axis="both", which="major", labelsize=25, labeltop=True)
ax2.tick_params(axis="both", which="minor", labelsize=25)
ax2.set_xlabel("Capsule Dimension", size=40)
ax2.set_ylabel("Capsule Dimension", size=40)
fig.tight_layout()
plt.show()




from collections import defaultdict
cluster_ids = hierarchy.fcluster(corr_linkage, 0.5, criterion='distance')
cluster_id_to_feature_ids = defaultdict(list)
for idx, cluster_id in enumerate(cluster_ids):
    cluster_id_to_feature_ids[cluster_id].append(idx)
selected_features = [v[0]+1 for v in cluster_id_to_feature_ids.values()]
selected_features<jupyter_output><empty_output><jupyter_text>Spearman's correlation (not for paper)<jupyter_code>from scipy.stats import spearmanr
spearman_corr = spearmanr(cat_corr[mask],caps_corr[mask],)[0]
spearman_corr = spearman_corr[:cat.shape[1],cat.shape[1]:]
plt.figure(figsize=(20,15))
sns.heatmap(spearman_corr, annot=True, cmap="icefire", xticklabels=dim_names, yticklabels=cat.columns.to_list())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1849603776766144, 'fraction_numerical': 0.026471084134210083, 'mean_word_length': 5.250790305584826, 'pattern_counts': {'":': 4, '<': 14, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '62592', 'license': 'MIT', 'n_tokens_mistral': 2608, 'n_tokens_neox': 2463, 'n_words': 345, 'path': 'notebooks/plot_correlations.ipynb', 'repo_name': 'galenvincent/encapZulate-1', 'size': '5931'}