Changeset 575 in lmdz_wrf for trunk/tools

Timestamp:

Jul 7, 2015, 11:52:59 AM (10 years ago)

Author:

lfita

Message:

Adding `DataSetSection_multivars', to slice a dataset according to values from a variable

File:

• trunk/tools/nc_var_tools.py (modified) (1 diff)

trunk/tools/nc_var_tools.py

@@ -16509 +16509 @@
     return 
 
-vals='20.5:36.0:west_east,south_north:XLONG_M,XLAT_M'
-ncf='/home/lluis/etudes/domains/medic051215_2km/geo_em.d01.nc'
-
-getvalues_lonlat(vals, ncf)
-quit()
+#vals='20.5:36.0:west_east,south_north:XLONG_M,XLAT_M'
+#ncf='/home/lluis/etudes/domains/medic051215_2km/geo_em.d01.nc'
+
+def put_variable_slice(inc, onc, vn, sliced):
+    """ Function to add a variable from a netcdf object to a new one following a slice
+      inc= input netcdf object
+      onc= output netcdf object
+      vn= variable name to include
+      sliced= dictioary with the dimension name as key and [beg, end, int] as values
+    """
+    fname = 'put_variable_slice'
+
+    if not searchInlist(inc.variables,vn):
+        print errormsg
+        print ' ' + fname + ": variable '" + vn + "' is not in file!!"
+        quit(-1)
+
+    ov = inc.variables[vn]
+
+    dnvs = ov.dimensions
+    varslice = []
+# Getting variable's dimensions
+    for dnv in dnvs:
+        if not searchInlist(onc.dimensions,dnv):
+            od = inc.dimensions[dnv]
+            if od.isunlimited():
+                dsize = None
+            else:
+                dsize = len(od)
+            newdim = onc.createDimension(dnv, dsize)
+        if searchInlist(sliced.keys(),dnv):
+            varslice.append(slice(sliced[dnv][0], sliced[dnv][1], sliced[dnv][2]))
+        else:
+            varslice.append(slice(0,len(inc.dimensions[dnv])))
+# Getting variable
+    varattrs = ov.ncattrs()
+    if searchInlist(varattrs, '_FillValue'):
+        varfil = ov._FillValue
+    else:
+        varfil = False
+    vartype = ov.dtype
+
+    newvar = onc.createVariable(vn, vartype, dnvs, fill_value=varfil)
+#    print 'Lluis shapes newvar:',newvar.shape,'slice:',varslice
+    newvar[:] = ov[tuple(varslice)]
+
+    for attrs in varattrs:
+        if not attrs == '_FillValue':
+            attrv = ov.getncattr(attrs)
+            attr = set_attribute(newvar, attrs, attrv)
+
+    return
+
+def DataSetSection_multivars(values, filen, varn):
+    """ Function to get a section (values along multiple variables) of a given data-set
+      values= [varn1],[beg1],[end1],[int1]@[...[[varnM],[begM],[endM],[intM]]]
+        [varni]: name of the variable ('WRFt', for WRF time varibale)
+        [begi],[endi],[inti]: beginning, end and interval along the dimension-axis
+          [endi] = -1, maximum value
+          [inti] = -1, all the values within the range
+          NOTE: dimensions without section by the variables are taken allover their size
+      filen= netCDF with the data-set
+      varn= ',' list of variables ('all', for all variables)
+    """
+    import numpy.ma as ma
+    fname = 'DataSetSection_multivars'
+
+# Variables not to check in the file
+    NOcheck = ['WRFt']
+
+    if values == 'h':
+        print fname + '_____________________________________________________________'
+        print DataSetSection_multivars.__doc__
+        quit()
+
+    Nvars = len(values.split('@'))
+
+# Slicing variables
+    varns = []
+    ofiletile = ''
+    slicevalS = ''
+
+    begvs = np.zeros((Nvars), dtype=int)
+    endvs = np.zeros((Nvars), dtype=int)
+    intvs = np.zeros((Nvars), dtype=int)
+
+    for ivar in range(Nvars):
+        val = values.split('@')[ivar]
+        varns.append(val.split(',')[0])
+        begvs[ivar] = np.float(val.split(',')[1])
+        endvs[ivar] = np.float(val.split(',')[2])
+        intvs[ivar] = np.float(val.split(',')[3])
+        ofiletile = ofiletile + '_' + varns[ivar] + '_B' + str(begvs[ivar]) + '-E' + \
+          str(endvs[ivar]) + '-I' + str(intvs[ivar])
+        if intvs[ivar] != -1:
+            slicevalS = slicevalS + varns[ivar] + ' (' + str(begvs[ivar]) + ',' +    \
+              str(endvs[ivar]) + ',' + str(intvs[ivar]) + '); '
+        else:
+            slicevalS = slicevalS + varns[ivar] + ' (' + str(begvs[ivar]) + ',' +    \
+              str(endvs[ivar]) + ',1); '
+
+    ofile = ofile=filen.split('.')[0] + ofiletile + '.nc'
+
+    nciobj = NetCDFFile(filen,'r')
+    ncoobj = NetCDFFile(ofile,'w')
+
+# Output variables
+##
+    if varn == 'all':
+        variables = nciobj.variables
+    else:
+        if varn.find(',') != -1:
+            variables = varn.split(',')
+        else:
+            variables = [varn]
+
+# Looking for limits due to the variables
+##
+    for ivar in range(Nvars):
+        print '  ' + fname + ": masking with '" + varns[ivar] + "':",  begvs[ivar],  \
+          ',', endvs[ivar], ',', intvs[ivar]
+        if not searchInlist(nciobj.variables,varns[ivar]) and not                    \
+          searchInlist(NOcheck,varns[ivar]):
+            print errormsg
+            print ' ' + fname + ": variable '" + varns[ivar] + "' is not in ' +      \
+              'input file!!"
+            quit(-1)
+
+        if varns[ivar] == 'WRFt':
+            ovar = nciobj.variables['Times'][:]
+            vals = []
+            for it in range(ovar.shape[0]):
+                vals.append(datetimeStr_conversion(ovar[it,:], 'WRFdatetime',        \
+                  'YmdHMS'))
+                vals[it] = int(vals[it])
+                begvs[ivar] = int(begvs[ivar])
+                endvs[ivar] = int(endvs[ivar])
+                intvs[ivar] = int(intvs[ivar])
+
+            vals = np.array(vals)
+            vardims = ['Time']
+        else:
+            ovar = nciobj.variables[varns[ivar]]
+            vals = ovar[:]
+            vardims = ovar.dimensions
+
+        if endvs[ivar] == -1: endvs[ivar] = np.max(vals)
+        maskinf = ma.masked_less(vals, begvs[ivar])
+        masksup = ma.masked_greater(vals, endvs[ivar])
+
+        if intvs[ivar] != -1:
+            print errormsg
+            print '  ' + fname + ': non-consecutive slices not ready!!'
+            quit(-1)
+        finalmask = maskinf + masksup
+        idn = 0
+
+        slicedims = {}
+        
+        for dn in vardims:
+            ddn = len(nciobj.dimensions[dn])
+            dinds = np.arange(ddn)
+            rightvals = np.where(finalmask.mask, False, True)
+            slicedims[dn]=[np.min(dinds[rightvals]), np.max(dinds[rightvals])+1, None]
+
+        print '  ' + fname + ': slicing variables with______'
+        print slicedims
+ 
+# Creating dimensions
+## 
+    print '  ' + fname + ': adding dimensions...'
+    for nd in slicedims.keys():
+        objdim = nciobj.dimensions[nd]
+        if objdim.isunlimited():
+            dimsize = None
+        else:
+            dimsize = slicedims[nd][1] - slicedims[nd][0] + 1
+        newdim = ncoobj.createDimension(nd, dimsize)
+    ncoobj.sync()
+
+    for var in variables:
+        put_variable_slice(nciobj, ncoobj, var, slicedims)
+
+# Global attributes
+## 
+    for attrs in nciobj.ncattrs():
+        attrv = nciobj.getncattr(attrs)
+        attr = set_attribute(ncoobj, attrs, attrv)
+
+    attr = set_attribute(ncoobj, 'sliced_variables', slicevalS)
+
+    nciobj.close()
+
+    ncoobj.sync()
+    ncoobj.close()
+
+    print '  ' + fname + ' succesfull creation of file "' + ofile + '" !!!'
+
+    return
+
+#DataSetSection_multivars('XTIME,4380,4800,-1','/home/lluis/PY/wrfout_d01_2001-11-11_00:00:00','all')
+#DataSetSection_multivars('WRFt,20011111060000,20011111180000,-1','/home/lluis/PY/wrfout_d01_2001-11-11_00:00:00','all')
+
+#quit()
 
 """operation to make: