NWChem/contrib/python/dplot.py
2021-10-30 02:11:52 +02:00

207 lines
6.9 KiB
Python

import math
import array
'''
A simple class and helper functions to read and compute with a
Gaussian 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 Gaussian 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