import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.transforms as transforms
from matplotlib.transforms import Affine2D
import math
from matplotlib.patheffects import RendererBase
import matplotlib.patheffects as PathEffects
import seaborn as sns
# Calculate the font size that match the unit length on x-axis or y-axis
def CalculateFontsize(xlim, ylim, ax, fig, rows, cols, unit_scale=1):
# Get axis limits
axXlim = ax.get_xlim()
axYlim = ax.get_ylim()
# Get figure dimensions in pixels
fig_width, fig_height = fig.get_size_inches() * fig.dpi
# Get axis dimensions in pixels
bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
ax_width, ax_height = bbox.width * fig.dpi, bbox.height * fig.dpi
# Calculate font size proportional to axis units
fontsize_x = unit_scale * ax_width / (axXlim[1] - axXlim[0]) / cols * (xlim[1] - xlim[0])
fontsize_y = 0.8 * unit_scale * ax_height / (axYlim[1] - axYlim[0]) / rows * (ylim[1] - ylim[0])
# Use the minimum of the two to keep the font size consistent
fontsize = min(fontsize_x, fontsize_y)
return fontsize
def GetStartEnd(msa, start, end):
length = len(msa[0])
if (start == None):
start = 0
elif (start < 0):
start = length + start
if (end == None):
end = length - 1
elif (end < 0):
end = length + end
return start, end
def GetColorMap(preset = None, msa=None, color_order = None, palette = None):
color_map = {}
if preset in ["dna", "nuc", "nucleotide"]:
color_map['A'] = [0,1,0]
color_map['C'] = [1,165/255,0]
color_map['G'] = [1,0,0]
color_map['T'] = [0.5,0.5,1]
else:
if (color_order != None):
for i,c in enumerate(color_order):
color_map[c] = sns.color_palette(palette)[i]
# The case some of the alphabet color is not specified in the order
for a in msa:
for c in a:
if c not in color_map:
size = len(color_map)
color_map[c] = sns.color_palette(palette)[size]
color_map['-'] = [1,1,1] # white
color_map['.'] = [1,1,1] # white
return color_map
def DrawMSA(msa, seq_names = None, start = None, end = None,
axlim = None, color_map = None, palette=None, ax=None, fig=None,
show_char=True):
# Get the canvas attributes.
ax = ax or plt.gca()
fig = fig or ax.get_figure()
renderer = fig.canvas.get_renderer()
height = len(msa)
length = len(msa[0])
# start, end: draw the [start,end] (both inclusive) region of the MSA
start, end = GetStartEnd(msa, start, end)
if (axlim == None):
fontsize = CalculateFontsize(ax.get_xlim(), ax.get_ylim(), ax, fig, height, end - start + 1)
else:
fontsize = CalculateFontsize(axlim[0], axlim[1], ax, fig, height, end - start + 1)
color_map = color_map or GetColorMap(msa=msa, color_order=None, palette=palette)
lengthUnit = 1 / (end - start + 1)
heightUnit = 1 / height
if (axlim != None):
lengthUnit = (axlim[0][1] - axlim[0][0]) / (end - start + 1)
heightUnit = (axlim[1][1] - axlim[1][0]) / height
for i, a in enumerate(msa):
for j,c in enumerate(a[start:end+1]):
linewidth = min(2,fontsize/50)
if show_char:
text = ax.text(x=(j + 0.5)*lengthUnit, y=(i+0.5) * heightUnit, s=c, color="black",
va="center_baseline", ha="center", fontsize=fontsize,
transform=ax.transAxes if axlim == None else ax.transData)
text.set_path_effects([PathEffects.withStroke(linewidth=linewidth,
foreground='w')])
ax.add_patch( patches.Rectangle(xy=(j * lengthUnit, i * heightUnit),
width = lengthUnit, height=heightUnit,
facecolor=color_map[c], linewidth=linewidth, edgecolor="white",
transform=ax.transAxes if axlim == None else ax.transData))
if (axlim == None):
ax.set_xlim(-0.5, end - start + 1 - 0.5)
ax.set_ylim(0-0.5, height-0.5)
# Set the x ticks adaptively at point easy to count
ticks = []
tickLabels = []
candidateSteps = [1,5,10,20,50,100,500,1000,5000,10000]
step = 1
for i,s in enumerate(candidateSteps):
if (s * 5 > end - start + 1):
if (i > 0):
step = candidateSteps[i - 1]
break
tickStart = (int)(start / step) * step + step - 1
if (tickStart != start):
ticks.append(0)
tickLabels.append(start+1)
for i in range(tickStart, end + 1, step):
if (i >= length):
break
ticks.append(i - start)
tickLabels.append(i+1)
if (tickLabels[-1] != min(length, end + 1)):
ticks.append(min(length - 1, end) - start)
tickLabels.append(min(length, end+1))
ax.set_xticks(ticks, tickLabels)
# Set the y
ticks = []
tickLabels = []
for i in range(height):
ticks.append(i)
tickLabels.append(i if seq_names == None else seq_names[i])
ax.set_yticks(ticks, tickLabels)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
return ax, color_map
def GetConsensus(msa, start = None, end = None):
start, end = GetStartEnd(msa, start, end)
consensus = []
consensusComposition = []
for j in range(start, end + 1):
composition = {}
if (j >= len(msa[0])):
consensus.append('')
consensusComposition.append({"":0})
continue
for i,a in enumerate(msa):
if (a[j] not in composition):
composition[a[j]] = 0
composition[a[j]] += 1
result = ""
maxCnt = 0
for c in composition:
if (composition[c] > maxCnt):
maxCnt = composition[c]
result = c
consensus.append(result)
consensusComposition.append(composition)
return consensus, consensusComposition
def SetConsensusAxTicks(consensus, consensusComposition, ax):
ticks = []
tickLabels = []
for i in range(len(consensus)):
ticks.append(i)
label = consensus[i]
for c in consensusComposition[i]:
if (c == label):
continue
if (consensusComposition[i][c] == consensusComposition[i][label]):
label = 'X'
break
tickLabels.append(label)
ax.set_xticks(ticks, tickLabels)
def DrawConsensusHisto(msa, color = [0, 0, 1], color_map = None, start = None, end = None, ax=None):
ax = ax or plt.gca()
consensus, consensusComposition = GetConsensus(msa, start, end)
binHeight = []
colors = []
for i in range(len(consensus)):
if (sum(consensusComposition[i].values()) == 0):
binHeight.append(0)
else:
binHeight.append( consensusComposition[i][consensus[i]] /
sum(consensusComposition[i].values()) )
colors.append([1-binHeight[i] * (1-color[0]),
1-binHeight[i] * (1-color[1]),
1-binHeight[i] * (1-color[2])]) # color-base in blue
ax.bar(x=list(range(0, len(consensus))), height=binHeight,
color = colors,
width=0.95)
ax.set(ylim=(0, 1))
ax.set(xlim=(-0.5, len(consensus)-0.5))
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.get_yaxis().set_ticks([])
SetConsensusAxTicks(consensus, consensusComposition, ax)
# Code from pyseqlogo for scale font to one direction
class Scale(RendererBase):
"""Scale alphabets using affine transformation"""
def __init__(self, sx, sy=None):
self._sx = sx
self._sy = sy
def draw_path(self, renderer, gc, tpath, affine, rgbFace):
affine = Affine2D().scale(self._sx, self._sy) + affine
renderer.draw_path(gc, tpath, affine, rgbFace)
def CalculateEntropy(count):
s = sum(count)
if (s == 0):
return 0
return sum([-c/s * math.log2(c/s) for c in count])
def DrawSeqLogo(msa, color_map, alphabet_size = None, start = None, end = None, ax=None):
ax = ax or plt.gca()
fig = ax.get_figure()
start, end = GetStartEnd(msa, start, end)
alphabet_size = len(color_map) - 2 # 2 for "-" and "."
consensus, consensusComposition = GetConsensus(msa, start, end)
# Definition from https://en.wikipedia.org/wiki/Sequence_logo
# Search for appropriate height.
r = []
adjuste = 1 / math.log(2) * (alphabet_size - 1) / (2 * len(msa[0]))
for i,c in enumerate(consensus):
entropy = CalculateEntropy(list(consensusComposition[i].values()))
r.append(math.log2(alphabet_size) - (entropy + adjuste))
ax.set_xlim(-0.5, end-start+1-0.5)
ax.set_ylim(0, max(r))
fontsize = CalculateFontsize(ax.get_xlim(), ax.get_ylim(),
ax, fig, 1, end - start + 1)
lengthUnit = 1
for i,c in enumerate(consensusComposition):
prevy = 0
totalCount = sum(list(c.values()))
if (totalCount == 0):
continue
for j,item in enumerate(sorted(c.items(), key=lambda x:x[1], reverse=True)):
k = item[0]
v = item[1]
text = ax.text(x=i * lengthUnit, y=prevy, s=k, fontsize=fontsize,
va="baseline", ha="center", color = (color_map[k] if (k not in ['.','-']) else [0,0,0]),
transform=ax.transData)
height = v / totalCount * r[i]
tbox = text.get_window_extent(text._renderer).transformed(ax.transData.inverted())
scale = height / (tbox.y1 - prevy)
#print(i, j, height, scale, tbox)
prevy = prevy + height
text.set_path_effects([Scale(1.0, scale)])
ax.axis('off')
SetConsensusAxTicks(consensus, consensusComposition, ax)
# Add the annotation for the sequence alignment
# annotations: a list of numbers like [['a',0,3]]: msa[0..3] (both inclusive) is annotated as the name 'a'
def DrawAnnotation(msa, annotations, color_map=None,start = None, end = None, ax=None):
ax = ax or plt.gca()
fig = ax.get_figure()
start, end = GetStartEnd(msa, start, end)
ax.set_xlim(start - 0.5, end + 0.5)
ax.set_ylim(0, 1)
fontsize = CalculateFontsize(ax.get_xlim(), ax.get_ylim(), ax, fig, 1, end - start + 1)
for a in annotations:
text = ax.text(x=(a[1]+a[2])/2, y=0.5, s=a[0], fontsize=fontsize,
va="center", ha="center", color="black", clip_on=True)
tbox = text.get_window_extent(text._renderer).transformed(ax.transData.inverted())
# Draw the bracket
ax.plot([a[1], a[1]], [0, 1], color="black", clip_on=True)
ax.plot([a[2], a[2]], [0, 1], color="black", clip_on=True)
ax.plot([a[1], tbox.x0], [0.5, 0.5], color="black", clip_on=True)
ax.plot([tbox.x1, a[2]], [0.5, 0.5], color="black", clip_on=True)
ax.axis('off')
# Draw multipanel MSA
def DrawComplexMSA(msa, panels=[], seq_names = None, panel_height_ratios=None, panel_params=None,
color_map=None, start=None, end=None, wrap=None, figsize=None):
color_map = color_map or GetColorMap(msa=msa)
start,end = GetStartEnd(msa, start,end)
wrap = wrap or (end - start + 1)
chunks = math.ceil((end - start + 1) / wrap)
height_ratios = None
if (panel_height_ratios is None):
panel_height_ratios = []
for p in panels:
if (p == DrawMSA):
panel_height_ratios.append(len(msa))
elif (p == DrawAnnotation):
panel_height_ratios.append(0.5)
else:
panel_height_ratios.append(1)
height_ratios = panel_height_ratios * chunks
fig,axes = plt.subplots(len(panels) * chunks, 1, constrained_layout=True,
figsize=figsize, height_ratios = height_ratios)
axidx = 0
for i in range(start, end + 1, wrap):
for j,func in enumerate(panels):
extraParam = {}
if (panel_params is not None):
extraParam = panel_params[j].copy()
if func is DrawMSA:
extraParam['seq_names'] = seq_names
func(msa, color_map = color_map, start=i, end=i+wrap-1,ax=axes[axidx], **extraParam)
axidx += 1
return axes