plot.py

'''
Some plot functions
Copyright (c) 2016 Peng Zhang <zhpn1024@163.com>
'''

import matplotlib
matplotlib.use('pdf')
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
matplotlib.rcParams['svg.fonttype'] = 'none'
from matplotlib.pylab import *

def plotTrans(t, ypos = 0, intv = None, r = [0.1, 0.3], color = 'blue',rid = -0.5):
  '''plot transcript
  '''
  plot([t.start,t.stop],[ypos,ypos],color=color)
  # arrows
  x, y = [], []
  if intv is None : intv = len(t) / 20
  for i in range(t.start+intv//2, t.stop-intv//3, intv):
    x.append(i)
    y.append(ypos)
  arr = '>'
  if t.is_reverse() : arr = "<"
  plot(x,y,'w'+arr)
  # blocks
  x = [[],[]]
  y = [[],[]]
  orf = t(start = t.thick_start, stop = t.thick_stop)
  for e in t.exons:
    for es in e - orf : # UTR 
      x[0].append(es.start)
      y[0].append(len(es))
    for ei in e.intersect(orf): # CDS
      x[1].append(ei.start)
      y[1].append(len(ei))
  bar(x[0],[r[0]*2]*len(x[0]),width=y[0],bottom=ypos-r[0],edgecolor=color,color=color, align='edge')
  bar(x[1],[r[1]*2]*len(x[1]),width=y[1],bottom=ypos-r[1],edgecolor=color,color=color, align='edge')
  text((t.start+t.stop)//2, ypos+rid, t.id)

def save(filename):
  savefig(filename, transparent=True)

def riboShow(ax, trans, cnts, start = 0, stop = -1, ymax = None, scale = 1, col = ['r','g','b'], title = '', showlegend = False, showframe = True, bottom = 0.8, height = 0.1):
  '''plot riboseq profile
  '''
  if stop < start : stop = trans.cdna_length()
  rlen = stop - start
  lx = [[], [], []]
  ly = [[], [], []]
  m = 1
  for j in range(rlen):
    p = j + start
    if cnts[j + start] > 0 : 
      i = p % 3
      lx[i].append(j)
      y = cnts[p] * scale
      if ymax is not None and ymax > 0 and y > ymax: y = ymax
      ly[i].append(y)
      if m < y : m = y
  ylab = 'Count'
  if scale != 1 : ylab = 'Scaled count'
  if ymax is None or ymax < 0 : ymax = m
  for i in range(3):
    ax.bar(lx[i], ly[i], color=col[i], width=1, edgecolor=col[i], log=False, alpha=0.4, label='Frame '+str(i+1), align='edge')
  
  [ax.spines[side].set_visible(False) for side in ('right','top','bottom')]
  ax.yaxis.set_ticks_position('left')
  ax.xaxis.set_ticks_position('bottom')
  ax.set_xlim((0, rlen))
  ax.set_ylim((0, ymax))
  ax.label_outer() # setp(ax.get_xticklabels(), visible=False)
  ax.set_ylabel(ylab)
  ax.set_title(title)
  if showlegend :
    try : ax.legend(loc='best', frameon=False)
    except : pass
  if not showframe : return
  fx = [[],[],[]]
  fy = [[],[],[]]
  fw = [[],[],[]]
  for i in range(3):
    last = False
    for j in range(i, stop, 3) :
      if j < start : continue
      top = True
      if cnts[j] <= 0 : top = False
      if j-1 >= 0 and cnts[j] <= cnts[j-1] : top = False
      if j+1 < stop and cnts[j] <= cnts[j+1] : top = False
      if top : ##
        if last : tw += 3
        else : tx, tw = j-start, 1
        last = True
      else : 
        if last and tw > 3 :
          fx[i].append(tx)
          fy[i].append(ymax * height)
          fw[i].append(tw)
        last = False
    if last and tw > 3 :
      fx[i].append(tx)
      fy[i].append(ymax * 0.1)
      fw[i].append(tw)
  for i in range(3):
    ax.bar(fx[i], fy[i], color=col[i], bottom = ymax * bottom,width=fw[i], alpha=0.2, linewidth = 0, align='edge')
  
def orfShow(ax, orfs, start = 0, stop = -1, col = ['r','g','b'], cds = [None, None], title = 'Potential ORFs in 3 reading frames', alt = True, morecds = None, morecdsbox = False, morecdslabel = None, markpept = None):
  '''plot possible ORFs
  '''
  if stop < start : stop = trans.cdna_length()
  rlen = stop - start
  # ORF in 3 frames
  lx = [[],[],[]]
  ly = [[],[],[]]
  orf_s = []
  for o in orfs:
    if not alt and len(o.starts) == 0 : continue
    if 0 <= o.stop <= start or o.start(alt=alt) >= stop : continue
    orf_s.append(o)
    o1 = o.start(alt=alt) - start
    if o1 < 0: o1 = 0
    lx[o.frame-1].append(o1) # (o.start(alt=alt) - start)
    if o.has_stop():
      o2 = min(stop, o.stop) - max(o.start(alt=alt), start)
    else:
      o2 = stop - max(o.start(), start)
    ly[o.frame-1].append(o2)
    #if o.has_stop(): ly[o.frame-1].append(o.stop-o.start(alt=alt))
    #else : ly[o.frame-1].append(rlen)
  for i in range(3):
    ax.bar(lx[i], [0.2]*len(lx[i]), color=col[i], bottom=2-i+0.4, width=ly[i], alpha=0.4, linewidth=0, align='edge')
  # annotated ORF
  if cds[0] is not None and not (cds[0]>stop or cds[1]< start): 
    i = cds[0] % 3
    newcds = [c - start for c in cds]
    if newcds[0] < 0: newcds[0] = 0
    if cds[1] > stop: newcds[1] = rlen
    ax.text(max(newcds[0],0), 2-i+0.8, 'Annotated ORF', color=col[i])
    ax.bar(newcds[0], [0.4] ,color=col[i], bottom=2-i+0.3, width=newcds[1]-newcds[0], alpha=0.3, edgecolor=col[i], linewidth=2, align='edge')
  if morecds is not None:
    for j, mc in enumerate(morecds):
      if mc[0]>stop or mc[1]< start: continue
      i = mc[0] % 3
      newcds = [c - start for c in mc]
      if newcds[0] < 0: newcds[0] = 0
      if mc[1] > stop: newcds[1] = rlen
      #ax.text(newcds[0], 2-i+0.8, 'Annotated ORF', color=col[i])
      ax.bar(newcds[0], [0.4] ,color=col[i], bottom=2-i+0.3, width=newcds[1]-newcds[0], alpha=0.3, edgecolor=col[i], linewidth=2, align='edge')
      if morecdsbox:
        ax.bar(newcds[0], [0.4] ,color='None', bottom=2-i+0.3, width=newcds[1]-newcds[0], alpha=0.8, fill=False, edgecolor=col[i], linewidth=2, align='edge')
      if morecdslabel is not None:
        ax.text(max(newcds[0],0), 2-i+0.8, morecdslabel[j], color=col[i])
  if markpept is not None:
    for j, mc in enumerate(markpept):
      if mc[0]>stop or mc[1]< start: continue
      i = mc[0] % 3
      newcds = [c - start for c in mc]
      if newcds[0] < 0: newcds[0] = 0
      if mc[1] > stop: newcds[1] = rlen
      ax.bar(newcds[0], [0.2], color='k', bottom=2-i+0.4, width=newcds[1]-newcds[0], alpha=0.8, linewidth=0, align='edge')
  # start & stop codons
  lx = [[],[],[]]
  ly = [[],[],[]]
  lz = [[],[],[]]
  for o in orf_s:
    lx[o.frame-1] += [s - start for s in o.starts if start<=s<stop]
    if alt : ly[o.frame-1] += [s - start for s in o.altstarts if start<=s<stop] # orf.altstarts
    if o.has_stop() and start<=o.stop-3<stop: lz[o.frame-1] += [o.stop - 3 - start]
  for i in range(3):
    ax.bar(ly[i], [0.4]*len(ly[i]), color='yellow', bottom=2-i+0.3, width=3, alpha=0.6, edgecolor='yellow', align='edge')
    ax.bar(lx[i], [0.4]*len(lx[i]), color='lime', bottom=2-i+0.3, width=3, alpha=1, edgecolor='lime', align='edge')
    ax.bar(lx[i], [0.04]*len(lx[i]), color='w', bottom=2-i+0.48, width=3, edgecolor='w', align='edge')
    ax.bar(lz[i], [0.4]*len(lz[i]), color='red', bottom=2-i+0.3, width=3, alpha=1, edgecolor='red', align='edge')
    ax.bar(lz[i], [0.04]*len(lz[i]), color='k', bottom=2-i+0.48, width=3, edgecolor='k', align='edge')
  ax.set_title(title)
  

def hapShow(hapdata, main_ax, tree_ax, size = (10, 5), dist = None, markgrp = None, title = None): # hapdata need to be well ordered
  mg = {}
  if markgrp is not None:
    for i, s in enumerate(markgrp.split(',')):
      mg[s] = 'C{}'.format(i)
  #print(markgrp, mg)
  #gs = GridSpec(1, 4, wspace=0.0)
  #main_ax = subplot(gs[0,1:4])
  main_ax.set_frame_on(False)
  main_ax.yaxis.set_visible(False)
  main_ax.set_xlabel('Variants')
  if title is not None: main_ax.set_title(title)
  #tree_ax = subplot(gs[0,0],sharey=main_ax)
  tree_ax.set_frame_on(False)
  tree_ax.xaxis.set_visible(False)
  tree_ax.set_ylabel('Haplotypes')
  #tree_ax.axis('off')
  #fp, axarr = subplots(1, 2, sharey=True, figsize= size)
  y = 0
  ybd = {} # y boundaries for groups
  hmerge = {} # local hap merge, degeneracy
  n = len(hapdata)
  for j, d in enumerate(hapdata):
    #print(d)
    grp, cnt, hap, dis = d
    y1 = y + cnt
    for i in range(1, len(grp)+1):
      g = grp[0:i]
      if g not in ybd: ybd[g] = [y, y1, []]
      else: ybd[g][1] = y1
    ybd[grp][2].append((y, y1, dis))
    p = hap.find('1')
    while p >= 0:
      p1 = hap.find('0', p)
      if p1 < 0: p1 = len(hap)
      width = p1 - p
      key = p, p1
      if key in hmerge and y+cnt <= hmerge[key]: pass # merged by previous block
      else:
        cntb = cnt
        for j1 in range(j+1, n):
          grp1, cnt1, hap1, dis1 = hapdata[j1]
          if hap1[p:p1].find('0') >= 0: break
          cntb += cnt1
        if cntb > cnt: 
          yb = y + cntb
          hmerge[key] = yb
        main_ax.bar(x=p, height=cntb, bottom=y, width=width, align='edge', color='black', alpha=1)
      p = hap.find('1', p1)
    y = y1

  glevels = {}
  #print(ybd)
  for g in ybd:
    l = len(g)
    if l not in glevels: glevels[l] = {}
    glevels[l][g] = ybd[g]
  if markgrp is None:
    for i, g in enumerate(glevels[1]):
      mg[g] = 'C{}'.format(i)

  for g in mg:
    if g in ybd: main_ax.bar(x=0, height=ybd[g][1]-ybd[g][0], bottom=ybd[g][0], width=len(hap), align='edge', color=mg[g], alpha=0.2)
  for l in range(4, 0, -1):
    #print(l)
    if l not in glevels: continue
    for g in sorted(glevels[l]):
      y, y1, leafs = glevels[l][g]
      color = 'black'
      if g in mg: color = mg[g]
      else:
        for i in range(1, len(g)+1):
          gup = g[0:i]
          if gup in mg:
            color = mg[gup]
            #break
      #ymean = (y + y1) / 2
      if dist is None: d0, d1 = l-1, l ### 
      else:
        d1 = dist[g]
        if len(g) > 1: d0 = dist[g[:-1]]
        else: d0 = 0
      for leaf in leafs:
        ly, ly1, ldis = leaf
        lymean = (ly + ly1) / 2
        if ly+1 < ly1: tree_ax.plot([ldis, ldis], [ly, ly1], color=color, alpha=0.3)
        tree_ax.plot([d1, ldis], [lymean, lymean], color=color, alpha=0.3)
        for i in range(1, len(g)+1):
          gup = g[0:i]
          if ybd[gup][0] == ly: ybd[gup][0] = lymean
          if ybd[gup][1] == ly1: ybd[gup][1] = lymean
      y, y1, leafs = glevels[l][g] # ybd[g]
      ymean = (y + y1) / 2
      if y+1 < y1: tree_ax.plot([d1, d1], [y, y1], color=color, alpha=0.3)
      tree_ax.plot([d0, d1], [ymean, ymean], color=color, alpha=0.3)
      if g in mg: tree_ax.text(d0, ymean, g, color=color, va='center')
      #print(l, ymean, g)
      for i in range(1, len(g)+1):
        gup = g[0:i]
        if ybd[gup][0] == y: ybd[gup][0] = ymean
        if ybd[gup][1] == y1: ybd[gup][1] = ymean