Changeset 704 in lmdz_wrf for trunk

Timestamp:

Mar 23, 2016, 11:57:41 AM (9 years ago)

Author:

lfita

Message:

Adding `plot_basin' to plot basins (mostly from ORCDHIEE files)

Location:

trunk/tools

Files:

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

trunk/tools/drawing.py

@@ -27 +27 @@
 ## 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
+## e.g. # drawing.py -o draw_basins -f routing.py -S 'Y|y|nav_lat|-1,X|x|nav_lon|-1:1@1,rainbow,9:basins,-:cyl,l:ORCDHIEE river-basins:pdf:basins_named' -v nav_lon,nav_lat,trip,basins
 
 main = 'drawing.py'
@@ -36 +37 @@
 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_2D_shad_line_time', 'draw_barbs', 'draw_basins',                             \
   'draw_lines', 'draw_lines_time', 'draw_Neighbourghood_evol',                       \
   'draw_points', 'draw_points_lonlat',                                               \
@@ -3166 +3167 @@
             '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]
+            '3rdvar'@[colorbar]@[varn]@[units]: 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)
@@ -3327 +3328 @@
 
     return
+
+
+def draw_basins(ncfile, values, varns):
+    """ Function to plot wind basins
+      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]: [colorbar]@[varn]@[units]: 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= [lon],[lat],[outflow],[basinID] ',' list of the name of the variables with the lon,lat, the outflow and the basin ID
+    """
+    fname = 'draw_basins'
+
+    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)])
+
+        ivar = ivar + 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]
+    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[3]]
+    colorv = ocolvec[:]
+    stdvn, lonvn, unitsvn = drw.var_3desc(ocolvec)
+    colorvals = colorvals + '@' + stdvn + '@' + unitsvn
+
+    lengthv = vecvals.split(',')[2]
+
+# Vector labels
+    windname = windlabels.split(',')[0]
+    windunits = windlabels.split(',')[1]
+
+# Vector angles
+    oflow = ofile.variables[varNs[2]]
+    angle = (oflow[:] - 1)*np.pi/4
+    xflow = np.where(oflow[:] < 9, np.float(lengthv)*np.sin(angle), 0.)
+    yflow = np.where(oflow[:] < 9, np.float(lengthv)*np.cos(angle), 0.)
+
+    drw.plot_basins(lonvals, latvals, xflow, yflow, freqv, colorvals, colorv,      \
+      lengthv, windname, windunits, mapvalues, gtit, kindfig, figuren)
+
+    of.close()
+
+    return
 #quit()
 
@@ -3397 +3554 @@
 elif oper == 'draw_barbs':
     draw_barbs(opts.ncfile, opts.values, opts.varname)
+elif oper == 'draw_basins':
+    draw_basins(opts.ncfile, opts.values, opts.varname)
 elif oper == 'draw_Neighbourghood_evol':
     draw_Neighbourghood_evol(opts.ncfile, opts.values, opts.varname)

trunk/tools/drawing_tools.py

@@ -3 +3 @@
 # From http://stackoverflow.com/questions/13336823/matplotlib-python-error
 import numpy as np
+from matplotlib.pylab import *
 import matplotlib as mpl
 mpl.use('Agg')
@@ -6079 +6080 @@
         '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]
+        '3rdvar'@[colorbar]@[varn]@[units]: 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:
@@ -6210 +6211 @@
     return
 
+def plot_basins(xvals,yvals,fvals,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
+      fvals= values for the flow (1-8: N, NE, E, ..., NW; 97: sub-basin; 98: small to sea; 99: large to sea)
+      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 (as '3rdvar@[basinsvar]@[varn]@[units]' from plot_vector)
+      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_basins'
+
+    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
+
+# flow direction
+    angle = (fvals[::yfreq,::xfreq] - 1)*np.pi/4
+    uvals = np.where(fvals[::yfreq,::xfreq] < 9, np.float(veclength)*np.sin(angle), 0.)
+    vvals = np.where(fvals[::yfreq,::xfreq] < 9, np.float(veclength)*np.cos(angle), 0.)
+
+# Colors
+    vcolor = vecoln.split('@')[0]
+
+    plt.rc('text', usetex=True)
+
+    ddx = np.abs(xvals[dy/2+1,dx/2] - xvals[dy/2,dx/2])
+    ddy = np.abs(yvals[dy/2,dx/2+1] - yvals[dy/2,dx/2])
+    fontcharac = {'family': 'serif', 'weight': 'normal', 'size': 3}
+
+# Setting up colors for each label
+#   From: http://stackoverflow.com/questions/15235630/matplotlib-pick-up-one-color-associated-to-one-value-in-a-colorbar
+    my_cmap = plt.cm.get_cmap(vcolor)
+    vcolmin = np.min(veccolor[::yfreq,::xfreq])
+    vcolmax = np.max(veccolor[::yfreq,::xfreq])
+    vcolmin = 0.
+    vcolmax = 2500.
+    norm = mpl.colors.Normalize(vcolmin, vcolmax)
+    print 'min col:',vcolmin,'max col:',vcolmax
+
+    xlabpos = []
+    ylabpos = []
+    labels = []
+    labcol = []
+    flow = []
+    flowvals = []
+    
+    for j in range(0,dy,yfreq):
+        for i in range(0,dx,xfreq):
+            if veccolor[j,i] != '--':
+                xlabpos.append(xvals[j,i])
+                ylabpos.append(yvals[j,i])
+                labels.append(int(veccolor[j,i]))
+                labcol.append(my_cmap(norm(veccolor[j,i])))
+                flowvals.append(fvals[j,i])
+
+    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 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()
+
+    vN = vecoln.split('@')[1]
+    vU = vecoln.split('@')[2]
+    cbar.set_label(vN + ' (' + units_lunits(vU) + ')')
+
+    for i in range(len(xlabpos)):
+        plt.text(xlabpos[i]+0.5*ddx, ylabpos[i]+0.95*ddy, labels[i],                  \
+          color=labcol[i], fontdict=fontcharac)
+
+# Sea-flow
+    for i in range(len(xlabpos)):
+        if flowvals[i] == 97:
+            plt.plot(xlabpos[i], ylabpos[i], 'x', color=labcol[i])
+        elif flowvals[i] == 98:
+            plt.plot(xlabpos[i], ylabpos[i], '*', color=labcol[i])
+        elif flowvals[i] == 99:
+            plt.plot(xlabpos[i], ylabpos[i], 'h', color=labcol[i])
+
+    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