Esempio n. 1
0
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
Esempio n. 2
0
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