dplot script from Robert Harrison

contrib/python/dplot.py

@@ -0,0 +1,207 @@
+import math
+import array
+
+'''
+A simple class and helper functions to read and compute with a
+Guassian cube file as produced by NWChem with the DPLOT module.
+It can read the file (or tabulate any function) and then compute the
+value at an arbitrary interior point using tri-linear interpolation.
+
+NWChem seems to produce the cube and atomic coordinates in atomic
+units.
+
+Written for Python 2.*.
+'''
+
+class Cube:
+    '''
+    3D uniform grid with bounding box that supports linear interpolation.
+
+    Data members of the class are Nx, Ny, Nz (the number of points in
+    each dimension), r0 and r1 the lower left and top right corners of
+    the volume, dx, dy, dz the increments in each dimension.
+
+    In the interpolation routine the volume axes are assumed aligned
+    with the corresponding Cartesian axes.
+    '''
+    def __init__(self,N,r0,r1,f=None):
+        '''
+        Makes a Cube of dimension N=(Nx,Ny,Nz) with bounding corners r0,r1
+        optionally initialized with function f(x,y,z).  If f is None, it
+        is initialized to zero.
+        '''
+        
+        self.d = array.array('d',[0.0]*N[0]*N[1]*N[2])
+        self.Nx, self.Ny, self.Nz = N[0], N[1], N[2]
+        self.r0, self.r1 = r0, r1
+        self.dx, self.dy, self.dz = (r1[0]-r0[0])/(N[0]-1), (r1[1]-r0[1])/(N[1]-1), (r1[2]-r0[2])/(N[2]-1)
+
+        # Used to fuzz numerical comparisons
+        self.eps = 1e-15*max(self.dx*self.Nx,self.dy*self.Ny,self.dz*self.Nz)
+
+        #print 'Cube:', self.r0, self.r1
+        #print '   N:', self.Nx, self.Ny, self.Nz
+        #print '   D:', self.dx, self.dy, self.dz
+        #print ' eps:', self.eps
+        if f:
+            for iz in range(self.Nz):
+                for iy in range(self.Ny):
+                    for ix in range(self.Nx):
+                        self[(ix,iy,iz)] = f(*self.get_coords(ix,iy,iz))
+
+    def __getitem__(self,ind):
+        ''' Gets/returns value indexed by index '''
+        ix,iy,iz = ind
+        return self.d[ix + self.Nx*(iy + self.Ny*iz)]
+
+    def __setitem__(self,ind,value):
+        ''' Assigns/sets value indexed by index '''
+        ix,iy,iz = ind
+        self.d[ix + self.Nx*(iy + self.Ny*iz)] = value
+
+    def get_coords(self,ix,iy,iz):
+        ''' Returns coordinates of index '''
+        return (self.r0[0]+ix*self.dx,self.r0[1]+iy*self.dy,self.r0[2]+iz*self.dz)
+
+    def containing_box(self,x,y,z):
+        ''' Returns lower left index of box containing coords '''
+        fac = 0.999999999999999
+        ix, iy, iz = int((x-self.r0[0])*fac/self.dx), int((y-self.r0[1])*fac/self.dy), int((z-self.r0[2])*fac/self.dz)
+        if ix<0 or ix>=(self.Nx-1) or iy<0 or iy>=(self.Ny-1) or iz<0 or iz>=(self.Nz-1):
+            print "Trying to find containing box for point out of bounds"
+            print "point", (x,y,z)
+            print "lower", self.r0
+            print "upper", self.r1
+            raise IndexError
+        return ix, iy, iz
+
+    def linear_1d(self,x,xlo,xhi,flo,fhi):
+        ''' Linear interpolation in 1D ... [xlo---x---xhi] '''
+
+        if (xlo-x)>self.eps or (x-xhi)>self.eps:
+            print "linear_1d: extrapolating!  (x, xlo, xhi):", x, xlo, xhi
+            raise IndexError
+        return flo + (fhi-flo)*(x-xlo)/(xhi-xlo)
+    
+    def interp(self,x,y,z):
+        ''' Linear interpolation within the 3D volume '''
+        ix, iy, iz = self.containing_box(x,y,z)
+        xlo,ylo,zlo = self.get_coords(ix,iy,iz)
+
+        v000 = self[(ix  , iy  , iz  )]
+        v100 = self[(ix+1, iy  , iz  )]
+        v010 = self[(ix  , iy+1, iz  )]
+        v001 = self[(ix  , iy  , iz+1)]
+        v110 = self[(ix+1, iy+1, iz  )]
+        v101 = self[(ix+1, iy  , iz+1)]
+        v011 = self[(ix  , iy+1, iz+1)]
+        v111 = self[(ix+1, iy+1, iz+1)]
+
+        # interp in z
+        v00 = self.linear_1d(z,zlo,zlo+self.dz,v000,v001)
+        v10 = self.linear_1d(z,zlo,zlo+self.dz,v100,v101)
+        v01 = self.linear_1d(z,zlo,zlo+self.dz,v010,v011)
+        v11 = self.linear_1d(z,zlo,zlo+self.dz,v110,v111)
+
+        # interp in y
+        v0 = self.linear_1d(y,ylo,ylo+self.dy,v00,v01)
+        v1 = self.linear_1d(y,ylo,ylo+self.dy,v10,v11)
+
+        # interp in x
+        return self.linear_1d(x,xlo,xlo+self.dx,v0,v1)
+        
+def read_i_f_f_f(f):
+    ''' Read line containing integer and three floats '''
+    line = f.readline().lstrip().rstrip().split()
+    return int(line[0]), float(line[1]), float(line[2]), float(line[3])
+
+def read_atom(f):
+    ''' Read line from Guassian cube file containing atomic info '''
+    line = f.readline().lstrip().rstrip().split()
+    return int(line[0]), (float(line[2]), float(line[3]), float(line[4]))
+
+def load_gaussian(filename):
+    '''
+    Returns tuple (cube,atoms) loaded from Gaussian cube file.
+
+    An atom in the list is the tuple (Z,(X,Y,Z)).
+
+    The volume axes are assumed aligned with the corresponding
+    Cartesian axes.
+    '''
+    f = open(filename,'r')
+    f.readline() # discard two comment lines
+    f.readline()
+
+    natoms, xlo, ylo, zlo = read_i_f_f_f(f)
+    Nx, dx, junk, junk = read_i_f_f_f(f)
+    Ny, junk, dy, junk = read_i_f_f_f(f)
+    Nz, junk, junk, dz = read_i_f_f_f(f)
+    
+    # Load geometry
+    atoms = []
+    for i in range(natoms): 
+        atoms.append(read_atom(f))
+
+    xhi = xlo + (Nx-1)*dx
+    yhi = ylo + (Ny-1)*dy
+    zhi = zlo + (Nz-1)*dz
+    
+    c = Cube((Nx,Ny,Nz),(xlo,ylo,zlo),(xhi,yhi,zhi))
+    
+    maxval = 0.0
+    for ix in range(Nx):
+        for iy in range(Ny):
+            line = f.readline().lstrip().rstrip().split()
+            n = 0
+            for iz in range(Nz):
+                try:
+                    value = float(line[n])
+                    maxval = max(maxval,abs(value))
+                    c[(ix,iy,iz)] = value
+                except:
+                    print line
+                    print n
+                    raise IndexError
+                n = n + 1
+                if n == 6:
+                    line = f.readline().lstrip().rstrip().split()
+                    n = 0
+
+    #print "maxval",maxval
+    return c, atoms
+
+if __name__ == "__main__":
+
+    def testfun(x,y,z):
+        return math.sin(x+y*0.7+z*0.5)
+        #return math.exp(-(x**2 + y**2 + z**2))
+        #return 0.1*x + 0.3*y + 0.7*z
+    
+    def test():
+        N = (50,60,70)
+        r0 = (-0.5,-0.5,-0.5)
+        r1 = (0.0,0.0,0.0)
+
+        c = Cube(N,r0,r1,testfun)
+
+        # print the interpolated function and error out along a line
+        dx,dy,dz = (r1[0]-r0[0])/10, (r1[1]-r0[1])/10, (r1[2]-r0[2])/10,
+        print "   x      y      z      exact     interp     error"
+        print "------ ------ ------ ---------- ---------- ----------"
+        for i in range(11):
+            x,y,z = r0[0]+i*dx, r0[1]+i*dy, r0[2]+i*dz
+            numeric = c.interp(x,y,z)
+            exact = testfun(x,y,z)
+            print "%6.2f %6.2f %6.2f %10.6f %10.6f %10.2e" % (x,y,z,exact,numeric,exact-numeric)
+    
+    test()
+    
+    #c, atoms = load_gaussian("dens.cube")
+    #print atoms
+    #print 'Cube:', c.r0, c.r1
+    #print '   N:', c.Nx, c.Ny, c.Nz
+    #print '   D:', c.dx, c.dy, c.dz
+    #print ' eps:', c.eps
+    
+