source: trunk/UTIL/PYTHON/planetoplot_v2/ppclass.py @ 910

Last change on this file since 910 was 910, checked in by aslmd, 12 years ago

PLANETOPLOT v2


  1. Spiga LMD/UPMC 24/03/2013

Contents


core


  • ppclass.py --> main class with pp() objects (need ppplot and ppcompute)
  • ppplot.py --> plot class (can be used independently, need ppcompute)
  • ppcompute.py --> computation class (can be used independently)

scripts


  • pp.py --> command line utility to use ppclass
  • pp_reload.py --> command line utility to load saved plot objects *.ppobj
  • example/* --> example scripts using ppclass

settings files


  • set_area.txt --> setting file: predefined areas for plotting (can be omitted)
  • set_back.txt --> setting file: predefined backgrounds for plotting (can be omitted)
  • set_multiplot.txt --> setting file: predefined coefficients for multiplots (can be omitted)
  • set_ppclass.txt --> setting file: predefined variables for x,y,z,t (can be omitted)
  • set_var.txt --> setting file: predefined colorbars, format, labels, etc... for variables (can be omitted)

documentation


  • README.TXT --> this README file

data


  • demo_data/* --> plot objects for a demonstration tour and customizing tests

Requirements


python + numpy + matplotlib + netCDF4

  • for mapping --> Basemap
  • for scientific computations --> scipy

[recommended: Enthought Python Distribution (free for academics)]

Installation


  • install required softwares and librairies in requirements
  • add planetoplot_v2 in your PYTHONPATH environment variable (and in your PATH to use pp.py)

Take a demo tour


pp_reload.py demo_data/*

Improvements compared to v1


  • code readability and structuration for future improvements
  • modularity (class formulation) + easy definition/addition of attributes
  • separation between data retrieval and plotting
  • versatility + command line (pp.py)

--> for quick inspection

+ interactive session (ipython)

--> for testing and exploring

+ scripts

--> for powerful and fully customized results

  • performance (overall and on large files) + memory consumption (only retrieve what is needed)
  • saving/loading plot objects in/from *.ppobj
  • plot aesthetics and customizing (see header in ppplot)
  • multiplot flexibility with .plotin attribute
  • easy definition of presets with set_*.txt files
  • function: one field vs. another one
  • emulation of + - / * operations on fields (between two fields or between a field and a int/float)
  • computations of min/max in addition to mean
  • simple inspection of dimensions+variables in a file (pp.py -f file)
  • verbose / non verbose mode

Acknowledgements


Thanks to A. Colaitis, T. Navarro, J. Leconte
for feedbacks and contributions on version 1

File size: 68.6 KB
Line 
1###############################################
2## PLANETOPLOT                               ##
3## --> PPCLASS                               ##
4## A generic and versatile Python module     ##
5## ... to read netCDF files and plot         ##
6###############################################
7## Author: Aymeric Spiga. 02-03/2013         ##
8###############################################
9# python built-in librairies
10import os
11import time as timelib
12import pickle
13# added librairies
14import numpy as np
15import netCDF4
16import matplotlib.pyplot as mpl
17# personal librairies
18import ppplot
19import ppcompute
20###############################################
21
22###################################
23#### HEADER                      ##
24#### ... executed when imported  ##
25###################################
26print "************************************"
27print "**** WELCOME TO PLANETOPLOT 2.0 ****"
28print "************************************"
29# where settings files are located...
30whereset = None
31whereset = ppcompute.findset(whereset)
32# ... we load user-defined automatic settings from set_ppclass.txt
33zefile = "set_ppclass.txt"
34glob_listx = [] ; glob_listy = [] ; glob_listz = [] ; glob_listt = []
35try: 
36    f = open(whereset+zefile, 'r') ; lines = f.readlines()
37    for stuff in lines[5].strip().split(';'): glob_listx.append(stuff)
38    for stuff in lines[8].strip().split(';'): glob_listy.append(stuff)
39    for stuff in lines[11].strip().split(';'): glob_listz.append(stuff)
40    for stuff in lines[14].strip().split(';'): glob_listt.append(stuff)
41except IOError: 
42    print "warning: "+zefile+" not in "+whereset+" ; no presets."
43
44##################################
45#### USEFUL GENERIC FUNCTIONS ####
46##################################
47
48# inspect variables and dimensions in a netCDF file
49def inspect(filename):
50    print "**** INSPECT FILE",filename
51    test = netCDF4.Dataset(filename)
52    print "**** VARIABLES: ",test.variables.keys()
53    for dim in test.dimensions.keys():
54        output = "**** DIMENSION: "+str(dim)+" "+str(len(test.dimensions[dim]))
55        try: output = output + " ----> "+str(test.variables[dim][0])+"  "+str(test.variables[dim][-1])
56        except: pass
57        print output ; output = ""
58
59# check a tab and exit if wrong. if just one string make it a list.
60# (if allownumber, convert this into a string).
61def checktab(tab,mess="",allownone=False,allownumber=False):
62    if tab is None: 
63      if not allownone:  print "pp.define: no "+mess ; exit()
64      else: pass
65    else:
66      if not isinstance(tab, list):
67        if isinstance(tab, str): 
68            tab = [tab]
69        elif (isinstance(tab, int) or isinstance(tab, float)) and allownumber: 
70            tab = [str(tab)] 
71        else: 
72            print "pp.define: "+mess+" should be either a string or a list of strings!" ; exit()
73      elif isinstance(tab, list):
74        if isinstance(tab[0],str): 
75            pass
76        elif (isinstance(tab[0], int) or isinstance(tab[0], float)) and allownumber:
77            for iii in range(len(tab)): tab[iii] = str(tab[iii])
78        else: 
79            print "pp.define: "+mess+" should be either a string or a list of strings!" ; exit()
80    return tab
81
82# determine which method is to be applied to a given dimension
83def findmethod(tab):
84    if tab is None:              output = "free"
85    elif tab[0,0] != tab[0,1]:   output = "comp"
86    else:                        output = "fixed"
87    return output
88
89# read what is given by the user (version of T. Navarro)
90def readslices(saxis):
91    if saxis == None:
92        zesaxis = None
93    else:
94        zesaxis = np.empty((len(saxis),2))
95        for i in range(len(saxis)):
96            a = separatenames(saxis[i])
97            if len(a) == 1:
98                zesaxis[i,:] = float(a[0])
99            else:
100                zesaxis[i,0] = float(a[0])
101                zesaxis[i,1] = float(a[1])         
102    return zesaxis
103# (needed by readslices)
104# look for comas in the input name to separate different names (files, variables,etc ..)
105def separatenames (name):
106    if name is None: names = None
107    else:
108      names = [] ; stop = 0 ; currentname = name
109      while stop == 0:
110        indexvir = currentname.find(',')
111        if indexvir == -1: stop = 1 ; name1 = currentname
112        else: name1 = currentname[0:indexvir]
113        names = np.concatenate((names,[name1]))
114        currentname = currentname[indexvir+1:len(currentname)]
115    return names
116
117#######################
118### THE MAIN OBJECT ###
119#######################
120class pp():
121
122    # print out a help string when help is invoked on the object
123    def __repr__(self):
124        whatprint = 'pp object. \"help(pp)\" for more information\n'
125        return whatprint
126
127    # default settings
128    # -- user can define settings by two methods.
129    # -- 1. yeah = pp(file="file.nc")
130    # -- 2. yeah = pp() ; yeah.file = "file.nc"
131    def __init__(self,file=None,var="notset",\
132                      filegoal=None,vargoal=None,\
133                      x=None,y=None,z=None,t=None,\
134                      stridex=1,stridey=1,\
135                      stridez=1,stridet=1,\
136                      compute="mean",\
137                      verbose=False,noproj=False,\
138                      superpose=False,\
139                      plotin=None,\
140                      forcedimplot=-1,\
141                      out="gui",\
142                      filename="myplot",\
143                      folder="./"):
144        self.request = None
145        self.nfin = 0 ; self.nvin = 0
146        self.nplotx = None ; self.nploty = None
147        self.nplotz = None ; self.nplott = None
148        self.status = "init"
149        self.fig = None ; self.subv = None ; self.subh = None 
150        self.n = 0 ; self.howmanyplots = 0
151        self.nplot = 0
152        self.p = None
153        ## what could be defined by the user
154        self.file = file
155        self.var = var
156        self.filegoal = filegoal
157        self.vargoal = vargoal
158        self.x = x ; self.y = y   ## if None, free dimension
159        self.z = z ; self.t = t   ## if None, free dimension
160        self.stridex = stridex ; self.stridey = stridey
161        self.stridez = stridez ; self.stridet = stridet
162        self.compute = compute
163        self.verbose = verbose
164        self.noproj = noproj
165        self.plotin = plotin
166        self.superpose = superpose
167        self.forcedimplot = forcedimplot
168        self.out = out
169        self.filename = filename
170        self.folder = folder
171
172    # print status
173    def printstatus(self):
174        print "**** Done step: " + self.status
175
176    #####################################################
177    # EMULATE OPERATORS + - * / ** << FOR PP() OBJECTS  #
178    #####################################################
179
180    # check the compatibility of two objects for operations
181    # --> if other is a pp class, test sizes and return isnum = False
182    # --> if other is an int or a float, return isnum = True
183    # --> otherwise, just print an error and exit
184    def checktwo(self,other):
185        if other.__class__.__name__ == "pp":
186          isnum = False
187          if self.status in ["init","defined"] or other.status in ["init","define"]: 
188             print "!! ERROR !! Please use .retrieve to get fields for plots with one of your pp operands." ; exit()
189          if self.nfin   != other.nfin   or \
190             self.nvin   != other.nvin   or \
191             self.nplott != other.nplott or \
192             self.nplotz != other.nploty or \
193             self.nploty != other.nploty or \
194             self.nplotx != other.nplotx :
195               print "!! ERROR !! The two operands do not have the same number of files, variables, t z y x requests."
196               exit()
197        elif isinstance(other,int) or isinstance(other,float):
198          isnum = True
199        else:
200          print "!! ERROR !! The operand is neither a pp class nor an integer or a float." ; exit()
201        return isnum
202
203    # define the operation + on two objects. or with an int/float.
204    def __add__(self,other):
205        if self.status == "definedplot":
206            self.status = "retrieved" # to release a warning so that .defineplot() is done before .makeplot()
207                                      # ... otherwise operated fields will not be displayed
208        isnum = self.checktwo(other)
209        for i in range(self.nfin):
210         for j in range(self.nvin):
211          for t in range(self.nplott):
212           for z in range(self.nplotz):
213            for y in range(self.nploty):
214             for x in range(self.nplotx):
215              obj = self.request[i][j][t][z][y][x]
216              if not isnum:   obj2 = other.request[i][j][t][z][y][x]
217              else:           obj2 = other
218              if "vector" in self.vargoal[j] + self.filegoal[i]:
219                  exit() # we do not operate on vectors yet.
220              else:
221                  obj = obj + obj2
222        return self
223
224    # define the operation - on two objects. or with an int/float.
225    def __sub__(self,other):
226        if self.status == "definedplot":
227            self.status = "retrieved" # to release a warning so that .defineplot() is done before .makeplot()
228                                      # ... otherwise operated fields will not be displayed       
229        isnum = self.checktwo(other)
230        for i in range(self.nfin):
231         for j in range(self.nvin):
232          for t in range(self.nplott):
233           for z in range(self.nplotz):
234            for y in range(self.nploty):
235             for x in range(self.nplotx):
236              obj = self.request[i][j][t][z][y][x]
237              if not isnum:   obj2 = other.request[i][j][t][z][y][x]
238              else:           obj2 = other
239              if "vector" in self.vargoal[j] + self.filegoal[i]:
240                  exit() # we do not operate on vectors yet.
241              else:
242                  obj = obj - obj2
243        return self
244
245    # define the operation * on two objects. or with an int/float.
246    def __mul__(self,other):
247        if self.status == "definedplot":
248            self.status = "retrieved" # to release a warning so that .defineplot() is done before .makeplot()
249                                      # ... otherwise operated fields will not be displayed
250        isnum = self.checktwo(other)
251        for i in range(self.nfin):
252         for j in range(self.nvin):
253          for t in range(self.nplott):
254           for z in range(self.nplotz):
255            for y in range(self.nploty):
256             for x in range(self.nplotx):
257              obj = self.request[i][j][t][z][y][x]
258              if not isnum:   obj2 = other.request[i][j][t][z][y][x]
259              else:           obj2 = other
260              if "vector" in self.vargoal[j] + self.filegoal[i]:
261                  exit() # we do not operate on vectors yet.
262              else:
263                  obj = obj * obj2
264        return self
265
266    # define the operation / on two objects. or with an int/float.
267    def __div__(self,other):
268        if self.status == "definedplot":
269            self.status = "retrieved" # to release a warning so that .defineplot() is done before .makeplot()
270                                      # ... otherwise operated fields will not be displayed
271        isnum = self.checktwo(other)
272        for i in range(self.nfin):
273         for j in range(self.nvin):
274          for t in range(self.nplott):
275           for z in range(self.nplotz):
276            for y in range(self.nploty):
277             for x in range(self.nplotx):
278              obj = self.request[i][j][t][z][y][x]
279              if not isnum:   obj2 = other.request[i][j][t][z][y][x]
280              else:           obj2 = other
281              if "vector" in self.vargoal[j] + self.filegoal[i]:
282                  exit() # we do not operate on vectors yet.
283              else:
284                  obj = obj / obj2
285        return self
286
287    # define the operation ** on one object.
288    def __pow__(self,num):
289        if self.status == "definedplot":
290            self.status = "retrieved" # to release a warning so that .defineplot() is done before .makeplot()
291                                      # ... otherwise operated fields will not be displayed
292        if isinstance(num,int) or isinstance(num,float):
293            for i in range(self.nfin):
294             for j in range(self.nvin):
295              for t in range(self.nplott):
296               for z in range(self.nplotz):
297                for y in range(self.nploty):
298                 for x in range(self.nplotx):
299                  obj  = self.request[i][j][t][z][y][x]
300                  if "vector" in self.vargoal[j] + self.filegoal[i]:
301                      exit() # we do not operate on vectors yet.
302                  else:
303                      obj = obj ** num
304        else:
305            print "!! ERROR !! To define a power, either an int or a float is needed." ; exit()
306        return self
307
308    # define the operation <<
309    # ... e.g. obj2 << obj1
310    # ... means: get init for pp object obj2 from another pp object obj1
311    # ... (this should solve the affectation trap)
312    def __lshift__(self,other):
313        if other.__class__.__name__ == "pp":
314            self.file = other.file
315            self.var = other.var
316            self.filegoal = other.filegoal
317            self.vargoal = other.vargoal
318            self.x = other.x ; self.y = other.y   ## if None, free dimension
319            self.z = other.z ; self.t = other.t   ## if None, free dimension
320            self.stridex = other.stridex ; self.stridey = other.stridey
321            self.stridez = other.stridez ; self.stridet = other.stridet
322            self.verbose = other.verbose
323            self.noproj = other.noproj
324            self.plotin = other.plotin
325            self.superpose = other.superpose
326            self.forcedimplot = other.forcedimplot
327            self.out = other.out
328            self.filename = other.filename
329            self.folder = other.folder
330        else:
331            print "!! ERROR !! argument must be a pp object." ; exit()
332
333    ##############################################################################################
334    # define method
335    # ---------
336    # ... (file and var are either one string or a vector of strings)
337    # ... the goal of define is to define a 2D array of onerequest() objects (see class below)
338    #     given the number of file, var, x, y, z, t asked by the user
339    # ... objectives for file or var are given through filegoal and vargoal
340    #     --> possible values: main contour vector
341    # ---------
342    # ... then onerequest() objects are being defined more precisely
343    #     by getting index_x index_y index_z index_t
344    #     and setting method_x method_y method_z method_t to either
345    #      - "free" for free dimensions (plot dimensions)
346    #      - "comp" for averages, max, min
347    #      - "fixed" for fixed dimensions (possibly several i.e. multislice)
348    ##############################################################################################
349    def define(self):
350        self.printstatus()
351        # initial check and get dimensions
352        self.file = checktab(self.file,mess="file")
353        self.nfin = len(self.file)
354        if self.verbose:
355            for i in range(self.nfin): inspect(self.file[i])
356        self.var = checktab(self.var,mess="var")
357        self.nvin = len(self.var)
358        # check goal tabs for files and variables
359        # ... default is to plot everything
360        if self.filegoal is None: self.filegoal = ["main"]*self.nfin
361        if self.vargoal is None:  self.vargoal  = ["main"]*self.nvin
362        self.filegoal = checktab(self.filegoal, mess="filegoal")
363        self.vargoal  = checktab(self.vargoal,  mess="vargoal")
364        if len(self.filegoal) != self.nfin:  print "!! ERROR !! filegoal must be the same size as file." ; exit()
365        if len(self.vargoal)  != self.nvin:  print "!! ERROR !! vargoal must be the same size as var." ; exit()
366        # variables: initial check
367        self.x = checktab(self.x,mess="x",allownone=True,allownumber=True)
368        self.y = checktab(self.y,mess="y",allownone=True,allownumber=True)
369        self.z = checktab(self.z,mess="z",allownone=True,allownumber=True)
370        self.t = checktab(self.t,mess="t",allownone=True,allownumber=True)
371        # for the moment not var- nor file- dependent.
372        # but this could be the case.
373        sx = readslices(self.x) ; sy = readslices(self.y)
374        sz = readslices(self.z) ; st = readslices(self.t)
375        # get methods
376        mx = findmethod(sx) ; my = findmethod(sy)
377        mz = findmethod(sz) ; mt = findmethod(st)
378        # get number of plots to be done
379        if mx == "fixed": self.nplotx = sx.size/2
380        else:             self.nplotx = 1
381        if my == "fixed": self.nploty = sy.size/2
382        else:             self.nploty = 1
383        if mz == "fixed": self.nplotz = sz.size/2
384        else:             self.nplotz = 1
385        if mt == "fixed": self.nplott = st.size/2
386        else:             self.nplott = 1
387        if self.verbose:  print "**** OK. Plots over x,y,z,t -->",self.nplotx,self.nploty,self.nplotz,self.nplott
388        # create the list of onerequest() objects
389        self.request = [[[[[[ \
390                       onerequest() \
391                       for x in range(self.nplotx)] for y in range(self.nploty)] \
392                       for z in range(self.nplotz)] for t in range(self.nplott)] \
393                       for j in range(self.nvin)]   for i in range(self.nfin)] 
394        # loop on onerequest() objects
395        for i in range(self.nfin):
396         for j in range(self.nvin):
397          for t in range(self.nplott):
398           for z in range(self.nplotz):
399            for y in range(self.nploty):
400             for x in range(self.nplotx):
401              obj = self.request[i][j][t][z][y][x]
402              # fill in names for files and variables
403              obj.verbose = self.verbose
404              obj.file = self.file[i]
405              obj.var = self.var[j]
406              # indicate the computation method
407              obj.compute = self.compute
408              # open the files (the same file might be opened several times but this is cheap)
409              obj.openfile()
410              ### get x,y,z,t dimensions from file
411              obj.getdim()
412              ### get methods
413              obj.method_x = mx ; obj.method_y = my
414              obj.method_z = mz ; obj.method_t = mt           
415              ### get index
416              obj.getindextime(st,t,self.stridet)
417              obj.getindexvert(sz,z,self.stridez)
418              obj.getindexhori(sx,sy,x,y,self.stridex,self.stridey)
419        # change status
420        self.status = "defined"
421
422    ##############################################################################################
423    # retrieve method
424    # --> for each element onerequest() in the array, get field .var from .f file
425    # --> see below the onerequest() class:
426    #        - only get what is needed for computing and plotting
427    #        - averages etc... are computed here
428    # --> RESULT: each onerequest() object has now its attribute .field filled
429    # --> if one wants to perform operations on fields, this should be done after retrieve()
430    ##############################################################################################
431    def retrieve(self):
432        self.printstatus()
433        # check if things were done OK before
434        if self.status != "defined": print "!! ERROR !! Please use .define() to define your pp object."
435        ## first get fields
436        ## ... only what is needed is extracted from the files
437        ## ... and averages are computed
438        for i in range(self.nfin):
439         for j in range(self.nvin):
440          for t in range(self.nplott):
441           for z in range(self.nplotz):
442            for y in range(self.nploty):
443             for x in range(self.nplotx):
444              obj = self.request[i][j][t][z][y][x]
445              obj.getfield()
446        # change status
447        self.status = "retrieved"
448
449    ##########################################################
450    # get: a shortcut method for the define + retrieve chain #
451    ##########################################################
452    def get(self):
453        self.define()
454        self.retrieve()
455
456    ########################################
457    # smooth: smooth the field in 1D or 2D #
458    ########################################
459    ## TBD: smooth not OK with masked array in the end of retrieve()
460    def smooth(self,window):
461        if self.verbose: 
462            print "!! WARNING !! Performing a smoothing with a window size",window
463            print "!! WARNING !! To come back to unsmoothed file, use .get() again"
464        for i in range(self.nfin):
465         for j in range(self.nvin):
466          for t in range(self.nplott):
467           for z in range(self.nplotz):
468            for y in range(self.nploty):
469             for x in range(self.nplotx):
470              obj = self.request[i][j][t][z][y][x]
471              if obj.field.ndim == 1:
472                  print "!! ERROR !! 1D smoothing not supported yet because reduces array sizes."
473                  exit()
474                  # TBD TBD TBD
475                  #obj.field = ppcompute.smooth1d(obj.field,window=window)
476              elif obj.field.ndim == 2:
477                  obj.field = ppcompute.smooth2d(obj.field,window=window)
478
479    ############################################################################################## 
480    # defineplot method
481    # --> defineplot first defines arrays of plot objects and set each of them
482    #     ... simple looping except cases where goal is not main (e.g. contour or vector)
483    # --> principle: each onerequest() object gives birth to a subplot
484    # --> defineplot vs. makeplot: defining plot and actually plotting it are clearly separated
485    # --> THE KEY OUPUT OF defineplot IS AN ARRAY self.p OF PLOT OBJECTS
486    # optional arguments
487    # --> extraplot: to indicate a number of plots to be added afterwards (use self.plotin)
488    # --> loadfile: to use self.p from a previously saved file
489    ##############################################################################################
490    def defineplot(self,extraplot=0,loadfile=None):
491        # -----------------------------------------------------
492        # LOAD MODE: load a self.p object. count plots from it.
493        # -----------------------------------------------------
494        if loadfile is not None:
495            try: filehandler = open(loadfile, 'r') ; self.p = pickle.load(filehandler) 
496            except IOError: print "!! ERROR !! Cannot find object file to load." ; exit()
497            self.status = "definedplot" ; self.plotin = None
498            self.nplot = len(self.p) ; self.howmanyplots = self.nplot
499            return
500        # -----------------------------------------------------
501        # REGULAR MODE
502        # -----------------------------------------------------
503        self.printstatus()
504        # check if things were done OK before
505        if self.status in ["init","defined"]: 
506            print "!! ERROR !! Please use .retrieve() to get fields for plots with your pp object." ; exit()
507        # check self.plotin (an existing fig on which to add plots afterwards)
508        if self.plotin.__class__.__name__ == "pp":
509            if self.plotin.fig is None:
510                self.plotin = None # this is an additional security in case
511                                   #   a pp object is given without figure opened yet.
512        elif self.plotin is not None:
513            print "!! ERROR !! plotin argument must be a pp object." ; exit()
514        # initialize the array of subplot objects
515        # either something new or attributes coming from plotin object
516        if self.plotin is None:  self.p = [ ]
517        else:                    self.p = self.plotin.p
518        # create an array of subplot objects
519        # ... in theory the order of looping can be changed without any harm
520        # ... the only important thing is to keep i,j,t,z,y,x resp. for file,var,t,z,y,x
521        count = 0
522        for i in range(self.nfin):
523         if self.filegoal[i] == "main": 
524          for j in range(self.nvin):
525           if self.vargoal[j] == "main":
526            for t in range(self.nplott):
527             for z in range(self.nplotz):
528              for y in range(self.nploty):
529               for x in range(self.nplotx):
530                # look at dimension and append the right plot object
531                obj = self.request[i][j][t][z][y][x]
532                dp = obj.dimplot
533                if dp == 1 or self.forcedimplot == 1:    plobj = ppplot.plot1d()
534                elif dp == 2 or self.forcedimplot == 2:  plobj = ppplot.plot2d()
535                elif dp == 0: print "**** OK. VALUES VALUES VALUES",obj.field
536                else:         print "!! ERROR !! 3D or 4D plots not supported" ; exit()
537                # load abscissa and ordinate in obj
538                obj.definecoord()
539                # we start to define things here before appending
540                # (convenient: could be overridden by the user before makeplot)
541                # ... the if loop is necessary so that we can loop above on the dp=0 case
542                if dp in [1,2]:
543                    # and define what to do in plobj
544                    plobj.invert = obj.invert_axes
545                    plobj.swap = obj.swap_axes
546                    # axis labels
547                    plobj.xlabel = obj.absclab ; plobj.ylabel = obj.ordilab
548                    # superpose or not (this is mostly for saving purpose)
549                    plobj.superpose = self.superpose
550                    # get title, colormaps, labels, etc.. from var
551                    plobj.var = obj.var
552                    plobj.define_from_var()
553                    # generic 1D/2D: load field and coord in plot object
554                    plobj.field = obj.field    # field to be plotted
555                    plobj.absc = obj.absc      # abscissa (or longitude)
556                    plobj.ordi = obj.ordi      # ordinate (or latitude)
557                                               # -- useless in 1D but not used anyway
558                    if dp == 2:
559                    # specific 2d plot stuff
560                        # -- light grey background for missing values
561                        if type(plobj.field).__name__ in 'MaskedArray': plobj.axisbg = '0.75'
562                        # -- define if it is a map or a plot
563                        plobj.mapmode = ( obj.method_x+obj.method_y == "freefree" \
564                                       and "grid points" not in obj.name_x \
565                                       and not self.noproj )
566                    # finally append
567                    self.p.append(plobj)
568                    count = count + 1
569        # self.nplot is number of plot to be defined in this call to defineplot()
570        # (because of self.plotin this might less than length of self.p)
571        self.nplot = count
572        # --- superimposed contours and vectors ---
573        # we have to start another loop because we need forward information
574        # TBD: there is probably a more flexible way to do that
575        count = 0
576        for i in range(self.nfin):
577         for j in range(self.nvin):
578          for t in range(self.nplott):
579           for z in range(self.nplotz):
580            for y in range(self.nploty):
581             for x in range(self.nplotx):
582              goal = self.vargoal[j] + self.filegoal[i]
583              obj = self.request[i][j][t][z][y][x]
584              if "mainmain" in goal and obj.dimplot == 2:
585                  # the plot object we consider in the loop
586                  pl = self.p[count]
587                  # -- see if there is a contour requested...
588                  # (we use try because we might be at the end of the list)
589                  found = 0
590                  try:    condvar = self.vargoal[j+1]
591                  except: condvar = "itisok"
592                  try:    condfile = self.filegoal[i+1]
593                  except: condfile = "itisok"
594                  # ... get contour
595                  ##########################################
596                  # NB: contour is expected to be right after main otherwise it is not displayed
597                  ##########################################
598                  if condvar == "contour":
599                      plobj.addcontour = self.request[i][j+1][t][z][y][x].field ; found += 1
600                  if condfile == "contour":
601                      plobj.addcontour = self.request[i+1][j][t][z][y][x].field ; found += 1
602                  # see if there is a vector requested...
603                  # (we use try because we might be at the end of the list)
604                  try:    condvar = self.vargoal[j+found+1]+self.vargoal[j+found+2]
605                  except: condvar = "itisok"
606                  try:    condfile = self.filegoal[i+found+1]+self.filegoal[i+found+2]
607                  except: condfile = "itisok"
608                  # ... get vector and go directly to the next iteration
609                  # (in some cases we would do this twice but this is cheap)
610                  if "vector" in condvar:
611                      plobj.addvecx = self.request[i][j+found+1][t][z][y][x].field
612                      plobj.addvecy = self.request[i][j+found+2][t][z][y][x].field
613                  if "vector" in condfile:
614                      plobj.addvecx = self.request[i+found+1][j][t][z][y][x].field
615                      plobj.addvecy = self.request[i+found+2][j][t][z][y][x].field
616                  count = count + 1
617        # COUNT PLOTS. if 0 just exit.
618        # self.howmanyplots is self.nplot + possible extraplots
619        self.howmanyplots = self.nplot + extraplot
620        if self.howmanyplots > 0:
621            if self.verbose: print "**** OK. expect %i plots" % (self.howmanyplots)
622        else:
623            exit() # because this means that we only had 0D values !
624        # final status
625        self.status = "definedplot"
626
627    ##############################################################################################
628    # makeplot method
629    # --> after defineplot and before makeplot, user-defined plot settings can be easily given
630    #     simply by modifying the attributes of each elements of self.p
631    # --> to change only one plot setting, no need to call defineplot again before makeplot
632    # --> in the end, the array self.p of plot objects is saved for easy and convenient replotting
633    # --> see practical examples in the folder 'examples'
634    ##############################################################################################
635    def makeplot(self):
636        self.printstatus()
637        # a few initial operations
638        # ------------------------
639        if "definedplot" not in self.status: 
640            print "!! ERROR !! Please use .defineplot() before .makeplot() can be used with your pp object." ; exit()
641        # open a figure and define subplots         
642        # ---------------------------------
643        if self.plotin is None: 
644            # start from scratch
645            self.fig = mpl.figure(figsize=(16,8))
646            self.subv,self.subh = ppplot.definesubplot(self.howmanyplots,self.fig) 
647            self.n = 0
648        else:
649            # start from an existing figure.
650            # extraplot must have been set in the call to the previous figure.
651            self.fig = self.plotin.fig
652            self.subv,self.subh = self.plotin.subv,self.plotin.subh
653            self.n = self.plotin.n
654            self.howmanyplots = self.plotin.howmanyplots
655        # LOOP on all subplots
656        # NB: cannot use 'for pl in self.p' if self.plotin not None
657        # --------------------
658        for count in range(self.nplot):
659            # the plot object we consider in the loop
660            pl = self.p[self.n]
661            # before making the plot, create a subplot. the first one is numbered 1 not 0.
662            # ... if pl.superpose, we use only one and only figure
663            # ... (and we have to be careful with not doing things several times)
664            if pl.superpose:
665                if self.n == 0: 
666                    self.fig.add_subplot(1,1,1,axisbg=pl.axisbg) # define one subplot (still needed for user-defined font sizes)
667                    sav = pl.xlabel,pl.ylabel,pl.title,pl.swaplab # save titles and labels
668                else: 
669                    pl.invert = False ; pl.lstyle = None # don't invert again axis
670                    pl.xlabel,pl.ylabel,pl.title,pl.swaplab = sav # set saved titles and labels
671            else:
672                self.fig.add_subplot(self.subv,self.subh,self.n+1,axisbg=pl.axisbg)
673            if self.verbose: print "**** Done subplot %i / %i " %( self.n+1,self.howmanyplots ) 
674            # finally make the plot
675            pl.make()
676            self.n = self.n+1 
677        # once completed show the plot (cannot show intermediate plotin)
678        print "**** Done step: makeplot"
679        if (self.n == self.howmanyplots): ppplot.save(mode=self.out,filename=self.filename,folder=self.folder)
680        # SAVE A PICKLE FILE WITH THE self.p ARRAY OF OBJECTS
681        if self.verbose: print "**** Saving session in "+self.filename + ".ppobj"
682        savfile = self.folder + "/" + self.filename + ".ppobj"
683        filehandler = open(savfile, 'w')
684        pickle.dump(self.p, filehandler)
685
686    ###########################################################
687    # plot: a shortcut method for the defineplot + plot chain #
688    ###########################################################
689    def plot(self,extraplot=0):
690        self.defineplot(extraplot=extraplot)
691        self.makeplot()
692
693    #######################################################
694    # getplot: a shortcut method for the get + plot chain #
695    #######################################################
696    def getplot(self,extraplot=0):
697        self.get()
698        self.plot(extraplot=extraplot)
699
700    ###################################################################
701    # getdefineplot: a shortcut method for the get + defineplot chain #
702    ###################################################################
703    def getdefineplot(self,extraplot=0):
704        self.get()
705        self.defineplot(extraplot=extraplot)
706
707    ##############################################################
708    # f: operation on two pp objects being on status 'definedplot'
709    # this allows for one field being function of another one
710    # e.g. u.f(v) means u will be displayed as a function of v
711    # ... no need to do defineplot after u.f(v), makeplot directly
712    ##############################################################
713    def f(self,other):
714        # preamble: for this operation to work, defineplot() must have been done
715        if self.status != "definedplot":
716            if self.verbose: print "!! WARNING !! performing defineplot on operand"
717            self.defineplot()
718        if other.status != "definedplot":
719            if self.verbose: print "!! WARNING !! performing defineplot on operand"
720            other.defineplot()
721        # check total number of plots
722        if self.howmanyplots != other.howmanyplots:
723               print "!! ERROR !! The two operands do not have the same number of subplots."
724               exit()
725        # and now operation.
726        count = 0
727        while count < self.howmanyplots:
728           sobj = self.p[count] ; oobj = other.p[count]
729           if sobj.field.ndim !=1 or oobj.field.ndim !=1:
730               if self.verbose: print "!! WARNING !! Flattening arrays because more than one-dimensional."
731               sobj.field = np.ravel(sobj.field)
732               oobj.field = np.ravel(oobj.field)
733           sobj.absc = oobj.field
734           sobj.xlabel = oobj.ylabel
735           if sobj.absc.size > sobj.field.size:
736               if self.verbose:
737                   print "!! WARNING !! Trying to define y=f(x) with x and y not at the same size.",sobj.absc.size,sobj.field.size
738                   print "!! WARNING !! Modifying x to fit y size but please check." 
739               sobj.absc = sobj.absc[0:sobj.field.size]
740           count = count + 1
741        return self
742
743    ###########################################################
744    # copyopt: get options from e.g. a parser
745    # ... allow for simple scripting and user-defined settings
746    # ... must be called between defineplot and makeplot
747    # REQUIRED: attributes of opt must be the same as in the pp object
748    ###########################################################
749    def getopt(self,opt):
750        self.superpose = opt.superpose
751        # -- if only one, or less than the number of plots --> we take the first one
752        # -- if as many as number of plots --> OK, each plot has its own setting
753        # (except a few cases such as trans)
754        for iii in range(self.howmanyplots):
755            ###
756            try: self.p[iii].trans = opt.trans
757            except: pass
758            ###
759            try: self.p[iii].div = opt.div
760            except: pass
761            ###
762            try: self.p[iii].logy = opt.logy
763            except: pass
764            ###
765            try: self.p[iii].colorb = opt.colorb[iii]
766            except: 
767                try: self.p[iii].colorb = opt.colorb[0]
768                except: pass
769            ###
770            try: self.p[iii].title = opt.title[iii]
771            except: 
772                try: self.p[iii].title = opt.title[0]
773                except: pass
774            ###
775            try: self.p[iii].xlabel = opt.xlabel[iii]
776            except: 
777                try: self.p[iii].xlabel = opt.xlabel[0]
778                except: pass
779            ###
780            try: self.p[iii].ylabel = opt.ylabel[iii]
781            except: 
782                try: self.p[iii].ylabel = opt.ylabel[0]
783                except: pass
784            ###
785            try: self.p[iii].lstyle = opt.lstyle[iii]
786            except: 
787                try: self.p[iii].lstyle = opt.lstyle[0]
788                except: pass
789            ###
790            try: self.p[iii].proj = opt.proj[iii]
791            except: 
792                try: self.p[iii].proj = opt.proj[0]
793                except: pass
794            ###
795            try: self.p[iii].back = opt.back[iii]
796            except: 
797                try: self.p[iii].back = opt.back[0]
798                except: pass
799            ###
800            try: self.p[iii].area = opt.area[iii]
801            except: 
802                try: self.p[iii].area = opt.area[0]
803                except: pass
804            ###
805            try: self.p[iii].blon = opt.blon[iii]
806            except: 
807                try: self.p[iii].blon = opt.blon[0]
808                except: pass
809            ###
810            try: self.p[iii].blat = opt.blat[iii]
811            except: 
812                try: self.p[iii].blat = opt.blat[0]
813                except: pass
814            ###
815            try: self.p[iii].vmin = opt.vmin[iii]
816            except: 
817                try: self.p[iii].vmin = opt.vmin[0]
818                except: pass
819            ###
820            try: self.p[iii].vmax = opt.vmax[iii]
821            except: 
822                try: self.p[iii].vmax = opt.vmax[0]
823                except: pass
824
825##########################################################
826### THE ONEREQUEST SUBOBJECT TO PP (ON WHICH IT LOOPS) ###
827##########################################################
828class onerequest():
829
830    # default settings. mostly initialized to diagnose problem, except dimplot, nplot, verbose, swap_axes, invert_axes
831    # -------------------------------
832    def __init__(self):
833        self.file  = '!! file: I am not set, damned !!'
834        self.f     = None
835        self.dim   = None
836        self.var   = '!! var: I am not set, damned !!'
837        self.index_x = [] ; self.index_y = [] ; self.index_z = [] ; self.index_t = []
838        self.index_x2d = [] ; self.index_y2d = []
839        self.method_x = '!! method_x: I am not set, damned !!'
840        self.method_y = '!! method_y: I am not set, damned !!'
841        self.method_z = '!! method_z: I am not set, damned !!'
842        self.method_t = '!! method_t: I am not set, damned !!'
843        self.field = None
844        self.name_x = None ; self.name_y = None ; self.name_z = None ; self.name_t = None
845        self.dim_x = None ; self.dim_y = None ; self.dim_z = None ; self.dim_t = None
846        self.field_x = None ; self.field_y = None ; self.field_z = None ; self.field_t = None
847        self.dimplot = 0
848        self.nplot = 1
849        self.absc = None ; self.ordi = None ; self.absclab = None ; self.ordilab = None
850        self.verbose = True
851        self.swap_axes = False ; self.invert_axes = False
852        self.compute = None
853
854    # redefine the + operation. works with objects or numbers.
855    # -------------------------------
856    def __add__(self,other):
857       if isinstance(other,int) or isinstance(other,float):
858           self.field = self.field + other
859       elif other.__class__.__name__ == "onerequest":
860           if self.field.shape != other.field.shape:
861              print "!! ERROR !! The two objects you want to add do not have the same shape.",self.field.shape,other.field.shape
862              exit()
863           else:
864              self.field = self.field + other.field
865              if self.file != other.file: self.file = self.file + "_+_" + other.file
866              if self.var != other.var: self.var = self.var + "_+_" + other.var
867       else:
868           print "!! ERROR !! Operation is being made on objects of the wrong type." ; exit()
869
870    # redefine the - operation. works with objects or numbers.
871    # -------------------------------
872    def __sub__(self,other):
873       if isinstance(other,int) or isinstance(other,float):
874           self.field = self.field - other
875       elif other.__class__.__name__ == "onerequest":
876           if self.field.shape != other.field.shape:
877              print "!! ERROR !! The two objects you want to add do not have the same shape.",self.field.shape,other.field.shape
878              exit()
879           else:
880              self.field = self.field - other.field
881              if self.file != other.file: self.file = self.file + "_-_" + other.file
882              if self.var != other.var: self.var = self.var + "_-_" + other.var
883       else:
884           print "!! ERROR !! Operation is being made on objects of the wrong type." ; exit()
885
886    # redefine the * operation. works with objects or numbers.
887    # -------------------------------
888    def __mul__(self,other):
889       if isinstance(other,int) or isinstance(other,float):
890           self.field = self.field * other
891       elif other.__class__.__name__ == "onerequest":
892           if self.field.shape != other.field.shape:
893              print "!! ERROR !! The two objects you want to add do not have the same shape.",self.field.shape,other.field.shape
894              exit()
895           else:
896              self.field = self.field * other.field
897              if self.file != other.file: self.file = self.file + "_*_" + other.file
898              if self.var != other.var: self.var = self.var + "_*_" + other.var
899       else:
900           print "!! ERROR !! Operation is being made on objects of the wrong type." ; exit()
901
902    # redefine the / operation. works with objects or numbers.
903    # -------------------------------
904    def __div__(self,other):
905       if isinstance(other,int) or isinstance(other,float):
906           self.field = self.field / other
907       elif other.__class__.__name__ == "onerequest":
908           if self.field.shape != other.field.shape:
909              print "!! ERROR !! The two objects you want to add do not have the same shape.",self.field.shape,other.field.shape
910              exit()
911           else:
912              self.field = self.field / other.field
913              if self.file != other.file: self.file = self.file + "_/_" + other.file
914              if self.var != other.var: self.var = self.var + "_/_" + other.var
915       else:
916           print "!! ERROR !! Operation is being made on objects of the wrong type." ; exit()
917
918    # redefine the power operation.
919    # -------------------------------
920    def __pow__(self,num):
921       if isinstance(num,int) or isinstance(num,float):
922           self.field = self.field ** num
923           self.file = self.file + "_^_"+str(num)
924           self.var = self.var + "_^_"+str(num)
925       else:
926           print "!! ERROR !! To define a power, either an int or a float is needed." ; exit()
927
928    # open a file. for now it is netcdf. TBD for other formats.
929    # check that self.var is inside.
930    # -------------------------------
931    def openfile(self):
932        if not os.path.exists(self.file): print '!! ERROR !! I could not find the following file: '+self.file ; exit()
933        if not os.path.isfile(self.file): print '!! ERROR !! This does not appear to be a file: '+self.file ; exit()
934        self.f = netCDF4.Dataset(self.file)
935        if self.verbose: print "**** OK. Opened file "+self.file
936        if self.var not in self.f.variables.keys(): 
937            print '!! ERROR !! File '+self.file+' does not contain variable: '+self.var
938            print '..... try instead with ',self.f.variables.keys() ; exit()
939
940    # get x,y,z,t dimensions from NETCDF file
941    # TBD: user could request for a specific altitude dimension
942    # TBD: staggered variables could request specific dimensions
943    # -------------------------------
944    def getdim(self):
945          # GET SIZES OF EACH DIMENSION
946          if self.verbose: print "**** OK. Found variable "+self.var
947          shape = self.f.variables[self.var].shape
948          self.dim = len(shape)
949          if self.dim == 1:
950              if self.verbose: print "**** OK. 1D field. I assume this varies with time."
951              self.dim_x = 1 ; self.dim_y = 1 ; self.dim_z = 1 ; self.dim_t = shape[0]
952          elif self.dim == 2:
953              if self.verbose: print "**** OK. 2D field. I assume this is not-time-varying lat-lon map."
954              self.dim_x = shape[1] ; self.dim_y = shape[0] ; self.dim_z = 1 ; self.dim_t = 1
955          elif self.dim == 3:
956              if self.verbose: print "**** OK. 3D field. I assume this is time-varying lat-lon map."
957              self.dim_x = shape[2] ; self.dim_y = shape[1] ; self.dim_z = 1 ; self.dim_t = shape[0]
958          elif self.dim == 4:
959              if self.verbose: print "**** OK. 4D field."
960              self.dim_x = shape[3] ; self.dim_y = shape[2] ; self.dim_z = shape[1] ; self.dim_t = shape[0]
961          # LONGITUDE. Try preset fields. If not present set grid points axis.
962          self.name_x = "nothing"
963          for c in glob_listx:
964            if c in self.f.variables.keys():
965             self.name_x = c
966          if self.name_x == "nothing":
967            self.field_x = np.array(range(self.dim_x))
968            self.name_x = "x grid points"
969          else:
970            self.field_x = self.f.variables[self.name_x]
971          # LATITUDE. Try preset fields. If not present set grid points axis.
972          self.name_y = "nothing"
973          for c in glob_listy:
974            if c in self.f.variables.keys():
975             self.name_y = c
976          if self.name_y == "nothing":
977            self.field_y = np.array(range(self.dim_y))
978            self.name_y = "y grid points"
979          else:
980            self.field_y = self.f.variables[self.name_y]
981          # ensure that lon and lat are 2D fields
982          # 1. simple 1D case (not time-varying)
983          if len(self.field_x.shape)*len(self.field_y.shape) == 1:
984               if self.verbose: print "**** OK. recasting lon and lat as 2D fields." 
985               [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y)
986          # 2. complex 3D case (time-varying, actually just copied over time axis)
987          elif len(self.field_x.shape)*len(self.field_y.shape) == 9:
988               if self.verbose: print "**** OK. reducing lon and lat as 2D fields. get rid of time."
989               self.field_x = self.field_x[0,:,:]
990               self.field_y = self.field_y[0,:,:]
991          # if xy axis are apparently undefined, set 2D grid points axis.
992          if "grid points" not in self.name_x:
993            if self.field_x.all() == self.field_x[0,0]:
994               print "!! WARNING !! xy axis look undefined. creating a non-dummy ones."
995               self.field_x = np.array(range(self.dim_x)) ; self.name_x = "x grid points"
996               self.field_y = np.array(range(self.dim_y)) ; self.name_y = "y grid points"
997               [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y)
998          if self.dim_x > 1: 
999               if self.verbose: print "**** OK. x axis %4.0f values [%5.1f,%5.1f]" % (self.dim_x,self.field_x.min(),self.field_x.max())
1000          if self.dim_y > 1: 
1001               if self.verbose: print "**** OK. y axis %4.0f values [%5.1f,%5.1f]" % (self.dim_y,self.field_y.min(),self.field_y.max())
1002          # ALTITUDE. Try preset fields. If not present set grid points axis.
1003          # WARNING: how do we do if several are available?
1004          self.name_z = "nothing"
1005          for c in glob_listz:
1006            if c in self.f.variables.keys():
1007             self.name_z = c
1008          if self.name_z == "nothing":
1009            self.field_z = np.array(range(self.dim_z))
1010            self.name_z = "z grid points"
1011          else:
1012            self.field_z = self.f.variables[self.name_z][:] # specify dimension
1013                                                            # TBD: have to check that this is not a 3D field
1014          if self.dim_z > 1: 
1015               if self.verbose: print "**** OK. z axis %4.0f values [%5.1f,%5.1f]" % (self.dim_z,self.field_z.min(),self.field_z.max())
1016          # TIME. Try preset fields.
1017          self.name_t = "nothing"
1018          for c in glob_listt:
1019            if c in self.f.dimensions.keys():
1020             self.name_t = c
1021          try:
1022            # speed up: only get first value, last one, and step.
1023            dafirst = self.f.variables[self.name_t][0]
1024            dalast = self.f.variables[self.name_t][self.dim_t-1]
1025            step = self.f.variables[self.name_t][1] - dafirst
1026            self.field_t = np.arange(dafirst,dalast,step)
1027            if dafirst == dalast: self.field_t = np.array([dafirst])
1028          except:
1029            # ... or if a problem encountered, define a simple time axis
1030            self.field_t = np.array(range(self.dim_t))
1031            self.name_t = "t grid points"
1032          if self.dim_t > 1: 
1033               if self.verbose: print "**** OK. t axis %4.0f values [%5.1f,%5.1f]" % (self.dim_t,self.field_t.min(),self.field_t.max())     
1034
1035    # get list of index to be retrieved for time axis
1036    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1037    # -------------------------------
1038    def getindextime(self,dalist,ind,stride):
1039        if self.method_t == "free": 
1040            self.index_t = np.arange(0,self.dim_t,stride)
1041            if self.dim_t > 1: 
1042                self.dimplot = self.dimplot + 1 
1043                if self.verbose: print "**** OK. t values. all."
1044            else:               
1045                self.method_t = "fixed"
1046                if self.verbose: print "**** OK. no t dimension."
1047        elif self.method_t == "comp":
1048            start = np.argmin( np.abs( self.field_t - dalist[ind][0] ) )
1049            stop = np.argmin( np.abs( self.field_t - dalist[ind][1] ) )
1050            self.index_t = np.arange(start,stop,stride)
1051            if self.verbose: print "**** OK. t values. comp over interval ",self.field_t[start],self.field_t[stop]," nvalues=",self.index_t.size
1052        elif self.method_t == "fixed":
1053            self.index_t.append( np.argmin( np.abs( self.field_t - dalist[ind][0] ) ))
1054            if self.verbose: print "**** OK. t values",self.field_t[self.index_t]
1055        else:
1056            print "!! ERROR !! method "+self.method_t+" not supported"
1057        self.index_t = np.array(self.index_t)
1058             
1059    # get list of index to be retrieved for vertical axis
1060    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1061    # -------------------------------
1062    def getindexvert(self,dalist,ind,stride):
1063        if self.method_z == "free": 
1064            self.index_z = np.arange(0,self.dim_z,stride)
1065            if self.dim_z > 1: 
1066                self.dimplot = self.dimplot + 1
1067                if self.verbose: print "**** OK. z values. all."
1068            else:               
1069                self.method_z = "fixed"
1070                if self.verbose: print "**** OK. no z dimension."
1071        elif self.method_z == "comp":
1072            start = np.argmin( np.abs( self.field_z - dalist[ind][0] ) )
1073            stop = np.argmin( np.abs( self.field_z - dalist[ind][1] ) )
1074            self.index_z = np.arange(start,stop,stride)
1075            if self.verbose: print "**** OK. z values. comp over interval",self.field_z[start],self.field_z[stop]," nvalues=",self.index_z.size
1076        elif self.method_z == "fixed":
1077            self.index_z.append( np.argmin( np.abs( self.field_z - dalist[ind][0] ) ))
1078            if self.verbose: print "**** OK. z values",self.field_z[self.index_z]
1079        else:
1080            if self.verbose: print "!! ERROR !! method "+self.method_z+" not supported"
1081        self.index_z = np.array(self.index_z)
1082
1083    # get list of index to be retrieved for horizontal grid
1084    # --> index_x and index_y are slices to be retrieved from NETCDF files
1085    # --> index_x2d and index_y2d are the actual (x,y) coordinates corresponding to each relevant point
1086    # [this is slightly more complicated because 2D arrays for lat-lon projection possibly irregular]
1087    # NB: to append index we use lists (the most convenient) then we convert into a numpy.array
1088    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1089    # -------------------------------
1090    def getindexhori(self,dalistx,dalisty,indx,indy,stridex,stridey):
1091        ## get what is the method over x and y axis
1092        test = self.method_x+self.method_y
1093        ## CASE 0, EASY CASES:
1094        ## - LAT IS FREE (we do here what must be done whatever LON case is)
1095        ## - LON IS FREE (we do here what must be done whatever LAT case is)
1096        ## - LAT IS COMP AND LON IS FREE
1097        ## - LON IS COMP AND LAT IS FREE
1098        if self.method_x == "free" or test in ["compfree","compcomp"]:
1099            self.index_x = range(0,self.dim_x,stridex)
1100            if self.dim_x > 1: 
1101                if self.method_x == "free": self.dimplot = self.dimplot + 1
1102                if self.verbose: print "**** OK. x values. all."
1103            else:               
1104                self.method_x = "fixed"
1105                if self.verbose: print "**** OK. no x dimension."
1106        if self.method_y == "free" or test in ["compfree","compcomp"]:
1107            self.index_y = range(0,self.dim_y,stridey)
1108            if self.dim_y > 1: 
1109                if self.method_y == "free": self.dimplot = self.dimplot + 1
1110                if self.verbose: print "**** OK. y values. all."
1111            else:               
1112                self.method_y = "fixed"
1113                if self.verbose: print "**** OK. no y dimension."
1114        ## CASE 0 above, this is just for continuity.
1115        if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]:
1116            self.index_x2d = self.index_x
1117            self.index_y2d = self.index_y
1118        ## AND NOW THE LITTLE BIT MORE COMPLICATED CASES
1119        ## CASE 1 LAT AND LON ARE FIXED
1120        elif test == "fixedfixed":
1121            idy,idx = np.unravel_index( np.argmin( ( self.field_x - dalistx[indx][0])**2 + (self.field_y - dalisty[indy][0])**2 ), self.field_x.shape ) 
1122                          #TBD: pb with staggered coord
1123            if idx not in self.index_x:  self.index_x.append(idx)
1124            if idy not in self.index_y:  self.index_y.append(idy)
1125            self.index_x2d.append(idx)
1126            self.index_y2d.append(idy)
1127        ## CASE 2 LON IS FIXED BUT NOT LAT
1128        elif test in ["fixedfree","fixedcomp"]:
1129            # find where are requested x values for each y on the free dimension
1130            # NB: this does not work for non-bijective cases e.g. polar stereographic
1131            for iy in range(self.dim_y):
1132              idx = np.argmin( np.abs( self.field_x[iy,:] - dalistx[indx][0] ) )
1133              # if comp is requested we select only indexes which yield values between requested min and max
1134              storeval = (self.method_y == "comp") and (self.field_y[iy,idx] > dalisty[indy][0]) and (self.field_y[iy,idx] < dalisty[indy][1])
1135              storeval = storeval or (self.method_y == "free")
1136              if storeval:
1137                  if idx not in self.index_x:  self.index_x.append(idx)
1138                  if iy not in self.index_y and self.method_y == "comp": self.index_y.append(iy)
1139                  if idx not in self.index_x2d or iy not in self.index_y2d:
1140                    self.index_x2d.append(idx)
1141                    self.index_y2d.append(iy)
1142        ## CASE 3 LAT IS FIXED BUT NOT LON
1143        elif test in ["freefixed","compfixed"]:
1144            # find where are requested y values for each x on the free dimension
1145            # NB: this does not work for non-bijective cases e.g. polar stereographic
1146            for ix in range(self.dim_x):
1147              idy = np.argmin( np.abs( self.field_y[:,ix] - dalisty[indy][0] ) )
1148              # if comp is requested we select only indexes which yield values between requested min and max
1149              storeval = (self.method_x == "comp") and (self.field_x[idy,ix] > dalistx[indx][0]) and (self.field_x[idy,ix] < dalistx[indx][1])
1150              storeval = storeval or (self.method_x == "free")
1151              if storeval:
1152                  if idy not in self.index_y:  self.index_y.append(idy)
1153                  if ix not in self.index_x and self.method_x == "comp": self.index_x.append(ix)
1154                  if ix not in self.index_x2d or idy not in self.index_y2d:
1155                    self.index_x2d.append(ix)
1156                    self.index_y2d.append(idy)
1157        ## check index tab
1158        if len(self.index_x) == 0 or len(self.index_y) == 0:
1159            print "!! ERROR !! no indices found. check prescribed latitudes or longitudes" ; exit()
1160        ## ensure the array is a numpy array for getfield to work
1161        self.index_x = np.array(self.index_x)
1162        self.index_y = np.array(self.index_y)
1163        self.index_x2d = np.array(self.index_x2d)
1164        self.index_y2d = np.array(self.index_y2d)
1165        ### print extrema
1166        printx = self.field_x[np.ix_(self.index_y2d, self.index_x2d)]
1167        printy = self.field_y[np.ix_(self.index_y2d, self.index_x2d)]
1168        if self.verbose: 
1169            print "**** OK. x values (min,max).", printx.min(),printx.max()
1170            print "**** OK. y values (min,max).", printy.min(),printy.max()
1171
1172    # get the field from the NETCDF file and perform averages
1173    # -------------------------------
1174    def getfield(self):
1175        ## first tell what is to be done
1176        if self.dimplot > 2:                       print "**** !! ERROR !! "+str(self.dimplot)+"D plots not supported!" ; exit()
1177        elif self.dimplot == 0 and self.verbose:   print "**** OK. 0D value requested."
1178        elif self.dimplot == 1 and self.verbose:   print "**** OK. 1D plot requested."
1179        elif self.verbose:                         print "**** OK. 2D section requested."
1180        # well, now get field from netcdf file
1181        # part below is necessary otherwise there is an index error below
1182        if self.index_x.size == 1: self.index_x = self.index_x[0]
1183        if self.index_y.size == 1: self.index_y = self.index_y[0]
1184        if self.index_z.size == 1: self.index_z = self.index_z[0]
1185        if self.index_t.size == 1: self.index_t = self.index_t[0]
1186        # then retrieve what is requested by user
1187        # each self.dim case corresponds to tests in the beginning of getdim.
1188        time0 = timelib.time()
1189        if self.verbose: print "**** OK. I am getting values from files. Please wait."
1190        if self.dim == 1: 
1191            nt = self.index_t.size ; nz = 1 ; ny = 1 ; nx = 1
1192            self.field = self.f.variables[self.var][self.index_t]
1193        elif self.dim == 2:
1194            nt = 1 ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size
1195            self.field = self.f.variables[self.var][self.index_y,self.index_x]
1196        elif self.dim == 3:
1197            nt = self.index_t.size ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size
1198            self.field = self.f.variables[self.var][self.index_t,self.index_y,self.index_x]
1199            # this is far faster than retrieving each term with a loop
1200        elif self.dim == 4:
1201            nt = self.index_t.size ; nz = self.index_z.size ; ny = self.index_y.size ; nx = self.index_x.size
1202            self.field = self.f.variables[self.var][self.index_t,self.index_z,self.index_y,self.index_x]
1203        else:
1204            print "!! ERROR !! field would have more than four dimensions ?" ; exit()
1205        # NB: ... always 4D array but possibly with "size 1" dimensions
1206        #     ... if one dimension is missing because 1D 2D or 3D requests, make it appear again
1207        self.field = np.reshape(self.field,(nt,nz,ny,nx))
1208        if self.verbose: print "**** OK. I got %7.1e values. This took me %6.4f seconds" % (nx*ny*nz*nt,timelib.time() - time0)
1209        if self.verbose: print "**** OK. I got var "+self.var+" with shape",self.field.shape
1210        # reduce coordinates to useful points
1211        # ... TBD: this should be ordered in the case of non-regular projections
1212        if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]:
1213          # we need 2D coordinates (free) or we get broadcast problem (comp) so we use np.ix
1214          self.field_x = self.field_x[np.ix_(self.index_y2d, self.index_x2d)]
1215          self.field_y = self.field_y[np.ix_(self.index_y2d, self.index_x2d)]
1216        else:
1217          # we are OK with 1D coordinates
1218          self.field_x = self.field_x[self.index_y2d, self.index_x2d]
1219          self.field_y = self.field_y[self.index_y2d, self.index_x2d]
1220        self.field_z = self.field_z[self.index_z]
1221        self.field_t = self.field_t[self.index_t]
1222        # now have to obtain the new indexes which correspond to the extracted self.field
1223        # for it to work with unique index, ensure that any index_* is a numpy array
1224        if not isinstance(self.index_x, np.ndarray): self.index_x = np.array([self.index_x])
1225        if not isinstance(self.index_y, np.ndarray): self.index_y = np.array([self.index_y])
1226        if not isinstance(self.index_z, np.ndarray): self.index_z = np.array([self.index_z])
1227        if not isinstance(self.index_t, np.ndarray): self.index_t = np.array([self.index_t])
1228        for val in self.index_x: self.index_x2d[np.where(self.index_x2d == val)] = np.where(self.index_x == val)[0]
1229        for val in self.index_y: self.index_y2d[np.where(self.index_y2d == val)] = np.where(self.index_y == val)[0]
1230        for val in self.index_z: self.index_z  [np.where(self.index_z   == val)] = np.where(self.index_z == val)[0]
1231        for val in self.index_t: self.index_t  [np.where(self.index_t   == val)] = np.where(self.index_t == val)[0]
1232               ##### VERY EXPENSIVE
1233               ## recast self.field with 2D horizontal arrays because we might have extracted
1234               ## more than what is to be actually plot or computed, in particular for comps on 2D lat,lon coordinates
1235               #self.field = self.field[np.ix_(self.index_t,self.index_z,self.index_y2d,self.index_x2d)]
1236               #(nt,nz,ny,nx) = self.field.shape       
1237        # extract relevant horizontal points
1238        # TBD: is compfree OK with computing on irregular grid?
1239        test = self.method_x + self.method_y
1240        if test in ["fixedfixed","freefree"]:
1241          pass
1242        elif test in ["fixedfree","fixedcomp"] or test in ["freefixed","compfixed"]: 
1243          time0 = timelib.time() 
1244          # prepare the loop on all relevant horizontal points
1245          if self.method_x in ["comp","free"]:   
1246              nnn = self.index_x2d.shape[0] ; what_I_am_supposed_to_do = "keepx"
1247          elif self.method_y in ["comp","free"]: 
1248              nnn = self.index_y2d.shape[0] ; what_I_am_supposed_to_do = "keepy" 
1249          # LOOP to extract only useful values over horizontal dimensions
1250          # only take diagonal terms, do not loop on all self.index_x2d*self.index_y2d
1251          # ... this method is fast enough, perhaps there is a faster way though
1252          # ... (for sure the method with np.diag is much slower)
1253          for iii in range(nnn):
1254           ix = self.index_x2d[iii] ; iy = self.index_y2d[iii]
1255           for iz in self.index_z:
1256            for it in self.index_t:
1257              if what_I_am_supposed_to_do == "keepx":    self.field[it,iz,0,ix] = self.field[it,iz,iy,ix]
1258              elif what_I_am_supposed_to_do == "keepy":  self.field[it,iz,iy,0] = self.field[it,iz,iy,ix]
1259          if self.verbose: print "**** OK. I got to pick the right values for your request. This took me %6.4f seconds" % (timelib.time() - time0)
1260          # we only keep the one value that was modified on the dimension which is not free
1261          if what_I_am_supposed_to_do == "keepx":     self.field = self.field[:,:,0,:] ; ny = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1262          elif what_I_am_supposed_to_do == "keepy":   self.field = self.field[:,:,:,0] ; nx = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1263        # make a mask in case there are non-NaN missing values. (what about NaN missing values?)
1264        # ... this is important for computations below (see ppcompute)
1265        masked = np.ma.masked_where(np.abs(self.field) > 1e25,self.field)
1266        if masked.mask.any() == True:
1267             if self.verbose: print "!! WARNING !! Values over +-1e25 are considered missing values."
1268             self.field = masked
1269             self.field.set_fill_value([np.NaN])
1270        # now ready to compute [TBD?? we would like to have e.g. mean over x,y and min over t]
1271        if self.method_t == "comp":
1272            if self.verbose: print "**** OK. Computing over t axis."
1273            if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=0)
1274            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=0)
1275            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=0)
1276            else: print "!! ERROR !! operation not supported." ; exit()
1277            nt = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1278        if self.method_z == "comp": 
1279            if self.verbose: print "**** OK. Computing over z axis."
1280            if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=1)
1281            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=1)
1282            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=1)
1283            nz = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1284        if self.method_y == "comp": 
1285            if self.verbose: print "**** OK. Computing over y axis."
1286            if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=2)
1287            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=2)
1288            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=2)
1289            ny = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1290            if self.field_x.ndim == 2: self.field_x = self.field_x[0,:] # TBD: this is OK for regular grid but not for irregular
1291        if self.method_x == "comp":
1292            if self.verbose: print "**** OK. Computing over x axis."
1293            if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=3)
1294            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=3)
1295            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=3)
1296            nx = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1297            if self.field_y.ndim == 2: self.field_y = self.field_y[:,0] # TBD: this is OK for regular grid but not for irregular
1298        # remove all dimensions with size 1 to prepare plot (and check the resulting dimension with dimplot)
1299        self.field = np.squeeze(self.field)
1300        if self.field.ndim != self.dimplot: 
1301            print "!! ERROR !! Problem: self.field is different than plot dimensions", self.field.ndim, self.dimplot ; exit()
1302        if self.verbose: 
1303            print "**** OK. Final shape for "+self.var+" after averaging and squeezing",self.field.shape
1304
1305    # define coordinates for plot
1306    # -------------------------------
1307    def definecoord(self):
1308        I_got_abs = False ; I_got_ord = False
1309        # here is the thing. time is usually taken as an abscissa so we start with time.
1310        if self.method_t ==  "free": 
1311            self.absc = self.field_t ; self.absclab = self.name_t
1312            I_got_abs = True
1313        # then we usually have x as an abscissa.
1314        if self.method_x == "free":
1315            if I_got_abs: 
1316                self.ordi = self.field_x ; self.ordilab = self.name_x
1317                I_got_ord = True
1318            else:         
1319                self.absc = self.field_x ; self.absclab = self.name_x
1320                I_got_abs = True
1321        # ... or we have y
1322        if self.method_y == "free":
1323            if I_got_abs:   
1324                self.ordi = self.field_y ; self.ordilab = self.name_y
1325                I_got_ord = True
1326            else:         
1327                self.absc = self.field_y ; self.absclab = self.name_y
1328                I_got_abs = True
1329        # ... and we end with z because it is usually not an abscissa (profiles).
1330        if self.method_z == "free":
1331            if self.field_z[0] > self.field_z[1]:
1332                self.invert_axes = True # the axis will be turned upside-down
1333            if I_got_abs: 
1334                self.ordi = self.field_z ; self.ordilab = self.name_z
1335                I_got_ord = True
1336            else:
1337                self.absc = self.field_z ; self.absclab = self.name_z
1338                I_got_abs = True
1339                self.swap_axes = True # says that altitude is not supposed to remain as an abscissa
1340        if I_got_abs and self.verbose: print "**** OK. abscissa:",self.absclab, self.absc.shape
1341        if I_got_ord and self.verbose: print "**** OK. ordinate:",self.ordilab, self.ordi.shape
Note: See TracBrowser for help on using the repository browser.