1#!/usr/bin/python
2
3""" psize class
4
5    Get dimensions and other information from a PQR file.
6
7    Originally written by Dave Sept
8    Additional APBS-specific features added by Nathan Baker
9    Ported to Python/Psize class by Todd Dolinsky and subsequently
10    hacked by Nathan Baker
11
12        ----------------------------
13
14    PDB2PQR -- An automated pipeline for the setup, execution, and analysis of
15    Poisson-Boltzmann electrostatics calculations
16
17    Copyright (c) 2002-2011, Jens Erik Nielsen, University College Dublin;
18    Nathan A. Baker, Battelle Memorial Institute, Developed at the Pacific
19    Northwest National Laboratory, operated by Battelle Memorial Institute,
20    Pacific Northwest Division for the U.S. Department Energy.;
21    Paul Czodrowski & Gerhard Klebe, University of Marburg.
22
23	All rights reserved.
24
25	Redistribution and use in source and binary forms, with or without modification,
26	are permitted provided that the following conditions are met:
27
28		* Redistributions of source code must retain the above copyright notice,
29		  this list of conditions and the following disclaimer.
30		* Redistributions in binary form must reproduce the above copyright notice,
31		  this list of conditions and the following disclaimer in the documentation
32		  and/or other materials provided with the distribution.
33        * Neither the names of University College Dublin, Battelle Memorial Institute,
34          Pacific Northwest National Laboratory, US Department of Energy, or University
35          of Marburg nor the names of its contributors may be used to endorse or promote
36          products derived from this software without specific prior written permission.
37
38	THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
39	ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
40	WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
41	IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
42	INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
43	BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
44	DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
45	LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
46	OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
47	OF THE POSSIBILITY OF SUCH DAMAGE.
48
49    ----------------------------
50"""
51
52__date__ = "4 June 2008"
53__author__ = "Dave Sept, Nathan Baker, Todd Dolinsky, Yong Huang"
54
55import string, sys, getopt
56from sys import stdout, stderr
57from math import log
58
59class Psize:
60    """Master class for parsing input files and suggesting settings"""
61    def __init__(self):
62        self.constants = {"cfac": 1.7, "fadd":20, "space": 0.50, "gmemfac": 200, "gmemceil": 400, "ofrac":0.1, "redfac": 0.25 }
63        self.minlen = [None, None, None]
64        self.maxlen = [None, None, None]
65        self.q = 0.0
66        self.gotatom = 0
67        self.gothet = 0
68        self.olen = [0.0, 0.0, 0.0]
69        self.cen = [0.0, 0.0, 0.0]
70        self.clen = [0.0, 0.0, 0.0]
71        self.flen = [0.0, 0.0, 0.0]
72        self.n = [0, 0, 0]
73        self.np = [0.0, 0.0, 0.0]
74        self.nsmall = [0,0,0]
75        self.nfocus = 0
76
77    def parseString(self, structure):
78        """ Parse the input structure as a string in PDB or PQR format """
79        lines = string.split(structure, "\n")
80        self.parseLines(lines)
81
82    def parseInput(self, filename):
83        """ Parse input structure file in PDB or PQR format """
84        with open(filename, 'rU') as f:
85            self.parseLines(f.readlines())
86
87    def parseLines(self, lines):
88        """ Parse the lines """
89        for line in lines:
90            if string.find(line,"ATOM") == 0:
91                subline = string.replace(line[30:], "-", " -")
92                words = string.split(subline)
93                if len(words) < 5:
94                    continue
95                self.gotatom += 1
96                self.q = self.q + float(words[3])
97                rad = float(words[4])
98                center = []
99                for word in words[0:3]:
100                    center.append(float(word))
101                for i in range(3):
102                    if self.minlen[i] == None or center[i]-rad < self.minlen[i]:
103                        self.minlen[i] = center[i]-rad
104                    if self.maxlen[i] == None or center[i]+rad > self.maxlen[i]:
105                        self.maxlen[i] = center[i]+rad
106            elif string.find(line, "HETATM") == 0:
107                self.gothet = self.gothet + 1
108                # Special handling for no ATOM entries in the pqr file, only HETATM entries
109                if self.gotatom == 0:
110                    subline = string.replace(line[30:], "-", " -")
111                    words = string.split(subline)
112                    if len(words) < 5:
113                        continue
114                    self.q = self.q + float(words[3])
115                    rad = float(words[4])
116                    center = []
117                    for word in words[0:3]:
118                        center.append(float(word))
119                    for i in range(3):
120                        if self.minlen[i] == None or center[i]-rad < self.minlen[i]:
121                            self.minlen[i] = center[i]-rad
122                        if self.maxlen[i] == None or center[i]+rad > self.maxlen[i]:
123                            self.maxlen[i] = center[i]+rad
124
125    def setConstant(self, name, value):
126        """ Set a constant to a value; returns 0 if constant not found """
127        try:
128            self.constants[name] = value
129            return 1
130        except KeyError:
131            return 0
132
133    def getConstant(self, name):
134        """ Get a constant value; raises KeyError if constant not found """
135        return self.constants[name]
136
137    def setLength(self, maxlen, minlen):
138        """ Compute molecule dimensions """
139        for i in range(3):
140            self.olen[i] = maxlen[i] - minlen[i]
141            if self.olen[i] < 0.1:
142                self.olen[i] = 0.1
143        return self.olen
144
145
146    def setCoarseGridDims(self, olen):
147        """ Compute coarse mesh dimensions """
148        for i in range(3):
149            self.clen[i] = self.constants["cfac"] * olen[i]
150        return self.clen
151
152    def setFineGridDims(self, olen, clen):
153        """ Compute fine mesh dimensions """
154        for i in range(3):
155            self.flen[i] = olen[i] + self.constants["fadd"]
156            if self.flen[i] > clen[i]:
157                #str = "WARNING: Fine length (%.2f) cannot be larger than coarse length (%.2f)\n" % (self.flen[i], clen[i])
158                #str = str + "         Setting fine grid length equal to coarse grid length\n"
159                #stdout.write(str)
160                self.flen[i] = clen[i]
161        return self.flen
162
163    def setCenter(self, maxlen, minlen):
164        """ Compute molecule center """
165        for i in range(3):
166            self.cen[i] = (maxlen[i] + minlen[i]) / 2
167        return self.cen
168
169
170    def setFineGridPoints(self, flen):
171        """ Compute mesh grid points, assuming 4 levels in MG hierarchy """
172        tn = [0,0,0]
173        for i in range(3):
174            tn[i] = int(flen[i]/self.constants["space"] + 0.5)
175            self.n[i] = 32*(int((tn[i] - 1) / 32.0 + 0.5)) + 1
176            if self.n[i] < 33:
177                self.n[i] = 33
178        return self.n
179
180    def setSmallest(self, n):
181        """ Compute parallel division in case memory requirement above ceiling
182        Find the smallest dimension and see if the number of grid points in
183        that dimension will fit below the memory ceiling
184        Reduce nsmall until an nsmall^3 domain will fit into memory """
185        nsmall = []
186        for i in range(3):
187            nsmall.append(n[i])
188        while 1:
189            nsmem = 200.0 * nsmall[0] * nsmall[1] * nsmall[2] / 1024 / 1024
190            if nsmem < self.constants["gmemceil"]: break
191            else:
192                i = nsmall.index(max(nsmall))
193                nsmall[i] = 32 * ((nsmall[i] - 1)/32 - 1) + 1
194                if nsmall <= 0:
195                    stdout.write("You picked a memory ceiling that is too small\n")
196                    sys.exit(0)
197
198        self.nsmall = nsmall
199        return nsmall
200
201    def setProcGrid(self, n, nsmall):
202        """ Calculate the number of processors required to span each
203        dimension """
204
205        zofac = 1 + 2 * self.constants["ofrac"]
206        for i in range(3):
207            self.np[i] = n[i]/float(nsmall[i])
208            if self.np[i] > 1: self.np[i] = int(zofac*n[1]/nsmall[i] + 1.0)
209        return self.np
210
211    def setFocus(self, flen, np, clen):
212        """ Calculate the number of levels of focusing required for each
213        processor subdomain """
214
215        nfoc = [0,0,0]
216        for i in range(3):
217            nfoc[i] = int(log((flen[i]/np[i])/clen[i])/log(self.constants["redfac"]) + 1.0)
218        nfocus = nfoc[0]
219        if nfoc[1] > nfocus: nfocus = nfoc[1]
220        if nfoc[2] > nfocus: nfocus = nfoc[2]
221        if nfocus > 0: nfocus = nfocus + 1
222        self.nfocus = nfocus
223
224    def setAll(self):
225        """ Set up all of the things calculated individually above """
226        maxlen = self.getMax()
227        minlen = self.getMin()
228        self.setLength(maxlen, minlen)
229        olen = self.getLength()
230
231        self.setCoarseGridDims(olen)
232        clen = self.getCoarseGridDims()
233
234        self.setFineGridDims(olen, clen)
235        flen = self.getFineGridDims()
236
237        self.setCenter(maxlen, minlen)
238        cen = self.getCenter()
239
240        self.setFineGridPoints(flen)
241        n = self.getFineGridPoints()
242
243        self.setSmallest(n)
244        nsmall = self.getSmallest()
245
246        self.setProcGrid(n, nsmall)
247        np = self.getProcGrid()
248
249        self.setFocus(flen, np, clen)
250        nfocus = self.getFocus()
251
252    def getMax(self): return self.maxlen
253    def getMin(self): return self.minlen
254    def getCharge(self): return self.q
255    def getLength(self): return self.olen
256    def getCoarseGridDims(self): return self.clen
257    def getFineGridDims(self): return self.flen
258    def getCenter(self): return self.cen
259    def getFineGridPoints(self): return self.n
260    def getSmallest(self): return self.nsmall
261    def getProcGrid(self): return self.np
262    def getFocus(self): return self.nfocus
263
264    def runPsize(self, filename):
265        """ Parse input PQR file and set parameters """
266        self.parseInput(filename)
267        self.setAll()
268
269    def printResults(self):
270        """ Return a string with the formatted results """
271
272        str = "\n"
273
274        if self.gotatom > 0:
275
276            maxlen = self.getMax()
277            minlen = self.getMin()
278            q = self.getCharge()
279            olen = self.getLength()
280            clen = self.getCoarseGridDims()
281            flen = self.getFineGridDims()
282            cen = self.getCenter()
283            n = self.getFineGridPoints()
284            nsmall = self.getSmallest()
285            np = self.getProcGrid()
286            nfocus = self.getFocus()
287
288            # Compute memory requirements
289
290            nsmem = 200.0 * nsmall[0] * nsmall[1] * nsmall[2] / 1024 / 1024
291            gmem = 200.0 * n[0] * n[1] * n[2] / 1024 / 1024
292
293            # Print the calculated entries
294            str = str + "################# MOLECULE INFO ####################\n"
295            str = str + "Number of ATOM entries = %i\n" % self.gotatom
296            str = str + "Number of HETATM entries (ignored) = %i\n" % self.gothet
297            str = str + "Total charge = %.3f e\n" % q
298            str = str + "Dimensions = %.3f x %.3f x %.3f A\n" % (olen[0], olen[1], olen[2])
299            str = str + "Center = %.3f x %.3f x %.3f A\n" % (cen[0], cen[1], cen[2])
300            str = str + "Lower corner = %.3f x %.3f x %.3f A\n" % (minlen[0], minlen[1], minlen[2])
301            str = str + "Upper corner = %.3f x %.3f x %.3f A\n" % (maxlen[0], maxlen[1], maxlen[2])
302
303            str = str + "\n"
304            str = str + "############## GENERAL CALCULATION INFO #############\n"
305            str = str + "Coarse grid dims = %.3f x %.3f x %.3f A\n" % (clen[0],
306clen[1], clen[2])
307            str = str + "Fine grid dims = %.3f x %.3f x %.3f A\n" % (flen[0], flen[1], flen[2])
308            str = str + "Num. fine grid pts. = %i x %i x %i\n" % (n[0], n[1], n[2])
309
310            if gmem > self.constants["gmemceil"]:
311                str = str + "Parallel solve required (%.3f MB > %.3f MB)\n" % (gmem, self.constants["gmemceil"])
312                str = str + "Total processors required = %i\n" % (np[0]*np[1]*np[2])
313                str = str + "Proc. grid = %i x %i x %i\n" % (np[0], np[1], np[2])
314                str = str + "Grid pts. on each proc. = %i x %i x %i\n" % (nsmall[0], nsmall[1], nsmall[2])
315                xglob = np[0]*round(nsmall[0]/(1 + 2*self.constants["ofrac"]) - .001)
316                yglob = np[1]*round(nsmall[1]/(1 + 2*self.constants["ofrac"]) - .001)
317                zglob = np[2]*round(nsmall[2]/(1 + 2*self.constants["ofrac"]) - .001)
318                if np[0] == 1: xglob = nsmall[0]
319                if np[1] == 1: yglob = nsmall[1]
320                if np[2] == 1: zglob = nsmall[2]
321                str = str + "Fine mesh spacing = %g x %g x %g A\n" % (flen[0]/(xglob-1), flen[1]/(yglob-1), flen[2]/(zglob-1))
322                str = str + "Estimated mem. required for parallel solve = %.3f MB/proc.\n" % nsmem
323                ntot = nsmall[0]*nsmall[1]*nsmall[2]
324
325            else:
326                str = str + "Fine mesh spacing = %g x %g x %g A\n" % (flen[0]/(n[0]-1), flen[1]/(n[1]-1), flen[2]/(n[2]-1))
327                str = str + "Estimated mem. required for sequential solve = %.3f MB\n" % gmem
328                ntot = n[0]*n[1]*n[2]
329
330            str = str + "Number of focusing operations = %i\n" % nfocus
331
332            str = str + "\n"
333            str = str + "################# ESTIMATED REQUIREMENTS ####################\n"
334            str = str + "Memory per processor                   = %.3f MB\n" % (200.0*ntot/1024/1024)
335            str = str + "Grid storage requirements (ASCII)      = %.3f MB\n" % (8.0*12*np[0]*np[1]*np[2]*ntot/1024/1024)
336            str = str + "\n"
337
338        else:
339            str = str + "No ATOM entires in file!\n\n"
340
341        return str
342
343def usage(rc):
344    """ Print usage information and exit with error code rc """
345    psize = Psize()
346    usage = "\n"
347    usage = usage + "Psize script\n"
348    usage = usage + "Usage: psize.py [opts] <filename>\n"
349    usage = usage + "Optional Arguments:\n"
350    usage = usage + "  --help               : Display this text\n"
351    usage = usage + "  --cfac=<value>       : Factor by which to expand mol dims to\n"
352    usage = usage + "                         get coarse grid dims\n"
353    usage = usage + "                         [default = %g]\n" % psize.getConstant("cfac")
354    usage = usage + "  --fadd=<value>       : Amount to add to mol dims to get fine\n"
355    usage = usage + "                         grid dims\n"
356    usage = usage + "                         [default = %g]\n" % psize.getConstant("fadd")
357    usage = usage + "  --space=<value>      : Desired fine mesh resolution\n"
358    usage = usage + "                         [default = %g]\n" % psize.getConstant("space")
359    usage = usage + "  --gmemfac=<value>    : Number of bytes per grid point required\n"
360    usage = usage + "                         for sequential MG calculation\n"
361    usage = usage + "                         [default = %g]\n" % psize.getConstant("gmemfac")
362    usage = usage + "  --gmemceil=<value>   : Max MB allowed for sequential MG\n"
363    usage = usage + "                         calculation.  Adjust this to force the\n"
364    usage = usage + "                         script to perform faster calculations (which\n"
365    usage = usage + "                         require more parallelism).\n"
366    usage = usage + "                         [default = %g]\n" % psize.getConstant("gmemceil")
367    usage = usage + "  --ofrac=<value>       : Overlap factor between mesh partitions\n"
368    usage = usage + "                         [default = %g]\n" % psize.getConstant("ofrac")
369    usage = usage + "  --redfac=<value>     : The maximum factor by which a domain\n"
370    usage = usage + "                         dimension can be reduced during focusing\n"
371    usage = usage + "                         [default = %g]\n" % psize.getConstant("redfac")
372
373
374    stderr.write(usage)
375    sys.exit(rc)
376
377def main():
378    filename = ""
379    shortOptList = "h"
380    longOptList = ["help", "cfac=", "fadd=", "space=", "gmemfac=", "gmemceil=", "ofrac=", "redfac="]
381    try:
382        opts, args = getopt.getopt(sys.argv[1:], shortOptList, longOptList)
383    except getopt.GetoptError, details:
384        stderr.write("Option error (%s)!\n" % details)
385        usage(2)
386    if len(args) != 1:
387        stderr.write("Invalid argument list!\n")
388        usage(2)
389    else:
390        filename = args[0]
391
392    psize = Psize()
393
394    for o, a in opts:
395        if o.lower() == "--help" or o == "-h":
396            usage(0)
397        if o.lower() == "--cfac":
398            psize.setConstant("cfac", float(a))
399        if o.lower() == "--fadd":
400            psize.setConstant("fadd", int(a))
401        if o.lower() == "--space":
402            psize.setConstant("space", float(a))
403        if o.lower() == "--gmemfac":
404            psize.setConstant("gmemfac", int(a))
405        if o.lower() == "--gmemceil":
406            psize.setConstant("gmemceil",  int(a))
407        if o.lower() == "--ofrac":
408            psize.setConstant("ofrac", float(a))
409        if o.lower() == "--redfac":
410            psize.setConstant("redfac", float(a))
411
412    psize.runPsize(filename)
413
414    stdout.write("# Constants used: \n");
415    for key in psize.constants.keys():
416        stdout.write("# \t%s: %s\n" % (key, psize.constants[key]))
417    stdout.write("# Run:\n")
418    stdout.write("#    `%s --help`\n" % sys.argv[0])
419    stdout.write("# for more information on these default values\n" )
420    stdout.write(psize.printResults())
421
422
423if __name__ == "__main__": main()
424