pelops

#!/usr/bin/env python3

import argparse
import hashlib
import itertools
import json
import re
import sys
import random
import math
import gzip
import itertools

## File Functions

def read_fasta(filename):
	"""read fasta file one entry at a time"""
	label = None
	seq = []

	fp = None
	if    filename == '-':         fp = sys.stdin
	elif filename.endswith('.gz'): fp = gzip.open(filename, 'rt')
	else:                          fp = open(filename)

	while True:
		line = fp.readline()
		if line == '': break
		line = line.rstrip()
		if line.startswith('>'):
			if len(seq) > 0:
				seq = ''.join(seq)
				yield(label, seq)
				label = line[1:]
				seq = []
			else:
				label = line[1:]
		else:
			seq.append(line)
	yield(label, ''.join(seq))
	fp.close()

## K-mer Functions

def count_kmers(seq, k, full=False, init=0):
	"""get a table of kmer counts for a sequence with optional initial counts"""
	if full:
		kcount = {}
		for t in itertools.product('ACGT', repeat=k):
			kcount[''.join(t)] = init
	else:
		kcount = {}

	for i in range(len(seq) -k +1):
		kmer = seq[i:i+k]
		if kmer not in kcount: kcount[kmer] = 0
		kcount[kmer] += 1

	return kcount

## Sequence Utilities

def anti_seq(seq):
	"""get the reverse-complement of a sequence"""
	comp = str.maketrans('ACGTRYMKWSBDHVacgtrymkwsbdhv',
		'TGCAYRKMWSVHDBtgcayrkmwsvhdb')
	anti = seq.translate(comp)[::-1]
	return anti

def shuffle_seq(seq):
	"""get a randomly shuffled copy of a sequence"""
	dna = list(seq)
	random.shuffle(dna)
	return ''.join(dna)

def frame_shift_seq(seq, n=0, p=0, insert=False):
	"""if n provided, insert / delete n random bases from sequence, else mutate with probability p"""
	nuc = ['A', 'C', 'G', 'T']
	dna = list(seq)
	if n > 0:
		for i in random.sample(range(len(seq)), n):
			if insert == True:
				if   seq[i] == 'A': dna.insert(i, random.choice('CGT'))
				elif seq[i] == 'C': dna.insert(i, random.choice('AGT'))
				elif seq[i] == 'G': dna.insert(i, random.choice('ACT'))
				else:               dna.insert(i, random.choice('ACGs'))
			else: dna.pop(i)
	elif p > 0:
		for i in range(len(seq)):
			if random.random() > p:
				if insert == True:
					if   seq[i] == 'A': dna.insert(i, random.choice('CGT'))
					elif seq[i] == 'C': dna.insert(i, random.choice('AGT'))
					elif seq[i] == 'G': dna.insert(i, random.choice('ACT'))
					else:               dna.insert(i, random.choice('ACGs'))
				else: dna.pop(i)
	return "".join(dna)

def nmutate_seq(seq, n=0, p=0):
	"""get a copy of a sequence, mutated at n positions"""
	assert(n <= len(seq))
	dna = list(seq)
	if n > 0:
		for i in random.sample(range(len(seq)), n):
			if   seq[i] == 'A': dna[i] = random.choice('CGT')
			elif seq[i] == 'C': dna[i] = random.choice('AGT')
			elif seq[i] == 'G': dna[i] = random.choice('ACT')
			else:               dna[i] = random.choice('ACG')
	elif p > 0:
		for i in range(len(seq)):
			if random.random() < p:
				if   seq[i] == 'A': dna[i] = random.choice('CGT')
				elif seq[i] == 'C': dna[i] = random.choice('AGT')
				elif seq[i] == 'G': dna[i] = random.choice('ACT')
				else:               dna[i] = random.choice('ACG')
	return dna


def ksample_seq(seq, k, flop=False):
	"""get a sequence via sampled kmers"""
	dna = ''
	while True:
		i = random.randint(0, len(seq) -k)
		kmer = seq[i:i+k]
		if flop and random.random() < 0.5: kmer = anti_seq(kmer)
		dna += kmer
		if len(dna) > len(seq): break
	return dna[:len(seq)]

def entropy_filter(seq, comp, threshold=0):
	"""filter a sequence based on entropy (need to calculate composition of seq)"""
	entropy = 0
	for i in seq:
		entropy += -1 * (comp[i] * math.log2(comp[i]))
	if entropy > threshold:
		return True
	else: 
		return False

def align_seq(seq1, seq2, aln_threshold):
	"""align 2 sequences and return True if they are similar enough"""
	score = 0
	for i, j in zip(seq1, seq2):
		if i == j:
			score += 1
	if score/len(seq1) >= aln_threshold: return True
	else: return False

def find_orthologs(input_file, kmer_len = 5, samples=1000, entropy_threshold=0, weights = "", order = 0):
	""" compares sequences in file and finds patches of similarity"""
	fasta_test = read_fasta(input_file)
	kmer_list = {}
	composition = {}
	for i in fasta_test:
		anti = anti_seq(i[1])
		seq = i[1]
		kmer_list[i[0]] = []

		# break up seq and anti into kmers by sampling??
		seq_kmer = []
		# for j in range(0, len(seq) - kmer_len + 1):
		for j in range(samples):
			loc = random.randint(0, len(seq) - kmer_len + 1)
			if random.random() < 0.5:
				seq_kmer.append(seq[loc:loc+kmer_len])
			else:
				seq_kmer.append(anti[loc:loc+kmer_len])

		# entropy filter
		# should this come before / after sampling?
		seq_len = len(seq)
		comp = {"A": (seq.count("A") + 1)/seq_len, 
				"C": (seq.count("C") + 1)/seq_len, 
				"G": (seq.count("G") + 1)/seq_len, 
				"T": (seq.count("T") + 1)/seq_len}

		num_removed = 0
		for k in seq_kmer:
			if not entropy_filter(k, comp, entropy_threshold): 
				seq_kmer.remove(k)
				num_removed += 1
		composition[i[0]] = comp
		print(f"{i[0]}: removed {num_removed} kmers")
		kmer_list[i[0]] += seq_kmer
		composition[i[0]] = comp
		if weights == "KL": # Kullback-Leibler divergence
			# more divergence = more weight
			# assuming that w/ similar seq compositions, similar kmers are expected
			kl = {}
			for key1, value1 in composition.items():
				kl_divergence = 0
				for key2, value2 in composition.items():
					if key1 != key2:
						for key3, value3 in value1.items():
							kl_divergence += value3 * math.log2(value3 / value2[key3])
						kl_divergence = kl_divergence / 2
						kl[(key1,key2)] = kl_divergence
					else:
						kl[(key1,key2)] = 0
		# align    
		score = {}
		temp = sum([i for i in kmer_list.values()], [])
		kmer_summary = {i:temp.count(i) for i in set(temp)}
		for key, value in kmer_list.items():
			score[key] = {}
			if order != 0: 
				z1 = [i for i in itertools.permutations(value, order)]
			for key2, value2 in kmer_list.items():
				if weights == "":
					score[key][key2] = math.log2(len(list(set(value) & set(value2)))) - math.log2(len(value) + len(value2))
				else:
					if weights == "KL":
						weight = (kl[(key, key2)] + kl[(key2, key)]) / 2
						if weight == 0: weight = 1 #???
						score[key][key2] = math.log2(weight) + math.log2(len(list(set(value) & set(value2)))) - math.log2(len(value) + len(value2))
					# Interactions (SUPER SLOW!!!)
					if order != 0: 
						z2 = [i for i in itertools.permutations(value2, order)]
						score[key][key2] = math.log2(weight) + math.log2(len(list(set(z1) & set(z2))) - math.log2(len(z1) + len(z2)))
		print(score)
		return score



#################
## Subcommands ##
#################

def compare(arg):
	print('comparing')

def testing(arg):
	"""internal testing"""

	if arg.kmers or arg.all:
		print('kmers', end=': ')
		seq = 'AAAAACCCCGGG'
		k = 3
		j = json.dumps(count_kmers(seq, k, full=True, init=1))
		if arg.verbose: print(seq, k, j)
		h = hashlib.md5(j.encode('utf-8')).hexdigest()
		assert(h == 'f985dbb89ede64a80574d55810c7e577')
		print('passed', h)

def benchmark(arg):
	import time
	seq = 'ACGT' * 99
	limit = 999

	# method 1: if-elif stack
	t0 = time.time()
	for i in range(limit):
		for nt in seq:
			if   nt == 'A': s = random.choice('CGT')
			elif nt == 'C': s = random.choice('AGT')
			elif nt == 'G': s = random.choice('ACT')
			else:           s = random.choice('ACG')

	# method 2: list comprehension
	t1 = time.time()
	alph = 'ACGT'
	for i in range(limit):
		for nt in seq:
			subs = [x for x in alph if x != nt]
			s = random.choice(subs)

	# method 3: dictionary
	t2 = time.time()
	sub = {'A': 'CGT', 'C': 'AGT', 'G': 'ACT',	'T': 'ACG'}
	for i in range(limit):
		for nt in seq:
			s = random.choice(sub[nt])

	t3 = time.time()
	print(t1 - t0)
	print(t2 - t1)
	print(t3 - t2)

#########
## CLI ##
#########

parser = argparse.ArgumentParser()
sub = parser.add_subparsers(required=True, help='sub-commands')

## compare sub-command ##
parse_cmp = sub.add_parser('compare', help='compare 2 sequences')
parse_cmp.add_argument('--file1', required=True, help='fasta file')
parse_cmp.add_argument('--file2', required=True, help='fasta file')
parse_cmp.add_argument('--kmin', required=False, type=int, default=5,
	help='minimum k-mer size [default %(default)i]')
parse_cmp.add_argument('--kmax', required=False, type=int, default=7,
	help='maximum k-mer size [default %(default)i]')
parse_cmp.add_argument('--single', required=False, action='store_true',
	help='use single-stranded k-mers [default double-stranded]')
parse_cmp.set_defaults(func=compare)

## test sub-command
parse_test = sub.add_parser('test', help='internal testing')
parse_test.add_argument('--all', required=False, action='store_true')
parse_test.add_argument('--kmers', required=False, action='store_true')
parse_test.add_argument('--verbose', required=False, action='store_true')
parse_test.set_defaults(func=testing)

## benchmark sub-command
parse_bench = sub.add_parser('bench', help='internal benchmarking')
parse_bench.set_defaults(func=benchmark)


## execute sub-command ##
try: arg = parser.parse_args()
except: sys.exit('pelops requires a sub-command, use --help for more info')
arg.func(arg)