sim_chr1.py

import sys
import os
import random
import copy
import scipy.stats
import shutil
from helpers import *

def main():

    orig_prefix = 'subset_302_homozygotes'
    new_prefix = orig_prefix + '_' + sys.argv[1]

    shuffle = False
    if len(sys.argv) > 3 and sys.argv[3] == 'shuffle':
        new_prefix += '_shuffled'
        shuffle = True

    print('choosing causal snp')
    # choose random snp from all snps that meet filter criteria (this is
    # more than just the tag snps)
    f_filtered = open(orig_prefix + '_filtered_inds_chr1.txt', 'r')
    snp_inds = [int(i) for i in f_filtered.readline().split()]
    f_filtered.close()
    random_snp_ind = random.choice(snp_inds) # index in original ped/map
    f_causal = open('output/' + new_prefix + '.causal.txt', 'w')
    f_causal.write(str(random_snp_ind) + '\n')
    f_causal.close()

    print('reading in genotypes and assigning phenotypes')
    # assign phenotypes based on chosen snp (alleles aren't converted to
    # 0s and 1s, and rows are strains)
    ped_lines = [line.split() for line in open(orig_prefix + '.ped', 'r').readlines()]
    num_strains = int(len(ped_lines))
    num_snps = int((len(ped_lines) - 6) / 2)
    # these genotypes are from the unfiltered ped file, so the indices in
    # snp_inds will correctly correspond to these rows
    alleles = []
    for strain in range(num_strains):
        alleles.append(ped_lines[strain][6::2][random_snp_ind])
    alleles_set = list(set(alleles))
    major = alleles_set[0]
    minor = alleles_set[1]
    if alleles.count(minor) > alleles.count(major):
        major = alleles_set[1]
        minor = alleles_set[0]
    maf = float(alleles.count(minor)) / len(alleles)
    phens = []

    h2 = float(sys.argv[2])
    fixed_effect = pow(float(h2 * (len(alleles)-1)) / ((h2-1) * len(alleles) * (maf-1) * maf), .5)

    for allele in alleles:
        mu = 0
        if allele == minor:
            mu = fixed_effect
        phens.append(random.gauss(mu, 1))

    # shuffle?
    if shuffle:
        random.shuffle(phens)

    # store the phenotypes so we can look at them later
    f_phen = open('output/' + new_prefix + '.phen.txt', 'w')
    for p in phens:
        f_phen.write(str(p) + '\n')
    f_phen.close()
    

    # create fam files
    phenotypes_to_fam(orig_prefix, new_prefix, phens)

    # association tests
    print('running association tests') 

    #####
    # basic whole-genome kinship matrix
    #####
    print('- whole genome K')

    os.symlink(orig_prefix + '_tag_chr1.bed', new_prefix + '.bed')
    os.symlink(orig_prefix + '_tag_chr1.bim', new_prefix + '.bim')
    gemma(new_prefix, 'output/' + orig_prefix, 'output/' + orig_prefix, '_whole_K')
    os.remove(new_prefix + '.bed')
    os.remove(new_prefix + '.bim')
    os.remove('output/' + new_prefix + '_whole_K.log.txt')


    #####
    # local kinship matrix by various window sizes
    #####
    print('- local K by windows')

    rsids = [int(line.split()[1][2:]) for line in open(orig_prefix + '_tag_chr1.map', 'r').readlines()]
    rsids.sort()
    tag_bim_lines = open(orig_prefix + '_tag_chr1.bim', 'r').readlines()

    windows = ['10', '100', '1000', '5000', '10000', '20000']
    for window in windows:
        f_assoc_local_window = open('output/' + new_prefix + '_local_window_' + window + '.assoc.txt', 'w')

        for i in range(len(rsids)):        
            rsid = 'rs' + str(rsids[i])
            orig_prefix_window = orig_prefix + '_' + rsid
            new_prefix_window =  orig_prefix_window + '_window_' + sys.argv[1]
            if shuffle:
                new_prefix_window += '_shuffled'
            os.symlink(new_prefix + '.fam', new_prefix_window + '.fam')
            os.symlink('bed/' + orig_prefix_window + '.bed', new_prefix_window + '.bed')
            f_bim = open(new_prefix_window + '.bim', 'w')
            f_bim.write(tag_bim_lines[i])
            f_bim.close()
        
            gemma(new_prefix_window, 'output/' + orig_prefix + '_' + window + '_eigenD_chr1/' + orig_prefix_window, 'output/' + orig_prefix + '_' + window + '_eigenU_chr1/' + orig_prefix_window)

            f_out_local = open('output/' + new_prefix_window + '.assoc.txt', 'r')
            f_out_local.readline()
            line = f_out_local.readline()
            while line != '':
                f_assoc_local_window.write(line)
                line = f_out_local.readline()
            f_out_local.close()

            os.remove('output/' + new_prefix_window + '.assoc.txt')
            os.remove('output/' + new_prefix_window + '.log.txt')
            os.remove(new_prefix_window + '.fam')
            os.remove(new_prefix_window + '.bim')
            os.remove(new_prefix_window + '.bed')
        
        f_assoc_local_window.close()

    #####
    # t-test
    #####
    print('- t-test')

    tag_inds = [int(x) for x in open(orig_prefix + '_tag_inds_chr1.txt', 'r').readline().split()]
    tag_inds.sort()
    num_markers = len(tag_inds)

    f_assoc_t = open('output/' + new_prefix + '_t.assoc.txt', 'w')
    for i in tag_inds:
        alleles = []
        for strain in range(num_strains):
            alleles.append(ped_lines[strain][6 + 2 * i])
        a1 = list(set(alleles))[0]
        phens_1 = []
        phens_2 = []
        for phen_ind in range(len(phens)):
            if alleles[phen_ind] == a1:
                phens_1.append(phens[phen_ind])
            else:
                phens_2.append(phens[phen_ind])

        p = scipy.stats.ttest_ind(phens_1, phens_2)[1]
        f_assoc_t.write(str(p) + '\n')
    f_assoc_t.close()

    
    os.remove(new_prefix + '.fam')

main()