Difference between revisions of "Plane Wizard"

Revision as of 20:40, 26 December 2016

Introduction

This wizard has a simple purpose - to draw a cgo plane that passes through three points picked by the user. Most of the wizard itself was copied from the measure wizard.

To use, put it in the same directory as the other wizards. This is not quality code, and there may be bugs, but it seems to work okay.

Original

import pymol
from pymol import cmd
from pymol.wizard import Wizard
from chempy import cpv
from pymol.cgo import *

def makePrimitive(cgo, name):
    az = cmd.get('auto_zoom', quiet=1)
    cmd.set('auto_zoom', 0, quiet=1)
    cmd.load_cgo(cgo, name)
    cmd.set('auto_zoom', az, quiet=1)

def point(p):
    x, y, z = p
    return [COLOR, 1, 1, 1, SPHERE, float(x), float(y), float(z), 0.5]

def line(p1, p2):
    x1, y1, z1 = p1
    x2, y2, z2 = p2
    return [CYLINDER, float(x1), float(y1), float(z1), float(x2), float(y2), float(z2), 0.25, 1, 1, 1, 1, 1, 1]

def plane(corner1, corner2, corner3, corner4, normal):
    planeObj = []
    planeObj.extend(point(corner1))
    planeObj.extend(point(corner2))
    planeObj.extend(point(corner3))
    planeObj.extend(point(corner4))
    planeObj.extend(line(corner1, corner2))
    planeObj.extend(line(corner2, corner3))
    planeObj.extend(line(corner3, corner4))
    planeObj.extend(line(corner4, corner1))

    planeObj.extend([COLOR, 0.8, 0.8, 0.8])
    planeObj.extend([BEGIN, TRIANGLE_STRIP])
    planeObj.append(NORMAL)
    planeObj.extend(normal)
    for corner in [corner1, corner2, corner3, corner4, corner1]:
        planeObj.append(VERTEX)
        planeObj.extend(corner)
    planeObj.append(END)
    return planeObj

def planeFromPoints(point1, point2, point3, facetSize):
    v1 = cpv.normalize(cpv.sub(point2, point1))
    v2 = cpv.normalize(cpv.sub(point3, point1))
    normal = cpv.cross_product(v1, v2)
    v2 = cpv.cross_product(normal, v1)
    x = cpv.scale(v1, facetSize)
    y = cpv.scale(v2, facetSize)
    center = point2
    corner1 = cpv.add(cpv.add(center, x), y)
    corner2 = cpv.sub(cpv.add(center, x), y)
    corner3 = cpv.sub(cpv.sub(center, x), y)
    corner4 = cpv.add(cpv.sub(center, x), y)
    return plane(corner1, corner2, corner3, corner4, normal)


class PlaneWizard(Wizard):

    def __init__(self):
        Wizard.__init__(self)

        # some attributes to do with picking
        self.pick_count = 0
        self.object_count = 0
        self.object_prefix = "pw"

        # the plane facet size (the 'radius' of the section of plane we show)
        self.facetSize = 5

        self.selection_mode = cmd.get_setting_legacy("mouse_selection_mode")
        cmd.set("mouse_selection_mode",0) # set selection mode to atomic
        cmd.deselect()

    def reset(self):
        cmd.delete(self.object_prefix + "*")
        cmd.delete("sele*")
        cmd.delete("_indicate*")
        cmd.unpick()
        cmd.refresh_wizard()

    def delete_all(self):
        cmd.delete("plane*")

    def cleanup(self):
        cmd.set("mouse_selection_mode",self.selection_mode) # restore selection mode
        self.reset()
        self.delete_all()

    def get_prompt(self):
        self.prompt = None
        if self.pick_count == 0:
            self.prompt = [ 'Please click on the first atom...']
        elif self.pick_count == 1:
            self.prompt = [ 'Please click on the second atom...' ]
        elif self.pick_count == 2:
            self.prompt = [ 'Please click on the third atom...' ]
        return self.prompt

    def do_select(self, name):
        # "edit" only this atom, and not others with the object prefix
        try:
            cmd.edit("%s and not %s*" % (name, self.object_prefix))
            self.do_pick(0)
        except pymol.CmdException, pmce:
            print pmce

    def pickNextAtom(self, atom_name):
        # transfer the click selection to a named selection
        cmd.select(atom_name, "(pk1)")

        # delete the click selection
        cmd.unpick()

        # using the magic of indicate, highlight stuff
        indicate_selection = "_indicate" + self.object_prefix
        cmd.select(indicate_selection, atom_name)
        cmd.enable(indicate_selection)

        self.pick_count += 1
        self.error = None

        # necessary to force update of the prompt
        cmd.refresh_wizard()

    def do_pick(self, picked_bond):

        # this shouldn't actually happen if going through the "do_select"
        if picked_bond:
            self.error = "Error: please select bonds, not atoms"
            print self.error
            return

        atom_name = self.object_prefix + str(self.pick_count)
        if self.pick_count < 2:
            self.pickNextAtom(atom_name)
        else:
            self.pickNextAtom(atom_name)

            point1 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "0"))
            point2 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "1"))
            point3 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "2"))
            plane = planeFromPoints(point1, point2, point3, self.facetSize)

            planeName = "plane-%02d" % self.object_count
            self.object_count += 1
            makePrimitive(plane, planeName)
            cmd.show("cgo", "plane*")

            self.pick_count = 0
            self.reset()

    def get_panel(self):
        return [
            [ 1, 'Plane Wizard',''],
            [ 2, 'Reset','cmd.get_wizard().reset()'],
            [ 2, 'Delete All Planes' , 'cmd.get_wizard().delete_all()'],
            [ 2, 'Done','cmd.set_wizard()'],
        ]

# create an instance

wiz = PlaneWizard()

# make this the active wizard

cmd.set_wizard(wiz)

Modification

Small modifications to the same code as above.

Gallery

Drawing a plane in a protein
Drawing 6 planes in a box
Drawing 6 planes in a box

plane.py

Make a plane.py file in the same directory where you are working

'''
Described at PyMOL wiki:
https://pymolwiki.org/index.php/Plane_Wizard

Authors : Troels Schwarz-Linnet
Date    : Dec 2016
Modified: From previous contributors. 
'''

import pymol
from pymol import cmd
from pymol.wizard import Wizard
from chempy import cpv
from pymol.cgo import COLOR, SPHERE, CYLINDER, BEGIN, TRIANGLE_STRIP, NORMAL, VERTEX, END, ALPHA

def makePrimitive(cgo, name):
    az = cmd.get('auto_zoom', quiet=1)
    cmd.set('auto_zoom', 0, quiet=1)
    cmd.load_cgo(cgo, name)
    cmd.set('auto_zoom', az, quiet=1)

def point(p):
    x, y, z = p
    return [COLOR, 1, 1, 1, SPHERE, float(x), float(y), float(z), 0.5]


def line(p1, p2):
    x1, y1, z1 = p1
    x2, y2, z2 = p2
    return [CYLINDER, float(x1), float(y1), float(z1), float(x2), float(y2), float(z2), 0.25, 1, 1, 1, 1, 1, 1]


def plane(corner1, corner2, corner3, corner4, normal, settings):
    planeObj = []
    planeObj.extend(point(corner1))
    planeObj.extend(point(corner2))
    planeObj.extend(point(corner3))
    planeObj.extend(point(corner4))
    planeObj.extend(line(corner1, corner2))
    planeObj.extend(line(corner2, corner3))
    planeObj.extend(line(corner3, corner4))
    planeObj.extend(line(corner4, corner1))

    # Make settings
    if 'ALPHA' in settings:
        planeObj.extend([ALPHA, settings['ALPHA']])

    if 'COLOR' in settings:
        planeObj.extend([COLOR, settings['COLOR'][0], settings['COLOR'][1], settings['COLOR'][2]])
    else:
        planeObj.extend([COLOR, 0.8, 0.8, 0.8]) # greyish

    planeObj.extend([BEGIN, TRIANGLE_STRIP])
    planeObj.append(NORMAL)

    if 'INVERT' in settings:
        if settings['INVERT']==True:
            planeObj.extend(cpv.negate(normal))
        else:
            planeObj.extend(normal)
    else:
        planeObj.extend(normal)


    for corner in [corner1, corner2, corner3, corner4, corner1]:
        planeObj.append(VERTEX)
        planeObj.extend(corner)
    planeObj.append(END)
    return planeObj


def planeFromPoints(p1, p2, p3, vm1=None, vm2=None, center=True, settings={}):
    v1 = cpv.sub(p1, p2)
    v2 = cpv.sub(p3, p2)
    normal = cpv.cross_product(v1, v2)

    if vm1!=None:
        v1 = cpv.scale(cpv.normalize(v1), vm1)
    if vm2!=None:
        v2 = cpv.scale(cpv.normalize(v2), vm2)

    centrum = p2
    if center:
        corner1 = cpv.add(cpv.add(centrum, v1), v2)
        corner2 = cpv.sub(cpv.add(centrum, v1), v2)
        corner3 = cpv.sub(cpv.sub(centrum, v1), v2)
        corner4 = cpv.add(cpv.sub(centrum, v1), v2)
    else:
        corner1 = cpv.add(cpv.add(centrum, v1), v2)
        corner2 = cpv.add(centrum, v1)
        corner3 = centrum
        corner4 = cpv.add(centrum, v2)

    return plane(corner1, corner2, corner3, corner4, normal, settings)


def print_info(name, coor1, coor2, coor3):
    cs1 = (map(float, [ '%.2f' % elem for elem in coor1 ]) )
    cs2 = (map(float, [ '%.2f' % elem for elem in coor2 ]) )
    cs3 = (map(float, [ '%.2f' % elem for elem in coor3 ]) )
    print("You can also use the function calls with these coordinates")
    print("make_plane_points(name='%s', l1=%s, l2=%s, l3=%s)"%(name, cs1, cs2, cs3))


def make_plane(name,a1='(pk1)',a2='(pk2)',a3='(pk3)', vm1=None, vm2=None, center=True, makepseudo=True, settings={}):
    """
    DESCRIPTION
    Create a CGO plane from three atomic coordinates

    USAGE
    make_plane name, a1, a2, a3

    where each atom is a standard PyMOL selection (defaults to pk1,pk2 and pk3)
    """
    # get coordinates from atom selections
    coor1 = cmd.get_model(a1).get_coord_list()[0]
    coor2 = cmd.get_model(a2).get_coord_list()[0]
    coor3 = cmd.get_model(a3).get_coord_list()[0]

    # Help with alternative
    print_info(name, coor1, coor2, coor3)

    # Get the plane
    plane = planeFromPoints(p1=coor1, p2=coor2, p3=coor3, vm1=vm1, vm2=vm2, center=center, settings=settings)
    makePrimitive(plane, name)
    #cmd.show("cgo", "plane*")

    if makepseudo:
        cmd.pseudoatom("%s_%s"%(name, "l1"), color="tv_blue", pos=coor1)
        cmd.pseudoatom("%s_%s"%(name, "l2"), color="tv_green", pos=coor2)
        cmd.pseudoatom("%s_%s"%(name, "l3"), color="red", pos=coor3)

# Extend function to be called inside pymol
cmd.extend("make_plane", make_plane)

def make_plane_points(name,l1=None,l2=None,l3=None, vm1=None, vm2=None, center=True, makepseudo=True, settings={}):
    """
    DESCRIPTION
    Create a CGO plane from three atomic coordinates

    USAGE
    make_plane name, l1, l2, l3

    where each xys is a list with floats of x,y,z coordinates
    """
    if l1==None or l2==None or l3==None:
        print("Please provide a list of xyz floats for each 3 positions")
        return
    if type(l1) is not list or type(l2) is not list or type(l3) is not list:
        print(type(l1),type(l2),type(l3))
        print("Please provide 3 list of xyz floats for each 3 positions")
        return

    plane = planeFromPoints(p1=l1, p2=l2, p3=l3, vm1=vm1, vm2=vm2, center=center, settings=settings)
    makePrimitive(plane, name)

    if makepseudo:
        cmd.pseudoatom("%s_%s"%(name, "l1"), color="tv_blue", pos=l1)
        cmd.pseudoatom("%s_%s"%(name, "l2"), color="tv_green", pos=l2)
        cmd.pseudoatom("%s_%s"%(name, "l3"), color="red", pos=l3)

# Extend function to be called inside pymol
cmd.extend("make_plane_points", make_plane_points)

class PlaneWizard(Wizard):
    def __init__(self):
        Wizard.__init__(self)

        # some attributes to do with picking
        self.pick_count = 0
        self.object_count = 0
        self.object_prefix = "pw"

        self.selection_mode = cmd.get_setting_legacy("mouse_selection_mode")
        cmd.set("mouse_selection_mode",0) # set selection mode to atomic
        cmd.deselect()

    def reset(self):
        cmd.delete(self.object_prefix + "*")
        cmd.delete("sele*")
        cmd.delete("_indicate*")
        cmd.unpick()
        cmd.refresh_wizard()

    def delete_all(self):
        cmd.delete("plane*")

    def cleanup(self):
        cmd.set("mouse_selection_mode",self.selection_mode) # restore selection mode
        self.reset()
        self.delete_all()

    def get_prompt(self):
        self.prompt = None
        if self.pick_count == 0:
            self.prompt = [ 'Please click on the first atom...']
        elif self.pick_count == 1:
            self.prompt = [ 'Please click on the second atom...' ]
        elif self.pick_count == 2:
            self.prompt = [ 'Please click on the third atom...' ]
        return self.prompt

    def do_select(self, name):
        # "edit" only this atom, and not others with the object prefix
        try:
            cmd.edit("%s and not %s*" % (name, self.object_prefix))
            self.do_pick(0)
        except pymol.CmdException, pmce:
            print pmce

    def pickNextAtom(self, atom_name):
        # transfer the click selection to a named selection
        cmd.select(atom_name, "(pk1)")

        # delete the click selection
        cmd.unpick()

        # using the magic of indicate, highlight stuff
        indicate_selection = "_indicate" + self.object_prefix
        cmd.select(indicate_selection, atom_name)
        cmd.enable(indicate_selection)

        self.pick_count += 1
        self.error = None

        # necessary to force update of the prompt
        cmd.refresh_wizard()

    def do_pick(self, picked_bond):

        # this shouldn't actually happen if going through the "do_select"
        if picked_bond:
            self.error = "Error: please select bonds, not atoms"
            print self.error
            return

        atom_name = self.object_prefix + str(self.pick_count)
        if self.pick_count < 2:
            self.pickNextAtom(atom_name)
        else:
            self.pickNextAtom(atom_name)

            point1 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "0"))
            point2 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "1"))
            point3 = cmd.get_atom_coords("(%s%s)" % (self.object_prefix, "2"))
            plane = planeFromPoints(point1, point2, point3)

            planeName = "plane-%02d" % self.object_count

            print_info(planeName, point1, point2, point3)

            self.object_count += 1
            makePrimitive(plane, planeName)
            cmd.show("cgo", "plane*")

            self.pick_count = 0
            self.reset()

    def get_panel(self):
        return [
            [ 1, 'Plane Wizard',''],
            [ 2, 'Reset','cmd.get_wizard().reset()'],
            [ 2, 'Delete All Planes' , 'cmd.get_wizard().delete_all()'],
            [ 2, 'Done','cmd.set_wizard()'],
        ]

Examples

Plane in a protein

reinitialize

fetch 1ubq, async=0
show_as cartoon, all 

import plane
# Start the wizard, and do manual picking
cmd.set_wizard(plane.PlaneWizard())

Or by selection

reinitialize

fetch 1ubq, async=0
show_as cartoon, all 

import plane

# Set alpha level and color
#violetpurple: 0.55, 0.25, 0.60 #yellow: 1.0, 1.0, 0.0, #blue: 0.0, 0.0, 1.0 #orange: 1.0, 0.5, 0.0 #forest: 0.2, 0.6, 0.2 #red: 1.0, 0.0, 0.0
dict = {'ALPHA':0.6, 'COLOR':[0.55, 0.25, 0.60], 'INVERT':True}
plane.make_plane(name='test', a1='/1ubq//A/24/CA', a2='/1ubq//A/29/CA', a3='/1ubq//A/40/CA', center=False, settings=dict)

Or by atom coordinates

reinitialize

fetch 1ubq, async=0
show_as cartoon, all 

import plane

# Or  make from atom coordinates
#plane.make_plane_points(name='test', l1=[35.03, 21.72, 17.07], l2=[37.47, 27.39, 10.67], l3=[37.74, 31.64, 23.71])

# Define plane, 10 angstrom in length
#plane.make_plane_points(name='p1', l1=[0.0, 10.0, 0.0], l2=[0.0, 0.0, 0.0], l3=[0.0, 0.0, 10.0], center=False, makepseudo=False)

Or make a color cube

reinitialize
#violetpurple: 0.55, 0.25, 0.60 #yellow: 1.0, 1.0, 0.0, #blue: 0.0, 0.0, 1.0 #orange: 1.0, 0.5, 0.0 #forest: 0.2, 0.6, 0.2 #red: 1.0, 0.0, 0.0

# YZ Plane, #purple
dict = {'ALPHA':0.4, 'COLOR':[0.55, 0.25, 0.60], 'INVERT':True}
plane.make_plane_points(name='p1', l1=[0.0, 10.0, 0.0], l2=[0.0, 0.0, 0.0], l3=[0.0, 0.0, 10.0], center=False, makepseudo=False, settings=dict)
# YZ Plane, shifted in X, #yellow
dict = {'ALPHA':0.4, 'COLOR':[1.0, 1.0, 0.0]}
plane.make_plane_points(name='p6', l1=[10.0, 10.0, 0.0], l2=[10.0, 0.0, 0.0], l3=[10.0, 0.0, 10.0], center=False, makepseudo=False, settings=dict)

# XZ Plane, blue
dict = {'ALPHA':0.4, 'COLOR':[0.0, 0.0, 1.0]}
plane.make_plane_points(name='p2', l1=[10.0, 0.0, 0.0], l2=[0.0, 0.0, 0.0], l3=[0.0, 0.0, 10.0], center=False, makepseudo=False, settings=dict)
# XZ Plane, shifted in Y, #orange
dict = {'ALPHA':0.4, 'COLOR':[1.0, 0.5, 0.0], 'INVERT':True}
plane.make_plane_points(name='p5', l1=[10.0, 10.0, 0.0], l2=[0.0, 10.0, 0.0], l3=[0.0, 10.0, 10.0], center=False, makepseudo=False, settings=dict)

# XY Plane, forest
dict = {'ALPHA':0.4, 'COLOR':[0.2, 0.6, 0.2], 'INVERT':True}
plane.make_plane_points(name='p4', l1=[10.0, 0.0, 0.0], l2=[0.0, 0.0, 0.0], l3=[0.0, 10.0, 0.0], center=False, makepseudo=False, settings=dict)
# XY Plane, shifted in Z, red
dict = {'ALPHA':0.4, 'COLOR':[1.0, 0.0, 0.0]}
plane.make_plane_points(name='p3', l1=[10.0, 0.0, 10.0], l2=[0.0, 0.0, 10.0], l3=[0.0, 10.0, 10.0], center=False, makepseudo=False, settings=dict)

zoom all