<jupyter_start><jupyter_text>Identifiability Test of nonlinear SSA on Synthetic Dataset<jupyter_code>%load_ext autoreload
%autoreload 2
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
import leap
import numpy as np
from leap.datasets.sim_dataset import DANS
from leap.modules.sssa import SSA
from leap.tools.utils import load_yaml
import random
import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib inline
use_cuda = True
device = torch.device("cuda:0" if use_cuda else "cpu")
latent_size = 4
content_size = 3
style_size = 1
data = DANS(directory = '/srv/data/ltcl/data/', 
            dataset='da_gau_5')
num_validation_samples = 4096
train_data, val_data = random_split(data, [len(data)-num_validation_samples, num_validation_samples])
train_loader = DataLoader(train_data, batch_size=25600, shuffle=True, pin_memory=True)
val_loader   = DataLoader(val_data, batch_size=16, shuffle=False, pin_memory=True)
cfg = load_yaml('../leap/configs/toy_gau_5.yaml')
model = SSA.load_from_checkpoint(
            checkpoint_path="/srv/data/ltcl/log/weiran/toy_gau_5/lightning_logs/version_31/checkpoints/epoch=62-step=16442.ckpt",
            input_dim=cfg['VAE']['INPUT_DIM'],
            c_dim=cfg['VAE']['CONTENT_DIM'],
            s_dim=cfg['VAE']['STYLE_DIM'],
            nclass=cfg['VAE']['NCLASS'],
            hidden_dim=cfg['VAE']['ENC']['HIDDEN_DIM'],
            bound=cfg['SPLINE']['BOUND'],
            count_bins=cfg['SPLINE']['BINS'],
            order=cfg['SPLINE']['ORDER'],
            beta=cfg['VAE']['BETA'],
            gamma=cfg['VAE']['GAMMA'],
            sigma=cfg['VAE']['SIGMA'],
            lr=cfg['VAE']['LR'],
            use_warm_start=cfg['SPLINE']['USE_WARM_START'],
            spline_pth=cfg['SPLINE']['PATH'],
            decoder_dist=cfg['VAE']['DEC']['DIST'],
            correlation=cfg['MCC']['CORR'])
model.eval()
model.to('cpu')
for batch in train_loader:
    break
batch_size = batch['x'].shape[0]
from scipy.stats import spearmanr<jupyter_output><empty_output><jupyter_text>Visualize style components<jupyter_code>A
B.shape
plt.scatter(A,B)
np.abs(spearmanr(B, A).correlation)
zs, mu, logvar = model.forward(batch)
mu = mu.view(batch_size, latent_size)
A = mu[:,-style_size:].detach().cpu().numpy()
B = batch['y'][:,-style_size:].detach().cpu().numpy()
C = np.zeros((style_size,style_size))
for i in range(style_size):
    C[i] = -np.abs(spearmanr(B, A).correlation[i,style_size:])
from scipy.optimize import linear_sum_assignment
row_ind, col_ind = linear_sum_assignment(C)
A = A[:, col_ind]
mask = np.ones(style_size)
for i in range(style_size):
    if np.corrcoef(B, A, rowvar=False)[i,style_size:][i] > 0:
        mask[i] = -1
print("Permutation:",col_ind)
print("Sign Flip:", mask)
fig = plt.figure(figsize=(2,2))
sns.heatmap(-C, vmin=0, vmax=1, annot=True, fmt=".2f", linewidths=.5, cbar=False, cmap='Greens')
plt.xlabel("Estimated latents ") 
plt.ylabel("True latents ") 
plt.title("MCC=%.3f"%np.abs(C[row_ind, col_ind]).mean());
A = mu[:,:].detach().cpu().numpy()
B = batch['y'][:,:].detach().cpu().numpy()
C = np.zeros((latent_size,latent_size))
fig, axs = plt.subplots(latent_size,latent_size, figsize=(4,4))
for i in range(latent_size):
    for j in range(latent_size):
        ax = axs[i, j]
        ax.scatter(B[:,i], A[:,j], s=2, color='b', alpha=0.25)
        ax.axis('off')
#     ax.set_xlabel('Ground truth latent')
#     ax.set_ylabel('Estimated latent')
#     ax.grid('..')
# plt.xlabel("True (4 content + 4 style)")
# plt.ylabel("Estimated (4 content + 4 style)")
fig.tight_layout()
fig, axs = plt.subplots(1,2, figsize=(2,1))
for i in range(style_size):
    ax = axs[i]
    ax.scatter(B[:,-style_size:], A[:,-style_size:], s=4, color='b', alpha=0.25)
    ax.axis('off')
#     ax.set_xlabel('Ground truth latent')
#     ax.set_ylabel('Estimated latent')
#     ax.grid('..')
fig.tight_layout()<jupyter_output><empty_output><jupyter_text>Visualize content components<jupyter_code>Y = mu[:,:content_size].detach().cpu().numpy()
X = batch['y'][:,:].detach().cpu().numpy()
fig = plt.figure(figsize=(4,4))
plt.scatter(Y[:,0], Y[:,1], s=1)
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.33, random_state=42)
rf = RandomForestRegressor()
rf.fit(X_train, y_train)
Y_pred = rf.predict(X)
print("R2: %.3f"%rf.score(X_test, y_test))
n_samples = 100
plt.figure(figsize=(3,3))
plt.scatter(Y[:n_samples,0], Y[:n_samples,1], marker="+", label='True', color='b')
plt.scatter(Y_pred[:n_samples,0], Y_pred[:n_samples,1], marker='*', label='Estimated', color='r')
plt.legend()
plt.ylabel('2nd dimension')
plt.xlabel('1st dimension')
np.random.seed(1)
Ncomp_s=4
Nsegment=10
Nsegment=10
modMat = np.random.uniform(0., 1, (Ncomp_s, Nsegment))
hidden=20
n_sample=1500
plt.plot(modMat.reshape(-1))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1816895715440582, 'fraction_numerical': 0.018997574777687955, 'mean_word_length': 5.07985257985258, 'pattern_counts': {'":': 0, '<': 13, '<?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': '483300', 'license': 'MIT', 'n_tokens_mistral': 1975, 'n_tokens_neox': 1913, 'n_words': 322, 'path': 'tests/test_ssa.ipynb', 'repo_name': 'weirayao/leap', 'size': '4948'}