<jupyter_start><jupyter_text>__Required files:__*  Part 1    `gs://phenotype_31063/hail/imputed/ukb31063.dosage.autosomes.mt`    `gs://nbaya/split/hapmap3_variants.tsv`    `gs://nbaya/split/hapmap3_variants_v2.tsv`    `gs://phenotype_31063/ukb31063.gwas_covariates.both_sexes.tsv`*  Part 2    `gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v1.mt` **(Deprecated)**    `gs://nbaya/split/ukb31063.hm3_variants.gwas_samples.mt` (Current)*  Part 3    `gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v2.mt` **(Deprecated)**    `gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_repart.mt` (Current)    `gs://nbaya/split/meta_split/ukb31063.hm3_variants.gwas_samples_+phen+desc+.tsv`*  Part 4    `gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_+phen+_grouped_v1.mt`*  Part 5    `gs://nbaya/split/meta_split/ukb31063.hm3_+phen+_gmt.mt`__Final output:__ `gs://nbaya/split/meta_split/ukb31063.hm3_+phen+_meta_split_final_s2.tsv.bgz`__Note:__The latest version of this is now a Python script submitted as a job to the cloud. When using this notebook, make sure to check code with the latest version to ensure that everything matches (most importantly, file names).  Import packages, set phenotype<jupyter_code>import hail as hl
import numpy as np
import pandas as pd

phen_dict = {
    '50':'height',
    '20160':'smoking',
    '2443':'diabetes',
    '23107':'legimp_r',
    '6138_100':'qual_none'
}

phen = '50'
desc = phen_dict[phen]
batch = '1'
n_chunks = 300
variant_set = 'hm3'
constant_sex_ratio = True<jupyter_output><empty_output><jupyter_text>Part 1Create filtered matrix table annotated with locus and alleles (Done)<jupyter_code>mt = hl.read_matrix_table('gs://phenotype_31063/hail/imputed/ukb31063.dosage.autosomes.mt')

mt.count()
if variant_set is 'hm3':
        variants = hl.import_table('gs://nbaya/split/hapmap3_variants.tsv')
elif variant_set is 'qc_pos':
        variants = hl.import_table('gs://ukb31063-mega-gwas/qc/ukb31063.gwas_variants.autosomes.tsv')
variants = variants.annotate(**hl.parse_variant(variants.v))
variants = variants.key_by('locus','alleles')
variants
# variants.write('gs://nbaya/split/hapmap3_variants_v2.tsv')
variants = hl.import_table('gs://nbaya/split/'+variant_set+'_variants.ht') # for hm3: import table hapmap3_variants.tsv'

mt = hl.read_matrix_table('gs://phenotype_31063/hail/imputed/ukb31063.dosage.autosomes.mt')
mt = mt.filter_rows(hl.is_defined(variants[mt.locus,mt.alleles])) #filter to variant_set variants

covs = hl.import_table('gs://phenotype_31063/ukb31063.gwas_covariates.both_sexes.tsv',
                             key='s', impute=True, types={'s': hl.tstr})

mt = mt.annotate_cols(**covs[mt.s])
mt = mt.filter_cols(hl.is_defined(mt.PC1), keep=True)
hl.read_table('gs://nbaya/split/qc_pos_variants.ht/')<jupyter_output><empty_output><jupyter_text>Filter to HapMap 3 variants<jupyter_code>mt = mt.filter_rows(hl.is_defined(variants[mt.locus,mt.alleles])) 
mt.count()
ht_samples = hl.import_table('gs://phenotype_31063/ukb31063.gwas_covariates.both_sexes.tsv',
                             key='s', impute=True, types={'s': hl.tstr})<jupyter_output><empty_output><jupyter_text>Annotate with covariates<jupyter_code>mt = mt.annotate_cols(**ht_samples[mt.s])
mt = mt.filter_cols(hl.is_defined(mt.PC1), keep=True) #Only keeps columns with PC1 defined
# mt.write('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v1.mt', overwrite=True) #Deprecated
mt.write('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_prerepart.mt', overwrite=True)<jupyter_output><empty_output><jupyter_text>Part 2Repartition (Done) Note: Task requires shuffle reading and writing. Do not use preemptible workers.<jupyter_code># mt = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v1.mt') #Deprecated
mt = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_prerepart.mt')


mt = mt.repartition(1000)

# mt.write('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v2.mt', overwrite=True) #Deprecated
mt.write('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_repart.mt', overwrite=True)<jupyter_output><empty_output><jupyter_text>Part 3Annotate with phenotype and label with group IDs<jupyter_code>mt0 = hl.read_matrix_table(f'gs://nbaya/split/ukb31063.{variant_set}_variants.gwas_samples_repart.mt')

if 'sim' in phen:
    print('\nReading simulated phenotype...')
    if variant_set == 'qc_pos':
        mt1 = hl.read_matrix_table('gs://nbaya/rg_sex/qc_pos.50_sim_inf_h2_0.485223.mt') #outdated
    elif variant_set == 'hm3':
#            mt1 = hl.read_matrix_table('gs://nbaya/rg_sex/50_sim_inf_h2_0.485223.mt')      
        sim_phen = phen.split('_sim')[0]
        if sim_phen == '50':
            phen_tb = hl.read_table(f'gs://nbaya/ldscsim/{variant_set}.phen_{sim_phen}.sim_h2_{0.485223}.ht')

else:
    print('\nReading UKB phenotype...')
#        mt0 = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_v2.mt') #old version

    if phen == '50_raw':
        phen_tb0 = hl.import_table('gs://nbaya/ukb31063.50_raw.tsv.bgz',missing='',impute=True,types={'s': hl.tstr}).rename({phen: 'phen'})
    elif phen == '50_raw_res':
        phen_tb0 = hl.read_table('gs://nbaya/split/50_raw_linreg.ht').rename({'res': 'phen'})
    else:
        phen_tb0 = hl.import_table('gs://phenotype_31063/ukb31063.phesant_phenotypes.both_sexes.tsv.bgz',
                                      missing='',impute=True,types={'"userId"': hl.tstr}).rename({ '"userId"': 's', '"'+phen+'"': 'phen'})
        phen_tb0 = phen_tb0.key_by('s')
    phen_tb = phen_tb0.select(phen_tb0['phen'])    
mt1 = mt0.annotate_cols(phen_str = hl.str(phen_tb[mt0.s]['phen']).replace('\"',''))
mt1 = mt1.filter_cols(mt1.phen_str == '',keep=False)

if phen_tb.phen.dtype == hl.dtype('bool'):
    mt1 = mt1.annotate_cols(phen = hl.bool(mt1.phen_str)).drop('phen_str')
else:
    mt1 = mt1.annotate_cols(phen = hl.float64(mt1.phen_str)).drop('phen_str')    
#Remove withdrawn samples
withdrawn = hl.import_table('gs://nbaya/w31063_20181016.csv',missing='',no_header=True)
withdrawn_set = set(withdrawn.f0.take(withdrawn.count()))
mt1 = mt1.filter_cols(hl.literal(withdrawn_set).contains(mt1['s']),keep=False) 
mt1 = mt1.key_cols_by('s')

n_samples = mt1.count_cols()
print(f'\n>>> N samples = {n_samples} <<<') #expect n samples to match n_non_missing from phenotypes.both_sexes.tsv, minus withdrawn samples.<jupyter_output>2019-04-11 11:54:46 Hail: INFO: Reading table with no type imputation
  Loading column 'f0' as type 'str' (type not specified)<jupyter_text>Approach 1 (does not work)<jupyter_code>mt2 = mt1.add_col_index()
n_per_sex = np.asarray([mt2.filter_cols(mt2.isFemale==i).count_cols() for i in range(2)])
print(n_per_sex)
group_size = (n_per_sex/n_chunks).astype(int)+1
randstate = np.random.RandomState(int(batch)) #seed with batch number
group_ids = []
for i in range(2):
    temp=np.ndarray.tolist(np.ndarray.flatten(np.asarray([range(n_chunks)]*group_size[i]))) 
    temp=temp[0:n_per_sex[i]]
    randstate.shuffle(temp)
    group_ids.append(temp)

cols = mt2.cols().key_by('s')
cols = cols.order_by(cols.isFemale)
group_ids_all = [i for x in group_ids for i in x] 
cols = cols.add_index().key_by('idx')
cols = cols.annotate(group_id = hl.literal(group_ids_all)[hl.int32(cols.idx)])
group_ids_all
cols.show()
cols = cols.key_by('s')
mt3 = mt2.annotate_cols(group_id = cols[mt2.s].group_id)
mt3.aggregate_cols(hl.agg.group_by(mt3.group_id,hl.agg.mean(mt3.isFemale)))<jupyter_output>2019-04-08 17:20:35 Hail: INFO: Coerced sorted dataset
2019-04-08 17:20:39 Hail: INFO: Ordering unsorted dataset with network shuffle
2019-04-08 17:20:54 Hail: INFO: Coerced sorted dataset
2019-04-08 17:22:45 Hail: INFO: Coerced sorted dataset<jupyter_text>Approach 2<jupyter_code>mt_ls = [mt1.filter_cols(mt1.isFemale==0), mt1.filter_cols(mt1.isFemale==1)]
mt_final = []
for mt_temp in mt_ls:
    n_samples = mt_temp.count_cols()
    print('\n>>> N samples = '+str(n_samples)+' <<<') #expect n samples to match n_non_missing from phenotypes.both_sexes.tsv, minus withdrawn samples.
    
    mt_temp2 = mt_temp.add_col_index()
    group_size = int(n_samples/n_chunks)+1     #the ideal number of samples in each group
    #list of group ids to be paired to each sample (Note: length of group_ids > # of cols in mt, but it doesn't affect the result)
    group_ids = np.ndarray.tolist(np.ndarray.flatten(np.asarray([range(n_chunks)]*group_size))) 
    group_ids = group_ids[0:n_samples]
    randstate = np.random.RandomState(int(batch)) #seed with batch number
    randstate.shuffle(group_ids)
    mt_final.append(mt_temp2.annotate_cols(group_id = hl.literal(group_ids)[hl.int32(mt_temp2.col_idx)])) #assign group ids # OLD VERSION

mt3 = mt_final[0].union_cols(mt_final[1])
mt3.aggregate_cols(hl.agg.group_by(mt3.group_id,hl.agg.mean(mt3.isFemale)))
group_size = int(n_samples/n_chunks)+1     #the ideal number of samples in each group
#list of group ids to be paired to each sample (Note: length of group_ids > # of cols in mt, but it doesn't affect the result)
group_ids = np.ndarray.tolist(np.ndarray.flatten(np.asarray([range(n_chunks)]*group_size))) 
group_ids = group_ids[0:n_samples]
randstate = np.random.RandomState(int(batch)) #seed with batch number
randstate.shuffle(group_ids)
mt3 = mt2.annotate_cols(group_id = hl.literal(group_ids)[hl.int32(mt2.col_idx)]) #assign group ids # OLD VERSION
ht_group_ids = mt3.select_cols(mt3.group_id).cols() #assign group ids
# Takes about 20 min with 100 workers, 0 preemptibles
# Requires shuffle reading and writing
mt3.write(f'gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_{phen}_grouped_v1_new.mt', overwrite=True)<jupyter_output>2018-10-02 16:24:56 Hail: INFO: Coerced sorted dataset
2018-10-02 17:16:30 Hail: INFO: wrote 1089172 items in 1000 partitions to gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_50_grouped_v1_new.mt<jupyter_text>Part 4Run linear regression (Done)<jupyter_code>mt = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_'+phen+'_grouped_batch_'+batch+'.mt')

mt = mt.rename({'dosage': 'x', 'phen': 'y'})

cov_list = [ mt['isFemale'], mt['age'], mt['age_squared'], mt['age_isFemale'],
            mt['age_squared_isFemale'] ]+ [mt['PC{:}'.format(i)] for i in range(1, 21)] 
gmt = (mt.group_cols_by(mt.group_id)
        .aggregate(linreg = hl.agg.linreg(y=mt.y, x = [mt.x, 1] + cov_list)))
# Takes about 1.7 hours (with 10 workers + 200 pre-emptible workers)
# First and only stage should have 1000 tasks
gmt.select_entries(beta=gmt.linreg.beta[0],standard_error=gmt.linreg.standard_error[0]).write('gs://nbaya/split/meta_split/ukb31063.hm3_'+phen+'_gmt_test.mt')<jupyter_output><empty_output><jupyter_text>Check linear regression resultsUsed only when constant_sex_ratio = True<jupyter_code>gmt = hl.read_matrix_table(f'gs://nbaya/split/meta_split/ukb31063.{variant_set}_{phen}_gmt{n_chunks}_constantsexratio_{constant_sex_ratio}_batch_{batch}.mt') #used for n=150 and other phenotypes
gmt = gmt.add_row_index()
gmt = gmt.key_rows_by('row_idx')
gmt = gmt.drop('locus','alleles','varid')
gmt = gmt.rename({'rsid': 'SNP'})
gmt.beta.show()<jupyter_output>+---------------+------------+----------+-----------+
| locus         | alleles    | group_id |      beta |
+---------------+------------+----------+-----------+
| locus<GRCh37> | array<str> |    int32 |   float64 |
+---------------+------------+----------+-----------+
| 1:754182      | ["A","G"]  |        0 |  1.82e-02 |
| 1:754182      | ["A","G"]  |        1 |  5.75e-02 |
| 1:754182      | ["A","G"]  |        2 |  7.73e-02 |
| 1:754182      | ["A","G"]  |        3 | -4.39e-02 |
| 1:754182      | ["A","G"]  |        4 |  6.18e-02 |
| 1:754182      | ["A","G"]  |        5 | -2.48e-02 |
| 1:754182      | ["A","G"]  |        6 | -3.53e-02 |
| 1:754182      | ["A","G"]  |        7 | -3.22e-02 |
| 1:754182      | ["A","G"]  |        8 | -9.15e-02 |
| 1:754182      | ["A","G"]  |        9 |  4.98e-02 |
+---------------+------------+----------+-----------+
showing top 10 rows<jupyter_text>Part 5Split 300 subgroups into two groups<jupyter_code>gmt = hl.read_matrix_table('gs://nbaya/split/meta_split/ukb31063.hm3_'+phen+'_gmt.mt')

gmt = gmt.add_row_index()
gmt = gmt.key_rows_by('row_idx')
gmt = gmt.drop('locus','alleles','varid')
gmt = gmt.rename({'rsid': 'SNP'})
n = 300 #number of subgroups
pi = ['A']*int(n/2) + ['B']*int(n/2)
np.random.shuffle(pi)
gmt = gmt.add_col_index()
gmt = gmt.annotate_cols(label = hl.literal(pi)[hl.int32(gmt.col_idx)])

# print(gmt.aggregate_cols(hl.agg.counter(gmt.label))) #determine the evenness of the split
mt = gmt.group_cols_by(gmt.label).aggregate(unnorm_meta_beta=hl.agg.sum(gmt.beta / gmt.standard_error ** 2), inv_se2 = hl.agg.sum(1 / gmt.standard_error ** 2))

#Make table of just the rsid column (already complete)
# rsid_mt = mt.drop( 'unnorm_meta_beta', 'inv_se2')
# rsid_tb = rsids_mt.make_table()
# rsid_tb.show()
# rsid_tb.export('gs://nbaya/split/meta_split/meta_split_rsid.tsv.bgz')

ht = mt.make_table()
# ht = ht.annotate(A_meta_beta = ht['A.unnorm_meta_beta'] / ht['A.inv_se2'],
#                    A_meta_se = hl.sqrt(1 / ht['A.inv_se2']), 
#                    B_meta_beta = ht['B.unnorm_meta_beta'] / ht['B.inv_se2'], 
#                    B_meta_se = hl.sqrt(1 / ht['B.inv_se2']))
ht = ht.annotate(A_Z = ht['A.unnorm_meta_beta'] / hl.sqrt(ht['A.inv_se2']),
                     B_Z = ht['B.unnorm_meta_beta'] / hl.sqrt(ht['B.inv_se2']))
ht = ht.drop('A.unnorm_meta_beta','B.unnorm_meta_beta','A.inv_se2','B.inv_se2').key_by('SNP')
ht.show()
variants = hl.import_table('gs://nbaya/rg_sex/'+phen+'_snps_alleles_N.tsv.gz',types={'N': hl.tint64})
variants = variants.key_by('SNP')
mt_all = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_'+phen+'_grouped_batch_'+batch+'.mt') #matrix table containing individual samples
N = mt_all.count_cols()
variants = variants.annotate(N = hl.int32(N/2))
variants.show()
print(N)
metaA = variants.annotate(Z = ht[variants.SNP].A_Z)
metaB = variants.annotate(Z = ht[variants.SNP].B_Z)
metaB.show()
metaA.aggregate(hl.agg.mean(metaA.Z))
metaB.aggregate(hl.agg.mean(metaB.Z))
metaA.export('gs://nbaya/rg_sex/'+phen+'_meta_A_s1.tsv.bgz')
metaB.export('gs://nbaya/rg_sex/'+phen+'_meta_B_s1.tsv.bgz')
ht.export('gs://nbaya/split/meta_split/ukb31063.hm3_'+phen+'_metasplit_final_s3.tsv.bgz')<jupyter_output><empty_output><jupyter_text>Check completed sumstats<jupyter_code>r_ls = []
for i in range(100):
#     metaA_path = f'gs://nbaya/rg_sex/{phen}_meta_A_n{n_chunks}_batch_{batch}_s{i}.tsv.bgz' 
#     metaB_path = f'gs://nbaya/rg_sex/{phen}_meta_B_n{n_chunks}_batch_{batch}_s{i}.tsv.bgz' 
    metaA_path = f'gs://nbaya/rg_sex/{variant_set}_{phen}_meta_A_n{n_chunks}_constantsexratio_{constant_sex_ratio}_batch_{batch}_s{i}.tsv.bgz'  #first used by constant_sex_ratio=True runs
    metaB_path = f'gs://nbaya/rg_sex/{variant_set}_{phen}_meta_B_n{n_chunks}_constantsexratio_{constant_sex_ratio}_batch_{batch}_s{i}.tsv.bgz'  #first used by constant_sex_ratio=True runs
    metaA = hl.import_table(metaA_path,impute=True,key='SNP')
    metaB = hl.import_table(metaB_path,impute=True,key='SNP')
    both = metaA.rename({'Z':'A_Z'})
    both = both.annotate(B_Z = metaB[both.SNP].Z)
    df = hl.Table.to_pandas(both)
    r = stats.pearsonr(df['A_Z'],df['B_Z'])[0]
    print(r)
    r_ls.append(r)
# old_r_ls = [x[0] for x in r_ls]
print(np.mean(old_r_ls))
print(np.std(old_r_ls))
new_r_ls = r_ls
print(np.mean(new_r_ls))
print(np.std(new_r_ls))<jupyter_output>0.6463037603727204
0.0017053361058531735
{'fraction_non_alphanumeric': 0.17461439588688946, 'fraction_numerical': 0.05, 'mean_word_length': 4.609589041095891, 'pattern_counts': {'":': 0, '<': 52, '<?xml version=': 0, '>': 54, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '22913', 'license': 'MIT', 'n_tokens_mistral': 6854, 'n_tokens_neox': 6267, 'n_words': 897, 'path': 'python/meta_split.ipynb', 'repo_name': 'nikbaya/split', 'size': '15560'}