<jupyter_start><jupyter_text>Analysis overviewBuilding PTRS.1. Split British data into 3 sets: training, test, validation.2. Train least squares predictor using British training data.3. Compute PTRS in all populations.  Load data<jupyter_code># set path to British data
hdf5_british = '/vol/bmd/yanyul/UKB/predicted_expression_tf2/ukb_imp_x_ctimp_Whole_Blood_British.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')
    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]
# import scipy.stats
# def inv_norm_col(mat):
#     return np.apply_along_axis(inv_norm_vec, 0, mat)
# def inv_norm_vec(vec, offset = 1):
#     rank = myrank(vec)
#     return scipy.stats.norm.ppf(rank / (len(rank) + offset), loc = 0, scale = 1)
# def myrank(vec):
#     argsort = np.argsort(vec)
#     ranks = np.empty_like(argsort)
#     ranks[argsort] = np.arange(len(vec))
#     return ranks + 1  # rank starts from 1
# y_norm = inv_norm_col(y)
# scipy.stats.describe(y_norm[:,2])
features
# 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)
# data_scheme.dataset = data_scheme.dataset.take(10)
data_scheme2, sample_size2 = util_hdf5.build_data_scheme(
    hdf5_british, 
    scheme_yaml, 
    batch_size = batch_size, 
    inv_norm_y = True
)

for u in data_scheme2.dataset.unbatch().batch(20):
#     print(u[0].shape)
    x, y = data_scheme2.get_data_matrix(u)
    break

x[:,-1]<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>least_square_solver = lib_LinearAlgebra.LeastSquaredEstimator(data_scheme, intercept = True, normalizer = True)
least_square_solver.solve(logging = logging, sample_size = sample_size - batch_size * 2)<jupyter_output>2020-02-02 10:52:17 AM  Progress 5%: 20480 / 348284
2020-02-02 10:52:18 AM  Progress 10%: 36864 / 348284
2020-02-02 10:52:18 AM  Progress 15%: 53248 / 348284
2020-02-02 10:52:19 AM  Progress 20%: 73728 / 348284
2020-02-02 10:52:20 AM  Progress 25%: 90112 / 348284
2020-02-02 10:52:21 AM  Progress 30%: 106496 / 348284
2020-02-02 10:52:22 AM  Progress 35%: 122880 / 348284
2020-02-02 10:52:22 AM  Progress 40%: 143360 / 348284
2020-02-02 10:52:23 AM  Progress 45%: 159744 / 348284
2020-02-02 10:52:24 AM  Progress 50%: 176128 / 348284
2020-02-02 10:52:25 AM  Progress 55%: 192512 / 348284
2020-02-02 10:52:26 AM  Progress 60%: 212992 / 348284
2020-02-02 10:52:27 AM  Progress 65%: 229376 / 348284
2020-02-02 10:52:27 AM  Progress 70%: 245760 / 348284
2020-02-02 10:52:28 AM  Progress 75%: 262144 / 348284
2020-02-02 10:52:29 AM  Progress 80%: 282624 / 348284
2020-02-02 10:52:30 AM  Progress 85%: 299008 / 348284
2020-02-02 10:52:31 AM  Progress 90%: 315392 / 348284
2020-02-02 10:52:32 AM  Progress 9[...]<jupyter_text>Testing  Loading other datasets<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>Calculate PTRSThis section is mainly for debugging.<jupyter_code># dataset = test_datasets['British_validation']
# dataset_in = test_datasets['British_insample']
# ptrs_british_valid = least_square_solver.predict_x(dataset)
# ptrs_british_in = least_square_solver.predict_x(dataset_in)
# least_square_solver.betahat
# fig, axe = plt.subplots(ncols = 3, nrows = 6, figsize = (10,20))
# for i in range(3):
#     for j in range(6):
#         k = j * 3 + i 
#         if k < 17:
#             axe[j][i].scatter(ptrs_british_valid['y'][:, k], ptrs_british_valid['y_pred_from_x'][:, k])
#             axe[j][i].title.set_text(features[trait_indice][k])
#             axe[j][i].scatter(ptrs_british_in['y'][:, k], ptrs_british_in['y_pred_from_x'][:, k])
#             axe[j][i].title.set_text(features[trait_indice][k])<jupyter_output><empty_output><jupyter_text>Calculate partial R2<jupyter_code>partial_r2 = {}
for i in test_datasets.keys():
    dataset = test_datasets[i]
    for ele in dataset:
        x, y = least_square_solver.data_scheme.get_data_matrix(ele)
        covar = x[:, -len(least_square_solver.data_scheme.covariate_indice) :]
        print('trait = {}, ncol(covar) = {}'.format(i, covar.shape[1]))
    out = least_square_solver.predict_x(dataset)
    partial_r2[i] = util_Stats.quick_partial_r2(covar, out['y'], np.expand_dims(out['y_pred_from_x'], axis = 2))
partial_r2['Chinese'].shape
def _pr2_format(ele, features, name):
    return pd.DataFrame({'partial_r2': ele[:, 0], 'trait': features, 'sample': name})
df = pd.DataFrame({'partial_r2': [], 'trait': [], 'sample': []})
for i in partial_r2.keys():
    df = pd.concat((df, _pr2_format(partial_r2[i], features[trait_indice], i)))
df<jupyter_output><empty_output><jupyter_text>Save partial R2<jupyter_code>df
df.to_csv('/vol/bmd/yanyul/UKB/ptrs-tf/models/partial_r2-least_square_model-ukb_imp_x_ctimp_Whole_Blood_British.tsv', sep = '\t', index = False)
plt.figure(figsize=(15,8))
sns.barplot(x = "trait", y = "partial_r2", hue = 'sample', data = df)<jupyter_output><empty_output><jupyter_text>Load REML result<jupyter_code>filename = '/vol/bmd/yanyul/UKB/gcta_regulability/reml_from_hail_martin_et_al_traits_x_ctimp_Whole_Blood_x_{pop}.tsv'
df_reml = pd.DataFrame({})
for pop in ['African', 'Chinese', 'Indian', 'British-test-1']:
    df_reml_i = pd.read_csv(filename.format(pop = pop), sep = '\t')
    if 'British' in pop:
        pop = 'British_test'
    df_reml_i['sample'] = pop
    df_reml = pd.concat((df_reml, df_reml_i))
plt.figure(figsize=(15,8))
sns.barplot(x = "trait", y = "h_sq", hue = 'sample', data = df_reml)<jupyter_output><empty_output><jupyter_text>REML vs Partial R2<jupyter_code>tmp = pd.read_table('/vol/bmd/yanyul/GitHub/ptrs-ukb/external_data/martin_et_al_2019ng_table_s6_trait_description.tsv', sep = '\t')
blood_traits = ['wbc', 'rbc', 'platelet', 'lymphocyte', 'monocyte', 'neutrophil', 'eosinophil', 'basophil'] 
df_join = df_reml.join(df.set_index(['trait', 'sample']), on = ['trait', 'sample'])
df_join['is_blood_cell_count'] = np.isin(df_join['trait'], blood_traits)
g = sns.FacetGrid(df_join, col = "is_blood_cell_count", hue = 'sample')  # 
g = g.map(plt.scatter, "h_sq", "partial_r2")
g.add_legend()
# plt.figure(figsize=(10,10))
# sns.scatterplot(x = "h_sq", y = "partial_r2", hue = 'trait', style = 'sample', data = df_join)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14184490224380003, 'fraction_numerical': 0.07991077286445349, 'mean_word_length': 4.813882532418002, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 32, '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': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '615098', 'license': 'MIT', 'n_tokens_mistral': 3440, 'n_tokens_neox': 2950, 'n_words': 563, 'path': 'notebook/first_look_at_ptrs_least_square.ipynb', 'repo_name': 'hakyimlab/ptrs-tf', 'size': '7621'}