Changeset 1493 in lmdz_wrf

Timestamp:

Apr 4, 2017, 6:52:22 PM (8 years ago)

Author:

lfita

Message:

Improving and flexibilizing `maskvar'

File:

• trunk/tools/nc_var_tools.py (modified) (11 diffs)

trunk/tools/nc_var_tools.py

@@ -3918 +3918 @@
 
 def maskvar(values, filename, varparameters):
-    """ Function to mask a variable using a mask. It is assumed that mask[...,dimM,dimJ,dimK] and 
-      var[...,dimM,dimJ,dimK] share the last dimensions
-    values= [maskfilename]:[maskvar]:[dimsmask]:[timename]
+    """ Function to mask a variable using another variable to mask it. Only need to share 
+      at least 1 dimension with the same size (no need same name)
+    values= [maskfilename]:[maskvarn]:[dimsmask]:[timename]
       [maskfilename]= file with the mask
-      [maskvar]= name of the variable with the mask
+      [maskvarn]= name of the variable with the mask
       [dimsmask]= ',' list of [dimn]|[sec] to use for slicing [maskvar]
         [sec]: section
@@ -3934 +3934 @@
     varparameters= [varn]|[dimsmask]
       [varn]: variable name
-      [samedimsmask]: ',' list of name of dimensions which fits [maskvar] slice
+      [samedimsmask]: ',' list of name of dimensions which fits [maskvarn] slice
     """
     fname='maskvar'
@@ -3943 +3943 @@
         quit()
 
+    expectargs = '[maskfilename]:[maskvarn]:[dimsmask]:[timename]'
+    gen.check_arguments(fname,values,expectargs,':')
+
+    maskfilename = values.split(':')[0]
+    maskvarn = values.split(':')[1]
+    dimsmask = values.split(':')[2]
+    maskvalue = values.split(':')[3]
+
+    varn = varparameters.split(':')[0]
+    samedimsmask = varparameters.split(':')[1].split(',')
+
     ofile = 'mask_' + varn + '.nc'
-
-    maskfilename = values.split(':')[0]
-    maskvar = values.split(':')[1]
-    dimsmask = values.split(':')[3]
-    maskvalue = values.split(':')[4]
-
-    varn = varparameters.splice(':')[0]
-    samedimsmask = varparameters.splice(':')[1].split(',')
 
     ncf = NetCDFFile(filename,'r')
@@ -3968 +3971 @@
     ncfmask = NetCDFFile(maskfilename,'r')
 
-    if not ncfmask.variables.has_key(maskvar):
+    if not ncfmask.variables.has_key(maskvarn):
         print errormsg
         print '    ' + fname + ": File  '" + maskfilename + "' does not have " +     \
-          " variable mask '" + maskvar + "' !!"
+          " variable mask '" + maskvarn + "' !!"
         print errormsg
         ncf.close()
@@ -3977 +3980 @@
         quit(-1)
 
-    varmaskobj = ncfmask.variables[maskvar]
+    varmaskobj = ncfmask.variables[maskvarn]
+    varmaskinf = variable_inf(varmaskobj)
+
     maskvalues, maskdims = slice_variable(varmaskobj, dimsmask.replace('|',':').     \
       replace(',','|'))
@@ -3988 +3993 @@
     testslicevar = []
     runshape = []
-    for dimn in varobj.dimensions.keys():
+    for dimn in varobj.dimensions:
         Ldim = len(ncf.dimensions[dimn])
         if gen.searchInlist(samedimsmask,dimn):
@@ -4022 +4027 @@
     elif varinf.dtype == type(np.float32(1.)) or varinf.dtype == type(np.float64(1.)):
         fillValue = gen.fillValueF64
-
+    else:
+        print errormsg
+        print '  ' + fname + ': type of variable:', varinf.dtype, 'not ready!!'
+        quit(-1)
+
+    maskval = gen.retype(maskvalue,varinf.dtype)
     newvar = ncfnew.createVariable(varn, varinf.dtype, varinf.dimns,                 \
       fill_value=fillValue)
@@ -4030 +4040 @@
 
     attr = newvar.setncattr('mask', 'variable masked using ' + values.split(':')[2] +\
-      ' of variable ' + maskvar + ' from file ' + maskfilename)
+      ' of variable ' + maskvarn + ' from file ' + maskfilename)
+    newvar[:] = varobj[:]
 
     ncfnew.sync()
@@ -4037 +4048 @@
     for varns in ncf.variables:
         if not varns == varn:
-            fvaradd(filename + ',' + varns, ofile)
+            try:
+                with gen.Capturing() as output:
+                    fvaradd(filename + ',' + varns, ofile)
+            except:
+                print errmsg
+                print 'fvaradd(' + filename + ', ' + varns + ',' + ofilen + ')'
+                for s1out in output: print s1out
+                quit(-1)
 
     ncfnew = NetCDFFile(ofile,'a')
@@ -4043 +4061 @@
 
     # Masking following that dimensions which are not coincident with the mask
-    maskval = gen.retype(maskvalue,varinf.dtype)
-    if runshape > 1:
-        idim = 0
-        runslice = list(runshape)
-        while idim <= runningsize:
-            slicevar = []
-            inosim = len(runshape)
-            for dimn in varobj.dimensions.keys():
-                Ldim = len(ncf.dimensions[dimn])
-                if gen.searchInlist(samedimsmask,dimn):
-                    slicevar.append(slice(0,Ldim,None))
-                else:
-                    slicevar.append()
-            testvar = varobj[tuple(testslicevar)]
-            
-
-    else:
-        maskedvals = ma.masked(varobj[:],maskval)
-
-    varobj[:] = maskedvals
-
-
-        
-
-
-
-    Nptmask = 1
-    for isdim in range(Nsamedims):
-        Nptmask = Nptmask * samedims[isdim]
-
-    Nvarloop = 1
-    varsecdims = []
-    for insdim in range(varinf.Ndims - Nsamedims):
-        Nvarloop = Nvarloop * varinf.dims[insdim]
-        varsecdims.append(0)
-
-    Nsecdims = len(varsecdims)
-
-    for iloop in range(Nvarloop):
-        varval = load_ncvariable_lastdims(ncf, varn, varsecdims, Nsamedims)
-        newvarval = np.where(maskvals == 1, varval, fillValue)
-        if np.sum(newvarval == fillValue) == Nptmask:
-            print errormsg
-            print '  ' + fname + ': all values got masked!!!'
-            print erromsg
-            quit(-1)
-
-        fill_ncvariable_lastdims(ncfnew,varn,varsecdims,Nsamedims,newvarval)
-
-        varsecdims[Nsecdims -1] = varsecdims[Nsecdims - 1] + 1
-        if varsecdims[Nsecdims - 1] > varinf.dims[Nsecdims - 1] - 1:
-            varsecdims[Nsecdims - 1] = 0
-            for jdim in range(Nsecdims - 1):
-                varsecdims[Nsecdims - 2 - jdim] = varsecdims[Nsecdims - 2 - jdim] + 1
-                if varsecdims[Nsecdims - 2 - jdim] <= varinf.dims[Nsecdims - 2 - jdim] - 1:
-                    break
-                else:
-                    varsecdims[Nsecdims - 2 - jdim] = 0
-                    
-    print 'variable: ', varinf.dims, ' and mask: ', varmaskinf.dims,' share: ', samedims, Nvarloop
+
+    slices = gen.provide_slices(varinf.dimns, varinf.dims, list(varmaskinf.dimns))
+
+    maskTOuse = ma.masked_equal(maskvalues, maskval)
+
+    for ir in range(runningsize):
+        islice = slices[ir]
+        varvals = varobj[tuple(islice)]
+        maskvarvals = np.where(maskTOuse.mask, fillValue, varvals)
+        newvar[tuple(islice)] = maskvarvals
 
     ncfnew.sync()
@@ -9267 +9235 @@
             dimnslice.append(vardims[idd])
 
-    print 'orig shape:',varobj[:].shape,'sliced:',varobj[tuple(varvalsdim)].shape
     varvalues = np.squeeze(varobj[tuple(varvalsdim)])
-    print 'varvalues shape:',varvalues.shape
 
     return varvalues, dimnslice