"""
Supplementary material for
Comments on the Voigt function approximation in the Astropy and SpectraPlot.com packages
Journal of Quantitative Spectroscopy and Radiative Transfer, in press, 2018
http://dx.doi.org/10.1016/j.jqsrt.2018.03.019

Franz Schreier;  DLR-IMF Oberpfaffenhofen

Execution of this script (e.g. %run mcLean.py from inside IPython) should produce the figures.
"""

from math import pi, sqrt, cos
recSqrtPi = 1./sqrt(pi)

import numpy as np
from scipy.special import wofz

# depending on your Python setup, the following import might not be required
from matplotlib.pyplot import figure, rc, subplot, plot, semilogy, xlabel, ylabel, legend, title, twinx, xticks, ylim, show, savefig

####################################################################################################################################

L32 = sqrt(32./sqrt(2.)) #  sqrt(N)/2^(1/4) = 4.1195342878142354

a32 = np.array([ 2.6837530678616557e+00,  # a0 = L/sqrtPi
                 2.5722534081245696e+00,   2.2635372999002676e+00,   1.8256696296324824e+00,   1.3455441692345453e+00,
                 9.0192548936480144e-01,   5.4601397206393498e-01,   2.9544451071508926e-01,   1.4060716226893769e-01,
                 5.7304403529837900e-02,   1.9006155784845689e-02,   4.5195411053501429e-03,   3.9259136070122748e-04,
                -2.4532980269928922e-04,  -1.3075449254548613e-04,  -2.1409619200870880e-05,   6.8210319440412389e-06,
                 4.4015317319048931e-06,   4.2558331390536872e-07,  -4.1840763666294341e-07,  -1.4813078891201116e-07,
                 2.2930439569075392e-08,   2.3797557105844622e-08,   8.1248960947953431e-10,  -3.2080150458594088e-09,
                -5.2310170266050247e-10,   4.1537465934749353e-10,   1.1658312885903929e-10,  -5.5441820344468828e-11,
                -2.1542618451370239e-11,   8.0314997274316680e-12,   3.7424975634801558e-12,  -1.3031797863050087e-12])

def weideman32a (x,y):
	""" Complex error function --- J.A.C. Weideman: SIAM J. Num. Anal. 31 1497-1518 (1994); equation (38.I) and table I. """
	iz  = 1j*x - y
	lpiz    = L32 + iz  # wL - y + x*complex(0.,1.)
	lmiz    = L32 - iz  # wL + y - x*complex(0.,1.)
	recLmiZ = 1.0 / lmiz
	Z       = lpiz * recLmiZ
	polynom = a32[1]+(a32[2]+(a32[3]+(a32[4]+(a32[5]+(a32[6]+(a32[7]+(a32[8]+(a32[9]+(a32[10]+(a32[11]+(a32[12]+(a32[13]+(a32[14]+(a32[15]+(a32[16]+(a32[17]+(a32[18]+(a32[19]+(a32[20]+(a32[21]+(a32[22]+(a32[23]+(a32[24]+(a32[25]+(a32[26]+(a32[27]+(a32[28]+(a32[29]+(a32[30]+(a32[31]+a32[32]*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z)*Z
	w       = (recSqrtPi  +  2.0 * polynom * recLmiZ)  *  recLmiZ
	return w


################################################################################################################################################################

def hum1wei32a (x,y):
 	""" Complex error function using Humlicek's and Weideman's rational approximation:
	    |x|+y>15:  Humlicek (JQSRT, 1982) rational approximation for region I;
	    else:      J.A.C. Weideman (SIAM-NA 1994); equation (38.I) and table I. """
	t = y - 1j*x
	w = t * recSqrtPi / (0.5 + t*t)  # Humlicek (1982) approx 1 for |x|+y>15
	# any points inside?
	mask = abs(x)+y<15.0
	if mask.any():
		np.place(w, mask, weideman32a(x,y))
	return w


####################################################################################################################################

def gauss (x):
	""" Gauss function as limit of Voigt function, i.e. with peak value=1.0 and normalized to sqrt(Pi). """
	# don't compute exponential for very large numbers > log(1e308)=709.
	# (in contrast to python.math.exp the numpy exp overflows!!!)
	return  np.exp(-np.where(x>100.,100.0,x)**2)


############################################################  Figure 1  ########################################################################################

def mcLean (x,y):
    """ Algorithm for the computation taken from
    McLean, A. B., Mitchell, C. E. J. & Swanston, D. M. Implementation of an
    efficient analytical approximation to the Voigt function for photoemission
    lineshape analysis. Journal of Electron Spectroscopy and Related Phenomena
    69, 125-132 (1994)
    """
    abcd = np.array([
        [-1.2150, -1.3509, -1.2150, -1.3509],  # A
        [1.2359, 0.3786, -1.2359, -0.3786],    # B
        [-0.3085, 0.5906, -0.3085, 0.5906],    # C
        [0.0210, -1.1858, -0.0210, 1.1858]])   # D
    A, B, C, D = abcd
    vgt = (C[0] * (y - A[0]) + D[0] * (x - B[0]))/(((y - A[0]) ** 2 + (x - B[0]) ** 2)) \
         +(C[1] * (y - A[1]) + D[1] * (x - B[1]))/(((y - A[1]) ** 2 + (x - B[1]) ** 2)) \
         +(C[2] * (y - A[2]) + D[2] * (x - B[2]))/(((y - A[2]) ** 2 + (x - B[2]) ** 2)) \
         +(C[3] * (y - A[3]) + D[3] * (x - B[3]))/(((y - A[3]) ** 2 + (x - B[3]) ** 2))

    return vgt


############################################################  Figure 2  ########################################################################################

x = np.linspace(0,10.,201)

rc('figure.subplot', top=0.975,bottom=0.04,left=0.125,right=0.86,hspace=0.15)
rc('axes',grid=0)
figure('wofz-mcLean', figsize=(5.,12))

y=1.0
kTest= mcLean(x,y);  khw = hum1wei32a(x,y).real;  kRef = wofz(x+1j*y).real
subplot(411)
plot (x, kRef, 'b--',  x, kTest, 'r')
xticks(visible=0);  ylabel('$K(x,y)$');  title ('y = %g' % y)
twinx();  semilogy (x, abs(kTest-kRef)/kRef, 'g', x, abs(khw-kRef)/kRef, 'c')
xticks(visible=0);  ylabel(r'$|\Delta K|/K_\mathrm{ref}$')

y=0.1
kTest= mcLean(x,y);  khw = hum1wei32a(x,y).real;  kRef = wofz(x+1j*y).real
subplot(412)
semilogy (x, kRef, 'b--',  x, kTest, 'r')
xticks(visible=0);  ylabel('$K(x,y)$');  title ('y = %g' % y)
twinx();  semilogy (x, abs(kTest-kRef)/kRef, 'g', x, abs(khw-kRef)/kRef, 'c')
xticks(visible=0);  ylabel(r'$|\Delta K|/K_{\rm ref}$')

y=0.01
kTest= mcLean(x,y);  khw = hum1wei32a(x,y).real;  kRef = wofz(x+1j*y).real
subplot(413)
semilogy (x, kRef, 'b--',  x, kTest, 'r')
xticks(visible=0);  ylabel('$K(x,y)$');  title ('y = %g' % y)
legend(['wofz', 'McLean'], loc=(0.55,0.45))
twinx();  semilogy (x, abs(kTest-kRef)/kRef, 'g', x, abs(khw-kRef)/kRef, 'c')
xticks(visible=0);  ylabel(r'$|\Delta K|/K_{\rm ref}$')

y=0.001
kTest= mcLean(x,y);  khw = hum1wei32a(x,y).real;  kRef = wofz(x+1j*y).real
subplot(414)
semilogy (x, kRef, 'b--',  x, kTest, 'r')
xlabel('$x$');  ylabel('$K(x,y)$');  title ('y = %g' % y)
twinx();  semilogy (x, abs(kTest-kRef)/kRef, 'g', x, abs(khw-kRef)/kRef, 'c')
legend(['error McLean', 'error Humlicek-Weideman'], loc=(0.25,0.55), fontsize='medium')
xlabel('$x$');  ylabel(r'$|\Delta K|/K_{\rm ref}$')

show()  # possibly not required (depending on your settings)
#savefig('wofz-mcLean.pdf')


################################################################################################################################################################

rc('figure.subplot', top=0.97,bottom=0.11,left=0.15,right=0.95)
figure('voigt-gauss', figsize=(5.,4.))
y=0.001
semilogy (x, hum1wei32a(x,y), 'r', label='Voigt:  $y$ = %.0e' % y)
y=0.0001
semilogy (x, hum1wei32a(x,y), 'g', label='Voigt:  $y$ = %.0e' % y)
y=0.00001
semilogy (x, hum1wei32a(x,y), 'b', label='Voigt:  $y$ = %.0e' % y)
semilogy (x, gauss(x), 'k:', label='Gauss')
xlabel('$x$');  ylabel('$K(x,y)$')
ylim(ymin=1e-10)
legend(loc=1)

show()
#savefig('voigt-gauss.pdf')