<jupyter_start><jupyter_text>Basic functions<jupyter_code>class gene:
    def __init__(self):
        self.sGeneSym       = ''
        self.sNMID          = ''
        self.sChrID         = ''
        self.nExons         = 0
        self.nExonStartlist = []
        self.nExonEndlist   = []
        self.sStrandDir     = ''
        self.nORF_5p_gidx   = 0
        self.nORF_3p_gidx   = 0
        self.nExonlen       = 0
        self.nU5len         = 0
        self.nU3len         = 0

    def parse_refflat(self,refline):
        sInfolist = refline.strip().replace(' ','\t').split('\t')
        self.sGeneSym = sInfolist[0].upper()
        self.sNMID    = sInfolist[1]
        self.sChrID   = sInfolist[2] ##chr1,,,,chrX,chrY for human
        self.sStrandDir   = sInfolist[3]
        self.nORF_5p_gidx = int(sInfolist[6])
        self.nORF_3p_gidx = int(sInfolist[7])
        self.nExons       = int(sInfolist[8])
        self.nExonStartlist = [int(i) for i in sInfolist[9].split(',') if i != '']
        self.nExonEndlist   = [int(i) for i in sInfolist[10].split(',') if i != '']
        assert (self.nExons == len(self.nExonStartlist)) and (self.nExons == len(self.nExonEndlist))
        self.nExonlen       = sum([end-start for start,end in zip(self.nExonStartlist, self.nExonEndlist)]) 
        
        tmp_exonlen = 0
        for start,end in zip(self.nExonStartlist, self.nExonEndlist):
            if start <= self.nORF_5p_gidx <  end:
                self.nU5len  = tmp_exonlen + (self.nORF_5p_gidx - start)
            if start <= self.nORF_3p_gidx   <= end:
                self.nU3len = self.nExonlen - (tmp_exonlen + (self.nORF_3p_gidx - start))
            tmp_exonlen += (end-start)
        if self.sStrandDir == '-':
            tmp_nU3len  = self.nU5len
            self.nU5len = self.nU3len
            self.nU3len = tmp_nU3len
#########################

def get_df_SARSCOV2_annot(exc_UTR = True):
    df = pd.read_csv(REFFLAT_SARSCOV2,sep='\t',header=None)
    df = df.iloc[:,[0,2,3,6,7]]
    df.columns = ['ORF','chromosome','strand','ORFstart','ORFend']
    df = df.set_index('ORF')
    if exc_UTR:
        df = df[~df.index.str.contains('UTR')]
    return df

def get_dict_refgenes(refdir):
    ref_lines  = [line.strip() for line in open(refdir)]
    chrdict = defaultdict(list)
    
    for line in ref_lines:
        refgene  = gene()
        refgene.parse_refflat(line)
        chrdict['%s%s'%(refgene.sChrID, refgene.sStrandDir)].append(refgene)
    chrdict = dict(chrdict)
    return chrdict

def get_readpos_arr(samplename, assay_prefix, norm_by_nh=True, pos_5end_range=(46,52), 
                    f2_5p_range=None, fetch_range=(0,265), offset_RPF = 12, junction_spanning_only=True, 
                    min_rlen=(32,31,30,29,28,27), as_RPM=False, regions_exclude = [], verbose=False):
    #offset_RPF: non-negative integer or 'readcenter'
    #min_rlen: non-negative integer or tuple 
    #    if tuple, applying minimum rlen for each position to the pos_5end_range, 
    #    assert len(minrlen) == pos_5end_range[1]-pos_5end_range[0])
    fname = BAMDIR_cov2mapped %(assay_prefix,samplename)
    bam = pysam.AlignmentFile(fname)
    arr_f1_3p = np.zeros(len(SARSCOV2_genome))
    arr_f2_5p = np.zeros(len(SARSCOV2_genome))
    arr_RPF   = np.zeros(len(SARSCOV2_genome))
    
    assigned_ct_raw     = 0
    assigned_ct_weighted= 0
    contig = bam.references[0]
    for read in bam.fetch(contig,fetch_range[0],fetch_range[1]):
        cigartuples = read.cigartuples
        pos_5end    = read.reference_start
        rlen        = read.infer_read_length()
        nh          = read.get_tag('NH')
        assigned    = False
        if norm_by_nh:
            norm_ct = 1/nh
        else:
            norm_ct = 1
            
        if not ((pos_5end_range[0]<= pos_5end) and (pos_5end< pos_5end_range[1])):
            continue
            
        if type(min_rlen) == int:
            if rlen<min_rlen:
                continue
        else:
            if rlen<min_rlen[pos_5end-pos_5end_range[0]]:
                continue
        reference_positions = read.get_reference_positions(full_length=True)    
        reference_positions = [pos for pos in reference_positions if pos != None]
        if offset_RPF == 'readcenter':
            offset = rlen//2
        else:
            offset = offset_RPF
            
        refpos = reference_positions[offset]
        
        contain_junction = (3 in [i[0] for i in cigartuples])
        cur_pos = pos_5end
        f1_detected = False
        f1_3p = None
        f2_5p = None
        
        if (not contain_junction):
            if junction_spanning_only:
                continue
            else: 
                pass
        else: ##determination of f1_3p (3' end of the first fragment) and f2_5p (5' end of the other fragments) for junction_spanning reads
            for idx,cigar_tp in enumerate(cigartuples):
                operation, length = cigar_tp
                if operation == 3: #skip
                    f1_detected = True
                    cur_pos += length
                else:
                    if operation == 4 or operation == 1:    # softclip or insertion
                        pass
                    elif operation == 0 or operation == 2:  # match/mismatch or deletion
                        cur_pos += length
                if operation == 0:
                    if not f1_detected: #junction 1st fragment
                        f1_3p = cur_pos

                    else:
                        f2_5p = cur_pos - length
                        break

            assert f1_3p != None
            assert f2_5p != None
            
            if f2_5p_range == None:
                arr_f2_5p[f2_5p] += norm_ct
            else:
                if (f2_5p_range[0]<= f2_5p) and (f2_5p< f2_5p_range[1]):
                    arr_f2_5p[f2_5p] += norm_ct
                else: ##roll back if the read doesn't meet f2_5p range
                    continue
            arr_f1_3p[f1_3p] += norm_ct
        
        arr_RPF[refpos]      += norm_ct
        assigned_ct_raw      += 1
        assigned_ct_weighted += norm_ct
    
    bam.close()
    arr_include = np.ones(len(arr_RPF))
    for s,e in regions_exclude:
        arr_include[s:e] = 0 
    arr_RPF = arr_RPF*arr_include
    
    total_mapped_reads = 0
    viral_mapped_reads = 0
    host_mapped_reads  = 0
    if as_RPM:
        #virus reads
        bam = pysam.AlignmentFile(fname)
        for read in bam.fetch():
            viral_mapped_reads += (1/(read.get_tag('NH')))
        bam.close()
        #human reads
        fname_host = BAMDIR_hostmapped %(assay_prefix,samplename)
        bam = pysam.AlignmentFile(fname_host)
        for read in bam.fetch():
            host_mapped_reads += (1/(read.get_tag('NH')))
        bam.close()
        total_mapped_reads = host_mapped_reads + viral_mapped_reads
        arr_f1_3p /= (total_mapped_reads / 1e+06)
        arr_f2_5p /= (total_mapped_reads / 1e+06)
        arr_RPF   /= (total_mapped_reads / 1e+06)
    
    report = [total_mapped_reads,viral_mapped_reads,host_mapped_reads, assigned_ct_raw, assigned_ct_weighted]
    report = pd.Series(report)
    report.index = ['total_mapped','viral_mapped','host_mapped','region_assigned(raw)','region_assigned(weighted)']
    if verbose:
        print(fname)
        display(report)
    return arr_f1_3p, arr_f2_5p, arr_RPF<jupyter_output><empty_output><jupyter_text>Figure 1e- QC analysis  Metagene analysis (which includes analysis of read enrichment at ORF start) (Figure 1e)<jupyter_code>def calc_metagene_array(samplename,assay_prefix='RPFSeq', genome='SARSCOV2', offset_RPF = 12, 
                        flanking=10, align='5p', nreads_cutoff = 50, verbose = True):
    if genome == 'SARSCOV2':
        fname = BAMDIR_cov2mapped %(assay_prefix,samplename)
        refdir= REFFLAT_SARSCOV2
    elif genome in ['hg38','chlSab2']:
        fname = BAMDIR_hostmapped %(assay_prefix,samplename)
        if genome == 'hg38':
            refdir = REFFLAT_hg38
        else:
            refdir = REFFLAT_chlSab2
    else:
        print('Wrong genome name:',genome)
        return None
    
    chr_ref_dict = get_dict_refgenes(refdir) #{'%s%s'%(refgene.sChrID, refgene.sStrandDir): [gene,,]}
    if genome == 'SARSCOV2': #modify annotation for flanking region
        for sChrID_strand,refgene_list in chr_ref_dict.items():
            for refgene in refgene_list:
                refgene.nExonStartlist[0]  -= flanking
                refgene.nExonEndlist[-1]   += flanking
                
                refgene.nExonlen += (flanking*2)
                refgene.nU5len    = flanking
                refgene.nU3len    = flanking
        
    bam      = pysam.AlignmentFile(fname)
    arr_RPF_ORF_list = []
    for sChrID_strand,refgene_list in chr_ref_dict.items():
        if verbose:
            print(sChrID_strand, time.ctime())
        sChrID     = sChrID_strand[:-1]
        sStrandDir = sChrID_strand[-1]
        gene_negstrand = (sStrandDir == '-')
        
        for refgene in refgene_list:
            if refgene.sGeneSym in ["5'UTR", "3'UTR", "ORF1b"]: #filter out those UTR annotations in SARSCOV2 refflat
                continue
            
            exon_s_list = refgene.nExonStartlist
            exon_e_list = refgene.nExonEndlist
            
            arr_RPF     = np.zeros(refgene.nExonlen,dtype=np.float32)
            offset_arr  = 0
            for exon_s, exon_e in zip(exon_s_list,exon_e_list):
                for read in bam.fetch(sChrID, exon_s,exon_e):
                    rlen        = read.infer_read_length()
                    nh          = read.get_tag('NH')
                    
                    if read.is_reverse != gene_negstrand:
                        continue
                    if offset_RPF == 'readcenter':
                        offset = rlen//2
                    else:
                        offset = offset_RPF
                    reference_positions = read.get_reference_positions(full_length=False)    
                    if None in reference_positions:
                        print('None in read reference position?')
                        print(reference_positions)
                        print(read.cigarstring)
                        return None
                    if gene_negstrand:
                        refpos   = reference_positions[-(offset+1)]
                    else:
                        refpos   = reference_positions[offset]
                    
                    nh = read.get_tag('NH')
                    if (exon_s <= refpos) and (refpos < exon_e):
                        arr_RPF[refpos - exon_s + offset_arr] += (1/nh)
                offset_arr += (exon_e - exon_s)
                
            assert offset_arr == len(arr_RPF)
            
            if sChrID == '-':
                arr_RPF = arr_RPF[::-1]
            
            flanking_5p = min(refgene.nU5len,flanking)
            arr_RPF_ORF_flank = np.concatenate([np.zeros(flanking-flanking_5p),
                                                arr_RPF[refgene.nU5len-flanking_5p:]])
            
            flanking_3p = min(refgene.nU3len,flanking)
            if (refgene.nU3len-flanking_3p)>0:
                arr_RPF_ORF_flank = arr_RPF_ORF_flank[:-(refgene.nU3len-flanking_3p)]
            arr_RPF_ORF_flank = np.concatenate([arr_RPF_ORF_flank,
                                                np.zeros(flanking-flanking_3p)])
            
            nreads_ORFmapped = arr_RPF_ORF_flank[flanking:-flanking].sum() #divide by ORF RPF signal sum
            if nreads_ORFmapped < nreads_cutoff:
                continue
            arr_RPF_ORF_flank /= nreads_ORFmapped
            
            len_ORF = len(arr_RPF) - refgene.nU5len - refgene.nU3len
            assert len_ORF+(2*flanking) == len(arr_RPF_ORF_flank)
            arr_RPF_ORF_list.append(arr_RPF_ORF_flank)
            
    max_ORF_flank_len = max([len(arr) for arr in arr_RPF_ORF_list])
    if align == '5p':
        arr_RPF_ORF_list = [np.concatenate([arr,np.full(max_ORF_flank_len-len(arr), np.nan)]) for arr in arr_RPF_ORF_list]
    elif align == '3p':
        arr_RPF_ORF_list = [np.concatenate([np.full(max_ORF_flank_len-len(arr), np.nan),arr]) for arr in arr_RPF_ORF_list]
        
    arr_mean_freq = np.array(arr_RPF_ORF_list)
    arr_mean_freq = np.nanmean(arr_mean_freq,axis=0)
    bam.close()
    return arr_mean_freq

def plot_metagene_analysis(hpi,assay_prefix='RPFSeq', sample_prefix_list=['RPF'],genome='SARSCOV2',
                           offset_RPF = 12, flanking = 10, align = '5p', nreads_cutoff = 100, 
                           verbose=False, plot_range = 50,ylim=None,OutFigname = '', ax=None,show_fig=True):
    if hpi == '48h':
        reps = 3
    else:
        reps = 2
        
    arr_merged = []
    for sample_prefix in sample_prefix_list:
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            arr_mean_freq = calc_metagene_array(samplename, assay_prefix=assay_prefix, genome=genome, 
                                                offset_RPF = offset_RPF, flanking=flanking, align=align, 
                                                nreads_cutoff = nreads_cutoff, verbose = verbose)
            if align == '5p':
                arr_merged.append(arr_mean_freq[:plot_range])
            elif align == '3p':
                arr_merged.append(arr_mean_freq[-plot_range:])
                
    arr_merged = np.array(arr_merged).mean(axis=0)*100
    if align == '5p':
        x   = np.arange(-flanking,plot_range-flanking)
        xticks      = [0,10,20,30]
    elif align == '3p':
        x   = np.arange(-(plot_range-flanking),flanking)
        xticks      = [-30,-20,-10,0]
        
    if ax == None:
        fig, ax = plt.subplots(figsize=(4,3))
    if genome == 'SARSCOV2': #red
        bar_colors = [['#ff6666','#fc8d59','#fdcc8a'][i%3] for i in x]
    else:                    #green
        bar_colors = [['#25ad62','#74c476','#bae4b3'][i%3] for i in x]
    ax.bar(x,arr_merged,width=1,color=bar_colors)
    ax.set_xticks(xticks)
    ax.set_xlim(x[0],x[-1])
    ax.set_xlabel('Relative position')
    ax.set_ylabel('Relative fraction (%)')
    if ylim!= None:
        ax.set_ylim(ylim)
    
    if OutFigname != '':
        plt.tight_layout()
        plt.savefig(FIGDIR %OutFigname)
        plt.close()
        
    elif show_fig:
        plt.tight_layout()
        plt.show()
        plt.close()
    
    return arr_merged

###################################
def plot_merged_main(hpi='4h',ylim_list=[[]],OutFigname=''):
    assay_prefix_list  = ['mRNASeq','RPFSeq','RPFSeq']
    sample_prefix_list_list = [['RPFpaired','QTIpaired'],['RPF'],['QTI']]
    genome_list        = ['hg38','SARSCOV2']
    
    fig, axes= plt.subplots(nrows=len(genome_list),
                            ncols=len(sample_prefix_list_list), 
                            figsize=(3*len(sample_prefix_list_list),2*len(genome_list)))
    df_source = pd.DataFrame(columns =['%s_%s' %(genome,sample_prefix_list[0]) for sample_prefix_list in sample_prefix_list_list 
                                       for genome in genome_list],
                             index = np.arange(-10,40))
    for idx_row,genome in enumerate(genome_list):
        for idx_col,sample_prefix_list in enumerate(sample_prefix_list_list):
            print(genome, sample_prefix_list, time.ctime())
            assay_prefix = assay_prefix_list[idx_col]
            ax   = axes[idx_row,idx_col]
            ylim = ylim_list[idx_row][idx_col]
            arr  = plot_metagene_analysis(hpi,assay_prefix=assay_prefix, sample_prefix_list=sample_prefix_list,
                                   genome=genome, offset_RPF = 12, flanking = 10, align = '5p', nreads_cutoff = 50, 
                                   verbose=False, plot_range = 50,ylim=ylim, ax=ax, show_fig = False)
            df_source['%s_%s' %(genome,sample_prefix_list[0])] = arr
    plt.tight_layout()
    plt.savefig(FIGDIR %(OutFigname))
    plt.close()
    df_source.to_csv(SOURCEDATADIR %(OutFigname.replace('pdf','tsv')), sep = '\t')
    return None
ylim_list_4h = [[(0.0,1.0),(0.0,6.0), (0.0,6.0) ], #hg38
                [(0.0,4.0),(0.0,20.0),(0.0,20.0)]] #SARSCOV2
plot_merged_main(hpi='4h', ylim_list = ylim_list_4h,
                 OutFigname='Fig1e_metagene_4h.pdf')
ylim_list_36h = [[(0.0,1.0),(0.0,4.0),(0.0,4.0) ], #hg38
                 [(0.0,4.0),(0.0,8.0),(0.0,8.0)]] #SARSCOV2
plot_merged_main(hpi='36h', ylim_list = ylim_list_36h,
                 OutFigname='Fig1e_metagene_36h.pdf')<jupyter_output>hg38 ['RPFpaired', 'QTIpaired'] Mon Sep 27 15:19:55 2021
hg38 ['RPF'] Mon Sep 27 15:28:17 2021
hg38 ['QTI'] Mon Sep 27 15:30:40 2021
SARSCOV2 ['RPFpaired', 'QTIpaired'] Mon Sep 27 15:31:47 2021
SARSCOV2 ['RPF'] Mon Sep 27 15:31:49 2021
SARSCOV2 ['QTI'] Mon Sep 27 15:31:49 2021
hg38 ['RPFpaired', 'QTIpaired'] Mon Sep 27 15:31:50 2021
hg38 ['RPF'] Mon Sep 27 15:35:24 2021
hg38 ['QTI'] Mon Sep 27 15:36:08 2021
SARSCOV2 ['RPFpaired', 'QTIpaired'] Mon Sep 27 15:36:46 2021
SARSCOV2 ['RPF'] Mon Sep 27 15:40:11 2021
SARSCOV2 ['QTI'] Mon Sep 27 15:40:21 2021<jupyter_text>Read length histogram (Supplementary Figure 1d)<jupyter_code>def hist_readlength_cmp(assay_prefix='RPFSeq',sample_prefix_list=['RPF'],hpi='48h',
                        genome_list = ['SARSCOV2','host'],
                        color_h = 'g', color_v = 'r',xlim=[18,40],
                        merge_reps = True,OutF='',OutFigname='', ax = None, show_fig = False):
    summary= []
    bins   = np.arange(*xlim)
    hist_group_keys = []
    hist_group_dict = {}
    reps = 3 if hpi == '48h' else 2
    
    for genome in genome_list:
        if genome == 'SARSCOV2':
            BAMDIR = BAMDIR_cov2mapped
        else:
            BAMDIR = BAMDIR_hostmapped
        
        for sample_prefix in sample_prefix_list:
            if merge_reps:
                hist_group_key = f'{genome}/{sample_prefix}_{hpi}'
                hist_group_keys.append(hist_group_key)
                hist_group_dict[hist_group_key] = []
            for rep in range(1, 1+reps):
                samplename = f'{sample_prefix}_{hpi}_rep{rep}'
                bam = pysam.AlignmentFile(BAMDIR %(assay_prefix,samplename))    #%(RPFSeq'','RPF_2h_rep1')
                rlen_list = []
                for read in bam.fetch():
                    rlen = read.infer_read_length()
                    rlen_list.append(rlen)
                bam.close()
                if not merge_reps:
                    hist_group_key = f'{genome}/{sample_prefix}_{hpi}_rep{rep}'
                    hist_group_keys.append(hist_group_key)
                    hist_group_dict[hist_group_key] = []
                
                hist_group_dict[hist_group_key] += rlen_list
    
    if ax == None:
        fig, ax = plt.subplots(figsize=(3,3))
    sub_df_source = pd.DataFrame(columns = hist_group_keys)
    for hist_group_key in hist_group_keys:
        rlen_list  = hist_group_dict[hist_group_key]
        genome, _  = hist_group_key.split('/')
        if genome == 'host':
            color = color_h
        else:
            color = color_v
        
        rlen_arr = np.array(rlen_list)
        n_reads = len(rlen_arr)
        
        summary.append([hist_group_key,n_reads,np.mean(rlen_arr),np.median(rlen_arr), np.std(rlen_arr)])        
        n, _, _ = ax.hist(rlen_arr,bins=bins,histtype='step',label = hist_group_key,
                          density=True,color=color,  align = 'left')
        sub_df_source[hist_group_key] = n
        
    ax.set_xlabel('Read length (nt)')
    ax.set_ylabel('Density')
    ax.set_xlim(*xlim)
    ax.set_xticks(np.arange(*xlim,2))
    ax.set_ylim(0,0.55)
    
    if merge_reps:
        if assay_prefix == 'mRNASeq':
            ax.legend(ncol=1,loc='upper left')
        else:
            ax.legend(ncol=1,loc='upper right') 
    else: 
        pass
    
    if OutFigname != '':
        plt.tight_layout()
        plt.savefig(FIGDIR %OutFigname)
        plt.close()
        
    elif show_fig:
        plt.show()
        display(summary_df)
        plt.close()
        
    summary_df = pd.DataFrame(summary)
    summary_df.columns = ['samplename','n','mean_len','median_len','std_len']
    summary_df = summary_df.set_index('samplename')
    summary_df.to_csv(RESULTDIR %OutF, sep = '\t')
    return sub_df_source

def draw_readhist_multipanel(OutFigname='',transposed=True):
    #(RPF, QTI) * (hg38, SARSCOV2) * (0~48h)
    hpi_list     = ['%dh' %i for i in [4,12,16,24,36,48]]+['48hCaco2']
    assay_prefix = 'RPFSeq'
    color_h = '#25ad62'
    color_v = '#ff6666'
    xlim = [20,33]  
    
    if transposed:
        fig, axes  = plt.subplots(ncols=len(hpi_list),nrows=2, figsize=(3*len(hpi_list),5))
    else:
        fig, axes  = plt.subplots(nrows=len(hpi_list),ncols=2, figsize=(6,2.5*len(hpi_list)))
        
    df_source = []
    for idx_row,hpi in enumerate(hpi_list):
        for idx_col,sample_prefix in enumerate(['RPF','QTI']):
            print(hpi,sample_prefix,time.ctime())
            if transposed:
                ax = axes[idx_col,idx_row]
            else:
                ax = axes[idx_row,idx_col]
            sub_df_source = hist_readlength_cmp(assay_prefix=assay_prefix,sample_prefix_list=[sample_prefix],hpi=hpi,
                                genome_list = ['SARSCOV2','host'],
                                color_h = color_h, color_v = color_v,xlim=xlim,
                                merge_reps = True, OutF=f'read_hist_summary_{sample_prefix}_{hpi}.tsv',
                                OutFigname='', ax=ax, show_fig=False)
            df_source.append(sub_df_source)
    df_source = pd.concat(df_source,axis=1)*100
    
    plt.tight_layout()
    plt.savefig(FIGDIR %(OutFigname))
    df_source.to_csv(SOURCEDATADIR %(OutFigname.replace('pdf','tsv')), sep = '\t')
    plt.close()
    return None
draw_readhist_multipanel(OutFigname = 'SFig1d_readhist_merged_transposed.pdf',transposed=True)<jupyter_output>4h RPF Mon Sep 27 15:40:27 2021
4h QTI Mon Sep 27 15:40:39 2021
12h RPF Mon Sep 27 15:40:45 2021
12h QTI Mon Sep 27 15:40:49 2021
16h RPF Mon Sep 27 15:40:51 2021
16h QTI Mon Sep 27 15:40:55 2021
24h RPF Mon Sep 27 15:40:56 2021
24h QTI Mon Sep 27 15:40:59 2021
36h RPF Mon Sep 27 15:41:00 2021
36h QTI Mon Sep 27 15:41:02 2021
48h RPF Mon Sep 27 15:41:03 2021
48h QTI Mon Sep 27 15:41:11 2021
48hCaco2 RPF Mon Sep 27 15:41:15 2021
48hCaco2 QTI Mon Sep 27 15:41:18 2021<jupyter_text>Triplet periodicity (Supplementary Figure 1e)<jupyter_code>def calc_triplet_periodicity(samplename,assay_prefix = 'RPFSeq',genome = 'hg38', offset=12, target_rlen = None,
                             offset_from_5p = True, verbose=False, return_codon1pos1 = False):
    
    if genome == 'SARSCOV2':
        BAMDIR = BAMDIR_cov2mapped
    else:
        BAMDIR = BAMDIR_hostmapped
        
    if genome == 'hg38':
        refdir = REFFLAT_hg38 
    elif genome == 'SARSCOV2':
        refdir = REFFLAT_SARSCOV2 
    elif genome == 'chlSab2':
        refdir = REFFLAT_chlSab2
    else:
        return None
    
    chrdict = get_dict_refgenes(refdir)
    
    bam = pysam.AlignmentFile(BAMDIR %(assay_prefix,samplename), 'rb')
    triplet_counts = np.zeros(3,dtype=np.float32)
    codon1pos1 = 0.0
    for sChrID_strand,refgene_list in chrdict.items():
        if verbose:
            print(sChrID_strand, time.ctime())
        sChrID     = sChrID_strand[:-1]
        sStrandDir = sChrID_strand[-1]
        gene_negstrand = (sStrandDir == '-')
        
        for refgene in refgene_list:
            if refgene.sGeneSym in ["5'UTR", "3'UTR"]: #filter out those UTR annotations in SARSCOV2 refflat
                continue
            len_ORF = refgene.nExonlen - refgene.nU5len - refgene.nU3len
            if len_ORF%3 != 0:
                continue
            exon_s_list = refgene.nExonStartlist
            exon_e_list = refgene.nExonEndlist
            arr_RPF     = np.zeros(refgene.nExonlen,dtype=np.float32)
            offset_arr  = 0
            for exon_s, exon_e in zip(exon_s_list,exon_e_list):
                for read in bam.fetch(sChrID, exon_s,exon_e):
                    if target_rlen != None:
                        rlen = read.infer_read_length()
                        if rlen!=target_rlen:
                            continue
                    if read.is_reverse != gene_negstrand:
                        continue
                    reference_positions = read.get_reference_positions(full_length=False)    
                    if None in reference_positions:
                        print('None in read reference position?')
                        print(reference_positions)
                        print(read.cigarstring)
                        return None
                    if gene_negstrand:
                        if offset_from_5p == True:
                            refpos   = reference_positions[-(offset+1)]
                        else:
                            refpos   = reference_positions[offset]
                    else:
                        if offset_from_5p == True:
                            refpos   = reference_positions[offset]
                        else:
                            refpos   = reference_positions[-(offset+1)]
                    #print(ribo_5p,read.infer_query_length(),read.is_reverse,reference_positions)
                    #return None
                    
                    nh = read.get_tag('NH')
                    if (exon_s <= refpos) and (refpos < exon_e):
                        arr_RPF[refpos-exon_s+offset_arr] += (1/nh)
                offset_arr += (exon_e - exon_s)
                
            assert offset_arr == len(arr_RPF)
            if sChrID == '-':
                arr_RPF = arr_RPF[::-1]
            arr_RPF_ORF = arr_RPF[refgene.nU5len:]
            if refgene.nU3len>0: 
                arr_RPF_ORF = arr_RPF_ORF[:-refgene.nU3len]
            codon1pos1 += arr_RPF_ORF[0]
            
            to_add = arr_RPF_ORF.reshape((len_ORF//3,3)).sum(axis=0)
            triplet_counts += to_add
            #print(refgene.sGeneSym,to_add/to_add.sum())

    triplet_periodicity = triplet_counts/(triplet_counts.sum())
    
    if verbose:
        print(samplename,genome, triplet_counts.sum(), triplet_periodicity, sep = '\t')
    bam.close()
    if return_codon1pos1:
        return triplet_counts, triplet_periodicity, codon1pos1
    else:
        return triplet_counts, triplet_periodicity


def calc_plot_periodicity(assay_prefix = 'RPFSeq', sample_prefix='RPF',load_precalc=False,
                          OutF='031521_RPF_triplet.merged.tsv', OutFigname=''):
    hpi_list = ['%dh' %i for i in [4,12,16,24,36,48]]+['48hCaco2']
    df_merged  = []
    if load_precalc:
        df_merged = pd.read_csv(RESULTDIR %(OutF),sep='\t')
    else:
        for hpi in hpi_list:
            print(hpi,time.ctime())
            for genome in ['hg38','SARSCOV2']:
                merged_counts = np.zeros(3,dtype=np.float32)
                reps = 3 if (hpi == '48h') else 2
                for rep in range(1, 1+reps):
                    samplename = f'{sample_prefix}_{hpi}_rep{rep}'
                    triplet_counts, triplet_periodicity = calc_triplet_periodicity(samplename, assay_prefix = assay_prefix, 
                                                                                   genome = genome, offset=12, target_rlen = None,
                                                                                   offset_from_5p = True, verbose=False, 
                                                                                   return_codon1pos1 = False)
                    merged_counts += triplet_counts

                df_merged.append([hpi,genome,*merged_counts,*(merged_counts/merged_counts.sum()*100)])

        df_merged  = pd.DataFrame(df_merged)
        df_merged.columns = ['hpi','genome','pos1','pos2','pos3','pct1','pct2','pct3']
        df_merged.to_csv(RESULTDIR %(OutF), sep='\t',index=False)   
    ###########
    df_merged = df_merged.set_index('genome')
    
    df_merged = df_merged[df_merged['hpi'].isin(hpi_list)]
    df_merged.index = df_merged.index.str.replace('hg38','host')
    df_merged.index = df_merged.index.str.replace('chlSab2','host')
    fig = plt.figure(figsize=(len(hpi_list)+2,4))
    ax = fig.add_subplot(111)
    x = np.arange(len(hpi_list))
    bar_width=0.12
    gap_hg38_SARSCOV2 = 0.06
    xpos_offset = -bar_width*2.5 - gap_hg38_SARSCOV2*0.5 
    colors = {}
    
    colors['host']     = ['#25ad62','#74c476','#bae4b3']
    colors['SARSCOV2'] = ['#ff6666','#fc8d59','#fdcc8a']
    
    for genome in ['host','SARSCOV2']:
        for i in range(3):
            pct = f'pct{i+1}'
            val = df_merged.loc[genome,pct]
            ax.bar(x+xpos_offset,val,width = bar_width,label = f'{genome} nt {i+1}', color= colors[genome][i])
            xpos_offset += bar_width
        xpos_offset += gap_hg38_SARSCOV2
    
    
    ax.set_xticks(x)
    ax.set_xticklabels(hpi_list)
    ax.set_xlabel('hpi')
    ax.set_xlim((-1,len(hpi_list)+1))
    ax.set_ylabel('Fraction of reads (%)')
    ax.set_ylim(top=70)
    ax.legend(loc='upper right')
    
    plt.tight_layout()
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
        
    plt.close()
    
    return None

calc_plot_periodicity(assay_prefix = 'RPFSeq', sample_prefix='RPF',load_precalc=False,
                          OutF='RPF_triplet.merged.tsv', OutFigname='SFig1e_triplet_RPF.pdf')
calc_plot_periodicity(assay_prefix = 'RPFSeq', sample_prefix='QTI',load_precalc=False,
                          OutF='QTI_triplet.merged.tsv', OutFigname='SFig1e_triplet_QTI.pdf')<jupyter_output>4h Mon Sep 27 15:41:21 2021
12h Mon Sep 27 15:43:28 2021
16h Mon Sep 27 15:44:29 2021
24h Mon Sep 27 15:45:28 2021
36h Mon Sep 27 15:46:21 2021
48h Mon Sep 27 15:47:11 2021
48hCaco2 Mon Sep 27 15:49:47 2021
4h Mon Sep 27 15:50:51 2021
12h Mon Sep 27 15:51:52 2021
16h Mon Sep 27 15:52:43 2021
24h Mon Sep 27 15:53:22 2021
36h Mon Sep 27 15:53:57 2021
48h Mon Sep 27 15:54:37 2021
48hCaco2 Mon Sep 27 15:56:12 2021<jupyter_text>Figure 2b- Quantification of viral ORF expression at mRNA-, RPF-, QTI-seq level- mRNA-seq: RPM of reads starting from the very first genomic location    - ORF10: aggregated all junctions within 100nt region upstream of ORF10- RPF-seq, QTI-seq, and TE level: read __Methods__ section of the paper  Plotting function<jupyter_code>def plot_timecourse_exp(exp_df, hpi_list =[], log10=True, ylim=None, psc = 0.0, grey_exceptS=False,
                        ORFs_to_highlight=[], include_3UTR=False, OutFigname = ''):
    
    global ORF_palette
    plt.style.use('default')
    
    fig = plt.figure(figsize=(4,3))
    ORFlist = list(exp_df.index)
    
    exp_df = exp_df[hpi_list]
    if include_3UTR:
        assert '3\'UTR' in ORFlist
    else:
        if '3\'UTR' in ORFlist:
            ORFlist.remove('3\'UTR')
    x = np.arange(len(exp_df.columns))
    for ORF in ORFlist:
        exp = exp_df.loc[ORF]
        
        if log10:
            exp = np.log10(exp+psc) #psc: pseudocount
        
        if ORF in ORFs_to_highlight:
            linewidth = 2
        else:
            linewidth = 1.5
            
        linestyle = '-'
        if ORF == 'S_after': ## for S_after (incorporating RPF signals from TIS-L)
            label    = 'S, including\nTIS-L reads'
        else:
            if ('S_after' in ORFlist) and (ORF == 'S'):
                linestyle='--'
            label    = ORF
        
        if grey_exceptS:
            color = '#e8e8e8'
            if ORF in ['S','S_after']:
                color = ORF_palette['S']
        elif ORF == 'ORF10_NmRNA': #ORF10_NmRNA: calculation of ORF10 TE by dividing N mRNA level
            linestyle='--'
            color = ORF_palette['ORF10']
        else:
            color = ORF_palette[ORF]
        plt.plot(x,exp,label=label,color=color,linewidth=linewidth,linestyle=linestyle)
        plt.text(x[-1],exp[-1],ORF)
    if ylim != None:
        plt.ylim(*ylim)
    plt.xlim(x.min(),x.max())
    plt.legend(loc='upper left')
    plt.xlabel('hpi')
    
    plt.xticks(x,exp_df.columns.str.replace('h',''))
    
    if log10:
        plt.ylabel('$\mathregular{log_{10}}$ (RPM+1)')
    
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
        
    else:
        plt.show()
    plt.close()
    return None<jupyter_output><empty_output><jupyter_text>mRNA-seq quantification<jupyter_code>def calc_timecourse_mRNAexp(hpi_list, assay_prefix='mRNASeq', as_RPM=True, ORF10_range=(29457,29557),
                           OutF = 'mRNA_quantification.tsv'):
    global CANONICAL_sgRNA_junc_dict

    exp_df = []
    for hpi in hpi_list:
        arr_f2_5p_list   = []
        ORF1a_quant_list = [] #gRNA amount
        reps = 3 if hpi == '48h' else 2
        
        for paired_assay in ['RPF','QTI']:
            for rep in range(1,reps+1):
                samplename = f'{paired_assay}paired_{hpi}_rep{rep}'
                arr_f1_3p, arr_f2_5p, arr_r5p = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                                                pos_5end_range=(0,1), f2_5p_range=None, fetch_range=(0,2),
                                                                offset_RPF = 0,junction_spanning_only=False,
                                                                min_rlen=80,as_RPM=as_RPM)
                arr_f2_5p_list.append(arr_f2_5p)
                
                ORF1a_quant = arr_r5p.sum() - arr_f2_5p.sum()
                ORF1a_quant_list.append(ORF1a_quant)

        avg_arr_f2_5p   = np.array(arr_f2_5p_list).mean(axis=0)
        sgORFs = ['S','ORF3a','E','M','ORF6','ORF7a','ORF7b','ORF8','N']
        result_list = [['ORF1a',np.mean(ORF1a_quant_list)]]
        for ORF in sgORFs:
            f2_5p = CANONICAL_sgRNA_junc_dict[ORF][1] #TRS junction position at body
            result_list.append([ORF,avg_arr_f2_5p[f2_5p]])
            
        #Manually added ORF10 noncanonical sgRNA quantification
        result_list.append(['ORF10',sum([avg_arr_f2_5p[f2_5p] for f2_5p in range(*ORF10_range)]) ] )
        result_df = pd.DataFrame(result_list)
        result_df.columns = ['ORF',hpi]
        result_df = result_df.set_index('ORF')
        print(hpi,time.ctime())
        exp_df.append(result_df)
    exp_df = pd.concat(exp_df,axis=1)
    display(exp_df)
    exp_df.to_csv(RESULTDIR %(OutF),sep='\t')
    return None

'''mRNA-seq quantification'''
hpi_list   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]+['48h']
calc_timecourse_mRNAexp(hpi_list, OutF = 'mRNA_quantification.tsv')

hpi_list_timecourse = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]
mRNAexp_df = pd.read_csv(RESULTDIR %('mRNA_quantification.tsv'), sep = '\t', index_col=0)
plot_timecourse_exp(mRNAexp_df,hpi_list = hpi_list_timecourse,log10=True,ylim=(0,3), psc = 1.0,
                    include_3UTR = False, OutFigname='Fig2b_mRNA.pdf')<jupyter_output>0h Mon Sep 27 15:58:06 2021
1h Mon Sep 27 15:59:18 2021
2h Mon Sep 27 16:00:28 2021
4h Mon Sep 27 16:01:23 2021
12h Mon Sep 27 16:02:03 2021
16h Mon Sep 27 16:02:33 2021
24h Mon Sep 27 16:03:05 2021
36h Mon Sep 27 16:03:39 2021
48h Mon Sep 27 16:05:24 2021<jupyter_text>RPF-, QTI-seq quantification, and PRF(programmed frameshift) of ORF1ab<jupyter_code>def measure_PRF_ORF1ab(samplename, assay_prefix = 'RPFSeq', offset_RPF=12, as_RPM=True):
    ORF1a_range_nonovl = [265,13467]   #real end idx of ORF1a= 13483 
    ORF1b_range_nonovl = [13483,21555] #real start idx of ORF1b= 13467 
    
    _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                            pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                            offset_RPF = offset_RPF,junction_spanning_only=False,
                                            min_rlen=0,as_RPM=as_RPM)
    
    density_ORF1a = arr_RPF[ORF1a_range_nonovl[0]:ORF1a_range_nonovl[1]].mean() #avg.RPM across the ORF
    density_ORF1b = arr_RPF[ORF1b_range_nonovl[0]:ORF1b_range_nonovl[1]].mean()
    density_ratio = density_ORF1b/density_ORF1a
    return density_ORF1a, density_ORF1b, density_ratio

def calc_timecourse_PRF_ORF1ab(hpi_list, assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                               as_RPM = True, OutF = 'ORF1ab_PRF.tsv'):
    PRF_df = []
    for hpi in hpi_list:
        reps = 3 if hpi == '48h' else 2
        avg_ORF1a, avg_ORF1b, avg_ratio = 0.0, 0.0, 0.0    
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            density_ORF1a,density_ORF1b,density_ratio = measure_PRF_ORF1ab(samplename, 
                                                                           assay_prefix = assay_prefix, 
                                                                           offset_RPF = offset_RPF, 
                                                                           as_RPM = as_RPM)
            avg_ORF1a += density_ORF1a
            avg_ORF1b += density_ORF1b
            
        avg_ORF1a = (avg_ORF1a*1000)/reps # -> RPKM
        avg_ORF1b = (avg_ORF1b*1000)/reps # -> RPKM
        avg_ratio = avg_ORF1b/avg_ORF1a
        PRF_df.append([hpi,avg_ORF1a,avg_ORF1b,avg_ratio])
        print(hpi,time.ctime())
        
    PRF_df = pd.DataFrame(PRF_df)
    PRF_df.columns = ['hpi','density_ORF1a','density_ORF1b','PRFratio']
    PRF_df = PRF_df.set_index('hpi')
    display(PRF_df)
    PRF_df.to_csv(RESULTDIR %(OutF),sep='\t')
    return None

def calc_timecourse_RPFexp(hpi_list, assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                           quant_range = (0,15), as_RPM = True, df_precalc_ORF1ab_PRFratio = None, 
                           OutF = 'RPF_quantification.tsv'):
    
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = False)
    df_SARSCOV2_annot = df_SARSCOV2_annot.drop(index='5\'UTR')
    exp_df = pd.DataFrame(index=df_SARSCOV2_annot.index, columns = hpi_list)
    for hpi in hpi_list:
        arr_RPF_list  = []
        reps = 3 if hpi == '48h' else 2
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                            pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                            offset_RPF = offset_RPF,junction_spanning_only=False,
                                            min_rlen=0,as_RPM=as_RPM)
            arr_RPF_list.append(arr_RPF)
        
        avg_arr_RPF = np.array(arr_RPF_list).mean(axis=0)
        
        for ORF,row in df_SARSCOV2_annot.iterrows():
            ORFstart = row['ORFstart']
            exp = avg_arr_RPF[ORFstart+quant_range[0] : ORFstart+quant_range[1]].sum()
            exp_df.at[ORF,hpi] = exp
            
        if type(df_precalc_ORF1ab_PRFratio) == pd.core.frame.DataFrame: #ORF1b estimation
            PRFratio = df_precalc_ORF1ab_PRFratio.at[hpi,'PRFratio']
            exp_df.at['ORF1b',hpi] = PRFratio*exp_df.at['ORF1a',hpi]
        print(hpi,time.ctime())
        
    display(exp_df)
    exp_df.to_csv(RESULTDIR %(OutF),sep='\t')        
    return None
'''RPF-seq quantification, wothout TIS-L S'''
#ORF1ab programmed frameshift ratio calculation
hpi_list   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]+['48h','48h0008U','24hVero','48hCaco2', '48hCalu3cured']
print('ORF1ab PRF ratio')
calc_timecourse_PRF_ORF1ab(hpi_list, assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                           as_RPM = True, OutF = 'RPF_ORF1ab_PRF.tsv')

#RPF RPM calculation
print('RPF expression')
hpi_list   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]+['48h','48h0008U','24hVero','48hCaco2','48hCalu3cured']
df_precalc_ORF1ab_PRFratio = pd.read_csv(RESULTDIR %('RPF_ORF1ab_PRF.tsv'), sep = '\t',index_col=0)
calc_timecourse_RPFexp(hpi_list, assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                       quant_range = (0,15), as_RPM = True, 
                       df_precalc_ORF1ab_PRFratio = df_precalc_ORF1ab_PRFratio, 
                       OutF = 'RPF_quantification_woTISL_S.tsv')
#Plotting
RPFexp_df = pd.read_csv(RESULTDIR %('RPF_quantification_woTISL_S.tsv'), sep = '\t', index_col=0)
hpi_list_timecourse   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]
plot_timecourse_exp(RPFexp_df,hpi_list = hpi_list_timecourse, log10=True,
                    ylim=(0.0,4.0),psc = 1.0, include_3UTR = True,
                    OutFigname='Fig2b_RPF_woTISL_S.pdf')     
'''QTI-seq quantification, without TIS-L S*'''
#QTI RPM calculation
print('QTI expression')
hpi_list   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]+['48h','48hCaco2', '48hCalu3cured']
df_precalc_ORF1ab_PRFratio = pd.read_csv(RESULTDIR %('RPF_ORF1ab_PRF.tsv'), sep = '\t', index_col = 0)
calc_timecourse_RPFexp(hpi_list, assay_prefix = 'RPFSeq',sample_prefix='QTI', offset_RPF = 12,
                       quant_range = (0,15), as_RPM = True, df_precalc_ORF1ab_PRFratio = df_precalc_ORF1ab_PRFratio, 
                       OutF = 'QTI_quantification_woTISL_S.tsv')

#Plotting
RPFexp_df = pd.read_csv(RESULTDIR %('QTI_quantification_woTISL_S.tsv'), sep = '\t', index_col=0)
hpi_list_timecourse   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]
plot_timecourse_exp(RPFexp_df,hpi_list = hpi_list_timecourse, log10=True,ylim=(0.0,4.0), 
                    psc = 1.0, include_3UTR = True,
                    OutFigname='Fig2b_QTI_woTISL_S_3UTR.pdf')<jupyter_output>QTI expression
0h Mon Sep 27 16:12:25 2021
1h Mon Sep 27 16:12:31 2021
2h Mon Sep 27 16:12:40 2021
4h Mon Sep 27 16:12:48 2021
12h Mon Sep 27 16:12:52 2021
16h Mon Sep 27 16:12:55 2021
24h Mon Sep 27 16:12:57 2021
36h Mon Sep 27 16:13:01 2021
48h Mon Sep 27 16:13:29 2021
48hCaco2 Mon Sep 27 16:13:33 2021
48hCalu3cured Mon Sep 27 16:13:33 2021<jupyter_text>Supplementary Figure 2j: ORF1ab PRF rate<jupyter_code>'''Plotting ORF1ab PRF rate, calculated above( calc_timecourse_PRF_ORF1ab() )'''
def plot_ORF1ab(hpi_list, InF='RPF_ORF1ab_PRF.tsv', OutFigname = ''):
    df_precalc_ORF1ab_PRFratio = pd.read_csv(RESULTDIR %(InF), 
                                             sep = '\t', index_col = 0)
    df_results = df_precalc_ORF1ab_PRFratio.loc[hpi_list]
    display(df_results)
    fig = plt.figure(figsize=(len(hpi_list)*0.5,3))
    ax = fig.add_subplot(111)
    x = np.arange(len(hpi_list))
    bar_width=0.8
    
    ax.bar(x,df_results['PRFratio']*100,width=bar_width)
    ax.set_xticks(x)
    ax.set_xticklabels(hpi_list)
    ax.set_xlabel('hpi')
    
    ax.set_ylabel('ORF1b/ORF1a (%)')
    ax.set_ylim(top=100)
    
    plt.tight_layout()
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
    plt.close()
    return None


hpi_list   = ['%dh' %i for i in [4,12,16,24,36,48]]+['48hCaco2']
plot_ORF1ab(hpi_list,InF='RPF_ORF1ab_PRF.tsv', 
            OutFigname = 'SFig2j_ORF1abPRF.pdf')<jupyter_output><empty_output><jupyter_text>TE (Translaion efficiency)<jupyter_code>def calc_timecourse_TE(InF_mRNA, InF_RPF, OutF_TE,hpi_list = [], psc = 1.0):
    mRNAexp_df = pd.read_csv(RESULTDIR %(InF_mRNA),sep='\t', index_col = 0)[hpi_list]
    RPFexp_df  = pd.read_csv(RESULTDIR %(InF_RPF), sep='\t', index_col = 0)[hpi_list]
    TE_df      = pd.DataFrame(columns = RPFexp_df.columns)
    
    for ORF,RPF_row in RPFexp_df.iterrows():
        if ORF == 'ORF1b':
            mRNA_row = mRNAexp_df.loc['ORF1a']
        elif ORF == 'S_after':
            mRNA_row = mRNAexp_df.loc['S']
        elif ORF == 'ORF7b': #ORF7a mRNA leaky scanning
            mRNA_row = mRNAexp_df.loc['ORF7a']
        elif ORF == '3\'UTR': #N mRNA leaky scanning
            mRNA_row = mRNAexp_df.loc['N']
        elif ORF == 'ORF10':
            mRNA_row = mRNAexp_df.loc[ORF]
            TE_row = (RPF_row+psc)/(mRNA_row+psc)
            TE_df.loc[ORF] = TE_row
            
            mRNA_row = mRNAexp_df.loc['N']
            TE_row = (RPF_row+psc)/(mRNA_row+psc)
            TE_df.loc['ORF10_NmRNA'] = TE_row
            continue
        else:
            mRNA_row = mRNAexp_df.loc[ORF]
            
        TE_row = (RPF_row+psc)/(mRNA_row+psc)
        TE_df.loc[ORF] = TE_row
            
    display(TE_df)
    TE_df.to_csv(RESULTDIR %(OutF_TE),sep='\t') 
    return None



hpi_list   = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]+['48h']
hpi_list_timecourse = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]
for assay in ['RPF','QTI']:
    print(assay)
    calc_timecourse_TE('mRNA_quantification.tsv', 
                       f'{assay}_quantification_woTISL_S.tsv', 
                       f'TE_{assay}_quantification_woTISL_S.tsv',
                       hpi_list = hpi_list, psc = 1.0)
    TE_df = pd.read_csv(RESULTDIR %(f'TE_{assay}_quantification_woTISL_S.tsv'), 
                        sep = '\t', index_col=0)
    plot_timecourse_exp(TE_df,hpi_list = hpi_list_timecourse, log10=True,ylim=(-2.5,1.5), psc = 0.0,
                        OutFigname=f'Fig2b_TE_{assay}_woTISL_S_3UTR.pdf')<jupyter_output>RPF<jupyter_text>Figure 3- Detection of TIS-L- Figure 3b, 3h, Extended Data Figures 3b, 4a: TIS-L fraction- Figure 3c, 3e, Extended Data Figures 3c, 4b-f: TIS-L long read split (for Sukjun)- Figure 3d, 3f, 3h, Extended Data Figures 3d, 4c, 4d: bar plots for TIS-L fraction or expression level- Figure 3g, Supplementary Figure 4e (timecourse expression after consideration of ORF S from TIS-L)  TIS-L enrichment bar plots (Figure 3b, 3h, Extended Data Figures 3b, 4a, 5c, 6a,b)<jupyter_code>def calc_plot_TISL_enrichment(hpi, assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                        figsize = (4,3), OutFigname = ''):
    idx_TISL = 58
    arr_RPF_list  = []
    reps = 3 if hpi == '48h' else 2
    for rep in range(1,reps+1):
        samplename = f'{sample_prefix}_{hpi}_rep{rep}'
        _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                        pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                        offset_RPF = offset_RPF,junction_spanning_only=False,
                                        min_rlen=0,as_RPM=as_RPM)

        arr_RPF_list.append(arr_RPF)

    avg_arr_RPF = np.array(arr_RPF_list).mean(axis=0)
    TISL_mapped = avg_arr_RPF[quant_range[0]:quant_range[1]].sum()
    total_mapped= avg_arr_RPF.sum()
    TISL_pct    = TISL_mapped/total_mapped*100
    report = [f'{sample_prefix}_{hpi}', total_mapped, TISL_pct]
    print(*report, sep = '\t')
    
    fig = plt.figure(figsize=figsize)
    arr_window = avg_arr_RPF[plot_range[0]:plot_range[1]]
    x   = np.arange(plot_range[0],plot_range[1])
    bar_colors = [['darkblue','cornflowerblue','lightskyblue'][(i-idx_TISL)%3] for i in x]
    plt.bar(x+1,arr_window,color=bar_colors)
    plt.ylabel('RPM')
    plt.text(70,arr_window[idx_TISL-plot_range[0]]/2, 
             '%.01f%% of total\nSARS-CoV-2\n%s-seq reads' %(TISL_pct,sample_prefix))
    plt.xlim(plot_range[0]+1,plot_range[1])
    
    plt.tight_layout()
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
        df_source = pd.DataFrame([arr_window], index=['RPM'],columns = np.arange(plot_range[0],plot_range[1])+1).T
        df_source.to_csv(SOURCEDATADIR %(OutFigname.replace('pdf','tsv')), sep = '\t')
    plt.close()
    
    return None
'''Fig.3b'''
calc_plot_TISL_enrichment('48h', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'Fig3b_RPF48h.pdf')

'''SFig.3b'''
calc_plot_TISL_enrichment('48h', assay_prefix = 'RPFSeq',sample_prefix='QTI', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig3b_QTI48h.pdf')

'''Fig.3h'''
calc_plot_TISL_enrichment('48h', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'Fig3h_RPF48h.pdf')
calc_plot_TISL_enrichment('48hCaco2', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'Fig3h_RPF48hCaco2.pdf')

'''SFig.4a'''
for hpi in ['16h','24h','36h']:
    for sample_prefix in ['RPF','QTI']:
        calc_plot_TISL_enrichment(hpi, assay_prefix = 'RPFSeq',sample_prefix=sample_prefix, offset_RPF = 12,
                                quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                                figsize = (4,3), OutFigname = f'SFig4a_{sample_prefix}{hpi}.pdf')

        
'''SFig.5c'''
calc_plot_TISL_enrichment('48hCalu3cured', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig5c_RPF48hCalu3cured.pdf')
calc_plot_TISL_enrichment('48hCalu3cured', assay_prefix = 'RPFSeq',sample_prefix='QTI', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig5c_QTI48hCalu3cured.pdf')
'''SFig.6'''
calc_plot_TISL_enrichment('48h', assay_prefix = 'RPFSeq',sample_prefix='QTI', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig6a_QTI48h.pdf')
calc_plot_TISL_enrichment('48hCaco2', assay_prefix = 'RPFSeq',sample_prefix='QTI', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (49,80), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig6a_QTI48hCaco2.pdf')

calc_plot_TISL_enrichment('24h', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig6b_RPF24hCalu3.pdf')
calc_plot_TISL_enrichment('24hVero', assay_prefix = 'RPFSeq',sample_prefix='RPF', offset_RPF = 12,
                        quant_range = (58,59), plot_range = (39,100), as_RPM = True, 
                        figsize = (4,3), OutFigname = 'SFig6b_RPF24hVero.pdf')<jupyter_output>RPF_48h	340550.6026858094	21.982331559945944
QTI_48h	421921.61884141556	10.870565508348424
RPF_48h	340550.6026858094	21.982331559945944
RPF_48hCaco2	97068.87031917248	7.116556154314048
RPF_16h	7594.026120825337	5.2092537432978485
QTI_16h	8150.962477589109	14.002747645685435
RPF_24h	77153.2697142527	8.579375504419998
QTI_24h	80499.15902695848	13.105030149565774
RPF_36h	192586.06863679376	6.45702079387735
QTI_36h	155506.56815578183	5.32983150090252
RPF_48hCalu3cured	52736.809839021116	11.187050845645382
QTI_48hCalu3cured	49326.75882488812	12.452554050461785
QTI_48h	421921.61884141556	10.870565508348424
QTI_48hCaco2	166457.8185324419	6.391328284080052
RPF_24h	77153.2697142527	8.579375504419998
RPF_24hVero	816914.4666214483	0.16243046524531893<jupyter_text>TIS-L long read split (Figure 3c,e, Supplementary Figures 3c,e,f, 4b, 6c)<jupyter_code>def get_TISL_reads(ORF,hpi,assay_prefix='RPFSeq', sample_prefix='RPF', offset_RPF=12,
                   pos_5end = 46, window_downstream = 6,min_rlen=[32,31,30,29,28,27]):
    #collect TISL reads with matched condition
    #pos_5end 46: 58(idx_TISl) - 12 (offset_RPF), (window_downstream = 6; 2 codons)
    #pos_5end 39: 58(idx-TISL) - 13 (offset_RPF) - 6 (window_downstream = 12; 4 codons)
    assert window_downstream == len(min_rlen)
    sgRNA_junc_dict = CANONICAL_sgRNA_junc_dict.copy()
    sgRNA_junc_dict['ORF1a'] = [75,75] ##no junction, CUG-translation gRNA ORF end
    
    reads_dict = {} #(read_5p,f1_3p,f2_5p,read_3p): count
    reps = 3 if hpi == '48h' else 2
    
    ORF_f2_5p = sgRNA_junc_dict[ORF][1]
    pos_5end_range = (pos_5end, pos_5end+window_downstream)
    for rep in range(1,reps+1):
        samplename = f'{sample_prefix}_{hpi}_rep{rep}'
        bam    = pysam.AlignmentFile(BAMDIR_cov2mapped %(assay_prefix,samplename))
        contig = bam.references[0]
        for read in bam.fetch(contig, pos_5end_range[0], pos_5end_range[1]):
            cigartuples = read.cigartuples
            pos_5end    = read.reference_start
            pos_3end    = read.reference_end
            rlen = read.infer_read_length()
            nh   = read.get_tag('NH')
            if not ((pos_5end_range[0]<= pos_5end) and (pos_5end< pos_5end_range[1])):
                continue
            if rlen<min_rlen[pos_5end-pos_5end_range[0]]:
                continue
            contain_junction = (3 in [i[0] for i in cigartuples])
            if (not contain_junction): #not a junction-spanning read
                if ORF == 'ORF1a':
                    try:
                        reads_dict[(pos_5end,*sgRNA_junc_dict['ORF1a'],pos_3end)] += (1/nh)
                    except KeyError:
                        reads_dict[(pos_5end,*sgRNA_junc_dict['ORF1a'],pos_3end)] = (1/nh)
                    continue
                else:
                    continue
            cur_pos = pos_5end
            f1_detected = False
            f1_3p = None
            f2_5p = None
            for idx,cigar_tp in enumerate(cigartuples):
                operation, length = cigar_tp
                if operation == 3: #skip
                    f1_detected = True
                    cur_pos += length
                else:
                    if operation == 4 or operation == 1:    # softclip or insertion
                        pass
                    elif operation == 0 or operation == 2:  # match/mismatch or deletion
                        cur_pos += length
                if operation == 0:
                    if not f1_detected: #junction 1st frag
                        f1_3p = cur_pos

                    else: #second fragment
                        f2_5p = cur_pos - length
                        break

            assert f1_3p != None
            assert f2_5p != None
            if f2_5p == ORF_f2_5p:
                try:
                    reads_dict[(pos_5end,f1_3p,f2_5p,pos_3end)] += (1/nh)
                except KeyError:
                    reads_dict[(pos_5end,f1_3p,f2_5p,pos_3end)]  = (1/nh)
            
        bam.close()
                    
    reads_df = [[*key,value] for key,value in reads_dict.items()]
    
    if len(reads_df) == 0:
        reads_df = pd.DataFrame(columns = ['read_5p','f1_3p','f2_5p','read_3p','count'])
    else:
        reads_df = pd.DataFrame(reads_df)
        reads_df.columns = ['read_5p','f1_3p','f2_5p','read_3p','count']
    reads_df = reads_df.sort_values('count', ascending=False)
    return reads_df

def plot_topreads_seqalign(ORF,reads_df,ax,n_allreads=0,window_start = 37,window_size=50,
                           n_topreads=3, splitspace=11):
    global SARSCOV2_genome
    idx_TISL      = 58
    idx_TRSLeader = 69
    len_TRSLeader = 6
    sgRNA_junc_dict = CANONICAL_sgRNA_junc_dict.copy()
    sgRNA_junc_dict['ORF1a'] = [75,75] ##no junction, CUG-translation gRNA ORF end
    
    total_n      = reads_df['count'].sum()
    junc_s, junc_e = sgRNA_junc_dict[ORF] #= f1_3p, f2_5p for read fragments
    
    top_reads_df = reads_df.iloc[:n_topreads]
    if ORF == 'ORF1a':
        seq_annot= (SARSCOV2_genome[window_start:junc_s]+
                    SARSCOV2_genome[junc_e:junc_e+window_size-(junc_s-window_start)+splitspace]
                   ).replace('T','U')
        coordinate = [i+1 for i in range(window_start,junc_s)]
        coordinate+= [i+1 for i in range(junc_e,junc_e+window_size-(junc_s-window_start)+splitspace)]
    else:
        seq_annot= (SARSCOV2_genome[window_start : idx_TRSLeader+5] +
                    ' ' + SARSCOV2_genome[junc_e+(idx_TRSLeader-junc_s)-5:
                                          junc_e+window_size-(junc_s-window_start)]
                   ).replace('T','U')
        coordinate = [i+1 for i in range(window_start,idx_TRSLeader+5)]
        coordinate+= [-1 for i in range(1)]+[i+1 for i in range(junc_e+(idx_TRSLeader-junc_s)-5,
                                                                junc_e+window_size-(junc_s-window_start))]
    seq_list = []
    n_reads_list = []
    
    for _,row in top_reads_df.iterrows():
        read_s,f1_3p, f2_5p,read_e = row.iloc[:4].astype(int)
        n_reads = row['count']
        
        if ORF == 'ORF1a':
            seq = SARSCOV2_genome[read_s:f1_3p] + SARSCOV2_genome[f2_5p:read_e]
            seq = seq+(' '*splitspace)
            seq = ' '*(read_s-window_start) + seq + ' '*(window_size-(len(seq)-splitspace)-(read_s-window_start))
            
        else:
            seq = SARSCOV2_genome[read_s:f1_3p] +(' '*splitspace)+ SARSCOV2_genome[f2_5p:read_e]
            seq = ' '*(read_s-window_start) + seq + ' '*(window_size-(len(seq)-splitspace)-(read_s-window_start))
        
        seq_list.append(seq.replace('T','U'))
        n_reads_list.append(n_reads)
        
    for ntidx,(nt,coor) in enumerate(zip(seq_annot,coordinate)):
        color = 'black'
        if (idx_TISL<=(ntidx+window_start) and 
            (ntidx+window_start)<(idx_TISL+3)): #TISL
            color = 'green'
        elif ((idx_TRSLeader+splitspace) <=(ntidx+window_start) and 
              (ntidx+window_start)<(idx_TRSLeader+splitspace+len_TRSLeader) and ORF!='ORF1a'): #TRS-Body
                color='red'
        elif (idx_TRSLeader <= (ntidx+window_start)   and 
              (ntidx+window_start)< (idx_TRSLeader+len_TRSLeader) and ORF=='ORF1a'): #TRS-Leader
            color='red'
            
        if nt != ' ':
            ax.text(ntidx/len(seq_annot)*0.9,0.1,nt,
                     fontsize=12,color=color)
        if coor%10 == 0:
            ax.text(ntidx/len(seq_annot)*0.9,0,coor,
                 fontsize=12,color='k')
    
    ax.axhline(1/(n_topreads+2),xmin=0,xmax=0.9,linestyle='-',color='k',linewidth=1)
    for rowidx,(seq,n_reads) in enumerate(zip(seq_list,n_reads_list)):
        for ntidx,nt in enumerate(seq):
            color = 'black'
            if (idx_TISL<=(ntidx+window_start) and 
                (ntidx+window_start)<(idx_TISL+3)): #TISL
                color = 'green'
            elif ((idx_TRSLeader+splitspace) <=(ntidx+window_start) and 
                  (ntidx+window_start)<(idx_TRSLeader+splitspace+len_TRSLeader) and ORF!='ORF1a'): #TRS-Body
                    color='red'
            elif (idx_TRSLeader <= (ntidx+window_start)   and 
                  (ntidx+window_start)< (idx_TRSLeader+len_TRSLeader) and ORF=='ORF1a'): #TRS-Leader
                color='red'
            if nt != ' ':
                ax.text(ntidx/len(seq)*0.9,(n_topreads-rowidx+1)/(n_topreads+2),nt,
                        fontsize=12,color=color)
            
        ax.text(0.9,(n_topreads-rowidx+1)/(n_topreads+2),np.ceil(n_reads).astype(np.int),
                 fontsize=12)
        
        if ORF == 'ORF1a':
            ax.text(int(len(seq)/2)/len(seq)*0.9,1/(n_topreads+2),'.',fontsize=12)
            ax.text(int(len(seq)/2)/len(seq)*0.9,1/(n_topreads+2)+0.05,'.',fontsize=12)
            ax.text(int(len(seq)/2)/len(seq)*0.9,1/(n_topreads+2)+0.1,'.',fontsize=12)
            ax.text(int(len(seq)/2)/len(seq)*0.9,1.1,'...',fontsize=15)
            ax.axhline(1,linestyle='--',color='r',linewidth=2)
        else:
            ax.text(int(len(seq)/4)/len(seq)*0.9,1/(n_topreads+2),'.',fontsize=12)
            ax.text(int(len(seq)/4)/len(seq)*0.9,1/(n_topreads+2)+0.05,'.',fontsize=12)
            ax.text(int(len(seq)/4)/len(seq)*0.9,1/(n_topreads+2)+0.1,'.',fontsize=12)
            ax.text(int(len(seq)/4*3)/len(seq)*0.9,1/(n_topreads+2),'.',fontsize=12)
            ax.text(int(len(seq)/4*3)/len(seq)*0.9,1/(n_topreads+2)+0.05,'.',fontsize=12)
            ax.text(int(len(seq)/4*3)/len(seq)*0.9,1/(n_topreads+2)+0.1,'.',fontsize=12)
            
        ax.text(0.9,1/(n_topreads+2),'.',fontsize=12)
        ax.text(0.9,1/(n_topreads+2)+0.05,'.',fontsize=12)
        ax.text(0.9,1/(n_topreads+2)+0.1,'.',fontsize=12)
    ax.text(1,0.4,'%s\n%d (%.01f%%)' %(ORF,np.ceil(total_n).astype(np.int),total_n/n_allreads*100),fontsize=12)
    ax.axis('off')
    return None

def plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '36h', 
                       offset_RPF = 12, pos_5end = 46, window_downstream = 6, 
                      n_display_sgRNA = 3, n_topreads = 3, splitspace=11, OutFigname=''):
    
    reps = 3 if hpi == '48h' else 2
    min_rlen = [78-i for i in range(pos_5end, pos_5end+window_downstream)]
    ORFlist  = ['ORF1a','S','ORF3a','E','M','ORF6','ORF7a','ORF7b','ORF8','N']
    
    ORF_reads_list = []
    for ORF in ORFlist:
        reads_df = get_TISL_reads(ORF,hpi,assay_prefix=assay_prefix, 
                                 sample_prefix=sample_prefix, offset_RPF=offset_RPF, pos_5end = pos_5end,
                                 window_downstream = window_downstream, min_rlen=min_rlen)
        
        ORF_reads_list.append([ORF,reads_df,reads_df['count'].sum()])
    total = sum([i[2] for i in ORF_reads_list])
    gRNA  = ORF_reads_list[0][2]
    
    ORF_reads_list_sorted = sorted(ORF_reads_list[1:],key = lambda x: x[2],reverse = True)
    ORF_reads_list_sorted = ORF_reads_list_sorted[:n_display_sgRNA]
    if not ('S' in [i[0] for i in ORF_reads_list_sorted]):
        ORF_reads_list_sorted.append(ORF_reads_list[1])
        
    ORF1a          = ORF_reads_list[0] 
    ORF_reads_list_sorted.append(ORF1a)
    
    if 'U' in hpi: #low RNAs conc.
        window_size = 60
        window_start= 35
        figwidth    = 12
    else:
        window_size = 40
        window_start= 45
        figwidth    = 9
        
    fig,axes = plt.subplots(nrows=len(ORF_reads_list_sorted),ncols=1,figsize=(figwidth,1.8*len(ORF_reads_list_sorted)))
    for ax,(ORF,reads_df,totalct) in zip(axes,ORF_reads_list_sorted):
        plot_topreads_seqalign(ORF,reads_df,ax,n_allreads=total,
                               window_start = window_start,window_size=window_size,
                               n_topreads=min(len(reads_df),n_topreads), splitspace=splitspace)
    gRNA_int = np.ceil(gRNA).astype(int)
    plt.suptitle('sgRNAs: %d (%.01f%%) gRNA: %d (%.01f%%)' %(total-gRNA_int,((total-gRNA)/total*100),
                                                             gRNA_int,gRNA/total*100))
    plt.tight_layout()
    
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
    else:
        plt.show()
    plt.close()
    return None
'''Fig.3c: RPF48h'''
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '48h', 
                       offset_RPF = 12, pos_5end = 46, window_downstream = 6, 
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='Fig3c_RPF48h_split.pdf')

'''Fig.3e: RPF48h, 0008U (low RNase I conc.)'''
##CAUTION: used 13nt as RPF offset, based on the previous analysis on triplet periodicity
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '48h0008U', 
                       offset_RPF = 13, pos_5end = 39, window_downstream = 12, 
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='Fig3e_RPF48h0008U_split.pdf')

'''SFig.3c: QTI48h'''
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'QTI',hpi = '48h', 
                       offset_RPF = 12, pos_5end = 46, window_downstream = 6,
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='SFig3c_QTI48h_split.pdf')

'''SFig. 3e,f: exact TIS-L reads'''
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '48h', 
                       offset_RPF = 12, pos_5end = 46, window_downstream = 1, 
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='SFig3e_RPF48h_exactTIS-L.pdf')
##CAUTION: used 13nt as RPF offset, based on the previous analysis on triplet periodicity
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '48h0008U', 
                       offset_RPF = 13, pos_5end = 45, window_downstream = 1, 
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='SFig3f_RPF48h0008U_exactTIS-L.pdf')

'''SFig.4b: RPF16-36h'''
for tmp_hpi in ['16h','24h','36h']:
    plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = tmp_hpi, 
                           offset_RPF = 12, pos_5end = 46, window_downstream = 6,
                          n_display_sgRNA = 3, n_topreads = 3, OutFigname=f'SFig4b_RPF{tmp_hpi}_split.pdf')

'''SFig.6c: Vero and Caco2'''
plot_CUGalignments(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '24hVero', 
                       offset_RPF = 12, pos_5end = 46, window_downstream = 6,
                      n_display_sgRNA = 3, n_topreads = 3, OutFigname='SFig6c_RPF24hVero_split.pdf')<jupyter_output><empty_output><jupyter_text>Bar plots for TIS-L fraction or expression level (Figure 3d,f,h, Supplementary Figures 3d, 4c,d, 5d, 6a,b)<jupyter_code>def calc_TISL_ORF_fraction(hpi='48h',assay_prefix='RPFSeq',sample_prefix='RPF',offset_RPF = 12, 
                             pos_5end = 46,window_downstream=6,OutF = 'RPF48h_fraction.tsv'):
    global CANONICAL_sgRNA_junc_dict
    idx_TISL = 58
    
    pos_5end_range = (pos_5end, pos_5end+window_downstream)
    min_rlen = [78-i for i in range(*pos_5end_range)]
    
    arr_f2_5p_list  = []
    arr_RPF_list    = []
    reps=3 if hpi == '48h' else 2
    for rep in range(1,reps+1):
        samplename = f'{sample_prefix}_{hpi}_rep{rep}'
        _, arr_f2_5p, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                                pos_5end_range=pos_5end_range, f2_5p_range=None, fetch_range=pos_5end_range,
                                                offset_RPF = offset_RPF,junction_spanning_only=False,
                                                min_rlen=min_rlen,as_RPM=False)
        arr_f2_5p_list.append(arr_f2_5p)
        arr_RPF_list.append(arr_RPF)
        
    sum_arr_f2_5p = np.array(arr_f2_5p_list).sum(axis=0)
    sum_arr_RPF   = np.array(arr_RPF_list).sum(axis=0)
    
    TISLreads_df  = [['ORF1a',sum_arr_RPF.sum()-sum_arr_f2_5p.sum()]]
    ORFs = ['S','ORF3a','E','M','ORF6','ORF7a','ORF7b','ORF8','N']
    for ORF in ORFs:
        f2_5p = CANONICAL_sgRNA_junc_dict[ORF][1]
        TISLreads_df.append([ORF,sum_arr_f2_5p[f2_5p]])
        
    TISLreads_df = pd.DataFrame(TISLreads_df)
    TISLreads_df.columns = ['ORF','n_TISLreads']
    TISLreads_df = TISLreads_df.set_index('ORF')
    TISLreads_df['Percentage'] = TISLreads_df['n_TISLreads']/(TISLreads_df['n_TISLreads'].sum())*100
    
    TISLreads_df.to_csv(RESULTDIR %OutF, sep = '\t')
    return None

#Calculation of prerequisite data: relative fraction of each ORF for TIS-L reads

#Calu-3
tmp_hpi_list = ['%dh' %i for i in [0,1,2,4,12,16,24,36,48]]
for tmp_hpi in tmp_hpi_list:
    for tmp_sample_prefix in ['RPF','QTI']:
        calc_TISL_ORF_fraction(hpi=tmp_hpi,assay_prefix='RPFSeq',sample_prefix=tmp_sample_prefix,offset_RPF = 12, 
                               pos_5end = 46,window_downstream=6,
                               OutF = f'{tmp_sample_prefix}{tmp_hpi}_fraction.tsv')
#Vero
for tmp_hpi in ['24hVero']:
    for tmp_sample_prefix in ['RPF']:
        calc_TISL_ORF_fraction(hpi=tmp_hpi,assay_prefix='RPFSeq',sample_prefix=tmp_sample_prefix,offset_RPF = 12, 
                               pos_5end = 46,window_downstream=6,
                               OutF = f'{tmp_sample_prefix}{tmp_hpi}_fraction.tsv')
#low Rnase conc.
for tmp_hpi in ['48h0008U']:
    for tmp_sample_prefix in ['RPF']:
        calc_TISL_ORF_fraction(hpi=tmp_hpi,assay_prefix='RPFSeq',sample_prefix=tmp_sample_prefix,offset_RPF = 13, 
                               pos_5end = 39, window_downstream=12,
                               OutF = f'{tmp_sample_prefix}{tmp_hpi}_fraction.tsv')
#Caco-2
for tmp_hpi in ['48hCaco2']:
    for tmp_sample_prefix in ['RPF','QTI']:
        calc_TISL_ORF_fraction(hpi=tmp_hpi,assay_prefix='RPFSeq',sample_prefix=tmp_sample_prefix,offset_RPF = 12, 
                               pos_5end = 46,window_downstream=6,
                               OutF = f'{tmp_sample_prefix}{tmp_hpi}_fraction.tsv')
#Calu-3, cured sample
for tmp_hpi in ['48hCalu3cured']:
    for tmp_sample_prefix in ['RPF','QTI']:
        calc_TISL_ORF_fraction(hpi=tmp_hpi,assay_prefix='RPFSeq',sample_prefix=tmp_sample_prefix,offset_RPF = 12, 
                               pos_5end = 46,window_downstream=6,
                               OutF = f'{tmp_sample_prefix}{tmp_hpi}_fraction.tsv')
#For Fig. 3d, 3h, etc (relative fractions)
def barplot_fraction(sample_prefix='RPF', hpi_list = ['48h','48h0008U'], label_list=['5U','0.008U'],
                    hatch_list=['','///'],ylim=(0,45),rotation_x=0,OutFigname=''):
    global ORF_palette
    merged_df = []
    
    for hpi in hpi_list:
        TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}{hpi}_fraction.tsv'),
                                   sep = '\t', index_col = 0)
        pct = TISLreads_df['Percentage']
        pct.name = hpi
        merged_df.append(pct)
        
    merged_df = pd.concat(merged_df, axis=1)
    colors    = [ORF_palette[ORF] for ORF in merged_df.index]
    
    fig,ax = plt.subplots(figsize=(6,4))
    width  = 0.8/len(hpi_list)
    x_offset_list = [-0.4+width*(i+1/2) for i in range(len(hpi_list))]
    
    for hpi,label,hatch,x_offset in zip(hpi_list,label_list,hatch_list,x_offset_list):
        ax.bar(np.arange(len(merged_df))+x_offset, merged_df[hpi], width=width, linewidth=0.5,
               edgecolor='k',hatch=hatch,color = colors,label = label)
    
    ax.set_xticks(np.arange(len(merged_df)))
    ax.set_xticklabels(merged_df.index,rotation=rotation_x)
    ax.set_xlim(-0.5,len(merged_df)-0.5)
    ax.set_ylim(*ylim)
    ax.set_ylabel('Relative fraction (%)')
    gRNA_text = 'gRNA\n'
    sgRNA_text= 'sgRNA\n'
    for hpi,label in zip(hpi_list,label_list):
        gRNA_text  += '%-8s %.01f%%\n' %(label+':',merged_df.loc['ORF1a',hpi])
        sgRNA_text += '%-8s %.01f%%\n' %(label+':',100-merged_df.loc['ORF1a',hpi])
    
    ax.text(-0.4,ylim[1]+3,gRNA_text)
    ax.text(4,ylim[1]+3,sgRNA_text)
    ax.axvline(0.5,color='k',linewidth=1.5,linestyle='--')
    
    plt.legend()
    plt.tight_layout()
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
    else:
        plt.show()
    plt.close()
    
    return None

'''Fig. 3d: RPF 48h (5U vs. 0.008U)'''
barplot_fraction(sample_prefix='RPF', hpi_list = ['48h','48h0008U'], label_list=['5U','0.008U'],
                 hatch_list=['','///'],ylim=(0,45),rotation_x=90,
                 OutFigname='Fig3d_RPF48h_5U0008U.pdf')

'''Fig. 3h right panel: Calu vs. Caco'''
barplot_fraction(sample_prefix='RPF', hpi_list = ['48h','48hCaco2'], label_list=['Calu-3','Caco-2'],
                 hatch_list=['','///'],ylim=(0,45),rotation_x=90,
                 OutFigname='Fig3h_CaluCaco.pdf')

'''SFig.3d : QTI 48h'''
barplot_fraction(sample_prefix='QTI', hpi_list = ['48h'], label_list=['QTI 48h'],
                 hatch_list=[''],ylim=(0,45),rotation_x=90,
                 OutFigname='SFig3d_QTI48h.pdf')

'''SFig. 4c: Calu RPF,QTI 16, 24, 36h'''
barplot_fraction(sample_prefix='RPF', hpi_list = ['16h','24h','36h'], label_list=['16 hpi','24 hpi','36 hpi'],
                 hatch_list=['','/','///'],ylim=(0,55),rotation_x=90,
                 OutFigname='SFig4c_RPF_162436h.pdf')
barplot_fraction(sample_prefix='QTI', hpi_list = ['16h','24h','36h'], label_list=['16 hpi','24 hpi','36 hpi'],
                 hatch_list=['','/','///'],ylim=(0,55),rotation_x=90,
                 OutFigname='SFig4c_QTI_162436h.pdf')

'''SFig.5d: before and after treatment of mycoplasma'''
barplot_fraction(sample_prefix='RPF', hpi_list = ['48h','48hCalu3cured'], label_list=['previous','Mycoplasma cured'],
                 hatch_list=['','/','///'],ylim=(0,60),rotation_x=90,
                 OutFigname='SFig5d_RPF48h_mycoplasma_cmp.pdf')
barplot_fraction(sample_prefix='QTI', hpi_list = ['48h','48hCalu3cured'], label_list=['previous','Mycoplasma cured'],
                 hatch_list=['','/','///'],ylim=(0,60),rotation_x=90,
                 OutFigname='SFig5d_QTI48h_mycoplasma_cmp.pdf')

'''SFig.6a,b: Calu vs. Caco (QTI),  Calu vs. Vero'''
barplot_fraction(sample_prefix='QTI', hpi_list = ['48h','48hCaco2'], label_list=['Calu-3','Caco-2'],
                 hatch_list=['','///'],ylim=(0,45),rotation_x=90,
                 OutFigname='SFig6a_CaluCaco_QTI.pdf')
barplot_fraction(sample_prefix='RPF', hpi_list = ['24h','24hVero'], label_list=['Calu-3','Vero'],
                 hatch_list=['','///'],ylim=(0,45),rotation_x=90,
                 OutFigname='SFig6b_CaluVero.pdf')
#For Fig. 3f, SFig. 4d (comparison b/w annotated ORF translation and ORF from TIS-L for each gene)
def calc_barplot_TISL_exp(sample_prefix='RPF', hpi = '36h',offset_RPF = 12,quant_range=(58,64),
                          hatch_list=['','///'],ylim=(0,5),rotation_x=0,OutFigname=''):
    global ORF_palette
    ORF_list = ['ORF1a','S','ORF3a','E','M','ORF6','ORF7a','ORF7b','ORF8','N','ORF10']
    
    exp_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}_quantification_woTISL_S.tsv'),
                         sep = '\t', index_col = 0)
    
    merged_df = pd.DataFrame(index=ORF_list)
    merged_df['annot_RPM'] = exp_df.loc[ORF_list,hpi]
    
    '''
    Calculate read enrichment at TISL and 
    distribute them in proportion to the precalculated relative fractions
    '''
    TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}{hpi}_fraction.tsv'),
                               sep = '\t', index_col = 0)
    TISL_pct = TISLreads_df['Percentage']
    arr_RPF_list  = []
    reps = 3 if hpi == '48h' else 2
    for rep in range(1,reps+1):
        samplename = f'{sample_prefix}_{hpi}_rep{rep}'
        _, _, arr_RPF = get_readpos_arr(samplename,'RPFSeq',norm_by_nh=True,
                                        pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=quant_range,
                                        offset_RPF = offset_RPF,junction_spanning_only=False,
                                        min_rlen=0,as_RPM=True)

        arr_RPF_list.append(arr_RPF)

    avg_arr_RPF = np.array(arr_RPF_list).mean(axis=0)
    TISL_mapped = avg_arr_RPF[quant_range[0]:quant_range[1]].sum()
    TISL_RPM    = TISL_pct*TISL_mapped/100
    merged_df['TISL_RPM'] = TISL_RPM
    merged_df   = merged_df.fillna(0.0)
    ######################################################    
    
    display(merged_df)
    
    colors    = [ORF_palette[ORF] for ORF in merged_df.index]
    cond_list = merged_df.columns
    label_list= ['Canonical AUG','TIS-L']
    df_source = pd.DataFrame(columns = label_list, index = merged_df.index)
    
    fig,ax = plt.subplots(figsize=(6,4))
    width  = 0.8/len(cond_list)
    x_offset_list = [-0.4+width*(i+1/2) for i in range(len(cond_list))]
    
    for cond,label,hatch,x_offset in zip(cond_list,label_list,hatch_list,x_offset_list):
        vals = np.log10(merged_df[cond]+1)
        ax.bar(np.arange(len(merged_df))+x_offset, vals, width=width, linewidth=0.5,
               edgecolor='k',hatch=hatch,color = colors,label = label)
        df_source[label] = vals
    
    ax.set_xticks(np.arange(len(merged_df)))
    ax.set_xticklabels(merged_df.index, rotation=rotation_x)
    ax.set_xlim(-0.5,len(merged_df)-0.5)
    ax.set_ylim(*ylim)
    ax.set_ylabel('log10(RPM+1)')
    
    ax.axvline(0.5,color='k',linewidth=1.5,linestyle='--')
    
    plt.legend()
    plt.tight_layout()
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
        df_source.to_csv(SOURCEDATADIR %(OutFigname.replace('pdf','tsv')), sep = '\t')
    else:
        plt.show()
    plt.close()
    
    return None

'''Fig.3f: RPF'''
calc_barplot_TISL_exp(sample_prefix='RPF', hpi = '36h',offset_RPF = 12,quant_range=(58,64),
                      hatch_list=['','///'],ylim=(0,5),rotation_x=90,
                      OutFigname='Fig3f_RPF_annot_TISL.pdf')
'''SFig.4d: QTI'''
calc_barplot_TISL_exp(sample_prefix='QTI', hpi = '36h',offset_RPF = 12,quant_range=(58,64),
                      hatch_list=['','///'],ylim=(0,5),rotation_x=90,
                      OutFigname='SFig4d_QTI_annot_TISL.pdf')
def calc_plot_timecourse_exp_TISL_S(sample_prefix='RPF',hpi_list=[], offset_RPF = 12,window_downstream=15,
                                    OutF = '', OutFigname=''):
    global ORF_palette, CANONICAL_sgRNA_junc_dict
    idx_TISL     = 58
    idx_f1_3p_S,idx_f2_5p_S  = CANONICAL_sgRNA_junc_dict['S']
    #print(idx_f1_3p_S,idx_f2_5p_S)
    exp_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}_quantification_woTISL_S.tsv'),
                         sep = '\t', index_col = 0)
    
    exp_df = exp_df[hpi_list]
    
    S_after_list = []
    
    
    for hpi in hpi_list:
        print(hpi,time.ctime())
        '''
        Calculate read enrichment at TISL and 
        distribute them in proportion to the precalculated relative fractions 
        (used 48h fraction for early time points)
        '''
        if hpi in ['0h','1h','2h','4h','12h','16h']:
            TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}48h_fraction.tsv'),
                                       sep = '\t', index_col = 0)
        else:
            TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}{hpi}_fraction.tsv'),
                                       sep = '\t', index_col = 0)
        TISL_pct = TISLreads_df.loc['S','Percentage']
        
        arr_RPF_list = []
        reps=3 if hpi == '48h' else 2
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            _, _, arr_RPF = get_readpos_arr(samplename,'RPFSeq',norm_by_nh=True,
                                            pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                            offset_RPF = offset_RPF,junction_spanning_only=False,
                                            min_rlen=0,as_RPM=True)

            arr_RPF_list.append(arr_RPF)

        avg_arr_RPF   = np.array(arr_RPF_list).mean(axis=0)
        window_leader = min(window_downstream,idx_f1_3p_S-idx_TISL)
        window_body   = window_downstream-window_leader
        
        S_TISL    = avg_arr_RPF[idx_TISL:idx_TISL+window_leader].sum()*TISL_pct/100
        S_TISL   += avg_arr_RPF[idx_f2_5p_S:idx_f2_5p_S+window_body].sum()
        
        
        S_after_list.append(exp_df.loc['S',hpi]+S_TISL)
    
    exp_df.loc['S_after'] = S_after_list
    display(exp_df)
    exp_df.to_csv(RESULTDIR %OutF, sep='\t')
    plot_timecourse_exp(exp_df, hpi_list=hpi_list, log10=True,ylim=(0,4),psc = 1.0,grey_exceptS=True,
                        ORFs_to_highlight=[], OutFigname = OutFigname)
    
    return None


'''Fig.3g, SFig.4e (left): RPF, QTI with TISL_S'''
tmp_hpi_list = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]

tmp_sample_prefix = 'RPF'
calc_plot_timecourse_exp_TISL_S(sample_prefix=tmp_sample_prefix,hpi_list=tmp_hpi_list, offset_RPF = 12,
                                window_downstream=15,
                                OutF = f'{tmp_sample_prefix}_quantification_TISL_S.tsv',
                                OutFigname=f'Fig3g_{tmp_sample_prefix}.pdf')
tmp_sample_prefix = 'QTI'
calc_plot_timecourse_exp_TISL_S(sample_prefix=tmp_sample_prefix,hpi_list=tmp_hpi_list, offset_RPF = 12,
                                window_downstream=15,
                                OutF = f'{tmp_sample_prefix}_quantification_TISL_S.tsv',
                                OutFigname=f'SFig4e_{tmp_sample_prefix}.pdf')
def calc_plot_timecourse_TE_TISL_S(InF_mRNA, InF_RPF, OutF_TE, hpi_list = [], OutFigname=''):
    calc_timecourse_TE(InF_mRNA, InF_RPF, OutF_TE, hpi_list = hpi_list, psc = 1.0)
    
    TE_df = pd.read_csv(RESULTDIR %(OutF_TE), sep = '\t', index_col=0)
    plot_timecourse_exp(TE_df, hpi_list=hpi_list, log10=True,ylim=(-1.5,1.5), psc = 0.0, grey_exceptS=True,
                        OutFigname=OutFigname)
    
    return None
'''Fig.3g, SFig.4e (right): RPF, QTI TE with TISL_S'''
tmp_hpi_list = ['%dh' %i for i in [0,1,2,4,12,16,24,36]]

tmp_sample_prefix = 'RPF'
calc_plot_timecourse_TE_TISL_S('mRNA_quantification.tsv',  
                               f'{tmp_sample_prefix}_quantification_TISL_S.tsv', 
                               f'TE_{tmp_sample_prefix}_quantification_TISL_S.tsv',
                               hpi_list = tmp_hpi_list,
                               OutFigname= f'Fig3g_TE_{tmp_sample_prefix}_TISL_S.pdf')
tmp_sample_prefix = 'QTI'
calc_plot_timecourse_TE_TISL_S('mRNA_quantification.tsv',  
                               f'{tmp_sample_prefix}_quantification_TISL_S.tsv', 
                               f'TE_{tmp_sample_prefix}_quantification_TISL_S.tsv',
                               hpi_list = tmp_hpi_list,
                               OutFigname= f'SFig4e_TE_{tmp_sample_prefix}_TISL_S.pdf')<jupyter_output><empty_output><jupyter_text>Comparison b/w Calu and Caco or cured Calu annotated SARS-CoV-2 viral ORF expression levels (Supplementary Figure 5a)<jupyter_code>def calc_barplot_cmp_annot_exp(sample_prefix='RPF', hpi_list = ['48h','48hCaco2'], label_list=['Calu-3','Caco-2'],
                    hatch_list=['','///'], ylim=(0,5),as_pct = False,rotation_x=0,OutFigname=''):
    global ORF_palette
    merged_df = []
    exp_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}_quantification_woTISL_S.tsv'),
                         sep = '\t', index_col = 0)
    for hpi in hpi_list:
        exp = exp_df[hpi]
        if as_pct:
            exp = exp/exp.sum()*100
        merged_df.append(exp)
        
    merged_df = pd.concat(merged_df, axis=1)
    colors    = [ORF_palette[ORF] for ORF in merged_df.index]
    
    fig,ax = plt.subplots(figsize=(6,4))
    width  = 0.8/len(hpi_list)
    x_offset_list = [-0.4+width*(i+1/2) for i in range(len(hpi_list))]
    
    for hpi,label,hatch,x_offset in zip(hpi_list,label_list,hatch_list,x_offset_list):
        if as_pct:
            ax.bar(np.arange(len(merged_df))+x_offset, merged_df[hpi], width=width, linewidth=0.5,
                   edgecolor='k',hatch=hatch,color = colors,label = label)
        else:
            ax.bar(np.arange(len(merged_df))+x_offset, np.log10(merged_df[hpi]+1), width=width, linewidth=0.5,
                   edgecolor='k',hatch=hatch,color = colors,label = label)
    
    ax.set_xticks(np.arange(len(merged_df)))
    ax.set_xticklabels(merged_df.index,rotation=rotation_x)
    ax.set_xlim(-0.5,len(merged_df)-0.5)
    ax.set_ylim(*ylim)
    if as_pct:
        ax.set_ylabel('Relative fraction (%)')
    else:
        ax.set_ylabel('log10(RPM+1)')
        
    ax.axvline(1.5,color='k',linewidth=1.5,linestyle='--')
    plt.legend()
    plt.tight_layout()
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
    else:
        plt.show()
    plt.close()
    
    return None

for sample_prefix in ['RPF','QTI']:
    calc_barplot_cmp_annot_exp(sample_prefix=sample_prefix, hpi_list = ['48h','48hCalu3cured'], 
                               label_list=['Before treatment','After treatment'],
                               hatch_list=['','///'],ylim=(0,80),rotation_x=90,as_pct = True,
                               OutFigname = f'SFig5a_{sample_prefix}48h_mycoplasma_annotORFs_cmp_pct.pdf')<jupyter_output><empty_output><jupyter_text>Figure 4- RPF-seq reads mapped on TIS-L and annotated ORFs- Figure 4b-e, Supplementary Figure 7a-f: RPF-seq- Supplementary Figure 8a-j: QTI-seq- Supplementary Figures 7g, 8k: ORF7b- Supplementary Figure 9: ORF10  Figure 4b-e, Supplementary Figure 7a-f: RPF-seq Supplementary Figure 8a-j: QTI-seq<jupyter_code>def plot_sgRNA(ORF = 'N',hpi_list = ['16h','24h','36h'], s = 49,
               assay_prefix = 'RPFSeq', sample_prefix = 'RPF', as_RPM = True, offset_RPF = 12,
               as_log10= True, ylim= 4, OutFigname=''):
    global CANONICAL_sgRNA_junc_dict, SARSCOV2_genome
    sgRNA_junc_dict = CANONICAL_sgRNA_junc_dict.copy()
    sgRNA_junc_dict['ORF1a'] = [75,75] ##no junction, CUG-translation gRNA ORF end
    idx_TISL = 58
    contig_len  =  len(SARSCOV2_genome)
    grid     = GridSpec(nrows=2*(1+len(hpi_list)),ncols=1)
    fig      = plt.figure(figsize=(10,1.5*(1+len(hpi_list))))
    thickness= 5
    
    junc_5p, junc_3p  = sgRNA_junc_dict[ORF]
    
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = True)
    ORFstart          = df_SARSCOV2_annot.at[ORF,'ORFstart']
    ORFstart_shifted  = (ORFstart-junc_3p)+(junc_5p-s)
    
    window_size       = 50*(ORFstart_shifted//50+1)
    if ORF in ['ORF3a','E']:
        window_size  += 50
    elif ORF == 'M':
        window_size  += 100
    
    if window_size == 50:
        ticklabel_width = 10
    else:
        ticklabel_width = 20
    
    e = ORFstart + (window_size-ORFstart_shifted)
    
    df_source = []
    for hpi_idx,hpi in enumerate(hpi_list):
        arr_RPF_list  = []
        if hpi == '48h':
            reps=3
        else:
            reps=2
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                            pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                            offset_RPF = offset_RPF,junction_spanning_only=False,
                                            min_rlen=0,as_RPM=as_RPM)

            arr_RPF_list.append(arr_RPF)
        
        avg_arr_RPF  = np.array(arr_RPF_list).mean(axis=0)
        
        TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}{hpi}_fraction.tsv'),
                            sep = '\t', index_col = 0)
        TISL_pct     = TISLreads_df.loc[ORF,'Percentage']
        arr_RPF_fused= np.concatenate((avg_arr_RPF[s:junc_5p]*TISL_pct/100, avg_arr_RPF[junc_3p:e]))
        
        if as_log10:
            psc = 1
            arr_RPF_fused = np.log10(arr_RPF_fused+psc)
        
        x = np.arange(len(arr_RPF_fused))
        tmp_xticklabels = np.arange(s,junc_5p)+1 #1-base coordinate
        xticks      = [i for i,nt in enumerate(tmp_xticklabels) if nt%ticklabel_width == 0]
        xticklabels = tmp_xticklabels[xticks]
        
        tmp_xticklabels = np.arange(junc_3p,e)+1 #1-base coordinate
        tmp_xticks      = [i for i,nt in enumerate(tmp_xticklabels) if nt%ticklabel_width == 0]
        tmp_xticklabels = tmp_xticklabels[tmp_xticks]
        xticks = xticks+[i+junc_5p-s for i in tmp_xticks]
        xticklabels= np.concatenate((xticklabels,tmp_xticklabels))

        bar_colors = [['darkblue','cornflowerblue','lightskyblue'][(i-ORFstart_shifted)%3] for i in x]

        ax1 = fig.add_subplot(grid[hpi_idx*2:hpi_idx*2+2,0])
        if ORF != 'ORF1a':
            ax1.axvline(junc_5p-s-0.5,linestyle='--',color='k')
            if hpi_idx == 0:
                #ax1.text(0,ylim+0.1,'Estimated RPM')
                #ax1.text(junc_5p-s,ylim+0.1,'Measured RPM')
                ax1.text(junc_5p-s,ylim+0.1,'Junction')
        ax1.bar(x,arr_RPF_fused,color = bar_colors)
        
        ax1.set_xlim((-0.5,len(arr_RPF_fused)-0.5))
        ax1.text(len(arr_RPF_fused)*0.95,ylim*0.9,hpi.replace('h', ' hpi'))
        ax1.set_xticks(xticks)
        if hpi_idx==(len(hpi_list)-1):
            ax1.set_xticklabels(xticklabels)
            #ax1.set_xlabel('Genomic position (nt)')
        else:
            ax1.set_xticklabels(['' for x in range(len(xticks))])

        if as_log10:
            #ax1.set_ylabel('$\mathregular{log_{10}}$ (Avg. RPM+1)')
            if ylim != None:
                ax1.set_ylim(0,ylim)
        else:
            ax1.set_ylabel('Avg. RPM')
        df_source.append(arr_RPF_fused)
        
    df_source = pd.DataFrame(df_source)
    df_source.index   = hpi_list
    df_source.columns = np.concatenate([np.arange(s,junc_5p), np.arange(junc_3p,e)])+1
    df_source = df_source.T
    df_source.to_csv(SOURCEDATADIR %OutFigname.replace('.pdf','.tsv'), sep = '\t')    

    ######################
    ax2 = fig.add_subplot(grid[len(hpi_list)*2:len(hpi_list)*2+2,0])

    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness*2,contig_len,3),255)),axis=0)
    annot_image = np.concatenate((annot_image[:,s:junc_5p],annot_image[:,junc_3p:e]),axis=1)
    annot_texts = get_annot_textlist(s,junc_5p,thickness)
    annot_texts+= [[i[0],i[1],i[2]+junc_5p-s] for i in get_annot_textlist(junc_3p,e,thickness)]
    for txt,ypos,xpos in annot_texts:
        if txt == '5\'UTR':
            color='white'
        else:
            color='k'
        ax2.text(xpos,ypos,txt,size=12,color=color)
    ax2.text(1,thickness/2,'Leader',size=12,color='white')


    trs_list = get_trs_b_list(s,junc_5p)
    trs_list+= [i+junc_5p-s for i in get_trs_b_list(junc_3p,e)]
    for trs_pos in trs_list:
        if ORF == 'ORF1a':
            TRS = 'TRS-L'
        else:
            TRS = 'TRS-B'
        ax2.text(trs_pos,int(thickness*1.2),TRS,size=12)
        annot_image[int(thickness*0.6):int(thickness*1.0),trs_pos:trs_pos+6] = (255,215,0)

    ax2.text(58-s,int(thickness*1.2),'CUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),58-s:61-s] = (255,87,51)

    ax2.text(ORFstart_shifted,int(thickness*1.2),'AUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),ORFstart_shifted:ORFstart_shifted+3] = (0,255,0)
    txt = SARSCOV2_genome[idx_TISL:junc_5p]+SARSCOV2_genome[junc_3p:]

    for i in range(0,len(txt),3):
        codon = txt[i:i+3]
        if codon in ['TGA','TAA','TAG']:
            putative_ORF_end = i+3
            break
    ax2.text(58-s,thickness*2.5,'ORF from TIS-L (%da.a)' %((putative_ORF_end-3)/3),size=12)
    annot_image[int(thickness*2):int(thickness*3),idx_TISL-s:idx_TISL-s+putative_ORF_end] = (3,181,148)

    ax2.imshow(annot_image,aspect='auto')
    if ORF != 'ORF1a':
        ax2.axvline(junc_5p-s-0.5,linestyle='--',color='k')
        
    ###nucleotide text###
    nt_fused  = (SARSCOV2_genome[s:junc_5p]+SARSCOV2_genome[junc_3p:e]).replace('T','U')
    nt_font_size = min(10,1000//window_size)
    for i in range(len(arr_RPF_fused)):
        ax2.text(i-0.5,0,nt_fused[i],fontsize=nt_font_size)
    #####################
    
    ax2.set_xlim((-0.5,len(arr_RPF_fused)-0.5))
    ax2.set_xticks([])
    ax2.set_yticks([0,2,14.5])
    #ax2.set_yticklabels(['','sgRNA',''])
    ax2.set_yticklabels([])
    ax2.grid(False)
    #####################
    plt.tight_layout()
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
    else:
        plt.show()
    plt.close()
    return None

def vis_cov2annotation(thickness = 5):
    global ORF_palette
    thick_unit = thickness//5
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = False)
    
    contig_len  =  29903
    arr = np.full((thickness,contig_len,3),255) #all white
    arr[0*thick_unit:5*thick_unit] = 0 ##black line at center
    
    for ORFname,row in df_SARSCOV2_annot.iterrows():
        ORF_s,ORF_e     = row['ORFstart'], row['ORFend']
        color   = ORF_palette[ORFname]
        color_int = []
        for i in range(1,len(color),2):
            color_int.append(int(color[i:i+2],16))
        thick_range=range(thickness)
        arr[thick_range,ORF_s:ORF_e] = color_int
    
    return arr

def get_annot_textlist(s,e,thickness):
    textlist = []
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = False)
    
    for ORFname,row in df_SARSCOV2_annot.iterrows():
        ORF_s,ORF_e     = row['ORFstart'], row['ORFend']
        if 'UTR' in ORFname:
            pass
        else:
            pass
        if s<=ORF_s and ORF_s<e:
            if ORF_e <= e: #whole ORF included in the window
                textlist.append([ORFname,thickness//2, (ORF_s+ORF_e)//2 -s-2]) #ORFname, textposition within window
            else:
                textlist.append([ORFname,thickness//2, (ORF_s+e)//2 - s-2])
        elif s<ORF_e and ORF_e<=e:
            textlist.append(    [ORFname,thickness//2, (s+ORF_e)//2 - s-2])
        else:
            pass
        
    return textlist

def get_trs_b_list(s,e):
    global SARSCOV2_genome
    cov2_fa = SARSCOV2_genome[:]
    cov2_fa = cov2_fa[s:e].upper().replace('T','U')
    trs_b_list = []
    for i in re.finditer('ACGAAC',cov2_fa):
        trs_b_list.append(i.start())
    return trs_b_list
#Main plots for Fig. 4 and SFigs. 7,8
tmp_ylim_dict = {'ORF1a':3, 'S': 3,'ORF3a':4,'E':2.5,'M':3,  'ORF6':4,'ORF7a':4,'ORF7b':2,'ORF8':3.5,'N':4}
for tmp_sample_prefix in ['RPF','QTI']:
    for ORF, tmp_ylim in tmp_ylim_dict.items():
        print(tmp_sample_prefix, ORF, time.ctime())
        plot_sgRNA(ORF = ORF, hpi_list = ['16h','24h','36h'], s = 49,
                   assay_prefix = 'RPFSeq', sample_prefix = tmp_sample_prefix, as_RPM = True, offset_RPF = 12,
                   as_log10= True, ylim = tmp_ylim, OutFigname=f'Fig4_{tmp_sample_prefix}_{ORF}.pdf')
def plot_ORF7b_cmp(assay_prefix = 'RPFSeq',sample_prefix = 'RPF',hpi = '36h', ylim=3.5, s = 49,
                   offset_RPF = 12, as_RPM=True, OutFigname = ''):
    global CANONICAL_sgRNA_junc_dict, SARSCOV2_genome
    ORF      = 'ORF7b'
    idx_TISL = 58
    sgRNA_junc_dict = CANONICAL_sgRNA_junc_dict.copy()
    junc_5p,junc_3p = sgRNA_junc_dict[ORF]
    
    contig_len  =  len(SARSCOV2_genome)
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = False)
    
    TISLreads_df = pd.read_csv(RESULTDIR %(f'{sample_prefix}{hpi}_fraction.tsv'),
                               sep = '\t', index_col = 0)
    TISL_pct     = TISLreads_df.loc[ORF,'Percentage']    
    
    grid = GridSpec(nrows=11,ncols=1)
    fig  = plt.figure(figsize=(10,6.67*0.75))
    
    as_log10=True
    
    ORFstart        = df_SARSCOV2_annot.at[ORF,'ORFstart']
    ORFstart_shifted= (ORFstart-junc_3p)+(junc_5p-s)
    window_size     = 50*(ORFstart_shifted//50+1)
    e               = ORFstart + (window_size-ORFstart_shifted)
    
    arr_RPF_list    = []
    reps            = 2
    for rep in range(1,reps+1):
        samplename    = f'{sample_prefix}_{hpi}_rep{rep}'
        _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                        pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                        offset_RPF = offset_RPF,junction_spanning_only=False,
                                        min_rlen=0,as_RPM=as_RPM)
        arr_RPF_list.append(arr_RPF)

    avg_arr_RPF  = np.array(arr_RPF_list).mean(axis=0)
    
    arr_RPF_raw     = avg_arr_RPF[e-window_size:e]
    arr_RPF_fused   = np.concatenate((avg_arr_RPF[s:junc_5p]*TISL_pct/100, avg_arr_RPF[junc_3p:e]))
    
    if window_size == 50:
        ticklabel_width = 10
    else:
        ticklabel_width = 20
    df_source = pd.DataFrame(columns = ['position_ORF7a_sgRNA','logRPM_ORF7a_sgRNA','position_ORF7b_sgRNA','logRPM_ORF7b_sgRNA'])
    df_source['position_ORF7a_sgRNA'] = np.arange(e-window_size, e)+1
    df_source['position_ORF7b_sgRNA'] = np.concatenate([np.arange(s,junc_5p),np.arange(junc_3p,e)])+1
    #######RPF, non-fused########
    arr = arr_RPF_raw
    if as_log10:
        psc = 1
        arr = np.log10(arr+psc)
    df_source['logRPM_ORF7a_sgRNA'] = arr
    thickness       = 5
    
    x = np.arange(len(arr))
    bar_colors = [['darkblue','cornflowerblue','lightskyblue'][(i-ORFstart_shifted)%3] for i in x]
    ax1 = fig.add_subplot(grid[0:3,0])
    ax1.bar(x,arr,color = bar_colors)
    ax1.set_xlim((-0.5,len(arr)-0.5))


    tmp_xticklabels = np.arange(e-window_size,e)+1
    xticks = [i for i,nt in enumerate(tmp_xticklabels) if nt%10 == 0]
    xticklabels = tmp_xticklabels[xticks]
    ax1.set_xticks(xticks)
    ax1.set_xticklabels(xticklabels)
    if as_log10:
        ax1.set_ylabel('$\mathregular{log_{10}}$ (Avg. RPM+1)')
    else:
        ax1.set_ylabel('Avg. RPM')
    if ylim != None:
        ax1.set_ylim(0,ylim)
    #########################
    #######Annotation of raw genome########
    ax2 = fig.add_subplot(grid[3:5,0])
    
    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness,contig_len,3),255)),axis=0)[:,e-window_size:e]
    annot_texts = get_annot_textlist(e-window_size,e,thickness)
    for txt,ypos,xpos in annot_texts:
        ax2.text(xpos,ypos,txt,size=12)
    ax2.text(ORFstart_shifted,int(thickness*1.2),'AUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),ORFstart_shifted:ORFstart_shifted+3] = (0,255,0)
    
    ###nucleotide text###
    nt_raw  = (SARSCOV2_genome[e-window_size:e]).replace('T','U')
    nt_font_size = min(10,1000//window_size)
    for i in range(len(arr_RPF_raw)):
        ax2.text(i-0.5,0,nt_raw[i],fontsize=nt_font_size)
    #####################
    
    ax2.imshow(annot_image,aspect='auto')

    ax2.set_xlim((-0.5,len(arr)-0.5))
    ax2.set_xticks([])
    ax2.set_yticks([0,2,9.5])
    ax2.set_yticklabels(['','',''])
    ax2.grid(False)
    #####################
    #######RPF, fused########
    arr = arr_RPF_fused
    if as_log10:
        psc = 1
        arr = np.log10(arr+psc)
    df_source['logRPM_ORF7b_sgRNA'] = arr
    thickness       = 5
    x = np.arange(len(arr))

    bar_colors = [['darkblue','cornflowerblue','lightskyblue'][(i-ORFstart_shifted)%3] for i in x]
    ax1 = fig.add_subplot(grid[5:8,0])
    ax1.bar(x,arr,color = bar_colors)
    ax1.set_xlim((-0.5,len(arr)-0.5))
    ax1.set_xticks(xticks)

    tmp_xticklabels = np.arange(s,junc_5p)+1 #1-base coordinate
    xticks      = [i for i,nt in enumerate(tmp_xticklabels) if nt%10 == 0]
    xticklabels = tmp_xticklabels[xticks]
    tmp_xticklabels = np.arange(junc_3p,e)+1 #1-base coordinate
    tmp_xticks      = [i for i,nt in enumerate(tmp_xticklabels) if nt%10 == 0]
    tmp_xticklabels = tmp_xticklabels[tmp_xticks]
    xticks = xticks+[i+junc_5p-s for i in tmp_xticks]
    xticklabels= np.concatenate((xticklabels,tmp_xticklabels))
    ax1.set_xticklabels(xticklabels)

    if ORF != 'ORF1a':
        ax1.axvline(junc_5p-s-0.5,linestyle='--',color='k')
        ax1.text(junc_5p-s,ylim+0.1,'Junction')
    if as_log10:
        ax1.set_ylabel('log10 (Avg. RPM+1)')
    else:
        ax1.set_ylabel('Avg. RPM')
    if ylim != None:
        ax1.set_ylim(0,ylim)
    ######################

    ax2 = fig.add_subplot(grid[8:11,0])

    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness*2,contig_len,3),255)),axis=0)
    annot_image = np.concatenate((annot_image[:,s:junc_5p],annot_image[:,junc_3p:e]),axis=1)
    annot_texts = get_annot_textlist(s,junc_5p,thickness)
    annot_texts+= [[i[0],i[1],i[2]+junc_5p-s] for i in get_annot_textlist(junc_3p,e,thickness)]
    for txt,ypos,xpos in annot_texts:
        ax2.text(xpos,ypos,txt,size=12)
    ax2.text(1,thickness/2,'Leader',size=12,color='white')


    trs_list = get_trs_b_list(s,junc_5p)
    trs_list+= [i+junc_5p-s for i in get_trs_b_list(junc_3p,e)]
    for trs_pos in trs_list:
        if ORF == 'ORF1a':
            TRS = 'TRS-L'
        else:
            TRS = 'TRS-B'
        ax2.text(trs_pos,int(thickness*1.6),TRS,size=12)
        annot_image[int(thickness*1):int(thickness*1.5),trs_pos:trs_pos+6] = (255,215,0)

    ax2.text(58-s,int(thickness*1.2),'CUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),58-s:61-s] = (255,87,51)

    txt = SARSCOV2_genome[idx_TISL:junc_5p]+SARSCOV2_genome[junc_3p:]

    for i in range(0,len(txt),3):
        codon = txt[i:i+3]
        if codon in ['TGA','TAA','TAG']:
            putative_ORF_end = i+3
            break

    ax2.text(58-s,thickness*2.5,'ORF from TIS-L (%da.a)' %((putative_ORF_end-3)/3),size=12)
    annot_image[int(thickness*2):int(thickness*3),idx_TISL-s:idx_TISL-s+putative_ORF_end] = (3,181,148)

    ax2.imshow(annot_image,aspect='auto')
    if ORF != 'ORF1a':
        ax2.axvline(junc_5p-s-0.5,linestyle='--',color='k')

    ###nucleotide text###
    nt_fused  = (SARSCOV2_genome[s:junc_5p]+SARSCOV2_genome[junc_3p:e]).replace('T','U')
    nt_font_size = min(10,1000//window_size)
    for i in range(len(arr_RPF_fused)):
        ax2.text(i-0.5,0,nt_fused[i],fontsize=nt_font_size)
    #####################
    
    ax2.set_xlim((-0.5,len(arr_RPF_fused)-0.5))
    ax2.set_xticks([])
    ax2.set_yticks([0,2,14.5])
    ax2.set_yticklabels(['','',''])
    ax2.grid(False)
    #####################
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
        df_source.to_csv(SOURCEDATADIR %OutFigname.replace('.pdf','.tsv'), sep = '\t', index = False)
    plt.close()
    return None

plot_ORF7b_cmp(sample_prefix = 'RPF',hpi = '36h', ylim=3.5, s = 49,
               OutFigname = 'SFig7g_RPF_ORF7b.pdf')
plot_ORF7b_cmp(sample_prefix = 'QTI',hpi = '36h', ylim=3.5, s = 49,
               OutFigname = 'SFig8k_QTI_ORF7b.pdf')<jupyter_output><empty_output><jupyter_text>ORF10 sgRNA production and ORF10 translation (SFig. 9)<jupyter_code>def get_junctionread_pairs(assay_prefix='mRNASeq',sample_prefix='RPFpaired',hpi='48h', rep=1,
                           norm_by_nh=True,
                           pos_5end_range=(0,265), f2_5p_range=None, fetch_range=(0,265),
                           min_rlen=(32,31,30,29,28,27),as_RPM=True):
    #Mainly for analyzing mRNA-seq data, but designed as it can be applied to RPF-seq data analysis
    samplename = f'{sample_prefix}_{hpi}_rep{rep}'
    fname = BAMDIR_cov2mapped %(assay_prefix,samplename)
    bam = pysam.AlignmentFile(fname)
    
    contig = bam.references[0]
    jpair_count_dict = defaultdict(float) #{(junction s, e): weighted count}
    for read in bam.fetch(contig,fetch_range[0],fetch_range[1]):
        cigartuples = read.cigartuples
        pos_5end    = read.reference_start
        rlen        = read.infer_read_length()
        nh          = read.get_tag('NH')
        assigned    = False
        if norm_by_nh:
            norm_ct = 1/nh
        else:
            norm_ct = 1
        ####read filtering####
        contain_junction = (3 in [i[0] for i in cigartuples])
        if not contain_junction:
            continue
            
        if not ((pos_5end_range[0]<= pos_5end) and (pos_5end< pos_5end_range[1])):
            continue
            
        if type(min_rlen) == int:
            if rlen<min_rlen:
                continue
        else:
            if rlen<min_rlen[pos_5end-pos_5end_range[0]]:
                continue
        
        #############################
        cur_pos = pos_5end
        f1_detected = False
        f1_3p = None
        f2_5p = None
        for idx,cigar_tp in enumerate(cigartuples):
            operation, length = cigar_tp
            if operation == 3: #skip
                f1_detected = True
                cur_pos += length
            else:
                if operation == 4 or operation == 1:    # softclip or insertion
                    pass
                elif operation == 0 or operation == 2:  # match/mismatch or deletion
                    cur_pos += length
            if operation == 0:
                if not f1_detected: #junction 1st frag
                    f1_3p = cur_pos
                else: #second fragment
                    f2_5p = cur_pos - length
                    break
        assert f1_3p != None
        assert f2_5p != None
        
        if f2_5p_range != None:
            if not ((f2_5p_range[0]<= f2_5p) and (f2_5p< f2_5p_range[1])):
                continue
        jpair_count_dict[(f1_3p,f2_5p)] += norm_ct
        
    bam.close()
    jpair_count_dict = dict(jpair_count_dict) #defaultdict to dict
    if as_RPM:
        viral_mapped_reads = 0
        host_mapped_reads  = 0
        #virus reads
        bam = pysam.AlignmentFile(fname)
        for read in bam.fetch():
            viral_mapped_reads += (1/(read.get_tag('NH')))
        bam.close()
        #human reads
        fname_hg38 = fname.replace('_SARSCOV2align','_hostalign')
        bam = pysam.AlignmentFile(fname_hg38)
        for read in bam.fetch():
            host_mapped_reads += (1/(read.get_tag('NH')))
        bam.close()
        total_mapped_reads = host_mapped_reads + viral_mapped_reads
        for jpair in jpair_count_dict: 
            jpair_count_dict[jpair] /= (total_mapped_reads / 1e+06)
        
    return jpair_count_dict

def junctionpair_analysis(assay_prefix='mRNASeq',sample_prefix_list=['RPFpaired','QTIpaired'],hpi='48h',
                      pos_5end_range=(0,265), f2_5p_range=(28240,28280), fetch_range=(28240,28280),
                      min_rlen=1):
    
    reps = 3 if hpi == '48h' else 2
    jpair_dict_list = []
    for sample_prefix in sample_prefix_list:
        for rep in range(1,reps+1):
            jpair_count_dict = get_junctionread_pairs(assay_prefix=assay_prefix,sample_prefix=sample_prefix,
                                                      hpi=hpi, rep=rep,norm_by_nh=True,
                                                      pos_5end_range=pos_5end_range, f2_5p_range=f2_5p_range, 
                                                      fetch_range=fetch_range, min_rlen=min_rlen, as_RPM=True)
            
            jpair_dict_list.append(jpair_count_dict)
    ndict = len(jpair_dict_list)
    total_jpairs = set([jpair for jpair_count_dict in jpair_dict_list for jpair in jpair_count_dict.keys() ])
    
    avg_jpair_count_list = []
    for jpair in total_jpairs:
        count = 0
        for jpair_count_dict in jpair_dict_list:
            try:
                count += jpair_count_dict[jpair]
            except KeyError:
                pass
        count /= ndict
        avg_jpair_count_list.append([*jpair,count]) #junction start, end, count
    avg_jpair_count_list = sorted(avg_jpair_count_list,key = lambda x: x[2],reverse = True)
    
    return avg_jpair_count_list

def plot_ORF10_junctions(hpi,f2_5p_range=(29457,29607),U5_display_range=(58,88),display_top=10,
                         OutFigname=''): #-100,+50
    canonical_AUG = 29557 #ORF10 start idx 
    idx_TISL = 58
    ticklabel_width = 20
    avg_jpair_count_list = junctionpair_analysis(assay_prefix='mRNASeq',sample_prefix_list=['RPFpaired','QTIpaired'],hpi=hpi,
                                                 pos_5end_range=(0,265), f2_5p_range=f2_5p_range, 
                                                 fetch_range=f2_5p_range,min_rlen=1)
    if len(avg_jpair_count_list) == 0:
        first_jpair= [np.nan, np.nan, np.nan]
        first_jpair_pct = np.nan
    else:
        total_jpairs    = sum([i[2] for i in avg_jpair_count_list]) #junction start, end, count
        first_jpair     = avg_jpair_count_list[0]
        first_jpair_pct = first_jpair[2]/total_jpairs*100
        
    TISL_frame_list = [0,0,0]
    for f1_3p,f2_5p,RPM in avg_jpair_count_list:
        dist_annotAUG_TISL = (canonical_AUG-f2_5p)+(f1_3p-idx_TISL)
        TISL_frame_list[dist_annotAUG_TISL%3] += RPM
    TISL_frame_list = [RPM/total_jpairs*100 for RPM in TISL_frame_list]
    
    
    
    f1_3p_arr = np.zeros(U5_display_range[1]-U5_display_range[0],dtype=np.float)
    f2_5p_arr = np.zeros(f2_5p_range[1]     -f2_5p_range[0],     dtype=np.float)
    
    rank_junc = 1
    for f1_3p,f2_5p,RPM in avg_jpair_count_list[:display_top]:
        f1_3p_arr[f1_3p-U5_display_range[0]] += RPM
        f2_5p_arr[f2_5p-f2_5p_range[0]]      += RPM
        rank_junc +=1
        
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,2, sharey='row', figsize=(8,4),
                                 gridspec_kw={'width_ratios': [len(f1_3p_arr),len(f2_5p_arr)],
                                              'height_ratios':[3,1]})
    
    ax1.bar( np.arange(*U5_display_range),f1_3p_arr)   
    ax1.spines['right'].set_visible(False)
    ax1.set_ylabel('RPM')
    ax1.set_xlim(U5_display_range[0]-0.5, U5_display_range[1]-0.5)
    tmp_xticklabels = np.arange(U5_display_range[0],U5_display_range[1])
    ##assigned value at the last 'position' because f1_3p is end index of [s,e)
    xticks      = [nt for nt in range(*U5_display_range) if nt%ticklabel_width == 0]
    xticklabels = xticks
    ax1.set_xticks(xticks)
    ax1.set_xticklabels(xticklabels)
    
    ax2.bar(np.arange(*f2_5p_range),  f2_5p_arr)
    ax2.spines['left'].set_visible(False)
    ax2.set_xlabel('Genomic position (nt)')
    ax2.set_xlim(f2_5p_range[0]-0.5,      f2_5p_range[1]-0.5)
    tmp_xticklabels = np.arange(f2_5p_range[0],f2_5p_range[1])+1 #1-base coordinate
    xticks      = [i for i,nt in enumerate(tmp_xticklabels) if nt%ticklabel_width == 0]
    xticklabels = tmp_xticklabels[xticks]
    xticks      = xticklabels-1
    
    ax2.set_xticks(xticks)
    ax2.set_xticklabels(xticklabels)
    
    ###Annotation: 5UTR###
    thickness   = 5
    contig_len  = len(SARSCOV2_genome)
    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness*2,contig_len,3),255)),axis=0)
    annot_image = annot_image[:,U5_display_range[0]:U5_display_range[1]]
    annot_texts = get_annot_textlist(U5_display_range[0],U5_display_range[1],thickness)
    for txt,ypos,xpos in annot_texts:
        if txt == '5\'UTR':
            color='white'
        else:
            color='k'
        ax3.text(xpos,ypos,txt,size=12,color=color)
    ax3.text(1,thickness/2,'Leader',size=12,color='white')


    trs_list = get_trs_b_list(U5_display_range[0], U5_display_range[1])
    
    for trs_pos in trs_list:
        TRS = 'TRS-L'
        ax3.text(trs_pos,int(thickness*1.2),TRS,size=12)
        annot_image[int(thickness*0.6):int(thickness*1.0),trs_pos:trs_pos+6] = (255,215,0)

    ax3.text(58-U5_display_range[0],int(thickness*1.2),'CUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),
                58-U5_display_range[0]:61-U5_display_range[0]] = (255,87,51)
    
    ax3.imshow(annot_image,aspect='auto')
        
    ###nucleotide text###
    nt_raw  = SARSCOV2_genome[U5_display_range[0]:U5_display_range[1]].replace('T','U')
    nt_font_size = min(10,1000//(f2_5p_range[1]-f2_5p_range[0]))
    for i in range(U5_display_range[1]-U5_display_range[0]):
        ax3.text(i-0.5,0,nt_raw[i],fontsize=nt_font_size)
    #####################
    
    ax3.set_xlim((-0.5,U5_display_range[1]-U5_display_range[0]-0.5))
    ax3.set_xticks([])
    ax3.set_yticks([0,2,14.5])
    ax3.set_yticklabels([])
    ax3.grid(False)
    #################################################
    ###Annotation: ORF10###
    thickness   = 5
    contig_len  = len(SARSCOV2_genome)
    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness*2,contig_len,3),255)),axis=0)
    annot_image = annot_image[:,f2_5p_range[0]:f2_5p_range[1]]
    annot_texts = get_annot_textlist(f2_5p_range[0],f2_5p_range[1],thickness)
    for txt,ypos,xpos in annot_texts:
        if txt == '5\'UTR':
            color='white'
        else:
            color='k'
        ax4.text(xpos,ypos,txt,size=12,color=color)
    
    ax4.imshow(annot_image,aspect='auto')
        
    ###nucleotide text###
    nt_raw  = SARSCOV2_genome[f2_5p_range[0]:f2_5p_range[1]].replace('T','U')
    nt_font_size = min(10,1000//(f2_5p_range[1]-f2_5p_range[0]))
    for i in range(f2_5p_range[1]-f2_5p_range[0]):
        ax4.text(i-0.5,0,nt_raw[i],fontsize=nt_font_size)
    #####################
    
    ax4.set_xlim((-0.5,f2_5p_range[1]-f2_5p_range[0]-0.5))
    ax4.set_xticks([])
    ax4.set_yticks([0,2,14.5])
    ax4.set_yticklabels([])
    ax4.grid(False)
    #################################################
    plt.tight_layout()
    for f1_3p,f2_5p,RPM in avg_jpair_count_list[:display_top]:
        # 1. Get transformation operators for axis and figure
        ax1tr = ax1.transData # ax1 -> Display
        ax2tr = ax2.transData # ax2 -> Display
        figtr = fig.transFigure.inverted() # Display -> Figure
        # 2. Transform arrow start point from axis 0 to figure coordinates
        ptB = figtr.transform(ax1tr.transform((f1_3p, f1_3p_arr[f1_3p-U5_display_range[0]]) ))
        # 3. Transform arrow end point from axis 1 to figure coordinates
        ptE = figtr.transform(ax2tr.transform((f2_5p, f2_5p_arr[f2_5p-f2_5p_range[0]])))
        # 4. Create the patch
        arrow = mpl.patches.FancyArrowPatch(
            ptB, ptE, transform=fig.transFigure,  # Place arrow in figure coord system
            fc = "k", connectionstyle="arc3,rad=-0.2", arrowstyle='simple', alpha = 1,
            mutation_scale = 20
        )
        # 5. Add patch to list of objects to draw onto the figure
        fig.patches.append(arrow)
    
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
    plt.close()
    return None
plot_ORF10_junctions('36h',f2_5p_range=(29457,29607),display_top=5, OutFigname='SFig9_ORF10sgRNA_36h.pdf')
def plot_ORF10_RPF(hpi_list = ['16h','24h','36h'],
               assay_prefix = 'RPFSeq', sample_prefix = 'RPF', as_RPM = True, offset_RPF = 12,
               as_log10= True, ylim= 2.5, OutFigname=''):
    global CANONICAL_sgRNA_junc_dict, SARSCOV2_genome
    ORF = 'ORF10'
    sgRNA_junc_dict = CANONICAL_sgRNA_junc_dict.copy()
    
    idx_TISL = 58
    contig_len  =  len(SARSCOV2_genome)
    grid     = GridSpec(nrows=11,ncols=1)
    fig      = plt.figure(figsize=(10,1.5*(1+len(hpi_list))))
    thickness= 5
    
    df_SARSCOV2_annot = get_df_SARSCOV2_annot(exc_UTR = True)
    ORFstart          = df_SARSCOV2_annot.at[ORF,'ORFstart']
    
    window_size       = 100
    ticklabel_width   = 20
    e                 = ORFstart + 50
    s                 = e-window_size
    ORFstart_shifted  = ORFstart-s
    
    df_source = pd.DataFrame(columns = hpi_list, index = np.arange(s,e)+1)
    
    for hpi_idx,hpi in enumerate(hpi_list):
        arr_RPF_list  = []
        reps = 3 if hpi == '48h' else 2
        for rep in range(1,reps+1):
            samplename = f'{sample_prefix}_{hpi}_rep{rep}'
            _, _, arr_RPF = get_readpos_arr(samplename,assay_prefix,norm_by_nh=True,
                                            pos_5end_range=(0,29903), f2_5p_range=None, fetch_range=(0,29903),
                                            offset_RPF = offset_RPF,junction_spanning_only=False,
                                            min_rlen=0,as_RPM=as_RPM)
            arr_RPF_list.append(arr_RPF)
        
        avg_arr_RPF  = np.array(arr_RPF_list).mean(axis=0)
        arr_RPF_raw  = avg_arr_RPF[e-window_size:e]
        df_source[hpi] = arr_RPF_raw
        if as_log10:
            psc = 1.0
            arr_RPF_raw = np.log10(arr_RPF_raw+psc)
        
        x = np.arange(len(arr_RPF_raw))
        tmp_xticklabels = np.arange(e-window_size,e)+1
        xticks = [i for i,nt in enumerate(tmp_xticklabels) if nt%ticklabel_width == 0]
        xticklabels = tmp_xticklabels[xticks]

        bar_colors = [['darkblue','cornflowerblue','lightskyblue'][(i-ORFstart_shifted)%3] for i in x]

        ax1 = fig.add_subplot(grid[hpi_idx*3:(hpi_idx+1)*3,0])
        
        ax1.bar(x,arr_RPF_raw,color = bar_colors)
        
        ax1.set_xlim((-0.5,len(arr_RPF_raw)-0.5))
        ax1.text(len(arr_RPF_raw)*0.95,ylim*0.9,hpi.replace('h', ' hpi'))
        ax1.set_xticks(xticks)
        if hpi_idx==(len(hpi_list)-1):
            ax1.set_xticklabels(xticklabels)
            #ax1.set_xlabel('Genomic position (nt)')
        else:
            ax1.set_xticklabels(['' for x in range(len(xticks))])

        if as_log10:
            if ylim != None:
                ax1.set_ylim(0,ylim)
        else:
            ax1.set_ylabel('Avg. RPM')
    ######################
    ax2 = fig.add_subplot(grid[9:11,0])

    annot_image = vis_cov2annotation(thickness)
    annot_image = np.concatenate((annot_image,np.full((thickness,contig_len,3),255)),axis=0)[:,e-window_size:e]
    annot_texts = get_annot_textlist(e-window_size,e,thickness)
    for txt,ypos,xpos in annot_texts:
        ax2.text(xpos,ypos,txt,size=12)
    trs_list = get_trs_b_list(e-window_size,e)
    for trs_pos in trs_list:
        if ORF == 'ORF1a':
            TRS = 'TRS-L'
        else:
            TRS = 'TRS-B'
        ax2.text(trs_pos,int(thickness*1.6),TRS,size=12)
        annot_image[int(thickness*1):int(thickness*1.5),trs_pos:trs_pos+6] = (255,215,0)
    ax2.text(ORFstart_shifted,int(thickness*1.2),'AUG',size=12)
    annot_image[int(thickness*0.6):int(thickness*1.0),ORFstart_shifted:ORFstart_shifted+3] = (0,255,0)
    
    ###nucleotide text###
    nt_raw  = (SARSCOV2_genome[e-window_size:e]).replace('T','U')
    nt_font_size = min(10,1000//window_size)
    for i in range(len(arr_RPF_raw)):
        ax2.text(i-0.5,0,nt_raw[i],fontsize=nt_font_size)
    #####################
    
    ax2.imshow(annot_image,aspect='auto')

    ax2.set_xlim((-0.5,len(arr_RPF_raw)-0.5))
    ax2.set_xticks([])
    ax2.set_yticks([0,2,9.5])
    ax2.set_yticklabels([])
    ax2.grid(False)
    #####################
    plt.tight_layout()
    if OutFigname != '':
        plt.savefig(FIGDIR %OutFigname)
        df_source.to_csv(SOURCEDATADIR %OutFigname.replace('.pdf','.tsv'), sep = '\t')
    else:
        plt.show()
    plt.close()
    return None


plot_ORF10_RPF(hpi_list = ['16h','24h','36h'],
               assay_prefix = 'RPFSeq', sample_prefix = 'RPF', as_RPM = True, offset_RPF = 12,
               as_log10= True, ylim= 2.5, OutFigname='SFig9b_RPF_ORF10.pdf')
plot_ORF10_RPF(hpi_list = ['16h','24h','36h'],
               assay_prefix = 'RPFSeq', sample_prefix = 'QTI', as_RPM = True, offset_RPF = 12,
               as_log10= True, ylim= 2.5, OutFigname='SFig9c_QTI_ORF10.pdf')
def triplet_periodicity_ORF(samplename,assay_prefix = 'RPFSeq',genome = 'SARSCOV2', ORF = 'ORF10',
                            offset=12, target_rlen = None,
                            offset_from_5p = True, verbose=False, return_codon1pos1 = False):
    
    if genome == 'SARSCOV2':
        BAMDIR = BAMDIR_cov2mapped
    else:
        BAMDIR = BAMDIR_hostmapped
        
    if genome == 'hg38':
        refdir = REFFLAT_hg38 
    elif genome == 'SARSCOV2':
        refdir = REFFLAT_SARSCOV2 
    elif genome == 'chlSab2':
        refdir = REFFLAT_chlSab2
    else:
        return None
    
    chrdict = get_dict_refgenes(refdir)
    
    bam = pysam.AlignmentFile(BAMDIR %(assay_prefix,samplename), 'rb')
    triplet_counts = np.zeros(3,dtype=np.float32)
    codon1pos1 = 0.0
    for sChrID_strand,refgene_list in chrdict.items():
        refgene_list = [refgene for refgene in refgene_list if refgene.sGeneSym == ORF]
        assert len(refgene_list)<2
        sChrID     = sChrID_strand[:-1]
        sStrandDir = sChrID_strand[-1]
        gene_negstrand = (sStrandDir == '-')
        
        for refgene in refgene_list:
            len_ORF = refgene.nExonlen - refgene.nU5len - refgene.nU3len
            if len_ORF%3 != 0:
                continue
            exon_s_list = refgene.nExonStartlist
            exon_e_list = refgene.nExonEndlist
            arr_RPF     = np.zeros(refgene.nExonlen,dtype=np.float32)
            offset_arr  = 0
            for exon_s, exon_e in zip(exon_s_list,exon_e_list):
                for read in bam.fetch(sChrID, exon_s,exon_e):
                    if target_rlen != None:
                        rlen = read.infer_read_length()
                        if rlen!=target_rlen:
                            continue
                    if read.is_reverse != gene_negstrand:
                        continue
                    reference_positions = read.get_reference_positions(full_length=False)    
                    if None in reference_positions:
                        print('None in read reference position?')
                        print(reference_positions)
                        print(read.cigarstring)
                        return None
                    if gene_negstrand:
                        if offset_from_5p == True:
                            refpos   = reference_positions[-(offset+1)]
                        else:
                            refpos   = reference_positions[offset]
                    else:
                        if offset_from_5p == True:
                            refpos   = reference_positions[offset]
                        else:
                            refpos   = reference_positions[-(offset+1)]
                    
                    nh = read.get_tag('NH')
                    if (exon_s <= refpos) and (refpos < exon_e):
                        arr_RPF[refpos-exon_s+offset_arr] += (1/nh)
                offset_arr += (exon_e - exon_s)
                
            assert offset_arr == len(arr_RPF)
            if sChrID == '-':
                arr_RPF = arr_RPF[::-1]
            arr_RPF_ORF = arr_RPF[refgene.nU5len:]
            if refgene.nU3len>0: 
                arr_RPF_ORF = arr_RPF_ORF[:-refgene.nU3len]
            codon1pos1 += arr_RPF_ORF[0]
            
            to_add = arr_RPF_ORF.reshape((len_ORF//3,3)).sum(axis=0)
            triplet_counts += to_add

    triplet_periodicity = triplet_counts/(triplet_counts.sum())
    
    if verbose:
        print(samplename,genome, triplet_counts.sum(), triplet_periodicity, sep = '\t')
    bam.close()
    if return_codon1pos1:
        return triplet_counts, triplet_periodicity, codon1pos1
    else:
        return triplet_counts, triplet_periodicity
    
def calc_plot_periodicity_ORF(assay_prefix = 'RPFSeq', sample_prefix='RPF', load_precalc=False, ORF='ORF10',
                              OutF='RPF_ORF10_triplet.merged.tsv', OutFigname=''):
    hpi_list = ['%dh' %i for i in [24,36,48]]+['48hCaco2'] ##read sum >100
    df_merged  = []
    if load_precalc:
        df_merged = pd.read_csv(RESULTDIR %(OutF),sep='\t')
    else:
        for hpi in hpi_list:
            print(hpi,time.ctime())
            for genome in ['SARSCOV2']:
                merged_counts = np.zeros(3,dtype=np.float32)
                reps = 3 if (hpi == '48h') else 2
                for rep in range(1, 1+reps):
                    samplename = f'{sample_prefix}_{hpi}_rep{rep}'
                    triplet_counts, triplet_periodicity = triplet_periodicity_ORF(samplename, assay_prefix = assay_prefix, ORF = ORF,
                                                                                  genome = genome, offset=12, target_rlen = None,
                                                                                  offset_from_5p = True, verbose=True, 
                                                                                  return_codon1pos1 = False)
                    merged_counts += triplet_counts

                df_merged.append([hpi,genome,*merged_counts,*(merged_counts/merged_counts.sum()*100)])

        df_merged  = pd.DataFrame(df_merged)
        df_merged.columns = ['hpi','genome','pos1','pos2','pos3','pct1','pct2','pct3']
        df_merged.to_csv(RESULTDIR %(OutF), sep='\t',index=False)   
    ###########
    df_merged = df_merged.set_index('genome')
    
    df_merged = df_merged[df_merged['hpi'].isin(hpi_list)]
    df_merged.index = df_merged.index.str.replace('hg38','host')
    df_merged.index = df_merged.index.str.replace('chlSab2','host')
    fig = plt.figure(figsize=(len(hpi_list)+2,4))
    ax = fig.add_subplot(111)
    x = np.arange(len(hpi_list))
    bar_width=0.8/3
    xpos_offset = -bar_width*1.5
    colors = {}
    
    colors['SARSCOV2'] = ['#ff6666','#fc8d59','#fdcc8a']
    
    for genome in ['SARSCOV2']:
        for i in range(3):
            pct = f'pct{i+1}'
            val = df_merged.loc[genome,pct]
            ax.bar(x+xpos_offset,val,width = bar_width,label = f'{genome} nt {i+1}', color= colors[genome][i])
            xpos_offset += bar_width
    
    ax.set_xticks(x)
    ax.set_xticklabels(hpi_list)
    ax.set_xlabel('hpi')
    ax.set_xlim((-1,len(hpi_list)))
    ax.set_ylabel('Fraction of reads (%)')
    ax.set_ylim(top=70)
    
    plt.tight_layout()
    if OutFigname == '':
        plt.show()
    else:
        plt.savefig(FIGDIR %OutFigname)
    plt.close()
    
    return None
calc_plot_periodicity_ORF(assay_prefix = 'RPFSeq', sample_prefix='RPF', load_precalc=False, ORF='ORF10',
                              OutF='RPF_ORF10_triplet.merged.tsv', OutFigname='SFig9d_ORF10_triplet_RPF.pdf')
calc_plot_periodicity_ORF(assay_prefix = 'RPFSeq', sample_prefix='QTI', load_precalc=False, ORF='ORF10',
                              OutF='QTI_ORF10_triplet.merged.tsv', OutFigname='SFig9d_ORF10_triplet_QTI.pdf')<jupyter_output>24h Mon Sep 27 16:30:41 2021
RPF_24h_rep1	SARSCOV2	141.0	[0.65957445 0.17730497 0.16312057]
RPF_24h_rep2	SARSCOV2	108.0	[0.6944444  0.18518518 0.12037037]
36h Mon Sep 27 16:30:41 2021
RPF_36h_rep1	SARSCOV2	309.0	[0.68608415 0.16828479 0.14563107]
RPF_36h_rep2	SARSCOV2	317.0	[0.55520505 0.24921136 0.1955836 ]
48h Mon Sep 27 16:30:41 2021
RPF_48h_rep1	SARSCOV2	315.0	[0.53968257 0.2825397  0.17777778]
RPF_48h_rep2	SARSCOV2	446.0	[0.5560538  0.25336322 0.19058296]
RPF_48h_rep3	SARSCOV2	332.0	[0.57831323 0.28915662 0.13253012]
48hCaco2 Mon Sep 27 16:30:42 2021
RPF_48hCaco2_rep1	SARSCOV2	224.0	[0.54910713 0.27232143 0.17857143]
RPF_48hCaco2_rep2	SARSCOV2	170.0	[0.5470588  0.24705882 0.20588236]
24h Mon Sep 27 16:30:42 2021
QTI_24h_rep1	SARSCOV2	66.0	[0.6515151  0.24242425 0.10606061]
QTI_24h_rep2	SARSCOV2	65.0	[0.52307695 0.23076923 0.24615385]
36h Mon Sep 27 16:30:42 2021
QTI_36h_rep1	SARSCOV2	97.0	[0.52577317 0.20618556 0.26804122]
QTI_36h_rep2	SARSCOV2	127.0	[0.43307087 0.2913386  0.27559[...]
{'fraction_non_alphanumeric': 0.131115257871124, 'fraction_numerical': 0.04970024218843054, 'mean_word_length': 2.4259909137898203, 'pattern_counts': {'":': 0, '<': 97, '<?xml version=': 0, '>': 62, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 122, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '516284', 'license': 'MIT', 'n_tokens_mistral': 52202, 'n_tokens_neox': 47466, 'n_words': 7066, 'path': 'codes/Kimetal_analysis_SARSCOV2part.ipynb', 'repo_name': 'DoyeonKim-BI/Kimetal_SARSCOV2translatome', 'size': '125935'}