Skip to content

Commit

Permalink
fix #1776
Browse files Browse the repository at this point in the history 
let user pass an axis argument, rather than always picking z level, that allows for non cf use
also pep8ed the file
  • Loading branch information
@doutriaux1
doutriaux1 committed Jun 10, 2016
1 parent 94c9b59 commit 9b4ce01
Showing 1 changed file with 164 additions and 128 deletions.
292 changes: 164 additions & 128 deletions Lib/vertical.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,19 @@
import cdms2
import numpy
import cdat_info
def reconstructPressureFromHybrid(ps,A,B,Po):


def reconstructPressureFromHybrid(ps, A, B, Po):
"""
Reconstruct the Pressure field on sigma levels, from the surface pressure
Input
Ps : Surface pressure
A,B,Po: Hybrid Convertion Coefficients, such as: p=B.ps+A.Po
Ps: surface pressure
B,A are 1D : sigma levels
Po and Ps must have same units
Output
Pressure field
Such as P=B*Ps+A*Po
Expand All @@ -23,200 +25,234 @@ def reconstructPressureFromHybrid(ps,A,B,Po):
P=reconstructPressureFromHybrid(ps,A,B,Po)
"""
# Compute the pressure for the sigma levels
cdat_info.pingPCMDIdb("cdat","cdutil.vertical.reconstructPressureFromHybrid")
ps,B=genutil.grower(ps,B)
ps,A=genutil.grower(ps,A)
p=ps*B
p=p+A*Po
cdat_info.pingPCMDIdb(
"cdat",
"cdutil.vertical.reconstructPressureFromHybrid")
ps, B = genutil.grower(ps, B)
ps, A = genutil.grower(ps, A)
p = ps * B
p = p + A * Po
p.setAxisList(ps.getAxisList())
p.id='P'
p.id = 'P'
try:
p.units=ps.units
p.units = ps.units
except:
pass
t=p.getTime()
pass
t = p.getTime()
if not t is None:
p=p(order='tz...')
p = p(order='tz...')
else:
p=p(order='z...')
p = p(order='z...')
return p

def linearInterpolation(A,I,levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000], status=None):


def linearInterpolation(
A, I, levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000,
30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000], status=None, axis='z'):
"""
Linear interpolation
to interpolate a field from some levels to another set of levels
Value below "surface" are masked
Input
A : array to interpolate
I : interpolation field (usually Pressure or depth) from TOP (level 0) to BOTTOM (last level), i.e P value going up with each level
levels : levels to interplate to (same units as I), default levels are:[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000]
axis: axis over which to do the linear interpolation, default is 'z', accepted: '1' '(myaxis)'
I and levels must have same units
Output
array on new levels (levels)
Examples:
A=interpolate(A,I,levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000])
"""
cdat_info.pingPCMDIdb("cdat","cdutil.vertical.linearInterpolation")

cdat_info.pingPCMDIdb("cdat", "cdutil.vertical.linearInterpolation")
try:
nlev=len(levels) # Number of pressure levels
nlev = len(levels) # Number of pressure levels
except:
nlev=1 # if only one level len(levels) would breaks
levels=[levels,]
order=A.getOrder()
A=A(order='z...')
I=I(order='z...')
sh=list(I.shape)
nsigma=sh[0] #number of sigma levels
sh[0]=nlev
t=MV2.zeros(sh,typecode=MV2.float32)
sh2=I[0].shape
prev=-1
for ilev in range(nlev): # loop through pressure levels
nlev = 1 # if only one level len(levels) would breaks
levels = [levels, ]
order = A.getOrder()
A = A(order='%s...' % axis)
I = I(order='%s...' % axis)
sh = list(I.shape)
nsigma = sh[0] # number of sigma levels
sh[0] = nlev
t = MV2.zeros(sh, typecode=MV2.float32)
sh2 = I[0].shape
prev = -1
for ilev in range(nlev): # loop through pressure levels
if status is not None:
prev=genutil.statusbar(ilev,nlev-1.,prev)
lev=levels[ilev] # get value for the level
Iabv=MV2.ones(sh2,MV2.float)
Aabv=-1*Iabv # Array on sigma level Above
Abel=-1*Iabv # Array on sigma level Below
Ibel=-1*Iabv # Pressure on sigma level Below
Iabv=-1*Iabv # Pressure on sigma level Above
Ieq=MV2.masked_equal(Iabv,-1) # Area where Pressure == levels
for i in range(1,nsigma): # loop from second sigma level to last one
a = MV2.greater_equal(I[i], lev) # Where is the pressure greater than lev
b = MV2.less_equal(I[i-1],lev) # Where is the pressure less than lev
prev = genutil.statusbar(ilev, nlev - 1., prev)
lev = levels[ilev] # get value for the level
Iabv = MV2.ones(sh2, MV2.float)
Aabv = -1 * Iabv # Array on sigma level Above
Abel = -1 * Iabv # Array on sigma level Below
Ibel = -1 * Iabv # Pressure on sigma level Below
Iabv = -1 * Iabv # Pressure on sigma level Above
Ieq = MV2.masked_equal(Iabv, -1) # Area where Pressure == levels
for i in range(1, nsigma): # loop from second sigma level to last one
a = MV2.greater_equal(
I[i],
lev) # Where is the pressure greater than lev
b = MV2.less_equal(
I[i - 1],
lev) # Where is the pressure less than lev
# Now looks if the pressure level is in between the 2 sigma levels
# If yes, sets Iabv, Ibel and Aabv, Abel
a=MV2.logical_and(a,b)
Iabv=MV2.where(a,I[i],Iabv) # Pressure on sigma level Above
Aabv=MV2.where(a,A[i],Aabv) # Array on sigma level Above
Ibel=MV2.where(a,I[i-1],Ibel) # Pressure on sigma level Below
Abel=MV2.where(a,A[i-1],Abel) # Array on sigma level Below
Ieq= MV2.where(MV2.equal(I[i],lev),A[i],Ieq)

val=MV2.masked_where(MV2.equal(Ibel,-1.),numpy.ones(Ibel.shape)*lev) # set to missing value if no data below lev if there is

tl=(val-Ibel)/(Iabv-Ibel)*(Aabv-Abel)+Abel # Interpolation
a = MV2.logical_and(a, b)
Iabv = MV2.where(a, I[i], Iabv) # Pressure on sigma level Above
Aabv = MV2.where(a, A[i], Aabv) # Array on sigma level Above
Ibel = MV2.where(
a,
I[i - 1],
Ibel) # Pressure on sigma level Below
Abel = MV2.where(a, A[i - 1], Abel) # Array on sigma level Below
Ieq = MV2.where(MV2.equal(I[i], lev), A[i], Ieq)

val = MV2.masked_where(
MV2.equal(Ibel, -1.), numpy.ones(Ibel.shape) * lev)
# set to missing value if no data below lev if
# there is

tl = (val - Ibel) / (Iabv - Ibel) * \
(Aabv - Abel) + Abel # Interpolation
if ((Ieq.mask is None) or (Ieq.mask is MV2.nomask)):
tl=Ieq
tl = Ieq
else:
tl=MV2.where(1-Ieq.mask,Ieq,tl)
t[ilev]=tl.astype(MV2.float32)
tl = MV2.where(1 - Ieq.mask, Ieq, tl)
t[ilev] = tl.astype(MV2.float32)

ax=A.getAxisList()
autobnds=cdms2.getAutoBounds()
ax = A.getAxisList()
autobnds = cdms2.getAutoBounds()
cdms2.setAutoBounds('off')
lvl=cdms2.createAxis(MV2.array(levels).filled())
lvl = cdms2.createAxis(MV2.array(levels).filled())
cdms2.setAutoBounds(autobnds)
try:
lvl.units=I.units
lvl.units = I.units
except:
pass
lvl.id='plev'
lvl.id = 'plev'

try:
t.units=I.units
t.units = I.units
except:
pass
ax[0]=lvl
pass

ax[0] = lvl
t.setAxisList(ax)
t.id=A.id
t.id = A.id
for att in A.listattributes():
setattr(t,att,getattr(A,att))
setattr(t, att, getattr(A, att))
return t(order=order)

def logLinearInterpolation(A,P,levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000],status=None):

def logLinearInterpolation(
A, P, levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000,
30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000], status=None, axis='z'):
"""
Log-linear interpolation
to convert a field from sigma levels to pressure levels
Value below surface are masked
Input
A : array on sigma levels
P : pressure field from TOP (level 0) to BOTTOM (last level)
A : array on sigma levels
P : pressure field from TOP (level 0) to BOTTOM (last level)
levels : pressure levels to interplate to (same units as P), default levels are:[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000]
axis: axis over which to do the linear interpolation, default is 'z', accepted: '1' '(myaxis)'
P and levels must have same units
Output
array on pressure levels (levels)
Examples:
A=logLinearInterpolation(A,P),levels=[100000, 92500, 85000, 70000, 60000, 50000, 40000, 30000, 25000, 20000, 15000, 10000, 7000, 5000, 3000, 2000, 1000])
"""
cdat_info.pingPCMDIdb("cdat","cdutil.vertical.logLinearInterpolation")

cdat_info.pingPCMDIdb("cdat", "cdutil.vertical.logLinearInterpolation")
try:
nlev=len(levels) # Number of pressure levels
nlev = len(levels) # Number of pressure levels
except:
nlev=1 # if only one level len(levels) would breaks
levels=[levels,]
order=A.getOrder()
A=A(order='z...')
P=P(order='z...')
sh=list(P.shape)
nsigma=sh[0] #number of sigma levels
sh[0]=nlev
t=MV2.zeros(sh,typecode=MV2.float32)
sh2=P[0].shape
prev=-1
for ilev in range(nlev): # loop through pressure levels
nlev = 1 # if only one level len(levels) would breaks
levels = [levels, ]
order = A.getOrder()
A = A(order='%s...' % axis)
P = P(order='%s...' % axis)
sh = list(P.shape)
nsigma = sh[0] # number of sigma levels
sh[0] = nlev
t = MV2.zeros(sh, typecode=MV2.float32)
sh2 = P[0].shape
prev = -1
for ilev in range(nlev): # loop through pressure levels
if status is not None:
prev=genutil.statusbar(ilev,nlev-1.,prev)
lev=levels[ilev] # get value for the level
Pabv=MV2.ones(sh2,MV2.float)
Aabv=-1*Pabv # Array on sigma level Above
Abel=-1*Pabv # Array on sigma level Below
Pbel=-1*Pabv # Pressure on sigma level Below
Pabv=-1*Pabv # Pressure on sigma level Above
Peq=MV2.masked_equal(Pabv,-1) # Area where Pressure == levels
for i in range(1,nsigma): # loop from second sigma level to last one
a=MV2.greater_equal(P[i], lev) # Where is the pressure greater than lev
b= MV2.less_equal(P[i-1],lev) # Where is the pressure less than lev
prev = genutil.statusbar(ilev, nlev - 1., prev)
lev = levels[ilev] # get value for the level
Pabv = MV2.ones(sh2, MV2.float)
Aabv = -1 * Pabv # Array on sigma level Above
Abel = -1 * Pabv # Array on sigma level Below
Pbel = -1 * Pabv # Pressure on sigma level Below
Pabv = -1 * Pabv # Pressure on sigma level Above
Peq = MV2.masked_equal(Pabv, -1) # Area where Pressure == levels
for i in range(1, nsigma): # loop from second sigma level to last one
a = MV2.greater_equal(
P[i],
lev) # Where is the pressure greater than lev
b = MV2.less_equal(
P[i - 1],
lev) # Where is the pressure less than lev
# Now looks if the pressure level is in between the 2 sigma levels
# If yes, sets Pabv, Pbel and Aabv, Abel
a=MV2.logical_and(a,b)
Pabv=MV2.where(a,P[i],Pabv) # Pressure on sigma level Above
Aabv=MV2.where(a,A[i],Aabv) # Array on sigma level Above
Pbel=MV2.where(a,P[i-1],Pbel) # Pressure on sigma level Below
Abel=MV2.where(a,A[i-1],Abel) # Array on sigma level Below
Peq= MV2.where(MV2.equal(P[i],lev),A[i],Peq)

val=MV2.masked_where(MV2.equal(Pbel,-1),numpy.ones(Pbel.shape)*lev) # set to missing value if no data below lev if there is

tl=MV2.log(val/Pbel)/MV2.log(Pabv/Pbel)*(Aabv-Abel)+Abel # Interpolation
a = MV2.logical_and(a, b)
Pabv = MV2.where(a, P[i], Pabv) # Pressure on sigma level Above
Aabv = MV2.where(a, A[i], Aabv) # Array on sigma level Above
Pbel = MV2.where(
a,
P[i - 1],
Pbel) # Pressure on sigma level Below
Abel = MV2.where(a, A[i - 1], Abel) # Array on sigma level Below
Peq = MV2.where(MV2.equal(P[i], lev), A[i], Peq)

val = MV2.masked_where(
MV2.equal(Pbel, -1), numpy.ones(Pbel.shape) * lev)
# set to missing value if no data below lev if
# there is

tl = MV2.log(
val / Pbel) / MV2.log(
Pabv / Pbel) * (
Aabv - Abel) + Abel # Interpolation
if ((Peq.mask is None) or (Peq.mask is MV2.nomask)):
tl=Peq
tl = Peq
else:
tl=MV2.where(1-Peq.mask,Peq,tl)
t[ilev]=tl.astype(MV2.float32)
ax=A.getAxisList()
autobnds=cdms2.getAutoBounds()
tl = MV2.where(1 - Peq.mask, Peq, tl)
t[ilev] = tl.astype(MV2.float32)

ax = A.getAxisList()
autobnds = cdms2.getAutoBounds()
cdms2.setAutoBounds('off')
lvl=cdms2.createAxis(MV2.array(levels).filled())
lvl = cdms2.createAxis(MV2.array(levels).filled())
cdms2.setAutoBounds(autobnds)
try:
lvl.units=P.units
lvl.units = P.units
except:
pass
lvl.id='plev'
lvl.id = 'plev'

try:
t.units=P.units
t.units = P.units
except:
pass
ax[0]=lvl
pass

ax[0] = lvl
t.setAxisList(ax)
t.id=A.id
t.id = A.id
for att in A.listattributes():
setattr(t,att,getattr(A,att))
setattr(t, att, getattr(A, att))
return t(order=order)
sigma2Pressure=logLinearInterpolation

sigma2Pressure = logLinearInterpolation

0 comments on commit 9b4ce01

Please sign in to comment.