source: lmdz_wrf/trunk/tools/multi_plot.py @ 1333

##### Python to plot multiple plots at the same time in a panel
# L. Fita, LMD September 2016
## e.g. #  multi_plot.py -f multi_plot.dat -o multi_plot -d True -s True
import os
import subprocess as sub
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import matplotlib.pyplot as plt
import nc_var_tools as ncvar
import drawing_tools as drw
import generic_tools as gen
import drawing as DrW
import matplotlib as mpl
from optparse import OptionParser

main='multi_plot.py'

errormsg='ERROR -- error -- ERROR -- error'
warnmsg='WARNING -- warning -- WARNING -- warning'

class values_figure(object):
    """ Class with the values to plot a figure
      self.files= list with the files to use
      self.vals= values to use for the plot
      self.vars= list with the variables to use
    """

    def __init__(self,files,vals,varns):
        self.files = files
        self.vals = vals
        self.vars = varns

####### ###### ##### #### ### ## #
fix_multiplot_conf = ['title','kind','panel']

parser = OptionParser()
parser.add_option("-d", "--debug", dest="debug", 
  help="debug ('true'/'false')", metavar="VALUE")
parser.add_option("-f", "--file_configuration", dest="fileconf", help="configuration file to use", 
  metavar="FILE")
parser.add_option("-o", "--output_name", dest="oname", 
  help="name of the output figure (without extension)", metavar="LABEL")
parser.add_option("-s", "--seeFigure", dest="seefig", 
  help="see resultan figure ('true'/'false')", metavar="VALUE")

(opts, args) = parser.parse_args()

#######    #######
## MAIN
    #######

if opts.debug is None:
    print warnmsg
    print '  ' + main + ": no debug provided!!"
    print "    give a value for the option '-d/--debug'"
    print "       debug ('true'/'false')"
    print " Assuming 'false'"
    dbg=False
else:
    dbg=gen.Str_Bool(opts.debug)

if opts.fileconf is None:
    print errormsg
    print '  ' + main + ": no configuration file provided!!"
    print "    give a value for the option '-f/--file_configuration'"
    print "      configuration file to use"
    quit(-1)

if not os.path.isfile(opts.fileconf):
    print errormsg
    print '  '+main+": configuration file '" + opts.fileconf + "' does not exist!!"
    quit(-1)

if opts.oname is None:
    print warnmsg
    print '  ' + main + ": no output figure name provided!!"
    print "    give a value for the option '-o/--output_name'"
    print "       name of the output figure (without extension)"
    print " Assuming 'multi_plot'"
    oname='multi_plot'
else:
    oname=opts.oname

if opts.seefig is None:
    print warnmsg
    print '  ' + main + ": no seeing figure provided!!"
    print "    give a value for the option '-s/--seeFigure'"
    print "      see resultan figure ('true'/'false')"
    print " Assuming 'false'"
    seefig=False
else:
    seefig=gen.Str_Bool(opts.seefig)


# Reading configuration file
# File strcture must be as follows
#    fig_title: Title of the figure ('!') for spaces  (None for no title)
#    fig_panel: [Ncol],[Nrow] panel of figures with Ncolumns and Nrows
#    fig_kind: kind of the output figure (None, png)
#    kind: [graphic_label] (from drawing.py)
#    file: [file1],[file2],...,[fileN] (',' list of files if required)
#    values: Specific values for this graphic (in same format as in [graphic_label]
#    variables: [var1],[var2],....,[varN] (',' list of variables if required)
# Figures will be draw consecutively as they appear in the file

# Dictionary with the figures to plot and their values kepts in a class
figures = {}

# List of the names of the figures as they appear in the file
figns = []

# List of the names of the variables as they appear in the file
varns = []

# Dictionary with the number of times that a given graphic appears in the panel
Nfigs = {}

oconf = open(opts.fileconf)
kindn = None
kindnow = None
variables = None
for line in oconf:
    if line[0:1] != '#' and len(line) > 1 and line.find(':') != -1:
        label = line.split(':')[0].replace(' ','')
        vals = line.split(':')[1].replace('\n','')
        if label == 'fig_title': figures['title'] = vals
        elif label == 'fig_panel': 
            vals0 = vals.replace(' ','').split(',')
            figures['panel']=[int(vals0[0]), int(vals0[1])]
        elif label == 'fig_kind': 
            if vals == 'None': figures['kind'] = png
            else: figures['kind'] = vals.replace(' ','')
        elif label == 'kind': kindn = vals.replace(' ','') + '@0'
        elif label == 'file': files = gen.str_list(vals.replace(' ',''),',')
        elif label == 'values': 
            Lvalues = len('values:')
            Lline = len(line)
            values = line[Lvalues:Lline - Lvalues + 1]
        elif label == 'variables': variables = gen.str_list(vals.replace(' ',''),',')
        else:
            print errormsg
            print '  ' + main + ": configuration label '" + label + "' not ready !!"
            quit(-1)

        if kindn is not None and kindn != kindnow and variables is not None:
            if kindnow is None: 
                kindnow = kindn.split('@')[0]
            else:
                kindnow = kindnow.split('@')[0]

            if dbg:
                print '    ' + main + ": adding in panel: '" + kindnow + "'... "
                print '      files:', files
                print '      values:', values
                print '      variables:', variables

            varns.append(':'.join(values))
            if Nfigs.has_key(kindnow):
                N = Nfigs[kindnow]
                Nfigs[kindnow] = N + 1
            else:
                Nfigs[kindnow] = 1

            labf = kindnow + '@' + str(Nfigs[kindnow])
            figns.append(labf)
            if dbg:
                print "    added as '" + labf + "'"

            figures[labf] = values_figure(files,values,variables)
                
            kindnow = kindn+'@-1'
            variables = None

oconf.close()

if dbg:
    print '  ' + main + " General figure configuration _______"
    for ff in fix_multiplot_conf:
        print '    ' + ff + ': ', figures[ff]

    for ff in figns:
        if not gen.searchInlist(fix_multiplot_conf,ff):
            print "    '" + ff + "' _______"
            gen.printing_class(figures[ff])

panelval = figures['panel']
if np.prod(panelval) != len(figns):
    print errormsg
    print '  ' + main + ': then number of plots=',len(figns),'not suits the panel=', \
      figures['panel']
    print '    figures to plot:',figns
    quit(-1)

#
## Plotting
#
plt.rc('text', usetex=True)

fig, axes = plt.subplots(panelval[0],panelval[1])

ifig = 1
for ff in figns:
    isubplot = gen.index_flatten_mat(ifig,list(panelval))
    isubplot = isubplot + 1
    if dbg:
        print '  ' + main + " adding in graph '" + ff + "' in ", isubplot, "..."
    graphinf = figures[ff]
    graphlab = ff.split('@')[0]
    plt.subplot(panelval[0],panelval[1],ifig)

    if graphlab == 'draw_2D_shad':
        Files = graphinf.files[0]
        Vals = graphinf.vals+':False'
        Vars = graphinf.vars[0]

        DrW.draw_2D_shad(Files,Vals,Vars)

    elif graphlab == 'direc_draw_2D_shad':
        Files = graphinf.files[0]
        if not os.path.isfile(Files):
            print errormsg
            print '  '+main+": netcdf file with data '" + Files + "' does not exist!!"
            quit(-1)

        Vals = graphinf.vals.split(':')

        ncf = NetCDFFile(Files,'r')
        ovar = ncf.variables[graphinf.vars[0]]

        dictslice = {}
        for dmn in Vals[1].split(','):
            dn = dmn.split('|')[0]
            dv = int(dmn.split('|')[1])
            dictslice[dn] = dv

        slicevar, slicedims = ncvar.SliceVarDict(ovar,dictslice)

        Vars = ovar[tuple(slicevar)]
        ncf.close()

        CFname = Vals[0].replace(' ','')
        dimxn = Vals[2]
        dimyn = Vals[3]
        colorbar = Vals[4]
        minv = np.float(Vals[5].split(',')[0])
        maxv = np.float(Vals[5].split(',')[1])
        titgraph = Vals[6].replace('|',' ')

        varvu = gen.variables_values(CFname)[5]
        dimxu = gen.variables_values(dimxn)[5]
        dimyu = gen.variables_values(dimyn)[5]

        # Plot itself
        plt.rc('text', usetex=True)
        plt.pcolormesh(Vars, cmap=plt.get_cmap(colorbar), vmin=minv, vmax=maxv)
        cbar = plt.colorbar()
        cbar.set_label(CFname + ' (' + drw.units_lunits(varvu) + ')')
        plt.title(titgraph)
        plt.xlabel(dimxn + ' (' + drw.units_lunits(dimxu) + ')')
        plt.ylabel(dimyn + ' (' + drw.units_lunits(dimyu) + ')')

    else:
        print errormsg
        print '  ' + main + ": kind of figure '" + graphlab + "' not ready !!"
    ifig = ifig+1

fig.suptitle(figures['title'].replace('!',' '))

figname = opts.oname + '.' + figures['kind']
sub.call('rm ' + figname + ' >& /dev/null', shell=True)
plt.savefig(figname)
print main + " succesfull drawing of figure '" + figname + "' !!"
if seefig:
    if figures['kind'] == 'pdf':
        sub.call('evince ' + figname + ' &', shell=True)
    else:
        sub.call('display ' + figname + ' &', shell=True)