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

Last change on this file since 1098 was 1093, checked in by aslmd, 11 years ago

PLANETOPLOT. fixed a bug in case lower case is used in set_var.txt

File size: 36.9 KB
Line 
1###############################################
2## PLANETOPLOT                               ##
3## --> PPPLOT                                ##
4###############################################
5## Author: Aymeric Spiga. 02-03/2013         ##
6###############################################
7# python built-in librairies
8import time as timelib
9# added librairies
10import numpy as np
11import matplotlib.pyplot as mpl
12from matplotlib.cm import get_cmap
13from mpl_toolkits.basemap import Basemap
14from matplotlib.ticker import FormatStrFormatter,MaxNLocator
15# personal librairies
16import ppcompute
17###############################################
18
19#################################
20# global variables and settings #
21#################################
22
23# matplotlib settings
24# http://matplotlib.org/users/customizing.html
25# -------------------------------
26mpl.rcParams['font.family'] = "serif"
27mpl.rcParams['axes.color_cycle'] = "b,r,g,k"
28mpl.rcParams['contour.negative_linestyle'] = "dashed" # ou "solid"
29mpl.rcParams['verbose.level'] = "silent"
30mpl.rcParams['lines.linewidth'] = 1.5
31mpl.rcParams['lines.markersize'] = 10
32mpl.rcParams['xtick.major.pad'] = 10
33mpl.rcParams['ytick.major.pad'] = 10
34
35# global variables
36# -------------------------------
37# - where settings files are located
38#   (None means planetoplot_v2 in PYTHONPATH)
39whereset = None
40# - some good default settings.
41# (bounds)
42how_many_sigma = 3.0 
43# (contours)
44ccol = 'black'
45cline = 0.55
46cline = 0.8
47# (vectors)
48widthvec = 0.002
49reducevec = 30.
50# (colorbar)
51zeorientation="vertical"
52zefrac = 0.05
53# (save figures)
54pad_inches_value=0.25
55# (negative mode)
56def_negative = False
57# (xkcd mode)
58def_xkcd = False
59###############################################
60
61### settings for 'negative-like' mode
62if def_negative:
63    mpl.rc('figure', facecolor='k', edgecolor='k')
64    mpl.rcParams['text.color'] = 'w'
65    mpl.rc('axes',labelcolor='w',facecolor='k',edgecolor='w')
66    mpl.rcParams['xtick.color'] = 'w'
67    mpl.rcParams['ytick.color'] = 'w'
68    mpl.rcParams['grid.color'] = 'w'
69    mpl.rc('savefig',facecolor='k',edgecolor='k')
70
71### settings for xkcd mode (only with matplotlib 1.3)
72### ... you must have Humori-Sans Font installed
73if def_xkcd:
74    mpl.xkcd()
75
76##########################
77# executed when imported #
78##########################
79###########################################
80# we load user-defined automatic settings #
81###########################################
82# initialize the warning variable about file not present...
83files_not_present = ""
84# ... and the whereset variable
85whereset = ppcompute.findset(whereset)
86
87# - variable settings
88# -------------------------------
89zefile = "set_var.txt"
90vf = {} ; vc = {} ; vl = {} ; vu = {}
91try: 
92    f = open(whereset+zefile, 'r')
93    for line in f:
94        if "#" in line: pass
95        else:
96            var, format, colorb, label, units = line.strip().split(';')
97            ind = var.strip().upper()
98            vf[ind] = format.strip()
99            vc[ind] = colorb.strip()
100            vl[ind] = label.strip()
101            vu[ind] = units.strip()
102    f.close()
103except IOError: 
104    files_not_present = files_not_present + zefile + " "
105
106## - file settings
107## -------------------------------
108#zefile = "set_file.txt"
109#prefix_t = {} ; suffix_t = {} ; name_t = {} ; proj_t = {} ; vecx_t = {} ; vecy_t = {}
110#try:
111#    f = open(whereset+zefile, 'r')
112#    for line in f:
113#        if "#" in line: pass
114#        else:
115#            prefix, suffix, name, proj, vecx, vecy = line.strip().split(';')
116#            ind = name.strip()
117#            prefix_t[ind] = prefix.strip()
118#            suffix_t[ind] = suffix.strip()
119#            #name_t[ind] = name.strip()
120#            proj_t[ind] = proj.strip()
121#            vecx_t[ind] = vecx.strip()
122#            vecy_t[ind] = vecy.strip()
123#    f.close()
124#except IOError:
125#    files_not_present = files_not_present + zefile + " "
126
127# - multiplot settings
128# -------------------------------
129zefile = "set_multiplot.txt"
130subv_t = {} ; subh_t = {} ; wspace_t = {} ; hspace_t = {} ; font_t = {}
131try:
132    f = open(whereset+zefile, 'r')
133    for line in f:
134        if "#" in line: pass
135        else:
136            num, subv, subh, wspace, hspace, font = line.strip().split(';')
137            ind = int(num.strip())
138            subv_t[ind] = int(subv.strip())
139            subh_t[ind] = int(subh.strip())
140            wspace_t[ind] = float(wspace.strip())
141            hspace_t[ind] = float(hspace.strip())
142            font_t[ind] = float(font.strip())
143    f.close()
144except IOError:
145    files_not_present = files_not_present + zefile + " "
146
147# - background settings
148# -------------------------------
149zefile = "set_back.txt"
150back = {}
151try:
152    f = open(whereset+zefile, 'r')
153    for line in f:
154        if "#" in line: pass
155        else:
156            name, link = line.strip().split(';')
157            ind = name.strip() 
158            back[ind] = link.strip()
159    f.close()
160except IOError:
161    files_not_present = files_not_present + zefile + " "
162
163# - area settings
164# -------------------------------
165zefile = "set_area.txt"
166area = {}
167try:
168    f = open(whereset+zefile, 'r')
169    for line in f:
170        if "#" in line: pass
171        else:
172            name, wlat1, wlat2, wlon1, wlon2 = line.strip().split(';')
173            area[name.strip()] = [[float(wlon1.strip()),float(wlon2.strip())],\
174                                  [float(wlat1.strip()),float(wlat2.strip())]]
175    f.close()
176except IOError:
177    files_not_present = files_not_present + zefile + " "
178# A note to the user about missing files
179if files_not_present != "":
180    print "warning: files "+files_not_present+" not in "+whereset+" ; those presets will be missing"
181
182# TBD: should change vector color with colormaps
183#"gist_heat":    "white",\
184#"hot":          "white",\
185#"gray":         "red",\
186
187####################
188# useful functions #
189####################
190
191# continuity with matplotlib
192def close():
193    mpl.close()
194
195# a function for predefined figure sizes
196def figuref(x=16,y=9):
197    fig = mpl.figure(figsize=(x,y))
198    return fig
199
200# a function to change color lines according to a color map (idea by A. Pottier)
201# ... two optional arguments: color maps + a number telling how much intervals are needed
202def rainbow(cb="jet",num=8):
203    ax = mpl.gca()
204    pal = mpl.cm.get_cmap(name=cb)
205    ax._get_lines.set_color_cycle([pal(i) for i in np.linspace(0,0.9,num)])
206
207# a function to define subplot
208# ... user can change settings in set_multiplot.txt read above
209# -------------------------------
210def definesubplot(numplot, fig):
211    try: 
212        mpl.rcParams['font.size'] = font_t[numplot]
213    except: 
214        mpl.rcParams['font.size'] = 18
215    try: 
216        fig.subplots_adjust(wspace = wspace_t[numplot], hspace = hspace_t[numplot])
217        subv, subh = subv_t[numplot],subh_t[numplot]
218    except: 
219        print "!! WARNING !! (ignore if superpose mode) no preset found from set_multiplot.txt, or this setting file was not found."
220        subv = 1 ; subh = numplot
221    return subv,subh
222
223# a function to calculate automatically bounds (or simply prescribe those)
224# -------------------------------
225def calculate_bounds(field,vmin=None,vmax=None):
226    # prescribed cases first
227    zevmin = vmin
228    zevmax = vmax
229    # computed cases
230    if zevmin is None or zevmax is None:
231       # select values
232       ind = np.where(field < 9e+35)
233       fieldcalc = field[ ind ] # field must be a numpy array
234       # calculate stdev and mean
235       dev = np.std(fieldcalc)*how_many_sigma
236       damean = ppcompute.mean(fieldcalc)
237       # fill min/max if needed
238       if vmin is None: zevmin = damean - dev
239       if vmax is None: zevmax = damean + dev
240       # special case: negative values with stddev while field is positive
241       if zevmin < 0. and ppcompute.min(fieldcalc) >= 0.: zevmin = 0.
242    # check that bounds are not too tight given the field
243    amin = ppcompute.min(field)
244    amax = ppcompute.max(field)
245    if np.abs(amin) < 1.e-15: 
246        cmin = 0.
247    else:
248        cmin = 100.*np.abs((amin - zevmin)/amin)
249    cmax = 100.*np.abs((amax - zevmax)/amax)
250    if cmin > 150. or cmax > 150.:
251        print "!! WARNING !! Bounds are a bit too tight. Might need to reconsider those."
252        print "!! WARNING !! --> actual",amin,amax,"adopted",zevmin,zevmax
253    return zevmin, zevmax   
254    #### treat vmin = vmax for continuity
255    #if vmin == vmax:  zevmin = damean - dev ; zevmax = damean + dev
256
257# a function to solve the problem with blank bounds !
258# -------------------------------
259def bounds(what_I_plot,zevmin,zevmax,miss=9e+35):
260    small_enough = 1.e-7
261    if zevmin < 0: what_I_plot[ what_I_plot < zevmin*(1.-small_enough) ] = zevmin*(1.-small_enough)
262    else:          what_I_plot[ what_I_plot < zevmin*(1.+small_enough) ] = zevmin*(1.+small_enough)
263    what_I_plot[ what_I_plot > miss  ] = -miss
264    what_I_plot[ what_I_plot > zevmax ] = zevmax*(1.-small_enough)
265    return what_I_plot
266
267# a function to change labels with modulo
268# ---------------------------------------
269def labelmodulo(ax,mod):
270    mpl.draw()
271    strtab = []
272    for tick in ax.get_xaxis().get_ticklabels():
273        onetick = tick.get_text()
274        if len(onetick) > 0:
275          num = float(onetick)
276          strtab.append(num % mod)
277    ax.get_xaxis().set_ticklabels(strtab)
278    return ax
279
280# a function to output an ascii file
281# ----------------------------------
282def writeascii (field=None,absc=None,name=None):
283    field = np.array(field)
284    absc = np.array(absc)
285    if name is None:
286        name = "prof"
287        for ttt in timelib.gmtime():
288            name = name + "_" + str(ttt)
289    if field is None or absc is None:
290        print "!! WARNING !! Not printing the file, incorrect field or absc."
291    else:
292        if field.ndim == 1:
293            myfile = open(name, 'w')
294            for ix in range(len(absc)):
295                myfile.write("%15.5e%15.5e\n" % ( absc[ix], field[ix] ))
296            myfile.close()
297        else:
298            print "!! WARNING !! Not printing the file, 2D fields not supported yet."
299    return
300
301# a generic function to show (GUI) or save a plot (PNG,EPS,PDF,...)
302# -------------------------------
303def save(mode="gui",filename="plot",folder="./",includedate=True,res=150,custom=False):
304    if mode != "nothing":
305      # a few settings
306      possiblesave = ['eps','ps','svg','png','jpg','pdf'] 
307      # now the main instructions
308      if mode == "gui": 
309          mpl.show()
310      elif mode in possiblesave:
311          ## name of plot
312          name = folder+'/'+filename
313          if includedate:
314              for ttt in timelib.gmtime():
315                  name = name + "_" + str(ttt)
316          name = name +"."+mode
317          ## save file
318          print "**** Saving file in "+mode+" format... Please wait."
319          if not custom:
320              # ... regular plots
321              mpl.savefig(name,dpi=res,pad_inches=pad_inches_value,bbox_inches='tight')
322          else:
323              # ... mapping mode, adapted space for labels etc...
324              mpl.savefig(name,dpi=res)
325      else:
326          print "!! ERROR !! File format not supported. Supported: ",possiblesave
327
328# a necessary addition for people used to matplotlib
329def show():
330    mpl.show()
331
332##################################
333# a generic class to make a plot #
334##################################
335class plot():
336
337    # print out a help string when help is invoked on the object
338    # -------------------------------
339    def __repr__(self):
340        whatprint = 'plot object. \"help(plot)\" for more information\n'
341        return whatprint
342
343    # default settings
344    # -- user can define settings by two methods.
345    # -- 1. yeah = plot2d(title="foo")
346    # -- 2. yeah = pp() ; yeah.title = "foo"
347    # -------------------------------
348    def __init__(self,\
349                 var=None,\
350                 f=None,\
351                 x=None,\
352                 xlabel="",\
353                 ylabel="",\
354                 div=20,\
355                 logx=False,\
356                 logy=False,\
357                 swap=False,\
358                 swaplab=True,\
359                 invert=False,\
360                 xcoeff=1.,\
361                 ycoeff=1.,\
362                 fmt=None,\
363                 colorbar="jet",\
364                 units="",\
365                 modx=None,\
366                 xmin=None,\
367                 ymin=None,\
368                 xmax=None,\
369                 ymax=None,\
370                 nxticks=10,\
371                 nyticks=10,\
372                 cbticks=None,\
373                 title=""):
374        ## what could be defined by the user
375        self.var = var
376        self.f = f
377        self.x = x
378        self.xlabel = xlabel
379        self.ylabel = ylabel
380        self.title = title
381        self.div = div
382        self.logx = logx
383        self.logy = logy
384        self.swap = swap
385        self.swaplab = swaplab # NB: swaplab only used if swap=True
386        self.invert = invert
387        self.xcoeff = xcoeff
388        self.ycoeff = ycoeff
389        self.fmt = fmt
390        self.units = units
391        self.colorbar = colorbar
392        self.modx = modx
393        self.xmin = xmin
394        self.ymin = ymin
395        self.xmax = xmax
396        self.ymax = ymax
397        self.nxticks = nxticks
398        self.nyticks = nyticks
399        self.cbticks = cbticks
400        ## other useful arguments
401        ## ... not used here in ppplot but need to be attached to plot object
402        self.axisbg = "white"
403        self.superpose = False
404
405    # check
406    # -------------------------------
407    def check(self):
408        if self.f is None: print "!! ERROR !! Please define a field to be plotted" ; exit()
409        else: self.f = np.array(self.f) # ensure this is a numpy array
410
411    # define_from_var
412    # ... this uses settings in set_var.txt
413    # -------------------------------
414    def define_from_var(self):
415        if self.var is not None:
416         if self.var.upper() in vl.keys():
417          self.title = vl[self.var.upper()]
418          self.units = vu[self.var.upper()]
419
420    # make
421    # this is generic to all plots
422    # -------------------------------
423    def make(self):
424        self.check()
425        # labels, title, etc...
426        mpl.xlabel(self.xlabel)
427        mpl.ylabel(self.ylabel)
428        if self.swap:
429         if self.swaplab:
430           mpl.xlabel(self.ylabel)
431           mpl.ylabel(self.xlabel)
432        mpl.title(self.title)
433        # if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
434        if type(self.f).__name__ in 'MaskedArray':
435            self.f[self.f.mask] = self.f.fill_value
436
437################################
438# a subclass to make a 1D plot #
439################################
440class plot1d(plot):
441
442    # print out a help string when help is invoked on the object
443    # -------------------------------
444    def __repr__(self):
445        whatprint = 'plot1d object. \"help(plot1d)\" for more information\n'
446        return whatprint
447
448    # default settings
449    # -- user can define settings by two methods.
450    # -- 1. yeah = plot1d(title="foo")
451    # -- 2. yeah = pp() ; yeah.title = "foo"
452    # -------------------------------
453    def __init__(self,\
454                 linestyle=None,\
455                 color=None,\
456                 marker='x',\
457                 legend=None):
458        ## get initialization from parent class
459        plot.__init__(self)
460        ## what could be defined by the user
461        self.linestyle = linestyle
462        self.color = color
463        self.marker = marker
464        self.legend = legend
465
466    # define_from_var
467    # ... this uses settings in set_var.txt
468    # -------------------------------
469    def define_from_var(self):
470        # get what is done in the parent class
471        plot.define_from_var(self)
472        # add specific stuff
473        if self.var is not None:
474         if self.var.upper() in vl.keys():
475          self.ylabel = vl[self.var.upper()] + " (" + vu[self.var.upper()] + ")"
476          self.title = ""
477          self.fmt = vf[self.var.upper()]
478
479    # make
480    # -------------------------------
481    def make(self):
482        # get what is done in the parent class
483        plot.make(self)
484        if self.fmt is None: self.fmt = '%.0f'
485        # add specific stuff
486        mpl.grid(color='grey')
487        if self.linestyle == "": self.linestyle = " " # to allow for no line at all with ""
488        # set dummy x axis if not defined
489        if self.x is None: 
490            self.x = np.array(range(self.f.shape[0]))
491            print "!! WARNING !! dummy coordinates on x axis"
492        else:
493            self.x = np.array(self.x) # ensure this is a numpy array
494        # swapping if requested
495        if self.swap:  x = self.f ; y = self.x
496        else:          x = self.x ; y = self.f
497        # coefficients on axis
498        x=x*self.xcoeff ; y=y*self.ycoeff
499        # check axis
500        if x.size != y.size:
501            print "!! ERROR !! x and y sizes don't match on 1D plot.", x.size, y.size
502            exit()
503        # make the 1D plot
504        # either request linestyle or let matplotlib decide
505        if self.linestyle is not None and self.color is not None:
506            mpl.plot(x,y,self.color+self.linestyle,marker=self.marker,label=self.legend)
507        elif self.color is not None:
508            mpl.plot(x,y,color=self.color,marker=self.marker,label=self.legend)
509        elif self.linestyle is not None:
510            mpl.plot(x,y,linestyle=self.linestyle,marker=self.marker,label=self.legend)
511        else:
512            mpl.plot(x,y,marker=self.marker,label=self.legend)
513        # AXES
514        ax = mpl.gca()
515        # make log axes and/or invert ordinate
516        # ... this must be after plot so that axis bounds are well-defined
517        # ... also inverting must be after making the thing logarithmic
518        if self.logx: mpl.xscale("log") # not mpl.semilogx() because excludes log on y
519        if self.logy: mpl.yscale("log") # not mpl.semilogy() because excludes log on x
520        if self.invert: ax.set_ylim(ax.get_ylim()[::-1])
521        if self.xmin is not None and self.xmax is not None:
522          if self.xmin > self.xmax:
523            ax.set_xlim(ax.get_xlim()[::-1])
524            self.xmin,self.xmax = self.xmax,self.xmin
525        # add a label for line(s)
526        if self.legend is not None:
527            if self.legend != "":
528                mpl.legend(loc="best",fancybox=True)
529        # format labels
530        if self.swap: ax.xaxis.set_major_formatter(FormatStrFormatter(self.fmt))
531        else: ax.yaxis.set_major_formatter(FormatStrFormatter(self.fmt))
532        # plot limits: ensure that no curve would be outside the window
533        # x-axis
534        x1, x2 = ax.get_xbound()
535        xmin = ppcompute.min(x)
536        xmax = ppcompute.max(x)
537        if xmin < x1: x1 = xmin
538        if xmax > x2: x2 = xmax
539        if self.xmin is not None: x1 = self.xmin
540        if self.xmax is not None: x2 = self.xmax
541        ax.set_xbound(lower=x1,upper=x2)
542        # y-axis
543        y1, y2 = ax.get_ybound()
544        ymin = ppcompute.min(y)
545        ymax = ppcompute.max(y)
546        if ymin < y1: y1 = ymin
547        if ymax > y2: y2 = ymax
548        if self.ymin is not None: y1 = self.ymin
549        if self.ymax is not None: y2 = self.ymax
550        ax.set_ybound(lower=y1,upper=y2)
551        ## set with .div the number of ticks.
552        if not self.logx:
553            ax.xaxis.set_major_locator(MaxNLocator(self.nxticks))
554        else:
555            print "!! WARNING. in logx mode, ticks are set automatically."
556        if not self.logy:
557            ax.yaxis.set_major_locator(MaxNLocator(self.nyticks))
558        else:
559            print "!! WARNING. in logy mode, ticks are set automatically."
560        ## specific modulo labels
561        if self.modx is not None:
562            ax = labelmodulo(ax,self.modx)
563
564    # makeshow = make + show
565    # -------------------------------
566    def makeshow(self):
567        self.make()
568        mpl.show()
569
570################################
571# a subclass to make a 2D plot #
572################################
573class plot2d(plot):
574
575    # print out a help string when help is invoked on the object
576    # -------------------------------
577    def __repr__(self):
578        whatprint = 'plot2d object. \"help(plot2d)\" for more information\n'
579        return whatprint
580
581    # default settings
582    # -- user can define settings by two methods.
583    # -- 1. yeah = plot2d(title="foo")
584    # -- 2. yeah = pp() ; yeah.title = "foo"
585    # -------------------------------
586    def __init__(self,\
587                 y=None,\
588                 mapmode=False,\
589                 proj="cyl",\
590                 back=None,\
591                 trans=1.0,\
592                 vx=None,\
593                 vy=None,\
594                 c=None,\
595                 blon=None,\
596                 blat=None,\
597                 area=None,\
598                 vmin=None,\
599                 vmax=None,\
600                 showcb=True,\
601                 wscale=None,\
602                 svx=1,\
603                 svy=1,\
604                 leftcorrect=False,\
605                 colorvec="black"):
606        ## get initialization from parent class
607        plot.__init__(self)
608        ## what could be defined by the user
609        self.y = y
610        self.mapmode = mapmode
611        self.proj = proj
612        self.back = back
613        self.trans = trans
614        self.vx = vx
615        self.vy = vy
616        self.colorvec = colorvec
617        self.c = c
618        self.blon = blon ; self.blat = blat
619        self.area = area
620        self.vmin = vmin ; self.vmax = vmax
621        self.showcb = showcb
622        self.wscale = wscale
623        self.svx = svx
624        self.svy = svy
625        self.leftcorrect = leftcorrect
626
627    # define_from_var
628    # ... this uses settings in set_var.txt
629    # -------------------------------
630    def define_from_var(self):
631        # get what is done in the parent class
632        plot.define_from_var(self)
633        # add specific stuff
634        if self.var is not None:
635         if self.var.upper() in vl.keys():
636          self.colorbar = vc[self.var.upper()]
637          self.fmt = vf[self.var.upper()]
638
639    # make
640    # -------------------------------
641    def make(self):
642        # get what is done in the parent class...
643        plot.make(self)
644        if self.fmt is None: self.fmt = "%.1e"
645        # ... then add specific stuff
646        ############################################################################################
647        ### PRE-SETTINGS
648        ############################################################################################
649        # set dummy xy axis if not defined
650        if self.x is None: 
651            self.x = np.array(range(self.f.shape[0]))
652            self.mapmode = False
653            print "!! WARNING !! dummy coordinates on x axis"
654        if self.y is None: 
655            self.y = np.array(range(self.f.shape[1]))
656            self.mapmode = False
657            print "!! WARNING !! dummy coordinates on y axis"
658        # check sizes
659        if self.c is not None:
660            if self.c.ndim != 2:
661                print "!! WARNING !! Contour is not a 2D field. No contour.",self.c.ndim
662                self.c = None
663        if self.f.ndim != 2:
664            print "!! ERROR !! Field is not two-dimensional" ; exit()
665        # transposing if necessary
666        shape = self.f.shape
667        if shape[0] != shape[1]:
668         if len(self.x) == shape[0] and len(self.y) == shape[1]:
669            print "!! WARNING !! Transposing axes"
670            self.f = np.transpose(self.f)
671        # bound field
672        zevmin, zevmax = calculate_bounds(self.f,vmin=self.vmin,vmax=self.vmax)
673        what_I_plot = bounds(self.f,zevmin,zevmax)
674        # define contour field levels. define color palette
675        ticks = self.div + 1
676        zelevels = np.linspace(zevmin,zevmax,ticks)
677        palette = get_cmap(name=self.colorbar)
678        # do the same thing for possible contourline entries
679        if self.c is not None:
680            # if masked array, set masked values to filled values (e.g. np.nan) for plotting purposes
681            if type(self.c).__name__ in 'MaskedArray':
682               self.c[self.c.mask] = self.c.fill_value
683            zevminc, zevmaxc = calculate_bounds(self.c)
684            what_I_contour = bounds(self.c,zevminc,zevmaxc)
685            ticks = self.div + 1
686            zelevelsc = np.linspace(zevminc,zevmaxc,ticks)
687        ############################################################################################
688        ### MAIN PLOT
689        ### NB: contour lines are done before contour shades otherwise colorar error
690        ############################################################################################
691        if not self.mapmode:
692            ## A SIMPLE 2D PLOT
693            ###################
694            # swapping if requested
695            if self.swap:  x = self.y ; y = self.x
696            else:          x = self.x ; y = self.y
697            # coefficients on axis
698            x=x*self.xcoeff ; y=y*self.ycoeff
699            # make shaded and line contours
700            if self.c is not None: 
701                objC = mpl.contour(x, y, what_I_contour, \
702                            zelevelsc, colors = ccol, linewidths = cline)
703                #mpl.clabel(objC, inline=1, fontsize="small",\
704                #             inline_spacing=1,fmt="%.0f")
705            mpl.contourf(x, y, \
706                         self.f, \
707                         zelevels, cmap=palette)
708            #mpl.pcolor(x,y,\
709            #             self.f, \
710            #             cmap=palette)
711            # make log axes and/or invert ordinate
712            ax = mpl.gca()
713            if self.logx: mpl.semilogx()
714            if self.logy: mpl.semilogy()
715            if self.invert: ax.set_ylim(ax.get_ylim()[::-1])
716            if self.xmin is not None and self.xmax is not None:
717              if self.xmin > self.xmax: 
718                ax.set_xlim(ax.get_xlim()[::-1])
719                self.xmin,self.xmax = self.xmax,self.xmin
720            if self.xmin is not None: ax.set_xbound(lower=self.xmin)
721            if self.xmax is not None: ax.set_xbound(upper=self.xmax)
722            if self.ymin is not None: ax.set_ybound(lower=self.ymin)
723            if self.ymax is not None: ax.set_ybound(upper=self.ymax)
724            # set the number of ticks
725            if not self.logx:
726                ax.xaxis.set_major_locator(MaxNLocator(self.nxticks))
727            else:
728                print "!! WARNING. in logx mode, ticks are set automatically."
729            if not self.logy:
730                ax.yaxis.set_major_locator(MaxNLocator(self.nyticks))
731            else:
732                print "!! WARNING. in logy mode, ticks are set automatically."
733            ## specific modulo labels
734            if self.modx is not None:
735                ax = labelmodulo(ax,self.modx)
736        else:
737            ## A 2D MAP USING PROJECTIONS (basemap)
738            #######################################
739            mpl.xlabel("") ; mpl.ylabel("")
740            # additional security in case self.proj is None here
741            # ... we set cylindrical projection (the simplest one)
742            if self.proj is None: self.proj = "cyl"
743            # get lon and lat in 2D version.
744            # (but first ensure we do have 2D coordinates)
745            if self.x.ndim == 1:        [self.x,self.y] = np.meshgrid(self.x,self.y)
746            elif self.x.ndim > 2:    print "!! ERROR !! lon and lat arrays must be 1D or 2D"
747            # get lat lon intervals and associated settings
748            wlon = [np.min(self.x),np.max(self.x)]
749            wlat = [np.min(self.y),np.max(self.y)]
750            # -- area presets are in set_area.txt
751            if self.area is not None:
752             if self.area in area.keys():
753                wlon, wlat = area[self.area]
754            # -- settings for meridians and parallels
755            steplon = int(abs(wlon[1]-wlon[0])/6.)
756            steplat = int(abs(wlat[1]-wlat[0])/3.)
757            #mertab = np.r_[wlon[0]:wlon[1]:steplon] ; merlab = [0,0,0,1]
758            #partab = np.r_[wlat[0]:wlat[1]:steplat] ; parlab = [1,0,0,0]
759            if np.abs(wlon[0]) < 180.1 and np.abs(wlon[1]) < 180.1:
760                mertab = np.r_[-180.:180.:steplon]
761            else:
762                mertab = np.r_[0.:360.:steplon]
763            merlab = [0,0,0,1]
764            partab = np.r_[-90.:90.:steplat] ; parlab = [1,0,0,0]
765            format = '%.1f'
766            # -- center of domain and bounding lats
767            lon_0 = 0.5*(wlon[0]+wlon[1])
768            lat_0 = 0.5*(wlat[0]+wlat[1])
769            # some tests, bug fixes, and good-looking settings
770            # ... cyl is good for global and regional
771            if self.proj == "cyl":
772                format = '%.0f'
773            # ... global projections
774            elif self.proj in ["ortho","moll","robin"]:
775                wlat[0] = None ; wlat[1] = None ; wlon[0] = None ; wlon[1] = None
776                lon_0 = np.ceil(lon_0) # reverse map if lon_0 is slightly below 180 with [0,360]
777                steplon = 30. ; steplat = 30.
778                if self.proj in ["moll"]: steplon = 60.
779                if self.proj in ["robin"]: steplon = 90.
780                mertab = np.r_[-360.:360.:steplon]
781                partab = np.r_[-90.:90.:steplat]
782                if self.proj == "ortho": 
783                    merlab = [0,0,0,0] ; parlab = [0,0,0,0]
784                    # in ortho projection, blon and blat can be used to set map center
785                    if self.blon is not None: lon_0 = self.blon
786                    if self.blat is not None: lat_0 = self.blat
787                elif self.proj == "moll":
788                    merlab = [0,0,0,0]
789                format = '%.0f'
790            # ... regional projections
791            elif self.proj in ["lcc","laea","merc"]:
792                if self.proj in ["lcc","laea"] and wlat[0] == -wlat[1]: 
793                    print "!! ERROR !! with Lambert lat1 must be different than lat2" ; exit()
794                if wlat[0] < -80. and wlat[1] > 80.:
795                    print "!! ERROR !! set an area (not global)" ; exit()
796                format = '%.0f'
797            elif self.proj in ["npstere","spstere"]:
798                # in polar projections, blat gives the bounding lat
799                # if not set, set something reasonable
800                if self.blat is None:   self.blat = 60.
801                # help the user who forgets self.blat would better be negative in spstere
802                # (this actually serves for the default setting just above)
803                if self.proj == "spstere" and self.blat > 0: self.blat = -self.blat
804                # labels
805                mertab = np.r_[-360.:360.:15.]
806                partab = np.r_[-90.:90.:5.]
807            # ... unsupported projections
808            else:
809                print "!! ERROR !! unsupported projection. supported: "+\
810                      "cyl, npstere, spstere, ortho, moll, robin, lcc, laea, merc"
811            # finally define projection
812            m = Basemap(projection=self.proj,\
813                        lat_0=lat_0,lon_0=lon_0,\
814                        boundinglat=self.blat,\
815                        llcrnrlat=wlat[0],urcrnrlat=wlat[1],\
816                        llcrnrlon=wlon[0],urcrnrlon=wlon[1])
817            # in some case need to translated to the left for colorbar + labels
818            # TBD: break stuff. a better solution should be found.
819            if self.leftcorrect:
820                ax = mpl.gca()
821                pos = ax.get_position().bounds
822                newpos = [0.,pos[1],pos[2],pos[3]]
823                ax.set_position(newpos)
824            # draw meridians and parallels
825            ft = int(mpl.rcParams['font.size']*3./4.)
826            zelatmax = 85.
827            m.drawmeridians(mertab,labels=merlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format,latmax=zelatmax)
828            m.drawparallels(partab,labels=parlab,color='grey',linewidth=0.75,fontsize=ft,fmt=format,latmax=zelatmax)
829            # define background (see set_back.txt)
830            if self.back is not None:
831              if self.back in back.keys():
832                 print "**** info: loading a background, please wait.",self.back
833                 if self.back not in ["coast","sea"]:
834                    try: m.warpimage(back[self.back],scale=0.75)
835                    except: print "!! ERROR !! no background image could be loaded. probably not connected to the internet?"
836                 elif self.back == "coast":
837                    m.drawcoastlines()
838                 elif self.back == "sea":
839                    m.drawlsmask(land_color='white',ocean_color='aqua')
840              else:
841                 print "!! ERROR !! requested background not defined. change name or fill in set_back.txt" ; exit()
842            # define x and y given the projection
843            x, y = m(self.x, self.y)
844            # contour field. first line contour then shaded contour.
845            if self.c is not None: 
846                objC2 = m.contour(x, y, what_I_contour, \
847                            zelevelsc, colors = ccol, linewidths = cline)
848                #mpl.clabel(objC2, inline=1, fontsize=10,manual=True,fmt='-%2.0f$^{\circ}$C',colors='r')
849            m.contourf(x, y, what_I_plot, zelevels, cmap = palette, alpha = self.trans, antialiased=True)
850        ############################################################################################
851        ### COLORBAR
852        ############################################################################################
853        if self.trans > 0. and self.showcb:
854            ## draw colorbar. settings are different with projections. or if not mapmode.
855            if not self.mapmode: orientation=zeorientation ; frac = 0.075 ; pad = 0.03 ; lu = 0.5
856            elif self.proj in ['moll']: orientation="horizontal" ; frac = 0.08 ; pad = 0.03 ; lu = 1.0
857            elif self.proj in ['cyl']: orientation="vertical" ; frac = 0.023 ; pad = 0.03 ; lu = 0.5
858            else: orientation = zeorientation ; frac = zefrac ; pad = 0.03 ; lu = 0.5
859            if self.cbticks is None:
860                self.cbticks = min([ticks/2+1,21])
861            zelevpal = np.linspace(zevmin,zevmax,num=self.cbticks)
862            zecb = mpl.colorbar(fraction=frac,pad=pad,\
863                                format=self.fmt,orientation=orientation,\
864                                ticks=zelevpal,\
865                                extend='neither',spacing='proportional')
866            if zeorientation == "horizontal": zecb.ax.set_xlabel(self.title) ; self.title = ""
867            # colorbar title --> units
868            if self.units not in ["dimless",""]:
869                zecb.ax.set_title("["+self.units+"]",fontsize=3.*mpl.rcParams['font.size']/4.,x=lu,y=1.025)
870
871        ############################################################################################
872        ### VECTORS. must be after the colorbar. we could also leave possibility for streamlines.
873        ############################################################################################
874        ### not expecting NaN in self.vx and self.vy. masked arrays is just enough.
875        if self.vx is not None and self.vy is not None: 
876                # vectors on map projection or simple 2D mapping
877                if self.mapmode: 
878                   [vecx,vecy] = m.rotate_vector(self.vx,self.vy,self.x,self.y) # for metwinds only ?
879                else:
880                   vecx,vecy = self.vx,self.vy
881                   if x.ndim < 2 and y.ndim < 2: x,y = np.meshgrid(x,y)
882                # reference vector is scaled
883                if self.wscale is None: 
884                    self.wscale = ppcompute.mean(np.sqrt(self.vx*self.vx+self.vy*self.vy))
885                # make vector field
886                if self.mapmode: 
887                    q = m.quiver( x[::self.svy,::self.svx],y[::self.svy,::self.svx],\
888                                  vecx[::self.svy,::self.svx],vecy[::self.svy,::self.svx],\
889                                  angles='xy',color=self.colorvec,pivot='tail',\
890                                  scale=self.wscale*reducevec,width=widthvec )
891                else:
892                    q = mpl.quiver( x[::self.svy,::self.svx],y[::self.svy,::self.svx],\
893                                    vecx[::self.svy,::self.svx],vecy[::self.svy,::self.svx],\
894                                    angles='xy',color=self.colorvec,pivot='tail',\
895                                    scale=self.wscale*reducevec,width=widthvec )
896                # make vector key.
897                #keyh = 1.025 ; keyv = 1.05 # upper right corner over colorbar
898                keyh = 0.97 ; keyv = 1.05
899                #keyh = -0.03 ; keyv = 1.08 # upper left corner
900                p = mpl.quiverkey(q,keyh,keyv,\
901                                  self.wscale,str(int(self.wscale)),\
902                                  fontproperties={'size': 'small'},\
903                                  color='black',labelpos='S',labelsep = 0.07)
904        ############################################################################################
905        ### TEXT. ANYWHERE. add_text.txt should be present with lines x ; y ; text ; color
906        ############################################################################################
907        try:
908            f = open("add_text.txt", 'r')
909            for line in f:
910              if "#" in line: pass
911              else:
912                  userx, usery, usert, userc = line.strip().split(';')
913                  userc = userc.strip()
914                  usert = usert.strip()
915                  userx = float(userx.strip())
916                  usery = float(usery.strip())
917                  if self.mapmode: userx,usery = m(userx,usery)
918                  mpl.text(userx,usery,usert,\
919                           color = userc,\
920                           horizontalalignment='center',\
921                           verticalalignment='center')
922            f.close()
923        except IOError:
924            pass
925
926    # makeshow = make + show
927    # -------------------------------
928    def makeshow(self):
929        self.make()
930        mpl.show()
Note: See TracBrowser for help on using the repository browser.