source: trunk/UTIL/PYTHON/planetoplot_v2/ppplot.py @ 911

###############################################
## PLANETOPLOT                               ##
## --> PPPLOT                                ##
###############################################
## Author: Aymeric Spiga. 02-03/2013         ##
###############################################
# python built-in librairies
import time as timelib
# added librairies
import numpy as np
import matplotlib.pyplot as mpl
from matplotlib.cm import get_cmap
from mpl_toolkits.basemap import Basemap
# personal librairies
import ppcompute
###############################################

#################################
# global variables and settings #
#################################

# matplotlib settings
# http://matplotlib.org/users/customizing.html
# -------------------------------
mpl.rcParams['font.family'] = "serif"
mpl.rcParams['axes.color_cycle'] = "r,g,b,k"
mpl.rcParams['contour.negative_linestyle'] = "dashed"
mpl.rcParams['verbose.level'] = "silent"
mpl.rcParams['lines.linewidth'] = 1.5
mpl.rcParams['lines.markersize'] = 10
mpl.rcParams['xtick.major.pad'] = 10
mpl.rcParams['ytick.major.pad'] = 10

# global variables
# -------------------------------
# - where settings files are located
#   (None means planetoplot_v2 in PYTHONPATH)
whereset = None
# - some good default settings.
# (bounds)
how_many_sigma = 3.0 
# (contours)
ccol = 'black'
cline = 0.55
# (vectors)
widthvec = 0.002
reducevec = 30.
# (colorbar)
zeorientation="vertical"
zefrac = 0.05
# (save figures)
pad_inches_value=0.25
###############################################

##########################
# executed when imported #
##########################
###########################################
# we load user-defined automatic settings #
###########################################
# initialize the warning variable about file not present...
files_not_present = ""
# ... and the whereset variable
whereset = ppcompute.findset(whereset)

# - variable settings
# -------------------------------
zefile = "set_var.txt"
vf = {} ; vc = {} ; vl = {}
try: 
    f = open(whereset+zefile, 'r')
    for line in f:
        if "#" in line: pass
        else:
            var, format, colorb, label = line.strip().split(';')
            ind = var.strip() 
            vf[ind] = format.strip()
            vc[ind] = colorb.strip()
            vl[ind] = label.strip()
    f.close()
except IOError: 
    files_not_present = files_not_present + zefile + " "

## - file settings
## -------------------------------
#zefile = "set_file.txt"
#prefix_t = {} ; suffix_t = {} ; name_t = {} ; proj_t = {} ; vecx_t = {} ; vecy_t = {}
#try:
#    f = open(whereset+zefile, 'r')
#    for line in f:
#        if "#" in line: pass
#        else:
#            prefix, suffix, name, proj, vecx, vecy = line.strip().split(';')
#            ind = name.strip() 
#            prefix_t[ind] = prefix.strip()
#            suffix_t[ind] = suffix.strip()
#            #name_t[ind] = name.strip()
#            proj_t[ind] = proj.strip()
#            vecx_t[ind] = vecx.strip()
#            vecy_t[ind] = vecy.strip()
#    f.close()
#except IOError:
#    files_not_present = files_not_present + zefile + " "

# - multiplot settings
# -------------------------------
zefile = "set_multiplot.txt"
subv_t = {} ; subh_t = {} ; wspace_t = {} ; hspace_t = {} ; font_t = {}
try:
    f = open(whereset+zefile, 'r')
    for line in f:
        if "#" in line: pass
        else:
            num, subv, subh, wspace, hspace, font = line.strip().split(';')
            ind = int(num.strip())
            subv_t[ind] = int(subv.strip())
            subh_t[ind] = int(subh.strip())
            wspace_t[ind] = float(wspace.strip())
            hspace_t[ind] = float(hspace.strip())
            font_t[ind] = float(font.strip())
    f.close()
except IOError:
    files_not_present = files_not_present + zefile + " "

# - background settings
# -------------------------------
zefile = "set_back.txt"
back = {}
try:
    f = open(whereset+zefile, 'r')
    for line in f:
        if "#" in line: pass
        else:
            name, link = line.strip().split(';')
            ind = name.strip() 
            back[ind] = link.strip()
    f.close()
except IOError:
    files_not_present = files_not_present + zefile + " "

# - area settings
# -------------------------------
zefile = "set_area.txt"
area = {}
try:
    f = open(whereset+zefile, 'r')
    for line in f:
        if "#" in line: pass
        else:
            name, wlat1, wlat2, wlon1, wlon2 = line.strip().split(';')
            area[name.strip()] = [[float(wlon1.strip()),float(wlon2.strip())],\
                                  [float(wlat1.strip()),float(wlat2.strip())]]
    f.close()
except IOError:
    files_not_present = files_not_present + zefile + " "
# A note to the user about missing files
if files_not_present != "":
    print "warning: files "+files_not_present+" not in "+whereset+" ; those presets will be missing"

# TBD: should change vector color with colormaps
#"gist_heat":    "white",\
#"hot":          "white",\
#"gray":         "red",\

####################
# useful functions #
####################

# a function to define subplot
# ... user can change settings in set_multiplot.txt read above
# -------------------------------
def definesubplot(numplot, fig):
    try: 
        mpl.rcParams['font.size'] = font_t[numplot]
    except: 
        mpl.rcParams['font.size'] = 18
    try: 
        fig.subplots_adjust(wspace = wspace_t[numplot], hspace = hspace_t[numplot])
        subv, subh = subv_t[numplot],subh_t[numplot]
    except: 
        print "!! WARNING !! no preset found from set_multiplot.txt, or this setting file was not found."
        subv = 1 ; subh = numplot
    return subv,subh

# a function to calculate automatically bounds (or simply prescribe those)
# -------------------------------
def calculate_bounds(field,vmin=None,vmax=None):
    # prescribed cases first
    zevmin = vmin
    zevmax = vmax
    # computed cases
    if zevmin is None or zevmax is None:
       # select values
       ind = np.where(field < 9e+35)
       fieldcalc = field[ ind ] # field must be a numpy array
       # calculate stdev and mean
       dev = np.std(fieldcalc)*how_many_sigma
       damean = ppcompute.mean(fieldcalc)
       # fill min/max if needed
       if vmin is None: zevmin = damean - dev
       if vmax is None: zevmax = damean + dev
       # special case: negative values with stddev while field is positive
       if zevmin < 0. and ppcompute.min(fieldcalc) > 0.: zevmin = 0.
    return zevmin, zevmax    
    #### treat vmin = vmax for continuity 
    #if vmin == vmax:  zevmin = damean - dev ; zevmax = damean + dev

# a function to solve the problem with blank bounds !
# -------------------------------
def bounds(what_I_plot,zevmin,zevmax,miss=9e+35):
    small_enough = 1.e-7
    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1.-small_enough) ] = zevmin*(1.-small_enough)
    else:          what_I_plot[ what_I_plot < zevmin*(1.+small_enough) ] = zevmin*(1.+small_enough)
    what_I_plot[ what_I_plot > miss  ] = -miss
    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1.-small_enough)
    return what_I_plot

# a generic function to show (GUI) or save a plot (PNG,EPS,PDF,...)
# -------------------------------
def save(mode="gui",filename="plot",folder="./",includedate=True,res=150):
    # a few settings
    possiblesave = ['eps','ps','svg','png','jpg','pdf'] 
    # now the main instructions
    if mode == "gui": 
        mpl.show()
    elif mode in possiblesave:
        ## name of plot
        name = folder+'/'+filename
        if includedate:
            for ttt in timelib.gmtime():
                name = name + "_" + str(ttt)
        name = name +"."+mode
        ## save file
        print "**** Saving file in "+mode+" format... Please wait."
        mpl.savefig(name,dpi=res,bbox_inches='tight',pad_inches=pad_inches_value)
    else:
        print "!! ERROR !! File format not supported. Supported: ",possiblesave

##################################
# a generic class to make a plot #
##################################
class plot():

    # print out a help string when help is invoked on the object
    # -------------------------------
    def __repr__(self):
        whatprint = 'plot object. \"help(plot)\" for more information\n'
        return whatprint

    # default settings
    # -- user can define settings by two methods. 
    # -- 1. yeah = plot2d(title="foo")
    # -- 2. yeah = pp() ; yeah.title = "foo"
    # -------------------------------
    def __init__(self,\
                 var=None,\
                 field=None,\
                 absc=None,\
                 xlabel="",\
                 ylabel="",\
                 div=20,\
                 logx=False,\
                 logy=False,\
                 swap=False,\
                 swaplab=True,\
                 invert=False,\
                 xcoeff=1.,\
                 ycoeff=1.,\
                 title=""):
        ## what could be defined by the user
        self.var = var
        self.field = field
        self.absc = absc
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.title = title
        self.div = div
        self.logx = logx
        self.logy = logy
        self.swap = swap
        self.swaplab = swaplab # NB: swaplab only used if swap=True
        self.invert = invert
        self.xcoeff = xcoeff
        self.ycoeff = ycoeff
        ## other useful arguments
        ## ... not used here in ppplot but need to be attached to plot object
        self.axisbg = "white"
        self.superpose = False

    # check
    # -------------------------------
    def check(self):
        if self.field is None: print "!! ERROR !! Please define a field to be plotted" ; exit()

    # define_from_var
    # ... this uses settings in set_var.txt
    # -------------------------------
    def define_from_var(self):
        if self.var is not None:
         if self.var.upper() in vl.keys():
          self.title = vl[self.var.upper()]

    # make
    # this is generic to all plots
    # -------------------------------
    def make(self):
        self.check()
        mpl.xlabel(self.xlabel)
        mpl.ylabel(self.ylabel)
        if self.swap:
         if self.swaplab:
           mpl.xlabel(self.ylabel)
           mpl.ylabel(self.xlabel)
        mpl.title(self.title)
        # if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
        if type(self.field).__name__ in 'MaskedArray':
            self.field[self.field.mask] = self.field.fill_value

################################
# a subclass to make a 1D plot #
################################
class plot1d(plot):

    # print out a help string when help is invoked on the object
    # -------------------------------
    def __repr__(self):
        whatprint = 'plot1d object. \"help(plot1d)\" for more information\n'
        return whatprint

    # default settings
    # -- user can define settings by two methods. 
    # -- 1. yeah = plot1d(title="foo")
    # -- 2. yeah = pp() ; yeah.title = "foo"
    # -------------------------------
    def __init__(self,\
                 lstyle='-',\
                 color='b',\
                 marker='x'):
        ## get initialization from parent class
        plot.__init__(self)
        ## what could be defined by the user
        self.lstyle = lstyle
        self.color = color
        self.marker = marker

    # define_from_var
    # ... this uses settings in set_var.txt
    # -------------------------------
    def define_from_var(self):
        # get what is done in the parent class
        plot.define_from_var(self)
        # add specific stuff
        if self.var is not None:
         if self.var.upper() in vl.keys():
          self.ylabel = vl[self.var.upper()]
          self.title = ""

    # make
    # -------------------------------
    def make(self):
        # get what is done in the parent class
        plot.make(self)
        # add specific stuff
        mpl.grid(color='grey')
        if self.lstyle == "": self.lstyle = " " # to allow for no line at all with ""
        # swapping if requested
        if self.swap:  x = self.field ; y = self.absc
        else:          x = self.absc ; y = self.field
        # coefficients on axis
        x=x*self.xcoeff ; y=y*self.ycoeff
        # check axis
        if x.size != y.size:
            print "!! ERROR !! x and y sizes don't match on 1D plot.", x.size, y.size
            exit()
        # make the 1D plot
        # either request linestyle or let matplotlib decide
        if self.lstyle is not None and self.color is not None:
            mpl.plot(x,y,self.color+self.lstyle,marker=self.marker)
        else:
            mpl.plot(x,y,marker=self.marker)
        # make log axes and/or invert ordinate
        # ... this must be after plot so that axis bounds are well-defined
        # ... also inverting must be after making the thing logarithmic
        if self.logx: mpl.semilogx()
        if self.logy: mpl.semilogy()
        if self.invert: ax = mpl.gca() ; ax.set_ylim(ax.get_ylim()[::-1])
        ## TBD: set with .div the number of ticks
        ## TBD: be able to control plot limits
        #ticks = self.div + 1
        #ax = mpl.gca()
        #ax.get_xaxis().set_ticks(np.linspace(ppcompute.min(x),ppcompute.max(x),ticks/2+1))

################################
# a subclass to make a 2D plot #
################################
class plot2d(plot):

    # print out a help string when help is invoked on the object
    # -------------------------------
    def __repr__(self):
        whatprint = 'plot2d object. \"help(plot2d)\" for more information\n'
        return whatprint

    # default settings
    # -- user can define settings by two methods. 
    # -- 1. yeah = plot2d(title="foo")
    # -- 2. yeah = pp() ; yeah.title = "foo"
    # -------------------------------
    def __init__(self,\
                 ordi=None,\
                 mapmode=False,\
                 proj="cyl",\
                 back=None,\
                 colorb="jet",\
                 trans=1.0,\
                 addvecx=None,\
                 addvecy=None,\
                 addcontour=None,\
                 fmt="%.2e",\
                 blon=None,\
                 blat=None,\
                 area=None,\
                 vmin=None,\
                 vmax=None,\
                 colorvec="black"):
        ## get initialization from parent class
        plot.__init__(self)
        ## what could be defined by the user
        self.ordi = ordi
        self.mapmode = mapmode
        self.proj = proj
        self.back = back
        self.colorb = colorb
        self.trans = trans
        self.addvecx = addvecx
        self.addvecy = addvecy
        self.colorvec = colorvec
        self.addcontour = addcontour
        self.fmt = fmt
        self.blon = blon ; self.blat = blat
        self.area = area
        self.vmin = vmin ; self.vmax = vmax

    # define_from_var
    # ... this uses settings in set_var.txt
    # -------------------------------
    def define_from_var(self):
        # get what is done in the parent class
        plot.define_from_var(self)
        # add specific stuff
        if self.var is not None:
         if self.var.upper() in vl.keys():
          self.colorb = vc[self.var.upper()]
          self.fmt = vf[self.var.upper()]

    # make
    # -------------------------------
    def make(self):
        # get what is done in the parent class...
        plot.make(self)
        # ... then add specific stuff
        ############################################################################################
        ### PRE-SETTINGS
        ############################################################################################
        # transposing if necessary
        shape = self.field.shape
        if shape[0] != shape[1]:
         if len(self.absc) == shape[0] and len(self.ordi) == shape[1]:
            print "!! WARNING !! Transposing axes"
            self.field = np.transpose(self.field)
        # bound field
        zevmin, zevmax = calculate_bounds(self.field,vmin=self.vmin,vmax=self.vmax)
        what_I_plot = bounds(self.field,zevmin,zevmax)
        # define contour field levels. define color palette
        ticks = self.div + 1
        zelevels = np.linspace(zevmin,zevmax,ticks)
        palette = get_cmap(name=self.colorb)
        # do the same thing for possible contourline entries
        if self.addcontour is not None:
            # if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
            if type(self.addcontour).__name__ in 'MaskedArray':
               self.addcontour[self.addcontour.mask] = self.addcontour.fill_value
            zevminc, zevmaxc = calculate_bounds(self.addcontour)
            what_I_contour = bounds(self.addcontour,zevminc,zevmaxc)
            ticks = self.div + 1
            zelevelsc = np.linspace(zevminc,zevmaxc,ticks)
        ############################################################################################
        ### MAIN PLOT
        ### NB: contour lines are done before contour shades otherwise colorar error
        ############################################################################################
        if not self.mapmode:
            ## A SIMPLE 2D PLOT
            ###################
            # swapping if requested
            if self.swap:  x = self.ordi ; y = self.absc
            else:          x = self.absc ; y = self.ordi
            # coefficients on axis
            x=x*self.xcoeff ; y=y*self.ycoeff
            # make shaded and line contours
            if self.addcontour is not None: 
                mpl.contour(x, y, what_I_contour, \
                            zelevelsc, colors = ccol, linewidths = cline)
            mpl.contourf(x, y, \
                         self.field, \
                         zelevels, cmap=palette)
            # make log axes and/or invert ordinate
            if self.logx: mpl.semilogx()
            if self.logy: mpl.semilogy()
            if self.invert: ax = mpl.gca() ; ax.set_ylim(ax.get_ylim()[::-1])
        else:
            ## A 2D MAP USING PROJECTIONS (basemap)
            #######################################
            mpl.xlabel("") ; mpl.ylabel("")
            # additional security in case self.proj is None here
            # ... we set cylindrical projection (the simplest one)
            if self.proj is None: self.proj = "cyl"
            # get lon and lat in 2D version.
            # (but first ensure we do have 2D coordinates)
            if self.absc.ndim == 1:     [self.absc,self.ordi] = np.meshgrid(self.absc,self.ordi)
            elif self.absc.ndim > 2:    print "!! ERROR !! lon and lat arrays must be 1D or 2D"
            # get lat lon intervals and associated settings
            wlon = [np.min(self.absc),np.max(self.absc)]
            wlat = [np.min(self.ordi),np.max(self.ordi)]
            # -- area presets are in set_area.txt
            if self.area is not None:
             if self.area in area.keys():
                wlon, wlat = area[self.area]
            # -- settings for meridians and parallels
            steplon = abs(wlon[1]-wlon[0])/6.
            steplat = abs(wlat[1]-wlat[0])/3.
            mertab = np.r_[wlon[0]:wlon[1]:steplon] ; merlab = [0,0,0,1]
            partab = np.r_[wlat[0]:wlat[1]:steplat] ; parlab = [1,0,0,0]
            format = '%.1f'
            # -- center of domain and bounding lats
            lon_0 = 0.5*(wlon[0]+wlon[1])
            lat_0 = 0.5*(wlat[0]+wlat[1])
            # some tests, bug fixes, and good-looking settings
            # ... cyl is good for global and regional
            if self.proj == "cyl":
                pass
            # ... global projections
            elif self.proj in ["ortho","moll","robin"]:
                wlat[0] = None ; wlat[1] = None ; wlon[0] = None ; wlon[1] = None
                steplon = 30. ; steplat = 30.
                if self.proj == "robin": steplon = 60.
                mertab = np.r_[-360.:360.:steplon]
                partab = np.r_[-90.:90.:steplat]
                if self.proj == "ortho": 
                    merlab = [0,0,0,0] ; parlab = [0,0,0,0]
                    # in ortho projection, blon and blat can be used to set map center
                    if self.blon is not None: lon_0 = self.blon
                    if self.blat is not None: lat_0 = self.blat
                elif self.proj == "moll":
                    merlab = [0,0,0,0]
                format = '%.0f'
            # ... regional projections
            elif self.proj in ["lcc","laea","merc"]:
                if self.proj in ["lcc","laea"] and wlat[0] == -wlat[1]: 
                    print "!! ERROR !! with Lambert lat1 must be different than lat2" ; exit()
                if wlat[0] < -80. and wlat[1] > 80.:
                    print "!! ERROR !! set an area (not global)" ; exit()
                format = '%.0f'
            elif self.proj in ["npstere","spstere"]:
                # in polar projections, blat gives the bounding lat
                # if not set, set something reasonable
                if self.blat is None:   self.blat = 60.
                # help the user who forgets self.blat would better be negative in spstere
                # (this actually serves for the default setting just above)
                if self.proj == "spstere" and self.blat > 0: self.blat = -self.blat
            # ... unsupported projections
            else:
                print "!! ERROR !! unsupported projection. supported: "+\
                      "cyl, npstere, spstere, ortho, moll, robin, lcc, laea, merc"
            # finally define projection
            m = Basemap(projection=self.proj,\
                        lat_0=lat_0,lon_0=lon_0,\
                        boundinglat=self.blat,\
                        llcrnrlat=wlat[0],urcrnrlat=wlat[1],\
                        llcrnrlon=wlon[0],urcrnrlon=wlon[1])
            # draw meridians and parallels
            ft = int(mpl.rcParams['font.size']*3./4.)
            m.drawmeridians(mertab,labels=merlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format)
            m.drawparallels(partab,labels=parlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format)
            # define background (see set_back.txt)
            if self.back is not None:
              if self.back in back.keys():
                 print "**** info: loading a background, please wait.",self.back
                 if back not in ["coast","sea"]:   m.warpimage(back[self.back],scale=0.75)
                 elif back == "coast":             m.drawcoastlines()
                 elif back == "sea":               m.drawlsmask(land_color='white',ocean_color='aqua')
              else:
                 print "!! ERROR !! requested background not defined. change name or fill in set_back.txt" ; exit()
            # define x and y given the projection
            x, y = m(self.absc, self.ordi)
            # contour field. first line contour then shaded contour.
            if self.addcontour is not None: 
                m.contour(x, y, what_I_contour, \
                            zelevelsc, colors = ccol, linewidths = cline)
            m.contourf(x, y, what_I_plot, zelevels, cmap = palette, alpha = self.trans)
        ############################################################################################
        ### COLORBAR
        ############################################################################################
        if self.trans > 0.:
            ## draw colorbar. settings are different with projections. or if not mapmode.
            if not self.mapmode: orientation=zeorientation ; frac = 0.075
            elif self.proj in ['moll']: orientation="horizontal" ; frac = 0.075
            elif self.proj in ['cyl']: orientation="vertical" ; frac = 0.023
            else: orientation = zeorientation ; frac = zefrac
            zelevpal = np.linspace(zevmin,zevmax,num=min([ticks/2+1,21]))
            zecb = mpl.colorbar(fraction=frac,pad=0.03,\
                                format=self.fmt,orientation=orientation,\
                                ticks=zelevpal,\
                                extend='neither',spacing='proportional') 
            if zeorientation == "horizontal": zecb.ax.set_xlabel(self.title) ; self.title = ""
        ############################################################################################
        ### VECTORS. must be after the colorbar. we could also leave possibility for streamlines.
        ############################################################################################
        ### not expecting NaN in self.addvecx and self.addvecy. masked arrays is just enough.
        stride = 3
        if self.addvecx is not None and self.addvecy is not None and self.mapmode:
                ## for metwinds only ???
                [vecx,vecy] = m.rotate_vector(self.addvecx,self.addvecy,self.absc,self.ordi) 
                # reference vector is scaled
                zescale = ppcompute.mean(np.sqrt(self.addvecx*self.addvecx+self.addvecy*self.addvecy))
                # make vector field
                q = m.quiver( x[::stride,::stride],y[::stride,::stride],\
                              vecx[::stride,::stride],vecy[::stride,::stride],\
                              angles='xy',color=self.colorvec,pivot='middle',\
                              scale=zescale*reducevec,width=widthvec )
                # make vector key. default is on upper left corner.
                p = mpl.quiverkey(q,-0.06,0.98,\
                                  zescale,str(int(zescale)),\
                                  color='black',labelpos='S',labelsep = 0.07)