atom2svg

atom2svg.py

# atom2svg.py by @biochem_fan  
# VERSION: 210224
# License: GPLv2 or later (ask me if this is inconvenient for you)
#
# LIMITATIONS:
# - This supports only a single chain. If you have more chains, shift residue IDs and merge chains.
#
# TODO:
# - Residue numbers (how to pack??)
#
# HISTORY:
# - Repaired Ile's abbreviation from L to I (pointed out by @QiuyeLi)

from __future__ import print_function
import sys
from collections import OrderedDict

# All values are in Angstrom, even thought the code says "px". They are scaled by 10 for rendering.
RADIUS = 1.2 # Radius of circles for sidechains
MARGIN = 5.0 # Margin of the whole figure
CB_FUDGE = 1.4 # Extend CA-CB vectors a bit so that sidechains stick out
# For other styles, edit the <style> tag below.

one_letter_code = {'ARG': 'R', 'HIS': 'H', 'LYS': 'K', 'ASP': 'D', 'GLU': 'E',
                   'SER': 'S', 'THR': 'T', 'ASN': 'N', 'GLN': 'Q', 'CYS': 'C',
                   'GLY': 'G', 'PRO': 'P', 'ALA': 'A', 'VAL': 'V', 'ILE': 'I',
                   'LEU': 'L', 'MET': 'M', 'PHE': 'F', 'TYR': 'Y', 'TRP': 'W'}

# a: acidic, b: basic, w: hydrophilic, n: hydrophobic, g: glycine, s: sulfer containing
residue_class = {'ARG': 'b', 'HIS': 'b', 'LYS': 'b', 'ASP': 'a', 'GLU': 'a',
                 'SER': 'w', 'THR': 'w', 'ASN': 'w', 'GLN': 'w', 'CYS': 's',
                 'GLY': 'g', 'PRO': 'p', 'ALA': 'n', 'VAL': 'n', 'ILE': 'n',
                 'LEU': 'n', 'MET': 's', 'PHE': 'n', 'TYR': 'n', 'TRP': 'n'}

if len(sys.argv) != 3:
    print("Usage: python atom2svg.py input.pdb output.svg")
    sys.exit(-1)

c_beta = OrderedDict()
c_alpha = OrderedDict() 
for line in open(sys.argv[1]):
    if line[0:4] != "ATOM": continue

    try:
        resn = line[17:20]
        name = line[12:16]
        x = float(line[30:38])
        y = float(line[38:46])
        resid = int(line[22:26])
        if name == " CA ":
            c_alpha[resid] = (x, y, resn)
        if name == " CB " or (name == " CA " and resn == "GLY"):
            c_beta[resid] = (x, y, resn)
    except:
        sys.stderr.write("Failed to parse line: %s" % line)

minx = -MARGIN + min([ca[0] for ca in c_alpha.values()])
miny = -MARGIN + min([ca[1] for ca in c_alpha.values()])
maxx =  MARGIN + max([ca[0] for ca in c_alpha.values()])
maxy =  MARGIN + max([ca[1] for ca in c_alpha.values()])

f = open(sys.argv[2], "w")

f.write('''<?xml version="1.0" standalone="no"?>
<svg viewBox="0 0 %d %d" version="1.1" xmlns="http://www.w3.org/2000/svg">
<style>
 line {stroke: black; stroke-width: 0.2px;}
 .mainchain {stroke-width: 0.4px;}
 .sidechain {stroke-width: 0.2px;}

 circle {stroke: black; stroke-width: 0.2px;}
 .a {fill: red;}
 .b {fill: skyblue;}
 .w {fill: green;}
 .g {fill: pink;}
 .n {fill: white;}
 .s {fill: yellow;}
 .p {fill: purple;}
 .aa {font-family: sans-serif; font-size: 1.76px; text-anchor: middle;}
</style>
<g transform="scale(10) translate(%f, %f)">
''' % (10 * (maxx - minx), 10 * (maxy - miny), -minx, -miny))

# Draw main chain
f.write("\n")
for resi in c_alpha.keys():
    if (resi - 1) not in c_alpha: continue 

    me = c_alpha[resi]
    prev = c_alpha[resi - 1]
    f.write(' <line x1="%f" y1="%f" x2="%f" y2="%f" class="mainchain" />\n' % (me[0], me[1], prev[0], prev[1]))

# Draw side chains
f.write("\n")
for resi, cb in c_beta.items():
    pos = cb

    if cb[2] != 'GLY':
        pos = (CB_FUDGE * cb[0] + (1 - CB_FUDGE) * c_alpha[resi][0], 
               CB_FUDGE * cb[1] + (1 - CB_FUDGE) * c_alpha[resi][1])
        f.write(' <line x1="%f" y1="%f" x2="%f" y2="%f" class="sidechain" />\n' % (pos[0], pos[1], c_alpha[resi][0], c_alpha[resi][1]))

    f.write(' <circle cx="%f" cy="%f" r="%f" class="%s" />\n' % (pos[0], pos[1], RADIUS, residue_class[cb[2]]))
    f.write(' <text x="%f" y="%f" class="aa">%s</text>\n\n' % (pos[0], pos[1] + RADIUS / 2.0, one_letter_code[cb[2]]))

f.write('</g>\n')
f.write('</svg>')