1#A* -------------------------------------------------------------------
2#B* This file contains source code for the PyMOL computer program
3#C* Copyright (c) Schrodinger, LLC.
4#D* -------------------------------------------------------------------
5#E* It is unlawful to modify or remove this copyright notice.
6#F* -------------------------------------------------------------------
7#G* Please see the accompanying LICENSE file for further information.
8#H* -------------------------------------------------------------------
9#I* Additional authors of this source file include:
10#-*
11#-*
12#-*
13#Z* -------------------------------------------------------------------
14
15cmd = __import__("sys").modules["pymol.cmd"]
16from . import util
17import traceback
18
19polar_contacts_suffix = "_pol_conts"
20default_polar_contacts = "polar_contacts"
21
22tmp_sele = "_p_tmp"
23
24prot_and_dna_sele = "(resn ALA+CYS+CYX+ASP+GLU+PHE+GLY+HIS+HID+HIE+HIP+HISE+HISD+HISP+ILE+LYS+LEU+MET+MSE+ASN+PRO+GLN+ARG+SER+THR+VAL+TRP+TYR+A+C+T+G+U+DA+DC+DT+DG+DU+DI)"
25wat_sele = "solvent"
26ion_sele = "(resn CA,HG,K,NA,ZN,MG,CL)"
27solv_sele = "("+wat_sele+"|"+ion_sele+")"
28lig_excl = "(resn MSE)"
29lig_sele = "((hetatm or not "+prot_and_dna_sele+") and not ("+solv_sele+"|"+ion_sele+"|"+lig_excl+"))"
30lig_and_solv_sele = "("+lig_sele+"|"+solv_sele+")"
31
32def get_sname_oname_dname(selection, _self=cmd):
33    '''
34    Get a named selection, object names and a distance (polar contacts) name
35    '''
36    _self.select(tmp_sele, selection)
37
38    if selection not in _self.get_object_list():
39        selection = ' '.join(_self.get_object_list(tmp_sele))
40
41    if selection and ' ' not in selection:
42        dname = selection + polar_contacts_suffix
43    else:
44        dname = default_polar_contacts
45
46    return tmp_sele, selection, dname
47
48def _prepare(selection,polar_contacts=None,_self=cmd):
49    cmd=_self
50    # this function should undo everything that is done by any preset function in this module
51    # (except for coloring)
52
53    s, selection, dname = get_sname_oname_dname(selection, _self=cmd)
54
55    cmd.cartoon("auto",s)
56    cmd.hide("everything",s)
57
58    cmd.set("two_sided_lighting",0) # global
59    cmd.unset("transparency",s)
60    cmd.unset("surface_quality", selection)
61    cmd.unset("surface_type",selection)
62    cmd.unset("sphere_scale",s)
63    cmd.unset_bond("stick_radius",s,s)
64    cmd.unset_bond("stick_color",s,s)
65    cmd.unset("cartoon_highlight_color",selection)
66    cmd.unset("cartoon_fancy_helices",selection)
67    cmd.unset("cartoon_smooth_loops",selection)
68    cmd.unset("cartoon_flat_sheets",selection)
69    cmd.unset("cartoon_side_chain_helper",selection)
70    if polar_contacts is None:
71        polar_contacts = dname
72        if polar_contacts in cmd.get_names('objects'):
73            cmd.delete(polar_contacts)
74
75    return s, selection, dname
76
77def simple(selection="(all)",_self=cmd):
78    cmd=_self
79    s, selection = _prepare(selection, _self=cmd)[:2]
80    util.cbc(s,_self=cmd)
81    cmd.show("ribbon",s)
82    cmd.show("lines","(byres (("+s+" & r. CYS+CYX & n. SG) & bound_to ("+s+" & r. CYS+CYX & n. SG))) & n. CA+CB+SG")
83    # try to show what covalent ligands are connected to...
84    cmd.show("sticks","("+lig_sele+" and ("+s+")) extend 2")
85    cmd.show("sticks","byres (("+lig_sele+" and ("+s+") and not resn ACE+NAC+NME+NH2) extend 1)")
86    cmd.hide("sticks","("+s+") and ((not rep sticks) extend 1)")
87    cmd.show("sticks","("+lig_sele+" and ("+s+")) extend 2")
88    # color by atom if lines or sticks are shown
89    util.cnc("(( rep lines or rep sticks or ("+lig_and_solv_sele+")) and ("+s+"))",_self=cmd)
90    cmd.show("nonbonded","("+lig_and_solv_sele+" and ("+s+"))")
91    cmd.show("lines","("+lig_and_solv_sele+" and ("+s+"))")
92    cmd.delete(s)
93
94def simple_no_solv(selection="(all)",_self=cmd):
95    cmd=_self
96    simple(selection,_self=_self)
97    s, selection = get_sname_oname_dname(selection, _self=_self)[:2]
98    cmd.hide("nonbonded","("+solv_sele+" and "+s+")")
99    cmd.delete(s)
100
101def ligands(selection="(all)",_self=cmd):
102    cmd=_self
103    try:
104        s, selection, polar_contacts = _prepare(selection, _self=cmd)
105        host = "_preset_host"
106        solvent = "_preset_solvent"
107        near_solvent = "_preset_solvent"
108        lig = "_preset_lig"
109        cmd.select(host,s+" and "+prot_and_dna_sele)
110        cmd.select(solvent,s+" and "+solv_sele)
111        cmd.select(lig,s+" and "+lig_sele)
112        cmd.select(near_solvent,s+" and ("+solvent+" within 4 of "+lig+")")
113
114        util.chainbow(host,_self=cmd)
115        util.cbc(lig,_self=cmd)
116        util.cbac("(("+s+") and not elem C)",_self=cmd)
117        cmd.hide("everything",s)
118        cmd.show("ribbon",host)
119        cmd.show("lines","("+s+" and byres ("+host+" within 5 of "+lig+"))")
120        cmd.show("sticks",lig)
121        cmd.show("sticks",solvent+" and neighbor "+lig)
122        cmd.show("lines","("+s+" and (rep lines extend 1) and "+lig+")")
123
124        if cmd.count_atoms(lig):
125            cmd.dist(polar_contacts,host+"|"+near_solvent,lig+"|"+near_solvent,
126                     mode=2,quiet=1,label=0,reset=1) # hbonds
127            if polar_contacts in cmd.get_names():
128                cmd.enable(polar_contacts)
129                cmd.hide("labels",polar_contacts)
130                cmd.show("dashes",polar_contacts)
131        else:
132            cmd.delete(polar_contacts)
133        cmd.show("nonbonded",lig+"|"+host+"|"+near_solvent)
134        if cmd.count_atoms(lig):
135            cmd.zoom(lig,3, animate=1)
136        cmd.delete(host)
137        cmd.delete(solvent)
138        cmd.delete(near_solvent)
139        cmd.delete(lig)
140    except:
141        traceback.print_exc()
142    cmd.delete(s)
143
144def ball_and_stick(selection="(all)",mode=1,_self=cmd):
145    cmd=_self
146    s, selection = _prepare(selection, _self=cmd)[:2]
147    if mode == 1:
148        cmd.hide("everything",s)
149        cmd.set_bond("stick_color","white",s,s)
150        cmd.set_bond("stick_radius","0.14",s,s)
151        cmd.set("sphere_scale","0.25",s)
152        cmd.show("sticks",s)
153        cmd.show("spheres",s)
154    elif mode == 2:
155        cmd.hide("everything",s)
156        cmd.set_bond("stick_color","white",s,s)
157        cmd.set_bond("stick_radius","-0.14",s,s)
158        cmd.set("stick_ball","1")
159        cmd.set("stick_ball_ratio",-1.0)
160        cmd.set("stick_ball_color","atomic")
161        cmd.show("sticks",s)
162    cmd.delete(s)
163
164def b_factor_putty(selection="(name CA+P)",_self=cmd):
165    cmd=_self
166    s, selection = _prepare(selection, _self=cmd)[:2]
167    cmd.select(s,"(name CA+P) and ("+selection+")")
168    cmd.show("cartoon",s)
169    cmd.set("cartoon_flat_sheets",0,selection)
170    cmd.cartoon("putty",s)
171    cmd.spectrum("b",selection=s)
172    cmd.delete(s)
173
174def ligand_cartoon(selection="(all)",_self=cmd):
175    cmd=_self
176    s, selection = ligand_sites(selection, _self)[:2]
177    cmd.set("cartoon_side_chain_helper",1,selection)
178    cmd.show("cartoon","rep ribbon")
179    cmd.hide("ribbon")
180    cmd.hide("surface")
181    cmd.delete(s)
182
183def ligand_sites(selection="(all)",_self=cmd):
184    cmd=_self
185    try:
186        s, selection, polar_contacts = _prepare(selection, _self=cmd)
187        host = "_preset_host"
188        solvent = "_preset_solvent"
189        near_solvent = "_preset_solvent"
190        lig = "_preset_lig"
191        cmd.select(host,s+" and "+prot_and_dna_sele)
192        cmd.select(solvent,s+" and "+solv_sele)
193        cmd.select(lig,s+" and "+lig_sele)
194        cmd.select(near_solvent,s+" and ("+solvent+" within 4 of "+lig+")")
195        cmd.flag("ignore",host,"clear")
196        cmd.flag("ignore",lig+"|"+solvent,"set")
197
198        util.chainbow(host,_self=cmd)
199        util.cbc(lig,_self=cmd)
200        util.cbac("(("+s+") and not elem C)",_self=cmd)
201        cmd.hide("everything",s)
202        cmd.show("ribbon",host)
203        cmd.show("lines","("+s+" and byres ("+host+" within 5 of "+lig+"))")
204        cmd.show("surface","("+s+" and ((rep lines expand 4) within 6 of "+lig+"))")
205        cmd.set("two_sided_lighting",1) # global setting
206        cmd.set("transparency",0,s)
207        cmd.set("surface_quality",0, selection)
208
209        cmd.show("sticks",lig)
210        cmd.show("sticks",solvent+" and neighbor "+lig)
211        cmd.show("lines","("+s+" and (rep lines extend 1) and "+lig+")")
212
213        if cmd.count_atoms(lig):
214            cmd.dist(polar_contacts,host+"|"+near_solvent,lig+"|"+near_solvent,mode=2,quiet=1,label=0,reset=1) # hbonds
215            if polar_contacts in cmd.get_names():
216                cmd.enable(polar_contacts)
217                cmd.hide("labels",polar_contacts)
218                cmd.show("dashes",polar_contacts)
219        else:
220            cmd.delete(polar_contacts)
221
222        cmd.show("nb_spheres",lig+"|"+host+"|"+near_solvent)
223        if cmd.count_atoms(lig):
224            cmd.zoom(lig,3, animate=1)
225        cmd.delete(host)
226        cmd.delete(solvent)
227        cmd.delete(near_solvent)
228        cmd.delete(lig)
229    finally:
230        pass
231
232    return s, selection, polar_contacts
233
234def ligand_sites_hq(selection="(all)",_self=cmd):
235    cmd=_self
236    s, selection = ligand_sites(selection, _self)[:2]
237    cmd.set("surface_quality","1",selection)
238    cmd.set("surface_type",0,selection)
239    cmd.delete(s)
240
241def ligand_sites_trans(selection="(all)",_self=cmd):
242    cmd=_self
243    s, selection = ligand_sites(selection, _self)[:2]
244    cmd.show("sticks",s+" and rep lines")
245    cmd.hide("lines",s+" and rep lines")
246    cmd.set("transparency","0.33",s)
247    cmd.set("surface_type",0,selection)
248    cmd.set("surface_quality",0,selection)
249    cmd.delete(s)
250
251def ligand_sites_trans_hq(selection="(all)",_self=cmd):
252    cmd=_self
253    s, selection = ligand_sites(selection, _self)[:2]
254    cmd.show("sticks",s+" and rep lines")
255    cmd.hide("lines",s+" and rep lines")
256    cmd.set("transparency","0.33",s)
257    cmd.set("surface_type",0,selection)
258    cmd.set("surface_quality",1,selection)
259    cmd.delete(s)
260
261def ligand_sites_mesh(selection="(all)",_self=cmd):
262    cmd=_self
263    s, selection = ligand_sites(selection, _self)[:2]
264    cmd.show("sticks",s+" and rep lines")
265    cmd.hide("lines",s+" and rep lines")
266    cmd.set("surface_type","2",selection)
267    cmd.set("surface_quality","0",selection)
268    cmd.delete(s)
269
270def ligand_sites_dots(selection="(all)",_self=cmd):
271    cmd=_self
272    s, selection = ligand_sites(selection, _self)[:2]
273    cmd.show("sticks",s+" and rep lines")
274    cmd.hide("lines",s+" and rep lines")
275    cmd.set("surface_type","1",selection)
276    cmd.set("surface_quality","1",selection)
277    cmd.delete(s)
278
279def technical(selection="(all)",_self=cmd):
280    cmd=_self
281    s, selection, polar_contacts = _prepare(selection, _self=cmd)
282    util.chainbow(s,_self=cmd)
283    util.cbc("("+lig_sele+" and ("+s+"))",_self=cmd)
284    util.cbac("(("+s+") and not elem C)",_self=cmd)
285    cmd.show("nonbonded",s)
286    cmd.show("lines","((("+s+") and not "+lig_sele+") extend 1)")
287    cmd.show("sticks","("+lig_sele+" and ("+s+"))")
288    cmd.show("ribbon",s)
289    cmd.dist(polar_contacts,s,s,mode=2,label=0,reset=1) # hbonds
290    if polar_contacts in cmd.get_names():
291        cmd.enable(polar_contacts)
292        cmd.set("dash_width",1.5,polar_contacts)
293        cmd.hide("labels",polar_contacts)
294        cmd.show("dashes",polar_contacts)
295    cmd.show("nonbonded","(("+lig_sele+"|resn HOH+WAT+H2O) and ("+s+"))")
296    cmd.delete(s)
297
298def pretty_solv(selection="(all)",_self=cmd):
299    cmd=_self
300    s, selection = _prepare(selection, _self=cmd)[:2]
301    cmd.dss(s,preserve=1)
302    cmd.cartoon("auto",s)
303    cmd.show("cartoon",s)
304    cmd.show("sticks","("+lig_sele+" and ("+s+"))")
305    cmd.show("nb_spheres","(("+lig_sele+"|resn HOH+WAT+H2O) and ("+s+"))")
306    util.cbc("("+lig_sele+" and ("+s+"))",_self=cmd)
307    util.cbac("("+lig_sele+" and ("+s+") and not elem C)",_self=cmd)
308    cmd.spectrum("count",selection="(elem C and ("+s+") and not "+lig_sele+")")
309    cmd.set("cartoon_highlight_color",-1,selection)
310    cmd.set("cartoon_fancy_helices",0,selection)
311    cmd.set("cartoon_smooth_loops",0,selection)
312    cmd.set("cartoon_flat_sheets",1,selection)
313    cmd.set("cartoon_side_chain_helper",0,selection)
314    cmd.delete(s)
315
316def pretty(selection,_self=cmd):
317    cmd=_self
318    pretty_solv(selection,_self)
319    s, selection = get_sname_oname_dname(selection, _self=_self)[:2]
320    cmd.hide("nb_spheres","("+s+" and "+lig_sele+"|resn HOH+WAT+H2O)")
321    cmd.delete(s)
322
323pretty_no_solv = pretty
324
325def pub_solv(selection="(all)",_self=cmd):
326    cmd=_self
327    pretty_solv(selection,_self)
328    s, selection = get_sname_oname_dname(selection, _self=_self)[:2]
329    cmd.set("cartoon_smooth_loops",1,selection)
330    cmd.set("cartoon_highlight_color","grey50",selection)
331    cmd.set("cartoon_fancy_helices",1,selection)
332    cmd.set("cartoon_flat_sheets",1,selection)
333    cmd.set("cartoon_side_chain_helper",0,selection)
334    cmd.delete(s)
335
336def publication(selection="(all)",_self=cmd):
337    cmd=_self
338    pub_solv(selection,_self)
339    s, selection = get_sname_oname_dname(selection, _self=_self)[:2]
340    cmd.hide("nb_spheres","(("+lig_sele+"|resn HOH+WAT+H2O) and "+s+")")
341    cmd.delete(s)
342
343pub_no_solv = publication
344
345def default(selection="(all)",_self=cmd):
346    cmd=_self
347    s, selection = _prepare(selection, _self=cmd)[:2]
348    cmd.show("lines",s)
349    cmd.show("nonbonded",s)
350    color=cmd.get_object_color_index(selection)
351    if color<0:
352        util.cbag(selection,_self=cmd)
353    else:
354        util.cnc(selection,_self=cmd)
355        cmd.color(str(color),"("+s+") and elem C")
356    cmd.delete(s)
357
358
359def interface(selection='*', _self=cmd):
360    '''
361    Protein-Protein interface preset, mimics the BioLuminate preset
362    '''
363    s = _prepare(selection, _self=_self)[0]
364
365    # temporary selection names
366    s_interface = _self.get_unused_name('_iface')
367
368    # interface atoms
369    _self.select(s_interface, '?%s & (%s)' % (s, ' '.join(
370        '((chain %s) around 4.5) ' % (chain)
371        for chain in _self.get_chains(s))), 0)
372
373    # Color by chain, non-carbon by element
374    util.cbc(s, _self=_self)
375    _self.color('atomic', '?%s & !(elem C)' % (s))
376
377    # Change everything to cartoons
378    _self.show_as('cartoon', s)
379
380    # interface residues as sticks
381    _self.show('sticks', 'byres ?' + s_interface)
382    _self.show('nb_spheres', '?' + s_interface)
383
384    # delete temporary selections
385    _self.delete(s_interface)
386    _self.delete(s)
387
388
389def classified(selection='*', _self=cmd):
390    '''
391    Equivalent of "auto_show_classified" setting. Sets representations
392    according to atom classification ("auto_classify_atoms"). Does not
393    change any colors or settings.
394    '''
395    s = _prepare(selection, _self=_self)[0]
396
397    _self.show_as('cartoon', 'polymer & %' + s)
398    _self.show_as('sticks', 'organic & %' + s)
399    _self.show_as('spheres', 'inorganic & %' + s)
400
401    _self.delete(s)
402