Simple Gene track plot from Ensembl mart export

DNA-helix.py

from matplotlib.collections import LineCollection
from matplotlib.colors import to_rgb
from matplotlib import pyplot as plt 
import numpy as np


def draw_DNA_helix(bp=50, major_groove=22, minor_groove=12, bp_per_turn=10.5, 
                   fig_height=1.5, colors=('red', 'black'), savefig=None, fig_format='png'):
    """
    
    major groove is 22A minor is 12A (B-DNA) https://doi.org/10.1038%2F287755a0
    
    :param 50 bp: number of base-pairs to draw
    :param 22 major_groove: in Angstroms, length of major groove in the helix.
    :param 12 minor_groove: in Angstroms, length of minor groove in the helix.
    :param 10.5 bp_per_turn: default is for B-DNA described by Watson and Crick
    :param ('red', 'black') colors: colors of each strand
    :param None savefig: path to store figure
    :param 'png' fig_format: format of the saved figure
    
    :returns: matplotlib axe object
    """

    length = bp / 10 * 2 * np.pi
    offset = (minor_groove + major_groove) / (2 * minor_groove)

    # create random base pairs
    xb = np.arange(0.1, length, 2 * np.pi / bp_per_turn)
    yu = np.sin(xb)
    yd = np.sin(xb + np.pi / offset)
    
    r1, g1, b1 = to_rgb(colors[0])
    r2, g2, b2 = to_rgb(colors[1])
            
    colors = []
    segments = []
    for xx, yb, yt in zip(xb, yu, yd):
        if yt > yb:
            colors.extend([(r1, g1, b1, 0.5), (r2, g2, b2, 0.5)])
            yt, yb = yb, yt
        else:
            colors.extend([(r2, g2, b2, 0.5), (r1, g1, b1, 0.5)])
        diff = (yb - yt) / 3
        if np.random.random() > 0.5:
            segments.append([[xx, yt], [xx, yb - diff]])
            segments.append([[xx, yt + 2 * diff], [xx, yb]])
        else:
            segments.append([[xx, yt], [xx, yb - diff * 2]])
            segments.append([[xx, yt + diff], [xx, yb]])

    lc0 = LineCollection(segments, linewidths=1, colors=colors)

    # create DNA strands
    xd = np.linspace(0, length, 100 * bp)
    y1 = np.sin(xd)
    y2 = np.sin(xd + np.pi / offset)

    # create + strand
    div = fig_height * 3
    lw1 = fig_height * (2 + np.sin(xd + 2 * np.pi / offset - 0.05 * fig_height))
    points1 = np.array([xd, y1]).T.reshape(-1, 1, 2)
    segments1 = np.concatenate([points1[:-1], points1[1:]], axis=1)
    lc1 = LineCollection(segments1, linewidths=lw1,
                         colors=[(r1, g1, b1, v / div) for v in lw1])

    # create - strand
    lw2 = fig_height * (2 + np.sin(xd + np.pi / offset / 2))
    points2 = np.array([xd, y2]).T.reshape(-1, 1, 2)
    segments2 = np.concatenate([points2[:-1], points2[1:]], axis=1)
    lc2 = LineCollection(segments2, linewidths=lw2, 
                         colors=[(r2, g2, b2, v / div) for v in lw2])

    # plot
    _, axe = plt.subplots(figsize=(1 + bp * 0.2, fig_height))

    axe.add_collection(lc0)
    axe.add_collection(lc1)
    axe.add_collection(lc2)

    axe.set_xlim(0, length)
    axe.set_ylim(-1.2, 1.2)

    _ = axe.axis('off')
    if not savefig is None:
        plt.savefig(savefig, format=fig_format)
    return axe

Gene_track.ipynb

example:

axe = draw_DNA_helix(bp=50)
plt.show()

resulting in: