def calcWF(options, geom, basis, Y):
"""
Calculate wavefunction, returns:
YY : complex wavefunction on regular grid
dstep : stepsize
origo : vector
nx, ny, nz : number of grid points
"""
xyz=N.array(geom.xyz[options.DeviceAtoms[0]-1:options.DeviceAtoms[1]])
# Size of cube
xmin, xmax = min(xyz[:, 0])-5.0, max(xyz[:, 0])+5.0
ymin, ymax = min(xyz[:, 1])-5.0, max(xyz[:, 1])+5.0
zmin, zmax = min(xyz[:, 2])-5.0, max(xyz[:, 2])+5.0
xl, yl, zl = xmax-xmin, ymax-ymin, zmax-zmin
dx, dy, dz = options.res, options.res, options.res
nx, ny, nz = int(xl/dx)+1, int(yl/dy)+1, int(zl/dz)+1
origo = N.array([xmin, ymin, zmin], N.float)
# Def cube
YY=N.zeros((nx, ny, nz), N.complex)
rx=N.array(range(nx), N.float)*dx+origo[0]
ry=N.array(range(ny), N.float)*dy+origo[1]
rz=N.array(range(nz), N.float)*dz+origo[2]
for ii, Yval in enumerate(Y):
if ii>0:# and ii%(int(len(Y)/10))==0:
SIO.printDone(ii, len(Y), 'Wavefunction')
rax, ray, raz=basis.xyz[ii, 0], basis.xyz[ii, 1], basis.xyz[ii, 2]
# Only calulate in subset
ixmin, ixmax = int((rax-origo[0]-basis.coff[ii])/dx), \
int((rax-origo[0]+basis.coff[ii])/dx)
iymin, iymax = int((ray-origo[1]-basis.coff[ii])/dy), \
int((ray-origo[1]+basis.coff[ii])/dy)
izmin, izmax = int((raz-origo[2]-basis.coff[ii])/dz), \
int((raz-origo[2]+basis.coff[ii])/dz)
ddx, ddy, ddz=rx[ixmin:ixmax]-rax, ry[iymin:iymax]-ray, rz[izmin:izmax]-raz
dr=N.sqrt(MM.outerAdd(ddx*ddx, ddy*ddy, ddz*ddz))
drho=N.sqrt(MM.outerAdd(ddx*ddx, ddy*ddy, 0*ddz))
imax=(basis.coff[ii]-2*basis.delta[ii])/basis.delta[ii]
ri=dr/basis.delta[ii]
ri=N.where(ri<imax, ri, imax)
ri=ri.astype(N.int)
costh = MM.outerAdd(0*ddx, 0*ddy, ddz)/dr
cosfi, sinfi = MM.outerAdd(ddx, 0*ddy, 0*ddz)/drho, MM.outerAdd(0*ddx, ddy, 0*ddz)/drho
# Numpy has changed the choose function to crap!
RR=N.take(basis.orb[ii], ri)
# Calculate spherical harmonics
l = basis.L[ii]
m = basis.M[ii]
if l==3:
print 'f-shell : l=%i, m=%i (NOT TESTED!!)'%(l, m)
thisSphHar = MM.sphericalHarmonics(l, m, costh, sinfi, cosfi)
YY[ixmin:ixmax, iymin:iymax, izmin:izmax]=YY[ixmin:ixmax, iymin:iymax, izmin:izmax]+\
RR*thisSphHar*Yval
print "Wave function norm on real space grid:", N.sum(YY.conjugate()*YY)*dx*dy*dz
return YY, options.res, origo, nx, ny, nz
def calcWF2(options, geom, DeviceAtoms, basis, Y, NN, Fold=True, k=[0, 0, 0], a=None, fn=''):
"""
Calculate wavefunction on real space mesh with optional folding
of the periodic boundary conditions. If Fold==True the vectors 'a'
will be choosen to be the periodic boundary condition vectors and
the real space wavefunction folded into the cell using the k-vector.
If Folded==False, you can choose any 'a' but folding by the PBC will
not be done.
Note: Folding assumes coupling only between N.N. cells
INPUT:
geom : MakeGeom structure for full geometry
DeviceAtoms : [first, last] numbering from 1 to N
basis : Basis struct for ONLY the device region
Y : Wavefunction for the device region
NN : [N1,N2,N3,minN3,maxN3] number of points along [a1, a2, a3]
only calculate from minN3 to maxN3
Fold : fold periodic boundary conditions
k : k-vector to use for folding [-0.5,0.5]
a : if Fold==True, uses geom.pbc, if false: [a1, a2, a3]
along these directions, i.e., da1=a1/N1 ...
RETURNS:
YY : complex wavefunction as N1, N2, N3 matrix
"""
xyz=N.array(geom.xyz[DeviceAtoms[0]-1:DeviceAtoms[1]])
N1, N2, N3, minN3, maxN3 = NN[0], NN[1], NN[2], NN[3], NN[4]
NN3 = maxN3-minN3+1
if Fold:
da1, da2, da3 = geom.pbc[0]/N1, geom.pbc[1]/N2, geom.pbc[2]/N3
else:
da1, da2, da3 = a[0]/N1, a[1]/N2, a[2]/N3
try:
NNY=Y.shape[1]
except:
NNY=1
# Def cube
YY=[N.zeros((N1, N2, NN3), N.complex) for ii in range(NNY)]
# Help indices
i1 = MM.outerAdd(N.arange(N1), N.zeros(N2), N.zeros(NN3))
i2 = MM.outerAdd(N.zeros(N1), N.arange(N2), N.zeros(NN3))
i3 = MM.outerAdd(N.zeros(N1), N.zeros(N2), N.arange(NN3)+minN3)
# Positions x,y,z for points in matrix form
rx = i1*da1[0]+i2*da2[0]+i3*da3[0]
ry = i1*da1[1]+i2*da2[1]+i3*da3[1]
rz = i1*da1[2]+i2*da2[2]+i3*da3[2]
if Fold:
pbc = N.array([da1*N1, da2*N2, da3*N3])
orig = da3*minN3
b = N.transpose(LA.inv(pbc))
olist = [orig, orig+pbc[0]+pbc[1]+pbc[2]]
pairs = N.array([[b[1], b[2]], [b[0], b[2]], [b[0], b[1]]])
pbcpairs = N.array([[pbc[1], pbc[2]], [pbc[0], pbc[2]], [pbc[0], pbc[1]]])
for iiatom in range(DeviceAtoms[1]-DeviceAtoms[0]+1):
if iiatom>0:
#Wavefunction percent done...
SIO.printDone(iiatom, DeviceAtoms[1]-DeviceAtoms[0]+1, 'Wave function')
if not Fold:
shiftvec = [0, 0, 0]
else: # include shift vectors if sphere cuts planes of PBC
basisindx = N.where(basis.ii==iiatom+DeviceAtoms[0])[0]
basisradius = N.max(basis.coff[basisindx])
minshift, maxshift = [0, 0, 0], [0, 0, 0]
for iithiso, thiso in enumerate(olist):
c=xyz[iiatom, :]-thiso
for ip in range(3):
n = N.cross(pbcpairs[ip, 0], pbcpairs[ip, 1])
n = n / LA.norm(n)
dist = N.abs(N.dot(n, N.dot(N.dot(pairs[ip], c), pbcpairs[ip])-c))
if dist < basisradius:
if iithiso==0:
minshift[ip]=-1
else:
maxshift[ip]=1
shiftvec = []
for ix in range(minshift[0], maxshift[0]+1):
for iy in range(minshift[1], maxshift[1]+1):
for iz in range(minshift[2], maxshift[2]+1):
shiftvec+=[[ix, iy, iz]]
#print(shiftvec)
for shift in shiftvec:
# Atom position shifted
vec=-N.dot(N.array(shift), geom.pbc)
phase = N.exp(-1.0j*(2*N.pi*N.dot(N.array(k), N.array(shift))))
# Difference and polar, cylinder distances
dx, dy, dz=rx-xyz[iiatom, 0]-vec[0], ry-xyz[iiatom, 1]-vec[1], rz-xyz[iiatom, 2]-vec[2]
dr2=dx*dx+dy*dy+dz*dz
drho2=dx*dx+dy*dy
basisindx = N.where(basis.ii==iiatom+DeviceAtoms[0])[0]
old_coff, old_delta = 0.0, 0.0
for basisorb in basisindx:
if not (basis.coff[basisorb]==old_coff and basis.delta[basisorb]==old_delta):
old_coff, old_delta = basis.coff[basisorb], basis.delta[basisorb]
indx = N.where(dr2<basis.coff[basisorb]**2) # Find points close to atom
idr, idrho=N.sqrt(dr2[indx]), N.sqrt(drho2[indx])
iri=(idr/basis.delta[basisorb]).astype(N.int)
idx, idy, idz = dx[indx], dy[indx], dz[indx]
costh = idz/idr
cosfi, sinfi = idx/idrho, idy/idrho
# Fix divide by zeros
indxRzero, indxRhozero = N.where(idr==0.0), N.where(idrho==0.0)
costh[indxRzero] = 1.0
cosfi[indxRhozero], sinfi[indxRhozero] = 1.0, 0.0
# Numpy has changed the choose function to crap!
RR=N.take(basis.orb[basisorb], iri)
# Calculate spherical harmonics
if len(idr)>0:
l = basis.L[basisorb]
m = basis.M[basisorb]
if l==3:
print('f-shell : l=%i, m=%i (NOT TESTED!!)'%(l, m))
thisSphHar = MM.sphericalHarmonics(l, m, costh, sinfi, cosfi)
for iy in range(NNY):
YY[iy][indx]=YY[iy][indx]+RR*thisSphHar*Y[basisorb, iy]*phase
N1, N2, N3, minN3, maxN3 = NN[0], NN[1], NN[2], NN[3], NN[4]
for ii in range(NNY):
writenetcdf2(geom, fn+'%i.nc'%ii, YY[ii], N1, N2, N3, minN3, maxN3, geom.pbc, options.DeviceAtoms)
return YY