# These are the van Genuchten (1980) equations
# The input is matric potential, psi and the hydraulic parameters.
# psi must be sent in as a numpy array.
# The pars variable is like a MATLAB structure.
import numpy as np

def thetaFun(psi,pars):
    if psi>=0.:
        Se = 1.
    else:
        Se=(1+abs(psi*pars['alpha'])**pars['n'])**(-pars['m'])
    return pars['thetaR']+(pars['thetaS']-pars['thetaR'])*Se

thetaFun=np.vectorize(thetaFun)

def CFun(psi,pars):
    if psi>=0.:
        Se=1.
    else:
        Se=(1+abs(psi*pars['alpha'])**pars['n'])**(-pars['m'])
    dSedh=pars['alpha']*pars['m']/(1-pars['m'])*Se**(1/pars['m'])*(1-Se**(1/pars['m']))**pars['m']
    return Se*pars['Ss']+(pars['thetaS']-pars['thetaR'])*dSedh

CFun = np.vectorize(CFun)

def KFun(psi,pars):
    if psi>=0.:
        Se=1.
    else:
        Se=(1+abs(psi*pars['alpha'])**pars['n'])**(-pars['m'])
    return pars['Ks']*Se**pars['neta']*(1-(1-Se**(1/pars['m']))**pars['m'])**2
  
KFun = np.vectorize(KFun)

def setpars():
  pars={}
  pars['thetaR']=float(raw_input("thetaR = "))
  pars['thetaS']=float(raw_input("thetaS = "))
  pars['alpha']=float(raw_input("alpha = "))
  pars['n']=float(raw_input("n = "))
  pars['m']=1-1/pars['n']
  pars['Ks']=float(raw_input("Ks = "))
  pars['neta']=float(raw_input("neta = "))
  pars['Ss']=float(raw_input("Ss = "))
  return pars
  
def PlotProps(pars):
  import numpy as np
  import pylab as pl
  import vanGenuchten as vg
  psi=np.linspace(-10,2,200)
  pl.figure
  pl.subplot(3,1,1)
  pl.plot(psi,vg.thetaFun(psi,pars))
  pl.ylabel(r'$\theta(\psi) [-]$')
  pl.subplot(3,1,2)
  pl.plot(psi,vg.CFun(psi,pars))
  pl.ylabel(r'$C(\psi) [1/m]$')
  pl.subplot(3,1,3)
  pl.plot(psi,vg.KFun(psi,pars))
  pl.xlabel(r'$\psi [m]$')
  pl.ylabel(r'$K(\psi) [m/d]$')
  #pl.show()
  
def HygieneSandstone():
  pars={}
  pars['thetaR']=0.153
  pars['thetaS']=0.25
  pars['alpha']=0.79
  pars['n']=10.4
  pars['m']=1-1/pars['n']
  pars['Ks']=1.08
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars
  
def TouchetSiltLoam():
  pars={}
  pars['thetaR']=0.19
  pars['thetaS']=0.469
  pars['alpha']=0.5
  pars['n']=7.09
  pars['m']=1-1/pars['n']
  pars['Ks']=3.03
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars

def SiltLoamGE3():
  pars={}
  pars['thetaR']=0.131
  pars['thetaS']=0.396
  pars['alpha']=0.423
  pars['n']=2.06
  pars['m']=1-1/pars['n']
  pars['Ks']=0.0496
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars
  
def GuelphLoamDrying():
  pars={}
  pars['thetaR']=0.218
  pars['thetaS']=0.520
  pars['alpha']=1.15
  pars['n']=2.03
  pars['m']=1-1/pars['n']
  pars['Ks']=0.316
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars
  
def GuelphLoamWetting():
  pars={}
  pars['thetaR']=0.218
  pars['thetaS']=0.434
  pars['alpha']=2.0
  pars['n']=2.76
  pars['m']=1-1/pars['n']
  pars['Ks']=0.316
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars
  
def BeitNetofaClay():
  pars={}
  pars['thetaR']=0.
  pars['thetaS']=0.446
  pars['alpha']=0.152
  pars['n']=1.17
  pars['m']=1-1/pars['n']
  pars['Ks']=0.00082
  pars['neta']=0.5
  pars['Ss']=0.000001
  return pars