Changeset 698 in lmdz_wrf

Timestamp:

Mar 10, 2016, 6:55:00 PM (9 years ago)

Author:

lfita

Message:

Adding 'draw_vectors' and 'plot_vector'

Location:

trunk/tools

Files:

• drawing.py (modified) (5 diffs)
• drawing_tools.py (modified) (5 diffs)

trunk/tools/drawing.py

@@ -26 +26 @@
 ## e.g. # drawing.py -o draw_ptZvals -f geo_v2_2012102123_RR1.nc -S 'pracc:lon,lat:o:80:2,42,7,47,:values!of!values:Blues:cyl,l:pdf' -v pr
 ## e.g. # drawing.py -f carteveg5km.nc -o draw_points_lonlat -S 'longitude:latitude:pdf:points!veget|type:green:.:0.5:None:0:legend'
+## e.g. # drawing.py -o draw_vectors -f wrfout_d01_2001-11-11_00:00:00 -S 'T|Time|Times|2,Y|south_north|XLAT|-1,X|west_east|XLONG|-1:3@3,wind@rainbow,9:10m wind,ms-1:cyl,l:WRF 10 m winds:pdf:winds' -v U10,V10
 
 main = 'drawing.py'
@@ -35 +36 @@
 namegraphics = ['create_movie', 'draw_2D_shad', 'draw_2D_shad_time',                 \
   'draw_2D_shad_cont', 'draw_2D_shad_cont_time', 'draw_2D_shad_line',                \
-  'draw_2D_shad_line_time', 'draw_barbs', 'draw_points', 'draw_points_lonlat',       \
+  'draw_2D_shad_line_time', 'draw_barbs',                                            \
+  'draw_lines', 'draw_lines_time', 'draw_Neighbourghood_evol',                       \
+  'draw_points', 'draw_points_lonlat',                                               \
   'draw_ptZvals', 'draw_timeSeries',                                                 \
   'draw_topo_geogrid',                                                               \
   'draw_topo_geogrid_boxes', 'draw_trajectories', 'draw_vals_trajectories',          \
-  'draw_lines', 'draw_lines_time', 'draw_Neighbourghood_evol', 'list_graphics',      \
-  'variable_values']
+  'draw_vectors',  'list_graphics', 'variable_values']
 
 def draw_2D_shad(ncfile, values, varn):
@@ -1295 +1297 @@
             uwvals = np.squeeze(np.array(ovarN[tuple(varslice)]))
         elif ivar == 3:
-            vwvals = np.squeeze(ovarN[tuple(varslice)])
+            vwvals = np.squeeze(ovarN[tuple(varslice)])            
 
         ivar = ivar + 1
@@ -3147 +3149 @@
 
 #draw_ptZvals('OBSnetcdf.nc', 'pracc:lon,lat:o:80:2,42,7,47,:values!of!values:Blues:cyl,l:pdf', 'pr')
+
+def draw_vectors(ncfile, values, varns):
+    """ Function to plot wind vectors
+      values= [dimname]|[vardimname]|[value]:[vecvals]:[windlabs]:[mapvalues]:
+        [gtit]:[kindfig]:[figuren]
+        'X/Y/Z/T'|[dimname]|[vardimname]|[value]: ',', list for each basic dimension '|' separated of:
+          [dimname]: name of the dimension in the file
+          [vardimname]: name of the variable with the values for the dimension in the file
+          [value]: which value of the given dimension (-1, all; [ibeg]@[iend], i-range beginning, end)
+          No value takes all the range of the dimension
+        [vecvals]= [frequency],[color],[length]
+          [frequency]: [xfreq]@[yfreq] frequency of values allong each axis ('None', all grid points; 
+            'auto', computed automatically to have 20 vectors along each axis)
+          [color]: color of the vectors
+            'singlecol'@[colorn]: all the vectors same color ('auto': for 'red')
+            'wind'@[colorbar]: color of the vectors according to wind speed sqrt(u^2+v^2) and given [colorbar]
+              all vectors the same length
+            '3rdvar'@[colorbar]: color of the vectors according to a 3rd variable (to be added at -v) and given [colorbar]
+              all vectors the same length
+          [length]: length of the wind vectors ('auto', for 9)
+        [windlabs]= [windname],[windunits]
+          [windname]: name of the wind variable in the graph
+          [windunits]: units of the wind variable in the graph ('None', for the value in the file)
+        [mapvalues]= map characteristics: [proj],[res]
+          see full documentation: http://matplotlib.org/basemap/
+          [proj]: projection
+            * 'cyl', cilindric
+            * 'lcc', lambert conformal
+          [res]: resolution:
+            * 'c', crude
+            * 'l', low
+            * 'i', intermediate
+            * 'h', high
+            * 'f', full
+        gtit= title of the graph ('|', for spaces)
+        kindfig= kind of figure
+        figuren= name of the figure
+      ncfile= file to use
+      varns= [uwind],[ywind] ',' list of the name of the variables with the u-wind,y-wind component
+    """
+    fname = 'draw_vectors'
+
+    if values == 'h':
+        print fname + '_____________________________________________________________'
+        print draw_vectors.__doc__
+        quit()
+
+    expectargs = '[X/Y/Z/T]|[dimname]|[vardimname]|[value]:[vecvals]:[windlabs]:' +  \
+      '[mapvalues]:[gtit]:[kindfig]:[figuren]'
+ 
+    drw.check_arguments(fname,values,expectargs,':')
+
+    dimvals = values.split(':')[0]
+    vecvals = values.split(':')[1]
+    windlabels = values.split(':')[2]
+    mapvalues = values.split(':')[3]
+    gtit = values.split(':')[4]
+    kindfig = values.split(':')[5]
+    figuren = values.split(':')[6]
+
+    of = NetCDFFile(ncfile,'r')
+
+    dims = {}
+    for dimv in dimvals.split(','):
+        dns = dimv.split('|')
+        dims[dns[0]] = [dns[1], dns[2], dns[3]]
+
+    varNs = []
+    for dn in dims.keys():
+        if dn == 'X':
+            varNs.append(dims[dn][1])
+            dimx = len(of.dimensions[dims[dn][0]])
+        elif dn == 'Y':
+            varNs.append(dims[dn][1])
+            dimy = len(of.dimensions[dims[dn][0]])
+
+    ivar = 0
+    for wvar in varns.split(','):
+        if not drw.searchInlist(of.variables.keys(), wvar):
+            print errormsg
+            print '  ' + fname + ": file does not have variable '" + wvar + "' !!"
+            quit(-1)
+        if ivar == 0:
+            varNs.append(wvar)
+        else:
+            varNs.append(wvar)
+
+    ivar = 0
+    for varN in varNs:
+        varslice = []
+
+        ovarN = of.variables[varN]
+        vard = ovarN.dimensions
+        for vdn in vard:
+            found = False
+            for dd in dims.keys():
+                if dims[dd][0] == vdn:
+                    if dims[dd][2].find('@') != -1:
+                        rvals = dims[dd][2].split('@')
+                        varslice.append(slice(int(rvals[0]), int(rvals[1])))
+                    elif dims[dd][2] == '-1':
+                        varslice.append(slice(0,len(of.dimensions[dims[dd][0]])))
+                    else:
+                        varslice.append(int(dims[dd][2]))
+
+                    found = True
+                    break
+            if not found:
+                varslice.append(slice(0,len(of.dimensions[dims[dd][0]])))
+
+        if varN == dims['X'][1]:
+            lonvals0 = np.squeeze(ovarN[tuple(varslice)])
+        elif varN == dims['Y'][1]:
+            latvals0 = np.squeeze(ovarN[tuple(varslice)])
+        elif ivar == 2:
+            uwvals = np.squeeze(np.array(ovarN[tuple(varslice)]))
+        elif ivar == 3:
+            vwvals = np.squeeze(ovarN[tuple(varslice)])
+
+        ivar = ivar + 1
+
+#    print 'Final shapes:',lonvals0.shape,':',latvals0.shape,':',uwvals.shape,':',
+#      vwvals.shape
+
+    if len(uwvals.shape) != 2 or len(vwvals.shape) != 2:
+        print errormsg
+        print '  ' + fname + ': wrong size of the wind fields! they must be ' +      \
+          '2-dimensional!'
+        print '    u-winds shape:',uwvals.shape,'dims:',of.variables[varNs[2]]
+        print '    v-winds shape:',vwvals.shape,'dims:',of.variables[varNs[3]]
+        print '      provide more values for their dimensions!!'
+        quit(-1)
+
+    if len(lonvals0.shape) == 1:
+        lonvals, latvals = np.meshgrid(lonvals0, latvals0)
+    else:
+        lonvals = lonvals0
+        latvals = latvals0
+
+# Vector values
+    if vecvals.split(',')[0] == 'None':
+        freqv = None
+    else:
+        freqv = vecvals.split(',')[0] 
+    colorvals = vecvals.split(',')[1]
+    coln = colorvals.split('@')[0]
+    colv = colorvals.split('@')[1]
+    if coln == 'singlecol':
+        colorv = colv
+    elif coln == 'wind':
+        colorv = np.sqrt(uwvals**2 + vwvals**2)
+    elif coln == '3rdvar':
+        if len(varn.split(',')) != 3:
+            print errormsg
+            print '  ' + fname + ": color of vectors should be according to '" +     \
+              coln + "' but a third varibale is not provided !!"
+            quit(-1)
+        ocolvec = of.variables[varNs[4]]
+        colorv = ocolvec[:]
+        stdvn, lonvn, unitsvn = drw.var_3desc(ocolvec)
+        colorvals = colorvals + '@' + stdvn + '@' + unitsvn
+    else:
+        print errormsg
+        print '  ' + fname + ": color type '" + coln + "' not ready !!"
+        quit(-1)
+
+    lengthv = vecvals.split(',')[2]
+
+# Vector labels
+    windname = windlabels.split(',')[0]
+    windunits = windlabels.split(',')[1]
+
+    drw.plot_vector(lonvals, latvals, uwvals, vwvals, freqv, colorvals, colorv,      \
+      lengthv, windname, windunits, mapvalues, gtit, kindfig, figuren)
+
+    of.close()
+
+    return
 #quit()
 
@@ -3239 +3419 @@
 elif oper == 'draw_vals_trajectories':
     draw_vals_trajectories(opts.ncfile, opts.values, opts.varname)
+elif oper == 'draw_vectors':
+    draw_vectors(opts.ncfile, opts.values, opts.varname)
 elif oper == 'list_graphics':
 # From: http://www.diveintopython.net/power_of_introspection/all_together.html

trunk/tools/drawing_tools.py

@@ -44 +44 @@
 # DegGradSec_deg:
 # intT2dt:
+# lonlat2D: Function to return lon, lat 2D matrices from any lon,lat matrix
 # lonlat_values:
 # date_CFtime:
@@ -1492 +1493 @@
         elif u == 'deg K': lu='$K$'
         elif u == 'degK': lu='$K$'
+        elif u == 'g/g': lu='$gg^{-1}$'
         elif u == 'hours': lu='$hour$'
         elif u == 'J/kg': lu='$Jkg^{-1}$'
@@ -1536 +1538 @@
         elif u == 'seconds': lu='$second$'
         elif u == '%': lu='\%'
+        elif u == '?': lu='-'
         else:
             print errormsg
@@ -5794 +5797 @@
 def plot_barbs(xvals,yvals,uvals,vvals,vecfreq,veccolor,veclength,windn,wuts,mapv,graphtit,kfig,figname):
     """ Function to plot wind barbs
-      xvals= values for the 'x-axis'
-      yvals= values for the 'y-axis'
+      xvals= x position of the values
+      yvals= y position of the values
+      uvals= values for the x-wind
+      uvals= values for the y-wind
       vecfreq= [xfreq],[yfreq] frequency of values allong each axis (None, all grid points; 
         'auto', computed automatically to have 20 vectors along each axis)
@@ -6061 +6066 @@
 
     return
+
+def plot_vector(xvals,yvals,uvals,vvals,vecfreq,vecoln,veccolor,veclength,windn,wuts,\
+  mapv,graphtit,kfig,figname):
+    """ Function to plot vectors
+      xvals= values for the x-axis
+      yvals= values for the y-axis
+      uvals= values for the x-wind
+      vvals= values for the y-wind
+      vecfreq= [xfreq],[yfreq] frequency of values allong each axis (None, all grid points; 
+        'auto', computed automatically to have 20 vectors along each axis)
+      veccoln= name for the color of the vectors
+        'singlecol'@[colorn]: all the vectors same color ('auto': for 'red')
+        'wind'@[colorbar]: color of the vectors according to wind speed sqrt(u^2+v^2) and given [colorbar]
+        '3rdvar'@[colorbar]: color of the vectors according to a 3rd variable (to be added at -v) and given [colorbar]
+      veccolor= color of the vectors
+      veclength= length of the wind vectors:
+        'singlecol': 'auto', for 9
+        'wind' and '3rdvar': 'auto' length as wind speed, otherwise fix length
+      windn= name of the wind variable in the graph
+      wuts= units of the wind variable in the graph
+      mapv= map characteristics: [proj],[res]
+        see full documentation: http://matplotlib.org/basemap/
+        [proj]: projection
+          * 'cyl', cilindric
+          * 'lcc', lambert conformal
+        [res]: resolution:
+          * 'c', crude
+          * 'l', low
+          * 'i', intermediate
+          * 'h', high
+          * 'f', full
+      graphtit= title of the graph ('|', for spaces)
+      kfig= kind of figure
+      figname= name of the figure
+    """
+    fname = 'plot_vector'
+
+    dx=xvals.shape[1]
+    dy=xvals.shape[0]
+
+# Frequency of vectors
+    if vecfreq is None:
+        xfreq = 1
+        yfreq = 1
+    elif vecfreq == 'auto':
+        xfreq = dx/20
+        yfreq = dy/20
+    else:
+        xfreq=int(vecfreq.split('@')[0])
+        yfreq=int(vecfreq.split('@')[1])
+
+# Vector length
+    if veclength == 'auto':
+        vlength = 9
+    else:
+        vlength = veclength
+
+# Colors
+    VecN = vecoln.split('@')[0]
+    if VecN == 'singlecol':
+        if veccolor == 'auto':
+            vcolor = "red"
+        else:
+            vcolor = veccolor
+    elif VecN == 'wind' or VecN == '3rdvar':
+        vcolor = vecoln.split('@')[1]
+
+    plt.rc('text', usetex=True)
+
+    if not mapv is None:
+        lon00 = np.where(xvals[:] < 0., 360. + xvals[:], xvals[:])
+        lat00 = yvals[:]
+
+        map_proj=mapv.split(',')[0]
+        map_res=mapv.split(',')[1]
+
+        nlon = np.min(xvals[::yfreq,::xfreq])
+        xlon = np.max(xvals[::yfreq,::xfreq])
+        nlat = np.min(yvals[::yfreq,::xfreq])
+        xlat = np.max(yvals[::yfreq,::xfreq])
+
+        lon2 = xvals[dy/2,dx/2]
+        lat2 = yvals[dy/2,dx/2]
+
+        print 'lon2:', lon2, 'lat2:', lat2, 'SW pt:', nlon, ',', nlat, 'NE pt:',     \
+          xlon, ',', xlat
+
+        if map_proj == 'cyl':
+            m = Basemap(projection=map_proj, llcrnrlon=nlon, llcrnrlat=nlat,         \
+              urcrnrlon=xlon, urcrnrlat= xlat, resolution=map_res)
+        elif map_proj == 'lcc':
+            m = Basemap(projection=map_proj, lat_0=lat2, lon_0=lon2, llcrnrlon=nlon, \
+              llcrnrlat=nlat, urcrnrlon=xlon, urcrnrlat= xlat, resolution=map_res)
+        else:
+            print errormsg
+            print '  ' + fname + ": projection '" + map_proj + "' not ready!!"
+            print '    projections available: cyl, lcc'
+            quit(-1)
+
+        m.drawcoastlines()
+
+        meridians = pretty_int(nlon,xlon,5)
+        m.drawmeridians(meridians,labels=[True,False,False,True],color="black")
+
+        parallels = pretty_int(nlat,xlat,5)
+        m.drawparallels(parallels,labels=[False,True,True,False],color="black")
+
+        plt.xlabel('W-E')
+        plt.ylabel('S-N')
+
+    if VecN == 'singlecol':
+        plt.quiver(xvals[::yfreq,::xfreq], yvals[::yfreq,::xfreq],                   \
+          uvals[::yfreq,::xfreq], vvals[::yfreq,::xfreq], color=vcolor, pivot='middle')
+        plt.annotate(windn.replace('_','\_') +' (' + units_lunits(wuts) + ')',       \
+          xy=(0.80,-0.15), xycoords='axes fraction', color=vcolor)
+    else:
+        if veclength != 'auto':
+            wind = np.sqrt(uvals**2 + vvals**2)
+            uvals = np.where(wind == 0., 0., uvals)
+            vvals = np.where(wind == 0., 0., vvals)
+            uvals = np.float(veclength)*uvals/wind
+            vvals = np.float(veclength)*vvals/wind
+
+        plt.quiver(xvals[::yfreq,::xfreq], yvals[::yfreq,::xfreq],                   \
+          uvals[::yfreq,::xfreq], vvals[::yfreq,::xfreq], 
+          veccolor[::yfreq,::xfreq], cmap=plt.get_cmap(vcolor), pivot='middle')
+        cbar = plt.colorbar()
+
+        if VecN == 'wind':
+            cbar.set_label('$\sqrt{u^{2} + v^{2}}$ (' + units_lunits(wuts) + ')')
+        else:
+            vN = vecoln.split('@')[2]
+            vU = vecoln.split('@')[3]
+            cbar.set_label(vN + ' (' + units_lunits(vU) + ')')
+
+        plt.annotate(windn.replace('_','\_') +' (' + units_lunits(wuts) + ')',       \
+          xy=(0.80,-0.15), xycoords='axes fraction', color='black')
+
+    plt.title(graphtit.replace('|',' ').replace('&','\&'))
+
+    output_kind(kfig, figname, True)
+
+    return
+
+def var_3desc(ovar):
+    """ Function to provide std_name, long_name and units from an object variable
+      ovar= object variable
+    """
+    fname = 'var_desc'
+    varattrs = ovar.ncattrs()
+
+    if searchInlist(varattrs,'std_name'):
+        stdn = ovar.getncattr('std_name')
+    else:
+        vvalues = variables_values(ovar._name)
+        stdn = vvalues[1]
+
+    if searchInlist(varattrs,'long_name'):
+        lonn = ovar.getncattr('long_name')
+    else:
+        lonn = vvalues[4].replace('|',' ')
+
+    if searchInlist(varattrs,'units'):
+        un = ovar.getncattr('units')
+    else:
+        un = vvalues[5]
+
+    return stdn, lonn, un