def parseIMGT(aligner_output, seq_file=None, no_parse=True, partial=False,
parse_scores=False, parse_regions=False, parse_junction=False,
out_args=default_out_args):
"""
Main for IMGT aligned sample sequences.
Arguments:
aligner_output : zipped file or unzipped folder output by IMGT.
seq_file : FASTA file input to IMGT (from which to get seqID).
no_parse : if ID is to be parsed for pRESTO output with default delimiters.
partial : If True put incomplete alignments in the pass file.
parse_scores : if True add alignment score fields to output file.
parse_regions : if True add FWR and CDR region fields to output file.
out_args : common output argument dictionary from parseCommonArgs.
Returns:
None
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'MakeDb'
log['ALIGNER'] = 'IMGT'
log['ALIGNER_OUTPUT'] = aligner_output
log['SEQ_FILE'] = os.path.basename(seq_file) if seq_file else ''
log['NO_PARSE'] = no_parse
log['PARTIAL'] = partial
log['SCORES'] = parse_scores
log['REGIONS'] = parse_regions
log['JUNCTION'] = parse_junction
printLog(log)
start_time = time()
printMessage('Loading sequence files', start_time=start_time, width=25)
# Extract IMGT files
temp_dir, imgt_files = extractIMGT(aligner_output)
# Count records in IMGT files
total_count = countDbFile(imgt_files['summary'])
# Get (parsed) IDs from fasta file submitted to IMGT
id_dict = getIDforIMGT(seq_file) if seq_file else {}
printMessage('Done', start_time=start_time, end=True, width=25)
# Parse IMGT output and write db
with open(imgt_files['summary'], 'r') as summary_handle, \
open(imgt_files['gapped'], 'r') as gapped_handle, \
open(imgt_files['ntseq'], 'r') as ntseq_handle, \
open(imgt_files['junction'], 'r') as junction_handle:
parse_iter = IMGTReader(summary_handle, gapped_handle, ntseq_handle, junction_handle,
parse_scores=parse_scores, parse_regions=parse_regions,
parse_junction=parse_junction)
file_prefix = getFilePrefix(aligner_output, out_args)
writeDb(parse_iter, parse_iter.fields, file_prefix, total_count, id_dict=id_dict,
no_parse=no_parse, partial=partial, out_args=out_args)
# Cleanup temp directory
temp_dir.cleanup()
return None
def parseIgBLAST(aligner_output, seq_file, repo, no_parse=True, partial=False,
parse_regions=False, parse_scores=False, parse_igblast_cdr3=False,
out_args=default_out_args):
"""
Main for IgBLAST aligned sample sequences.
Arguments:
aligner_output : IgBLAST output file to process.
seq_file : fasta file input to IgBlast (from which to get sequence).
repo : folder with germline repertoire files.
no_parse : if ID is to be parsed for pRESTO output with default delimiters.
partial : If True put incomplete alignments in the pass file.
parse_regions : if True add FWR and CDR fields to output file.
parse_scores : if True add alignment score fields to output file.
parse_igblast_cdr3 : if True parse CDR3 sequences generated by IgBLAST
out_args : common output argument dictionary from parseCommonArgs.
Returns:
None
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'MakeDB'
log['ALIGNER'] = 'IgBlast'
log['ALIGNER_OUTPUT'] = os.path.basename(aligner_output)
log['SEQ_FILE'] = os.path.basename(seq_file)
log['NO_PARSE'] = no_parse
log['PARTIAL'] = partial
log['SCORES'] = parse_scores
log['REGIONS'] = parse_regions
printLog(log)
start_time = time()
printMessage('Loading sequence files', start_time=start_time, width=25)
# Count records in sequence file
total_count = countSeqFile(seq_file)
# Get input sequence dictionary
seq_dict = getSeqDict(seq_file)
# Create germline repo dictionary
repo_dict = readRepo(repo)
printMessage('Done', start_time=start_time, end=True, width=25)
# Parse and write output
with open(aligner_output, 'r') as f:
parse_iter = IgBLASTReader(f, seq_dict, repo_dict,
parse_scores=parse_scores, parse_regions=parse_regions,
parse_igblast_cdr3=parse_igblast_cdr3)
file_prefix = getFilePrefix(aligner_output, out_args)
writeDb(parse_iter, parse_iter.fields, file_prefix, total_count,
no_parse=no_parse, partial=partial, out_args=out_args)
return None
def sortDbFile(db_file, field, numeric=False, descend=False,
out_args=default_out_args):
"""
Sorts records by values in an annotation field
Arguments:
db_file = the database filename
field = the field name to sort by
numeric = if True sort field numerically;
if False sort field alphabetically
descend = if True sort in descending order;
if False sort in ascending order
out_args = common output argument dictionary from parseCommonArgs
Returns:
the output file name
"""
log = OrderedDict()
log['START'] = 'ParseDb'
log['COMMAND'] = 'sort'
log['FILE'] = os.path.basename(db_file)
log['FIELD'] = field
log['NUMERIC'] = numeric
printLog(log)
# Open file handles
db_iter = readDbFile(db_file, ig=False)
pass_handle = getOutputHandle(db_file, out_label='parse-sort', out_dir=out_args['out_dir'],
out_name=out_args['out_name'], out_type='tab')
pass_writer = getDbWriter(pass_handle, db_file)
# Store all records in a dictionary
start_time = time()
printMessage("Indexing: Running", start_time=start_time)
db_dict = {i:r for i, r in enumerate(db_iter)}
result_count = len(db_dict)
# Sort db_dict by field values
tag_dict = {k:v[field] for k, v in db_dict.items()}
if numeric: tag_dict = {k:float(v or 0) for k, v in tag_dict.items()}
sorted_keys = sorted(tag_dict, key=tag_dict.get, reverse=descend)
printMessage("Indexing: Done", start_time=start_time, end=True)
# Iterate over records
start_time = time()
rec_count = 0
for key in sorted_keys:
# Print progress for previous iteration
printProgress(rec_count, result_count, 0.05, start_time)
rec_count += 1
# Write records
pass_writer.writerow(db_dict[key])
# Print counts
printProgress(rec_count, result_count, 0.05, start_time)
log = OrderedDict()
log['OUTPUT'] = os.path.basename(pass_handle.name)
log['RECORDS'] = rec_count
log['END'] = 'ParseDb'
printLog(log)
# Close file handles
pass_handle.close()
return pass_handle.name
def sortDbFile(db_file,
field,
numeric=False,
descend=False,
out_file=None,
out_args=default_out_args):
"""
Sorts records by values in an annotation field
Arguments:
db_file : the database filename
field : the field name to sort by
numeric : if True sort field numerically;
if False sort field alphabetically
descend : if True sort in descending order;
if False sort in ascending order
out_file : output file name. Automatically generated from the input file if None.
out_args : common output argument dictionary from parseCommonArgs
Returns:
str : output file name
"""
log = OrderedDict()
log['START'] = 'ParseDb'
log['COMMAND'] = 'sort'
log['FILE'] = os.path.basename(db_file)
log['FIELD'] = field
log['NUMERIC'] = numeric
printLog(log)
# Open input
db_handle = open(db_file, 'rt')
db_iter = TSVReader(db_handle)
out_fields = db_iter.fields
__, __, out_args['out_type'] = splitName(db_file)
# Open output
if out_file is not None:
pass_handle = open(out_file, 'w')
else:
pass_handle = getOutputHandle(db_file,
out_label='parse-sort',
out_dir=out_args['out_dir'],
out_name=out_args['out_name'],
out_type=out_args['out_type'])
pass_writer = TSVWriter(pass_handle, out_fields)
# Store all records in a dictionary
start_time = time()
printMessage("Indexing: Running", start_time=start_time)
db_dict = {i: r for i, r in enumerate(db_iter)}
result_count = len(db_dict)
# Sort db_dict by field values
tag_dict = {k: v[field] for k, v in db_dict.items()}
if numeric: tag_dict = {k: float(v or 0) for k, v in tag_dict.items()}
sorted_keys = sorted(tag_dict, key=tag_dict.get, reverse=descend)
printMessage("Indexing: Done", start_time=start_time, end=True)
# Iterate over records
start_time = time()
rec_count = 0
for key in sorted_keys:
# Print progress for previous iteration
printProgress(rec_count, result_count, 0.05, start_time=start_time)
rec_count += 1
# Write records
pass_writer.writeDict(db_dict[key])
# Print counts
printProgress(rec_count, result_count, 0.05, start_time=start_time)
log = OrderedDict()
log['OUTPUT'] = os.path.basename(pass_handle.name)
log['RECORDS'] = rec_count
log['END'] = 'ParseDb'
printLog(log)
# Close file handles
pass_handle.close()
db_handle.close()
return pass_handle.name
def convertToGenbank(db_file, inference=None, db_xref=None, molecule=default_molecule,
product=default_product, features=None, c_field=None, label=None,
count_field=None, index_field=None, allow_stop=False,
asis_id=False, asis_calls=False, allele_delim=default_allele_delim,
build_asn=False, asn_template=None, tbl2asn_exec=default_tbl2asn_exec,
format=default_format, out_file=None,
out_args=default_out_args):
"""
Builds GenBank submission fasta and table files
Arguments:
db_file : the database file name.
inference : reference alignment tool.
db_xref : reference database link.
molecule : source molecule (eg, "mRNA", "genomic DNA")
product : Product (protein) name.
features : dictionary of sample features (BioSample attributes) to add to the description of each record.
c_field : column containing the C region gene call.
label : a string to use as a label for the ID. if None do not add a field label.
count_field : field name to populate the AIRR_READ_COUNT note.
index_field : field name to populate the AIRR_CELL_INDEX note.
allow_stop : if True retain records with junctions having stop codons.
asis_id : if True use the original sequence ID for the output IDs.
asis_calls : if True do not parse gene calls for IMGT nomenclature.
allele_delim : delimiter separating the gene name from the allele number when asis_calls=True.
build_asn : if True run tbl2asn on the generated .tbl and .fsa files.
asn_template : template file (.sbt) to pass to tbl2asn.
tbl2asn_exec : name of or path to the tbl2asn executable.
format : input and output format.
out_file : output file name without extension. Automatically generated from the input file if None.
out_args : common output argument dictionary from parseCommonArgs.
Returns:
tuple : the output (feature table, fasta) file names.
"""
log = OrderedDict()
log['START'] = 'ConvertDb'
log['COMMAND'] = 'genbank'
log['FILE'] = os.path.basename(db_file)
printLog(log)
# Define format operators
try:
reader, __, schema = getFormatOperators(format)
except ValueError:
printError('Invalid format %s.' % format)
# Open input
db_handle = open(db_file, 'rt')
db_iter = reader(db_handle)
# Check for required columns
try:
required = ['sequence_input',
'v_call', 'd_call', 'j_call',
'v_seq_start', 'd_seq_start', 'j_seq_start']
checkFields(required, db_iter.fields, schema=schema)
except LookupError as e:
printError(e)
# Open output
if out_file is not None:
out_name, __ = os.path.splitext(out_file)
fsa_handle = open('%s.fsa' % out_name, 'w')
tbl_handle = open('%s.tbl' % out_name, 'w')
else:
fsa_handle = getOutputHandle(db_file, out_label='genbank', out_dir=out_args['out_dir'],
out_name=out_args['out_name'], out_type='fsa')
tbl_handle = getOutputHandle(db_file, out_label='genbank', out_dir=out_args['out_dir'],
out_name=out_args['out_name'], out_type='tbl')
# Count records
result_count = countDbFile(db_file)
# Define writer
writer = csv.writer(tbl_handle, delimiter='\t', quoting=csv.QUOTE_NONE)
# Iterate over records
start_time = time()
rec_count, pass_count, fail_count = 0, 0, 0
for rec in db_iter:
# Print progress for previous iteration
printProgress(rec_count, result_count, 0.05, start_time=start_time)
rec_count += 1
# Extract table dictionary
name = None if asis_id else rec_count
seq = makeGenbankSequence(rec, name=name, label=label, count_field=count_field, index_field=index_field,
molecule=molecule, features=features)
tbl = makeGenbankFeatures(rec, start=seq['start'], end=seq['end'], product=product,
db_xref=db_xref, inference=inference, c_field=c_field,
allow_stop=allow_stop, asis_calls=asis_calls, allele_delim=allele_delim)
if tbl is not None:
pass_count +=1
# Write table
writer.writerow(['>Features', seq['record'].id])
for feature, qualifiers in tbl.items():
writer.writerow(feature)
if qualifiers:
for x in qualifiers:
writer.writerow(list(chain(['', '', ''], x)))
# Write sequence
SeqIO.write(seq['record'], fsa_handle, 'fasta')
else:
fail_count += 1
# Final progress bar
printProgress(rec_count, result_count, 0.05, start_time=start_time)
# Run tbl2asn
if build_asn:
start_time = time()
printMessage('Running tbl2asn', start_time=start_time, width=25)
result = runASN(fsa_handle.name, template=asn_template, exec=tbl2asn_exec)
printMessage('Done', start_time=start_time, end=True, width=25)
# Print ending console log
log = OrderedDict()
log['OUTPUT_TBL'] = os.path.basename(tbl_handle.name)
log['OUTPUT_FSA'] = os.path.basename(fsa_handle.name)
log['RECORDS'] = rec_count
log['PASS'] = pass_count
log['FAIL'] = fail_count
log['END'] = 'ConvertDb'
printLog(log)
# Close file handles
tbl_handle.close()
fsa_handle.close()
db_handle.close()
return (tbl_handle.name, fsa_handle.name)
def pairSeq(seq_file_1, seq_file_2, fields_1=None, fields_2=None,
coord_type=default_coord_type,
out_args=default_out_args):
"""
Generates consensus sequences
Arguments:
seq_file_1 = the file containing the grouped sequences and annotations
seq_file_2 = the file to assign annotations to from seq_file_1
fields_1 = list of annotations in seq_file_1 records to copy to seq_file_2 records;
if None do not copy any annotations
fields_2 = list of annotations in seq_file_2 records to copy to seq_file_1 records;
if None do not copy any annotations
coord_type = the sequence header format
out_args = common output argument dictionary from parseCommonArgs
Returns:
a list of tuples holding successfully paired filenames for (seq_file_1, seq_file_2)
"""
# Define private functions
def _key_func(x):
return getCoordKey(x, coord_type=coord_type, delimiter=out_args['delimiter'])
log = OrderedDict()
log['START'] = 'PairSeq'
log['FILE1'] = os.path.basename(seq_file_1)
log['FILE2'] = os.path.basename(seq_file_2)
log['FIELDS_1'] = ','.join(fields_1) if fields_1 is not None else None
log['FIELDS_2'] = ','.join(fields_2) if fields_2 is not None else None
log['COORD_TYPE'] = coord_type
printLog(log)
# Define output type
if out_args['out_type'] is None:
out_type_1 = getFileType(seq_file_1)
out_type_2 = getFileType(seq_file_2)
else:
out_type_1 = out_type_2 = out_args['out_type']
# Define output name
if out_args['out_name'] is None:
out_name_1 = out_name_2 = None
else:
out_name_1 = '%s-1' % out_args['out_name']
out_name_2 = '%s-2' % out_args['out_name']
# Open and count files
start_time = time()
printMessage("Indexing files", start_time=start_time)
# Index file 1
seq_count_1 = countSeqFile(seq_file_1)
seq_dict_1 = readSeqFile(seq_file_1, index=True, key_func=_key_func)
# Define file 2 iterator
seq_count_2 = countSeqFile(seq_file_2)
seq_iter_2 = readSeqFile(seq_file_2, index=False)
printMessage("Done", start_time=start_time, end=True)
# Open output file handles
pass_handle_1 = getOutputHandle(seq_file_1, 'pair-pass', out_args['out_dir'],
out_name=out_name_1, out_type=out_type_1)
pass_handle_2 = getOutputHandle(seq_file_2, 'pair-pass', out_args['out_dir'],
out_name=out_name_2, out_type=out_type_2)
if out_args['failed']:
fail_handle_1 = getOutputHandle(seq_file_1, 'pair-fail', out_dir=out_args['out_dir'],
out_name=out_name_1, out_type=out_type_1)
fail_handle_2 = getOutputHandle(seq_file_2, 'pair-fail', out_dir=out_args['out_dir'],
out_name=out_name_2, out_type=out_type_2)
pass_keys = list()
# Iterate over pairs and write to output files
start_time = time()
rec_count = pair_count = 0
for seq_2 in seq_iter_2:
# Print progress for previous iteration
printProgress(rec_count, seq_count_2, 0.05, start_time)
rec_count += 1
# Check for file 2 mate pair in file 1
coord_2 = getCoordKey(seq_2.id, coord_type=coord_type,
delimiter=out_args['delimiter'])
seq_1 = seq_dict_1.get(coord_2, None)
if seq_1 is not None:
# Record paired keys
pair_count += 1
if fields_1 is not None or fields_2 is not None:
ann_1 = parseAnnotation(seq_1.description,
delimiter=out_args['delimiter'])
ann_2 = parseAnnotation(seq_2.description,
delimiter=out_args['delimiter'])
# Prepend annotations from seq_1 to seq_2
if fields_1 is not None:
copy_ann = OrderedDict([(k, v) for k, v in ann_1.items() \
if k in fields_1])
merge_ann = mergeAnnotation(ann_2, copy_ann, prepend=True,
delimiter=out_args['delimiter'])
seq_2.id = flattenAnnotation(merge_ann,
delimiter=out_args['delimiter'])
seq_2.description = ''
# Append annotations from seq_2 to seq_1
if fields_2 is not None:
copy_ann = OrderedDict([(k, v) for k, v in ann_2.items() \
if k in fields_2])
merge_ann = mergeAnnotation(ann_1, copy_ann, prepend=False,
delimiter=out_args['delimiter'])
seq_1.id = flattenAnnotation(merge_ann,
delimiter=out_args['delimiter'])
seq_1.description = ''
# Write paired records
SeqIO.write(seq_1, pass_handle_1, out_type_1)
SeqIO.write(seq_2, pass_handle_2, out_type_2)
# Write unpaired file 2 records and updated paired key list for finding unpaired file 1 records
if out_args['failed']:
if seq_1 is not None: pass_keys.append(coord_2)
else: SeqIO.write(seq_2, fail_handle_2, out_type_2)
# Print final progress
printProgress(rec_count, seq_count_2, 0.05, start_time)
# Find and write unpaired file 1 records
if out_args['failed']:
start_time = time()
printMessage("Finding unpaired", start_time=start_time)
# Find file 1 unpaired keys
pass_keys = set(pass_keys)
unpaired = set(seq_dict_1).difference(pass_keys)
# Write unpaired file 1 records
for k in unpaired: SeqIO.write(seq_dict_1[k], fail_handle_1, out_type_1)
printMessage("Done", start_time=start_time, end=True)
# Print log
log = OrderedDict()
log['OUTPUT1'] = os.path.basename(pass_handle_1.name)
log['OUTPUT2'] = os.path.basename(pass_handle_2.name)
log['SEQUENCES1'] = seq_count_1
log['SEQUENCES2'] = seq_count_2
log['PASS'] = pair_count
log['END'] = 'PairSeq'
printLog(log)
# Close file handles
pass_handle_1.close()
pass_handle_2.close()
return [(pass_handle_1.name, pass_handle_2.name)]
def samplePairSeqFile(seq_file_1,
seq_file_2,
max_count,
field=None,
values=None,
coord_type=default_coord,
out_args=default_out_args):
"""
Samples from paired-end sequence files
Arguments:
seq_file_1 : filename of the first paired-end sequence file
seq_file_2 : filename of the second paired-end sequence file
max_count : a list of the maximum number of sequences to sample
field : the annotation field to check for required values
values : a list of annotation values that a sample must contain one of
coord_type : the sequence header format
out_args : common output argument dictionary from parseCommonArgs
Returns:
list: seq_file_1 and seq_file_2 output file names
"""
# Sequence index key function
def _key_func(x):
return getCoordKey(x,
coord_type=coord_type,
delimiter=out_args['delimiter'])
# Function to sample from two lists of sequence indices
def _sample_list(n, index_1, index_2):
key_set = set(index_1).intersection(index_2)
max_n = len(key_set)
return random.sample(key_set, min(n, max_n))
# Function to sample from two dictionaries of grouped sequence indices
def _sample_dict(n, index_1, index_2):
group_set = set(index_1.keys()).intersection(index_2.keys())
sample_list = []
for k in group_set:
key_set = set(index_1[k]).intersection(index_2[k])
max_n = len(key_set)
sample_list.extend(random.sample(key_set, min(n, max_n)))
return sample_list
# Print console log
log = OrderedDict()
log['START'] = 'SplitSeq'
log['COMMAND'] = 'samplepair'
log['FILE1'] = os.path.basename(seq_file_1)
log['FILE2'] = os.path.basename(seq_file_2)
log['MAX_COUNTS'] = ','.join([str(x) for x in max_count])
log['FIELD'] = field
log['VALUES'] = ','.join(values) if values else None
printLog(log)
# Define output type
in_type_1 = getFileType(seq_file_1)
in_type_2 = getFileType(seq_file_2)
if out_args['out_type'] is None:
out_type_1 = in_type_1
out_type_2 = in_type_2
else:
out_type_1 = out_type_2 = out_args['out_type']
# Define output name
if out_args['out_name'] is None:
out_name_1 = out_name_2 = None
else:
out_name_1 = '%s-1' % out_args['out_name']
out_name_2 = '%s-2' % out_args['out_name']
# Index input files
start_time = time()
printMessage('Reading files', start_time=start_time, width=25)
seq_dict_1 = readSeqFile(seq_file_1, index=True, key_func=_key_func)
seq_dict_2 = readSeqFile(seq_file_2, index=True, key_func=_key_func)
# Subset keys to those meeting field/value criteria
if field is not None and values is not None:
_sample = _sample_list
printMessage('Subsetting by annotation',
start_time=start_time,
width=25)
seq_index_1 = subsetSeqIndex(seq_dict_1,
field,
values,
delimiter=out_args['delimiter'])
seq_index_2 = subsetSeqIndex(seq_dict_2,
field,
values,
delimiter=out_args['delimiter'])
elif field is not None and values is None:
_sample = _sample_dict
printMessage('Indexing by annotation', start_time=start_time, width=25)
seq_index_1 = indexSeqSets(seq_dict_1,
field,
delimiter=out_args['delimiter'])
seq_index_2 = indexSeqSets(seq_dict_2,
field,
delimiter=out_args['delimiter'])
else:
_sample = _sample_list
seq_index_1 = list(seq_dict_1.keys())
seq_index_2 = list(seq_dict_2.keys())
printMessage('Done', start_time=start_time, end=True, width=25)
# Generate sample set for each value in max_count
out_files = []
for i, n in enumerate(max_count):
start_time = time()
printMessage('Sampling n=%i' % n, start_time=start_time, width=25)
# Sample
sample_keys = _sample(n, seq_index_1, seq_index_2)
sample_count = len(sample_keys)
# Open file handles
out_handle_1 = getOutputHandle(seq_file_1,
'sample%i-n%i' % (i + 1, sample_count),
out_dir=out_args['out_dir'],
out_name=out_name_1,
out_type=out_type_1)
out_handle_2 = getOutputHandle(seq_file_2,
'sample%i-n%i' % (i + 1, sample_count),
out_dir=out_args['out_dir'],
out_name=out_name_2,
out_type=out_type_2)
out_files.append((out_handle_1.name, out_handle_2.name))
for k in sample_keys:
SeqIO.write(seq_dict_1[k], out_handle_1, out_type_1)
SeqIO.write(seq_dict_2[k], out_handle_2, out_type_2)
printMessage('Done', start_time=start_time, end=True, width=25)
# Print log for iteration
log = OrderedDict()
log['MAX_COUNT'] = n
log['SAMPLED'] = sample_count
log['OUTPUT1'] = os.path.basename(out_files[i][0])
log['OUTPUT2'] = os.path.basename(out_files[i][1])
printLog(log)
# Close file handles
out_handle_1.close()
out_handle_2.close()
# Print log
log = OrderedDict()
log['END'] = 'SplitSeq'
printLog(log)
return out_files
def sampleSeqFile(seq_file,
max_count,
field=None,
values=None,
out_args=default_out_args):
"""
Samples from a sequence file
Arguments:
seq_file : filename of the sequence file to sample from
max_count : a list of the maximum number of sequences to sample
field : the annotation field to check for required values
values : a list of annotation values that a sample must contain one of
out_args : common output argument dictionary from parseCommonArgs
Returns:
str: output file name
"""
# Function to sample from a list of sequence indices
def _sample_list(n, index_list):
max_n = len(index_list)
r = random.sample(range(max_n), n) if n < max_n else range(max_n)
return [index_list[x] for x in r]
# Function to sample from a dictionary of grouped sequence indices
def _sample_dict(n, index_dict):
sample_list = []
for v in index_dict.values():
max_n = len(v)
r = random.sample(range(max_n), n) if n < max_n else range(max_n)
sample_list.extend([v[x] for x in r])
return sample_list
# Print console log
log = OrderedDict()
log['START'] = 'SplitSeq'
log['COMMAND'] = 'sample'
log['FILE'] = os.path.basename(seq_file)
log['MAX_COUNTS'] = ','.join([str(x) for x in max_count])
log['FIELD'] = field
log['VALUES'] = ','.join(values) if values else None
printLog(log)
# Read input files and open output files
start_time = time()
printMessage('Reading files', start_time=start_time, width=25)
in_type = getFileType(seq_file)
seq_dict = readSeqFile(seq_file, index=True)
if out_args['out_type'] is None: out_args['out_type'] = in_type
# Generate subset of records
if field is not None and values is not None:
_sample = _sample_list
printMessage('Subsetting by annotation',
start_time=start_time,
width=25)
seq_index = subsetSeqIndex(seq_dict,
field,
values,
delimiter=out_args['delimiter'])
elif field is not None and values is None:
_sample = _sample_dict
printMessage('Indexing by annotation', start_time=start_time, width=25)
seq_index = indexSeqSets(seq_dict,
field,
delimiter=out_args['delimiter'])
else:
_sample = _sample_list
seq_index = [k for k in seq_dict]
printMessage('Done', start_time=start_time, end=True, width=25)
# Generate sample set for each value in max_count
out_files = []
for i, n in enumerate(max_count):
start_time = time()
printMessage('Sampling n=%i' % n, start_time=start_time, width=25)
# Sample from records
sample_keys = _sample(n, seq_index)
sample_count = len(sample_keys)
# Write sampled sequences to files
with getOutputHandle(seq_file,
'sample%i-n%i' % (i + 1, sample_count),
out_dir=out_args['out_dir'],
out_name=out_args['out_name'],
out_type=out_args['out_type']) as out_handle:
for k in sample_keys:
SeqIO.write(seq_dict[k], out_handle, out_args['out_type'])
out_files.append(out_handle.name)
printMessage('Done', start_time=start_time, end=True, width=25)
# Print log for iteration
log = OrderedDict()
log['MAX_COUNT'] = n
log['SAMPLED'] = sample_count
log['OUTPUT'] = os.path.basename(out_files[i])
printLog(log)
# Print log
log = OrderedDict()
log['END'] = 'SplitSeq'
printLog(log)
return out_files
def createGermlines(db_file, references, seq_field=default_seq_field, v_field=default_v_field,
d_field=default_d_field, j_field=default_j_field,
cloned=False, clone_field=default_clone_field, germ_types=default_germ_types,
format=default_format, out_file=None, out_args=default_out_args):
"""
Write germline sequences to tab-delimited database file
Arguments:
db_file : input tab-delimited database file.
references : folders and/or files containing germline repertoire data in FASTA format.
seq_field : field in which to look for sequence.
v_field : field in which to look for V call.
d_field : field in which to look for D call.
j_field : field in which to look for J call.
cloned : if True build germlines by clone, otherwise build individual germlines.
clone_field : field containing clone identifiers; ignored if cloned=False.
germ_types : list of germline sequence types to be output from the set of 'full', 'dmask', 'vonly', 'regions'
format : input and output format.
out_file : output file name. Automatically generated from the input file if None.
out_args : arguments for output preferences.
Returns:
dict: names of the 'pass' and 'fail' output files.
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'CreateGermlines'
log['FILE'] = os.path.basename(db_file)
log['GERM_TYPES'] = ','.join(germ_types)
log['SEQ_FIELD'] = seq_field
log['V_FIELD'] = v_field
log['D_FIELD'] = d_field
log['J_FIELD'] = j_field
log['CLONED'] = cloned
if cloned: log['CLONE_FIELD'] = clone_field
printLog(log)
# Define format operators
try:
reader, writer, schema = getFormatOperators(format)
except ValueError:
printError('Invalid format %s' % format)
out_args['out_type'] = schema.out_type
# TODO: this won't work for AIRR necessarily
# Define output germline fields
germline_fields = OrderedDict()
seq_type = seq_field.split('_')[-1]
if 'full' in germ_types: germline_fields['full'] = 'germline_' + seq_type
if 'dmask' in germ_types: germline_fields['dmask'] = 'germline_' + seq_type + '_d_mask'
if 'vonly' in germ_types: germline_fields['vonly'] = 'germline_' + seq_type + '_v_region'
if 'regions' in germ_types: germline_fields['regions'] = 'germline_regions'
if cloned:
germline_fields['v'] = 'germline_v_call'
germline_fields['d'] = 'germline_d_call'
germline_fields['j'] = 'germline_j_call'
out_fields = getDbFields(db_file,
add=[schema.fromReceptor(f) for f in germline_fields.values()],
reader=reader)
# Get repertoire and open Db reader
reference_dict = readGermlines(references)
db_handle = open(db_file, 'rt')
db_iter = reader(db_handle)
# Check for required columns
try:
required = ['v_germ_start_imgt', 'd_germ_start', 'j_germ_start',
'np1_length', 'np2_length']
checkFields(required, db_iter.fields, schema=schema)
except LookupError as e:
printError(e)
# Check for IMGT-gaps in germlines
if all('...' not in x for x in reference_dict.values()):
printWarning('Germline reference sequences do not appear to contain IMGT-numbering spacers. Results may be incorrect.')
# Count input
total_count = countDbFile(db_file)
# Check for existence of fields
for f in [v_field, d_field, j_field, seq_field]:
if f not in db_iter.fields:
printError('%s field does not exist in input database file.' % f)
# Translate to Receptor attribute names
v_field = schema.toReceptor(v_field)
d_field = schema.toReceptor(d_field)
j_field = schema.toReceptor(j_field)
seq_field = schema.toReceptor(seq_field)
clone_field = schema.toReceptor(clone_field)
# Define Receptor iterator
if cloned:
start_time = time()
printMessage('Sorting by clone', start_time=start_time, width=20)
sorted_records = sorted(db_iter, key=lambda x: x.getField(clone_field))
printMessage('Done', start_time=start_time, end=True, width=20)
receptor_iter = groupby(sorted_records, lambda x: x.getField(clone_field))
else:
receptor_iter = ((x.sequence_id, [x]) for x in db_iter)
# Define log handle
if out_args['log_file'] is None:
log_handle = None
else:
log_handle = open(out_args['log_file'], 'w')
# Initialize handles, writers and counters
pass_handle, pass_writer = None, None
fail_handle, fail_writer = None, None
rec_count, pass_count, fail_count = 0, 0, 0
start_time = time()
# Iterate over rows
for key, records in receptor_iter:
# Print progress
printProgress(rec_count, total_count, 0.05, start_time=start_time)
# Define iteration variables
records = list(records)
rec_log = OrderedDict([('ID', key)])
rec_count += len(records)
# Build germline for records
if len(records) == 1:
germ_log, germlines, genes = buildGermline(records[0], reference_dict, seq_field=seq_field, v_field=v_field,
d_field=d_field, j_field=j_field)
else:
germ_log, germlines, genes = buildClonalGermline(records, reference_dict, seq_field=seq_field, v_field=v_field,
d_field=d_field, j_field=j_field)
rec_log.update(germ_log)
# Write row to pass or fail file
if germlines is not None:
pass_count += len(records)
# Add germlines to Receptor record
annotations = {}
if 'full' in germ_types: annotations[germline_fields['full']] = germlines['full']
if 'dmask' in germ_types: annotations[germline_fields['dmask']] = germlines['dmask']
if 'vonly' in germ_types: annotations[germline_fields['vonly']] = germlines['vonly']
if 'regions' in germ_types: annotations[germline_fields['regions']] = germlines['regions']
if cloned:
annotations[germline_fields['v']] = genes['v']
annotations[germline_fields['d']] = genes['d']
annotations[germline_fields['j']] = genes['j']
# Write records
try:
for r in records:
r.setDict(annotations)
pass_writer.writeReceptor(r)
except AttributeError:
# Create output file handle and writer
if out_file is not None:
pass_handle = open(out_file, 'w')
else:
pass_handle = getOutputHandle(db_file,
out_label='germ-pass',
out_dir=out_args['out_dir'],
out_name=out_args['out_name'],
out_type=out_args['out_type'])
pass_writer = writer(pass_handle, fields=out_fields)
for r in records:
r.setDict(annotations)
pass_writer.writeReceptor(r)
else:
fail_count += len(records)
if out_args['failed']:
try:
fail_writer.writeReceptor(records)
except AttributeError:
fail_handle = getOutputHandle(db_file,
out_label='germ-fail',
out_dir=out_args['out_dir'],
out_name=out_args['out_name'],
out_type=out_args['out_type'])
fail_writer = writer(fail_handle, fields=out_fields)
fail_writer.writeReceptor(records)
# Write log
printLog(rec_log, handle=log_handle)
# Print log
printProgress(rec_count, total_count, 0.05, start_time=start_time)
log = OrderedDict()
log['OUTPUT'] = os.path.basename(pass_handle.name) if pass_handle is not None else None
log['RECORDS'] = rec_count
log['PASS'] = pass_count
log['FAIL'] = fail_count
log['END'] = 'CreateGermlines'
printLog(log)
# Close file handles
db_handle.close()
output = {'pass': None, 'fail': None}
if pass_handle is not None:
output['pass'] = pass_handle.name
pass_handle.close()
if fail_handle is not None:
output['fail'] = fail_handle.name
fail_handle.close()
if log_handle is not None:
log_handle.close()
return output
def clusterBarcodes(seq_file,
ident=default_cluster_ident,
length_ratio=default_length_ratio,
barcode_field=default_barcode_field,
cluster_field=default_cluster_field,
cluster_prefix=default_cluster_prefix,
cluster_tool=default_cluster_tool,
cluster_exec=default_cluster_exec,
out_file=None,
out_args=default_out_args,
nproc=None):
"""
Performs clustering on sets of sequences
Arguments:
seq_file : the sample sequence file name.
ident : the identity threshold for clustering sequences.
length_ratio : minimum short/long length ratio allowed within a cluster.
barcode_field : the annotation field containing barcode sequences.
cluster_field : the name of the output cluster field.
cluster_prefix : string defining a prefix for the cluster identifier.
seq_start : the start position to trim sequences at before clustering.
seq_end : the end position to trim sequences at before clustering.
cluster_tool : the clustering tool to use; one of cd-hit or usearch.
cluster_exec : the path to the executable for usearch.
out_file : output file name. Automatically generated from the input file if None.
out_args : output arguments.
nproc : the number of processQueue processes;
if None defaults to the number of CPUs.
Returns:
str: the clustered output file name
"""
# Function to modify SeqRecord header with cluster identifier
def _header(seq,
cluster,
field=cluster_field,
prefix=cluster_prefix,
delimiter=out_args['delimiter']):
label = '%s%i' % (prefix, cluster)
header = parseAnnotation(seq.description, delimiter=delimiter)
header = mergeAnnotation(header, {field: label}, delimiter=delimiter)
seq.id = seq.name = flattenAnnotation(header, delimiter=delimiter)
seq.description = ''
return seq
# Function to extract to make SeqRecord object from a barcode annotation
def _barcode(seq, field=barcode_field, delimiter=out_args['delimiter']):
header = parseAnnotation(seq.description, delimiter=delimiter)
return SeqRecord(Seq(header[field]), id=seq.id)
# Print parameter info
log = OrderedDict()
log['START'] = 'ClusterSets'
log['COMMAND'] = 'barcode'
log['FILE'] = os.path.basename(seq_file)
log['IDENTITY'] = ident
log['BARCODE_FIELD'] = barcode_field
log['CLUSTER_FIELD'] = cluster_field
log['CLUSTER_PREFIX'] = cluster_prefix
log['CLUSTER_TOOL'] = cluster_tool
log['NPROC'] = nproc
printLog(log)
# Set cluster tool
try:
cluster_func = map_cluster_tool.get(cluster_tool)
except:
printError('Invalid clustering tool %s.' % cluster_tool)
# Check the minimum identity
if ident < min_cluster_ident[cluster_tool]:
printError('Minimum identity %s too low for clustering tool %s.' %
(str(ident), cluster_tool))
# Count sequence file and parse into a list of SeqRecords
result_count = countSeqFile(seq_file)
barcode_iter = (_barcode(x) for x in readSeqFile(seq_file))
# Perform clustering
start_time = time()
printMessage('Running %s' % cluster_tool, start_time=start_time, width=25)
cluster_dict = cluster_func(barcode_iter,
ident=ident,
length_ratio=length_ratio,
seq_start=0,
seq_end=None,
threads=nproc,
cluster_exec=cluster_exec)
printMessage('Done', start_time=start_time, end=True, width=25)
# Determine file type
if out_args['out_type'] is None:
out_args['out_type'] = getFileType(seq_file)
# Open output file handles
if out_file is not None:
pass_handle = open(out_file, 'w')
else:
pass_handle = getOutputHandle(seq_file,
'cluster-pass',
out_dir=out_args['out_dir'],
out_name=out_args['out_name'],
out_type=out_args['out_type'])
# Open indexed sequence file
seq_dict = readSeqFile(seq_file, index=True)
# Iterate over sequence records and update header with cluster annotation
start_time = time()
rec_count = pass_count = 0
for cluster, id_list in cluster_dict.items():
printProgress(rec_count, result_count, 0.05, start_time=start_time)
rec_count += len(id_list)
# TODO: make a generator. Figure out how to get pass_count updated
# Define output sequences
seq_output = [_header(seq_dict[x], cluster) for x in id_list]
# Write output
pass_count += len(seq_output)
SeqIO.write(seq_output, pass_handle, out_args['out_type'])
# Update progress
printProgress(rec_count, result_count, 0.05, start_time=start_time)
# Print log
log = OrderedDict()
log['OUTPUT'] = os.path.basename(pass_handle.name)
log['CLUSTERS'] = len(cluster_dict)
log['SEQUENCES'] = result_count
log['PASS'] = pass_count
log['FAIL'] = rec_count - pass_count
log['END'] = 'ClusterSets'
printLog(log)
# Close handles
pass_handle.close()
return pass_handle.name
def parseIHMM(aligner_file, seq_file, repo, cellranger_file=None, partial=False, asis_id=True,
extended=False, format=default_format, out_file=None, out_args=default_out_args):
"""
Main for iHMMuneAlign aligned sample sequences.
Arguments:
aligner_file : iHMMune-Align output file to process.
seq_file : fasta file input to iHMMuneAlign (from which to get sequence).
repo : folder with germline repertoire files.
partial : If True put incomplete alignments in the pass file.
asis_id : if ID is to be parsed for pRESTO output with default delimiters.
extended : if True parse alignment scores, FWR and CDR region fields.
format : output format. One of 'changeo' or 'airr'.
out_file : output file name. Automatically generated from the input file if None.
out_args : common output argument dictionary from parseCommonArgs.
Returns:
dict : names of the 'pass' and 'fail' output files.
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'MakeDB'
log['COMMAND'] = 'ihmm'
log['ALIGNER_FILE'] = os.path.basename(aligner_file)
log['SEQ_FILE'] = os.path.basename(seq_file)
log['ASIS_ID'] = asis_id
log['PARTIAL'] = partial
log['EXTENDED'] = extended
printLog(log)
start_time = time()
printMessage('Loading files', start_time=start_time, width=20)
# Count records in sequence file
total_count = countSeqFile(seq_file)
# Get input sequence dictionary
seq_dict = getSeqDict(seq_file)
# Create germline repo dictionary
references = readGermlines(repo)
# Load supplementary annotation table
if cellranger_file is not None:
f = cellranger_extended if extended else cellranger_base
annotations = readCellRanger(cellranger_file, fields=f)
else:
annotations = None
printMessage('Done', start_time=start_time, end=True, width=20)
# Check for IMGT-gaps in germlines
if all('...' not in x for x in references.values()):
printWarning('Germline reference sequences do not appear to contain IMGT-numbering spacers. Results may be incorrect.')
# Define format operators
try:
__, writer, schema = getFormatOperators(format)
except ValueError:
printError('Invalid format %s.' % format)
out_args['out_type'] = schema.out_type
# Define output fields
fields = list(schema.required)
if extended:
custom = IHMMuneReader.customFields(scores=True, regions=True, schema=schema)
fields.extend(custom)
# Parse and write output
with open(aligner_file, 'r') as f:
parse_iter = IHMMuneReader(f, seq_dict, references)
germ_iter = (addGermline(x, references) for x in parse_iter)
output = writeDb(germ_iter, fields=fields, aligner_file=aligner_file, total_count=total_count,
annotations=annotations, asis_id=asis_id, partial=partial,
writer=writer, out_file=out_file, out_args=out_args)
return output
def parseIgBLAST(aligner_file, seq_file, repo, amino_acid=False, cellranger_file=None, partial=False,
asis_id=True, asis_calls=False, extended=False, regions='default',
format='changeo', out_file=None, out_args=default_out_args):
"""
Main for IgBLAST aligned sample sequences.
Arguments:
aligner_file (str): IgBLAST output file to process.
seq_file (str): fasta file input to IgBlast (from which to get sequence).
repo (str): folder with germline repertoire files.
amino_acid (bool): if True then the IgBLAST output files are results from igblastp. igblastn is assumed if False.
partial : If True put incomplete alignments in the pass file.
asis_id (bool): if ID is to be parsed for pRESTO output with default delimiters.
asis_calls (bool): if True do not parse gene calls for allele names.
extended (bool): if True add alignment scores, FWR regions, and CDR regions to the output.
regions (str): name of the IMGT FWR/CDR region definitions to use.
format (str): output format. one of 'changeo' or 'airr'.
out_file (str): output file name. Automatically generated from the input file if None.
out_args (dict): common output argument dictionary from parseCommonArgs.
Returns:
dict : names of the 'pass' and 'fail' output files.
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'MakeDB'
log['COMMAND'] = 'igblast-aa' if amino_acid else 'igblast'
log['ALIGNER_FILE'] = os.path.basename(aligner_file)
log['SEQ_FILE'] = os.path.basename(seq_file)
log['ASIS_ID'] = asis_id
log['ASIS_CALLS'] = asis_calls
log['PARTIAL'] = partial
log['EXTENDED'] = extended
printLog(log)
# Set amino acid conditions
if amino_acid:
format = '%s-aa' % format
parser = IgBLASTReaderAA
else:
parser = IgBLASTReader
# Start
start_time = time()
printMessage('Loading files', start_time=start_time, width=20)
# Count records in sequence file
total_count = countSeqFile(seq_file)
# Get input sequence dictionary
seq_dict = getSeqDict(seq_file)
# Create germline repo dictionary
references = readGermlines(repo, asis=asis_calls)
# Load supplementary annotation table
if cellranger_file is not None:
f = cellranger_extended if extended else cellranger_base
annotations = readCellRanger(cellranger_file, fields=f)
else:
annotations = None
printMessage('Done', start_time=start_time, end=True, width=20)
# Check for IMGT-gaps in germlines
if all('...' not in x for x in references.values()):
printWarning('Germline reference sequences do not appear to contain IMGT-numbering spacers. Results may be incorrect.')
# Define format operators
try:
__, writer, schema = getFormatOperators(format)
except ValueError:
printError('Invalid format %s.' % format)
out_args['out_type'] = schema.out_type
# Define output fields
fields = list(schema.required)
if extended:
custom = parser.customFields(schema=schema)
fields.extend(custom)
# Parse and write output
with open(aligner_file, 'r') as f:
parse_iter = parser(f, seq_dict, references, regions=regions, asis_calls=asis_calls)
germ_iter = (addGermline(x, references, amino_acid=amino_acid) for x in parse_iter)
output = writeDb(germ_iter, fields=fields, aligner_file=aligner_file, total_count=total_count,
annotations=annotations, amino_acid=amino_acid, partial=partial, asis_id=asis_id,
regions=regions, writer=writer, out_file=out_file, out_args=out_args)
return output
def parseIMGT(aligner_file, seq_file=None, repo=None, cellranger_file=None, partial=False, asis_id=True,
extended=False, format=default_format, out_file=None, out_args=default_out_args):
"""
Main for IMGT aligned sample sequences.
Arguments:
aligner_file : zipped file or unzipped folder output by IMGT.
seq_file : FASTA file input to IMGT (from which to get seqID).
repo : folder with germline repertoire files.
partial : If True put incomplete alignments in the pass file.
asis_id : if ID is to be parsed for pRESTO output with default delimiters.
extended : if True add alignment score, FWR, CDR and junction fields to output file.
format : output format. one of 'changeo' or 'airr'.
out_file : output file name. Automatically generated from the input file if None.
out_args : common output argument dictionary from parseCommonArgs.
Returns:
dict : names of the 'pass' and 'fail' output files.
"""
# Print parameter info
log = OrderedDict()
log['START'] = 'MakeDb'
log['COMMAND'] = 'imgt'
log['ALIGNER_FILE'] = aligner_file
log['SEQ_FILE'] = os.path.basename(seq_file) if seq_file else ''
log['ASIS_ID'] = asis_id
log['PARTIAL'] = partial
log['EXTENDED'] = extended
printLog(log)
start_time = time()
printMessage('Loading files', start_time=start_time, width=20)
# Extract IMGT files
temp_dir, imgt_files = extractIMGT(aligner_file)
# Count records in IMGT files
total_count = countDbFile(imgt_files['summary'])
# Get (parsed) IDs from fasta file submitted to IMGT
id_dict = getIDforIMGT(seq_file) if seq_file else {}
# Load supplementary annotation table
if cellranger_file is not None:
f = cellranger_extended if extended else cellranger_base
annotations = readCellRanger(cellranger_file, fields=f)
else:
annotations = None
printMessage('Done', start_time=start_time, end=True, width=20)
# Define format operators
try:
__, writer, schema = getFormatOperators(format)
except ValueError:
printError('Invalid format %s.' % format)
out_args['out_type'] = schema.out_type
# Define output fields
fields = list(schema.required)
if extended:
custom = IMGTReader.customFields(scores=True, regions=True, junction=True, schema=schema)
fields.extend(custom)
# Parse IMGT output and write db
with open(imgt_files['summary'], 'r') as summary_handle, \
open(imgt_files['gapped'], 'r') as gapped_handle, \
open(imgt_files['ntseq'], 'r') as ntseq_handle, \
open(imgt_files['junction'], 'r') as junction_handle:
# Open parser
parse_iter = IMGTReader(summary_handle, gapped_handle, ntseq_handle, junction_handle)
# Add germline sequence
if repo is None:
germ_iter = parse_iter
else:
references = readGermlines(repo)
# Check for IMGT-gaps in germlines
if all('...' not in x for x in references.values()):
printWarning('Germline reference sequences do not appear to contain IMGT-numbering spacers. Results may be incorrect.')
germ_iter = (addGermline(x, references) for x in parse_iter)
# Write db
output = writeDb(germ_iter, fields=fields, aligner_file=aligner_file, total_count=total_count,
annotations=annotations, id_dict=id_dict, asis_id=asis_id, partial=partial,
writer=writer, out_file=out_file, out_args=out_args)
# Cleanup temp directory
temp_dir.cleanup()
return output
def pairSeq(seq_file_1,
seq_file_2,
fields_1=None,
fields_2=None,
action=None,
coord_type=default_coord,
out_args=default_out_args):
"""
Syncronized paired end files and copies annotations between them
Arguments:
seq_file_1 : the file containing the grouped sequences and annotations.
seq_file_2 : the file to assign annotations to from seq_file_1.
fields_1 : list of annotations in seq_file_1 records to copy to seq_file_2 records;
if None do not copy any annotations.
fields_2 : list of annotations in seq_file_2 records to copy to seq_file_1 records;
if None do not copy any annotations.
action : the collapse action to take on all copied annotation if they already exist in the
target header.
coord_type : the sequence header format.
out_args : common output argument dictionary from parseCommonArgs.
Returns:
list: a list of tuples holding successfully paired filenames for (seq_file_1, seq_file_2).
"""
# Define private functions
def _key_func(x):
return getCoordKey(x,
coord_type=coord_type,
delimiter=out_args['delimiter'])
log = OrderedDict()
log['START'] = 'PairSeq'
log['FILE1'] = os.path.basename(seq_file_1)
log['FILE2'] = os.path.basename(seq_file_2)
log['FIELDS_1'] = ','.join(fields_1) if fields_1 is not None else None
log['FIELDS_2'] = ','.join(fields_2) if fields_2 is not None else None
log['COORD_TYPE'] = coord_type
printLog(log)
# Define output type
if out_args['out_type'] is None:
out_type_1 = getFileType(seq_file_1)
out_type_2 = getFileType(seq_file_2)
else:
out_type_1 = out_type_2 = out_args['out_type']
# Define output name
if out_args['out_name'] is None:
out_name_1 = out_name_2 = None
else:
out_name_1 = '%s-1' % out_args['out_name']
out_name_2 = '%s-2' % out_args['out_name']
# Open and count files
start_time = time()
printMessage("Indexing files", start_time=start_time)
# Index file 1
seq_count_1 = countSeqFile(seq_file_1)
seq_dict_1 = readSeqFile(seq_file_1, index=True, key_func=_key_func)
# Define file 2 iterator
seq_count_2 = countSeqFile(seq_file_2)
seq_iter_2 = readSeqFile(seq_file_2, index=False)
printMessage("Done", start_time=start_time, end=True)
# Open output file handles
pass_handle_1 = getOutputHandle(seq_file_1,
'pair-pass',
out_args['out_dir'],
out_name=out_name_1,
out_type=out_type_1)
pass_handle_2 = getOutputHandle(seq_file_2,
'pair-pass',
out_args['out_dir'],
out_name=out_name_2,
out_type=out_type_2)
if out_args['failed']:
fail_handle_1 = getOutputHandle(seq_file_1,
'pair-fail',
out_dir=out_args['out_dir'],
out_name=out_name_1,
out_type=out_type_1)
fail_handle_2 = getOutputHandle(seq_file_2,
'pair-fail',
out_dir=out_args['out_dir'],
out_name=out_name_2,
out_type=out_type_2)
pass_keys = list()
# Iterate over pairs and write to output files
start_time = time()
rec_count = pair_count = 0
for seq_2 in seq_iter_2:
# Print progress for previous iteration
printProgress(rec_count, seq_count_2, 0.05, start_time=start_time)
rec_count += 1
# Check for file 2 mate pair in file 1
coord_2 = getCoordKey(seq_2.id,
coord_type=coord_type,
delimiter=out_args['delimiter'])
seq_1 = seq_dict_1.get(coord_2, None)
if seq_1 is not None:
# Record paired keys
pair_count += 1
if fields_1 is not None or fields_2 is not None:
ann_1 = parseAnnotation(seq_1.description,
delimiter=out_args['delimiter'])
ann_2 = parseAnnotation(seq_2.description,
delimiter=out_args['delimiter'])
# Prepend annotations from seq_1 to seq_2
if fields_1 is not None:
copy_ann = OrderedDict([(k, v) for k, v in ann_1.items() \
if k in fields_1])
merge_ann = mergeAnnotation(
ann_2,
copy_ann,
prepend=True,
delimiter=out_args['delimiter'])
# Collapse if necessary
if action is not None:
merge_ann = collapseAnnotation(
merge_ann,
action,
fields=fields_1,
delimiter=out_args['delimiter'])
# Flatten
seq_2.id = flattenAnnotation(
merge_ann, delimiter=out_args['delimiter'])
seq_2.description = ''
# Append annotations from seq_2 to seq_1
if fields_2 is not None:
copy_ann = OrderedDict([(k, v) for k, v in ann_2.items() \
if k in fields_2])
merge_ann = mergeAnnotation(
ann_1,
copy_ann,
prepend=False,
delimiter=out_args['delimiter'])
# Collapse if necessary
if action is not None:
merge_ann = collapseAnnotation(
merge_ann,
action,
fields=fields_2,
delimiter=out_args['delimiter'])
# Flatten
seq_1.id = flattenAnnotation(
merge_ann, delimiter=out_args['delimiter'])
seq_1.description = ''
# Write paired records
SeqIO.write(seq_1, pass_handle_1, out_type_1)
SeqIO.write(seq_2, pass_handle_2, out_type_2)
# Write unpaired file 2 records and updated paired key list for finding unpaired file 1 records
if out_args['failed']:
if seq_1 is not None: pass_keys.append(coord_2)
else: SeqIO.write(seq_2, fail_handle_2, out_type_2)
# Print final progress
printProgress(rec_count, seq_count_2, 0.05, start_time=start_time)
# Find and write unpaired file 1 records
if out_args['failed']:
start_time = time()
printMessage("Finding unpaired", start_time=start_time)
# Find file 1 unpaired keys
pass_keys = set(pass_keys)
unpaired = set(seq_dict_1).difference(pass_keys)
# Write unpaired file 1 records
for k in unpaired:
SeqIO.write(seq_dict_1[k], fail_handle_1, out_type_1)
printMessage("Done", start_time=start_time, end=True)
# Print log
log = OrderedDict()
log['OUTPUT1'] = os.path.basename(pass_handle_1.name)
log['OUTPUT2'] = os.path.basename(pass_handle_2.name)
log['SEQUENCES1'] = seq_count_1
log['SEQUENCES2'] = seq_count_2
log['PASS'] = pair_count
log['END'] = 'PairSeq'
printLog(log)
# Close file handles
pass_handle_1.close()
pass_handle_2.close()
return [(pass_handle_1.name, pass_handle_2.name)]