tractor-response/simulation/Snakefile at main · bogdanlab/tractor-response · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from os.path import join
import os
import yaml
import experiment
import admix
import itertools
import numpy as np
import pandas as pd
config = {
    "N_TEST_CHUNK": 100
}

RAW_DATA_DIR = "raw_geno"

GENO_PREFIX_LIST = [
    "EUR_0.2_AFR_0.8_7_80000"
]

PARAM_LIST = [param for param in itertools.product([1.0, 1.15, 1.2, 1.25], [0.0, 0.5])]

rule all:
    input:
        expand("out/single_snp_test_tractor/{geno_prefix}/{case_control_param}/{sim_param}/chunk_{chunk_i}.csv.gz",
            geno_prefix=GENO_PREFIX_LIST,
            case_control_param=["0.1_1.0", "0.1_2.5"],
            sim_param=['_'.join([str(i) for i in p]) for p in PARAM_LIST],
            chunk_i=np.arange(config["N_TEST_CHUNK"])),
        expand("out/single_snp_test_tractor/{geno_prefix}/{case_control_param}/{sim_param}/summary.csv.gz",
            geno_prefix=GENO_PREFIX_LIST,
            case_control_param=["0.1_1.0", "0.1_2.5"],
            sim_param=['_'.join([str(i) for i in p]) for p in PARAM_LIST])


rule format_data:
    resources:
        mem_gb=8,
        time_min=20
    input:
        legend = join(RAW_DATA_DIR, "legend.txt"),
        anc = join(RAW_DATA_DIR, "{sim_prefix}/admix.hanc"),
        phgeno = join(RAW_DATA_DIR, "{sim_prefix}/admix.phgeno"),
    output:
        anc = "data/geno/{sim_prefix}/anc.npy",
        phgeno = "data/geno/{sim_prefix}/phgeno.npy",
        legend = "data/geno/{sim_prefix}/legend.csv"
    run:
        import numpy as np
        import pandas as pd

        legend = pd.read_csv(input.legend, delim_whitespace=True)
        ancestry = admix.read_int_mat(input.anc)
        phgeno = admix.read_int_mat(input.phgeno).T

        np.save(output.anc, ancestry)
        np.save(output.phgeno, phgeno)
        legend.to_csv(output.legend, index=False)


rule single_snp_test_tractor:
    resources:
        mem_gb=5,
        time_min=180
    input:
        anc = "data/geno/{geno_prefix}/anc.npy",
        phgeno = "data/geno/{geno_prefix}/phgeno.npy",
        legend = "data/geno/{geno_prefix}/legend.csv"
    output:
        "out/single_snp_test_tractor/{geno_prefix}/{case_prevalence}_{control_ratio}/{odds_ratio}_{anc_effect}/chunk_{chunk_i}.csv.gz"
    run:
        import numpy as np
        # read phgeno, anc
        anc = np.load(input.anc)
        phgeno = np.load(input.phgeno)
        # calculating global ancestry
        global_ancestry = anc.reshape((anc.shape[0] // 2, anc.shape[1] * 2)).mean(axis=1)
        print(global_ancestry)
        print(np.mean(global_ancestry), np.std(global_ancestry))
        assert np.all(anc.shape == phgeno.shape)

        n_snp = anc.shape[1]
        chunk_index = np.array_split(np.arange(n_snp), config["N_TEST_CHUNK"])[int(wildcards.chunk_i)]
        print("chunk_index", chunk_index)
        phgeno = phgeno[:, chunk_index]
        anc = anc[:, chunk_index]
        print(f"case_prevalence: {wildcards.case_prevalence}, control_ratio: {wildcards.control_ratio}")
        print(f"odds_ratio: {wildcards.odds_ratio}, anc_effect: {wildcards.anc_effect}")
        rls = experiment.single_snp_test_tractor(phgeno=phgeno,
                                                 anc=anc,
                                                 theta=global_ancestry,
                                                 odds_ratio=float(wildcards.odds_ratio),
                                                 anc_effect=float(wildcards.anc_effect),
                                                 case_prevalence=float(wildcards.case_prevalence),
                                                 control_ratio=float(wildcards.control_ratio),
                                                 seed=int(wildcards.chunk_i), n_sim=50)

        # save
        rls.to_csv(output[0], index=False)

rule single_snp_test_summary:
    resources:
        mem_gb=8,
        time_min=10
    input:
        expand("out/single_snp_test_tractor/{{prefix}}/chunk_{chunk_i}.csv.gz",
                chunk_i=np.arange(config["N_TEST_CHUNK"])),
    output:
        "out/single_snp_test_tractor/{prefix}/summary.csv.gz"
    run:
        import numpy as np

        rls_list = []
        total_n_snp = 0
        for chunk_i in range(config["N_TEST_CHUNK"]):
            chunk_rls = pd.read_csv(input[chunk_i])
            chunk_n_snp = len(chunk_rls["SNP_I"].unique())
            chunk_rls["SNP_I"] += total_n_snp
            total_n_snp += chunk_n_snp
            rls_list.append(chunk_rls)
        print(f"total_n_snp: {total_n_snp}")
        pd.concat(rls_list).to_csv(output[0], index=False)