<jupyter_start><jupyter_text>Quantification of GAIN imputation performance<jupyter_code>m_rate = 0.5
x_test_extended, cat_covs_test_extended, num_covs_test_extended, mask_extended = get_test_extended(x_test,
                                                                                          cat_covs_test,
                                                                                          num_covs_test,
                                                                                          tissues_test,
                                                                                          n = 171*2,
                                                                                          m_low=m_rate, m_high=m_rate)

# Get predictions
x_test_extended_ = x_test_extended * mask_extended
# Evaluate data
print('Generating synthetic data...')
x_gen_extended = predict(x=x_test_extended_,
                         cc=cat_covs_test_extended,
                         nc=num_covs_test_extended,
                         mask=mask_extended,
                         gen=gen)
x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended_
# Compute R^2
r_sq = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)
r_sq
# Save scores
df = pd.DataFrame(data=r_sq, index=tissues_dict_inv, columns=symbols)
df.to_csv('R2_tissue_genes.csv')
df.head()
counter_train = Counter(tissues_train)
nb_train = np.array([counter_train[t] for t in np.arange(nb_tissues)])
fig=plt.figure(figsize=(18, 6))
plt.tight_layout()
box_plot_scores(r_sq, nb_train, [t.replace('_', ' ') for t in tissues_dict_inv])
plt.savefig('r2_scores_horizontal.pdf', bbox_inches = 'tight', pad_inches = 0);
np.mean(r_sq)
mse_gain = compute_mse_score(x_test_extended, x_gen_extended_, mask_extended)
mse_gain
np.mean(mse_gain)
def box_plot_scores_vertical(r_sq, nb_train, tissues_dict_inv):
    colors = plt.cm.rainbow(nb_train.flatten())
    sns.boxplot(y=np.array([[t]*r_sq.shape[1] for t in tissues_dict_inv]).flatten(), x=r_sq.flatten(),
                sym='')
    ax = plt.gca()

    cmap = plt.cm.summer # Greens #Blues # viridis
    colors = cmap(nb_train.flatten())
    for i, a in enumerate(ax.artists):
        # Change the appearance of that box
        a.set_facecolor(colors[i])
        # a.set_edgecolor('black')
        # a.set_linewidth(3)

    norm = plt.Normalize(nb_train.min(), nb_train.max())
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    
    
    plt.colorbar(sm, fraction=0.03, pad=0.04)
    # plt.xticks(rotation=90)

    # ax.figure.colorbar(colors)
    ax.set_xlim([0, 1])
    plt.title('$R^2$ imputation scores per tissue')
    return plt.gca()
counter_train = Counter(tissues_train)
nb_train = np.array([counter_train[t] for t in np.arange(nb_tissues)])
fig=plt.figure(figsize=(7, 13))
plt.tight_layout()
box_plot_scores_vertical(r_sq, nb_train, [t.replace('_', ' ') for t in tissues_dict_inv])
plt.savefig('r2_scores_vertical.pdf', bbox_inches = 'tight', pad_inches = 0);<jupyter_output><empty_output><jupyter_text>Baselines  Median/mean imputation<jupyter_code>nb_train, nb_genes = x_train.shape
tissues_train = cat_covs_train[:, 2]
gene_medians = []
gene_means = []

for t in range(nb_tissues):
    idxs_tissue = np.argwhere(tissues_train == t).ravel()
    medians = np.median(x_train[idxs_tissue], axis=0)
    means = np.mean(x_train[idxs_tissue], axis=0)
    gene_medians.append(medians)
    gene_means.append(means)
    
gene_medians = np.array(gene_medians)
gene_means = np.array(gene_means)
gene_means
gene_medians
x_test_extended_r = x_test_extended.reshape((nb_tissues, -1, nb_genes))  # Ground truth
x_test_extended_r_ = x_test_extended_.reshape((nb_tissues, -1, nb_genes))  # Ground truth with missing values
mask_r = np.copy(mask_extended.reshape((nb_tissues, -1, nb_genes)))
assert (x_test_extended_r_ == (x_test_extended_r * mask_r)).all()

x_test_median_imp = x_test_extended_r * mask_r + gene_medians[:, None, :] * (1 - mask_r) 
x_test_mean_imp = x_test_extended_r * mask_r + gene_means[:, None, :] * (1 - mask_r) 
r_sq_medians = compute_r2_score(x_test_extended, x_test_median_imp.reshape((-1, nb_genes)), mask_extended)
r_sq_means = compute_r2_score(x_test_extended, x_test_mean_imp.reshape((-1, nb_genes)), mask_extended)
r_sq_medians
np.mean(r_sq_medians)
np.mean(r_sq_means)
np.mean(r_sq_means, axis=-1)
mse_median = compute_mse_score(x_test_extended, x_test_median_imp.reshape((-1, nb_genes)), mask_extended)
np.sqrt(np.mean(mse_median))
r_sqs_mean = []
r_sqs_median = []
rmses_mean = []
rmses_median = []

for n in range(10):
    print(n+1)
    x_test_extended, cat_covs_test_extended, num_covs_test_extended, mask_extended = get_test_extended(x_test,
                                                                                              cat_covs_test,
                                                                                              num_covs_test,
                                                                                              tissues_test,
                                                                                              n = 171*2,
                                                                                              m_high=m_rate,
                                                                                              m_low=m_rate)
    x_test_extended_ = x_test_extended * mask_extended

    x_test_extended_r = x_test_extended.reshape((nb_tissues, -1, nb_genes))  # Ground truth
    x_test_extended_r_ = x_test_extended_.reshape((nb_tissues, -1, nb_genes))  # Ground truth with missing values
    mask_r = np.copy(mask_extended.reshape((nb_tissues, -1, nb_genes)))
    assert (x_test_extended_r_ == (x_test_extended_r * mask_r)).all()

    x_test_median_imp = x_test_extended_r * mask_r + gene_medians[:, None, :] * (1 - mask_r) 
    x_test_mean_imp = x_test_extended_r * mask_r + gene_means[:, None, :] * (1 - mask_r) 
    
    r_sq_medians = compute_r2_score(x_test_extended, x_test_median_imp.reshape((-1, nb_genes)), mask_extended)
    r_sq_means = compute_r2_score(x_test_extended, x_test_mean_imp.reshape((-1, nb_genes)), mask_extended)
    mse_medians = compute_mse_score(x_test_extended, x_test_median_imp.reshape((-1, nb_genes)), mask_extended)
    mse_means = compute_mse_score(x_test_extended, x_test_mean_imp.reshape((-1, nb_genes)), mask_extended)

    print('R^2 Mean: ', np.mean(r_sq_means))
    print('R^2 Median: ', np.mean(r_sq_medians))
    rmses_mean.append(mse_means)
    rmses_median.append(mse_medians)
    r_sqs_mean.append(r_sq_means)
    r_sqs_median.append(r_sq_medians)
np.std(r_sqs_mean)
r_sqs_mean_ = np.array(r_sqs_mean)
np.mean(r_sqs_mean_)
np.mean(np.std(r_sqs_mean_, axis=0))
rmses_mean_ = np.sqrt(np.array(rmses_median))
np.mean(rmses_mean_)
np.mean(np.std(rmses_mean_, axis=0))<jupyter_output><empty_output><jupyter_text>Whole blood surrogate<jupyter_code>blood_idx = tissues_dict['Whole_Blood']
blood_idxs = np.argwhere(tissues_test == blood_idx).ravel()
surrogate_expr = x_test[blood_idxs, :]
patients_with_blood_tissue = sampl_ids_test[blood_idxs]

blood_idxs
x_test_blood_extended = [[] for _ in range(nb_tissues)]
x_test_blood_surrogate = [[] for _ in range(nb_tissues)]

for i, idx in enumerate(blood_idxs):
    patient_id = sampl_ids_test[idx]
    blood_expr = x_test[idx, :]  # surrogate_expr[i, :]
    samples_patient = [j for j, s in enumerate(sampl_ids_test) if s == patient_id]
    
    for s in samples_patient:
        t = tissues_test[s]
        expr = x_test[s, :]
        x_test_blood_extended[t].append(expr)
        x_test_blood_surrogate[t].append(blood_expr)    
n = 171*2
x_surrogate_extended = np.empty((nb_tissues * n, nb_genes), dtype=np.float32)
x_test_tissue_extended = np.empty((nb_tissues * n, nb_genes), dtype=np.float32)
for t in range(nb_tissues):
    x_t = np.array(x_test_blood_extended[t])  # Tissue expression that we want to impute
    x_s = np.array(x_test_blood_surrogate[t])  # Blood expression
    if np.all(x_t == x_s):  # Imputing whole blood
        assert t == blood_idx
    x_surrogate_extended[t*n: (t+1)*n, :] = np.resize(x_s, [n, nb_genes])
    x_test_tissue_extended[t*n: (t+1)*n, :] = np.resize(x_t, [n, nb_genes])
assert np.all(x_test_tissue_extended[n] == x_test_blood_extended[1][0])
assert np.all(x_test_tissue_extended[10*n+2] == x_test_blood_extended[10][2])
np.resize([[0, 7], [1, 1], [2, 2]], [6, 2])
x_test_tissue_extended_r = x_test_tissue_extended.reshape((nb_tissues, -1, nb_genes))  # Ground truth
x_surrogate_extended_r = x_surrogate_extended.reshape((nb_tissues, -1, nb_genes))  # Blood surrogate

i = 2
j = 3
assert np.all(x_test_tissue_extended_r[i, j, :] == x_test_tissue_extended[i*n + j])
n = 10
m_rate = 0.5

r_sqs_surrogate = []
rmses_surrogate = []
for n in range(10):
    print(n+1)
    
    # Generate mask
    mask_extended, _, _ = get_mask_hint_b(x_test_tissue_extended.shape[0], nb_genes, m_low=m_rate, m_high=m_rate)
    mask_extended_r = mask_extended.reshape((nb_tissues, -1, nb_genes))
    
    # Perform blood surrogate imputation and compute scores
    x_test_surrogate_imp = x_test_tissue_extended_r * mask_extended_r + x_surrogate_extended_r * (1 - mask_extended_r)
    r_sq_surrogate = compute_r2_score(x_test_tissue_extended, x_test_surrogate_imp.reshape((-1, nb_genes)), mask_extended)
    mse_surrogate = compute_mse_score(x_test_tissue_extended, x_test_surrogate_imp.reshape((-1, nb_genes)), mask_extended)
    
    # Append scores
    r_sqs_surrogate.append(r_sq_surrogate)
    rmses_surrogate.append(np.sqrt(mse_surrogate))
rmses_surrogate_ = np.array(rmses_surrogate)
means = np.nanmean(rmses_surrogate_)
stds = np.nanmean(np.std(rmses_surrogate_, axis=0))
r_sqs_surrogate_ = np.array(r_sqs_surrogate)
np.mean(r_sqs_surrogate_)
np.nanmean(np.std(r_sqs_surrogate_, axis=0))
np.array(r_sqs_surrogate).shape
np.std(np.mean(r_sqs_surrogate, axis=(1, 2)))
np.std(rmses_surrogate)
mask_r = np.copy(mask_extended.reshape((nb_tissues, -1, nb_genes)))

x_test_surrogate_imp = x_test_tissue_extended_r * mask_r + x_surrogate_extended_r * (1 - mask_r) 
r_sq_surrogate = compute_r2_score(x_test_tissue_extended, x_test_surrogate_imp.reshape((-1, nb_genes)), mask_extended)
r_sq_surrogate
mse_surrogate = compute_mse_score(x_test_tissue_extended, x_test_surrogate_imp.reshape((-1, nb_genes)), mask_extended)
np.mean(r_sq_surrogate)
np.mean(mse_surrogate)
plt.figure(figsize=(16, 5))
sns.distplot(np.mean(r_sq_surrogate, axis=0));
df = pd.DataFrame(data=r_sq_surrogate, index=tissues_dict_inv, columns=symbols)
df.to_csv('R2_tissue_genes_blood_surrogate.csv')
df.head()<jupyter_output><empty_output><jupyter_text>GAIN without adversarial loss<jupyter_code>gen_mse = tfk.models.load_model(MODELS_DIR + 'gen_gain_mse.h5')

print('Generating synthetic data...')
x_gen = predict(x=x_test_,
                cc=cat_covs_test,
                nc=num_covs_test,
                mask=mask,
                gen=gen_mse)

gamma_dx_dz = gamma_coef(x_test, x_gen)
gamma_dx_dz_upper_bound = gamma_coef(x_test, x_train[:x_test.shape[0], :])
print('Gamma(Dx, Dz): {:.3f}/{:.3f}'.format(gamma_dx_dz, gamma_dx_dz_upper_bound))

imp_mse = np.sum((1-mask)*(x_gen - x_test)**2)/np.sum(1-mask)
imp_mse
"""
x_test_extended, cat_covs_test_extended, num_covs_test_extended, mask_extended = get_test_extended(x_test,
                                                                                          cat_covs_test,
                                                                                          num_covs_test,
                                                                                          tissues_test,
                                                                                          n = 171*2)

# Get predictions
x_test_extended_ = x_test_extended * mask_extended
"""


# Evaluate data
print('Generating synthetic data...')
x_gen_extended = predict(x=x_test_extended_,
                         cc=cat_covs_test_extended,
                         nc=num_covs_test_extended,
                         mask=mask_extended,
                         gen=gen_mse)
x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended_
# Compute R^2
r_sq_gain_mse = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)
r_sq_gain_mse
np.mean(r_sq_gain_mse)
x_test_extended.shape
nb_tissues=49
counter_train = Counter(tissues_train)
nb_train = np.array([counter_train[t] for t in np.arange(nb_tissues)])
fig=plt.figure(figsize=(18, 6))
box_plot_scores(r_sq_gain_mse, nb_train, tissues_dict_inv);
mse_gain_mse = compute_mse_score(x_test_extended, x_gen_extended_.reshape((-1, nb_genes)), mask_extended)
mse_gain_mse
np.mean(mse_gain_mse)
# Compute intervals
m_rate = 0.5

r_sqs = []
r_sqs_mse = []
rmses = []
rmses_mse = []

for n in range(10):
    x_test_extended, cat_covs_test_extended, num_covs_test_extended, mask_extended = get_test_extended(x_test,
                                                                                              cat_covs_test,
                                                                                              num_covs_test,
                                                                                              tissues_test,
                                                                                              n = 171*2,
                                                                                              m_high=m_rate,
                                                                                              m_low=m_rate)

    # Get predictions
    x_test_extended_ = x_test_extended * mask_extended

    # Evaluate data for GAIN
    print('Generating synthetic data...')
    x_gen_extended = predict(x=x_test_extended_,
                             cc=cat_covs_test_extended,
                             nc=num_covs_test_extended,
                             mask=mask_extended,
                             gen=gen)
    x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended_
    
    r_sq_gain = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)
    r_sqs.append(r_sq_gain)
    
    mse_gain = compute_mse_score(x_test_extended, x_gen_extended_.reshape((-1, nb_genes)), mask_extended)
    rmses.append(mse_gain)
    
    print('R^2 GAIN: ', np.mean(r_sq_gain))
    
    # Evaluate data for GAIN_MSE
    print('Generating synthetic data...')
    x_gen_extended = predict(x=x_test_extended_,
                             cc=cat_covs_test_extended,
                             nc=num_covs_test_extended,
                             mask=mask_extended,
                             gen=gen_mse)
    x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended_
    
    r_sq_gain_mse = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)
    r_sqs_mse.append(r_sq_gain_mse)
    
    mse_gain_mse = compute_mse_score(x_test_extended, x_gen_extended_.reshape((-1, nb_genes)), mask_extended)
    rmses_mse.append(mse_gain_mse)
    
    print('R^2 GAIN MSE: ', np.mean(r_sq_gain_mse))

np.std(r_sq_gain, axis=-1)
np.std(rmses)
r_sqs_ = np.array(r_sqs_mse)
means = np.nanmean(r_sqs_)
stds = np.nanmean(np.std(r_sqs_, axis=0))
means
stds
rmses_ = np.array(np.sqrt(rmses_))
means = np.nanmean(rmses_)
stds = np.nanmean(np.std(rmses_, axis=0))
means
stds<jupyter_output><empty_output><jupyter_text>R score for different levels of missing data<jupyter_code>m_rate = 0.5
mask_extended, _, _ = get_mask_hint_b(x_test_extended.shape[0], nb_genes, m_low=m_rate, m_high=m_rate)
# mask_extended = np.repeat(mask_extended, x_test_tissue_extended.shape[0], axis=0)

x_gen_extended = predict(x=x_test_extended * mask_extended,
                         cc=cat_covs_test_extended,
                         nc=num_covs_test_extended,
                         mask=mask_extended,
                         gen=gen)
x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended * mask_extended
r_sq_gain = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)

scores = []
xs = np.append(np.arange(0, 1, 0.05), 0.99)
print(xs)

for m_rate in xs:
    mask_extended, _, _ = get_mask_hint_b(x_test_extended.shape[0], nb_genes, m_low=m_rate, m_high=m_rate)
    # mask_extended = np.repeat(mask_extended, x_test_tissue_extended.shape[0], axis=0)

    x_gen_extended = predict(x=x_test_extended * mask_extended,
                             cc=cat_covs_test_extended,
                             nc=num_covs_test_extended,
                             mask=mask_extended,
                             gen=gen)
    x_gen_extended_ = x_gen_extended * (1 - mask_extended) + x_test_extended * mask_extended
    s = compute_r2_score(x_test_extended, x_gen_extended_, mask_extended)
    # s = compute_mse_score(x_test_extended, x_gen_extended_.reshape((-1, nb_genes)), mask_extended)
    
    scores.append(s)
    print('Rate: {}. Score: {}'.format(m_rate, np.mean(s)))
len(scores)
means = np.array([np.mean(s) for s in scores])
stds = np.array([np.std(np.mean(s, axis=0)) for s in scores])

plt.figure(figsize=(12, 5))
t = 2
xs_ = xs[t:]
means_ = means[t:]
stds_ = stds[t:]
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df)
ax.fill_between(xs_, (means_-stds_), (means_+stds_), color='b', alpha=.1);
#plt.xlim(0, 1);<jupyter_output><empty_output><jupyter_text>Save data<jupyter_code># Save data to CSV for methods implemented in R
x_test_nan = x_test_
x_test_nan[mask == 0] = np.nan

x = np.concatenate((x_train, x_test_nan), axis=0)
num_covs = np.concatenate((num_covs_train, num_covs_test), axis=0)
cat_covs = np.concatenate((cat_covs_train, cat_covs_test), axis=0)
x_ = np.concatenate((x, num_covs, cat_covs), axis=-1)
df_csv = pd.DataFrame(data=x_, index=sampl_ids)
df_csv.astype({12558: 'int32', 12559: 'int32', 12560: 'int32'})
df_csv.to_csv('data/to_impute.csv')

# Save test data to CSV
df_test = pd.DataFrame(data=x_test, index=sampl_ids_test, columns=symbols)
tissues_test = cat_covs_test[:, 2]
df_test['tissue'] = tissues_test
df_test.head()
df_test.to_csv('observed_test.csv')

# Save imputed test data to CSV
x_gen_ = (1 - mask) * x_gen + mask * x_test
df_imputed = pd.DataFrame(data=x_gen_, index=sampl_ids_test, columns=symbols)
tissues_test = cat_covs_test[:, 2]
df_imputed['tissue'] = tissues_test
df_imputed.to_csv('imputed_test.csv')

# Save mask to CSV
df_mask = pd.DataFrame(data=mask, index=sampl_ids_test, columns=symbols)
tissues_test = cat_covs_test[:, 2]
df_mask.to_csv('mask.csv')
# Save whole blood data to CSV for methods implemented in R
blood_idx = tissues_dict['Whole_Blood']
blood_idxs = np.argwhere(tissues_test == blood_idx).ravel()
blood_expr_test = x_test[blood_idxs, :]
blood_cat_covs_test = cat_covs_test[blood_idxs, :]
blood_num_covs_test = num_covs_test[blood_idxs, :]
blood_sample_ids_test = sampl_ids_test[blood_idxs]

blood_idxs = np.argwhere(tissues_train == blood_idx).ravel()
blood_expr_train = x_train[blood_idxs, :]
blood_cat_covs_train = cat_covs_train[blood_idxs, :]
blood_num_covs_train = num_covs_train[blood_idxs, :]
blood_sample_ids_train = sampl_ids_train[blood_idxs]

# Save to impute data for blood tissue
m_rate = 0.5
nb_samples, nb_genes = blood_expr_test.shape
mask_blood, _, _ = get_mask_hint_b(nb_samples, nb_genes, m_low=m_rate, m_high=m_rate)
x_test_nan_test = blood_expr_test
x_test_nan_test[mask_blood == 0] = np.nan

x = np.concatenate((blood_expr_train, x_test_nan_test), axis=0)
num_covs = np.concatenate((blood_num_covs_train, blood_num_covs_test), axis=0)
cat_covs = np.concatenate((blood_cat_covs_train, blood_cat_covs_test), axis=0)
blood_sample_ids = np.concatenate((blood_sample_ids_train, blood_sample_ids_test), axis=0)
x_ = np.concatenate((x, num_covs, cat_covs), axis=-1)
df_csv = pd.DataFrame(data=x_, index=blood_sample_ids)
df_csv.astype({12558: 'int32', 12559: 'int32', 12560: 'int32'})
df_csv.to_csv('data/to_impute_blood.csv')

# Save mask to CSV
df_mask = pd.DataFrame(data=mask_blood, index=blood_sample_ids_test, columns=symbols)
df_mask.to_csv('mask_blood.csv')<jupyter_output><empty_output><jupyter_text>Tissue embeddings<jupyter_code>gen.layers[6].get_weights()[0].shape
embeddings = gen.layers[6].get_weights()[0]
emb_2d = tsne_2d(embeddings)
tissues_dict_inv
embeddings
plt.figure(figsize=(12, 12))
ax = plt.gca()
x1, x2 = emb_2d.T
ax.scatter(x1, x2)

for i, txt in enumerate(tissues_dict_inv):
    ax.annotate(txt, (x1[i], x2[i]),
                textcoords="offset points", # how to position the text
                xytext=(0,10), # distance from text to points (x,y)
                ha='center')
    
plt.axis('off');
emb_2d_ = tsne_2d(embeddings, perplexity=40)
plt.figure(figsize=(12, 12))
ax = plt.gca()
x1, x2 = emb_2d_.T
ax.scatter(x1, x2)

for i, txt in enumerate(tissues_dict_inv):
    ax.annotate(txt, (x1[i], x2[i]))
    
plt.axis('off');<jupyter_output><empty_output><jupyter_text>Evaluate on TCGA<jupyter_code>def scores_missing_rate(x_gt, cc_gt, nc_gt, gen, repeats=10):
    scores = []
    xs = np.arange(0.05, 0.9, 0.05)
    
    for m_rate in xs:
        for r in range(repeats):
            mask, _, _ = get_mask_hint_b(x_gt.shape[0], nb_genes, m_low=m_rate, m_high=m_rate)
            x_gt_ = x_gt * mask

            x_gen = predict(x=x_gt_,
                          cc=cc_gt,
                          nc=nc_gt,
                          mask=mask,
                          gen=gen)

            r_sq = r2_scores_(x_gt, x_gen, mask)
            scores.append(r_sq)

        print('Missing rate: {:.2f}. R^2: {:.2f}'.format(m_rate, np.nanmean(r_sq)))
    scores = np.reshape(scores, (-1, repeats, r_sq.shape[0]))
    
    return xs, scores

def r2_scores_(x_gt, x_pred, mask):
    mask_r = np.copy(mask)
    mask_r[mask_r == 1] = np.nan

    gene_means = np.mean(x_gt, axis=0)  # Shape=(nb_genes,)
    mask_r[:, gene_means == 0] = np.nan  # Discard genes with 0 variance
    ss_res = np.nansum((1-mask_r)*(x_gt - x_pred) ** 2, axis=0)
    ss_tot = np.nansum((1-mask_r)*(x_gt - gene_means) ** 2, axis=0)
    r_sq = 1 - ss_res/(ss_tot)
    return r_sq

def get_TCGA(TCGA_TISSUE = 'LUAD'):
    # TCGA data
    if TCGA_TISSUE == 'LAML':
        path = '/local/scratch/USER/tcga/AML/TCGA-LAML.htseq_fpkm.tsv'  # /LungAdenocarcinoma/TCGA-LUAD.htseq_fpkm.tsv'  # 'AML/TCGA-LAML.htseq_fpkm.tsv'
        clinical = '/local/scratch/USER/tcga/AML/LAML_clinicalMatrix'
        t = 'Whole_Blood'
    elif TCGA_TISSUE == 'BRCA':
        path = '/local/scratch/USER/tcga/BreastCancer/TCGA-BRCA.htseq_fpkm.tsv'  # /LungAdenocarcinoma/TCGA-LUAD.htseq_fpkm.tsv'  # 'AML/TCGA-LAML.htseq_fpkm.tsv'
        clinical = '/local/scratch/USER/tcga/BreastCancer/BRCA_clinicalMatrix'
        t = 'Breast_Mammary_Tissue'
    else:
        path = '/local/scratch/USER/tcga/LungAdenocarcinoma/TCGA-LUAD.htseq_fpkm.tsv'
        clinical = '/local/scratch/USER/tcga/LungAdenocarcinoma/LUAD_clinicalMatrix'
        t = 'Lung'
    t_idx = np.argwhere(tissues_dict_inv == t).ravel()[0]
    df = pd.read_csv(path, delimiter='\t')
    df = df.set_index('Ensembl_ID')
    
    # Transform gene symbols
    gene_symbols, ENSEMBL_found = ENSEMBL_to_gene_symbols(df.index)
    df = df.loc[ENSEMBL_found]
    df = df.rename(index=dict(zip(df.index, gene_symbols)))
    df = df.loc[symbols]
    df = df.groupby('Ensembl_ID', group_keys=False).apply(lambda x: x[x.sum(axis=1) == np.max(x.sum(axis=1))])  # Remove duplicates, keep max

    # Standardize data
    x_TCGA = df.values.T
    x_TCGA = (x_TCGA - np.mean(x_TCGA, axis=0))/(np.std(x_TCGA, axis=0) + 1e-7)

    # Process covariates
    clinical_df = pd.read_csv(clinical, delimiter='\t')
    idxs = [np.argwhere(s[:-1] == clinical_df['sampleID']).ravel()[0] for s in df.columns]
    gender = np.array([0 if g == 'MALE' else 1 for g in clinical_df.iloc[idxs]['gender']])
    age = clinical_df.iloc[idxs]['age_at_initial_pathologic_diagnosis'].values
    mean_age = 52.7763  # Mean age GTEx
    std_age = 12.9351  # Std age GTEx
    age = (age - mean_age)/std_age
    cc_TCGA = np.zeros((x_TCGA.shape[0], 3)) + stats.mode(cat_covs).mode
    cc_TCGA[:, 0] = gender
    cc_TCGA[:, 2] = t_idx
    nc_TCGA = age
    
    return x_TCGA, cc_TCGA, nc_TCGA, t, t_idx

def plot_r2_missingrate(xs, scores, label=None):
    means = np.nanmean(scores, axis=(1, 2)) #  np.array([np.nanmean(s) for s in scores_GTEx])
    stds = np.nanmean(np.nanstd(scores, axis=1), axis=-1)
    
    xs_ = xs
    means_ = means
    stds_ = stds
    df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
    ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label=label)
    ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);


def plot_r2_missingrate_(xs, scores_GTEx, scores_TCGA, gtex_tissue, tcga_tissue):
    means_GTEx = np.nanmean(scores_GTEx, axis=(1, 2)) #  np.array([np.nanmean(s) for s in scores_GTEx])
    stds_GTEx = np.nanmean(np.nanstd(scores_GTEx, axis=1), axis=-1)
    means_TCGA = np.nanmean(scores_TCGA, axis=(1, 2))
    stds_TCGA = np.nanmean(np.nanstd(scores_TCGA, axis=1), axis=-1)
    
    xs_ = xs
    means_ = means_GTEx
    stds_ = stds_GTEx
    df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
    ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='GTEx {}'.format(gtex_tissue))
    ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

    means_ = means_TCGA
    stds_ = stds_TCGA
    df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
    ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='TCGA {}'.format(tcga_tissue))
    ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);
x_TCGA, cc_TCGA, nc_TCGA, t, t_idx = get_TCGA('LUAD')
x_TCGA.shape
x_TCGA, cc_TCGA, nc_TCGA, t, t_idx = get_TCGA(TCGA_TISSUE)

mask, _, _ = get_mask_hint_b(x_TCGA.shape[0], nb_genes, m_low=0.5, m_high=0.5)
x_TCGA_ = x_TCGA * mask

x_gen_TCGA = predict(x=x_TCGA_,
                  cc=cc_TCGA,
                  nc=nc_TCGA,
                  mask=mask,
                  gen=gen)

r_sq_TCGA = r2_scores_(x_TCGA, x_gen_TCGA, mask)
np.nanmean(r_sq_TCGA)
idxs_GTEx = np.argwhere(tissues_test == t_idx).ravel()
x_GTEx = x_test[idxs_GTEx]
cc_GTEx = cat_covs_test[idxs_GTEx]
nc_GTEx = num_covs_test[idxs_GTEx]
xs, scores_GTEx = scores_missing_rate(x_GTEx, cc_GTEx, nc_GTEx, gen)
_, scores_TCGA = scores_missing_rate(x_TCGA, cc_TCGA, nc_TCGA, gen)
plt.figure(figsize=(18, 4))

sg = []
st = []
for i, tcga_tissue in enumerate(['LAML', 'BRCA', 'LUAD']):
    x_TCGA, cc_TCGA, nc_TCGA, t, t_idx = get_TCGA(tcga_tissue)
    
    idxs_GTEx = np.argwhere(tissues_test == t_idx).ravel()
    x_GTEx = x_test[idxs_GTEx]
    cc_GTEx = cat_covs_test[idxs_GTEx]
    nc_GTEx = num_covs_test[idxs_GTEx]

    print('GTEx {}'.format(t))
    xs, scores_GTEx = scores_missing_rate(x_GTEx, cc_GTEx, nc_GTEx, gen)
    print('Missing rate 0.5, R^2: {} ({})'.format(np.nanmean(scores_GTEx[9]), np.nanmean(np.nanstd(scores_GTEx[9], axis=0))))
    print('TCGA {}'.format(tcga_tissue))
    _, scores_TCGA = scores_missing_rate(x_TCGA, cc_TCGA, nc_TCGA, gen)
    print('Missing rate 0.5, R^2: {} ({})'.format(np.nanmean(scores_TCGA[9]), np.nanmean(np.nanstd(scores_TCGA[9], axis=0))))

    sg.append(scores_GTEx)
    st.append(scores_TCGA)
plt.figure(figsize=(15, 5))
plt.tight_layout()
plt.subplot(1, 2, 1)
plot_r2_missingrate(xs, sg[0], 'Whole Blood')
plot_r2_missingrate(xs, sg[1], 'Breast Mammary Tissue')
plot_r2_missingrate(xs, sg[2], 'Lung')
plt.title('GTEx')

plt.subplot(1, 2, 2)
plot_r2_missingrate(xs, st[0], 'LAML')
plot_r2_missingrate(xs, st[1], 'BRCA')
plot_r2_missingrate(xs, st[2], 'LUAD')

plt.title('TCGA')
plt.savefig('r2_missing_rate.pdf', bbox_inches = 'tight', pad_inches = 0);
np.nanmean(st[0][9])  #.shape
np.nanmean(np.nanstd(sg[0][9], axis=0))
means_GTEx = np.nanmean(scores_GTEx, axis=(1, 2)) #  np.array([np.nanmean(s) for s in scores_GTEx])
stds_GTEx = np.nanmean(np.nanstd(scores_GTEx, axis=1), axis=-1)
means_TCGA = np.nanmean(scores_TCGA, axis=(1, 2))
stds_TCGA = np.nanmean(np.nanstd(scores_TCGA, axis=1), axis=-1)

plt.figure(figsize=(18, 4))
ax=plt.subplot(1, 3, 1)
xs_ = xs
means_ = means_GTEx
stds_ = stds_GTEx
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='GTEx {}'.format(t))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

means_ = means_TCGA
stds_ = stds_TCGA
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='TCGA {}'.format(TCGA_TISSUE))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

plt.subplot(1, 3, 2)
xs_ = xs
means_ = means_GTEx
stds_ = stds_GTEx
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='GTEx {}'.format(t))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

means_ = means_TCGA
stds_ = stds_TCGA
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='TCGA {}'.format(TCGA_TISSUE))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

plt.subplot(1, 3, 3)
xs_ = xs
means_ = means_GTEx
stds_ = stds_GTEx
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='GTEx {}'.format(t))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);

means_ = means_TCGA
stds_ = stds_TCGA
df = pd.DataFrame(data=np.array([xs_, means_, stds_]).T, columns=['$1 -$missing rate', '$R^2$', 'std'])
ax = sns.lineplot(x='$1 -$missing rate', y='$R^2$', data=df, label='TCGA {}'.format(TCGA_TISSUE))
ax.fill_between(xs_, (means_-stds_), (means_+stds_), alpha=.1);<jupyter_output><empty_output><jupyter_text>Single-cell data<jupyter_code>path = '/local/scratch/USER/scRNA/GSM3244108_TS543.UMItab.txt'
df_sc = pd.read_csv(path, delimiter='\t')
df_sc
df_sc.shape
df_sc.head()
del df_sc['ENS_ID']
df_sc = df_sc.set_index('SYMBOL')
df_sc_ = df_sc.loc[symbols]
df_sc_ = df_sc_.fillna(0.)
df_sc_
df_sc_.index[np.argwhere(df_sc_.index.duplicated()).ravel()]
df_sc_ = df_sc_.groupby('SYMBOL', group_keys=False).apply(lambda x: x[x.sum(axis=1) == np.max(x.sum(axis=1))])
df_sc_ = df_sc_.loc[~df_sc_.index.duplicated(keep='first')]  # If still some duplicates remain (e.g. all are 0), drop all and keep first
df_sc_.index[np.argwhere(df_sc_.index.duplicated()).ravel()]
assert (df_sc_.index.values == symbols).all()
x_sc = df_sc_.values.T
x_sc = (x_sc - np.mean(x_sc, axis=0))/(np.std(x_sc, axis=0) + 1e-7)
x_sc.shape
np.max(x_sc)
mean_age = 52.7763  # Mean age GTEx
cc_sc = np.zeros((x_sc.shape[0], 3)) + stats.mode(cat_covs).mode
nc_sc = np.zeros((x_sc.shape[0])) #  + mean_age
mask, _, _ = get_mask_hint_b(x_sc.shape[0], nb_genes, m_low=0.5, m_high=0.5)
x_gen_sc = predict(x=x_sc,
                  cc=cc_sc,
                  nc=nc_sc,
                  mask=mask,
                  gen=gen)
x_gen_sc
r_sq_sc = r2_scores_(x_sc, x_gen_sc, mask)
np.nanmean(r_sq_sc)<jupyter_output>/auto/homes/USER/adversarial-gtex/gtexenv/lib/python3.5/site-packages/ipykernel_launcher.py:32: RuntimeWarning: invalid value encountered in true_divide<jupyter_text>R squared distributions<jupyter_code>tissues_test
sampl_ids
genes_r_sq = np.mean(r_sq, axis=0)
plt.figure(figsize=(14, 5))
plt.title('Distribution of genes R^2')
sns.distplot(genes_r_sq);
r_sq_genes = genes_r_sq
idxs = np.argwhere(r_sq_genes > 0.8).flatten()
symbols[idxs]
idxs = np.argwhere(r_sq_genes < 0.2).flatten()
symbols[idxs]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13547159215001092, 'fraction_numerical': 0.01628654332158123, 'mean_word_length': 3.3792867878125428, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 85, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '332457', 'license': 'MIT', 'n_tokens_mistral': 13331, 'n_tokens_neox': 12461, 'n_words': 1959, 'path': 'gtex_gain_analysis.ipynb', 'repo_name': 'gamazonlab/GAIN-GTEx', 'size': '32051'}