Using py3Dmol

in Lessons
#software

In this blog post, I want to introduce a very useful tool: py3Dmol. py3Dmol is python wrapper for the 3Dmol javascript library, which can be used to create visualizations of molecular systems. I find this tool incredibly useful because it means we can visualize a system inside a jupyter notebook. This means that when you're setting up a system, or analyzing results about it, you don't have to leave the notebook, which means you get quick feedback, and can easily share results. For this reason, I've added it to the PyBigDFT library as the InlineVisualizor class.

Unfortunately, the documentation is lacking. I think the developers are mainly concerned with making it easy for you to show a molecule on a website. But when using it from a jupyter notebook, you're inevitably going to do more complex things. In this blog post, I want to show some uses of it. 

import py3Dmol

Basics

The first thing you need to know about using py3Dmol is that you should always think about structure in terms of a file format. That is the most straight forward way to use its API. So let's try reading in a PDB file...

with open("1crn.pdb") as ifile:
    system = "".join([x for x in ifile])

Now let's try visualizing this system.

view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(system)
view.setStyle({'model': -1}, {"cartoon": {'color': 'spectrum'}})
view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

Notice that the way parameters work when setting the style is that you pass a dictionary. This is because the javascript API itself is based on json.

Highlighting

Now one thing I often want to do is to highlight a particular part of the system I am studying. To do that, we can manually modify the style of individual atoms.

view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(system)

i = 0
for line in system.split("\n"):
    split = line.split()
    if len(split) == 0 or split[0] != "ATOM":
        continue
    if split[3] == "PRO":
        color = "yellow"
    else:
        color = "black"
    idx = int(split[1])

    view.setStyle({'model': -1, 'serial': i+1}, {"cartoon": {'color': color}})
    i += 1
view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

This kind of access is also important when you want to mix multiple styles.

ion = "ATOM    1  NE   NE     1      10.047  14.099   3.625  0.1812  1.5500       NE"

system2 = system + ion

view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(system2)
view.setStyle({'model': -1}, {"cartoon": {'color': 'spectrum'}})

i = 0
for line in system2.split("\n"):
    split = line.split()
    if len(split) == 0 or split[0] != "ATOM":
        continue
    if split[-1] == "NE":
        view.setStyle({'model': -1, 'serial': i+1}, {"sphere": {'color': "pink"}})
    i += 1

view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

Datastructures

Up until now, the code I've been presenting can be described as, at best, pretty hacky. This is because we don't have a datastructure for our molecules. So let's build one. But let's not just build any old structure, instead I want to create a constructor that gives us maximum flexibility when dealing with py3Dmol.

In my opinion, the best approach is to break down a system into two levels: molecules and atoms. We will consider molecules to be lists of atoms. And atoms will be dictionaries which can be flexible about what kind of data they store. For this blog post, we just need to write a molecule and atom that can convert to and from pdb.

class Atom(dict):
    def __init__(self, line):
        self["type"] = line[0:6].strip()
        self["idx"] = line[6:11].strip()
        self["name"] = line[12:16].strip()
        self["resname"] = line[17:20].strip()
        self["resid"] = int(int(line[22:26]))
        self["x"] = float(line[30:38])
        self["y"] = float(line[38:46])
        self["z"] = float(line[46:54])
        self["sym"] = line[76:78].strip()

    def __str__(self):
        line = list(" " * 80)

        line[0:6] = self["type"].ljust(6)
        line[6:11] = self["idx"].ljust(5)
        line[12:16] = self["name"].ljust(4)
        line[17:20] = self["resname"].ljust(3)
        line[22:26] = str(self["resid"]).ljust(4)
        line[30:38] = str(self["x"]).rjust(8)
        line[38:46] = str(self["y"]).rjust(8)
        line[46:54] = str(self["z"]).rjust(8)
        line[76:78] = self["sym"].rjust(2)
        return "".join(line) + "\n"

class Molecule(list):
    def __init__(self, file):
        for line in file:
            if "ATOM" in line or "HETATM" in line:
                self.append(Atom(line))

    def __str__(self):
        outstr = ""
        for at in self:
            outstr += str(at)

        return outstr

with open("1crn.pdb") as ifile:
    mol = Molecule(ifile)

Now we can verify that we can do a basic visualization.

view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(str(mol))
view.setStyle({'model': -1}, {"cartoon": {'color': 'spectrum'}})
view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

Now let's get to the real meat: how do we ensure flexibility? The answer lies in the dictionary representation of the atoms. What I suggest is that we add arbitrary pymol commands at an atomic level:

for at in mol:
    if at["resname"] == "PRO":
        at["pymol"] = {"cartoon": {'color': "yellow"}}
    elif at["resname"] == "GLY":
        at["pymol"] = {"cartoon": {'color': 'blue'}}

Then we can access that when we are performing our visualization.

view = py3Dmol.view(width=400, height=300)
view.addModelsAsFrames(str(mol))
for i, at in enumerate(mol):
    default = {"cartoon": {'color': 'black'}}
    view.setStyle({'model': -1, 'serial': i+1}, at.get("pymol", default))

view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

Let's finish this demonstration with a movie.

molecules = []
for i in range(25):
    with open("1crn-"+str(i)+".pdb") as ifile:
        molecules.append(Molecule(ifile))

for mol in molecules:
    for at in mol:
        if at["resname"] == "PRO":
            at["pymol"] = {"stick": {'color': "yellow"}}
        elif at["resname"] == "GLY":
            at["pymol"] = {"stick": {'color': 'blue'}}

The main thing you need to know about animiations is simply that you have to combine your file format into a multimodel format. 

view = py3Dmol.view(width=400, height=300)

models = ""
for i, mol in enumerate(molecules):
    models += "MODEL " + str(i) + "\n"
    models += str(mol)
    models += "ENDMDL\n"
view.addModelsAsFrames(models)

for i, at in enumerate(molecules[0]):
    default = {"stick": {'color': 'black'}}
    view.setStyle({'model': -1, 'serial': i+1}, at.get("pymol", default))

view.zoomTo()
view.animate({'loop': "forward"})
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol