Script Tutorial

Introduction

One of the more powerful features of PyMOL is that it supports Python scripting. That gives you the power of using all the Python libraries. You can even use the Python API to write programs in other languages and then send the results back into PyMOL (this is what cealign does). The PyMOLWiki has a rather extensive script_library full of useful scripts (feel free to add your own).

General Scripts

Scripting follows a simple recipe, in PyMOL.

To write them:

1. Write the function, let's call it doSimpleThing, in a Python file, let's call the file pyProgram.py.
2. Add the following command to the end of the pyProgram.py file

   cmd.extend(doSimpleThing,doSimpleThing)
   

To use them:

1. simply import the script into PyMOL:

   run /home/userName/path/toscript/pyProgram.py
   

That's it. Your script can, through Python, import any modules you need and also edit modify objects in PyMOL.

Getting PyMOL Data into your Script

To get PyMOL data into your script you will need to somehow get access to the PyMOL objects and pull out the data. For example, if you want the atomic coordinates of a selection of alpha carbon atoms your Python function may do something like this (all PyMOL functions are referenced in the See Also section, below):

# Import PyMOL's stored module.  This will allow us with a 
# way to pull out the PyMOL data and modify it in our script.
# See below.
from pymol import stored

def functionName( userSelection ):
    # this array will be used to hold the coordinates.  It
    # has access to PyMOL objects and, we have access to it.
    stored.alphaCabons = []

    # let's just get the alpha carbons, so make the
    # selection just for them
    userSelection = userSelection + " and n. CA"

    # iterate over state 1, or the userSelection -- this just means
    # for each item in the selection do what the next parameter says.
    # And, that is to append the (x,y,z) coordinates to the stored.alphaCarbon
    # array.
    cmd.iterate_state(1, selector.process(userSelection), "stored.alphaCarbons.append([x,y,z])")

    # stored.alphaCarbons now has the data you want.

    ... do something to your coordinates ...

Getting Data From your Script into PyMOL

Usually this step is easier. To get your data into PyMOL, it's usually through modifying some object, rotating a molecule, for example. To do that, you can use the alter or alter_state commands. Let's say for example, that we have translated the molecular coordinates from the last example by some vector (we moved the alpha carbons). Now, we want to make the change and see it in PyMOL. To write the coordinates back we do:

# we need to know which PyMOL object to modify.  There could be many molecules and objects
# in the session, and we don't want to ruin them.  The following line, gets the object
# name from PyMOL
objName = cmd.identify(sel2,1)[0][0]

# Now, we alter each (x,y,z) array for the object, by popping out the values
# in stored.alphaCarbons.  PyMOL should now reflect the changed coordinates.
cmd.alter_state(1,objName,"(x,y,z)=stored.alphaCarbons.pop(0)")

Example

Here's a script I wrote for cealign. It takes two selections of equal length and computes the optimal overlap, and aligns them. This is Kabsch.

def optAlign( sel1, sel2 ):
        """
        optAlign performs the Kabsch alignment algorithm upon the alpha-carbons of two selections.
        Example:   optAlign MOL1 and i. 20-40, MOL2 and i. 102-122
        Example 2: optAlign 1GGZ and i. 4-146 and n. CA, 1CLL and i. 4-146 and n. CA

        Two RMSDs are returned.  One comes from the Kabsch algorithm and the other from
        PyMol based upon your selections.

        By default, this program will optimally align the ALPHA CARBONS of the selections provided.
        To turn off this feature remove the lines between the commented "REMOVE ALPHA CARBONS" below.

        @param sel1: First PyMol selection with N-atoms
        @param sel2: Second PyMol selection with N-atoms
        """

        # make the lists for holding coordinates
        # partial lists
        stored.sel1 = []
        stored.sel2 = []
        # full lists
        stored.mol1 = []
        stored.mol2 = []

        # -- CUT HERE
        sel1 = sel1 + " and N. CA"
        sel2 = sel2 + " and N. CA"
        # -- CUT HERE

        # Get the selected coordinates.  We
        # align these coords.
        cmd.iterate_state(1, selector.process(sel1), "stored.sel1.append([x,y,z])")
        cmd.iterate_state(1, selector.process(sel2), "stored.sel2.append([x,y,z])")

        # get molecule name
        mol1 = cmd.identify(sel1,1)[0][0]
        mol2 = cmd.identify(sel2,1)[0][0]

        # Get all molecule coords.  We do this because
        # we have to rotate the whole molcule, not just
        # the aligned selection
        cmd.iterate_state(1, mol1, "stored.mol1.append([x,y,z])")
        cmd.iterate_state(1, mol2, "stored.mol2.append([x,y,z])")

        # check for consistency
        assert( len(stored.sel1) == len(stored.sel2))
        L = len(stored.sel1)
        assert( L > 0 )

        # must alway center the two proteins to avoid
        # affine transformations.  Center the two proteins
        # to their selections.
        COM1 = numpy.sum(stored.sel1,axis=0) / float(L)
        COM2 = numpy.sum(stored.sel2,axis=0) / float(L)
        stored.sel1 = stored.sel1 - COM1
        stored.sel2 = stored.sel2 - COM2

        # Initial residual, see Kabsch.
        E0 = numpy.sum( numpy.sum(stored.sel1 * stored.sel1,axis=0),axis=0) + numpy.sum( numpy.sum(stored.sel2 * stored.sel2,axis=0)
,axis=0)

        #
        # This beautiful step provides the answer.  V and Wt are the orthonormal
        # bases that when multiplied by each other give us the rotation matrix, U.
        # S, (Sigma, from SVD) provides us with the error!  Isn't SVD great!
       # HACK COMMENT:
        # Numpy has some strangeness with returning floats from
        # its calculation.  So, I made a rather silly work around,
        # but, it seems to work!  I cast the float to a string,
        # then back to a float; this works.  See below.

        # We already have our solution, in the results from SVD.
        # we just need to check for reflections and then produce
        # the rotation.  V and Wt are orthonormal, so their det's
        # are +/-1.0 (and thus products are +/- 1.0 ).
        reflect = float(str(float(numpy.linalg.det(V) * numpy.linalg.det(Wt))))

        if reflect == -1.0:
                S[-1] = -S[-1]
                V[:,-1] = -V[:,-1]

        RMSD = E0 - (2.0 * sum(S))
        RMSD = numpy.sqrt(abs(RMSD / L))

        #U is simply V*Wt
        U = numpy.dot(V, Wt)

        # rotate and translate the molecule
        stored.sel2 = numpy.dot((stored.mol2 - COM2), U) + COM1
        stored.sel2 = stored.sel2.tolist()

        # let PyMol know about the changes to the coordinates
        cmd.alter_state(1,mol2,"(x,y,z)=stored.sel2.pop(0)")

        print "RMSD=%f" % RMSD

        # make the alignment OBVIOUS
        # comment out below to leave molecule representation unchanged
        cmd.orient(sel1 + " and " + sel2)

cmd.extend("optAlign", optAlign)

Advanced Scripts

Other Languages

Python is a great language, but sometimes we have libraries built in other languages, or Python's math is just too slow to be useful. So, we typically export our PyMOL data to the other language, do the math/problem, and import the changes back into PyMOL. This is shown below using the Python API and C. (This code comes from cealign.)

This is more advanced scripting, and requires some knowledge of the Python API.

Getting Your Data From Python/PyMOL into C

First, you need to let C know that you're interfacing with Python. To do that include the Python header file:

#include <Python.h>

Getting Your Data from C back into Python/PyMOL

Notes

I'll finish this soon. Tree

Example

See the source code for cealign.

See Also

stored, iterate_state, identify.