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

Last change on this file since 1030 was 959, checked in by lfita, 8 years ago

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