castle/Scripts/makeIndices.py at master · jmerkin/castle · 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
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
###############################################################################
# David Morgens
# 04/28/2016
################################################################################
# Import neccessary modules

import argparse
import sys
import csv
import subprocess
import shlex
import os


##############################################################################
# Initiates argument parser

parser = argparse.ArgumentParser('Make index for alignment')

parser.add_argument('oligo_file', help='Input oligo file; csv format', type=str)

parser.add_argument('short_name', help='The screen type for reference', type=str)

parser.add_argument('full_name', help='Name for output files', type=str)

# Optional arguments: base trimming of fasta
parser.add_argument('-s','--strim', help='Trim bases from start; default is 0',
                        default=0, type=int)

parser.add_argument('-e', '--etrim', help='Trim bases from end; default is 0',
                        default=0, type=int)

parser.add_argument('-n', '--nums', help='Number of oligos to output',
                        default=2, type=int)

parser.add_argument('-o', '--override', help='Flag to override existing indexes',
                        action='store_true')

parser.add_argument('-t', '--test', help="Flag to not run bowtie",
                        action='store_false')

args = parser.parse_args()


##############################################################################
#

index_file = os.path.join('Indices', 'screen_type_index.txt')
index = []

# Check whether screen type and name already exist
with open(index_file, 'r') as index_open:

    index_csv = csv.reader(index_open, delimiter='\t')

    for line in index_csv:

        name, location = line

        if name == args.short_name or location == args.full_name:

            if not args.override:
                sys.exit('Error: Screen name taken')

            else:
                print('Warning: Previous screen overwritten')

        else:
            index.append((name, location))


##############################################################################
# Convert oligo file into bowtie-compatible fasta

with open(args.oligo_file, 'r') as oligo_file:

    oligo_csv = csv.reader(oligo_file, delimiter=',')
    oligo_list = []

    if args.etrim > 0:
        for line in oligo_csv:
            oligo = line[1][args.strim: -args.etrim]
            oligo_list.append(['>' + line[0], oligo])

    else:
        for line in oligo_csv:
            oligo = line[1][args.strim: ]
            oligo_list.append(['>' + line[0], oligo])

if args.test:

    for t in range(args.nums):
        oligo = oligo_list[t][1]
        print('Sample oligo: ' + oligo)

    for t in range(args.nums):
        oligo = oligo_list[t][1]
        print('Sample length: ' + str(len(oligo)))

    sys.exit('Warning: Files not created.\n' +
                'Use -t or --test to create files')


##############################################################################
#

fasta_location = os.path.join('Indices', 'temp_fasta.fna')

with open(fasta_location, 'w') as fasta_open:

    fasta_csv = csv.writer(fasta_open, delimiter='\n')

    for i in oligo_list:
        fasta_csv.writerow(i)


##############################################################################
# Call bowtie-build to build new index

try:
    subprocess.check_call('bowtie-build ' + fasta_location + ' ' +
                                            os.path.join('Indices', args.full_name),
                                                                    shell=True)
except:
    sys.exit()

# delete fasta files
try:
    subprocess.check_call('rm ' + fasta_location, shell=True)
except:
    sys.exit()


##############################################################################
# Add new index to screentype_index file

index_name = os.path.join('Indices', args.full_name)
index.append([args.short_name, index_name])
index.sort(key=lambda x: x[0])

with open(index_file, 'w') as index_open:

    index_csv = csv.writer(index_open, delimiter='\t')

    for ind in index:
        index_csv.writerow(ind)


##############################################################################