genmut

#!/usr/bin/env python3

from readfasta import read_record
from longestorf import longestorf
import argparse
import json
import os

parser = argparse.ArgumentParser(description='Simulate splice-site-generating point mutations')
parser.add_argument('fasta', type=str, metavar='<fasta>', help='input path to fasta file')
parser.add_argument('gff', type=str, metavar='<gff>', help='input path to gff file')
parser.add_argument('-m', type=float, metavar='<float>', required=False, default='1.0',
	help='input Manhattan distance threshold for selecting mutant isoforms (default: 1.0)')
parser.add_argument('-p', type=float, metavar='<float>', required=False, default='0.01',
	help='input probability threshold for selecting mutant isoforms (default: 0.01)')
parser.add_argument('-o', type=str, metavar='<out_dir>', required=False, help='output dir')
arg = parser.parse_args()

if arg.o:
	os.makedirs(arg.o, exist_ok=True)
	OUTDIR = os.path.abspath(arg.o)
else: OUTDIR = os.path.abspath('.')
base_name = None

#######################################################
## Simulate de novo splice site generating mutaitons ##
#######################################################

for idn, seq in read_record(arg.fasta):
	file_name = os.path.basename(arg.fasta)
	base_name = os.path.splitext(file_name)[0]
	count = 1
	for i in range(0, len(seq)-1, 2):
		# donor site
			cur_id  = None
			cur_seq = None
			if seq[i] == 'G' and seq[i+1] != 'T':
				cur_id  = f'{idn} {i+1} d {seq[i+1]}~T'
				cur_seq = f'{seq[:i+1]}T{seq[i+2:]}'
			if seq[i] != 'G' and seq[i+1] == 'T':
				cur_id  = f'{idn} {i} d {seq[i]}~G'
				cur_seq = f'{seq[:i]}G{seq[i+1:]}'
		# acceptor site
			if seq[i] == 'A' and seq[i+1] != 'G':
				cur_id  = f'{idn} {i+1} a {seq[i+1]}~G'
				cur_seq = f'{seq[:i+1]}G{seq[i+2:]}'
			if seq[i] != 'A' and seq[i+1] == 'G':
				cur_id  = f'{idn} {i} a {seq[i]}~A'
				cur_seq = f'{seq[:i]}A{seq[i+1:]}'
		# out
			if cur_id is not None:
				os.makedirs(f'{OUTDIR}/mutants/{base_name}', exist_ok=True)
				with open(f'{OUTDIR}/mutants/{base_name}/mut_{count}_{base_name}.fa', 'w') as fh:
					fh.write(f'>{cur_id}\n')
					for i in range(0, len(cur_seq), 80):
						fh.write(cur_seq[i:i+80])
				count+=1

###############################
## Predict splicing patterns ##
###############################

MODELS_DIR = os.path.abspath('data')
MODELS = f'-dpwm {MODELS_DIR}/donor.pwm -apwm {MODELS_DIR}/acceptor.pwm\
	-emm {MODELS_DIR}/exon.mm -imm {MODELS_DIR}/intron.mm\
	-elen {MODELS_DIR}/exon.len -ilen {MODELS_DIR}/intron.len'

for mutant in os.listdir(f'{OUTDIR}/mutants/{base_name}'):
	os.makedirs(f'{OUTDIR}/mutantiso/{base_name}', exist_ok=True)
	mut_base = os.path.splitext(mutant)[0]
	os.system(f'./isoformer_mut {MODELS} -fa {OUTDIR}/mutants/{base_name}/{mutant} \
		-introns {arg.gff} > \
		{OUTDIR}/mutantiso/{base_name}/{mut_base}.isoform')
os.makedirs(f'{OUTDIR}/ogiso', exist_ok=True)
os.system(f'./isoformer {MODELS} -fa {arg.fasta} -introns {arg.gff} > \
	{OUTDIR}/ogiso/{base_name}.isoform')
	
##############################################################
## Compare splicing patterns between mutants and annotation ##
##############################################################

os.makedirs(f'{OUTDIR}/cmpiso', exist_ok=True)
for mutantiso in os.listdir(f'{OUTDIR}/mutantiso/{base_name}'):
	os.system(f'./cmpiso {OUTDIR}/ogiso/{base_name}.isoform \
		{OUTDIR}/mutantiso/{base_name}/{mutantiso} >> {OUTDIR}/cmpiso/{base_name}.cmp')

true_mutants = {}
true_mutants_sum = {}
with open(f'{OUTDIR}/cmpiso/{base_name}.cmp') as fh:
	for line in fh.readlines():
		dist = float(line.split()[0])
		mut  = line.split()[1]
		gene = mut.split('_')[2]
		true_mut = False
		if dist > arg.m:
			with open(f'{OUTDIR}/mutantiso/{base_name}/{mut}.isoform') as mf:
				idn = mf.readline().split()
				pos = int(idn[5])
				typ = idn[6]
				for line in mf.readlines():
					if line.startswith('#'): continue
					if len(line.split()) < 8: continue
					if line.split()[2] != 'intron': continue
					if typ == 'd' and int(line.split()[3]) == pos+1: true_mut = True
					if typ == 'a' and int(line.split()[4]) == pos+1: true_mut = True
			if true_mut:
				if gene not in true_mutants: true_mutants[gene] = {}
				if gene not in true_mutants_sum: true_mutants_sum[gene] = 0
				true_mutants_sum[gene] += 1
				true_mutants[gene][mut] = str(dist)	

os.makedirs(f'{OUTDIR}/jsonout', exist_ok=True)
with open(f'{OUTDIR}/jsonout/{base_name}.json', 'w') as fh:
	fh.write(json.dumps(true_mutants, indent=4))
with open(f'{OUTDIR}/jsonout/{base_name}_summary.json', 'w') as fh:
	true_mutants_sum = dict(sorted(true_mutants_sum.items(), key=lambda item: item[1], reverse=True))
	fh.write(json.dumps(true_mutants_sum, indent=4))

#########################################
## Select isoform and output orf fasta ##
#########################################
def get_cdss(exons):
	rnaseq = ''
	orf    = ''
	cdss   = []
	for exon in exons:
		rnaseq += mutseq[exon[0]:exon[1]+1]
	#print(f'>rnaseq\n{rnaseq}')
	
	for i in range(0, len(rnaseq)-2):
		start = (rnaseq[i:i+3] == 'ATG')
		if start:
			cur_orf = ''
			for j in range(i, len(rnaseq)-2, 3):
				cur_orf += rnaseq[j:j+3]
				if rnaseq[j:j+3] in stop_codons:
					if len(cur_orf) > len(orf):
						orf = cur_orf
						orf_pos = [i,j+3]
					break
	tot = 0
	exon_start = None
	orf_start = orf_pos[0]
	for idx, exon in enumerate(exons):
		exlen = exon[1]-exon[0]
		tot += exlen
		if tot > orf_pos[0]:
			exon_start = idx
			break
		orf_start -= exlen + 1
	
	r = len(orf)
	for idx, exon in enumerate(exons[exon_start:]):
		if idx == 0:
			if exon[0] + orf_start + r > exon[1]:
				cdss.append([exon[0]+orf_start,exon[1]])
				r = r - (exon[1] - (exon[0] + orf_start)) -1
			else:
				cdss.append([exon[0]+orf_start,exon[0]+orf_start+r-1])
				break
		else:
			if exon[0] + r > exon[1]:
				cdss.append(exon)
				r = r - (exon[1] - exon[0]) -1
			else:
				cdss.append([exon[0],exon[0]+r-1])
				break
	return cdss

stop_codons = ['TAA', 'TAG', 'TGA']
for mutant in true_mutants:
	for mutiso in true_mutants[mutant]:
		mutseq = ''
		for idn, seq in read_record(f'{OUTDIR}/mutants/{base_name}/{mutiso}.fa'): mutseq += seq
		mut_pos = None
		mut_typ = None
		isoforms = []
		with open(f'{OUTDIR}/mutantiso/{base_name}/{mutiso}.isoform') as fh:
			isoform = []
			while True:
				line = fh.readline()
				if line == '': break
				line = line.rstrip()
				if line.startswith('#'):
					if line.startswith('# name'):
						mut_pos = int(line.split()[5])
						mut_typ = line.split()[6]
					continue
				fields = line.split()
				if len(fields) > 8 and fields[2] == 'mRNA' and float(fields[5]) > arg.p:
					isoform.append(line)
				if len(fields) < 1 and len(isoform) > 1:
					isoforms.append(isoform)
					isoform = []
				if len(isoform) > 0 and fields[2] != 'mRNA': isoform.append(line)
		
		ff_id = None
		with open(f'{OUTDIR}/mutants/{base_name}/{mutiso}.fa') as fh:
			ff_id = fh.readline().rstrip()
		
		os.makedirs(f'{OUTDIR}/selected_iso', exist_ok=True)
		for idx, isoform in enumerate(isoforms):
			exons = []
			introns = []
			prob = None
			for line in isoform:
				fields = line.split()
				if fields[2] == 'mRNA': prob = float(fields[5])
				if fields[2] == 'exon':
					exons.append([int(fields[3])-1, int(fields[4])-1])
				if fields[2] == 'intron':
					introns.append(int(fields[3])-1)
					introns.append(int(fields[4])-1)
			
			if mut_pos in introns:
				os.makedirs(f'{OUTDIR}/selected_iso/{base_name}', exist_ok=True)
				mut_cdss = get_cdss(exons)
				with open(f'{OUTDIR}/selected_iso/{base_name}/{mutiso}_{idx}.fa', 'w') as fh:
					fh.write(f'{ff_id} ')
					for cds in mut_cdss: fh.write(str(cds))
					fh.write(f' {prob}\n')
					for cds in mut_cdss: fh.write(f'{mutseq[cds[0]:cds[1]+1]}')

########################################################
## Translate selected isoforms and annotated isoforms ##
########################################################

os.makedirs(f'{OUTDIR}/selected_pep', exist_ok=True)
os.makedirs(f'{OUTDIR}/annotated_pep', exist_ok=True)
if os.path.isdir(f'{OUTDIR}/selected_iso/{base_name}'):
	if not os.path.isfile(f'{OUTDIR}/annotated_pep/{base_name}.fa'):
		cdss = []
		name = None
		with open(arg.gff) as fh:
			while True:
				line = fh.readline()
				if line == '': break
				f = line.split()
				if f[2] == 'mRNA' and name == None:
					for att in f[8].split(';'):
						if 'ID=' in att:
							name = att.split('=')[1]
							break
		with open(arg.gff) as fh:
			while True:
				line = fh.readline()
				if line == '': break
				f = line.split()
				if f[1] != 'WormBase': continue
				if f[2] == 'CDS' and name != None and name in f[8] : cdss.append((int(f[3])-1, int(f[4])))
		cdss.sort()
		orf = ''
		for idn, seq in read_record(arg.fasta):
			for cds in cdss: orf += seq[cds[0]:cds[1]]
			with open(f'{OUTDIR}/annotated_pep/{base_name}.fa', 'w') as fh:
				protein = longestorf(orf)
				fh.write(f'>{idn}\n{protein}\n')
	for orf in os.listdir(f'{OUTDIR}/selected_iso/{base_name}'):
		os.makedirs(f'{OUTDIR}/selected_pep/{base_name}', exist_ok=True)
		with open(f'{OUTDIR}/selected_pep/{base_name}/{orf}', 'w') as fh:
			for idn, orf in read_record(f'{OUTDIR}/selected_iso/{base_name}/{orf}'):
				protein = longestorf(orf)
				fh.write(f'>{idn}\n{protein}\n')