source: lmdz_wrf/trunk/tools/diagnostics.py @ 422

Last change on this file since 422 was 413, checked in by lfita, 10 years ago

Changing to 'e.g.'

File size: 35.9 KB
Line 
1# Pthon script to comput diagnostics
2# L. Fita, LMD. CNR, UPMC-Jussieu, Paris, France
3# File diagnostics.inf provides the combination of variables to get the desired diagnostic
4#
5## e.g. # diagnostics.py -d 'Time@time,bottom_top@ZNU,south_north@XLAT,west_east@XLONG' -v 'clt|CLDFRA,cllmh|CLDFRA@WRFp,RAINTOT|RAINC@RAINNC@XTIME' -f WRF_LMDZ/NPv31/wrfout_d01_1980-03-01_00:00:00
6## e.g. # diagnostics.py -f /home/lluis/PY/diagnostics.inf -d variable_combo -v WRFprc
7
8from optparse import OptionParser
9import numpy as np
10from netCDF4 import Dataset as NetCDFFile
11import os
12import re
13import nc_var_tools as ncvar
14
15main = 'diagnostics.py'
16errormsg = 'ERROR -- error -- ERROR -- error'
17warnmsg = 'WARNING -- warning -- WARNING -- warning'
18
19# Gneral information
20##
21def reduce_spaces(string):
22    """ Function to give words of a line of text removing any extra space
23    """
24    values = string.replace('\n','').split(' ')
25    vals = []
26    for val in values:
27         if len(val) > 0:
28             vals.append(val)
29
30    return vals
31
32def variable_combo(varn,combofile):
33    """ Function to provide variables combination from a given variable name
34      varn= name of the variable
35      combofile= ASCII file with the combination of variables
36        [varn] [combo]
37          [combo]: '@' separated list of variables to use to generate [varn]
38            [WRFdt] to get WRF time-step (from general attributes)
39    >>> variable_combo('WRFprls','/home/lluis/PY/diagnostics.inf')
40    deaccum@RAINNC@XTIME@prnc
41    """
42    fname = 'variable_combo'
43
44    if varn == 'h':
45        print fname + '_____________________________________________________________'
46        print variable_combo.__doc__
47        quit()
48
49    if not os.path.isfile(combofile):
50        print errormsg
51        print '  ' + fname + ": file with combinations '" + combofile +              \
52          "' does not exist!!"
53        quit(-1)
54
55    objf = open(combofile, 'r')
56
57    found = False
58    for line in objf:
59        linevals = reduce_spaces(line)
60        varnf = linevals[0]
61        combo = linevals[1].replace('\n','')
62        if varn == varnf: 
63            found = True
64            break
65
66    if not found:
67        print errormsg
68        print '  ' + fname + ": variable '" + varn + "' not found in '" + combofile +\
69          "' !!"
70        combo='ERROR'
71
72    objf.close()
73
74    return combo
75
76# Mathematical operators
77##
78def compute_deaccum(varv, dimns, dimvns):
79    """ Function to compute the deaccumulation of a variable
80    compute_deaccum(varv, dimnames, dimvns)
81      [varv]= values to deaccum (assuming [t,])
82      [dimns]= list of the name of the dimensions of the [varv]
83      [dimvns]= list of the name of the variables with the values of the
84        dimensions of [varv]
85    """
86    fname = 'compute_deaccum'
87
88    deacdims = dimns[:]
89    deacvdims = dimvns[:]
90
91    slicei = []
92    slicee = []
93
94    Ndims = len(varv.shape)
95    for iid in range(0,Ndims):
96        slicei.append(slice(0,varv.shape[iid]))
97        slicee.append(slice(0,varv.shape[iid]))
98
99    slicee[0] = np.arange(varv.shape[0])
100    slicei[0] = np.arange(varv.shape[0])
101    slicei[0][1:varv.shape[0]] = np.arange(varv.shape[0]-1)
102
103    vari = varv[tuple(slicei)]
104    vare = varv[tuple(slicee)]
105
106    deac = vare - vari
107
108    return deac, deacdims, deacvdims
109
110def derivate_centered(var,dim,dimv):
111    """ Function to compute the centered derivate of a given field
112      centered derivate(n) = (var(n-1) + var(n+1))/(2*dn).
113    [var]= variable
114    [dim]= which dimension to compute the derivate
115    [dimv]= dimension values (can be of different dimension of [var])
116    >>> derivate_centered(np.arange(16).reshape(4,4)*1.,1,1.)
117    [[  0.   1.   2.   0.]
118     [  0.   5.   6.   0.]
119     [  0.   9.  10.   0.]
120     [  0.  13.  14.   0.]]
121    """
122
123    fname = 'derivate_centered'
124   
125    vark = var.dtype
126
127    if hasattr(dimv, "__len__"):
128# Assuming that the last dimensions of var [..., N, M] are the same of dimv [N, M]
129        if len(var.shape) != len(dimv.shape):
130            dimvals = np.zeros((var.shape), dtype=vark)
131            if len(var.shape) - len(dimv.shape) == 1:
132                for iz in range(var.shape[0]):
133                    dimvals[iz,] = dimv
134            elif len(var.shape) - len(dimv.shape) == 2:
135                for it in range(var.shape[0]):
136                    for iz in range(var.shape[1]):
137                        dimvals[it,iz,] = dimv
138            else:
139                print errormsg
140                print '  ' + fname + ': dimension difference between variable',      \
141                  var.shape,'and variable with dimension values',dimv.shape,         \
142                  ' not ready !!!'
143                quit(-1)
144        else:
145            dimvals = dimv
146    else:
147# dimension values are identical everywhere!
148# from: http://stackoverflow.com/questions/16807011/python-how-to-identify-if-a-variable-is-an-array-or-a-scalar   
149        dimvals = np.ones((var.shape), dtype=vark)*dimv
150
151    derivate = np.zeros((var.shape), dtype=vark)
152    if dim > len(var.shape) - 1:
153        print errormsg
154        print '  ' + fname + ': dimension',dim,' too big for given variable of ' +   \
155          'shape:', var.shape,'!!!'
156        quit(-1)
157
158    slicebef = []
159    sliceaft = []
160    sliceder = []
161
162    for id in range(len(var.shape)):
163        if id == dim:
164            slicebef.append(slice(0,var.shape[id]-2))
165            sliceaft.append(slice(2,var.shape[id]))
166            sliceder.append(slice(1,var.shape[id]-1))
167        else:
168            slicebef.append(slice(0,var.shape[id]))
169            sliceaft.append(slice(0,var.shape[id]))
170            sliceder.append(slice(0,var.shape[id]))
171
172    if hasattr(dimv, "__len__"):
173        derivate[tuple(sliceder)] = (var[tuple(slicebef)] + var[tuple(sliceaft)])/   \
174          ((dimvals[tuple(sliceaft)] - dimvals[tuple(slicebef)]))
175        print (dimvals[tuple(sliceaft)] - dimvals[tuple(slicebef)])
176    else:
177        derivate[tuple(sliceder)] = (var[tuple(slicebef)] + var[tuple(sliceaft)])/   \
178          (2.*dimv)
179
180#    print 'before________'
181#    print var[tuple(slicebef)]
182
183#    print 'after________'
184#    print var[tuple(sliceaft)]
185
186    return derivate
187
188def rotational_z(Vx,Vy,pos):
189    """ z-component of the rotatinoal of horizontal vectorial field
190    \/ x (Vx,Vy,Vz) = \/xVy - \/yVx
191    [Vx]= Variable component x
192    [Vy]=  Variable component y
193    [pos]= poisition of the grid points
194    >>> rotational_z(np.arange(16).reshape(4,4)*1., np.arange(16).reshape(4,4)*1., 1.)
195    [[  0.   1.   2.   0.]
196     [ -4.   0.   0.  -7.]
197     [ -8.   0.   0. -11.]
198     [  0.  13.  14.   0.]]
199    """
200
201    fname =  'rotational_z'
202
203    ndims = len(Vx.shape)
204    rot1 = derivate_centered(Vy,ndims-1,pos)
205    rot2 = derivate_centered(Vx,ndims-2,pos)
206
207    rot = rot1 - rot2
208
209    return rot
210
211# Diagnostics
212##
213
214def var_clt(cfra):
215    """ Function to compute the total cloud fraction following 'newmicro.F90' from
216      LMDZ using 1D vertical column values
217      [cldfra]= cloud fraction values (assuming [[t],z,y,x])
218    """
219    ZEPSEC=1.0E-12
220
221    fname = 'var_clt'
222
223    zclear = 1.
224    zcloud = 0.
225
226    dz = cfra.shape[0]
227    for iz in range(dz):
228        zclear =zclear*(1.-np.max([cfra[iz],zcloud]))/(1.-np.min([zcloud,1.-ZEPSEC]))
229        clt = 1. - zclear
230        zcloud = cfra[iz]
231
232    return clt
233
234def compute_clt(cldfra, dimns, dimvns):
235    """ Function to compute the total cloud fraction following 'newmicro.F90' from
236      LMDZ
237    compute_clt(cldfra, dimnames)
238      [cldfra]= cloud fraction values (assuming [[t],z,y,x])
239      [dimns]= list of the name of the dimensions of [cldfra]
240      [dimvns]= list of the name of the variables with the values of the
241        dimensions of [cldfra]
242    """
243    fname = 'compute_clt'
244
245    cltdims = dimns[:]
246    cltvdims = dimvns[:]
247
248    if len(cldfra.shape) == 4:
249        clt = np.zeros((cldfra.shape[0],cldfra.shape[2],cldfra.shape[3]),            \
250          dtype=np.float)
251        dx = cldfra.shape[3]
252        dy = cldfra.shape[2]
253        dz = cldfra.shape[1]
254        dt = cldfra.shape[0]
255        cltdims.pop(1)
256        cltvdims.pop(1)
257
258        for it in range(dt):
259            for ix in range(dx):
260                for iy in range(dy):
261                    zclear = 1.
262                    zcloud = 0.
263                    ncvar.percendone(it*dx*dy + ix*dy + iy, dx*dy*dt, 5, 'diagnosted')
264                    clt[it,iy,ix] = var_clt(cldfra[it,:,iy,ix])
265
266    else:
267        clt = np.zeros((cldfra.shape[1],cldfra.shape[2]), dtype=np.float)
268        dx = cldfra.shape[2]
269        dy = cldfra.shape[1]
270        dy = cldfra.shape[0]
271        cltdims.pop(0)
272        cltvdims.pop(0)
273        for ix in range(dx):
274            for iy in range(dy):
275                zclear = 1.
276                zcloud = 0.
277                ncvar.percendone(ix*dy + iy, dx*dy*dt, 5, 'diagnosted')
278                clt[iy,ix] = var_clt(cldfra[:,iy,ix])
279
280    return clt, cltdims, cltvdims
281
282def var_cllmh(cfra, p):
283    """ Fcuntion to compute cllmh on a 1D column
284    """
285
286    fname = 'var_cllmh'
287
288    ZEPSEC =1.0E-12
289    prmhc = 440.*100.
290    prmlc = 680.*100.
291
292    zclearl = 1.
293    zcloudl = 0.
294    zclearm = 1.
295    zcloudm = 0.
296    zclearh = 1.
297    zcloudh = 0.
298
299    dvz = cfra.shape[0]
300
301    cllmh = np.ones((3), dtype=np.float)
302
303    for iz in range(dvz):
304        if p[iz] < prmhc:
305            cllmh[2] = cllmh[2]*(1.-np.max([cfra[iz], zcloudh]))/(1.-                \
306              np.min([zcloudh,1.-ZEPSEC]))
307            zcloudh = cfra[iz]
308        elif p[iz] >= prmhc and p[iz] < prmlc:
309            cllmh[1] = cllmh[1]*(1.-np.max([cfra[iz], zcloudm]))/(1.-                \
310              np.min([zcloudm,1.-ZEPSEC]))
311            zcloudm = cfra[iz]
312        elif p[iz] >= prmlc:
313            cllmh[0] = cllmh[0]*(1.-np.max([cfra[iz], zcloudl]))/(1.-                \
314              np.min([zcloudl,1.-ZEPSEC]))
315            zcloudl = cfra[iz]
316
317    cllmh = 1.- cllmh
318
319    return cllmh
320
321def compute_cllmh(cldfra, pres, dimns, dimvns):
322    """ Function to compute cllmh: low/medium/hight cloud fraction following newmicro.F90 from LMDZ
323    compute_clt(cldfra, pres, dimns, dimvns)
324      [cldfra]= cloud fraction values (assuming [[t],z,y,x])
325      [pres] = pressure field
326      [dimns]= list of the name of the dimensions of [cldfra]
327      [dimvns]= list of the name of the variables with the values of the
328        dimensions of [cldfra]
329    """
330    fname = 'compute_cllmh'
331
332    cllmhdims = dimns[:]
333    cllmhvdims = dimvns[:]
334
335    if len(cldfra.shape) == 4:
336        dx = cldfra.shape[3]
337        dy = cldfra.shape[2]
338        dz = cldfra.shape[1]
339        dt = cldfra.shape[0]
340        cllmhdims.pop(1)
341        cllmhvdims.pop(1)
342
343        cllmh = np.ones(tuple([3, dt, dy, dx]), dtype=np.float)
344
345        for it in range(dt):
346            for ix in range(dx):
347                for iy in range(dy):
348                    ncvar.percendone(it*dx*dy + ix*dy + iy, dx*dy*dt, 5, 'diagnosted')
349                    cllmh[:,it,iy,ix] = var_cllmh(cldfra[it,:,iy,ix], pres[it,:,iy,ix])
350       
351    else:
352        dx = cldfra.shape[2]
353        dy = cldfra.shape[1]
354        dz = cldfra.shape[0]
355        cllmhdims.pop(0)
356        cllmhvdims.pop(0)
357
358        cllmh = np.ones(tuple([3, dy, dx]), dtype=np.float)
359
360        for ix in range(dx):
361            for iy in range(dy):
362                ncvar.percendone(ix*dy + iy,dx*dy, 5, 'diagnosted')
363                cllmh[:,iy,ix] = var_cllmh(cldfra[:,iy,ix], pres[:,iy,ix])
364
365    return cllmh, cllmhdims, cllmhvdims
366
367def var_virtualTemp (temp,rmix):
368    """ This function returns virtual temperature in K,
369      temp: temperature [K]
370      rmix: mixing ratio in [kgkg-1]
371    """
372
373    fname = 'var_virtualTemp'
374
375    virtual=temp*(0.622+rmix)/(0.622*(1.+rmix))
376
377    return virtual
378
379
380def var_mslp(pres, psfc, ter, tk, qv):
381    """ Function to compute mslp on a 1D column
382    """
383
384    fname = 'var_mslp'
385
386    N = 1.0
387    expon=287.04*.0065/9.81
388    pref = 40000.
389
390# First find where about 400 hPa is located
391    dz=len(pres) 
392
393    kref = -1
394    pinc = pres[0] - pres[dz-1]
395
396    if pinc < 0.:
397        for iz in range(1,dz):
398            if pres[iz-1] >= pref and pres[iz] < pref: 
399                kref = iz
400                break
401    else:
402        for iz in range(dz-1):
403            if pres[iz] >= pref and pres[iz+1] < pref: 
404                kref = iz
405                break
406
407    if kref == -1:
408        print errormsg
409        print '  ' + fname + ': no reference pressure:',pref,'found!!'
410        print '    values:',pres[:]
411        quit(-1)
412
413    mslp = 0.
414
415# We are below both the ground and the lowest data level.
416
417# First, find the model level that is closest to a "target" pressure
418# level, where the "target" pressure is delta-p less that the local
419# value of a horizontally smoothed surface pressure field.  We use
420# delta-p = 150 hPa here. A standard lapse rate temperature profile
421# passing through the temperature at this model level will be used
422# to define the temperature profile below ground.  This is similar
423# to the Benjamin and Miller (1990) method, using 
424# 700 hPa everywhere for the "target" pressure.
425
426# ptarget = psfc - 15000.
427    ptarget = 70000.
428    dpmin=1.e4
429    kupper = 0
430    if pinc > 0.:
431        for iz in range(dz-1,0,-1):
432            kupper = iz
433            dp=np.abs( pres[iz] - ptarget )
434            if dp < dpmin: exit
435            dpmin = np.min([dpmin, dp])
436    else:
437        for iz in range(dz):
438            kupper = iz
439            dp=np.abs( pres[iz] - ptarget )
440            if dp < dpmin: exit
441            dpmin = np.min([dpmin, dp])
442
443    pbot=np.max([pres[0], psfc])
444#    zbot=0.
445
446#    tbotextrap=tk(i,j,kupper,itt)*(pbot/pres_field(i,j,kupper,itt))**expon
447#    tvbotextrap=virtual(tbotextrap,qv(i,j,1,itt))
448
449#    data_out(i,j,itt,1) = (zbot+tvbotextrap/.0065*(1.-(interp_levels(1)/pbot)**expon))
450    tbotextrap = tk[kupper]*(psfc/ptarget)**expon
451    tvbotextrap = var_virtualTemp(tbotextrap, qv[kupper])
452    mslp = psfc*( (tvbotextrap+0.0065*ter)/tvbotextrap)**(1./expon)
453
454    return mslp
455
456def compute_mslp(pressure, psurface, terrain, temperature, qvapor, dimns, dimvns):
457    """ Function to compute mslp: mean sea level pressure following p_interp.F90 from WRF
458    var_mslp(pres, ter, tk, qv, dimns, dimvns)
459      [pressure]= pressure field [Pa] (assuming [[t],z,y,x])
460      [psurface]= surface pressure field [Pa]
461      [terrain]= topography [m]
462      [temperature]= temperature [K]
463      [qvapor]= water vapour mixing ratio [kgkg-1]
464      [dimns]= list of the name of the dimensions of [cldfra]
465      [dimvns]= list of the name of the variables with the values of the
466        dimensions of [pres]
467    """
468
469    fname = 'compute_mslp'
470
471    mslpdims = list(dimns[:])
472    mslpvdims = list(dimvns[:])
473
474    if len(pressure.shape) == 4:
475        mslpdims.pop(1)
476        mslpvdims.pop(1)
477    else:
478        mslpdims.pop(0)
479        mslpvdims.pop(0)
480
481    if len(pressure.shape) == 4:
482        dx = pressure.shape[3]
483        dy = pressure.shape[2]
484        dz = pressure.shape[1]
485        dt = pressure.shape[0]
486
487        mslpv = np.zeros(tuple([dt, dy, dx]), dtype=np.float)
488
489# Terrain... to 2D !
490        terval = np.zeros(tuple([dy, dx]), dtype=np.float)
491        if len(terrain.shape) == 3:
492            terval = terrain[0,:,:]
493        else:
494            terval = terrain
495
496        for ix in range(dx):
497            for iy in range(dy):
498                if terval[iy,ix] > 0.:
499                    for it in range(dt):
500                        mslpv[it,iy,ix] = var_mslp(pressure[it,:,iy,ix],             \
501                          psurface[it,iy,ix], terval[iy,ix], temperature[it,:,iy,ix],\
502                          qvapor[it,:,iy,ix])
503
504                        ncvar.percendone(it*dx*dy + ix*dy + iy, dx*dy*dt, 5, 'diagnosted')
505                else:
506                    mslpv[:,iy,ix] = psurface[:,iy,ix]
507
508    else:
509        dx = pressure.shape[2]
510        dy = pressure.shape[1]
511        dz = pressure.shape[0]
512
513        mslpv = np.zeros(tuple([dy, dx]), dtype=np.float)
514
515# Terrain... to 2D !
516        terval = np.zeros(tuple([dy, dx]), dtype=np.float)
517        if len(terrain.shape) == 3:
518            terval = terrain[0,:,:]
519        else:
520            terval = terrain
521
522        for ix in range(dx):
523            for iy in range(dy):
524                ncvar.percendone(ix*dy + iy,dx*dy, 5, 'diagnosted')
525                if terval[iy,ix] > 0.:
526                    mslpv[iy,ix] = var_mslp(pressure[:,iy,ix], psurface[iy,ix],          \
527                      terval[iy,ix], temperature[:,iy,ix], qvapor[:,iy,ix])
528                else:
529                    mslpv[iy,ix] = psfc[iy,ix]
530
531    return mslpv, mslpdims, mslpvdims
532
533def compute_prw(dens, q, dimns, dimvns):
534    """ Function to compute water vapour path (prw)
535      [dens] = density [in kgkg-1] (assuming [t],z,y,x)
536      [q] = mixing ratio in [kgkg-1] (assuming [t],z,y,x)
537      [dimns]= list of the name of the dimensions of [q]
538      [dimvns]= list of the name of the variables with the values of the
539        dimensions of [q]
540    """
541    fname = 'compute_prw'
542
543    prwdims = dimns[:]
544    prwvdims = dimvns[:]
545
546    if len(q.shape) == 4:
547        prwdims.pop(1)
548        prwvdims.pop(1)
549    else:
550        prwdims.pop(0)
551        prwvdims.pop(0)
552
553    data1 = dens*q
554    prw = np.sum(data1, axis=1)
555
556    return prw, prwdims, prwvdims
557
558def compute_rh(p, t, q, dimns, dimvns):
559    """ Function to compute relative humidity following 'Tetens' equation (T,P) ...'
560      [t]= temperature (assuming [[t],z,y,x] in [K])
561      [p] = pressure field (assuming in [hPa])
562      [q] = mixing ratio in [kgkg-1]
563      [dimns]= list of the name of the dimensions of [t]
564      [dimvns]= list of the name of the variables with the values of the
565        dimensions of [t]
566    """
567    fname = 'compute_rh'
568
569    rhdims = dimns[:]
570    rhvdims = dimvns[:]
571
572    data1 = 10.*0.6112*np.exp(17.67*(t-273.16)/(t-29.65))
573    data2 = 0.622*data1/(0.01*p-(1.-0.622)*data1)
574
575    rh = q/data2
576
577    return rh, rhdims, rhvdims
578
579def turbulence_var(varv, dimvn, dimn):
580    """ Function to compute the Taylor's decomposition turbulence term from a a given variable
581      x*=<x^2>_t-(<X>_t)^2
582    turbulence_var(varv,dimn)
583      varv= values of the variable
584      dimvn= names of the dimension of the variable
585      dimn= names of the dimensions (as a dictionary with 'X', 'Y', 'Z', 'T')
586    >>> turbulence_var(np.arange((27)).reshape(3,3,3),['time','y','x'],{'T':'time', 'Y':'y', 'X':'x'})
587    [[ 54.  54.  54.]
588     [ 54.  54.  54.]
589     [ 54.  54.  54.]]
590    """
591    fname = 'turbulence_varv'
592
593    timedimid = dimvn.index(dimn['T'])
594
595    varv2 = varv*varv
596
597    vartmean = np.mean(varv, axis=timedimid)
598    var2tmean = np.mean(varv2, axis=timedimid)
599
600    varvturb = var2tmean - (vartmean*vartmean)
601
602    return varvturb
603
604def compute_turbulence(v, dimns, dimvns):
605    """ Function to compute the rubulence term of the Taylor's decomposition ...'
606      x*=<x^2>_t-(<X>_t)^2
607      [v]= variable (assuming [[t],z,y,x])
608      [dimns]= list of the name of the dimensions of [v]
609      [dimvns]= list of the name of the variables with the values of the
610        dimensions of [v]
611    """
612    fname = 'compute_turbulence'
613
614    turbdims = dimns[:]
615    turbvdims = dimvns[:]
616
617    turbdims.pop(0)
618    turbvdims.pop(0)
619
620    v2 = v*v
621
622    vartmean = np.mean(v, axis=0)
623    var2tmean = np.mean(v2, axis=0)
624
625    turb = var2tmean - (vartmean*vartmean)
626
627    return turb, turbdims, turbvdims
628
629def timeunits_seconds(dtu):
630    """ Function to transform a time units to seconds
631    timeunits_seconds(timeuv)
632      [dtu]= time units value to transform in seconds
633    """
634    fname='timunits_seconds'
635
636    if dtu == 'years':
637        times = 365.*24.*3600.
638    elif dtu == 'weeks':
639        times = 7.*24.*3600.
640    elif dtu == 'days':
641        times = 24.*3600.
642    elif dtu == 'hours':
643        times = 3600.
644    elif dtu == 'minutes':
645        times = 60.
646    elif dtu == 'seconds':
647        times = 1.
648    elif dtu == 'miliseconds':
649        times = 1./1000.
650    else:
651        print errormsg
652        print '  ' + fname  + ": time units '" + dtu + "' not ready !!"
653        quit(-1)
654
655    return times
656
657####### ###### ##### #### ### ## #
658comboinf="\nIF -d 'variable_combo', provides information of the combination to obtain -v [varn] with the ASCII file with the combinations as -f [combofile]"
659
660parser = OptionParser()
661parser.add_option("-f", "--netCDF_file", dest="ncfile", help="file to use", metavar="FILE")
662parser.add_option("-d", "--dimensions", dest="dimns", 
663  help="[dimxn]@[dxvn],[dimyn]@[dxvn],[...,[dimtn]@[dxvn]], ',' list with the couples [dimDn]@[dDvn], [dimDn], name of the dimension D and name of the variable [dDvn] with the values of the dimension" + comboinf, 
664  metavar="LABELS")
665parser.add_option("-v", "--variables", dest="varns", 
666  help=" [varn1]|[var11]@[...[varN1]],[...,[varnM]|[var1M]@[...[varLM]]] ',' list of variables to compute [varnK] and its necessary ones [var1K]...[varPK]", metavar="VALUES")
667
668(opts, args) = parser.parse_args()
669
670#######    #######
671## MAIN
672    #######
673availdiags = ['ACRAINTOT', 'clt', 'cllmh', 'deaccum', 'LMDZrh', 'mslp', 'RAINTOT',   \
674  'rvors', 'turbulence', 'WRFrvors']
675
676# Variables not to check
677NONcheckingvars = ['cllmh', 'deaccum', 'WRFbils', 'WRFdens', 'WRFgeop', 'WRFp',      \
678  'WRFpos', 'WRFprc', 'WRFprls', 'WRFrh', 'LMDZrh', 'LMDZrhs', 'WRFrhs', 'WRFrvors', \
679  'WRFt']
680
681ofile = 'diagnostics.nc'
682
683dimns = opts.dimns
684varns = opts.varns
685
686# Special method. knowing variable combination
687##
688if opts.dimns == 'variable_combo':
689    print warnmsg
690    print '  ' + main + ': knowing variable combination !!!'
691    combination = variable_combo(opts.varns,opts.ncfile)
692    print '     COMBO: ' + combination
693    quit(-1)
694
695if not os.path.isfile(opts.ncfile):
696    print errormsg
697    print '   ' + main + ": file '" + opts.ncfile + "' does not exist !!"
698    quit(-1)
699
700ncobj = NetCDFFile(opts.ncfile, 'r')
701
702# File creation
703newnc = NetCDFFile(ofile,'w')
704
705# dimensions
706dimvalues = dimns.split(',')
707dnames = []
708dvnames = []
709
710for dimval in dimvalues:
711    dnames.append(dimval.split('@')[0])
712    dvnames.append(dimval.split('@')[1])
713
714# diagnostics to compute
715diags = varns.split(',')
716Ndiags = len(diags)
717
718# Looking for specific variables that might be use in more than one diagnostic
719WRFp_compute = False
720WRFt_compute = False
721WRFrh_compute = False
722WRFght_compute = False
723WRFdens_compute = False
724WRFpos_compute = False
725
726for idiag in range(Ndiags):
727    if diags[idiag].split('|')[1].find('@') == -1:
728        depvars = diags[idiag].split('|')[1]
729        if depvars == 'WRFp': WRFp_compute = True
730        if depvars == 'WRFt': WRFt_compute = True
731        if depvars == 'WRFrh': WRFrh_compute = True
732        if depvars == 'WRFght': WRFght_compute = True
733        if depvars == 'WRFdens': WRFdens_compute = True
734        if depvars == 'WRFpos': WRFpos_compute = True
735
736    else:
737        depvars = diags[idiag].split('|')[1].split('@')
738        if ncvar.searchInlist(depvars, 'WRFp'): WRFp_compute = True
739        if ncvar.searchInlist(depvars, 'WRFt'): WRFt_compute = True
740        if ncvar.searchInlist(depvars, 'WRFrh'): WRFrh_compute = True
741        if ncvar.searchInlist(depvars, 'WRFght'): WRFght_compute = True
742        if ncvar.searchInlist(depvars, 'WRFdens'): WRFdens_compute = True
743        if ncvar.searchInlist(depvars, 'WRFpos'): WRFpos_compute = True
744
745if WRFp_compute:
746    print '  ' + main + ': Retrieving pressure value from WRF as P + PB'
747    dimv = ncobj.variables['P'].shape
748    WRFp = ncobj.variables['P'][:] + ncobj.variables['PB'][:]
749
750if WRFght_compute:
751    print '    ' + main + ': computing geopotential height from WRF as PH + PHB ...' 
752    WRFght = ncobj.variables['PH'][:] + ncobj.variables['PHB'][:]
753
754if WRFrh_compute:
755    print '    ' + main + ": computing relative humidity from WRF as 'Tetens'" +     \
756      ' equation (T,P) ...'
757    p0=100000.
758    p=ncobj.variables['P'][:] + ncobj.variables['PB'][:]
759    tk = (ncobj.variables['T'][:] + 300.)*(p/p0)**(2./7.)
760    qv = ncobj.variables['QVAPOR'][:]
761
762    data1 = 10.*0.6112*np.exp(17.67*(tk-273.16)/(tk-29.65))
763    data2 = 0.622*data1/(0.01*p-(1.-0.622)*data1)
764
765    WRFrh = qv/data2
766
767if WRFt_compute:
768    print '    ' + main + ': computing temperature from WRF as inv_potT(T + 300) ...'
769    p0=100000.
770    p=ncobj.variables['P'][:] + ncobj.variables['PB'][:]
771
772    WRFt = (ncobj.variables['T'][:] + 300.)*(p/p0)**(2./7.)
773
774if WRFdens_compute:
775    print '    ' + main + ': computing air density from WRF as ((MU + MUB) * ' +     \
776      'DNW)/g ...'
777    grav = 9.81
778
779# Just we need in in absolute values: Size of the central grid cell
780##    dxval = ncobj.getncattr('DX')
781##    dyval = ncobj.getncattr('DY')
782##    mapfac = ncobj.variables['MAPFAC_M'][:]
783##    area = dxval*dyval*mapfac
784
785    mu = (ncobj.variables['MU'][:] + ncobj.variables['MUB'][:])
786    dnw = ncobj.variables['DNW'][:]
787
788    WRFdens = np.zeros((mu.shape[0], dnw.shape[1], mu.shape[1], mu.shape[2]),        \
789      dtype=np.float)
790    levval = np.zeros((mu.shape[1], mu.shape[2]), dtype=np.float)
791
792    for it in range(mu.shape[0]):
793        for iz in range(dnw.shape[1]):
794            levval.fill(np.abs(dnw[it,iz]))
795            WRFdens[it,iz,:,:] = levval
796            WRFdens[it,iz,:,:] = mu[it,:,:]*WRFdens[it,iz,:,:]/grav
797
798if WRFpos_compute:
799# WRF positions from the lowest-leftest corner of the matrix
800    print '    ' + main + ': computing position from MAPFAC_M as sqrt(DY*j**2 + ' +  \
801      'DX*x**2)*MAPFAC_M ...'
802
803    mapfac = ncobj.variables['MAPFAC_M'][:]
804
805    distx = np.float(ncobj.getncattr('DX'))
806    disty = np.float(ncobj.getncattr('DY'))
807
808    print 'distx:',distx,'disty:',disty
809
810    dx = mapfac.shape[2]
811    dy = mapfac.shape[1]
812    dt = mapfac.shape[0]
813
814    WRFpos = np.zeros((dt, dy, dx), dtype=np.float)
815
816    for i in range(1,dx):
817        WRFpos[0,0,i] = distx*i/mapfac[0,0,i]
818    for j in range(1,dy):
819        i=0
820        WRFpos[0,j,i] = WRFpos[0,j-1,i] + disty/mapfac[0,j,i]
821        for i in range(1,dx):
822#            WRFpos[0,j,i] = np.sqrt((disty*j)**2. + (distx*i)**2.)/mapfac[0,j,i]
823#            WRFpos[0,j,i] = np.sqrt((disty*j)**2. + (distx*i)**2.)
824             WRFpos[0,j,i] = WRFpos[0,j,i-1] + distx/mapfac[0,j,i]
825
826    for it in range(1,dt):
827        WRFpos[it,:,:] = WRFpos[0,:,:]
828
829### ## #
830# Going for the diagnostics
831### ## #
832print '  ' + main + ' ...'
833
834for idiag in range(Ndiags):
835    print '    diagnostic:',diags[idiag]
836    diag = diags[idiag].split('|')[0]
837    depvars = diags[idiag].split('|')[1].split('@')
838    if diags[idiag].split('|')[1].find('@') != -1:
839        depvars = diags[idiag].split('|')[1].split('@')
840        if depvars[0] == 'deaccum': diag='deaccum'
841        for depv in depvars:
842            if not ncobj.variables.has_key(depv) and not                             \
843              ncvar.searchInlist(NONcheckingvars, depv) and depvars[0] != 'deaccum':
844                print errormsg
845                print '  ' + main + ": file '" + opts.ncfile +                       \
846                  "' does not have variable '" + depv + "' !!"
847                quit(-1)
848    else:
849        depvars = diags[idiag].split('|')[1]
850        if not ncobj.variables.has_key(depvars) and not                              \
851          ncvar.searchInlist(NONcheckingvars, depvars) and depvars[0] != 'deaccum':
852            print errormsg
853            print '  ' + main + ": file '" + opts.ncfile +                           \
854              "' does not have variable '" + depvars + "' !!"
855            quit(-1)
856
857    print "\n    Computing '" + diag + "' from: ", depvars, '...'
858
859# acraintot: accumulated total precipitation from WRF RAINC, RAINNC
860    if diag == 'ACRAINTOT':
861           
862        var0 = ncobj.variables[depvars[0]]
863        var1 = ncobj.variables[depvars[1]]
864        diagout = var0[:] + var1[:]
865
866        dnamesvar = var0.dimensions
867        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
868
869        ncvar.insert_variable(ncobj, 'acpr', diagout, dnamesvar, dvnamesvar, newnc)
870
871# cllmh with cldfra, pres
872    elif diag == 'cllmh':
873           
874        var0 = ncobj.variables[depvars[0]]
875        if depvars[1] == 'WRFp':
876            var1 = WRFp
877        else:
878            var01 = ncobj.variables[depvars[1]]
879            if len(size(var1.shape)) < len(size(var0.shape)):
880                var1 = np.brodcast_arrays(var01,var0)[0]
881            else:
882                var1 = var01
883
884        diagout, diagoutd, diagoutvd = compute_cllmh(var0,var1,dnames,dvnames)
885        ncvar.insert_variable(ncobj, 'cll', diagout[0,:], diagoutd, diagoutvd, newnc)
886        ncvar.insert_variable(ncobj, 'clm', diagout[1,:], diagoutd, diagoutvd, newnc)
887        ncvar.insert_variable(ncobj, 'clh', diagout[2,:], diagoutd, diagoutvd, newnc)
888
889# clt with cldfra
890    elif diag == 'clt':
891           
892        var0 = ncobj.variables[depvars]
893        diagout, diagoutd, diagoutvd = compute_clt(var0,dnames,dvnames)
894        ncvar.insert_variable(ncobj, 'clt', diagout, diagoutd, diagoutvd, newnc)
895
896# deaccum: deacumulation of any variable as (Variable, time [as [tunits]
897#   from/since ....], newvarname)
898    elif diag == 'deaccum':
899
900        var0 = ncobj.variables[depvars[1]]
901        var1 = ncobj.variables[depvars[2]]
902
903        dnamesvar = var0.dimensions
904        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
905
906        diagout, diagoutd, diagoutvd = compute_deaccum(var0,dnamesvar,dvnamesvar)
907
908# Transforming to a flux
909        if depvars[2] == 'XTIME':
910            dtimeunits = var1.getncattr('description')
911            tunits = dtimeunits.split(' ')[0]
912        else:
913            dtimeunits = var1.getncattr('units')
914            tunits = dtimeunits.split(' ')[0]
915
916        dtime = (var1[1] - var1[0])*timeunits_seconds(tunits)
917        ncvar.insert_variable(ncobj, depvars[3], diagout/dtime, diagoutd, diagoutvd, newnc)
918
919# LMDZrh (pres, t, r)
920    elif diag == 'LMDZrh':
921           
922        var0 = ncobj.variables[depvars[0]][:]
923        var1 = ncobj.variables[depvars[1]][:]
924        var2 = ncobj.variables[depvars[2]][:]
925
926        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnames,dvnames)
927        ncvar.insert_variable(ncobj, 'hus', diagout, diagoutd, diagoutvd, newnc)
928
929# LMDZrhs (psol, t2m, q2m)
930    elif diag == 'LMDZrhs':
931           
932        var0 = ncobj.variables[depvars[0]][:]
933        var1 = ncobj.variables[depvars[1]][:]
934        var2 = ncobj.variables[depvars[2]][:]
935
936        dnamesvar = ncobj.variables[depvars[0]].dimensions
937        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
938
939        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnamesvar,dvnamesvar)
940
941        ncvar.insert_variable(ncobj, 'huss', diagout, diagoutd, diagoutvd, newnc)
942
943# mslp: mean sea level pressure (pres, psfc, terrain, temp, qv)
944    elif diag == 'mslp' or diag == 'WRFmslp':
945           
946        var1 = ncobj.variables[depvars[1]][:]
947        var2 = ncobj.variables[depvars[2]][:]
948        var4 = ncobj.variables[depvars[4]][:]
949
950        if diag == 'WRFmslp':
951            var0 = WRFp
952            var3 = WRFt
953            dnamesvar = ncobj.variables['P'].dimensions
954        else:
955            var0 = ncobj.variables[depvars[0]][:]
956            var3 = ncobj.variables[depvars[3]][:]
957            dnamesvar = ncobj.variables[depvars[0]].dimensions
958
959        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
960
961        diagout, diagoutd, diagoutvd = compute_mslp(var0, var1, var2, var3, var4,    \
962          dnamesvar, dvnamesvar)
963
964        ncvar.insert_variable(ncobj, 'psl', diagout, diagoutd, diagoutvd, newnc)
965
966# raintot: instantaneous total precipitation from WRF as (RAINC + RAINC) / dTime
967    elif diag == 'RAINTOT':
968
969        var0 = ncobj.variables[depvars[0]]
970        var1 = ncobj.variables[depvars[1]]
971        var2 = ncobj.variables[depvars[2]]
972
973        var = var0[:] + var1[:]
974
975        dnamesvar = var0.dimensions
976        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
977
978        diagout, diagoutd, diagoutvd = compute_deaccum(var,dnamesvar,dvnamesvar)
979
980# Transforming to a flux
981        dtimeunits = var2.getncattr('units')
982        tunits = dtimeunits.split(' ')[0]
983
984        dtime = (var2[1] - var2[0])*timeunits_seconds(tunits)
985        ncvar.insert_variable(ncobj, 'pr', diagout/dtime, diagoutd, diagoutvd, newnc)
986
987# turbulence (var)
988    elif diag == 'turbulence':
989
990        var0 = ncobj.variables[depvars][:]
991
992        dnamesvar = list(ncobj.variables[depvars].dimensions)
993        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
994
995        diagout, diagoutd, diagoutvd = compute_turbulence(var0,dnamesvar,dvnamesvar)
996        valsvar = ncvar.variables_values(depvars)
997
998        ncvar.insert_variable(ncobj, valsvar[0] + 'turb', diagout, diagoutd, 
999          diagoutvd, newnc)
1000        varobj = newnc.variables[valsvar[0] + 'turb']
1001        attrv = varobj.long_name
1002        attr = varobj.delncattr('long_name')
1003        newattr = ncvar.set_attribute(varobj, 'long_name', attrv +                   \
1004          " Taylor decomposition turbulence term")
1005
1006# WRFbils fom WRF as HFX + LH
1007    elif diag == 'WRFbils':
1008           
1009        var0 = ncobj.variables[depvars[0]][:]
1010        var1 = ncobj.variables[depvars[1]][:]
1011
1012        diagout = var0 + var1
1013
1014        ncvar.insert_variable(ncobj, 'bils', diagout, dnames, dvnames, newnc)
1015
1016# WRFp pressure from WRF as P + PB
1017    elif diag == 'WRFp':
1018           
1019        diagout = WRFp
1020
1021        ncvar.insert_variable(ncobj, 'pres', diagout, dnames, dvnames, newnc)
1022
1023# WRFpos
1024    elif diag == 'WRFpos':
1025           
1026        dnamesvar = ncobj.variables['MAPFAC_M'].dimensions
1027        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1028
1029        ncvar.insert_variable(ncobj, 'WRFpos', WRFpos, dnamesvar, dvnamesvar, newnc)
1030
1031# WRFprw WRF water vapour path WRFdens, QVAPOR
1032    elif diag == 'WRFprw':
1033           
1034        var0 = WRFdens
1035        var1 = ncobj.variables[depvars[1]]
1036
1037        dnamesvar = list(var1.dimensions)
1038        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1039
1040        diagout, diagoutd, diagoutvd = compute_prw(var0, var1, dnamesvar,dvnamesvar)
1041
1042        ncvar.insert_variable(ncobj, 'prw', diagout, diagoutd, diagoutvd, newnc)
1043
1044# WRFrh (P, T, QVAPOR)
1045    elif diag == 'WRFrh':
1046           
1047        dnamesvar = list(ncobj.variables[depvars[2]].dimensions)
1048        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1049
1050        ncvar.insert_variable(ncobj, 'hus', WRFrh, dnames, dvnames, newnc)
1051
1052# WRFrhs (PSFC, T2, Q2)
1053    elif diag == 'WRFrhs':
1054           
1055        var0 = ncobj.variables[depvars[0]][:]
1056        var1 = ncobj.variables[depvars[1]][:]
1057        var2 = ncobj.variables[depvars[2]][:]
1058
1059        dnamesvar = list(ncobj.variables[depvars[2]].dimensions)
1060        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1061
1062        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnamesvar,dvnamesvar)
1063        ncvar.insert_variable(ncobj, 'huss', diagout, diagoutd, diagoutvd, newnc)
1064
1065# rvors (u10, v10, WRFpos)
1066    elif diag == 'WRFrvors':
1067           
1068        var0 = ncobj.variables[depvars[0]]
1069        var1 = ncobj.variables[depvars[1]]
1070
1071        diagout = rotational_z(var0, var1, distx)
1072
1073        dnamesvar = ncobj.variables[depvars[0]].dimensions
1074        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1075
1076        ncvar.insert_variable(ncobj, 'rvors', diagout, dnamesvar, dvnamesvar, newnc)
1077
1078# wss (u10, v10)
1079    elif diag == 'wss':
1080           
1081        var0 = ncobj.variables[depvars[0]][:]
1082        var1 = ncobj.variables[depvars[1]][:]
1083
1084        diagout = np.sqrt(var0*var0 + var1*var1)
1085
1086        dnamesvar = ncobj.variables[depvars[0]].dimensions
1087        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1088
1089        print 'dnamesvar',dnamesvar
1090        print 'dnames',dnames
1091        print 'dvnames',dvnames
1092        print 'dvnamesvar',dvnamesvar
1093
1094        ncvar.insert_variable(ncobj, 'wss', diagout, dnamesvar, dvnamesvar, newnc)
1095
1096    else:
1097        print errormsg
1098        print '  ' + main + ": diagnostic '" + diag + "' not ready!!!"
1099        print '    available diagnostics: ', availdiags
1100        quit(-1)
1101
1102    newnc.sync()
1103
1104#   end of diagnostics
1105
1106# Global attributes
1107##
1108atvar = ncvar.set_attribute(newnc, 'program', 'diagnostics.py')
1109atvar = ncvar.set_attribute(newnc, 'version', '1.0')
1110atvar = ncvar.set_attribute(newnc, 'author', 'Fita Borrell, Lluis')
1111atvar = ncvar.set_attribute(newnc, 'institution', 'Laboratoire Meteorologie ' +      \
1112  'Dynamique')
1113atvar = ncvar.set_attribute(newnc, 'university', 'Universite Pierre et Marie ' +     \
1114  'Curie -- Jussieu')
1115atvar = ncvar.set_attribute(newnc, 'centre', 'Centre national de la recherche ' +    \
1116  'scientifique')
1117atvar = ncvar.set_attribute(newnc, 'city', 'Paris')
1118atvar = ncvar.set_attribute(newnc, 'original_file', opts.ncfile)
1119
1120gorigattrs = ncobj.ncattrs()
1121
1122for attr in gorigattrs:
1123    attrv = ncobj.getncattr(attr)
1124    atvar = ncvar.set_attribute(newnc, attr, attrv)
1125
1126ncobj.close()
1127newnc.close()
1128
1129print '\n' + main + ': successfull writting of diagnostics file "' + ofile + '" !!!'
Note: See TracBrowser for help on using the repository browser.