Port of Quantum Espresso's setlocq routine

setlocq.py

#------------------------------------------------------------------------------
# Function: setlocq
#------------------------------------------------------------------------------
#
# Description:
# Port of Quantum Espresso's 'setlocq' routine (PHonon/PH/setlocq.f90)
#
#------------------------------------------------------------------------------
#
# What does it do:
#  * Compute fourier transform of local potential for a given q-point
#  * If available, uses gpaw's erf function instead of math.erf
#  * Also ports Quantum Espresso's 'simpson' routine (flib/simpsn.f90)
#
#------------------------------------------------------------------------------
#
# Dependency:
# * numpy
# * gpaw (optional)
#
#------------------------------------------------------------------------------
#
# Author: P. R. Vaidyanathan (aditya95sriram <at> gmail <dot> com)
# Date: 27th June 2017

import numpy as np
import math

# try using the gpaw error-function, fall back to math module's erf
try:
    import gpaw
except ImportError:
    print "using error function from 'math' module (less accurate than gpaw)"
    erf = math.erf
else:
    erf = gpaw.utilities.erf

def setlocq(xq, mesh, msh, rab, r, vloc_at, zp, tpiba2, ngm, g, omega):    
    """ This routine computes the Fourier transform of the local
    part of the pseudopotential in the q+G vectors.

    The local pseudopotential of the US case is always in
    numerical form, expressed in Ry units.
    
    Assuming all arrays are dtype=np.double or np.float64
    
    parameters 
    integer:ngm           - number of g vectors
            mesh          - dimensions of mesh
            msh           - mesh points for radial integration
    double: xq(3)         - the q point
            zp            - valence pseudocharge
            rab(mesh)     - the derivative of mesh points
            r(mesh)       - the mesh points
            vloc_at(mesh) - the pseudo on the radial
            tpiba2        - 2 pi / alat
            omega         - the volume of the unit cell
            g(ngm,3)      - the g vectors coordinates

    returns
            vloc(ngm)     - the fourier transform of the potential
    """
        
    # slightly more accurate than fortran double precision (selected_real_kind(14,200))
    DP = np.double 
    
    # constants
    e2 = DP(2.0) # square of the electron charge
    pi = DP('3.14159265358979323846')
    fpi = 4*pi
    
    eps = DP('1e-8')
    # auxiliary variables
    aux = np.empty(mesh)
    aux1 = np.empty(mesh)
    # the erf function
    qe_erf = gpaw.utilities.erf
    
    # Pseudopotentials in numerical form (Vnl(lloc) contain the local part)
    # in order to perform the Fourier transform, a term erf(r)/r is
    # subtracted in real space and added again in G space
    
    # first the G=0 term
    aux = r * ( r*vloc_at + zp*e2)
    
    vloc0 = simpson(msh, aux, rab)
    
    #   here the G<>0 terms, we first compute the part of the integrand func
    #   indipendent of |G| in real space
    #
    aux1 = r * vloc_at + zp*e2*qe_erf(r)
    fac = zp*e2 / tpiba2

    #    and here we perform the integral, after multiplying for the |G|
    #    dependent  part
    #
    vloc = np.empty(ngm)
    pre_g2a = np.sum(np.square(xq + g), axis=1)
    for ig in range(ngm):
        g2a = pre_g2a[ig]
        if (g2a < eps):
            vloc[ig] = vloc0
        else:
            gx = math.sqrt(g2a * tpiba2)
            aux = aux1 * np.sin(gx * r) / gx
            vlcp = simpson(msh, aux, rab)
            #     here we add the analytic fourier transform of the erf function
            vlcp = vlcp - fac * math.exp( -g2a * tpiba2 * 0.25 ) / g2a
            vloc[ig] = vlcp
    vloc = vloc * fpi / omega
    return vloc
    
def simpson(mesh, func, rab):
    """simpson's rule integration. On input:
       mesh = the number of grid points (should be odd)
       func(i)= function to be integrated
       rab(i) = r(i) * dr(i)/di * di
     For the logarithmic grid not including r=0 :
       r(i) = r_0*exp((i-1)*dx) ==> rab(i)=r(i)*dx
     For the logarithmic grid including r=0 :
       r(i) = a(exp((i-1)*dx)-1) ==> rab(i)=(r(i)+a)*dx
     Output in asum = \sum_i c_i f(i)*rab(i) = \int_0^\infty f(r) dr
     where c_i are alternativaly 2/3, 4/3 except c_1 = c_mesh = 1/3
     
    parameters 
    integer:mesh
    double: rab(mesh)
            func(mesh)
            
    returns
            asum
    """

    # equivalent to fortran double precision ( > selected_real_kind(14,200))
    DP = np.double 
    
    asum = DP(0)
    r12 = DP(1)/3
    
    #precompute func(i)*rab(i)*r12
    a = func * rab * r12
    
    for i in range(1, mesh, 2):
        asum = asum + a[i-1] + 4*a[i] + a[i+1]
    return asum