def calc_sample_coverage(args):
"""Counts the total number of non-gap/ambiguous characters for
each sample per contig.
"""
mvf = MultiVariantFile(args.mvf, 'read')
data = {}
# data_order = []
# Set up sample indices
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
sample_labels = mvf.get_sample_ids(indices=sample_indices)
# Set up contig ids
if args.contig_ids is not None:
contig_indices = mvf.get_contig_indices(args.contig_ids[0].split(","))
elif args.contig_labels is not None:
contig_indices = mvf.get_contig_indices(
labels=args.contig_labels[0].split(","))
else:
contig_indices = None
for contig, _, allelesets in mvf.iterentries(contig_indices=contig_indices,
subset=sample_indices,
decode=True):
if contig not in data:
data[contig] = dict((x, 0) for x in sample_labels)
data[contig]['contig'] = contig
for j, elem in enumerate(sample_indices):
data[contig][sample_labels[elem]] += int(
allelesets[0][j] not in 'Xx-')
outfile = OutputFile(path=args.out,
headers=(["contig"] +
[sample_labels[x] for x in sample_indices]))
for contig in data:
outfile.write_entry(data[contig])
return ''
def mvf2phy(args):
"""Main method"""
mvf = MultiVariantFile(args.mvf, 'read')
if (mvf.flavor in ("dna", "rna") and args.output_data == "prot") or (
mvf.flavor == "prot" and args.output_data in ("dna", "rna")):
raise RuntimeError(
"--outdput-data {} incompatiable with '{}' flavor mvf".format(
args.output_data, mvf.flavor))
max_region_coord = dict((x, None) for x in mvf.get_contig_ids())
if args.regions is not None:
_, max_region_coord, _ = parse_regions_arg(args.regions,
mvf.get_contig_ids())
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
sample_labels = mvf.get_sample_ids(indices=sample_indices)
skipcontig = ''
tmp_files = dict((fn,
open("{}-{}.tmp".format(fn, randint(1000000, 9999999)),
'w+', args.buffer)) for fn in sample_labels)
labelwritten = dict.fromkeys(sample_labels, False)
current_contig_id = None
current_contig_start = 1
current_contig_end = 1
if args.partition is True:
partprefix = "PROT" if args.output_data == "prot" else "DNA"
partitionfile = open("{}.part".format(args.out), 'w')
for contig, _, allelesets in mvf.iterentries(
contig_ids=(mvf.get_contig_ids()
if args.regions is None else max_region_coord[:]),
decode=True):
if contig == skipcontig:
continue
if contig not in max_region_coord:
skipcontig = contig[:]
continue
if current_contig_id is None:
current_contig_id = contig[:]
elif contig != current_contig_id:
if args.partition is True:
if current_contig_end > current_contig_start:
partitionfile.write("{}, {} = {}-{}\n".format(
partprefix,
mvf.get_contig_labels(ids=current_contig_id),
current_contig_start, current_contig_end - 1))
current_contig_id = contig[:]
# reset start as one position after end of last
current_contig_start = current_contig_end
current_contig_end = current_contig_end + 1
for col, label in zip(sample_indices, sample_labels):
if not labelwritten[label]:
if args.label_type == 'long':
tmp_files[label].write("{}{}".format(
label[:100], " " * (100 - len(label[:100]))))
elif args.label_type == 'short':
tmp_files[label].write("{}{}".format(
label[:20], " " * (20 - len(label[:20]))))
labelwritten[label] = True
if mvf.flavor == 'dna':
tmp_files[label].write(allelesets[0][col] == 'X' and 'N'
or allelesets[0][col])
if label == sample_labels[0]:
current_contig_end += 1
elif ((mvf.flavor == 'codon' and args.output_data == 'prot')
or (mvf.flavor == 'prot')):
tmp_files[label].write(allelesets[0][col])
if label == sample_labels[0]:
current_contig_end += 1
elif mvf.flavor == 'codon':
codon = [
"N" if allelesets[x][col] == 'X' else allelesets[x][col]
for x in (1, 2, 3)
]
tmp_files[label].write(''.join(codon))
if label == sample_labels[0]:
current_contig_end += 3
first_file = True
totalseqlen = 0
with open(args.out, 'w') as outfile:
for filehandler in tmp_files.values():
# read first file to establish sequence length for phylip header
if first_file is True:
filehandler.seek(0, 0)
buff = filehandler.read(args.buffer)
while buff != '':
if " " in buff:
totalseqlen += len(buff.strip().split(" ")[-1])
else:
totalseqlen += len(buff.strip())
buff = filehandler.read(args.buffer)
outfile.write("{} {}\n".format(len(sample_labels),
totalseqlen))
first_file = False
filehandler.seek(0, 0)
buff = filehandler.read(args.buffer)
while buff != '':
if first_file is True:
outfile.write("{} {}\n".format(len(sample_labels),
len(buff.split()[1])))
first_file = False
outfile.write(buff)
buff = filehandler.read(args.buffer)
outfile.write("\n")
filehandler.close()
os.remove(os.path.join(args.temp_dir, filehandler.name))
if args.partition is True:
if current_contig_end > current_contig_start:
partitionfile.write("{}, {} = {}-{}\n".format(
partprefix, mvf.get_contig_labels(ids=current_contig_id),
current_contig_start, current_contig_end - 1))
partitionfile.close()
return ''
def mvf2fastagene(args):
"""Main method"""
args.qprint("Indexing MVF")
mvf = MultiVariantFile(args.mvf, 'read', contigindex=True)
if (mvf.flavor in ("dna", "rna") and args.output_data == "prot") or (
mvf.flavor == "prot" and args.output_data in ("dna", "rna")):
raise RuntimeError(
"--output-data {} incompatiable with '{}' flavor mvf".format(
args.output_data, mvf.flavor))
if args.output_data is None:
raise RuntimeError("--output-data required")
sample_labels = mvf.get_sample_ids()
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
args.qprint("Beginning Entries.")
if not os.path.exists(args.output_dir):
args.qprint("Output Directory Created: {}".format(args.output_dir))
os.mkdir(args.output_dir)
else:
args.qprint("Output Directory Exists Already: {}".format(
args.output_dir))
write_buffer = {}
for targetcontig in mvf.get_contig_indices():
contiglabel = mvf.get_contig_labels(indices=targetcontig)[0]
args.qprint("Reading Contig {}: {}".format(targetcontig, contiglabel))
write_buffer = dict((x, []) for x in sample_labels)
data_in_buffer = False
for _, _, allelesets in mvf.itercontigentries(targetcontig,
decode=True):
for col, label in zip(sample_indices, sample_labels):
if mvf.flavor == 'dna':
write_buffer[label].append('N' if allelesets[0][col] ==
'X' else allelesets[0][col])
data_in_buffer = True
elif mvf.flavor in ('codon', 'prot') and (args.output_data
== 'prot'):
write_buffer[label].append(allelesets[0][col])
data_in_buffer = True
elif mvf.flavor == 'codon' and args.output_data == 'dna':
if args.choose_allele == 'random1':
codon = [
'N' if allelesets[x][col] == 'X' else
(MLIB.randomnuc(allelesets[x][col]) if
(allelesets[x][col]
in MLIB.validchars['dnaambig23']) else
allelesets[x][col]) for x in (1, 2, 3)
]
else:
codon = [
'N' if allelesets[x][col] == 'X' else
allelesets[x][col] for x in (1, 2, 3)
]
write_buffer[label].append(''.join(codon))
data_in_buffer = True
if data_in_buffer:
args.qprint("Writing Align")
with open(os.path.join(args.output_dir, contiglabel + ".fa"),
'w') as outfile:
for label in write_buffer:
if (mvf.flavor == 'codon'
and args.output_data in ('dna', 'prot')):
if ((mvf.contig_data[targetcontig].get('strand', '+')
== '-') and (args.ignore_strand is False)):
entryseq = ''.join(write_buffer[label][::-1])
else:
entryseq = ''.join(write_buffer[label])
else:
entryseq = ''.join(write_buffer[label])
outfile.write(">{}\n{}\n".format(label, entryseq))
outfile.write("\b")
return ''
def mvf2fasta(args):
"""Main method"""
mvf = MultiVariantFile(args.mvf, 'read')
if (mvf.flavor in ("dna", "rna") and args.output_data == "prot") or (
mvf.flavor == "prot" and args.output_data in ("dna", "rna")):
raise RuntimeError(
"--output-data {} incompatiable with '{}' flavor mvf".format(
args.output_data, mvf.flavor))
regions, max_region_coord, regionlabel = parse_regions_arg(
args.regions, mvf.contig_data)
sample_labels = mvf.get_sample_ids()
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
skipcontig = None
tmp_files = dict(
(fname, tempfile.NamedTemporaryFile(mode='w+', prefix=fname))
for fname in sample_labels)
labelwritten = dict.fromkeys(sample_labels, False)
write_buffer = {}
current_contig = None
data_written = False
args.qprint("Regions determined. Reading entries.")
for contig, pos, allelesets in mvf.iterentries(
contig_indices=mvf.get_contig_indices(
ids=list(max_region_coord.keys())),
decode=True):
if current_contig is None:
current_contig = mvf.get_contig_indices(ids=contig)
if contig == skipcontig:
continue
if (contig not in max_region_coord) or (
max_region_coord[contig] is not None
and pos > max_region_coord[contig]):
skipcontig = contig[:]
continue
inregion = False
for rcontig, rstart, rstop, _ in regions[contig]:
if contig == rcontig:
if rstart is None or pos >= rstart:
if rstop is None or pos <= rstop:
inregion = True
break
if inregion is False:
continue
for col, label in zip(sample_indices, sample_labels):
if not labelwritten[label]:
if args.label_type == 'long':
xlabel = "{} region={}".format(label, regionlabel)
elif args.label_type == 'short':
xlabel = "{}".format(label)
tmp_files[label].write(">{}\n".format(xlabel))
labelwritten[label] = True
if mvf.flavor == 'dna':
tmp_files[label].write("N" if allelesets[0][col] ==
'X' else allelesets[0][col])
data_written = True
elif mvf.flavor in ('codon', 'prot') and (args.output_data
== 'prot'):
tmp_files[label].write(allelesets[0][col])
data_written = True
elif mvf.flavor == 'codon' and args.output_data == 'dna':
codon = [
"N" if allelesets[x][col] == 'X' else allelesets[x][col]
for x in (1, 2, 3)
]
if not args.gene_mode:
tmp_files[label].write(''.join(codon))
data_written = True
else:
if contig != current_contig:
if mvf.metadata['contigs'][current_contig].get(
'strand', "+") == '-':
write_buffer[label] = write_buffer[label][::-1]
tmp_files[label].write(''.join(write_buffer[label]))
data_written = True
if label not in write_buffer:
write_buffer[label] = []
write_buffer[label].append(''.join(codon))
if args.gene_mode and current_contig != contig:
write_buffer = {}
current_contig = contig[:]
if write_buffer:
for label in write_buffer:
if mvf.metadata['contigs'][current_contig].get('strand',
"+") == '-':
write_buffer[label] = write_buffer[label][::-1]
tmp_files[label].write(''.join(write_buffer[label]))
data_written = True
write_buffer = {}
if data_written is False:
print("ERROR NO DATA WRITTEN")
with open(args.out, 'w') as outfile:
for filehandler in tmp_files.values():
filehandler.seek(0, 0)
buff = filehandler.read(args.buffer)
while buff:
outfile.write(buff)
buff = filehandler.read(args.buffer)
outfile.write("\n")
filehandler.close()
return ''
def calc_all_character_count_per_sample(args):
"""Count the number of and relative rate of certain bases
spatially along chromosomes
"""
args.qprint("Running CalcAllCharacterCountPerSample")
mvf = MultiVariantFile(args.mvf, 'read')
current_contig = None
current_position = 0
data_in_buffer = False
# Set up sample indices
sample_labels = mvf.get_sample_ids()
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
# Set up contig ids
if args.contig_ids is not None:
contig_ids = args.contig_ids[0].split(",")
elif args.contig_labels is not None:
contig_ids = mvf.get_contig_ids(
labels=args.contig_labels[0].split(","))
else:
contig_ids = None
data = dict((i, {}) for i in sample_indices)
data_characters = [{} for i in sample_indices]
for contig, pos, allelesets in mvf.iterentries(decode=False,
contig_ids=contig_ids):
# Check Minimum Site Coverage
if check_mincoverage(args.mincoverage, allelesets[0]) is False:
continue
if current_contig is None:
current_contig = contig[:]
if args.windowsize > 0:
while pos > current_position + args.windowsize - 1:
current_position += args.windowsize
# Check if windows are specified.
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
args.qprint("Processing contig {}".format(current_contig))
for i in sample_indices:
data[i][(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data[i][(current_contig,
current_position)].update(data_characters[i])
if contig != current_contig:
current_contig = contig[:]
current_position = 0
else:
current_position += (0 if args.windowsize == -1 else
args.windowsize)
data_characters = [{} for i in sample_indices]
data_in_buffer = False
alleles = allelesets[0]
if len(alleles) == 1:
for i in sample_indices:
data_characters[i][alleles[0]] = (
data_characters[i].get(alleles[0], 0) + 1)
else:
alleles = mvf.decode(alleles)
for i in sample_indices:
data_characters[i][alleles[i]] = (
data_characters[i].get(alleles[i], 0) + 1)
data_in_buffer = True
if data_in_buffer:
for i in sample_indices:
data[i][(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data[i][(current_contig,
current_position)].update(data_characters[i])
# WRITE OUTPUT
all_chars = set([])
for sampleid in data:
for window in data[sampleid]:
all_chars.update([
x for x in data[sampleid][window]
if x not in ('contig', 'position')
])
headers = ['contig', 'position']
headers.extend(list(sorted(all_chars)))
outfile = OutputFile(path=args.out, headers=headers)
for sampleid in sample_indices:
outfile.write("#{}\n".format(sample_labels[sampleid]))
sorted_entries = [(data[sampleid][k]['contig'],
data[sampleid][k]['position'], k)
for k in data[sampleid]]
for _, _, k in sorted_entries:
outfile.write_entry(data[sampleid][k], defaultvalue='0')
return ''
def calc_pairwise_distances(args):
"""Count the pairwise nucleotide distance between
combinations of samples in a window
"""
args.qprint("Running CalcPairwiseDistances")
mvf = MultiVariantFile(args.mvf, 'read')
args.qprint("Input MVF: Read")
data = {}
data_order = []
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
sample_labels = mvf.get_sample_ids(indices=sample_indices)
args.qprint("Calculating for sample columns: {}".format(
list(sample_indices)))
current_contig = None
current_position = 0
data_in_buffer = False
sample_pairs = [tuple(x) for x in combinations(sample_indices, 2)]
base_matches = dict((x, {}) for x in sample_pairs)
all_match = {}
if mvf.flavor == 'dna':
allele_frames = (0, )
args.data_type = 'dna'
elif mvf.flavor == 'prot':
allele_frames = (0, )
args.data_type = 'dna'
elif mvf.flavor == 'codon':
if args.data_type == 'prot':
allele_frames = (0, )
else:
allele_frames = (1, 2, 3)
args.data_type = 'dna'
args.qprint("MVF flavor is: {}".format(mvf.flavor))
args.qprint("Data type is: {}".format(args.data_type))
args.qprint("Ambiguous mode: {}".format(args.ambig))
args.qprint("Processing MVF Records")
pwdistance_function = get_pairwise_function(args.data_type, args.ambig)
if args.emit_counts:
outfile_emitcounts = open(args.out + ".pairwisecounts", 'w')
for contig, pos, allelesets in mvf.iterentries(decode=None):
# Check Minimum Site Coverage
if check_mincoverage(args.mincoverage, allelesets[0]) is False:
continue
# Establish first contig
if current_contig is None:
current_contig = contig[:]
if args.windowsize > 0:
while pos > current_position + args.windowsize - 1:
current_position += args.windowsize
# Check if windows are specified.
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
data[(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data_order.append((current_contig, current_position))
all_diff, all_total = pwdistance_function(all_match)
for samplepair in base_matches:
ndiff, ntotal = pwdistance_function(base_matches[samplepair])
taxa = "{};{}".format(sample_labels[samplepair[0]],
sample_labels[samplepair[1]])
data[(current_contig, current_position)].update({
'{};ndiff'.format(taxa):
ndiff + all_diff,
'{};ntotal'.format(taxa):
ntotal + all_total,
'{};dist'.format(taxa):
zerodiv(ndiff + all_diff, ntotal + all_total)
})
if contig != current_contig:
current_contig = contig[:]
current_position = 0
if args.windowsize > 0:
while pos > current_position + args.windowsize - 1:
current_position += args.windowsize
else:
current_position += args.windowsize
if args.emit_counts:
args.qprint("Writing Full Count Table")
for p0, p1 in base_matches:
outfile_emitcounts.write("#{}\t{}\t{}\t{}\n{}\n".format(
p0, p1, current_position, current_contig, "\n".join([
"{} {}".format(x,
(base_matches[(p0, p1)].get(x, 0) +
all_match.get(x, 0)))
for x in set(base_matches[(p0,
p1)]).union(all_match)
])))
base_matches = dict((x, {}) for x in sample_pairs)
all_match = {}
data_in_buffer = False
for iframe in allele_frames:
alleles = allelesets[iframe]
if len(alleles) == 1:
all_match["{0}{0}".format(alleles)] = (
all_match.get("{0}{0}".format(alleles), 0) + 1)
data_in_buffer = True
continue
if alleles[1] == '+':
if alleles[2] in 'X-':
continue
samplepair = (0, int(alleles[3:]))
if any(x not in sample_indices for x in samplepair):
continue
basepair = "{0}{1}".format(alleles[0], alleles[2])
base_matches[samplepair][basepair] = (
base_matches[samplepair].get(basepair, 0) + 1)
data_in_buffer = True
continue
alleles = mvf.decode(alleles)
valid_positions = [
i for i, x in enumerate(alleles)
if x not in 'X-' and i in sample_indices
]
assert len(alleles) == 4
assert alleles[0] not in 'X-', alleles
assert alleles[1] not in 'X-', alleles
for i, j in combinations(valid_positions, 2):
samplepair = (i, j)
basepair = "{0}{1}".format(alleles[i], alleles[j])
base_matches[samplepair][basepair] = (
base_matches[samplepair].get(basepair, 0) + 1)
data_in_buffer = True
# print(base_matches)
if data_in_buffer is True:
print(sum(base_matches[samplepair].values()), base_matches[samplepair],
samplepair)
print(sum(all_match.values()), all_match)
print(sum(base_matches[samplepair].values()) + sum(all_match.values()))
# Check whether, windows, contigs, or total
if args.windowsize == 0:
current_contig = 'TOTAL'
current_position = 0
elif args.windowsize == -1:
current_position = 0
data[(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data_order.append((current_contig, current_position))
# print("All match")
all_diff, all_total = pwdistance_function(all_match)
print(all_diff, all_total)
for samplepair in base_matches:
ndiff, ntotal = pwdistance_function(base_matches[samplepair])
taxa = "{};{}".format(sample_labels[samplepair[0]],
sample_labels[samplepair[1]])
data[(current_contig, current_position)].update({
'{};ndiff'.format(taxa):
ndiff + all_diff,
'{};ntotal'.format(taxa):
ntotal + all_total,
'{};dist'.format(taxa):
zerodiv(ndiff + all_diff, ntotal + all_total)
})
if args.emit_counts:
args.qprint("Writing Full Count Table")
for p0, p1 in base_matches:
outfile_emitcounts.write("#{}\t{}\t{}\t{}\n{}\n".format(
p0, p1, current_position, current_contig, "\n".join([
"{} {}".format(x, (base_matches[(p0, p1)].get(x, 0) +
all_match.get(x, 0)))
for x in set(base_matches[(p0, p1)]).union(all_match)
])))
args.qprint("Writing Output")
headers = ['contig', 'position']
for samplepair in sample_pairs:
headers.extend([
'{};{};{}'.format(sample_labels[samplepair[0]],
sample_labels[samplepair[1]], x)
for x in ('ndiff', 'ntotal', 'dist')
])
outfile = OutputFile(path=args.out, headers=headers)
for okey in data_order:
outfile.write_entry(data[okey])
if args.emit_counts:
outfile_emitcounts.close()
return ''
def calc_pattern_count(args):
"""Count biallelic patterns spatially along
chromosomes (e.g,, for use in DFOIL or Dstats
http://www.github.com/jbpease/dfoil).
The last sample specified will determine the 'A'
versus 'B' allele.
"""
mvf = MultiVariantFile(args.mvf, 'read')
data = {}
current_contig = None
current_position = 0
sitepatterns = {}
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
nsamples = len(sample_indices)
for contig, pos, allelesets in mvf.iterentries(decode=True,
subset=sample_indices):
alleles = allelesets[0]
# Check Minimum Site Coverage
if check_mincoverage(args.mincoverage, alleles) is False:
continue
# Establish first contig
if current_contig is None:
current_contig = contig[:]
if args.windowsize > 0:
while pos > current_position + args.windowsize - 1:
current_position += args.windowsize
# Check if windows are specified.
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
data[(current_contig,
current_position)] = dict([('contig', current_contig),
('position', current_position)])
data[(current_contig, current_position)].update(sitepatterns)
sitepatterns = {}
if contig != current_contig:
current_position = 0
current_contig = contig[:]
else:
current_position += (0 if args.windowsize == -1 else
args.windowsize)
if set(alleles) - set("ACGT"):
continue
if len(set(alleles)) > 2:
continue
pattern = ''.join(
['A' if x == alleles[-1] else 'B' for x in alleles[:-1]]) + 'A'
sitepatterns[pattern] = sitepatterns.get(pattern, 0) + 1
if sitepatterns:
data[(current_contig,
current_position)] = dict([('contig', current_contig),
('position', current_position)])
data[(current_contig, current_position)].update(sitepatterns)
# WRITE OUTPUT
headers = ['contig', 'position']
headers.extend(
[MLIB.abpattern(x, nsamples) for x in range(0, 2**nsamples, 2)])
outfile = OutputFile(path=args.out, headers=headers)
outfile.write("#{}\n".format(",".join(mvf.get_sample_ids(sample_indices))))
sorted_entries = sorted([(data[k]['contig'], data[k]['position'], k)
for k in data])
for _, _, k in sorted_entries:
outfile.write_entry(data[k])
# WRITE LIST OUTPUT
if args.output_lists is True:
sorted_entries = sorted([(data[k]['contig'], data[k]['position'], k)
for k in data])
total_counts = {}
for contig, pos, k in sorted_entries:
outfilepath = "{}-{}-{}.counts.list".format(args.out, contig, pos)
with open(outfilepath, 'w') as outfile:
outfile.write("pattern,count\n")
for pattern, pcount in sorted(data[k].items()):
if pattern in ['contig', 'position']:
continue
outfile.write("{},{}\n".format(pattern, pcount))
total_counts[pattern] = (total_counts.get(pattern, 0) +
pcount)
outfilepath = "{}-TOTAL.counts.list".format(args.out)
with open(outfilepath, 'w') as outfile:
outfile.write("pattern,count\n")
for pattern, pcount in sorted(total_counts.items()):
if pattern in ['contig', 'position']:
continue
outfile.write("{},{}\n".format(pattern, pcount))
return ''
def calc_character_count(args):
"""Count the number of and relative rate of certain bases
spatially along chromosomes
"""
mvf = MultiVariantFile(args.mvf, 'read')
data = {}
current_contig = None
current_position = 0
all_match = 0
all_total = 0
data_in_buffer = False
# Set up base matching from special words
data_order = []
def proc_special_word(argx):
if argx == 'dna':
argx = MLIB.validchars['dna']
elif argx == 'dnaambig2':
argx = MLIB.validchars['dna+ambig2']
elif argx == 'dnaambig3':
argx = MLIB.validchars['dna+ambig3']
elif argx == 'dnaambigall':
argx = MLIB.validchars['dna+ambigall']
elif argx == 'prot':
argx = MLIB.validchars['amino']
return argx
args.base_match = proc_special_word(args.base_match)
args.base_total = proc_special_word(args.base_total)
# Set up sample indices
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
sample_labels = mvf.get_sample_ids(indices=sample_indices)
# Set up contig ids
if args.contig_ids is not None:
contig_indices = mvf.get_contig_indices(
ids=args.contig_ids[0].split(","))
elif args.contig_labels is not None:
contig_indices = mvf.get_contig_indices(
labels=args.contig_labels[0].split(","))
else:
contig_indices = None
match_counts = dict().fromkeys([sample_labels[i] for i in sample_indices],
0)
total_counts = dict().fromkeys([sample_labels[i] for i in sample_indices],
0)
for contig, pos, allelesets in mvf.iterentries(
decode=False, contig_indices=contig_indices):
# Check Minimum Site Coverage
if check_mincoverage(args.mincoverage, allelesets[0]) is False:
continue
# if contig not in contig_ids:
# continue
# Establish first contig
if current_contig is None:
current_contig = contig[:]
if args.windowsize > 0:
while pos > current_position + args.windowsize - 1:
current_position += args.windowsize
# Check if windows are specified.
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
data[(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data_order.append((current_contig, current_position))
for k in match_counts:
data[(current_contig, current_position)].update([
(k + '.match', match_counts[k] + all_match),
(k + '.total', total_counts[k] + all_total),
(k + '.prop', ((float(match_counts[k] + all_match) /
float(total_counts[k] + all_total))
if total_counts[k] + all_total > 0 else 0))
])
if contig != current_contig:
current_contig = contig[:]
current_position = 0
else:
current_position += (0 if args.windowsize == -1 else
args.windowsize)
match_counts = dict().fromkeys(
[sample_labels[i] for i in sample_indices], 0)
total_counts = dict().fromkeys(
[sample_labels[i] for i in sample_indices], 0)
all_total = 0
all_match = 0
data_in_buffer = False
else:
alleles = allelesets[0]
if len(alleles) == 1:
if args.base_match is None:
all_match += 1
elif alleles in args.base_match:
all_match += 1
if args.base_total is None:
all_total += 1
elif alleles in args.base_total:
all_total += 1
else:
alleles = mvf.decode(alleles)
for i in sample_indices:
if args.base_match is None:
match_counts[sample_labels[i]] += 1
elif alleles[i] in args.base_match:
match_counts[sample_labels[i]] += 1
if args.base_total is None:
total_counts[sample_labels[i]] += 1
elif alleles[i] in args.base_total:
total_counts[sample_labels[i]] += 1
data_in_buffer = True
if data_in_buffer:
data[(current_contig, current_position)] = {
'contig': current_contig,
'position': current_position
}
data_order.append((current_contig, current_position))
for k in match_counts:
data[(current_contig, current_position)].update([
(k + '.match', match_counts[k] + all_match),
(k + '.total', total_counts[k] + all_total),
(k + '.prop', ((float(match_counts[k] + all_match) /
float(total_counts[k] + all_total))
if total_counts[k] + all_total > 0 else 0))
])
# WRITE OUTPUT
headers = ['contig', 'position']
for label in sample_labels:
headers.extend([label + x for x in ('.match', '.total', '.prop')])
outfile = OutputFile(path=args.out, headers=headers)
for okey in data_order:
outfile.write_entry(data[okey])
return ''
def calc_group_unique_allele_window(args):
"""Count the number of and relative rate of uniquely held alleles
spatially along chromosomes (i.e. Lineage-specific rates)"""
args.qprint("Running InferGroupSpecificAllele")
data = {}
mvf = MultiVariantFile(args.mvf, 'read')
if mvf.flavor != 'codon':
raise RuntimeError(
"\n=====================\nERROR: MVF is not codon flavor!")
ncol = mvf.metadata['ncol']
args.qprint("Input MVF read with {} columns.".format(ncol))
annotations = {}
coordinates = {}
labels = mvf.get_sample_ids()[:]
current_contig = None
current_position = 0
counts = Counter()
totals = Counter()
args.start_contig = (
args.start_contig if args.start_contig is not None else 0)
args.end_contig = (
args.end_contig if args.end_contig is not None else 100000000000)
if args.output_align is True:
outputalign = []
if args.gff is not None:
annotations, coordinates = (parse_gff_analysis(args.gff))
if args.allele_groups is not None:
args.allele_groups = procarg_allelegroups(
args.allele_groups, mvf)
if args.species_groups is None:
args.species_groups = args.allele_groups
else:
args.species_groups = procarg_speciesgroups(
args.species_groups, mvf)
fieldtags = [
'likelihood', 'bgdnds0', 'bgdnds1', 'bgdnds2a', 'bgdnds2b',
'fgdnds0', 'fgdnds1', 'fgdnds2a', 'fgdnds2b', 'dndstree',
'errorstate']
if args.branch_lrt is not None:
with open(args.branch_lrt, 'w') as branchlrt:
genealign = []
branchlrt.write("\t".join(
['contig', 'ntaxa', 'alignlength', 'lrtscore'] +
["null.{}".format(x) for x in fieldtags] +
["test.{}".format(x) for x in fieldtags] +
['tree']) + "\n")
groups = args.allele_groups.values()
if args.species_groups is not None:
speciesgroups = args.species_groups.values()
allsets = set([])
for group in groups:
allsets.update(group)
allsets = list(sorted(allsets))
speciesnames = args.species_groups.keys()
speciesrev = {}
if args.species_groups is not None:
for species in args.species_groups:
speciesrev.update([
(x, species) for x in args.species_groups[species]])
if args.mincoverage is not None:
if args.mincoverage < len(groups) * 2:
raise RuntimeError("""
Error: InferGroupSpecificAllele:
--mincoverage cannot be lower than the twice the number
of specified groups in --allele-groups
""")
genealign = []
args.qprint("Parameter Check Complete.")
args.qprint("Number of Groups Specified: {}".format(len(groups)))
for group in groups:
args.qprint(group)
args.qprint([labels[x] for x in group])
if not(group):
raise RuntimeError(
"Group is Empty! Check group labels/indicies specified.")
args.qprint("Processing Entries.")
for contig, pos, allelesets in mvf.iterentries(decode=False):
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
xkey = (current_contig, current_position,)
data[xkey] = counts.copy()
data[xkey].update([
('contig', (mvf.get_contig_labels(ids=current_contig)
if args.use_labels is True else
current_contig)),
('position', current_position),
('nonsynyonymous_changes',
counts.get('nonsynonymous_changes', 0) or 0),
('synyonymous_changes',
counts.get('synonymous_changes', 0) or 0)
])
data[xkey].update([
('ns_ratio', (float(data[xkey].get(
'nonsynonymous_changes', 0)) / (
data[xkey].get('synonymous_changes', 1.0)))),
('annotation',
annotations.get(data[xkey]['contig'], '.')),
('coordinates',
coordinates.get(data[xkey]['contig'], '.'))
])
if genealign:
if (args.end_contig >= int(current_contig)) and (
args.start_contig <= int(current_contig)):
(pamlnull, pamltest, tree) = paml_branchsite(
genealign, labels[:],
species=speciesnames,
speciesrev=speciesrev,
codemlpath=args.codeml_path,
raxmlpath=args.raxml_path,
pamltmp=args.paml_tmp,
target=args.target,
targetspec=args.num_target_species,
allsampletrees=args.all_sample_trees,
outgroup=args.outgroup)
lrtscore = -1
if (pamlnull.get('likelihood', -1) != -1 and
pamltest.get('likelihood', -1) != -1):
lrtscore = 2 * (pamltest['likelihood'] -
pamlnull['likelihood'])
with open(args.branch_lrt, 'a') as branchlrt:
branchlrt.write("\t".join([str(x) for x in [
data[xkey]['contig'],
len(genealign),
len(genealign[0]) * 3,
lrtscore] + [
pamlnull.get(y, -1) for y in fieldtags] + [
pamltest.get(y, -1) for y in fieldtags] + [
str(tree).rstrip()]]) + "\n")
genealign = None
totals.add('genes_total')
if counts.get('total_codons', 0) > 0:
totals.add('genes_tested')
if counts.get('total_nsyn_codons', 0) > 0:
totals.add('genes_with_nsyn')
if contig != current_contig:
current_contig = contig[:]
current_position = 0
elif args.windowsize > 0:
current_position += args.windowsize
counts = Counter()
proteins = allelesets[0]
codons = allelesets[1:4]
if len(proteins) == 1 and all(len(x) == 1 for x in codons):
if proteins == '*' or ''.join(codons) in MLIB.stop_codons:
continue
counts.add('total_codons')
totals.add('total_codons')
if args.output_align is True:
if not outputalign:
outputalign = [[''.join(codons)]
for x in range(mvf.metadata['ncol'])]
else:
for subalign in outputalign:
subalign.append(''.join(codons))
if args.branch_lrt is not None:
if not genealign:
genealign = [[''.join(codons)]
for x in range(ncol)]
else:
for subalign in genealign:
subalign.append(''.join(codons))
continue
if len(proteins) > 1:
if allelesets[0][1] == '+':
continue
proteins = mvf.decode(proteins)
if args.mincoverage is not None:
if sum([int(x not in 'X-') for x in proteins]) < (
args.mincoverage):
continue
species_groups = [[proteins[i] for i in x
if proteins[i] not in '-X']
for x in speciesgroups]
if any(not x for x in species_groups):
continue
xcodons = [mvf.decode(x) for x in codons]
codons = [''.join(x) for x in zip(*xcodons)]
if any(codons[x] in MLIB.stop_codons for x in allsets):
continue
if any(any(x != species_groups[0][0] for x in y)
for y in species_groups):
totals.add('total_nsyn_codons')
counts.add('total_nsyn_codons')
totals.add('total_codons')
totals.add('tested_codons')
counts.add('total_codons')
totals.add('variable_codons',
val=int(sum([int(len(set(x) - set('X-')) > 1)
for x in xcodons]) > 0))
if args.output_align is not None:
if not outputalign:
outputalign = [[x] for x in codons]
else:
for j, subalign in enumerate(outputalign):
subalign.append(codons[j])
if args.branch_lrt is not None:
if not genealign:
genealign = [[x] for x in codons]
else:
for j, codon in enumerate(codons):
genealign[j].append(codon)
nonsyn_change = False
synon_change = False
codon_groups = [
set(codons[i] for i in x if '-' not in codons[i] and
'X' not in codons[i])
for x in groups]
protein_groups = None
for i, grp in enumerate(codon_groups):
if any(base in codon for base in 'RYWKMS'
for codon in grp):
codon_groups[i] = hapgroup(grp)
if all(grp1.isdisjoint(grp0) for grp0, grp1 in
combinations(codon_groups, 2)):
protein_groups = [set(MLIB.codon_tables['full'][''.join(x)]
for x in codon_groups[i])
for i, grp in enumerate(codon_groups)]
if all(grp1.isdisjoint(grp0) for grp0, grp1 in
combinations(protein_groups, 2)):
nonsyn_change = True
elif all(grp1 == grp0 for grp0, grp1 in combinations(
protein_groups, 2)):
synon_change = True
if nonsyn_change:
if args.verbose is True:
print('NON', contig, pos, allelesets, codon_groups,
protein_groups, groups, mvf.get_contig_labels(
ids=contig))
counts.add('nonsynonymous_changes')
totals.add('nonsynonymous_changes')
elif synon_change:
if args.verbose is True:
print('SYN', contig, pos, allelesets, codon_groups,
protein_groups, groups, mvf.get_contig_labels(
ids=contig))
counts.add('synonymous_changes')
totals.add('synonymous_changes')
args.totals = totals
# WRITE OUTPUT
headers = ["contig", "position", "nonsynonymous_changes",
"synonymous_changes", "ns_ratio",
"nonsynonymous_total", "synonymous_total",
"pvalue",
"total_codons", "annotation", "coordinates"]
if args.windowsize == -1:
headers.remove('position')
if args.chi_test is None:
headers.remove('pvalue')
outfile = OutputFile(path=args.out, headers=headers)
sorted_entries = sorted([
(data[k]['ns_ratio'], k)
for k in data
if data[k].get('nonsynonymous_changes', 0) > 0],
reverse=True)
for _, k in sorted_entries:
outfile.write_entry(data[k])
with open(args.out + '.total', 'w') as totalfile:
for entry in args.totals.iter_sorted():
totalfile.write(entry)
if args.output_align is not None:
with open(args.output_align, 'w') as alignfile:
alignfile.write(
"\n".join([">{}\n{}".format(mvf.metadata['labels'][i],
''.join(outputalign[i]))
for i in range(len(outputalign))]))
return ''
def calc_dstat_combinations(args):
"""Calculate genome-wide D-statstics for
all possible trio combinations of samples
and outgroups specified.
"""
mvf = MultiVariantFile(args.mvf, 'read')
data = {}
sample_labels = mvf.get_sample_ids()
if args.outgroup_indices is not None:
outgroup_indices = [
int(x) for x in args.outgroup_indices[0].split(",")
]
elif args.outgroup_labels is not None:
outgroup_indices = mvf.get_sample_indices(
ids=args.outgroup_labels[0].split(","))
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
if args.contig_ids is not None:
contig_ids = args.contig_ids[0].split(",")
elif args.contig_labels is not None:
contig_ids = mvf.get_contig_ids(
labels=args.contig_labels[0].split(","))
else:
contig_ids = None
if any(x in outgroup_indices for x in sample_indices):
raise RuntimeError("Sample and Outgroup column lists cannot overlap.")
for contig, _, allelesets in mvf:
if contig not in contig_ids:
continue
alleles = mvf.decode(allelesets[0])
for i, j, k in combinations(sample_indices, 3):
for outgroup in outgroup_indices:
subset = [alleles[x] for x in [i, j, k, outgroup]]
if any(x not in 'ATGC' for x in subset):
continue
if subset[-1] not in subset[:3]:
continue
if len(set(subset)) != 2:
continue
# [ABBA, BABA, BBAA]
val = (0 + 1 * (subset[0] == subset[3]) + 2 *
(subset[1] == subset[3]) + 4 * (subset[2] == subset[3]))
if val in (1, 2):
val -= 1
elif val == 4:
val = 2
else:
continue
tetrad = (i, j, k, outgroup)
if tetrad not in data:
data[tetrad] = {}
if contig not in data[tetrad]:
data[tetrad][contig] = [0, 0, 0]
data[tetrad][contig][val] += 1
# WRITE OUTPUT
headers = ['sample0', 'sample1', 'sample2', "outgroup"]
for xcontig in contig_ids:
headers.extend([
'{}:abba'.format(xcontig), '{}:baba'.format(xcontig),
'{}:bbaa'.format(xcontig), '{}:D'.format(xcontig)
])
outfile = OutputFile(path=args.out, headers=headers)
for i, j, k in combinations(sample_indices, 3):
for outgroup in outgroup_indices:
tetrad = tuple([i, j, k, outgroup])
if tetrad not in data:
continue
entry = dict(('sample{}'.format(i), sample_labels[x])
for i, x in enumerate(tetrad[:3]))
entry['outgroup'] = sample_labels[outgroup]
for contig in contig_ids:
if contig not in data[tetrad]:
entry.update(dict().fromkeys([
'{}:abba'.format(contig), '{}:baba'.format(contig),
'{}:bbaa'.format(contig), '{}:D'.format(contig)
], '0'))
else:
[abba, baba, bbaa] = data[tetrad][contig]
if abba > baba and abba > bbaa:
dstat = zerodiv(baba - bbaa, baba + bbaa)
elif baba > bbaa and baba > abba:
dstat = zerodiv(abba - bbaa, abba + bbaa)
else:
dstat = zerodiv(abba - baba, abba + baba)
entry.update([('{}:abba'.format(contig), abba),
('{}:baba'.format(contig), baba),
('{}:bbaa'.format(contig), bbaa),
('{}:D'.format(contig), dstat)])
outfile.write_entry(entry)
return ''
def filter_mvf(args):
"""Main method"""
args.qprint("Running FilterMVF")
if args.more_help is True:
modulehelp()
sys.exit()
if args.mvf is None and args.test is None:
raise RuntimeError("No input file specified with --mvf")
if args.out is None and args.test is None:
raise RuntimeError("No output file specified with --out")
# Establish Input MVF
if args.test is not None:
ncol = args.test_nchar or len(args.test.split()[1])
else:
mvf = MultiVariantFile(args.mvf, 'read')
ncol = mvf.metadata['ncol']
args.qprint("Input MVF read with {} columns.".format(ncol))
# Create Actionset
if args.labels:
for i in range(len(args.actions)):
action = args.actions[i]
arr = action.split(':')
if arr[0] in ('collapsepriority', 'collapsemerge'):
arr[1] = ','.join([
str(mvf.sample_id_to_index[x])
for x in arr[1].split(',')])
if arr[0] in ('columns', 'allelegroup',
'notmultigroup', 'reqsample'):
for j in range(1, len(arr)):
arr[j] = ','.join([
str(mvf.sample_id_to_index[x])
for x in arr[j].split(',')])
args.actions[i] = ':'.join(arr)
removed_columns = set([])
for i in range(len(args.actions)):
action = args.actions[i]
arr = action.split(':')
if arr[0] in ('collapsepriority', 'collapsemerge'):
tmp_arr = arr[1][:]
arr[1] = ','.join([
str(int(x) - len([y for y in removed_columns if y < int(x)]))
for x in arr[1].split(',')])
removed_columns.update([int(x) for x in tmp_arr.split(',')[1:]])
print(arr)
print(removed_columns)
if arr[0] in ('columns', 'allelegroup',
'notmultigroup', 'reqsample'):
for j in range(1, len(arr)):
arr[j] = ','.join([
str(int(x) - len([y for y in removed_columns if y < int(x)]))
for x in arr[j].split(',')])
args.actions[i] = ':'.join(arr)
actionset = build_actionset(args.actions, ncol)
args.qprint("Actions established.")
args.qprint(actionset)
# TESTING MODE
if args.test:
loc, alleles = args.test.split()
linefail = False
transformed = False
# invar = invariant (single character)
# refvar (all different than reference, two chars)
# onecov (single coverage, + is second character)
# onevar (one variable base, + is third character)
# full = full alleles (all chars)
if args.verbose:
print(alleles)
linetype = get_linetype(alleles)
sys.stdout.write("MVF Encoding type '{}' detected\n".format(linetype))
for actionname, actiontype, actionfunc, actionarg in actionset:
sys.stdout.write("Applying action {} ({}): ".format(
actionname, actiontype))
if actiontype == 'filter':
if not actionfunc(alleles, linetype):
linefail = True
sys.stdout.write("Filter Fail\n")
break
sys.stdout.write("Filter Pass\n")
elif actiontype == 'transform':
transformed = True
alleles = actionfunc(alleles, linetype)
linetype = get_linetype(alleles)
if linetype == 'empty':
linefail = True
sys.stdout.write("Transform removed all alleles\n")
break
sys.stdout.write("Transform result {}\n".format(alleles))
elif actiontype == 'location':
loc = loc.split(':')
loc[1] = int(loc[1])
if actionfunc(loc) is False:
linefail = True
sys.stdout.write("Location Fail\n")
break
sys.stdout.write("Location Pass\n")
if linefail is False:
if transformed:
if linetype == 'full':
alleles = encode_mvfstring(alleles)
if alleles:
test_output = "{}\t{}\n".format(loc, alleles)
sys.stdout.write("Final output = {}\n".format(
test_output))
else:
sys.stdout.write("Transform removed all alleles\n")
else:
sys.stdout.write("No changes applied\n")
sys.stdout.write("Final output = {}\n".format(args.test))
sys.exit()
# MAIN MODE
# Set up file handler
outmvf = MultiVariantFile(args.out, 'write', overwrite=args.overwrite)
outmvf.copy_headers_from(mvf)
removed_indices = set([])
# reprocess header if actions are used that filter columns
if any(x == y[0] for x in ('columns', 'collapsepriority', 'collapsemerge')
for y in actionset):
for actionname, actiontype, actionfunc, actionarg in actionset:
if actionname == 'columns':
if args.labels:
oldindices = [outmvf.sample_id_to_index[int(x)]
for x in actionarg[0]]
else:
oldindices = [int(x) for x in actionarg[0]]
elif actionname in ('collapsepriority', 'collapsemerge'):
actionarg[0] = [x - len([y for y in removed_indices if y < x])
for x in actionarg[0]]
oldindices = [x for x in outmvf.sample_indices
if x not in actionarg[0][1:]]
outmvf.sample_ids = outmvf.get_sample_ids(oldindices)
outmvf.sample_data = dict(
(i, outmvf.sample_data[oldindices[i]])
for i, _ in enumerate(oldindices))
if actionname in ('collapsepriority', 'collapsemerge'):
if len(actionarg) == 2:
outmvf.sample_data[actionarg[0][0]]['id'] = actionarg[1][0]
outmvf.sample_ids[actionarg[0][0]] = actionarg[1][0]
outmvf.sample_indices = list(range(len(oldindices)))
outmvf.metadata['ncol'] = len(outmvf.sample_indices)
outmvf.notes.append(args.command_string)
outmvf.write_data(outmvf.get_header())
args.qprint("Output MVF established.")
# End header editing
linebuffer = []
nbuffer = 0
args.qprint("Processing Entries.")
write_total = 0
for chrom, pos, allelesets in mvf.iterentries(decode=False):
linefail = False
transformed = False
# invar = invariant (single character)
# refvar (all different than reference, two chars)
# onecov (single coverage, + is second character)
# onevar (one variable base, + is third character)
# full = full alleles (all chars)
alleles = allelesets[0]
linetype = get_linetype(alleles)
if linetype == 'empty':
continue
if args.verbose is True:
sys.stdout.write(" {} {} ".format(alleles, linetype))
for actionname, actiontype, actionfunc, _ in actionset:
if actiontype == 'filter':
linefail = not actionfunc(alleles, linetype)
elif actiontype == 'transform':
transformed = True
alleles = actionfunc(alleles, linetype)
linetype = get_linetype(alleles)
linefail = linetype == 'empty'
elif actiontype == 'location':
linefail = not actionfunc([chrom, pos])
if linefail:
break
if linefail is False:
if transformed:
if linetype == 'full':
alleles = mvf.encode(alleles)
if not alleles:
linefail = True
nbuffer += 1
linebuffer.append((chrom, pos, (alleles,)))
if args.verbose:
sys.stdout.write("{}\n".format(alleles))
if nbuffer == args.line_buffer:
write_total += args.line_buffer
args.qprint("{} entries written. Total written: {}.".format(
args.line_buffer, write_total))
outmvf.write_entries(linebuffer)
linebuffer = []
nbuffer = 0
elif args.verbose:
sys.stdout.write("FAIL\n")
if linebuffer:
outmvf.write_entries(linebuffer)
write_total += len(linebuffer)
args.qprint("{} entries written. Total written: {}.".format(
args.line_buffer, write_total))
linebuffer = []
return ''
def plot_chromoplot(args):
"""Main method"""
args.qprint("Running PlotChromoplot")
pallette = Pallette()
if args.colors is not None:
pallette.basecolors = args.colors
# Establish MVF and parse chromosome information
args.qprint("Reading MVF...")
mvf = MultiVariantFile(args.mvf, 'read')
args.qprint("Parsing headers...")
if args.contig_ids is not None:
contigids = args.contig_ids[0].split(",")
elif args.contig_labels is not None:
contigids = mvf.get_contig_ids(labels=args.contig_labels[0].split(","))
else:
contigids = mvf.get_contig_ids()
args.qprint("Plotting chromoplot for contigs: {}".format(
",".join(contigids)))
sample_labels = mvf.get_sample_ids()
if args.sample_indices is not None:
sample_indices = [int(x) for x in args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
assert len(sample_indices) >= 3
if args.outgroup_indices is not None:
outgroup_indices = [
int(x) for x in args.outgroup_indices[0].split(",")
]
elif args.outgroup_labels is not None:
outgroup_indices = mvf.get_sample_indices(
ids=args.outgroup_labels[0].split(","))
assert len(outgroup_indices) >= 1
quartets = [(x, y, z, outgroup)
for x, y, z in combinations(sample_indices, 3)
for outgroup in outgroup_indices]
# Begin iterations
for quartet_indices in quartets:
quartet_labels = [sample_labels[x] for x in quartet_indices]
if args.quiet is False:
print("Beginning quartet {}".format(",".join(quartet_labels)))
params = {
'contigs': [[
contigid, mvf.contig_data[contigid]['label'],
mvf.contig_data[contigid]['length']
] for contigid in contigids],
'outpath':
((args.out_prefix if args.out_prefix is not None else '')
or '_'.join(quartet_labels)) + ".png",
'labels':
quartet_labels,
'indices':
quartet_indices,
'windowsize':
args.windowsize,
'majority':
args.majority,
'infotrack':
args.info_track,
'yscale':
args.yscale,
'xscale':
args.xscale,
'quiet':
args.quiet,
'plottype':
args.plot_type
}
chromoplot = Chromoplot(params=params, pallette=pallette)
current_contig = ''
for contig, pos, allelesets in mvf.iterentries(subset=quartet_indices,
decode=True,
contig_ids=contigids):
if contig != current_contig:
if args.quiet is False:
print("Starting contig {}".format(contig))
current_contig = contig[:]
alleles = allelesets[0]
if '-' in alleles:
site_code = 'gap'
elif any(x not in 'ATGCatgc' for x in alleles):
site_code = 'ambiguous'
elif alleles[3] not in alleles[:3]:
site_code = 'nonpolar'
elif len(set(alleles)) > 2:
site_code = 'triallelic'
else:
site_code = sum([
2**(3 - j) * (alleles[j] != alleles[3]) for j in range(3)
])
chromoplot.add_data(str(contig), int(pos // args.windowsize),
site_code)
contig = ''
current_contig = ''
if not args.quiet:
print("Writing image...")
chromoplot.plot_chromoplot()
if not args.quiet:
print("Writing log...")
chromoplot.write_total_log()
return ''
def merge_mvf(args):
"""Main method"""
args.qprint("Running MergeMVF")
if any(fpath.endswith('.gz') for fpath in args.mvf):
print("WARNING! Running MergeMVF with gzipped input files is "
"extremely slow and strongly discouraged.")
concatmvf = MultiVariantFile(args.out, 'write', overwrite=args.overwrite)
# Copy the first file's metadata
args.qprint("Reading First File and Establishing Output")
if args.main_header_file:
if args.main_header_file not in args.mvf:
raise RuntimeError("{} not found in files".format(
args.main_header_file))
args.main_header_file = args.mvf.index(args.main_header_file)
else:
args.main_header_file = 0
first_mvf = MultiVariantFile(args.mvf[args.main_header_file], 'read')
concatmvf.copy_header(first_mvf)
# Open each MVF file, read headers to make unified header
transformers = []
mvfmetadata = []
inputfiles = []
for mvfname in args.mvf:
args.qprint("Reading headers from {}".format(mvfname))
# This will create a dictionary of samples{old:new}, contigs{old:new}
args.qprint("Processing Headers and Indexing: {}".format(mvfname))
transformer = MvfTransformer()
mvf = MultiVariantFile(mvfname,
'read',
contigindex=(not args.skip_index))
if args.skip_index:
mvf.read_index_file()
mvf.reset_max_contig()
mvfmetadata.append(mvf.metadata)
for i, sid in enumerate(mvf.get_sample_ids()):
if sid not in concatmvf.get_sample_ids():
new_sindex = concatmvf.max_sample_index + 0
concatmvf.max_sample_index += 1
concatmvf.sample_indices.append(new_sindex)
concatmvf.sample_ids.append(sid)
concatmvf.sample_data[new_sindex] = {}
concatmvf.sample_data[new_sindex]['id'] = sid
concatmvf.sample_id_to_index[sid] = new_sindex
transformer.set_label(i, concatmvf.sample_id_to_index[sid])
for cindex in mvf.contig_indices:
if (mvf.contig_data[cindex]['label']
not in concatmvf.contig_label_to_index):
new_cindex = (mvf.contig_data[cindex]['id']
if mvf.contig_data[cindex]['id']
not in concatmvf.contig_ids else
concatmvf.get_next_contig_index())
concatmvf.contig_data[new_cindex] = (
mvf.contig_data[cindex].copy())
else:
new_cindex = concatmvf.contig_label_to_index[
mvf.contig_data[cindex]['label']]
transformer.set_contig(cindex, new_cindex)
transformers.append(transformer)
inputfiles.append(mvf)
# Write output header
args.qprint("Writing headers to merge output")
concatmvf.reset_max_sample()
concatmvf.notes.append(args.command_string)
concatmvf.write_data(concatmvf.get_header())
# Now loop through each file
blank_entry = '-' * len(concatmvf.sample_indices)
for cons_contig in concatmvf.contig_indices:
contig_merged_entries = {}
args.qprint("Merging Contig Index: {}".format(cons_contig))
for ifile, mvffile in enumerate(inputfiles):
if cons_contig not in transformers[ifile].contigs:
continue
localcontig = transformers[ifile].contigs[cons_contig]
if 'idx' not in mvffile.contig_data[localcontig]:
print("not found")
continue
for _, pos, allelesets in mvffile.itercontigentries(localcontig,
decode=True):
if pos not in contig_merged_entries:
contig_merged_entries[pos] = blank_entry[:]
for j, base in enumerate(allelesets[0]):
xcoord = transformers[ifile].labels_rev[j]
if contig_merged_entries[pos][xcoord] != '-':
if contig_merged_entries[pos][xcoord] == base:
continue
if base in '-X':
continue
raise RuntimeError(
("Merging columns have two different bases: "
"{} {} {}").format(
pos, contig_merged_entries[pos][xcoord],
base))
contig_merged_entries[pos] = (
contig_merged_entries[pos][:xcoord] + base +
contig_merged_entries[pos][xcoord + 1:])
if contig_merged_entries:
concatmvf.write_entries(
((cons_contig, coord, (entry, ))
for coord, entry in sorted(contig_merged_entries.items())),
encoded=False)
args.qprint("Entries written for contig {}: {}".format(
cons_contig, len(contig_merged_entries)))
return ''
def concatenate_mvf(args):
"""Main method"""
args.qprint("Running ConcatenateMVF")
concatmvf = MultiVariantFile(args.out, 'write', overwrite=args.overwrite)
# Copy the first file's metadata
if args.main_header_file:
if args.main_header_file not in args.mvf:
raise RuntimeError("{} not found in files".format(
args.main_header_file))
args.main_header_file = args.mvf.index(args.main_header_file)
else:
args.main_header_file = 0
first_mvf = MultiVariantFile(args.mvf[args.main_header_file], 'read')
concatmvf.metadata = first_mvf.metadata.copy()
# Open each MVF file, read headers to make unified header
transformers = []
for mvfname in args.mvf:
# This will create a dictionary of samples{old:new}, contigs{old:new}
transformer = MvfTransformer()
mvf = MultiVariantFile(mvfname, 'read')
mvf.reset_max_contig()
for i, label in enumerate(mvf.get_sample_ids()):
if label not in concatmvf.get_sample_ids():
concatmvf.metadata['labels'].append(label)
concatmvf.metadata['samples'][
concatmvf.metadata['labels'].index(label)] = {
'label': label
}
if concatmvf.metadata['labels'].index(label) != i:
transformer.set_label(
i, concatmvf.metadata['labels'].index(label))
for contigid, contigdata in iter(mvf.metadata['contigs'].items()):
if contigdata['label'] not in [
concatmvf.metadata['contigs'][x]['label']
for x in concatmvf.metadata['contigs']
]:
newid = (contigid
if contigid not in concatmvf.metadata['contigs'] else
concatmvf.get_next_contig_id())
concatmvf.metadata['contigs'][newid] = contigdata
else:
for concatid, concatdata in (
concatmvf.metadata['contigs'].items()):
if contigdata['label'] == concatdata['label']:
newid = concatid
break
if newid != contigid:
transformer.set_contig(contigid, newid)
transformers.append(transformer)
# Write output header
concatmvf.notes.append(args.command_string)
concatmvf.write_data(concatmvf.get_header())
# Now loop through each file
entries = []
nentries = 0
for ifile, mvfname in enumerate(args.mvf):
args.qprint("Processing {} ...\n".format(mvfname))
transformer = transformers[ifile]
mvf = MultiVariantFile(mvfname, 'read')
for contigid, pos, allelesets in mvf.iterentries(decode=False):
if transformer.labels:
allelesets = [mvf.decode(x) for x in allelesets]
for j, alleles in enumerate(allelesets):
allelesets[j] = concatmvf.encode(''.join([
x in transformer.labels
and alleles[transformer.labels[x]] or alleles[x]
for x in range(len(alleles))
]))
if transformer.contigs:
contigid = (transformer['contigs'][contigid] if contigid
in transformer['contigs'] else contigid)
entries.append((contigid, pos, allelesets))
nentries += 1
if nentries == args.line_buffer:
concatmvf.write_entries(entries)
entries = []
nentries = 0
if entries:
concatmvf.write_entries(entries)
entries = []
nentries = 0
return ''
def infer_window_tree(args):
"""Main method"""
args.qprint("Running InferTree")
# ESTABLISH FILE OBJECTS
mvf = MultiVariantFile(args.mvf, 'read')
args.qprint("Read MVF File: {}".format(args.mvf))
# Set up contig ids
if args.contig_ids is not None:
contig_ids = args.contig_ids[0].split(",")
elif args.contig_labels is not None:
contig_ids = mvf.get_contig_ids(
labels=args.contig_labels[0].split(","))
else:
contig_ids = mvf.get_contig_ids()
treefile = OutputFile(
args.out,
headers=['contig', 'windowstart', 'windowsize', 'tree',
'topology', 'topoid',
# 'templabels', ### USED FOR DEBUGGING ###
'alignlength', 'aligndepth', 'status'])
topofile = OutputFile(args.out + '.counts',
headers=['rank', 'topology', 'count'])
if args.sample_indices is not None:
sample_indices = [int(x) for x in
args.sample_indices[0].split(",")]
elif args.sample_labels is not None:
sample_indices = mvf.get_sample_indices(
ids=args.sample_labels[0].split(","))
else:
sample_indices = mvf.get_sample_indices()
if not os.path.exists(args.temp_dir):
os.mkdir(args.temp_dir)
os.chdir(args.temp_dir)
# SETUP PARAMS
main_labels = mvf.get_sample_ids(sample_indices)
if args.choose_allele in ['randomboth', 'majorminor']:
main_labels = [label + x for x in ['a', 'b'] for label in main_labels]
params = {
'bootstrap': args.bootstrap,
'chooseallele': args.choose_allele,
'collapse_polytomies': args.collapse_polytomies,
'duplicateseq': args.duplicate_seq,
'engine': args.engine,
'engine_path': args.engine_path,
'engine_opts': args.engine_opts,
'mindepth': args.min_depth,
'minseqcoverage': args.min_seq_coverage,
'minsites': args.min_sites,
'model': args.model,
'outgroups': (args.raxml_outgroups
if args.raxml_outgroups is not None
else None),
'rootwith': (args.root_with.split(',')
if args.root_with is not None
else []),
'tempdir': args.temp_dir,
'tempprefix': args.temp_prefix,
'windowsize': args.windowsize,
}
# DEFAULT MODEL
if params['model'] is None:
if params['engine'] == 'raxml':
params['model'] = 'GTRGAMMA'
elif params['engine'] == 'raxml-ng':
params['model'] = "GTR+G"
# WINDOW START INTERATION
verify_raxml(params)
args.qprint("RAxML Found.")
current_contig = None
current_position = 0
window_data = None
# skip_contig = False
topo_ids = {}
topo_counts = {}
args.qprint("Prcocessing Records")
windowsizename = "window size={}".format(args.windowsize)
if windowsizename == "window size=-1":
windowsizename = "whole contig"
elif windowsizename == "window size=0":
windowsizename = "whole genome"
window_data = WindowData(window_params={
'contigname': 'all',
"windowstart": 0,
"windowsize": 0,
"labels": main_labels[:]})
for contig, pos, allelesets in mvf.iterentries(
contig_ids=contig_ids, subset=sample_indices,
no_invariant=False, no_ambig=False, no_gap=False, decode=True):
# if current_contig == contig:
# if skip_contig is True:
# args.qprint("Skipping contig: {}".format(current_contig))
# continue
if not same_window((current_contig, current_position),
(contig, pos), args.windowsize):
# skip_contig = False
if window_data is not None:
args.qprint(("Making tree for {} "
"at contig {} position {}").format(
windowsizename,
current_contig,
current_position))
entry = window_data.maketree_raxml(params)
if entry['status'] != 'ok':
if args.output_empty:
treefile.write_entry(entry)
# if args.windowsize != -1:
# skip_contig = True
args.qprint(
"TREE REJECTED with error code: {} ({})".format(
entry['status'], entry.get('comment', "None")))
else:
args.qprint("Tree completed.")
topo = entry["topology"]
topo_counts[topo] = topo_counts.get(topo, 0) + 1
if topo not in topo_ids:
topo_ids[topo] = (max(topo_ids.values()) + 1
if topo_ids else 0)
entry["topoid"] = topo_ids[topo]
treefile.write_entry(entry)
current_position = current_position + args.windowsize if (
contig == current_contig and args.windowsize > 0) else 0
current_contig = contig[:]
window_data = None
window_data = WindowData(window_params={
'contigname': (mvf.get_contig_labels(ids=current_contig) if
args.output_contig_labels is not None else
current_contig[:]),
"windowstart": ('-1' if args.windowsize == -1
else current_position + 0),
"windowsize": args.windowsize,
"labels": main_labels[:]})
# ADD ALLELES
if mvf.flavor == 'dna':
if args.choose_allele != 'none':
allelesets[0] = hapsplit(allelesets[0], args.choose_allele)
window_data.append_alleles(allelesets[0], mindepth=args.min_depth)
elif mvf.flavor == 'codon':
for i in (1, 2, 3):
if args.choose_allele != 'none':
allelesets[i] = hapsplit(allelesets[i], args.choose_allele)
window_data.append_alleles(allelesets[i], mindepth=args.min_depth)
# LAST LOOP
if window_data:
entry = window_data.maketree_raxml(params)
if entry['status'] != 'ok':
if args.output_empty:
treefile.write_entry(entry)
else:
topo = entry["topology"]
topo_counts[topo] = topo_counts.get(topo, 0) + 1
if topo not in topo_ids:
topo_ids[topo] = (
max(topo_ids.values()) + 1 if topo_ids else 0)
entry["topoid"] = topo_ids[topo]
treefile.write_entry(entry)
window_data = None
# END WINDOW ITERATION
topo_list = sorted([(v, k) for k, v in topo_counts.items()],
reverse=True)
for rank, [value, topo] in enumerate(topo_list):
topofile.write_entry({'rank': rank, 'count': value, 'topology': topo})
return ''
