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

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

Fixing variable name
File size: 43.0 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_OMEGAw(omega, p, t, dimns, dimvns):
535    """ Function to transform OMEGA [Pas-1] to velocities [ms-1]
536      tacking: https://www.ncl.ucar.edu/Document/Functions/Contributed/omega_to_w.shtml
537      [omega] = vertical velocity [in ms-1] (assuming [t],z,y,x)
538      [p] = pressure in [Pa] (assuming [t],z,y,x)
539      [t] = temperature in [K] (assuming [t],z,y,x)
540      [dimns]= list of the name of the dimensions of [q]
541      [dimvns]= list of the name of the variables with the values of the
542        dimensions of [q]
543    """
544    fname = 'compute_OMEGAw'
545
546    rgas = 287.058            # J/(kg-K) => m2/(s2 K)
547    g    = 9.80665            # m/s2
548
549    wdims = dimns[:]
550    wvdims = dimvns[:]
551
552    rho  = p/(rgas*t)         # density => kg/m3
553    w    = -omega/(rho*g)     
554
555    return w, wdims, wvdims
556
557def compute_prw(dens, q, dimns, dimvns):
558    """ Function to compute water vapour path (prw)
559      [dens] = density [in kgkg-1] (assuming [t],z,y,x)
560      [q] = mixing ratio in [kgkg-1] (assuming [t],z,y,x)
561      [dimns]= list of the name of the dimensions of [q]
562      [dimvns]= list of the name of the variables with the values of the
563        dimensions of [q]
564    """
565    fname = 'compute_prw'
566
567    prwdims = dimns[:]
568    prwvdims = dimvns[:]
569
570    if len(q.shape) == 4:
571        prwdims.pop(1)
572        prwvdims.pop(1)
573    else:
574        prwdims.pop(0)
575        prwvdims.pop(0)
576
577    data1 = dens*q
578    prw = np.sum(data1, axis=1)
579
580    return prw, prwdims, prwvdims
581
582def compute_rh(p, t, q, dimns, dimvns):
583    """ Function to compute relative humidity following 'Tetens' equation (T,P) ...'
584      [t]= temperature (assuming [[t],z,y,x] in [K])
585      [p] = pressure field (assuming in [hPa])
586      [q] = mixing ratio in [kgkg-1]
587      [dimns]= list of the name of the dimensions of [t]
588      [dimvns]= list of the name of the variables with the values of the
589        dimensions of [t]
590    """
591    fname = 'compute_rh'
592
593    rhdims = dimns[:]
594    rhvdims = dimvns[:]
595
596    data1 = 10.*0.6112*np.exp(17.67*(t-273.16)/(t-29.65))
597    data2 = 0.622*data1/(0.01*p-(1.-0.622)*data1)
598
599    rh = q/data2
600
601    return rh, rhdims, rhvdims
602
603def compute_td(p, temp, qv, dimns, dimvns):
604    """ Function to compute the dew point temperature
605      [p]= pressure [Pa]
606      [temp]= temperature [C]
607      [qv]= mixing ratio [kgkg-1]
608      [dimns]= list of the name of the dimensions of [p]
609      [dimvns]= list of the name of the variables with the values of the
610        dimensions of [p]
611    """
612    fname = 'compute_td'
613
614#    print '    ' + fname + ': computing dew-point temperature from TS as t and Tetens...'
615# tacking from: http://en.wikipedia.org/wiki/Dew_point
616    tk = temp
617    data1 = 10.*0.6112*np.exp(17.67*(tk-273.16)/(tk-29.65))
618    data2 = 0.622*data1/(0.01*p-(1.-0.622)*data1)
619
620    rh = qv/data2
621               
622    pa = rh * data1
623    td = 257.44*np.log(pa/6.1121)/(18.678-np.log(pa/6.1121))
624
625    tddims = dimns[:]
626    tdvdims = dimvns[:]
627
628    return td, tddims, tdvdims
629
630def turbulence_var(varv, dimvn, dimn):
631    """ Function to compute the Taylor's decomposition turbulence term from a a given variable
632      x*=<x^2>_t-(<X>_t)^2
633    turbulence_var(varv,dimn)
634      varv= values of the variable
635      dimvn= names of the dimension of the variable
636      dimn= names of the dimensions (as a dictionary with 'X', 'Y', 'Z', 'T')
637    >>> turbulence_var(np.arange((27)).reshape(3,3,3),['time','y','x'],{'T':'time', 'Y':'y', 'X':'x'})
638    [[ 54.  54.  54.]
639     [ 54.  54.  54.]
640     [ 54.  54.  54.]]
641    """
642    fname = 'turbulence_varv'
643
644    timedimid = dimvn.index(dimn['T'])
645
646    varv2 = varv*varv
647
648    vartmean = np.mean(varv, axis=timedimid)
649    var2tmean = np.mean(varv2, axis=timedimid)
650
651    varvturb = var2tmean - (vartmean*vartmean)
652
653    return varvturb
654
655def compute_turbulence(v, dimns, dimvns):
656    """ Function to compute the rubulence term of the Taylor's decomposition ...'
657      x*=<x^2>_t-(<X>_t)^2
658      [v]= variable (assuming [[t],z,y,x])
659      [dimns]= list of the name of the dimensions of [v]
660      [dimvns]= list of the name of the variables with the values of the
661        dimensions of [v]
662    """
663    fname = 'compute_turbulence'
664
665    turbdims = dimns[:]
666    turbvdims = dimvns[:]
667
668    turbdims.pop(0)
669    turbvdims.pop(0)
670
671    v2 = v*v
672
673    vartmean = np.mean(v, axis=0)
674    var2tmean = np.mean(v2, axis=0)
675
676    turb = var2tmean - (vartmean*vartmean)
677
678    return turb, turbdims, turbvdims
679
680def compute_wds(u, v, dimns, dimvns):
681    """ Function to compute the wind direction
682      [u]= W-E wind direction [ms-1, knot, ...]
683      [v]= N-S wind direction [ms-1, knot, ...]
684      [dimns]= list of the name of the dimensions of [u]
685      [dimvns]= list of the name of the variables with the values of the
686        dimensions of [u]
687    """
688    fname = 'compute_wds'
689
690#    print '    ' + fname + ': computing wind direction as ATAN2(v,u) ...'
691    theta = np.arctan2(v,u)
692    theta = np.where(theta < 0., theta + 2.*np.pi, theta)
693
694    wds = 360.*theta/(2.*np.pi)
695
696    wdsdims = dimns[:]
697    wdsvdims = dimvns[:]
698
699    return wds, wdsdims, wdsvdims
700
701def compute_wss(u, v, dimns, dimvns):
702    """ Function to compute the wind speed
703      [u]= W-E wind direction [ms-1, knot, ...]
704      [v]= N-S wind direction [ms-1, knot, ...]
705      [dimns]= list of the name of the dimensions of [u]
706      [dimvns]= list of the name of the variables with the values of the
707        dimensions of [u]
708    """
709    fname = 'compute_wss'
710
711#    print '    ' + fname + ': computing wind speed as SQRT(v**2 + u**2) ...'
712    wss = np.sqrt(u*u + v*v)
713
714    wssdims = dimns[:]
715    wssvdims = dimvns[:]
716
717    return wss, wssdims, wssvdims
718
719def timeunits_seconds(dtu):
720    """ Function to transform a time units to seconds
721    timeunits_seconds(timeuv)
722      [dtu]= time units value to transform in seconds
723    """
724    fname='timunits_seconds'
725
726    if dtu == 'years':
727        times = 365.*24.*3600.
728    elif dtu == 'weeks':
729        times = 7.*24.*3600.
730    elif dtu == 'days':
731        times = 24.*3600.
732    elif dtu == 'hours':
733        times = 3600.
734    elif dtu == 'minutes':
735        times = 60.
736    elif dtu == 'seconds':
737        times = 1.
738    elif dtu == 'miliseconds':
739        times = 1./1000.
740    else:
741        print errormsg
742        print '  ' + fname  + ": time units '" + dtu + "' not ready !!"
743        quit(-1)
744
745    return times
746
747####### ###### ##### #### ### ## #
748comboinf="\nIF -d 'variable_combo', provides information of the combination to obtain -v [varn] with the ASCII file with the combinations as -f [combofile]"
749
750parser = OptionParser()
751parser.add_option("-f", "--netCDF_file", dest="ncfile", help="file to use", metavar="FILE")
752parser.add_option("-d", "--dimensions", dest="dimns", 
753  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, 
754  metavar="LABELS")
755parser.add_option("-v", "--variables", dest="varns", 
756  help=" [varn1]|[var11]@[...[varN1]],[...,[varnM]|[var1M]@[...[varLM]]] ',' list of variables to compute [varnK] and its necessary ones [var1K]...[varPK]", metavar="VALUES")
757
758(opts, args) = parser.parse_args()
759
760#######    #######
761## MAIN
762    #######
763availdiags = ['ACRAINTOT', 'clt', 'cllmh', 'deaccum', 'LMDZrh', 'mslp', 'RAINTOT',   \
764  'rvors', 'td', 'turbulence', 'WRFrvors', 'wds', 'wss']
765
766# Variables not to check
767NONcheckingvars = ['cllmh', 'deaccum', 'TSrhs', 'TStd', 'TSwds', 'TSwss', 'WRFbils', \
768  'WRFdens', 'WRFgeop',                                                              \
769  'WRFp', 'WRFtd',                                                                   \
770  'WRFpos', 'WRFprc', 'WRFprls', 'WRFrh', 'LMDZrh', 'LMDZrhs', 'WRFrhs', 'WRFrvors', \
771  'WRFt', 'WRFtime', 'WRFwds', 'WRFwss']
772
773ofile = 'diagnostics.nc'
774
775dimns = opts.dimns
776varns = opts.varns
777
778# Special method. knowing variable combination
779##
780if opts.dimns == 'variable_combo':
781    print warnmsg
782    print '  ' + main + ': knowing variable combination !!!'
783    combination = variable_combo(opts.varns,opts.ncfile)
784    print '     COMBO: ' + combination
785    quit(-1)
786
787if not os.path.isfile(opts.ncfile):
788    print errormsg
789    print '   ' + main + ": file '" + opts.ncfile + "' does not exist !!"
790    quit(-1)
791
792ncobj = NetCDFFile(opts.ncfile, 'r')
793
794# File creation
795newnc = NetCDFFile(ofile,'w')
796
797# dimensions
798dimvalues = dimns.split(',')
799dnames = []
800dvnames = []
801
802for dimval in dimvalues:
803    dnames.append(dimval.split('@')[0])
804    dvnames.append(dimval.split('@')[1])
805
806# diagnostics to compute
807diags = varns.split(',')
808Ndiags = len(diags)
809
810# Looking for specific variables that might be use in more than one diagnostic
811WRFp_compute = False
812WRFt_compute = False
813WRFrh_compute = False
814WRFght_compute = False
815WRFdens_compute = False
816WRFpos_compute = False
817
818for idiag in range(Ndiags):
819    if diags[idiag].split('|')[1].find('@') == -1:
820        depvars = diags[idiag].split('|')[1]
821        if depvars == 'WRFp': WRFp_compute = True
822        if depvars == 'WRFt': WRFt_compute = True
823        if depvars == 'WRFrh': WRFrh_compute = True
824        if depvars == 'WRFght': WRFght_compute = True
825        if depvars == 'WRFdens': WRFdens_compute = True
826        if depvars == 'WRFpos': WRFpos_compute = True
827
828    else:
829        depvars = diags[idiag].split('|')[1].split('@')
830        if ncvar.searchInlist(depvars, 'WRFp'): WRFp_compute = True
831        if ncvar.searchInlist(depvars, 'WRFt'): WRFt_compute = True
832        if ncvar.searchInlist(depvars, 'WRFrh'): WRFrh_compute = True
833        if ncvar.searchInlist(depvars, 'WRFght'): WRFght_compute = True
834        if ncvar.searchInlist(depvars, 'WRFdens'): WRFdens_compute = True
835        if ncvar.searchInlist(depvars, 'WRFpos'): WRFpos_compute = True
836
837if WRFp_compute:
838    print '  ' + main + ': Retrieving pressure value from WRF as P + PB'
839    dimv = ncobj.variables['P'].shape
840    WRFp = ncobj.variables['P'][:] + ncobj.variables['PB'][:]
841
842if WRFght_compute:
843    print '    ' + main + ': computing geopotential height from WRF as PH + PHB ...' 
844    WRFght = ncobj.variables['PH'][:] + ncobj.variables['PHB'][:]
845
846if WRFrh_compute:
847    print '    ' + main + ": computing relative humidity from WRF as 'Tetens'" +     \
848      ' equation (T,P) ...'
849    p0=100000.
850    p=ncobj.variables['P'][:] + ncobj.variables['PB'][:]
851    tk = (ncobj.variables['T'][:] + 300.)*(p/p0)**(2./7.)
852    qv = ncobj.variables['QVAPOR'][:]
853
854    data1 = 10.*0.6112*np.exp(17.67*(tk-273.16)/(tk-29.65))
855    data2 = 0.622*data1/(0.01*p-(1.-0.622)*data1)
856
857    WRFrh = qv/data2
858
859if WRFt_compute:
860    print '    ' + main + ': computing temperature from WRF as inv_potT(T + 300) ...'
861    p0=100000.
862    p=ncobj.variables['P'][:] + ncobj.variables['PB'][:]
863
864    WRFt = (ncobj.variables['T'][:] + 300.)*(p/p0)**(2./7.)
865
866if WRFdens_compute:
867    print '    ' + main + ': computing air density from WRF as ((MU + MUB) * ' +     \
868      'DNW)/g ...'
869    grav = 9.81
870
871# Just we need in in absolute values: Size of the central grid cell
872##    dxval = ncobj.getncattr('DX')
873##    dyval = ncobj.getncattr('DY')
874##    mapfac = ncobj.variables['MAPFAC_M'][:]
875##    area = dxval*dyval*mapfac
876
877    mu = (ncobj.variables['MU'][:] + ncobj.variables['MUB'][:])
878    dnw = ncobj.variables['DNW'][:]
879
880    WRFdens = np.zeros((mu.shape[0], dnw.shape[1], mu.shape[1], mu.shape[2]),        \
881      dtype=np.float)
882    levval = np.zeros((mu.shape[1], mu.shape[2]), dtype=np.float)
883
884    for it in range(mu.shape[0]):
885        for iz in range(dnw.shape[1]):
886            levval.fill(np.abs(dnw[it,iz]))
887            WRFdens[it,iz,:,:] = levval
888            WRFdens[it,iz,:,:] = mu[it,:,:]*WRFdens[it,iz,:,:]/grav
889
890if WRFpos_compute:
891# WRF positions from the lowest-leftest corner of the matrix
892    print '    ' + main + ': computing position from MAPFAC_M as sqrt(DY*j**2 + ' +  \
893      'DX*x**2)*MAPFAC_M ...'
894
895    mapfac = ncobj.variables['MAPFAC_M'][:]
896
897    distx = np.float(ncobj.getncattr('DX'))
898    disty = np.float(ncobj.getncattr('DY'))
899
900    print 'distx:',distx,'disty:',disty
901
902    dx = mapfac.shape[2]
903    dy = mapfac.shape[1]
904    dt = mapfac.shape[0]
905
906    WRFpos = np.zeros((dt, dy, dx), dtype=np.float)
907
908    for i in range(1,dx):
909        WRFpos[0,0,i] = distx*i/mapfac[0,0,i]
910    for j in range(1,dy):
911        i=0
912        WRFpos[0,j,i] = WRFpos[0,j-1,i] + disty/mapfac[0,j,i]
913        for i in range(1,dx):
914#            WRFpos[0,j,i] = np.sqrt((disty*j)**2. + (distx*i)**2.)/mapfac[0,j,i]
915#            WRFpos[0,j,i] = np.sqrt((disty*j)**2. + (distx*i)**2.)
916             WRFpos[0,j,i] = WRFpos[0,j,i-1] + distx/mapfac[0,j,i]
917
918    for it in range(1,dt):
919        WRFpos[it,:,:] = WRFpos[0,:,:]
920
921### ## #
922# Going for the diagnostics
923### ## #
924print '  ' + main + ' ...'
925
926for idiag in range(Ndiags):
927    print '    diagnostic:',diags[idiag]
928    diag = diags[idiag].split('|')[0]
929    depvars = diags[idiag].split('|')[1].split('@')
930    if diags[idiag].split('|')[1].find('@') != -1:
931        depvars = diags[idiag].split('|')[1].split('@')
932        if depvars[0] == 'deaccum': diag='deaccum'
933        for depv in depvars:
934            if not ncobj.variables.has_key(depv) and not                             \
935              ncvar.searchInlist(NONcheckingvars, depv) and depvars[0] != 'deaccum':
936                print errormsg
937                print '  ' + main + ": file '" + opts.ncfile +                       \
938                  "' does not have variable '" + depv + "' !!"
939                quit(-1)
940    else:
941        depvars = diags[idiag].split('|')[1]
942        if not ncobj.variables.has_key(depvars) and not                              \
943          ncvar.searchInlist(NONcheckingvars, depvars) and depvars[0] != 'deaccum':
944            print errormsg
945            print '  ' + main + ": file '" + opts.ncfile +                           \
946              "' does not have variable '" + depvars + "' !!"
947            quit(-1)
948
949    print "\n    Computing '" + diag + "' from: ", depvars, '...'
950
951# acraintot: accumulated total precipitation from WRF RAINC, RAINNC
952    if diag == 'ACRAINTOT':
953           
954        var0 = ncobj.variables[depvars[0]]
955        var1 = ncobj.variables[depvars[1]]
956        diagout = var0[:] + var1[:]
957
958        dnamesvar = var0.dimensions
959        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
960
961        ncvar.insert_variable(ncobj, 'acpr', diagout, dnamesvar, dvnamesvar, newnc)
962
963# cllmh with cldfra, pres
964    elif diag == 'cllmh':
965           
966        var0 = ncobj.variables[depvars[0]]
967        if depvars[1] == 'WRFp':
968            var1 = WRFp
969        else:
970            var01 = ncobj.variables[depvars[1]]
971            if len(size(var1.shape)) < len(size(var0.shape)):
972                var1 = np.brodcast_arrays(var01,var0)[0]
973            else:
974                var1 = var01
975
976        diagout, diagoutd, diagoutvd = compute_cllmh(var0,var1,dnames,dvnames)
977        ncvar.insert_variable(ncobj, 'cll', diagout[0,:], diagoutd, diagoutvd, newnc)
978        ncvar.insert_variable(ncobj, 'clm', diagout[1,:], diagoutd, diagoutvd, newnc)
979        ncvar.insert_variable(ncobj, 'clh', diagout[2,:], diagoutd, diagoutvd, newnc)
980
981# clt with cldfra
982    elif diag == 'clt':
983           
984        var0 = ncobj.variables[depvars]
985        diagout, diagoutd, diagoutvd = compute_clt(var0,dnames,dvnames)
986        ncvar.insert_variable(ncobj, 'clt', diagout, diagoutd, diagoutvd, newnc)
987
988# deaccum: deacumulation of any variable as (Variable, time [as [tunits]
989#   from/since ....], newvarname)
990    elif diag == 'deaccum':
991
992        var0 = ncobj.variables[depvars[1]]
993        var1 = ncobj.variables[depvars[2]]
994
995        dnamesvar = var0.dimensions
996        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
997
998        diagout, diagoutd, diagoutvd = compute_deaccum(var0,dnamesvar,dvnamesvar)
999
1000# Transforming to a flux
1001        if depvars[2] == 'XTIME':
1002            dtimeunits = var1.getncattr('description')
1003            tunits = dtimeunits.split(' ')[0]
1004        else:
1005            dtimeunits = var1.getncattr('units')
1006            tunits = dtimeunits.split(' ')[0]
1007
1008        dtime = (var1[1] - var1[0])*timeunits_seconds(tunits)
1009        ncvar.insert_variable(ncobj, depvars[3], diagout/dtime, diagoutd, diagoutvd, newnc)
1010
1011# LMDZrh (pres, t, r)
1012    elif diag == 'LMDZrh':
1013           
1014        var0 = ncobj.variables[depvars[0]][:]
1015        var1 = ncobj.variables[depvars[1]][:]
1016        var2 = ncobj.variables[depvars[2]][:]
1017
1018        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnames,dvnames)
1019        ncvar.insert_variable(ncobj, 'hus', diagout, diagoutd, diagoutvd, newnc)
1020
1021# LMDZrhs (psol, t2m, q2m)
1022    elif diag == 'LMDZrhs':
1023           
1024        var0 = ncobj.variables[depvars[0]][:]
1025        var1 = ncobj.variables[depvars[1]][:]
1026        var2 = ncobj.variables[depvars[2]][:]
1027
1028        dnamesvar = ncobj.variables[depvars[0]].dimensions
1029        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1030
1031        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnamesvar,dvnamesvar)
1032
1033        ncvar.insert_variable(ncobj, 'huss', diagout, diagoutd, diagoutvd, newnc)
1034
1035# mslp: mean sea level pressure (pres, psfc, terrain, temp, qv)
1036    elif diag == 'mslp' or diag == 'WRFmslp':
1037           
1038        var1 = ncobj.variables[depvars[1]][:]
1039        var2 = ncobj.variables[depvars[2]][:]
1040        var4 = ncobj.variables[depvars[4]][:]
1041
1042        if diag == 'WRFmslp':
1043            var0 = WRFp
1044            var3 = WRFt
1045            dnamesvar = ncobj.variables['P'].dimensions
1046        else:
1047            var0 = ncobj.variables[depvars[0]][:]
1048            var3 = ncobj.variables[depvars[3]][:]
1049            dnamesvar = ncobj.variables[depvars[0]].dimensions
1050
1051        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1052
1053        diagout, diagoutd, diagoutvd = compute_mslp(var0, var1, var2, var3, var4,    \
1054          dnamesvar, dvnamesvar)
1055
1056        ncvar.insert_variable(ncobj, 'psl', diagout, diagoutd, diagoutvd, newnc)
1057
1058# OMEGAw (omega, p, t) from NCL formulation (https://www.ncl.ucar.edu/Document/Functions/Contributed/omega_to_w.shtml)
1059    elif diag == 'OMEGAw':
1060           
1061        var0 = ncobj.variables[depvars[0]][:]
1062        var1 = ncobj.variables[depvars[1]][:]
1063        var2 = ncobj.variables[depvars[2]][:]
1064
1065        dnamesvar = ncobj.variables[depvars[0]].dimensions
1066        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1067
1068        diagout, diagoutd, diagoutvd = compute_OMEGAw(var0,var1,var2,dnamesvar,dvnamesvar)
1069
1070        ncvar.insert_variable(ncobj, 'wa', diagout, diagoutd, diagoutvd, newnc)
1071
1072# raintot: instantaneous total precipitation from WRF as (RAINC + RAINC) / dTime
1073    elif diag == 'RAINTOT':
1074
1075        var0 = ncobj.variables[depvars[0]]
1076        var1 = ncobj.variables[depvars[1]]
1077        if depvars[2] != 'WRFtime':
1078            var2 = ncobj.variables[depvars[2]]
1079
1080        var = var0[:] + var1[:]
1081
1082        dnamesvar = var0.dimensions
1083        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1084
1085        diagout, diagoutd, diagoutvd = compute_deaccum(var,dnamesvar,dvnamesvar)
1086
1087# Transforming to a flux
1088        if var2.shape[0] > 0:
1089            if depvars[2] != 'WRFtime':
1090                dtimeunits = var2.getncattr('units')
1091                tunits = dtimeunits.split(' ')[0]
1092   
1093                dtime = (var2[1] - var2[0])*timeunits_seconds(tunits)
1094            else:
1095                var2 = ncobj.variables['Times']
1096                time1 = var2[0,:]
1097                time2 = var2[1,:]
1098                tmf1 = ''
1099                tmf2 = ''
1100                for ic in range(len(time1)):
1101                    tmf1 = tmf1 + time1[ic]
1102                    tmf2 = tmf2 + time2[ic]
1103                dtdate1 = dt.datetime.strptime(tmf1,"%Y-%m-%d_%H:%M:%S")
1104                dtdate2 = dt.datetime.strptime(tmf2,"%Y-%m-%d_%H:%M:%S")
1105                diffdate12 = dtdate2 - dtdate1
1106                dtime = diffdate12.total_seconds()
1107                print 'dtime:',dtime
1108        else:
1109            print warnmsg
1110            print '  ' + fname + ": only 1 time-step for '" + diag + "' !!"
1111            print '    leaving a zero value!'
1112            diagout = var0*0.
1113            dtime=1.
1114
1115        ncvar.insert_variable(ncobj, 'pr', diagout/dtime, diagoutd, diagoutvd, newnc)
1116
1117# rhs (psfc, t, q) from TimeSeries files
1118    elif diag == 'TSrhs':
1119           
1120        p0=100000.
1121        var0 = ncobj.variables[depvars[0]][:]
1122        var1 = (ncobj.variables[depvars[1]][:])*(var0/p0)**(2./7.)
1123        var2 = ncobj.variables[depvars[2]][:]
1124
1125        dnamesvar = ncobj.variables[depvars[0]].dimensions
1126        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1127
1128        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnamesvar,dvnamesvar)
1129
1130        ncvar.insert_variable(ncobj, 'huss', diagout, diagoutd, diagoutvd, newnc)
1131
1132# td (psfc, t, q) from TimeSeries files
1133    elif diag == 'TStd' or diag == 'td':
1134           
1135        var0 = ncobj.variables[depvars[0]][:]
1136        var1 = ncobj.variables[depvars[1]][:] - 273.15
1137        var2 = ncobj.variables[depvars[2]][:]
1138
1139        dnamesvar = ncobj.variables[depvars[0]].dimensions
1140        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1141
1142        diagout, diagoutd, diagoutvd = compute_td(var0,var1,var2,dnamesvar,dvnamesvar)
1143
1144        ncvar.insert_variable(ncobj, 'tds', diagout, diagoutd, diagoutvd, newnc)
1145
1146# td (psfc, t, q) from TimeSeries files
1147    elif diag == 'TStdC' or diag == 'tdC':
1148           
1149        var0 = ncobj.variables[depvars[0]][:]
1150# Temperature is already in degrees Celsius
1151        var1 = ncobj.variables[depvars[1]][:]
1152        var2 = ncobj.variables[depvars[2]][:]
1153
1154        dnamesvar = ncobj.variables[depvars[0]].dimensions
1155        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1156
1157        diagout, diagoutd, diagoutvd = compute_td(var0,var1,var2,dnamesvar,dvnamesvar)
1158
1159        ncvar.insert_variable(ncobj, 'tds', diagout, diagoutd, diagoutvd, newnc)
1160
1161# wds (u, v)
1162    elif diag == 'TSwds' or diag == 'wds' :
1163 
1164        var0 = ncobj.variables[depvars[0]][:]
1165        var1 = ncobj.variables[depvars[1]][:]
1166
1167        dnamesvar = ncobj.variables[depvars[0]].dimensions
1168        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1169
1170        diagout, diagoutd, diagoutvd = compute_wds(var0,var1,dnamesvar,dvnamesvar)
1171
1172        ncvar.insert_variable(ncobj, 'wds', diagout, diagoutd, diagoutvd, newnc)
1173
1174# wss (u, v)
1175    elif diag == 'TSwss' or diag == 'wss':
1176           
1177        var0 = ncobj.variables[depvars[0]][:]
1178        var1 = ncobj.variables[depvars[1]][:]
1179
1180        dnamesvar = ncobj.variables[depvars[0]].dimensions
1181        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1182
1183        diagout, diagoutd, diagoutvd = compute_wss(var0,var1,dnamesvar,dvnamesvar)
1184
1185        ncvar.insert_variable(ncobj, 'wss', diagout, diagoutd, diagoutvd, newnc)
1186
1187# turbulence (var)
1188    elif diag == 'turbulence':
1189
1190        var0 = ncobj.variables[depvars][:]
1191
1192        dnamesvar = list(ncobj.variables[depvars].dimensions)
1193        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1194
1195        diagout, diagoutd, diagoutvd = compute_turbulence(var0,dnamesvar,dvnamesvar)
1196        valsvar = ncvar.variables_values(depvars)
1197
1198        ncvar.insert_variable(ncobj, valsvar[0] + 'turb', diagout, diagoutd, 
1199          diagoutvd, newnc)
1200        varobj = newnc.variables[valsvar[0] + 'turb']
1201        attrv = varobj.long_name
1202        attr = varobj.delncattr('long_name')
1203        newattr = ncvar.set_attribute(varobj, 'long_name', attrv +                   \
1204          " Taylor decomposition turbulence term")
1205
1206# WRFbils fom WRF as HFX + LH
1207    elif diag == 'WRFbils':
1208           
1209        var0 = ncobj.variables[depvars[0]][:]
1210        var1 = ncobj.variables[depvars[1]][:]
1211
1212        diagout = var0 + var1
1213
1214        ncvar.insert_variable(ncobj, 'bils', diagout, dnames, dvnames, newnc)
1215
1216# WRFp pressure from WRF as P + PB
1217    elif diag == 'WRFp':
1218           
1219        diagout = WRFp
1220
1221        ncvar.insert_variable(ncobj, 'pres', diagout, dnames, dvnames, newnc)
1222
1223# WRFpos
1224    elif diag == 'WRFpos':
1225           
1226        dnamesvar = ncobj.variables['MAPFAC_M'].dimensions
1227        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1228
1229        ncvar.insert_variable(ncobj, 'WRFpos', WRFpos, dnamesvar, dvnamesvar, newnc)
1230
1231# WRFprw WRF water vapour path WRFdens, QVAPOR
1232    elif diag == 'WRFprw':
1233           
1234        var0 = WRFdens
1235        var1 = ncobj.variables[depvars[1]]
1236
1237        dnamesvar = list(var1.dimensions)
1238        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1239
1240        diagout, diagoutd, diagoutvd = compute_prw(var0, var1, dnamesvar,dvnamesvar)
1241
1242        ncvar.insert_variable(ncobj, 'prw', diagout, diagoutd, diagoutvd, newnc)
1243
1244# WRFrh (P, T, QVAPOR)
1245    elif diag == 'WRFrh':
1246           
1247        dnamesvar = list(ncobj.variables[depvars[2]].dimensions)
1248        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1249
1250        ncvar.insert_variable(ncobj, 'hus', WRFrh, dnames, dvnames, newnc)
1251
1252# WRFrhs (PSFC, T2, Q2)
1253    elif diag == 'WRFrhs':
1254           
1255        var0 = ncobj.variables[depvars[0]][:]
1256        var1 = ncobj.variables[depvars[1]][:]
1257        var2 = ncobj.variables[depvars[2]][:]
1258
1259        dnamesvar = list(ncobj.variables[depvars[2]].dimensions)
1260        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1261
1262        diagout, diagoutd, diagoutvd = compute_rh(var0,var1,var2,dnamesvar,dvnamesvar)
1263        ncvar.insert_variable(ncobj, 'huss', diagout, diagoutd, diagoutvd, newnc)
1264
1265# rvors (u10, v10, WRFpos)
1266    elif diag == 'WRFrvors':
1267           
1268        var0 = ncobj.variables[depvars[0]]
1269        var1 = ncobj.variables[depvars[1]]
1270
1271        diagout = rotational_z(var0, var1, distx)
1272
1273        dnamesvar = ncobj.variables[depvars[0]].dimensions
1274        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1275
1276        ncvar.insert_variable(ncobj, 'rvors', diagout, dnamesvar, dvnamesvar, newnc)
1277
1278# wss (u10, v10)
1279    elif diag == 'wss':
1280           
1281        var0 = ncobj.variables[depvars[0]][:]
1282        var1 = ncobj.variables[depvars[1]][:]
1283
1284        diagout = np.sqrt(var0*var0 + var1*var1)
1285
1286        dnamesvar = ncobj.variables[depvars[0]].dimensions
1287        dvnamesvar = ncvar.var_dim_dimv(dnamesvar,dnames,dvnames)
1288
1289        print 'dnamesvar',dnamesvar
1290        print 'dnames',dnames
1291        print 'dvnames',dvnames
1292        print 'dvnamesvar',dvnamesvar
1293
1294        ncvar.insert_variable(ncobj, 'wss', diagout, dnamesvar, dvnamesvar, newnc)
1295
1296    else:
1297        print errormsg
1298        print '  ' + main + ": diagnostic '" + diag + "' not ready!!!"
1299        print '    available diagnostics: ', availdiags
1300        quit(-1)
1301
1302    newnc.sync()
1303
1304#   end of diagnostics
1305
1306# Global attributes
1307##
1308atvar = ncvar.set_attribute(newnc, 'program', 'diagnostics.py')
1309atvar = ncvar.set_attribute(newnc, 'version', '1.0')
1310atvar = ncvar.set_attribute(newnc, 'author', 'Fita Borrell, Lluis')
1311atvar = ncvar.set_attribute(newnc, 'institution', 'Laboratoire Meteorologie ' +      \
1312  'Dynamique')
1313atvar = ncvar.set_attribute(newnc, 'university', 'Universite Pierre et Marie ' +     \
1314  'Curie -- Jussieu')
1315atvar = ncvar.set_attribute(newnc, 'centre', 'Centre national de la recherche ' +    \
1316  'scientifique')
1317atvar = ncvar.set_attribute(newnc, 'city', 'Paris')
1318atvar = ncvar.set_attribute(newnc, 'original_file', opts.ncfile)
1319
1320gorigattrs = ncobj.ncattrs()
1321
1322for attr in gorigattrs:
1323    attrv = ncobj.getncattr(attr)
1324    atvar = ncvar.set_attribute(newnc, attr, attrv)
1325
1326ncobj.close()
1327newnc.close()
1328
1329print '\n' + main + ': successfull writting of diagnostics file "' + ofile + '" !!!'
Note: See TracBrowser for help on using the repository browser.