From d0418738bba6412560d37f01c26c869ab04dc39c Mon Sep 17 00:00:00 2001 From: Charles Doutriaux Date: Fri, 10 Jun 2016 08:43:08 -0700 Subject: [PATCH] fix #1776 let user pass an axis argument, rather than always picking z level, that allows for non cf use also pep8ed the file --- Packages/cdutil/Lib/vertical.py | 292 ++++++++++++++++++-------------- testing/cdutil/test_vert.py | 28 +-- 2 files changed, 179 insertions(+), 141 deletions(-) diff --git a/Packages/cdutil/Lib/vertical.py b/Packages/cdutil/Lib/vertical.py index 9d6df0ae8f..167b7d4cde 100644 --- a/Packages/cdutil/Lib/vertical.py +++ b/Packages/cdutil/Lib/vertical.py @@ -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 @@ -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 diff --git a/testing/cdutil/test_vert.py b/testing/cdutil/test_vert.py index 591ac61306..54f9c625a0 100644 --- a/testing/cdutil/test_vert.py +++ b/testing/cdutil/test_vert.py @@ -2,22 +2,24 @@ # Adapted for numpy/ma/cdms2 by convertcdms.py -import cdutil,cdat_info +import cdutil +import cdat_info +import numpy import cdms2 import os bg = 0 -f = cdms2.open(os.path.join(cdat_info.get_sampledata_path(),'vertical.nc')) -Ps=f('PS') -U=f('U') -B=f('hybm') -A=f('hyam') -Po=f('variable_2') -P=cdutil.reconstructPressureFromHybrid(Ps,A,B,Po) +f = cdms2.open(os.path.join(cdat_info.get_sampledata_path(), 'vertical.nc')) +Ps = f('PS') +U = f('U') +B = f('hybm') +A = f('hyam') +Po = f('variable_2') +P = cdutil.reconstructPressureFromHybrid(Ps, A, B, Po) -U2=cdutil.logLinearInterpolation(U,P) - -#x=vcs.init() -#x.plot(U2,bg=bg) -#raw_input() +U2 = cdutil.logLinearInterpolation(U, P) +U2b = cdutil.logLinearInterpolation(U, P, axis='0') +assert(numpy.ma.allclose(U2, U2b)) +U2b = cdutil.logLinearInterpolation(U, P, axis='(lev)') +assert(numpy.ma.allclose(U2, U2b))