<s> import sys <EOL> import os <EOL> import shutil <EOL> import glob <EOL> import logging <EOL> import tempfile <EOL> import subprocess <EOL> import pkg_resources <EOL> import zlib <EOL> import subprocess <EOL> from pbcore . util . ToolRunner import PBMultiToolRunner <EOL> from pbcore . io import FastaReader <EOL> from pbtools . pbtranscript . icedagcon . c_aligngraph import * <EOL> from pbtools . pbtranscript . icedagcon . c_utils import construct_aln_graph_from_fasta <EOL> from pbtools . pbtranscript . icedagcon . c_utils import sorted_nodes <EOL> from pbtools . pbtranscript . icedagcon . c_utils import best_template_by_blasr <EOL> from pbtools . pbtranscript . icedagcon . c_utils import clustering_read <EOL> from pbtools . pbtranscript . icedagcon . c_utils import get_subset_reads <EOL> from pbtools . pbtranscript . icedagcon . c_utils import read_node_vector <EOL> from pbtools . pbtranscript . icedagcon . c_utils import detect_missing <EOL> from pbtools . pbtranscript . icedagcon . c_utils import mark_lower_case_base <EOL> try : <EOL> __p4revision__ = "<STR_LIT>" <EOL> __p4change__ = "<STR_LIT>" <EOL> revNum = int ( __p4revision__ . strip ( "<STR_LIT:$>" ) . split ( "<STR_LIT:U+0020>" ) [ <NUM_LIT:1> ] . strip ( "<STR_LIT:#>" ) ) <EOL> changeNum = int ( __p4change__ . strip ( "<STR_LIT:$>" ) . split ( "<STR_LIT::>" ) [ - <NUM_LIT:1> ] ) <EOL> __version__ = "<STR_LIT>" % ( pkg_resources . require ( "<STR_LIT>" ) [ <NUM_LIT:0> ] . version , revNum , changeNum ) <EOL> except : <EOL> __version__ = "<STR_LIT>" <EOL> rmap = dict ( zip ( "<STR_LIT>" , "<STR_LIT>" ) ) <EOL> def normalize_fasta ( fastaFile , refFile , outFile , nproc ) : <EOL> f = FastaReader ( fastaFile ) <EOL> recs = [ ] <EOL> with open ( outFile , "<STR_LIT:w>" ) as of : <EOL> for r in f : <EOL> r_id = "<STR_LIT:%s>" % hex ( zlib . adler32 ( r . name + r . sequence ) & <NUM_LIT> ) <EOL> print >> of , "<STR_LIT:>>" + r_id <EOL> seq = r . sequence . upper ( ) <EOL> print >> of , seq <EOL> os . system ( "<STR_LIT>" % ( nproc , outFile , refFile , refFile ) ) <EOL> output = subprocess . check_output ( "<STR_LIT>" % refFile , shell = True ) <EOL> direction = { } <EOL> output = output . strip ( ) . split ( "<STR_LIT:\n>" ) <EOL> for l in output : <EOL> l = l . strip ( ) . split ( ) <EOL> rId = l [ <NUM_LIT:0> ] . split ( "<STR_LIT:/>" ) [ <NUM_LIT:0> ] <EOL> if l [ <NUM_LIT:2> ] != l [ <NUM_LIT:3> ] : <EOL> direction [ rId ] = "<STR_LIT:->" <EOL> else : <EOL> direction [ rId ] = "<STR_LIT:+>" <EOL> os . system ( "<STR_LIT>" % refFile ) <EOL> f = FastaReader ( outFile ) <EOL> outData = [ ] <EOL> for r in f : <EOL> r_id = "<STR_LIT:%s>" % r . name <EOL> outData . append ( "<STR_LIT:>>" + r_id ) <EOL> seq = r . sequence . upper ( ) <EOL> if direction != None : <EOL> if direction . get ( r_id , "<STR_LIT:+>" ) != "<STR_LIT:+>" : <EOL> seq = "<STR_LIT>" . join ( [ rmap [ c ] for c in seq [ : : - <NUM_LIT:1> ] ] ) <EOL> outData . append ( seq ) <EOL> with open ( outFile , "<STR_LIT:w>" ) as of : <EOL> print >> of , "<STR_LIT:\n>" . join ( outData ) <EOL> def get_consensus ( read_fn , init_ref , consensus_fn , consens_seq_name , <EOL> hp_correction = False , <EOL> min_iteration = <NUM_LIT:4> , <EOL> max_num_reads = <NUM_LIT> , <EOL> entropy_th = <NUM_LIT> , <EOL> min_cov = <NUM_LIT:8> , <EOL> max_cov = <NUM_LIT> , <EOL> nproc = <NUM_LIT:4> , <EOL> mark_lower_case = False , <EOL> use_read_id = False ) : <EOL> g = construct_aln_graph_from_fasta ( read_fn , init_ref , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc ) <EOL> s , c = g . generate_consensus ( min_cov = min_cov ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if min_iteration > <NUM_LIT:1> : <EOL> for i in range ( min_iteration - <NUM_LIT:2> ) : <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc , use_read_id = use_read_id ) <EOL> s , c = g . generate_consensus ( min_cov = min_cov ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if hp_correction : <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc ) <EOL> s = detect_missing ( g , entropy_th = entropy_th ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = True , max_cov = max_cov , nproc = nproc ) <EOL> s , c = g . generate_consensus ( min_cov = min_cov ) <EOL> if mark_lower_case : <EOL> s = mark_lower_case_base ( g , entropy_th = entropy_th ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s <EOL> return s <EOL> def get_all_consensus ( read_fn , init_ref , consensus_fn , consens_seq_name , <EOL> hp_correction = False , <EOL> min_iteration = <NUM_LIT:4> , <EOL> max_num_reads = <NUM_LIT> , <EOL> entropy_th = <NUM_LIT> , <EOL> min_cov = <NUM_LIT:8> , <EOL> max_cov = <NUM_LIT> , <EOL> nproc = <NUM_LIT:4> ) : <EOL> g = construct_aln_graph_from_fasta ( read_fn , init_ref , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc ) <EOL> s , c = g . generate_consensus ( min_cov = <NUM_LIT:0> ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if min_iteration > <NUM_LIT:1> : <EOL> for i in range ( min_iteration - <NUM_LIT:2> ) : <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc ) <EOL> s , c = g . generate_consensus ( min_cov = <NUM_LIT:0> ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if hp_correction : <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = False , max_cov = max_cov , nproc = nproc ) <EOL> s = detect_missing ( g , entropy_th = entropy_th ) <EOL> with open ( consensus_fn , "<STR_LIT:w>" ) as f : <EOL> print >> f , "<STR_LIT:>>" + consens_seq_name <EOL> print >> f , s . upper ( ) <EOL> if len ( s ) < <NUM_LIT:100> : <EOL> return s <EOL> g = construct_aln_graph_from_fasta ( read_fn , consensus_fn , max_num_reads = max_num_reads , remove_in_del = True , max_cov = max_cov , nproc = nproc ) <EOL> all_consensus = g . generate_all_consensus ( min_cov = min_cov ) <EOL> return all_consensus <EOL> def output_dag_info ( aln_graph , out_file_name ) : <EOL> with open ( out_file_name , "<STR_LIT:w>" ) as f : <EOL> read_ids = [ ] <EOL> read_ids_set = set ( ) <EOL> for n in sorted_nodes ( aln_graph ) : <EOL> for r , p in n . get_info ( ) : <EOL> if r not in read_ids_set : <EOL> read_ids . append ( r ) <EOL> read_ids_set . add ( r ) <EOL> read_id_to_pos = dict ( ( ( x [ <NUM_LIT:1> ] , x [ <NUM_LIT:0> ] ) for x in enumerate ( list ( read_ids ) ) ) ) <EOL> for read_id in read_ids : <EOL> pileup_pos = read_id_to_pos [ read_id ] <EOL> print >> f , "<STR_LIT:\t>" . join ( [ "<STR_LIT:R>" , "<STR_LIT>" % pileup_pos , read_id ] ) <EOL> backbone_node_to_pos = aln_graph . get_backbone_node_to_pos ( ) <EOL> ne , hne = aln_graph . get_high_entropy_nodes ( coverage_th = <NUM_LIT:0> ) <EOL> node_to_entropy = dict ( [ ( v [ <NUM_LIT:1> ] , "<STR_LIT:\t>" . join ( [ str ( x ) for x in [ v [ <NUM_LIT:2> ] ] + list ( v [ <NUM_LIT:3> ] ) ] ) ) for v in ne ] ) <EOL> data = [ ] <EOL> consensus_pos = <NUM_LIT:0> <EOL> consensus_path = aln_graph . get_consensus_path ( ) <EOL> for n in sorted_nodes ( aln_graph ) : <EOL> s = [ "<STR_LIT:.>" ] * len ( read_id_to_pos ) <EOL> for r , pos in n . get_info ( ) : <EOL> s [ read_id_to_pos [ r ] ] = n . get_base ( ) <EOL> if n . get_base ( ) not in [ "<STR_LIT:B>" , "<STR_LIT:E>" ] : <EOL> bpos = backbone_node_to_pos [ n . get_backbone_node ( ) ] <EOL> entropy_th = <NUM_LIT:0> <EOL> ent = node_to_entropy [ n ] if n in node_to_entropy else <NUM_LIT:0> <EOL> if n . get_base ( ) not in [ "<STR_LIT:B>" , "<STR_LIT:E>" ] and ent >= entropy_th : <EOL> data . append ( ( n . get_ID ( ) , consensus_pos , backbone_node_to_pos [ n . get_backbone_node ( ) ] , "<STR_LIT:+>" if n in consensus_path else "<STR_LIT:->" , "<STR_LIT:+>" if n . get_is_backbone ( ) == True else "<STR_LIT:->" , n . get_base ( ) , "<STR_LIT>" . join ( s ) , n . get_weight ( ) , n . get_backbone_node ( ) . get_coverage ( ) , node_to_entropy [ n ] if n in node_to_entropy else "<STR_LIT:->" ) ) <EOL> if n in consensus_path : <EOL> consensus_pos += <NUM_LIT:1> <EOL> for l in data : <EOL> print >> f , "<STR_LIT:N>" + "<STR_LIT:\t>" + "<STR_LIT:\t>" . join ( [ str ( c ) for c in l ] ) <EOL> aln_edges = aln_graph . get_edges ( ) <EOL> for e_id in aln_edges : <EOL> e = aln_edges [ e_id ] <EOL> print >> f , "<STR_LIT:\t>" . join ( [ "<STR_LIT:E>" , "<STR_LIT>" % e_id , "<STR_LIT>" % e . get_in_node ( ) . get_ID ( ) , "<STR_LIT>" % e . get_out_node ( ) . get_ID ( ) , "<STR_LIT>" % e . get_count ( ) ] ) <EOL> def generate_consensus ( input_fasta_name , <EOL> ref_fasta_name , <EOL> prefix , <EOL> consensus_name , <EOL> hpFix , <EOL> min_iteration , <EOL> max_num_reads , <EOL> entropy_th , <EOL> dump_dag_info , <EOL> min_cov , <EOL> max_cov , <EOL> mark_lower_case , <EOL> nproc ) : <EOL> normalize_fasta ( input_fasta_name , ref_fasta_name , "<STR_LIT>" % prefix , nproc ) <EOL> s = get_consensus ( "<STR_LIT>" % prefix , <EOL> ref_fasta_name , <EOL> "<STR_LIT>" % prefix , <EOL> consensus_name , <EOL> hp_correction = hpFix , <EOL> min_iteration = min_iteration , <EOL> max_num_reads = max_num_reads , <EOL> entropy_th = entropy_th , <EOL> min_cov = min_cov , <EOL> max_cov = max_cov , <EOL> nproc = nproc , <EOL> mark_lower_case = mark_lower_case ) <EOL> if dump_dag_info == True : <EOL> output_dag_info ( g , "<STR_LIT>" % prefix ) <EOL> return s <EOL> def generate_all_consensus ( input_fasta_name , <EOL> ref_fasta_name , <EOL> prefix , <EOL> consensus_name , <EOL> hpFix , <EOL> min_iteration , <EOL> max_num_reads , <EOL> entropy_th , <EOL> dump_dag_info , <EOL> min_cov , <EOL> max_cov , <EOL> mark_lower_case , <EOL> nproc ) : <EOL> normalize_fasta ( input_fasta_name , ref_fasta_name , "<STR_LIT>" % prefix , nproc ) <EOL> s = get_all_consensus ( "<STR_LIT>" % prefix , <EOL> ref_fasta_name , <EOL> "<STR_LIT>" % prefix , <EOL> consensus_name , <EOL> hp_correction = hpFix , <EOL> min_iteration = min_iteration , <EOL> max_num_reads = max_num_reads , <EOL> entropy_th = entropy_th , <EOL> min_cov = min_cov , <EOL> max_cov = max_cov , <EOL> nproc = nproc ) <EOL> return s </s>
