vaspwfc.py 文件源码

def wfc_r(self, ispin=1, ikpt=1, iband=1,
                    gvec=None, ngrid=None, norm=False,
                    gamma=False):
        '''
        Obtain the pseudo-wavefunction of the specified KS states in real space
        by performing FT transform on the reciprocal space planewave
        coefficients.  The 3D FT grid size is determined by ngrid, which
        defaults to self._ngrid if not given.  Gvectors of the KS states is used
        to put 1D planewave coefficients back to 3D grid.
        '''
        self.checkIndex(ispin, ikpt, iband)

        if ngrid is None:
            ngrid = self._ngrid.copy()
        else:
            ngrid = np.array(ngrid, dtype=int)
            assert ngrid.shape == (3,)
            assert np.alltrue(ngrid >= self._ngrid), \
                    "Minium FT grid size: (%d, %d, %d)" % \
                    (self._ngrid[0], self._ngrid[1], self._ngrid[2])
        if gvec is None:
            gvec = self.gvectors(ikpt, gamma)

        if gamma:
            phi_k = np.zeros((ngrid[0], ngrid[1], ngrid[2]/2 + 1), dtype=np.complex128)
        else:
            phi_k = np.zeros(ngrid, dtype=np.complex128)

        gvec %= ngrid[np.newaxis,:]
        phi_k[gvec[:,0], gvec[:,1], gvec[:,2]] = self.readBandCoeff(ispin, ikpt, iband, norm)

        if gamma:
            # add some components that are excluded and perform c2r FFT
            for ii in range(ngrid[0]):
                for jj in range(ngrid[1]):
                    fx = ii if ii < ngrid[0] / 2 + 1 else ii - ngrid[0]
                    fy = ii if ii < ngrid[1] / 2 + 1 else ii - ngrid[1]
                    if (fy > 0) or (fy == 0 and fx >= 0):
                        continue
                    phi_k[ii,jj,0] = phi_k[-ii,-jj,0].conjugate()
            phi_k /= np.sqrt(2.)
            phi_k[0,0,0] *= np.sqrt(2.)
            return np.fft.irfftn(phi_k, s=ngrid)
        else:
            # perform complex2complex FFT
            return ifftn(phi_k)