1# TODO
2# * default settings for color and rep
3# * make the final viewing step a function
4# *   setup_map(name,levels=,colors=,reps=)
5
6from __future__ import print_function
7
8import pymol
9from pymol import headering
10import Pmw
11
12class PyMOLMapLoad:
13
14    def __init__(self,parent,app,f):
15        self._parent   = parent
16        self._app      = app
17        self._fileName = f
18        self._fileData = None
19
20        # model state
21        self._amplitudes = None
22        self._phases     = None
23        self._weights    = None
24        self._min_res    = None
25        self._max_res    = None
26        self._fofc       = None
27        self._name_prefix= None
28
29        # reflection file header data
30        if f[-3:] in ("MTZ", "mtz"):
31            self._fileData = headering.MTZHeader(f)
32        elif f[-3:] in ("CIF", "cif"):
33            self._fileData = headering.CIFHeader(f)
34        elif f[-3:] in ("CNS", "cns", "hkl", "HKL"):
35            self._fileData = headering.CNSHeader(f)
36
37    def pack_and_show(self):
38        self.pack()
39        return self.show()
40
41    def pack(self):
42        # MAIN DIALOG
43        self._d = Pmw.Dialog(self._parent,
44                             buttons = ("OK", "Cancel", "Help"),
45                             defaultbutton = "OK",
46                             title = "PyMOL Map Generation",
47                             command = self.run)
48        self._d.geometry("+%d+%d" % (self._app.winfo_reqwidth(), self._app.winfo_reqheight()))
49        self._d.withdraw()
50        self._d.protocol('WM_DELETE_WINDOW', self.quit)
51
52        #
53        # COLUMN LABEL GROUP
54        #
55        self._col_gp = Pmw.Group(self._d.interior(),
56                                 tag_text="Column Labels",)
57        self._col_gp.pack(fill='x', expand='yes')
58
59        defaultListHeight = 125
60        FCols = []
61        FCols.extend(self._fileData.getColumnsOfType("F"))
62        FCols.extend(self._fileData.getColumnsOfType("G"))
63        if not len(FCols): FCols = [ "" ]
64        self._ampl_chooser = Pmw.ComboBox(self._col_gp.interior(),
65                                          label_text = "Amplitudes",
66                                          labelpos = "nw",
67                                          selectioncommand = self.set_amplitudes,
68                                          scrolledlist_items = FCols,
69                                          dropdown = 1,
70                                          listheight=defaultListHeight,
71                                          sticky='ew')
72        self._ampl_chooser.pack(fill='both',expand=1,padx=7,pady=4)
73        _FC, _PC, _looksLike = self._fileData.guessCols("FoFc")
74        _2FC, _2PC, _looksLike = self._fileData.guessCols("2FoFc")
75        # be nice and choose the most appropriate col
76        if _2FC!=None:
77            if _2FC in FCols:
78                self._ampl_chooser.selectitem(_2FC)
79        elif _FC!=None:
80            if _FC in FCols:
81                self._ampl_chooser.selectitem(_FC)
82        else:
83            self._ampl_chooser.selectitem(FCols[0])
84
85        PCols = []
86        PCols.extend(self._fileData.getColumnsOfType("P"))
87        if not len(PCols): PCols = [ "" ]
88        self._phase_chooser = Pmw.ComboBox(self._col_gp.interior(),
89                                          label_text = "Phases",
90                                          labelpos = "nw",
91                                          selectioncommand = self.set_phases,
92                                          scrolledlist_items = PCols,
93                                          dropdown = 1,
94                                          listheight=defaultListHeight)
95        self._phase_chooser.pack(fill='both', expand=1,padx=7,pady=4)
96        # be nice and choose the most appropriate col
97        if _2PC!=None:
98            if _2PC in PCols:
99                self._phase_chooser.selectitem(PCols.index(_2PC))
100        elif _PC!=None:
101            if _PC in PCols:
102                self._phase_chooser.selectitem(PCols.index(_PC))
103        else:
104            self._phase_chooser.selectitem(PCols[0])
105
106        WCols = [ "None", ]
107        WCols.extend(self._fileData.getColumnsOfType("W"))
108        WCols.extend(self._fileData.getColumnsOfType("Q"))
109        self._wt_chooser = Pmw.ComboBox(self._col_gp.interior(),
110                                          label_text = "Weights",
111                                          labelpos = "nw",
112                                          selectioncommand = self.set_weights,
113                                          scrolledlist_items = WCols,
114                                          dropdown = 1,
115                                          listheight=defaultListHeight)
116        self._wt_chooser.pack(fill='both', expand=1,padx=7,pady=4)
117        self._wt_chooser.selectitem("None")
118
119        #
120        # INPUT OPTIONS GROUP
121        #
122        self._input_gp = Pmw.Group(self._d.interior(),
123                                   tag_text="Input Options",)
124        self._input_gp.pack(fill='both', expand='yes')
125
126        if self._fileData.reso_min!=None:
127            default_min_res = float("%3.5f"%float(self._fileData.reso_min))
128        else:
129            default_min_res = ""
130        if self._fileData.reso_max!=None:
131            default_max_res = float("%3.5f"%float(self._fileData.reso_max))
132        else:
133            default_max_res = ""
134        self._min_res_fld = Pmw.EntryField(self._input_gp.interior(),
135                                           labelpos="wn",
136                                           label_text="Min. Resolution",
137                                           value = default_min_res,
138                                           validate = { "validator" : 'real' },
139                                           entry_width=7,
140                                           modifiedcommand=self.set_min_res,
141                                           command = self.set_min_res)
142        self._min_res_fld.grid(row=1,column=0,rowspan=2,sticky='ew',pady=4)
143
144        self._max_res_fld = Pmw.EntryField(self._input_gp.interior(),
145                                           labelpos="wn",
146                                           label_text = "Max Resolution",
147                                           value = default_max_res,
148                                           validate = { "validator" : 'real' },
149                                           entry_width=7,
150                                           modifiedcommand=self.set_max_res,
151                                           command = self.set_max_res)
152        self._max_res_fld.grid(row=1,column=1,rowspan=2,sticky='ew',pady=4)
153
154
155        #
156        # MAP OPTIONS GROUP
157        #
158        self._options_gp = Pmw.Group(self._d.interior(),
159                                     tag_text="Map Options",)
160        self._options_gp.pack(fill='x', expand='yes')
161
162        self._name_prefix_fld = Pmw.EntryField(self._options_gp.interior(),
163                                               labelpos="wn",
164                                               label_text = "New Map Name Prefix",
165                                               value = "",
166                                               validate = { "validator" : 'alphanumeric' },
167                                               entry_width=20,
168                                               modifiedcommand=self.set_name_prefix,
169                                               command = self.set_name_prefix)
170        self._name_prefix_fld.pack(fill="x", expand=0, anchor='w')
171
172        self._fofc_chooser = Pmw.RadioSelect(self._options_gp.interior(),
173                                             command = self.set_fofc,
174                                             buttontype="checkbutton",)
175        self._fofc_chooser.add("FoFc")
176        self._fofc_chooser.pack(fill="none", expand=0, anchor="w")
177
178    def show(self):
179        self._d.show()
180    def quit(self):
181        if __name__=="__main__":
182            # TODO--remove me; use for development only!
183            self._parent.destroy()
184        else:
185            # TODO -- use only this in release
186            self._d.destroy()
187
188
189
190    # UI SETTERS
191    def set_amplitudes(self,arg):
192        self._amplitudes = arg
193    def set_phases(self,arg):
194        self._phases = arg
195    def set_weights(self,arg):
196        self._weights = arg
197    def set_min_res(self):
198        self._min_res = self._min_res_fld.getvalue()
199    def set_max_res(self):
200        self._max_res = self._max_res_fld.getvalue()
201    def set_fofc(self,arg,state):
202        self._fofc = state
203    def set_name_prefix(self):
204       self._name_prefix = self._name_prefix_fld.getvalue()
205
206    def update_state(self):
207        # grab all values
208        self._amplitudes = self._ampl_chooser.get()
209        self._phases     = self._phase_chooser.get()
210        self._weights    = self._wt_chooser.get()
211        self._min_res    = self._min_res_fld.getvalue()
212        self._max_res    = self._max_res_fld.getvalue()
213        self._fofc       = len(self._fofc_chooser.getvalue())>0
214        self._name_prefix= self._name_prefix_fld.getvalue()
215
216    def report_state(self):
217        print("Here is the state of the box")
218        print("Amplitudes:\t%s" % self._amplitudes)
219        print("Phases    :\t%s" % self._phases)
220        print("Weights   :\t%s" % self._weights)
221        print("Min Res   :\t%s" % self._min_res)
222        print("Max Res   :\t%s" % self._max_res)
223        print("FoFc      :\t%s" % str(self._fofc))
224        print("Name Prefix :\t'%s'" % self._name_prefix)
225
226    def show_help(self,msg=None,title=None):
227        # TODO -- CHANGE THE HELP TEXT
228        if msg==None:
229            helpText = pymol.cmd.map_generate.__doc__
230        else:
231            helpText = msg
232
233        if title==None:
234            title="PyMOL Map Loading Help"
235
236        h = Pmw.TextDialog(self._parent,
237                           title=title,)
238        h.insert("end", helpText)
239        h.configure(text_state='disabled')
240
241
242    def run(self,action):
243        if action=="OK":
244            self.update_state()
245            #self.report_state()
246
247            if self._name_prefix==None or self._name_prefix=="":
248                # grep the dataset name from amplitudes
249                if '/' in self._amplitudes:
250                    pfx = self._amplitudes.split('/')
251                    if len(pfx)>=2:
252                        pfx = pfx[1]
253                else:
254                    pfx = self._amplitudes
255            else:
256                pfx = self._name_prefix
257
258            # to ensure a clean name
259            pfx = pymol.cmd.get_unused_name(pfx)
260
261            if not len(self._amplitudes):
262                missing_ampl = """
263To synthesize a map from reflection data you need to specify at
264leastone column for amplitudes and one column for phases. The
265amplitudes column name was blank, and therefore PyMOL cannot create
266the map.  Please select an amplitude column name from the file and try
267again.
268               """
269                self.show_help(missing_ampl,"Missing Amplitudes Column Name")
270                return None
271
272            if not len(self._phases):
273                missing_phases = """
274To synthesize a map from reflection data you need to specify at least
275one column for amplitudes and one column for phases. The phases column
276name was blank, and therefore PyMOL cannot create the map.  Please
277select an amplitude column name from the file and try again.
278               """
279                self.show_help(missing_phases, "Missing Phases Column Name")
280                return None
281
282            try:
283                r = pymol.cmd.map_generate(pfx, self._fileName,
284                                       self._amplitudes, self._phases, self._weights,
285                                       self._min_res, self._max_res, 1, 1)
286            except pymol.CmdException as e:
287                print(e)
288                return None
289
290            if r==None or r=="None" or r=="":
291                print(" MapLoad-Error: PyMOL could not load the MTZ file '%s' due to an unspecified error." % self._fileName)
292                print(" MapLoad-Error: This typically occurs with bad data or blank column names. Please try again")
293                print(" MapLoad-Error: or contact 'help@schrodinger.com' for more information.")
294                return None
295            skin     = pymol._ext_gui.skin
296            try:
297                pymol.cmd.set("suspend_updates", 1)
298                if self._fofc:
299                    toShow = pymol.cmd.get_setting_text("default_fofc_map_rep")
300                    if toShow=="isosurface":
301                        pymol.cmd.isosurface(pymol.cmd.get_unused_name(r+"-srf"),
302                                             pfx, level=1.0)
303                    elif toShow=="isomesh":
304                        meshName=pymol.cmd.get_unused_name(r+"-msh3")
305                        pymol.cmd.isomesh(meshName, pfx, level=3.0)
306                        pymol.cmd.color("green", meshName)
307
308                        meshName=pymol.cmd.get_unused_name(r+"-msh-3")
309                        pymol.cmd.isomesh(meshName, pfx, level=-3.0)
310                        pymol.cmd.color("red", meshName)
311                    else:
312                        # setup volume view
313                        volName = pymol.cmd.get_unused_name(r+"-vol")
314                        pymol.cmd.volume(volName, pfx, "fofc")
315                        # if you don't do this, PyMOL will crash
316                        # when it tries to load the panel
317                else:
318                    toShow = pymol.cmd.get_setting_text("default_2fofc_map_rep")
319                    if toShow=="isosurface":
320                        surfName=pymol.cmd.get_unused_name(r+"-srf")
321                        pymol.cmd.isosurface(surfName, pfx, level=1.0)
322                        pymol.cmd.color("blue", surfName)
323                    elif toShow=="isomesh":
324                        meshName=pymol.cmd.get_unused_name(r+"-msh")
325                        pymol.cmd.isomesh(meshName, pfx, level=1.0)
326                        pymol.cmd.color("blue", meshName)
327                    else:
328                        # setup volume view
329                        volName = pymol.cmd.get_unused_name(r+"-vol")
330                        pymol.cmd.volume(volName, pfx, "2fofc")
331                        # if you don't do this, PyMOL will crash
332                        # when it tries to load the panel
333
334            except:
335                pass
336            finally:
337                pymol.cmd.set("suspend_updates", 0)
338
339            if r!=None:
340                # setting?
341                if pymol.cmd.get_setting_boolean("autoclose_dialogs"):
342                    self.quit()
343
344        elif action=="Cancel":
345            self.quit()
346        elif action=="Help":
347            self.show_help()
348
349
350
351if __name__=="__main__":
352    try:
353        import Tkinter as TK
354    except ImportError:
355        import tkinter as TK
356
357    a = TK.Tk()
358    t = PyMOLMapLoad(a,None)
359    t.pack_and_show()
360    a.mainloop()
361