seqxtant-legacy

#!/usr/bin/env python3

import argparse
import json
import os
import re
import subprocess
import sys
import csv
## Database Functions

def read_db(db):
	with open(db) as fp:
		return json.load(fp)

def write_db(db, d):
	with open(db, 'w') as fp:
		fp.write(json.dumps(d, indent=4))

def run(arg, cmd):
	if arg.verbose: print(cmd, file=sys.stderr, end=' ...')
	ret = subprocess.run(cmd, shell=True).returncode
	if ret != 0:
		raise Exception(f'{cmd} failed with status {ret}')
	if arg.verbose: print(f' done', file=sys.stderr)

## Subcommand Functions

def create(arg, env):
	db = f'{env}/seqxtant.json'
	if os.path.isfile(db):
		raise Exception(f'seqxtant databse already exists at {db}')
	write_db(db, {'state': 'ready', 'genomes': []})

def status(arg, env):
	db = f'{env}/seqxtant.json'
	if not os.path.isfile(db):
		raise Exception(f'{db} is not a seqxtant database location')
	print(f'Data location: {db}')
	d = read_db(db)
	print(f'Database status: {d["state"]}')
	print(f'Genomes: {len(d["genomes"])}')
	for genome in d['genomes']:
		print(f'\t{genome}')

def add_genome(arg, env):
	db = f'{env}/seqxtant.json'
	d = read_db(db)
	if d['state'] == 'busy' and not arg.force:
		raise Exception(f'cannot add genome, {db} is currently busy')
	if arg.genome in d['genomes'] and not arg.force:
		raise Exception(f'cannot add genome, {arg.genome} already present')

	# copy files to build directory under new names
	d['state'] = 'busy'
	write_db(db, d)
	dna = f'{arg.env}/{arg.genome}.fa'
	msk = f'{arg.env}/{arg.genome}.masked.fa'
	gff  = f'{arg.env}/{arg.genome}.gff3'
	
	# genome
	if arg.fasta.endswith('.gz'):
		run(arg, f'cp -f {arg.fasta} {dna}.gz')
		run(arg, f'gunzip -f {dna}.gz')
	else:
		run(arg, f'cp -f {arg.fasta} {dna}')
	
	# masked genome
	if arg.masked.endswith('.gz'):
		run(arg, f'cp -f {arg.masked} {msk}.gz')
		run(arg, f'gunzip -f {msk}.gz')
	else:
		run(arg, f'cp -f {arg.masked} {msk}')
	
	# gff3
	if arg.gff3.endswith('.gz'):
		run(arg, f'cp -f {arg.gff3} {gff}.gz')
		run(arg, f'gunzip -f {gff}.gz')
	else:
		run(arg, f'cp -f {arg.gff3} {gff}')
	
	# make blast databases
	write_db(db, d)
	run(arg, f'formatdb -i {dna} -p F -o T')
	run(arg, f'formatdb -i {msk} -p F -o T')
	
	# done
	if arg.genome not in d['genomes']: d['genomes'].append(arg.genome)
	d['state'] = 'ready'
	write_db(db, d)

def validate_genome(arg, env):
	db = f'{env}/seqxtant.json'
	d = read_db(db)
	
	if arg.genome not in d['genomes']:
		raise Exception(f'{arg.genome} not found in database')
	
	ff = f'--fasta {arg.env}/{arg.genome}.fa'
	g3 = f'--gff {arg.env}/{arg.genome}.gff3'
	ht = f'--html {arg.html} --tables {arg.tables}'
	
	run(arg, f'calfo --title {arg.genome} {ff} {g3} {ht}')

def cbe(s):
	m = re.search('(\S+):(\d+)\-(\d+)', s)
	if not m: raise Exception('location string malformed')
	c = m.group(1)
	b = int(m.group(2))
	e = int(m.group(3))
	return c, b, e

def store_fasta(arg, env, masked=False):
	c, b, e = cbe(arg.location)
	if masked:
		infile = f'{env}/{arg.genome}.masked.fa'
		outfile = f'{arg.odir}/query.masked.fa'
	else:
		infile = f'{env}/{arg.genome}.fa'
		outfile = f'{arg.odir}/query.fa'
	cmd = f'fastacmd -d {infile} -p F -s {c} -L {b},{e}'
	proc = subprocess.run(cmd, shell=True, capture_output=True)
	if proc.returncode != 0: raise Exception(f'fastacmd failed: {proc.stderr}')
	lines = proc.stdout.decode().split('\n')
	lines[0] = f'>{arg.genome}-{arg.location}'
	with open(outfile, 'w') as fp:
		for line in lines:
			fp.write(line)
			fp.write('\n')

def run_blast(arg, env, target):
	qf = f'{arg.odir}/query.masked.fa'
	db = f'{env}/{target}.fa'
	of = f'{arg.odir}/{target}.blastn'
	params = (
		'-r 1 -q -1',     # +1/-1 scoring scheme, default 1/-3 too stringent
		'-G 2 -E 2',      # gap penalty is constant -2
		'-W 10',          # wordsize
		'-e 1e-10',       # e-value
		'-b 100 -v 100',  # limit the number of blast alignments
		'-m 8',
	)
	opts = ' '.join(params)
	cmd = f'blastall -p blastn -d {db} -i {qf} -o {of} {opts} -a {arg.cpus}'
	run(arg, cmd)
	
	# read HSPs
	hsps = []
	with open(of) as fp:
		for line in fp:
			if line.startswith('#'): continue
			qid, sid, pct, alen, mis, go, qb, qe, sb, se, e, s = line.split()
			if sb < se:
				st = '+'
			else:
				st = '-'
				sb, se = se, sb
			hsps.append([
				sid, int(qb), int(qe), int(sb), int(se), st, float(s)
			])
	
	return hsps

def score_algo(combo):
	# scoring system parameters
	dist_penalty = 0.5
	overlap_penalty = 2
	connect_reward = 20
	# get sum of blast scores
	score = sum(each[6] for each in combo)
	
	# get the distance penalty, by the max projection of the distance onto one coordinate axis
	score -= (dist_penalty * 
		sum(max(combo[i+1][1]-combo[i][2], combo[i+1][3]-combo[i][4]) 
		for i in range(len(combo)-1) 
		if max(combo[i+1][1]-combo[i][2], combo[i+1][3]-combo[i][4]) > 0))

	
	# get the overlap penalty
	score += (overlap_penalty * 
		sum(min(combo[i+1][1]-combo[i][2], combo[i+1][3]-combo[i][4]) 
		for i in range(len(combo)-1) 
		if min(combo[i+1][1]-combo[i][2], combo[i+1][3]-combo[i][4]) < 0))
		

	
	# get connect reward
	score += connect_reward * (len(combo)-1)
	return score



def score_sys(hsps):
	# sort hsps by coordinate and blast score respectively
	hsps.sort(key = lambda x: (x[1], x[3]))
	hsps_order = [each for each in hsps]
	hsps_order.sort(key = lambda x: x[6])
	combos = []

	while hsps != []:
		combo = [hsps_order[-1]]

		idx = hsps.index(hsps_order[-1])
		score = hsps_order[-1][6]
		# forward
		for i in range(len(hsps[idx+1:])):
			new_combo = combo + [hsps[idx+1+i]]
			new_score = score_algo(new_combo)

			if new_score > score:
				combo = new_combo
				score = new_score

		# backward
		# maybe we should design a penalty for backward connection? (first intron
		# intends to be the longest one in plant genome)
		for i in range(len(hsps[:idx])):
			new_combo = [hsps[idx-1-i]] + combo
			new_score = score_algo(new_combo)

			if new_score > score:
				combo = new_combo
				score = new_score

		# delete used alignments

		for each in combo:
			hsps.remove(each)
			hsps_order.remove(each)
			
		combos.append(combo)
	return combos

def score(hsps):
	dic_hsps = {}
	dic_combos = {}
	for hsp in hsps:
		if 	hsp[0]+hsp[5] not in dic_hsps: dic_hsps[hsp[0]+hsp[5]] = [hsp]
		else: dic_hsps[hsp[0]+hsp[5]] += [hsp]
	for sid_st in dic_hsps:
		dic_combos[sid_st] = score_sys(dic_hsps[sid_st])
	return dic_combos

def cluster(arg, env):
	db = f'{env}/seqxtant.json'
	d = read_db(db)
	if d['state'] == 'busy': raise Exception(f'database is busy')
	
	store_fasta(arg, env,)
	store_fasta(arg, env, masked=True)
	
	if arg.limit: targets = arg.limit
	else:         targets = d['genomes']
	


# This is where we are stopped
	for target in targets:
		
		if arg.genome == target and not arg.paralogs: continue
		of = f'{arg.odir}/{target}.gff'

		hsps = run_blast(arg, env, target)
		dic_hsps = score(hsps)

		# write gff file
		f = open(of, "w")
		writer = csv.writer(f, delimiter='\t', lineterminator='\n')

		for sid_st in dic_hsps:

			f.write(f"## {sid_st[:-1]} {sid_st[-1]}\n")
			for region in dic_hsps[sid_st]:
				writer.writerow((
					region[0][0], "prediction", "region", 
					region[0][3], region[-1][4], ".", 
					region[0][5], ".", 
					"ID=homoregion"
				))

				writer.writerows(
					(algn[0], "prediction", "alignment", 
					algn[3], algn[4], ".", algn[5], 
					".", "ID=alignment") 
					for algn in region
				)
		f.close()	
'''
		for sid_st in dic_hsps:
			print('\n\n', sid_st)
			for each in dic_hsps[sid_st]:
				print(each, '\n')
'''
		
		# figure out how they are organized
		# get annotation from same region (and change coordinates?)
		# create a sub-location pair of files (fa, gff)
	
	# some kind of report on clustering
	# send to multiple align?
	

## CLI top-level ##
parser = argparse.ArgumentParser(description='genomic homology multi-mapper')
parser.add_argument('--env',
	help='set data location instead of $SEQXTANT')
parser.add_argument('--verbose', action='store_true',
	help='see behind the scenes messages')
subparsers = parser.add_subparsers(required=True,
	help='sub-commands')

### status sub-command
parse_status = subparsers.add_parser('status',
	help='get information about the database')
parse_status.set_defaults(func=status)

### create sub-command
parse_create = subparsers.add_parser('create',
	help='create a new database')
parse_create.set_defaults(func=create)

### add sub-command
parse_add = subparsers.add_parser('add',
	help='add new genome to database')
parse_add.add_argument('--genome', required=True,
	help='genome name, e.g. C.elegans')
parse_add.add_argument('--fasta', required=True,
	help='genome in fasta')
parse_add.add_argument('--masked', required=True,
	help='hard-masked genome in fasta')
parse_add.add_argument('--gff3', required=True,
	help='annotation in gff3')
parse_add.add_argument('--force', action='store_true',
	help='force overwrite of previous data')
parse_add.set_defaults(func=add_genome)

### validate sub-command
parse_validate = subparsers.add_parser('validate',
	help='validate genome to check for errors')
parse_validate.add_argument('--genome', required=True,
	help='genome name, e.g. C.elegans')
parse_validate.add_argument('--html', required=True,
	help='path to html output file')
parse_validate.add_argument('--tables', required=True,
	help='path to output directory')
parse_validate.set_defaults(func=validate_genome)

### cluster sub-command
parse_find = subparsers.add_parser('cluster',
	help='cluster related sequences')
parse_find.add_argument('--genome', required=True,
	help='genome name')
parse_find.add_argument('--location', required=True,
	help='chrom:begin-end')
parse_find.add_argument('--odir', required=True,
	help='path to output directory')
parse_find.add_argument('--limit', nargs='+', required=False,
	help='limit to specific genomes')
parse_find.add_argument('--cpus', type=int, required=False, default=1,
	help='number of cpus to use [%(default)i]')
parse_find.add_argument('--paralogs', action='store_true',
	help='include paralogs from host genome')
parse_find.set_defaults(func=cluster)

### sub-command checkpoint
try:
	arg = parser.parse_args()
except:
	print('seqxtant requires a sub-command, use --help for more info')
	sys.exit(1)

## Database location ##
if arg.env: env = arg.env
else:       env = os.getenv('SEQXTANT')
if not env: raise Exception('you must set SEQXTANT or use --env option')

## Run subcommand ##
arg.func(arg, env)