#-------------------------------------------------------------------------------
# Name:        plot-expifun-theis.py
# Purpose:     Plot Theis well function W(u) with formula from Srivastava(1998)
# Source:      Practical Approximations of the Well Function
#              Rajesh Srivastava and Amado Guzman-Guzman
#              Ground Water Vol.36 No.5 september-october 1998
# Created:     15-01-2012
# Copyright:   (c) Grondwaterformules.nl
#-------------------------------------------------------------------------------
#!/usr/bin/env python
from scipy import *
from pylab import *

def fastWu(u):
    # 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 plotWu(u,Wu):

    # 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+yticklines:
        line.set_visible(True)

    # 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('W(u)')
    xlabel('u')

    # Plot the graph
    graph.loglog(u,Wu,color=orange,lw=2.0)

    # Filename
    filename = "TheisWu.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__':

    u = logspace(-15,1)
    Wu = fastWu(u)
    plotWu(u,Wu)