<jupyter_start><jupyter_text>Comments go hereHere we take two random subsets of individuals and run two multi-trait GWASs.The idea is to test the usage of `linear_regression_rows` in hail 0.2 in which `y` can be list of lists.  Before running the thingWe take a detour to create individual subsetting list.<jupyter_code>nsubset = 2
if not os.path.isfile('output/training_0.txt'): 
    def draw_subset(ntotal, ntarget, seed):
        if ntotal < ntarget:
            return None
        else:
            np.random.seed(seed)
            return np.random.choice(ntotal, ntarget, replace = False)
    def save_as_table(vec, f):
        o = open(f, 'w')
        for i in vec:
            o.write(i + '\n')
        o.close()
    indiv_list = np.array(mt.s.collect())
    subset_size = 200
    for i in range(nsubset):
        subset_idx = draw_subset(indiv_list.shape[0], subset_size, seed = i)
        training = indiv_list[subset_idx]
        save_as_table(training, f'output/training_{i}.txt')
        test = indiv_list[np.where(np.isin(np.array(range(indiv_list.shape[0])), subset_idx, invert = True))]
        save_as_table(test, f'output/test_{i}.txt')<jupyter_output><empty_output><jupyter_text>OK, we can start to work on the GWAS pipeline<jupyter_code># loading phenotypes
pheno_full_table = pd.read_csv('output/indiv_pheno.tsv', sep = '\t')
target_trait_list = [ f'trait_{i}' for i in range(12) ]
my_pheno_list = []
for i in range(nsubset):
    pheno_set1 = pd.read_csv(f'output/training_{i}.txt', sep = '\t', header = None)
    pheno_subset = pheno_full_table.loc[pheno_full_table['s'].isin(pheno_set1[0].to_list())]
    pheno_subset = pheno_subset[['s'] + target_trait_list]
    pheno_subset_ht = hl.Table.from_pandas(pheno_subset, key = 's')
    my_pheno_list.append(pheno_subset_ht)
# loading covariates
type_dic = { f'covar_{i}': hl.tfloat for i in range(5) }
covar = hl.import_table('output/indiv_covar.tsv', types = type_dic, key = 's')
# looks good, let's run GWAS
# annotate phenotypes
anno_pheno_expr = {
    f'subset_{i}': my_pheno_list[i][mt.s] for i in range(len(my_pheno_list))
}
mt = mt.annotate_cols(**anno_pheno_expr)
# annotate covariates
mt = mt.annotate_cols(covariate = covar[mt.s])
# and run linear_regression_rows
pheno_list = [ [ mt[f'subset_{i}'][j] for j in mt[f'subset_{i}'] ] for i in range(nsubset) ]
pheno_names = [ [ f'subset_{i}_x_{j}' for j in mt[f'subset_{i}'] ] for i in range(nsubset) ]
covar_list = [ mt.covariate[i] for i in list(mt.covariate.keys()) ]
gwas_out = hl.linear_regression_rows(
    y = pheno_list,
    x = mt.GT.n_alt_alleles(),
    covariates = covar_list + [1]
)
gwas_out = gwas_out.annotate_globals(phenotypes = pheno_names)
gwas_out.describe()<jupyter_output>----------------------------------------
Global fields:
    'phenotypes': array<array<str>> 
----------------------------------------
Row fields:
    'locus': locus<GRCh37> 
    'alleles': array<str> 
    'n': array<int32> 
    'sum_x': array<float64> 
    'y_transpose_x': array<array<float64>> 
    'beta': array<array<float64>> 
    'standard_error': array<array<float64>> 
    't_stat': array<array<float64>> 
    'p_value': array<array<float64>> 
----------------------------------------
Key: ['locus', 'alleles']
----------------------------------------<jupyter_text>OK, the last step is to format results into tables by traitHere I mostly follow routine in [here](https://github.com/Nealelab/UK_Biobank_GWAS/blob/95ac260a5d4cf9c40effff13fa33fb95ed825e2a/0.2/run_regressions.biomarkers.pyL54) and [here](https://github.com/Nealelab/UK_Biobank_GWAS/blob/95ac260a5d4cf9c40effff13fa33fb95ed825e2a/0.2/export_results.biomarkers.py)<jupyter_code>gwas_out = gwas_out.annotate( 
    variant = hl.delimit(
        hl.array([
            gwas_out['locus'].contig,
            hl.str(gwas_out['locus'].position),
            gwas_out['alleles'][0],
            gwas_out['alleles'][1]
        ]), 
    delimiter = ':')
)
gwas_out = gwas_out.key_by('variant')
## Hey, this repartition is important
## in the sense that it avoids the unnecessary and repeated sorting caused by key_by
gwas_out = gwas_out.repartition(8)
gwas_out = gwas_out.cache()

phenotypes = gwas_out['phenotypes'].collect()[0]
for i, subset in enumerate(phenotypes):
    for j, trait in enumerate(subset):
#         trait_name = trait.split('_x_')[-1]
        ht_export = gwas_out.annotate(
            n_complete_samples = gwas_out['n'][i],
            AC = gwas_out['sum_x'][i],
            ytx = gwas_out['y_transpose_x'][i][j],
            beta = gwas_out['beta'][i][j],
            se = gwas_out['standard_error'][i][j],
            tstat = gwas_out['t_stat'][i][j],
            pval = gwas_out['p_value'][i][j])
        ht_export = ht_export.annotate(
            AF = ht_export['AC'] / (2 * ht_export['n_complete_samples']))
        ht_export = ht_export.annotate(
            minor_AF = hl.cond(ht_export['AF'] <= 0.5, ht_export['AF'], 1.0 - ht_export['AF']),
            minor_allele = hl.cond(ht_export['AF'] <= 0.5, ht_export['alleles'][1], ht_export['alleles'][0]))
        ht_export = ht_export.annotate(
            low_confidence_variant=ht_export['minor_AF'] < 0.001)
        ht_export = ht_export.select(
            'minor_allele',
            'minor_AF',
            'low_confidence_variant',
            'n_complete_samples',
            'AC',
            'ytx',
            'beta',
            'se',
            'tstat',
            'pval')
        ht_export.export(f'output/gwas_{trait}.tsv')
# # gwas_out.annotate_globals()
# phenotypes = gwas_out['phenotypes'].collect()[0]
# phenotypes
# gwas_out['n'][i].show()
# gwas_out['y_transpose_x'][i][j].show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16836555360281194, 'fraction_numerical': 0.017047451669595783, 'mean_word_length': 3.11794500723589, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 28, 'https://': 2, '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': '586653', 'license': 'MIT', 'n_tokens_mistral': 2180, 'n_tokens_neox': 2078, 'n_words': 376, 'path': 'notebook/run_multi_trait_gwas.ipynb', 'repo_name': 'liangyy/toy-example-of-hail-using-1kg', 'size': '5690'}