<jupyter_start><jupyter_text>OverviewWe've trained a sequence of elastic net models (with different $\lambda$'s and $\alpha$'s) and here we'd like to compute the PTRS on held out Britishs and other populations. So, for each $\lambda$ and $\alpha$ pair, we have one PTRS nsample x ntrait matrix. In total, we will have $$nsample \times (ntrait = 17) \times (nlambda = 50) \times (nalpha = 3)$$ PTRS tensor for each population.  Load predicted expression First load British and then load other populations.<jupyter_code>population = 'British'

# set path to British data
hdf5_british = f'/vol/bmd/yanyul/UKB/predicted_expression_tf2/ukb_imp_x_ctimp_Whole_Blood_{population}.hdf5'

# data scheme specifying which are traits and covariates
scheme_yaml = '../misc_files/data_scheme.yaml'

# loading names of traits/covariates
# the order is matched with the data being loaded
feature_dic = util_hdf5.read_yaml(scheme_yaml)
with h5py.File(hdf5_british, 'r') as f:
    features = f['columns_y'][:].astype('str')
    sample_size = f['y'].shape[0]
    y = f['y'][:]
covar_indice = np.where(np.isin(features, feature_dic['covar_names']))[0]
trait_indice = np.where(np.isin(features, feature_dic['outcome_names']))[0]

# load data_scheme for training
batch_size = 2 ** 12
print(f'batch_size in British set is {batch_size}')
data_scheme, sample_size = util_hdf5.build_data_scheme(
    hdf5_british, 
    scheme_yaml, 
    batch_size = batch_size, 
    inv_norm_y = True
)

# set validation and test set as the first and second batch
dataset_valid = data_scheme.dataset.take(1)
data_scheme.dataset = data_scheme.dataset.skip(1)
dataset_test = data_scheme.dataset.take(1)
data_scheme.dataset = data_scheme.dataset.skip(1)
dataset_insample = data_scheme.dataset.take(1)<jupyter_output>batch_size in British set is 4096<jupyter_text>Load other populations.<jupyter_code>test_datasets = {
    'British_validation': dataset_valid,
    'British_test': dataset_test,
    'British_insample': dataset_insample,
    'Chinese': None,
    'Indian': None,
    'African': None
}
batch_size_here = 8096
for i in test_datasets.keys():
    if 'British' not in i:
        filename = f'/vol/bmd/yanyul/UKB/predicted_expression_tf2/ukb_imp_x_ctimp_Whole_Blood_{i}.hdf5'
        data_scheme, sample_size = util_hdf5.build_data_scheme(
            filename, 
            scheme_yaml, 
            batch_size = batch_size_here, 
            inv_norm_y = True
        )
        test_datasets[i] = data_scheme.dataset<jupyter_output><empty_output><jupyter_text>Load models<jupyter_code>alpha_list = [0.1, 0.5, 0.9]
model_list = {}
for alpha in alpha_list:
    filename = f'/vol/bmd/yanyul/UKB/ptrs-tf/models/elastic_net_alpha_{alpha}_British.hdf5'
    model_list[alpha] = lib_LinearAlgebra.ElasticNetEstimator('', None, minimal_load = True)
    model_list[alpha].minimal_load(filename)
model_list[0.1].lambda_seq<jupyter_output><empty_output><jupyter_text>Calculate PTRS<jupyter_code>out = model_list[0.1].predict_x(test_datasets['British_validation'], model_list[0.1].beta_hat_path)
# out = o1
fig, aes = plt.subplots(nrows = 3, ncols = 3, figsize = (15, 10))
seq = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 16, 19]
for i in range(3):
    for j in range(3):
        idx = seq[(i * 3 + j)] + 45
        if idx < len(model_list[alpha].lambda_seq[0]):
            for k in range(1):
                aes[i][j].scatter(out['y'][:,k], out['y_pred_from_x'][:, k, idx])
            aes[i][j].set_title(
                'lambda = ' + "{:.3E} cor = {:.3E}".format(
                    model_list[alpha].lambda_seq[0][idx], 
                    np.corrcoef(out['y'][:, 0], out['y_pred_from_x'][:, 0, idx])[0,1]
                ) # + '\n' +
#                 'lambda = ' + "{:.3E} cor = {:.3E}".format(
#                     model_list[alpha].lambda_seq[1][idx], 
#                     np.corrcoef(out['y'][:, 1], out['y_pred_from_x'][:, 1, idx])[0,1]
#                 )
            )<jupyter_output><empty_output><jupyter_text>Calculate partial R2<jupyter_code>partial_r2 = {}
for alpha in alpha_list:
    partial_r2[alpha] = {}
    model_i = model_list[alpha]
    for i in test_datasets.keys():
        dataset = test_datasets[i]
        for ele in dataset:
            x, y = model_i.data_scheme.get_data_matrix(ele)
            covar = x[:, -len(model_i.data_scheme.covariate_indice) :]
            print('alpha = {}, trait = {}, ncol(covar) = {}'.format(alpha, i, covar.shape[1]))
        out = model_i.predict_x(dataset, model_i.beta_hat_path)
        partial_r2[alpha][i] = util_Stats.quick_partial_r2(covar, out['y'], out['y_pred_from_x'])
def _pr2_format(ele, features, name, alpha, lambda_):
    nlambda = lambda_.shape[1]
    ntrait = lambda_.shape[0]
    ele_seq = np.reshape(ele, (nlambda * ntrait), order = 'C')
    lambda_seq = np.reshape(lambda_, (nlambda * ntrait), order = 'C')
    f_seq = np.repeat(features, nlambda)
    return pd.DataFrame({'partial_r2': ele_seq, 'trait': f_seq, 'sample': name, 'alpha': alpha, 'lambda': lambda_seq})
df = pd.DataFrame({'partial_r2': [], 'trait': [], 'sample': [], 'alpha': [], 'lambda': []})
for alpha in alpha_list:
    model_i = model_list[alpha]
    lambda_i = np.array(model_i.lambda_seq)
    for i in partial_r2[alpha].keys():
        df = pd.concat((df, _pr2_format(partial_r2[alpha][i], features[trait_indice], i, alpha, lambda_i)))
df.to_csv('/vol/bmd/yanyul/UKB/ptrs-tf/models/partial_r2-elastic_net_British.tsv', sep = '\t', index = False)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14442013129102846, 'fraction_numerical': 0.01987600291757841, 'mean_word_length': 3.7862129144851657, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 20, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '51804', 'license': 'MIT', 'n_tokens_mistral': 2084, 'n_tokens_neox': 1978, 'n_words': 408, 'path': 'notebook/results_on_ptrs_elastic_net.ipynb', 'repo_name': 'hakyimlab/ptrs-tf', 'size': '5484'}