#-------------------------------------------------------------------------------
# Name:        wellfun-Hantush.py
# Purpose:     Calculate Hantush Well function
# Source:      Kees Maas en Ed Veling, 2010
#              Een snelle benadering van de formule van Hantush
#              Stromingen 16 (2010) nummer 1 p61-69
# Created:     7 januari 2012
# Copyright:   (c) Grondwaterformules.nl
#-------------------------------------------------------------------------------
#!/usr/bin/env python

from pylab import *
from scipy import *
from scipy import special
from scipy.special import kn as BesselK
#from scipy.special import expn as ExpInt
from math import sqrt as sqrt
from math import cosh as cosh
from math import log as ln # in python: log(x)=Ln(x)

def DeGlee(r,Q0,kD,c):
    labda = sqrt(kD*c)
    s = Q0/(2*pi*kD)*BesselK(0,r/labda)
    q = Q0*r/labda*BesselK(1,r/labda)
    return s,q,labda

def Hantush(r,t,Q0,kD,c,S):
    labda = sqrt(kD*c)
    rho = r/labda
    tau = ln(2*labda/r*t/(c*S))
    F = func_F(rho,tau)
    s = Q0/(4*pi*kD)*F
    return s

def ExpInt(n,u):

    # n wordt niet gebruikt

    # Fast approximation for Wu according to equation 7a and 7b from Srivastava(1998)
    gamma = 0.5772 # Euler-Macheroni constant

    # Wu for u<1
    u0  = where(u<1.0,u,1) # u=1 is just a dummy to make ln() work on all values
    Wu0 = log(exp(-gamma)/u0) + 0.9653*u - 0.1690*u0**2

    #Wu for u>=1
    u1 = where(u>=1.0,u,1) # u=1 is just a dummy to make ln() work on all values
    Wu1 = 1/(u1*exp(u1))*(u1+0.3575)/(u1+1.280)

    # combine Wu0 and Wu1
    Wu = where(u<1.0,Wu0,Wu1)
    return Wu

def func_F(rho,tau):
# func_F is a fast approximation of Hantush well function W

    # Calculate parameter w
    w = (ExpInt(1,rho)-BesselK(0,rho))/(ExpInt(1,rho)-ExpInt(1,rho/2))

    # Calculate F(rho,tau)
    if tau <= 0:
        F = w*ExpInt(1,rho/2*e**(-tau))-(w-1)*ExpInt(1,rho*cosh(tau))
    else:
        F = 2*BesselK(0,rho) - w*ExpInt(1,rho/2*e**tau) + (w-1)*ExpInt(1,rho*cosh(tau))

    # return calculated value of F(rho,tau)
    return F

def plot_s(r,t,s_stat,s):

    # Define colors
    orange = '#FF9900'
    blue   = '#2F64B2'

    # Create plots
    graph = subplot(111)

    # Set ticklines
    xticklines = graph.get_xticklines()
    yticklines = graph.get_yticklines()

    for line in xticklines+yticklines:
        line.set_linewidth(3)

    for line in xticklines:
        line.set_visible(False)

    # Set gridlines
    gridlines = graph.get_xgridlines()
    gridlines.extend( graph.get_ygridlines() )
    for line in gridlines:
        line.set_linestyle('-')
    grid(True)

    # Set ticklabels
    xticklabels = graph.get_xticklabels()
    yticklabels = graph.get_yticklabels()

    for label in xticklabels+yticklabels:
        label.set_color('k')
        label.set_fontsize('medium')
        label.set_visible(True)

    # Create annotation
    ylabel('verlaging (m)')
    xlabel('afstand tot de put (m)')

    # Plot lowring in time (hantush)
    for ti in range(len(t)):
        si = s[ti,:]
        graph.plot(r, -si, color=blue, lw=1.0)

    # Plot stationary lowering s_stat (De Glee)
    graph.plot(r, -s_stat, color=orange, lw=2.0)

    # Set axes limits
    # If you see a plot without lines, delete or comment next line:
    graph.set_ylim(-3, 0)

    # Filename
    filename = "Hantush.png"

    # Save figure to file
    fig = figure(1)
    mydpi = 100.0
    fig.set_dpi(mydpi)
    fig.set_size_inches(6.0,4.0)
    fig.savefig(filename,dpi=mydpi, facecolor='w', edgecolor='w')

    # Show the plot on screen
    show()


if __name__ == '__main__':

    r = arange(0.1,1000.0,0.1) # r is een numpy array
    #r = array([0.1,10.0,100.0,500.0,1000.0])
    day = 1440.0
    t = array([10.0/day,60/day,180/day,720/day,1,2])
    Q0 = 1000000.0/365.0
    kD = 800.0
    c = 1500.0
    S = 0.001

    # Calculate stationary lowering (De Glee)
    s_stat,q_stat,labda = DeGlee(r,Q0,kD,c)

    # Calculate lowering in time (Hantush)
    print "Calculation Hantush well function started..."
    s = zeros((len(t),len(r)))
    for ti in range(len(t)):
        for ri in range(len(r)):
            s[ti,ri] = Hantush(r[ri],t[ti],Q0,kD,c,S)
    print "Calculation Hantush well function finished..."

    # Plot lowering of groundwater
    plot_s(r,t,s_stat,s)