-
Notifications
You must be signed in to change notification settings - Fork 0
/
phase-01.py
162 lines (137 loc) · 4.45 KB
/
phase-01.py
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
"""
Phasing NGS and array datasets with shapeit or eagle.
Author: Gennady Khvorykh, info@inzilico.com
Only autosomes are phased. They are processed in parallel.
The input is vcf.gz file. The reference panel and recombination map are applied.
The genomic assembly is GRCh38.
Date created: January 23, 2024
"""
import os, sys
import argparse
import glob
import time
from subprocess import run
from multiprocessing import Pool
from natsort import natsorted
# === Functions ===
def check_file(x):
if not os.path.isfile(x):
print(x, "doesn't exist")
sys.exit(1)
def load_resources(res_file):
res = {}
with open(res_file, "r") as f:
for line in f:
values = line.strip().split(",")
res[values[0]] = values[1]
return res
def index(x):
if not os.path.isfile(x + ".csi"):
cmd = f"{res['bcftools']} index --threads {cpu} {x}"
run(cmd, shell = True)
def phase_shapeit(i):
"""
Phase a chromosome with SHAPEIT tool
"""
print("Processing chromosome", i)
map_file = f"{res['map38']}/plink.chr{i}.GRCh38.map.shapeit.txt"
ref_file = f"{res['ref1kg38']}/chr{i}.vcf.gz"
phased_file = f"chr{i}.phased.bcf"
cmd = f"{res['shapeit']} --input {input_file} --region chr{i} --reference {ref_file} \
--map {map_file} --thread {cpu} --output {phased_file} --log chr{i}.log"
p = run(cmd, shell = True)
if p.returncode != 0:
raise Exception(f'Phase chromosome {i} invalid result: {p.returncode}')
check_file(phased_file)
def phase_eagle(i):
"""
Phase a chromosome with Eagle tool
"""
print("Processing chromosome", i)
map_file = f"{res['eagle']}/tables/genetic_map_hg38_withX.txt.gz"
output_prefix = f"chr{i}"
ref_file = f"{res['ref1kg38']}/{output_prefix}.vcf.gz"
cmd = f"{res['eagle']}/eagle --vcfRef={ref_file} \
--geneticMapFile={map_file} --chrom={i} \
--vcfOutFormat=u --outPrefix={output_prefix}.phased \
--vcfTarget={input_file}"
p = run(cmd, shell=True, capture_output=True, text=True)
print(p.stdout, file=open(f'{output_prefix}.log', 'w'))
print("\n", p.stderr, file=open(f'{output_prefix}.log', 'a'))
# === End of Functions ===
# Parse command line arguments
parser = argparse.ArgumentParser(description="A wrapper script to phase genotypes")
parser.add_argument(
"vcfgz_file", help="path/to/filename.vcf.gz with testing genetic variants"
)
parser.add_argument(
"-r", "--res_file", required=True, help="path/to/file.csv with paths to tools"
)
parser.add_argument("-c", "--cpu", default=22, help="The number of CPUs (default: 22)")
parser.add_argument("-t", "--tool", default="shapeit", help="The tool to phase the genotypes: shapeit, eagle (default: shapeit)")
args = parser.parse_args()
# Initiate variables
vcfgz_file = args.vcfgz_file
res_file = args.res_file
tool = args.tool
cpu = args.cpu
# Check input
check_file(vcfgz_file)
# Load resources
res = load_resources(res_file)
# Get time
t1 = time.time()
# Copy input file into working folder
input_file = vcfgz_file
output_file = os.path.basename(input_file)
if not os.path.isfile(output_file):
cmd = f"rsync -avzh {input_file} {output_file}"
print(cmd)
run(cmd, shell = True)
check_file(output_file)
# Index input file
index(output_file)
# Process further the copy of input file in the working folder
input_file = output_file
# Loop over the chromosomes and phase them
if tool == "shapeit":
with Pool(cpu) as pool:
pool.map(phase_shapeit, range(1, 23))
elif tool == "eagle":
with Pool(cpu) as pool:
pool.map(phase_eagle, range(1, 23))
# for i in range(22, 23):
# phase_eagle(i)
else:
print("Unknown tool:", tool)
sys.exit(1)
# Glob phased files
bcfs = glob.glob("chr*.phased.bcf")
# Sort phased files
bcfs = natsorted(bcfs)
# Concatenate phased bcf files
bcf_arg = ' '.join(bcfs)
output_file = input_file.replace(".vcf.gz", ".phased.vcf.gz")
cmd = f"{res['bcftools']} concat --write-index --threads {cpu} -Oz -o {output_file} {bcf_arg}"
p = run(cmd, shell = True)
if p.returncode != 0:
raise Exception(f'Concatenate bcf failed: {p.returncode}')
# Glob log files
logs = glob.glob("chr*.log")
logs = natsorted(logs)
# Concatenate log files
log_arg = ' '.join(logs)
log_file = input_file.replace(".vcf.gz", ".log")
cmd = f"cat {log_arg} > {log_file}"
p = run(cmd, shell = True)
if p.returncode != 0:
raise Exception(f'Concatenate log failed: {p.returncode}')
# Clean
for bcf in bcfs:
os.remove(bcf)
for log in logs:
os.remove(log)
# Show time elapsed
dur = time.time()-t1
dur = time.strftime("%H:%M:%S", time.gmtime(dur))
print("\nTime spent: ", dur, file = open(log_file, 'a'))