DevelopmentActivities/CMIP6/DevelopmentsCMIP6/soil_physic: Revised_HydroVert.py

import numpy as np
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.cm as cm
from mpl_toolkits.basemap import Basemap
from matplotlib.patches import Polygon

def DiagThick(dd):
    nbx=len(dd)-1
    hh=np.zeros(nbx, float)    
    for i in range(0,nbx):
        hh[i]=(dd[i]+dd[i+1])/2
    hcum=np.zeros(nbx, float)
    hcum[0]=hh[0]
    for i in range(1,nbx):
        hcum[i]=hcum[i-1]+hh[i]
    #
    return hh, hcum
#
# Compute layer thickness
#
def CompThickness(x):
    nbx=len(x)
    dx=0.0
    for i in range(1,nbx):
        dx = np.append(dx,(x[i-1]+x[i])/2.0)
    # Last layer
    dx=np.append(dx,(x[nbx-1]-dx[nbx-1])*2+dx[nbx-1])
    #
    tx=[]
    for i in range(len(dx)-1):
        tx = np.append(tx,dx[i+1]-dx[i])
    #
    return dx,tx
#
# Reading booleans
#
def str2bool(v):
  return v.lower() in ("yes", "true", "t", "1")
#
# Ploting function
#
def discretizationplot(wlev, tlev, wdepth_max, tsol_max, overalltitle):
    #
    # Compute thicknesses
    #
    wy,wz=CompThickness(wlev)
    ty,tz=CompThickness(tlev)
    #
    x=np.array([0,1])

    f, (ax1, ax2) = plt.subplots(1, 2, sharey=False)

    zmax=max([np.max(wz),np.max(tz)])
    cNorm = mpl.colors.Normalize(vmin=0, vmax=zmax)
    cmap = plt.get_cmap('Dark2')    

    ztmp=np.reshape(np.repeat(wz,2), (len(wy)-1, len(x)) )
    ax1.set_yscale('symlog')
    ax1.yaxis.set_major_formatter(mpl.ticker.ScalarFormatter())
    ax1.yaxis.set_major_formatter(mpl.ticker.FormatStrFormatter('%d'))
    ax1.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(lambda x, pos: str(int(round(x)))))
#    ax1.set_yticks([-0.02, -0.2, -2])
    datamap=ax1.pcolor(x,-wy, ztmp, zorder=0, norm=cNorm, cmap=cmap)
    # Add tine lines to show the middle of the layer where SM is positioned
    for i in range(len(wlev)):
        ax1.plot(x,np.array([-wlev[i],-wlev[i]]), color='k', linewidth=0.4)
#
    ax1.text(0.5, -wdepth_max, "depth_Wmax", fontweight='bold',ha='center',
             va='center')
    ax1.set_xlabel("Min thickness="+str(round(np.min(wz),5))+"[m]")
    ax1.set_title("nslm="+str(len(wlev)))
#
    ztmp=np.reshape(np.repeat(tz, 2), (len(ty)-1,len(x)))
    datamap = ax2.pcolor(x,-ty, ztmp, zorder=0, norm=cNorm, cmap=cmap)
    # Add tine lines to show the middle of the layer where T is positioned
    for i in range(len(tlev)):
        ax2.plot(x,np.array([-tlev[i],-tlev[i]]), color='k', linewidth=0.4)

    ax2.text(0.5, -tsol_max, "depth_max", fontweight='bold',ha='center',
                 va='center')
    ax2.set_title("ngrnd="+str(len(tlev)))
    ax2.set_xlabel("Max thickness="+str(round(np.max(tz),5))+"[m]")
#
    cb=plt.colorbar(datamap, cmap=cmap, shrink=0.5, norm=cNorm)
    cb.set_label("Layer Thickness [m]")
#
    plt.suptitle(overalltitle)
#
    pdf.savefig()
#
############################################################################################
#
# As a reminder, the levels of the hydrology as in the code today 
#
depth_Wmax=2.0
nslm=11
zz=np.zeros(nslm, float)
dz=np.zeros(nslm+1, float)
thickness=np.zeros(nslm, float)
depth=np.zeros(nslm+1, float)

for i in range(1,nslm):
    zz[i] = depth_Wmax*((2**i)-1)/((2**(nslm-1))-1)
    dz[i] = zz[i]-zz[i-1]

hh,hcum=DiagThick(dz)
print "== Old depth of the hydrology for reference =========================="
print "zz = depth of nodes"
print zz
print "dz = internode distance"
print dz
print "hh = soil layer thickness"
print hh
print "hcum = depth of soil layer bottom"
print hcum
print "======================================================================"
print "  "
print "  "
#
###############################################################################
#
# Ask the user for input to configure the vertical levels
#
# New Methode for computing levels with the new parameters to control.
#
# All units in meters
#
# Maximum depth of the soil thermodynamics
depth_max = 8.0
try :
    depth_max=float(raw_input("Maximum depth of soil : depth_max [8.0m] >"))
except ValueError:
    depth_max = 8.0
#
# Maximum depth of soil moisture
depth_Wmax=2.0
try :
    depth_Wmax=float(raw_input("Depth of hydrology (smaller than depth_max) : depth_Wmax [2.0m] >"))
except ValueError:
    depth_Wmax=2.0
if ( depth_Wmax > depth_max) :
    print "ERROR : depth_Wmax needs to be smaller than depth_max."
    print "Correction : depth_Wmax set to depth_max"
    depth_Wmax = depth_max
#
# Thickness of first Layer 
depth_toptickness = 9.77517107e-04
try :
    depth_toptickness=float(raw_input("Depth of top hydrology layer : depth_toptickness [9.77517107e-04 m] >"))
except ValueError:
    depth_toptickness = 9.77517107e-04
# The constant layer thickness starts at :
depth_cstthickness=depth_Wmax
try :
    depth_cstthickness=float(raw_input("Depth at which constant layer thickness start (smaller than depth_Wmax/2)  : depth_cstthickness [depth_Wmax] >"))
except ValueError:
    depth_cstthickness = depth_Wmax
if ( (depth_cstthickness != depth_Wmax) & (depth_cstthickness > depth_Wmax/2.0) ) :
    print "ERROR : depth_cstthickness needs to be smaller than depth_Wmax/2.0"
    print "Correction : Constant thickness disabled"
    depth_cstthickness = depth_Wmax
#
# Do we refine the layers as we reach the depth_Wmax
refinebottom = False
try :
    refinebottom=str2bool(raw_input("Refinement at the bottom (symetrical to the top refinement): refinebottom [False] >"))
except ValueError:
    refinebottom=False
#
# Depth at which we resume geometrical increases
depth_geom=depth_Wmax
try :
    depth_geom=float(raw_input("Depth at which geometrical increase starts again (larger than depth_Wmax)  : depth_geom [depth_Wmax] >"))
except ValueError:
    depth_geom=depth_Wmax
if ( depth_geom < depth_Wmax ) :
    print "ERROR : depth_geom needs to be larger than depth_Wmax"
    print "Correction : setting depth_geom to depth_Wmax"

    depth_geom=depth_Wmax
#
# Ratio of geometrical serie below depth_geom 
ratio_geom_below=2
try :
    ratio_geom_below=float(raw_input("Ratio of geometrical serie below depth_geom : ratio_geom_below [2] >"))
except ValueError:
    ratio_geom_below=2
#
###################################################################################################
#
#  Start computing the new levels
#
# It is coded FORTRAN style so that translations is easier
#
nltmp=500
ztmp=np.zeros(nltmp, float)
dtmp=np.zeros(nltmp+1, float)
for i in range(1,nltmp):
    if ( ztmp[i-1] < depth_cstthickness ) :
        ztmp[i] = depth_toptickness*2.0*((2**i)-1)
        dtmp[i] = ztmp[i]-ztmp[i-1]
        igeo=i
        zgeoend=ztmp[i]
        dgeoend=dtmp[i]
    else :
        ztmp[i] = ztmp[i-1]+dtmp[i-1]
        dtmp[i] = dtmp[i-1]
#
nbrefine=0
drefine=np.zeros(nltmp, float)
if (refinebottom) :
    #
    # Compute parameters for the constant increment interval before refining,
    # If needed !
    cstint=depth_Wmax-(2.0*zgeoend)
    nbcst=max(int(np.modf(cstint/dgeoend)[1]),0)
    if ( nbcst > 0 ) :
        dcst=cstint/nbcst
    else :
        dcst=dgeoend
    #
    # If we have to add constant increments
    #
    if ( nbcst > 0 ) :
        # Add linear levels
        for i in range(igeo+1,igeo+nbcst+1) :
            ztmp[i] = ztmp[i-1]+dcst
            dtmp[i] = dcst
        # Refine the levels toward the bottom
        for i in range(igeo+nbcst+1,2*igeo+nbcst+1) :
            irev=(2*igeo+nbcst+1)-i
            ztmp[i] = ztmp[i-1]+dtmp[irev]
            dtmp[i] = ztmp[i]-ztmp[i-1]
            drefine[nbrefine]=dtmp[i]
            nbrefine=nbrefine+1
    # Without constant increments
    else :
        igeo=np.argmin(np.abs(ztmp-(depth_Wmax/2.0)))
        ztmp[igeo]=depth_Wmax/2.0
        dtmp[igeo] = ztmp[igeo]-ztmp[igeo-1]
        for i in range(igeo+1,2*igeo+1) :
            irev=(2*igeo+1)-i
            ztmp[i] = ztmp[i-1]+dtmp[irev]
            dtmp[i] = ztmp[i]-ztmp[i-1]
            drefine[nbrefine]=dtmp[i]
            nbrefine=nbrefine+1
#
nslm=np.argmin(np.abs(ztmp-depth_Wmax))+1
#
# Extract now the values we need to reach depth_Wmax
#
zz=ztmp[0:nslm]
dz=dtmp[0:nslm+1]
dz[nslm]=0.0
hh,hcum = DiagThick(dz)
print "== New values ========================================================"
print "Number of layers in hydrol = ", nslm
print "zz = depth of nodes"
print zz
print "dz = internode distance"
print dz
print "hh = soil layer thickness"
print hh
print "hcum = depth of soil layer bottom"
print hcum
print "======================================================================"
print "  "
print "  "
#
# Extension for the thermodynamics below depth_Wmax
#
ztmp=np.zeros(nltmp, float)
# Exception for the first layer. It is at the top in the hydrology and in
# the middle for temlperature.
ztmp[0]=depth_toptickness/2
for i in range(1,nslm):
    ztmp[i]=zz[i]
# 
# Exception for the last layer where again temperature needs to be in
# the middle of the layer. Also add a layer with the same thickness
ztmp[nslm-1]=ztmp[nslm-1]-hh[nslm-1]/2.0
ztmp[nslm]=ztmp[nslm-1]+hh[nslm-1]
nslm=nslm+1
#
# If we have created a refined region at the bottom of the soil moisture. We need to
# unwinde it for temperature below depth_Wmax
#
if ( nbrefine > 1 ) :
    for i in range(nslm,nslm+nbrefine,1):
        iref=(nbrefine-1)-(i-nslm)
        ztmp[i]=ztmp[i-1]+drefine[iref]
#
# New nslm
#
nslm=nslm+nbrefine
#
for i in range(nslm,nltmp):
    if ( ztmp[i-1] < depth_geom ) :
        ratio = 1.0
    else :
        ratio = ratio_geom_below
    ztmp[i]=ztmp[i-1]+ratio*(ztmp[i-1]-ztmp[i-2])
#        
ngrnd=np.argmin(np.abs(ztmp-depth_max))+1
#
tdepth=np.zeros(ngrnd, float)
tdepth=ztmp[0:ngrnd]
print "Number of layers in thermosoil = ", ngrnd
print "tdepth=",tdepth
print "======================================================================"
#
# Do the plot
#
pdf = PdfPages('VerticalLayers.pdf')
wdepth=np.copy(zz)
wdepth[0]=(zz[1]-zz[0])/2
discretizationplot(wdepth, tdepth, depth_Wmax, depth_max, "Mixing water and temperature")
pdf.close()