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

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

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