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

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

UTIL PYTHON planetoplot_v2

RENAMED VARIABLES TO BE SIMPLER or COMPLIANT WITH MATPLOTLIB

colorb >> colorbar
lstyle >> linestyle
label >> legend

stridevecx >> svx
stridevecy >> svy
stridex >> sx
stridey >> sy
stridez >> sz
stridet >> st

(in ppplot only)
field >> f
absc >> x
ordi >> y
addvecx >> vx
addvecy >> vy
addontour >> c

ALSO added (same reason of compatibility)

  • the possibility to call ppplot.show() as in matplotlib
  • a shortcut method for a ppplot object: the so-called .makeshow()
File size: 93.8 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###################################
26# where settings files are located...
27# ... this can be hardcoded here
28whereset = None
29whereset = ppcompute.findset(whereset)
30# ... we load user-defined automatic settings from set_ppclass.txt
31zefile = "set_ppclass.txt"
32glob_listx = [] ; glob_listy = [] ; glob_listz = [] ; glob_listt = []
33glob_listarea = []
34try: 
35    f = open(whereset+zefile, 'r') ; lines = f.readlines()
36    for stuff in lines[5].strip().split(';'): glob_listx.append(stuff)
37    for stuff in lines[8].strip().split(';'): glob_listy.append(stuff)
38    for stuff in lines[11].strip().split(';'): glob_listz.append(stuff)
39    for stuff in lines[14].strip().split(';'): glob_listt.append(stuff)
40    for stuff in lines[17].strip().split(';'): glob_listarea.append(stuff)
41except IOError: 
42    print "PPCLASS 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
104# look for comas in the input name to separate different names (files, variables,etc ..)
105# (needed by readslices)
106def separatenames (name):
107    if name is None: names = None
108    else:
109      names = [] ; stop = 0 ; currentname = name
110      while stop == 0:
111        indexvir = currentname.find(',')
112        if indexvir == -1: stop = 1 ; name1 = currentname
113        else: name1 = currentname[0:indexvir]
114        names = np.concatenate((names,[name1]))
115        currentname = currentname[indexvir+1:len(currentname)]
116    return names
117
118#######################
119### THE MAIN OBJECT ###
120#######################
121class pp():
122
123    # print out a help string when help is invoked on the object
124    def __repr__(self):
125        whatprint = 'pp object. \"help(pp)\" for more information\n'
126        return whatprint
127
128    # default settings
129    # -- user can define settings by two methods.
130    # -- 1. yeah = pp(file="file.nc")
131    # -- 2. yeah = pp() ; yeah.file = "file.nc"
132    def __init__(self,file=None,var="notset",\
133                      filegoal=None,vargoal=None,\
134                      x=None,y=None,z=None,t=None,\
135                      sx=1,sy=1,\
136                      sz=1,st=1,\
137                      svx=1,\
138                      svy=1,\
139                      compute="mean",\
140                      verbose=False,\
141                      quiet=False,\
142                      noproj=False,\
143                      superpose=False,\
144                      plotin=None,\
145                      forcedimplot=-1,\
146                      out="gui",\
147                      filename="myplot",\
148                      folder="./",\
149                      includedate=True,\
150                      res=150.,\
151                      xlabel=None,ylabel=None,\
152                      xcoeff=None,ycoeff=None,\
153                      nxticks=10,nyticks=10,\
154                      proj=None,\
155                      vmin=None,vmax=None,\
156                      div=None,\
157                      colorbar=None,\
158                      linestyle=None,\
159                      marker=None,\
160                      color=None,\
161                      legend=None,\
162                      changetime=None,\
163                      units=None,\
164                      savtxt=False,\
165                      modx=None,\
166                      xp=16,yp=8,\
167                      missing=1.e25,\
168                      title=None):
169        self.request = None
170        self.nrequest = 0
171        self.nfin = 0 ; self.nvin = 0
172        self.nplotx = None ; self.nploty = None
173        self.nplotz = None ; self.nplott = None
174        self.status = "init"
175        self.fig = None ; self.subv = None ; self.subh = None 
176        self.n = 0 ; self.howmanyplots = 0
177        self.nplot = 0
178        self.p = None
179        self.customplot = False
180        self.f = None
181        self.l = None
182        ## what could be defined by the user
183        self.file = file
184        self.var = var
185        self.filegoal = filegoal
186        self.vargoal = vargoal
187        self.x = x ; self.y = y   ## if None, free dimension
188        self.z = z ; self.t = t   ## if None, free dimension
189        self.sx = sx ; self.sy = sy
190        self.sz = sz ; self.st = st
191        self.svx = svx
192        self.svy = svy
193        self.compute = compute
194        self.verbose = verbose
195        self.quiet = quiet
196        self.noproj = noproj
197        self.plotin = plotin
198        self.superpose = superpose
199        self.forcedimplot = forcedimplot
200        self.out = out
201        self.filename = filename
202        self.res = res
203        self.folder = folder
204        self.includedate = includedate
205        self.changetime = changetime
206        self.savtxt = savtxt
207        self.modx = modx
208        self.missing = missing
209        ## here are user-defined plot settings
210        ## -- if not None, valid on all plots in the pp() objects
211        self.xlabel = xlabel ; self.xcoeff = xcoeff
212        self.ylabel = ylabel ; self.ycoeff = ycoeff
213        self.proj = proj
214        self.vmin = vmin ; self.vmax = vmax
215        self.div = div
216        self.colorbar = colorbar
217        self.linestyle = linestyle
218        self.marker = marker
219        self.color = color
220        self.legend = legend
221        self.units = units
222        self.title = title
223        self.xp = xp ; self.yp = yp
224        self.nxticks = nxticks ; self.nyticks = nyticks
225
226    # print status
227    def printstatus(self):
228      if not self.quiet:
229        if self.filename == "THIS_IS_A_CLONE":
230            pass
231        else:
232            print "**** PPCLASS. Done step: " + self.status
233
234    # print attributes
235    def printme(self):
236        for k, v in vars(self).items():
237            print k,v
238
239    #####################################################
240    # EMULATE OPERATORS + - * / ** << FOR PP() OBJECTS  #
241    #####################################################
242
243    # define the operation <<
244    # ... e.g. obj2 << obj1
245    # ... means: get init for pp object obj2 from another pp object obj1
246    # ... (this should solve the affectation trap obj2 = obj1)
247    def __lshift__(self,other):
248        if other.__class__.__name__ == "pp":
249            self.file = other.file
250            self.var = other.var
251            self.filegoal = other.filegoal
252            self.vargoal = other.vargoal
253            self.x = other.x ; self.y = other.y   ## if None, free dimension
254            self.z = other.z ; self.t = other.t   ## if None, free dimension
255            self.sx = other.sx ; self.sy = other.sy
256            self.sz = other.sz ; self.st = other.st
257            self.verbose = other.verbose
258            self.noproj = other.noproj
259            self.plotin = other.plotin
260            self.superpose = other.superpose
261            self.forcedimplot = other.forcedimplot
262            self.out = other.out
263            self.filename = other.filename
264            self.folder = other.folder
265            self.xlabel = other.xlabel ; self.xcoeff = other.xcoeff
266            self.ylabel = other.ylabel ; self.ycoeff = other.ycoeff
267            self.proj = other.proj
268            self.vmin = other.vmin ; self.vmax = other.vmax
269            self.div = other.div
270            self.colorbar = other.colorbar
271            self.linestyle = other.linestyle
272            self.marker = other.marker
273            self.color = other.color
274            self.legend = other.legend
275            self.units = other.units
276            self.title = other.title
277            self.includedate = other.includedate
278            self.changetime = other.changetime
279            self.savtxt = other.savtxt
280            self.modx = other.modx
281            self.xp = other.xp ; self.yp = other.yp
282            self.missing = other.missing
283            self.nxticks = other.nxticks ; self.nyticks = other.nyticks
284        else:
285            print "!! ERROR !! argument must be a pp object." ; exit()
286
287    # check the compatibility of two objects for operations
288    # --> if other is a pp class, test sizes and return isnum = False
289    # --> if other is an int or a float, return isnum = True
290    # --> otherwise, just print an error and exit
291    def checktwo(self,other):
292        if other.__class__.__name__ == "pp":
293          isnum = False
294          if self.status in ["init","defined"] or other.status in ["init","define"]: 
295             print "!! ERROR !! Please use .retrieve to get fields for plots with one of your pp operands." ; exit()
296          if self.nfin   != other.nfin   or \
297             self.nvin   != other.nvin   or \
298             self.nplott != other.nplott or \
299             self.nplotz != other.nploty or \
300             self.nploty != other.nploty or \
301             self.nplotx != other.nplotx :
302               print "!! ERROR !! The two operands do not have the same number of files, variables, t z y x requests."
303               exit()
304        elif isinstance(other,int) or isinstance(other,float):
305          isnum = True
306        else:
307          print "!! ERROR !! The operand is neither a pp class nor an integer or a float." ; exit()
308        return isnum
309
310    # define a selective copy of a pp() object for operations
311    # ... copy.copy() is not conservative (still acts like a pointer)
312    # ... copy.deepcopy() does not work with netCDF objects
313    # so what is done here is a copy of everything except
314    # (to avoid sharing with self and therefore modifying self through operations)
315    # - request attribute of pp() object
316    # - field attribute of the onerequest() objects
317    def selective_copy(self):
318        if self.status in ["init","defined"]:
319            print "!! ERROR !! Please use .retrieve to get fields for the object you want to copy from." ; exit()
320        the_clone = pp()
321        for k, v in vars(self).items():
322           if k != "request":
323               setattr(the_clone,k,v)
324        the_clone.verbose = False
325        the_clone.filename = "THIS_IS_A_CLONE" # trick to avoid additional outputs
326        the_clone.define()
327        for i in range(self.nfin):
328         for j in range(self.nvin):
329          for t in range(self.nplott):
330           for z in range(self.nplotz):
331            for y in range(self.nploty):
332             for x in range(self.nplotx):
333              obj_ref = self.request[i][j][t][z][y][x]
334              obj = the_clone.request[i][j][t][z][y][x]
335              for k, v in vars(obj_ref).items():
336               if k != "field":
337                setattr(obj,k,v)
338        the_clone.status = "retrieved"
339        the_clone.filename = self.filename
340        return the_clone
341
342    # define the operation + on two objects. or with an int/float.
343    # ... with selective_copy the self object is not modified.
344    def __add__(self,other):
345        isnum = self.checktwo(other)
346        the_clone = self.selective_copy()
347        for i in range(self.nfin):
348         for j in range(self.nvin):
349          for t in range(self.nplott):
350           for z in range(self.nplotz):
351            for y in range(self.nploty):
352             for x in range(self.nplotx):
353              obj = the_clone.request[i][j][t][z][y][x]
354              obj_ref = self.request[i][j][t][z][y][x]
355              if not isnum:   
356                  ope = other.request[i][j][t][z][y][x].field
357                  if ope.ndim == 0: 
358                    ope = float(ope) # if no dimension then this means that ope is a single value (not to be kept as an array)
359                  elif obj_ref.field.shape != ope.shape:
360                    print "!! ERROR !! The two fields for operation do not have the same shape.",obj_ref.field.shape,ope.shape
361                    exit()
362              else:           
363                  ope = other
364              goal = self.vargoal[j] + self.filegoal[i]
365              if ("vector" in goal) or ("contour" in goal):
366                  if self.verbose: print "!! WARNING !! No operation was made on contours and vectors. This can be debatted actually."
367                  obj.field = obj_ref.field
368              else:
369                  obj.field = obj_ref.field + ope
370        return the_clone
371
372    # define the operation - on two objects. or with an int/float.
373    # ... with selective_copy the self object is not modified.
374    def __sub__(self,other):
375        isnum = self.checktwo(other)
376        the_clone = self.selective_copy()
377        for i in range(self.nfin):
378         for j in range(self.nvin):
379          for t in range(self.nplott):
380           for z in range(self.nplotz):
381            for y in range(self.nploty):
382             for x in range(self.nplotx):
383              obj = the_clone.request[i][j][t][z][y][x]
384              obj_ref = self.request[i][j][t][z][y][x]
385              if not isnum:
386                  ope = other.request[i][j][t][z][y][x].field
387                  if ope.ndim == 0: 
388                    ope = float(ope) # if no dimension then this means that ope is a single value (not to be kept as an array)
389                  elif obj_ref.field.shape != ope.shape:
390                    print "!! ERROR !! The two fields for operation do not have the same shape.",obj_ref.field.shape,ope.shape
391                    exit()
392              else:
393                  ope = other
394              goal = self.vargoal[j] + self.filegoal[i]
395              if ("vector" in goal) or ("contour" in goal):
396                  if self.verbose: print "!! WARNING !! No operation was made on contours and vectors. This can be debatted actually."
397                  obj.field = obj_ref.field
398              else:
399                  obj.field = obj_ref.field - ope
400        return the_clone
401
402    # define the operation * on two objects. or with an int/float.
403    # ... with selective_copy the self object is not modified.
404    def __mul__(self,other):
405        isnum = self.checktwo(other)
406        the_clone = self.selective_copy()
407        for i in range(self.nfin):
408         for j in range(self.nvin):
409          for t in range(self.nplott):
410           for z in range(self.nplotz):
411            for y in range(self.nploty):
412             for x in range(self.nplotx):
413              obj = the_clone.request[i][j][t][z][y][x]
414              obj_ref = self.request[i][j][t][z][y][x]
415              if not isnum:
416                  ope = other.request[i][j][t][z][y][x].field
417                  if ope.ndim == 0: 
418                    ope = float(ope) # if no dimension then this means that ope is a single value (not to be kept as an array)
419                  elif obj_ref.field.shape != ope.shape:
420                    print "!! ERROR !! The two fields for operation do not have the same shape.",obj_ref.field.shape,ope.shape
421                    exit()
422              else:
423                  ope = other
424              goal = self.vargoal[j] + self.filegoal[i]
425              if ("vector" in goal) or ("contour" in goal):
426                  if self.verbose: print "!! WARNING !! No operation was made on contours and vectors. This can be debatted actually."
427                  obj.field = obj_ref.field
428              else:
429                  obj.field = obj_ref.field * ope
430        return the_clone
431
432    # define the operation / on two objects. or with an int/float.
433    # ... with selective_copy the self object is not modified.
434    def __div__(self,other):
435        isnum = self.checktwo(other)
436        the_clone = self.selective_copy()
437        for i in range(self.nfin):
438         for j in range(self.nvin):
439          for t in range(self.nplott):
440           for z in range(self.nplotz):
441            for y in range(self.nploty):
442             for x in range(self.nplotx):
443              obj = the_clone.request[i][j][t][z][y][x]
444              obj_ref = self.request[i][j][t][z][y][x]
445              if not isnum:
446                  ope = other.request[i][j][t][z][y][x].field
447                  if ope.ndim == 0: 
448                    ope = float(ope) # if no dimension then this means that ope is a single value (not to be kept as an array)
449                  elif obj_ref.field.shape != ope.shape:
450                    print "!! ERROR !! The two fields for operation do not have the same shape.",obj_ref.field.shape,ope.shape
451                    exit()
452              else:
453                  ope = other
454              goal = self.vargoal[j] + self.filegoal[i]
455              if ("vector" in goal) or ("contour" in goal):
456                  if self.verbose: print "!! WARNING !! No operation was made on contours and vectors. This can be debatted actually."
457                  obj.field = obj_ref.field
458              else:
459                  obj.field = obj_ref.field / ope
460        return the_clone
461
462    # define the reverse operation float/int + object
463    def __radd__(self,other):
464        isnum = self.checktwo(other)
465        if not isnum: print "!! ERROR !! Operand should be a number" ; exit()
466        return self.__add__(other)
467
468    # define the reverse operation float/int - object
469    def __rsub__(self,other):
470        isnum = self.checktwo(other)
471        if not isnum: print "!! ERROR !! Operand should be a number" ; exit()
472        return self.__sub__(other)
473
474    # define the reverse operation float/int * object
475    def __rmul__(self,other):
476        isnum = self.checktwo(other)
477        if not isnum: print "!! ERROR !! Operand should be a number" ; exit()
478        return self.__mul__(other)
479
480    # define the reverse operation float/int / object
481    def __rdiv__(self,other):
482        isnum = self.checktwo(other)
483        if not isnum: print "!! ERROR !! Operand should be a number" ; exit()
484        return self.__div__(other)
485
486    # define the operation ** on one object.
487    # ... with selective_copy the self object is not modified.
488    def __pow__(self,num):
489        the_clone = self.selective_copy()
490        if isinstance(num,int) or isinstance(num,float):
491            for i in range(self.nfin):
492             for j in range(self.nvin):
493              for t in range(self.nplott):
494               for z in range(self.nplotz):
495                for y in range(self.nploty):
496                 for x in range(self.nplotx):
497                  obj  = the_clone.request[i][j][t][z][y][x]
498                  obj_ref = self.request[i][j][t][z][y][x]
499                  goal = self.vargoal[j] + self.filegoal[i]
500                  if ("vector" in goal) or ("contour" in goal):
501                      if self.verbose: print "!! WARNING !! No operation was made on contours and vectors. This can be debatted actually."
502                      obj.field = obj_ref.field
503                  else:
504                      obj.field = obj_ref.field ** num
505        else:
506            print "!! ERROR !! To define a power, either an int or a float is needed." ; exit()
507        return the_clone
508
509    ### TBD: reverse power? for exponentials?
510
511    ##############################################################################################
512    # define method
513    # ---------
514    # ... (file and var are either one string or a vector of strings)
515    # ... the goal of define is to define a 2D array of onerequest() objects (see class below)
516    #     given the number of file, var, x, y, z, t asked by the user
517    # ... objectives for file or var are given through filegoal and vargoal
518    #     --> possible values: main contour vector
519    # ---------
520    # ... then onerequest() objects are being defined more precisely
521    #     by getting index_x index_y index_z index_t
522    #     and setting method_x method_y method_z method_t to either
523    #      - "free" for free dimensions (plot dimensions)
524    #      - "comp" for averages, max, min
525    #      - "fixed" for fixed dimensions (possibly several i.e. multislice)
526    ##############################################################################################
527    def define(self):
528        self.printstatus()
529        # initial check and get dimensions
530        self.file = checktab(self.file,mess="file")
531        self.nfin = len(self.file)
532        if self.verbose:
533            for i in range(self.nfin): inspect(self.file[i])
534        self.var = checktab(self.var,mess="var")
535        self.nvin = len(self.var)
536        # check goal tabs for files and variables
537        # ... default is to plot everything
538        if self.filegoal is None: self.filegoal = ["main"]*self.nfin
539        if self.vargoal is None:  self.vargoal  = ["main"]*self.nvin
540        self.filegoal = checktab(self.filegoal, mess="filegoal")
541        self.vargoal  = checktab(self.vargoal,  mess="vargoal")
542        if len(self.filegoal) != self.nfin:  print "!! ERROR !! filegoal must be the same size as file." ; exit()
543        if len(self.vargoal)  != self.nvin:  print "!! ERROR !! vargoal must be the same size as var." ; exit()
544        # variables: initial check
545        self.x = checktab(self.x,mess="x",allownone=True,allownumber=True)
546        self.y = checktab(self.y,mess="y",allownone=True,allownumber=True)
547        self.z = checktab(self.z,mess="z",allownone=True,allownumber=True)
548        self.t = checktab(self.t,mess="t",allownone=True,allownumber=True)
549        # for the moment not var- nor file- dependent.
550        # but this could be the case.
551        sx = readslices(self.x) ; sy = readslices(self.y)
552        sz = readslices(self.z) ; st = readslices(self.t)
553        # get methods
554        mx = findmethod(sx) ; my = findmethod(sy)
555        mz = findmethod(sz) ; mt = findmethod(st)
556        # get number of plots to be done
557        if mx in ["fixed","comp"]: self.nplotx = sx.size/2
558        else:                      self.nplotx = 1
559        if my in ["fixed","comp"]: self.nploty = sy.size/2
560        else:                      self.nploty = 1
561        if mz in ["fixed","comp"]: self.nplotz = sz.size/2
562        else:                      self.nplotz = 1
563        if mt in ["fixed","comp"]: self.nplott = st.size/2
564        else:                      self.nplott = 1
565        if self.verbose:  print "**** OK. Plots over x,y,z,t -->",self.nplotx,self.nploty,self.nplotz,self.nplott
566        # create the list of onerequest() objects
567        self.request = [[[[[[ \
568                       onerequest() \
569                       for x in range(self.nplotx)] for y in range(self.nploty)] \
570                       for z in range(self.nplotz)] for t in range(self.nplott)] \
571                       for j in range(self.nvin)]   for i in range(self.nfin)] 
572        # store how many onerequest() objects are in self.request
573        self.nrequest = self.nfin*self.nvin*self.nplotx*self.nploty*self.nplotz*self.nplott
574        # loop on onerequest() objects
575        for i in range(self.nfin):
576         for j in range(self.nvin):
577          for t in range(self.nplott):
578           for z in range(self.nplotz):
579            for y in range(self.nploty):
580             for x in range(self.nplotx):
581              obj = self.request[i][j][t][z][y][x]
582              # fill in names for files and variables
583              obj.verbose = self.verbose
584              obj.file = self.file[i]
585              obj.var = self.var[j]
586              # get methods
587              obj.method_x = mx ; obj.method_y = my
588              obj.method_z = mz ; obj.method_t = mt
589              # indicate the computation method
590              obj.compute = self.compute
591              # open the files (the same file might be opened several times but this is cheap)
592              obj.openfile()
593              ### get x,y,z,t dimensions from file
594              obj.getdim()
595              ### possible time axis change
596              obj.changetime = self.changetime
597              obj.performtimechange()
598              # get strides
599              obj.sx = self.sx ; obj.sy = self.sy
600              obj.sz = self.sz ; obj.st = self.st
601              ### get index
602              obj.getindextime(dalist=st,ind=t)
603              obj.getindexvert(dalist=sz,ind=z)
604              obj.getindexhori(dalistx=sx,dalisty=sy,indx=x,indy=y)
605              # missing value
606              obj.missing = self.missing
607        # change status
608        self.status = "defined"
609        return self
610
611    ##############################################################################################
612    # retrieve method
613    # --> for each element onerequest() in the array, get field .var from .f file
614    # --> see below the onerequest() class:
615    #        - only get what is needed for computing and plotting
616    #        - averages etc... are computed here
617    # --> RESULT: each onerequest() object has now its attribute .field filled
618    # --> if one wants to perform operations on fields, this should be done after retrieve()
619    ##############################################################################################
620    def retrieve(self):
621        self.printstatus()
622        # check if things were done OK before
623        if self.status != "defined": print "!! ERROR !! Please use .define() to define your pp object." ; exit()
624        ## create the list of f() and l() objects
625        ## --> so that the user can easily access values (and labels for easy exploration)
626        ## --> see example easy_get_field
627        self.f = [ [] for iii in range(self.nrequest) ]
628        self.l = [ [] for iii in range(self.nrequest) ]
629        count = 0
630        ## first get fields
631        ## ... only what is needed is extracted from the files
632        ## ... and computations are performed
633        for i in range(self.nfin):
634         for j in range(self.nvin):
635          for t in range(self.nplott):
636           for z in range(self.nplotz):
637            for y in range(self.nploty):
638             for x in range(self.nplotx):
639              obj = self.request[i][j][t][z][y][x]
640              obj.getfield()
641              if self.compute != "nothing":
642                  obj.computations()
643              # save fields in self.f for the user
644              self.f[count] = obj.field
645              # save a legend in self.l for the user
646              self.l[count] = "_"
647              if self.nfin > 1:   self.l[count] = self.l[count] + "f=#"+str(int(i+1))+'_'
648              if self.nvin > 1:   self.l[count] = self.l[count] + "v="+obj.var+'_'
649              if self.nplotx > 1: self.l[count] = self.l[count] + "x="+str(self.x[x])+'_'
650              if self.nploty > 1: self.l[count] = self.l[count] + "y="+str(self.y[y])+'_'
651              if self.nplotz > 1: self.l[count] = self.l[count] + "z="+str(self.z[z])+'_'
652              if self.nplott > 1: self.l[count] = self.l[count] + "t="+str(self.t[t])+'_'
653              # close the file
654              obj.closefile()
655              count = count + 1
656        ## make it simple: self.f is simply the data array if self.nrequest=1
657        if self.nrequest == 1: self.f = self.f[0]
658        # change status
659        self.status = "retrieved"
660        return self
661
662    ##########################################################
663    # get: a shortcut method for the define + retrieve chain #
664    ##########################################################
665    def get(self):
666        self.define()
667        self.retrieve()
668        return self 
669
670    ###########################################################
671    # getf: a shortcut method for the define + retrieve chain #
672    #       ... in which the output is self.f                 #
673    #       ... and the ppclass is kept quiet                 #
674    ###########################################################
675    def getf(self):
676        self.quiet = True
677        self.get()
678        return self.f
679
680    ############################################################
681    # getfl: a shortcut method for the define + retrieve chain #
682    #       ... in which the output is self.f, self.l          #
683    #       ... and the ppclass is kept quiet                 #
684    ############################################################
685    def getfl(self):
686        self.quiet = True
687        self.get()
688        return self.f,self.l
689
690    ########################################
691    # smooth: smooth the field in 1D or 2D #
692    ########################################
693    ## TBD: smooth not OK with masked array in the end of retrieve()
694    def smooth(self,window):
695        if self.verbose: 
696            print "!! WARNING !! Performing a smoothing with a window size",window
697            print "!! WARNING !! To come back to unsmoothed file, use .get() again"
698        for i in range(self.nfin):
699         for j in range(self.nvin):
700          for t in range(self.nplott):
701           for z in range(self.nplotz):
702            for y in range(self.nploty):
703             for x in range(self.nplotx):
704              obj = self.request[i][j][t][z][y][x]
705              if obj.field.ndim == 1:
706                  print "!! ERROR !! 1D smoothing not supported yet because reduces array sizes."
707                  exit()
708                  # TBD TBD TBD
709                  #obj.field = ppcompute.smooth1d(obj.field,window=window)
710              elif obj.field.ndim == 2:
711                  obj.field = ppcompute.smooth2d(obj.field,window=window)
712
713    ############################################################################################## 
714    # defineplot method
715    # --> defineplot first defines arrays of plot objects and set each of them
716    #     ... simple looping except cases where goal is not main (e.g. contour or vector)
717    # --> principle: each onerequest() object gives birth to a subplot
718    # --> defineplot vs. makeplot: defining plot and actually plotting it are clearly separated
719    # --> THE KEY OUPUT OF defineplot IS AN ARRAY self.p OF PLOT OBJECTS
720    # optional arguments
721    # --> extraplot: to indicate a number of plots to be added afterwards (use self.plotin)
722    # --> loadfile: to use self.p from a previously saved file
723    ##############################################################################################
724    def defineplot(self,extraplot=0,loadfile=None):
725        # -----------------------------------------------------
726        # LOAD MODE: load a self.p object. count plots from it.
727        # -----------------------------------------------------
728        if loadfile is not None:
729            try: filehandler = open(loadfile, 'r') ; self.p = pickle.load(filehandler) 
730            except IOError: print "!! ERROR !! Cannot find object file to load." ; exit()
731            self.status = "definedplot" ; self.plotin = None
732            self.nplot = len(self.p) ; self.howmanyplots = self.nplot
733            ## [BUG FIX: apparently info about missing values is not saved correctly]
734            for count in range(self.nplot):
735              pl = self.p[count]
736              masked = np.ma.masked_where(np.abs(pl.f) > self.missing,pl.f)
737              pl.f = masked ; pl.f[pl.f.mask] = np.NaN
738            return #self?
739        # -----------------------------------------------------
740        # REGULAR MODE
741        # -----------------------------------------------------
742        self.printstatus()
743        # check if things were done OK before
744        if self.status in ["init","defined"]: 
745            print "!! ERROR !! Please use .retrieve() to get fields for plots with your pp object." ; exit()
746        # check self.plotin (an existing fig on which to add plots afterwards)
747        if self.plotin.__class__.__name__ == "pp":
748            if self.plotin.fig is None:
749                self.plotin = None # this is an additional security in case
750                                   #   a pp object is given without figure opened yet.
751        elif self.plotin is not None:
752            print "!! ERROR !! plotin argument must be a pp object." ; exit()
753        # initialize the array of subplot objects
754        # either something new or attributes coming from plotin object
755        if self.plotin is None:  self.p = [ ]
756        else:                    self.p = self.plotin.p
757        # create an array of subplot objects
758        # ... in theory the order of looping can be changed without any harm
759        # ... the only important thing is to keep i,j,t,z,y,x resp. for file,var,t,z,y,x
760        count = 0
761        for i in range(self.nfin):
762         if self.filegoal[i] == "main": 
763          for j in range(self.nvin):
764           if self.vargoal[j] == "main":
765            for t in range(self.nplott):
766             for z in range(self.nplotz):
767              for y in range(self.nploty):
768               for x in range(self.nplotx):
769                # look at dimension and append the right plot object
770                obj = self.request[i][j][t][z][y][x]
771                dp = obj.dimplot
772                if dp == 1 or self.forcedimplot == 1:    plobj = ppplot.plot1d()
773                elif dp == 2 or self.forcedimplot == 2:  plobj = ppplot.plot2d()
774                elif dp == 0: print "**** OK. VALUES VALUES VALUES",obj.field
775                else:         print "!! ERROR !! 3D or 4D plots not supported" ; exit()
776                # load abscissa and ordinate in obj
777                obj.definecoord()
778                # we start to define things here before appending
779                # (convenient: could be overridden by the user before makeplot)
780                # ... the if loop is necessary so that we can loop above on the dp=0 case
781                if dp in [1,2]:
782                    # and define what to do in plobj
783                    plobj.invert = obj.invert_axes
784                    plobj.swap = obj.swap_axes
785                    # axis labels
786                    plobj.xlabel = obj.absclab ; plobj.ylabel = obj.ordilab
787                    # superpose or not (this is mostly for saving purpose)
788                    plobj.superpose = self.superpose
789                    # get title, colormaps, labels, etc.. from var
790                    plobj.var = obj.var
791                    plobj.define_from_var()
792                    # generic 1D/2D: load field and coord in plot object
793                    plobj.f = obj.field    # field to be plotted
794                    plobj.x = obj.absc      # abscissa (or longitude)
795                    plobj.y = obj.ordi      # ordinate (or latitude)
796                                               # -- useless in 1D but not used anyway
797                    # specific 1D plot stuff
798                    if dp == 1:
799                        # -- a default legend
800                        plobj.legend = ""
801                        if self.nfin > 1: plobj.legend = plobj.legend + " file #"+str(i+1)
802                        if self.nvin > 1: plobj.legend = plobj.legend + " var "+plobj.var
803                        if self.nplott > 1: plobj.legend = plobj.legend + " t="+str(self.t[t])
804                        if self.nplotz > 1: plobj.legend = plobj.legend + " z="+str(self.z[z])
805                        if self.nploty > 1: plobj.legend = plobj.legend + " y="+str(self.y[y])
806                        if self.nplotx > 1: plobj.legend = plobj.legend + " x="+str(self.x[x])
807                    # specific 2d plot stuff
808                    if dp == 2:
809                        # -- light grey background for missing values
810                        if type(plobj.f).__name__ in 'MaskedArray': plobj.axisbg = '0.75'
811                        # -- define if it is a map or a plot
812                        plobj.mapmode = ( obj.method_x+obj.method_y == "freefree" \
813                                       and "grid points" not in obj.name_x \
814                                       and not self.noproj )
815                    # possible user-defined plot settings shared by all plots
816                    if self.div is not None: plobj.div = self.div
817                    if self.xlabel is not None: plobj.xlabel = self.xlabel
818                    if self.xcoeff is not None: plobj.xcoeff = self.xcoeff
819                    if self.ylabel is not None: plobj.ylabel = self.ylabel
820                    if self.ycoeff is not None: plobj.ycoeff = self.ycoeff
821                    if self.title is not None: plobj.title = self.title
822                    if self.units is not None: plobj.units = self.units
823                    if self.colorbar is not None: plobj.colorbar = self.colorbar
824                    if self.modx is not None: plobj.modx = self.modx
825                    if self.nxticks is not None: plobj.nxticks = self.nxticks
826                    if self.nyticks is not None: plobj.nyticks = self.nyticks
827                    # -- 1D specific
828                    if dp == 1:
829                        if self.linestyle is not None: plobj.linestyle = self.linestyle
830                        if self.marker is not None: plobj.marker = self.marker
831                        if self.color is not None: plobj.color = self.color
832                        if self.legend is not None: plobj.legend = self.legend
833                    # -- 2D specific
834                    elif dp == 2:
835                        if self.proj is not None and not self.noproj: plobj.proj = self.proj
836                        if self.vmin is not None: plobj.vmin = self.vmin
837                        if self.vmax is not None: plobj.vmax = self.vmax
838                        plobj.svx = self.svx
839                        plobj.svy = self.svy
840                    # finally append plot object
841                    self.p.append(plobj)
842                    count = count + 1
843        # self.nplot is number of plot to be defined in this call to defineplot()
844        # (because of self.plotin this might less than length of self.p)
845        self.nplot = count
846        # --- superimposed contours and vectors ---
847        # we have to start another loop because we need forward information
848        # TBD: there is probably a more flexible way to do that
849        count = 0
850        for i in range(self.nfin):
851         for j in range(self.nvin):
852          for t in range(self.nplott):
853           for z in range(self.nplotz):
854            for y in range(self.nploty):
855             for x in range(self.nplotx):
856              goal = self.vargoal[j] + self.filegoal[i]
857              obj = self.request[i][j][t][z][y][x]
858              if "mainmain" in goal and obj.dimplot == 2:
859                  # the plot object we consider in the loop
860                  pl = self.p[count]
861                  # -- see if there is a contour requested...
862                  # (we use try because we might be at the end of the list)
863                  found = 0
864                  try:    condvar = self.vargoal[j+1]
865                  except: condvar = "itisok"
866                  try:    condfile = self.filegoal[i+1]
867                  except: condfile = "itisok"
868                  # ... get contour
869                  ##########################################
870                  # NB: contour is expected to be right after main otherwise it is not displayed
871                  ##########################################
872                  if condvar == "contour":
873                      plobj.c = self.request[i][j+1][t][z][y][x].field ; found += 1
874                  if condfile == "contour":
875                      plobj.c = self.request[i+1][j][t][z][y][x].field ; found += 1
876                  # see if there is a vector requested...
877                  # (we use try because we might be at the end of the list)
878                  try:    condvar = self.vargoal[j+found+1]+self.vargoal[j+found+2]
879                  except: condvar = "itisok"
880                  try:    condfile = self.filegoal[i+found+1]+self.filegoal[i+found+2]
881                  except: condfile = "itisok"
882                  # ... get vector and go directly to the next iteration
883                  # (in some cases we would do this twice but this is cheap)
884                  if "vector" in condvar:
885                      plobj.vx = self.request[i][j+found+1][t][z][y][x].field
886                      plobj.vy = self.request[i][j+found+2][t][z][y][x].field
887                  if "vector" in condfile:
888                      plobj.vx = self.request[i+found+1][j][t][z][y][x].field
889                      plobj.vy = self.request[i+found+2][j][t][z][y][x].field
890                  count = count + 1
891        # COUNT PLOTS. if 0 just exit.
892        # self.howmanyplots is self.nplot + possible extraplots
893        self.howmanyplots = self.nplot + extraplot
894        if self.howmanyplots > 0:
895            if self.verbose: print "**** OK. expect %i plots" % (self.howmanyplots)
896        else:
897            pass # because this means that we only had 0D values !
898        # final status
899        self.status = "definedplot"
900        return self
901
902    ##############################################################################################
903    # makeplot method
904    # --> after defineplot and before makeplot, user-defined plot settings can be easily given
905    #     simply by modifying the attributes of each elements of self.p
906    # --> to change only one plot setting, no need to call defineplot again before makeplot
907    # --> in the end, the array self.p of plot objects is saved for easy and convenient replotting
908    # --> see practical examples in the folder 'examples'
909    ##############################################################################################
910    def makeplot(self):
911      if self.howmanyplots > 0:
912        self.printstatus()
913        # a few initial operations
914        # ------------------------
915        if "definedplot" not in self.status: 
916            print "!! ERROR !! Please use .defineplot() before .makeplot() can be used with your pp object." ; exit()
917        # open a figure and define subplots         
918        # ---------------------------------
919        if self.plotin is None: 
920            # start from scratch
921            self.fig = ppplot.figuref(x=self.xp,y=self.yp)
922            self.subv,self.subh = ppplot.definesubplot(self.howmanyplots,self.fig) 
923            self.n = 0
924            ## adapted space for labels etc
925            ## ... except for ortho because there is no label anyway
926            self.customplot = self.p[0].f.ndim == 2 \
927                        and self.p[0].mapmode == True \
928                        and self.p[0].proj not in ["ortho"]
929            if self.customplot:
930                margin = 0.07
931                self.fig.subplots_adjust(left=margin,right=1-margin,bottom=margin,top=1-margin)
932        else:
933            # start from an existing figure.
934            # extraplot must have been set in the call to the previous figure.
935            self.fig = self.plotin.fig
936            self.subv,self.subh = self.plotin.subv,self.plotin.subh
937            self.n = self.plotin.n
938            self.howmanyplots = self.plotin.howmanyplots
939            self.customplot = self.plotin.customplot
940        # LOOP on all subplots
941        # NB: cannot use 'for pl in self.p' if self.plotin not None
942        # --------------------
943        for count in range(self.nplot):
944            # the plot object we consider in the loop
945            pl = self.p[self.n]
946            # before making the plot, create a subplot. the first one is numbered 1 not 0.
947            # ... if pl.superpose, we use only one and only figure
948            # ... (and we have to be careful with not doing things several times)
949            if pl.superpose:
950                if self.n == 0: 
951                    self.fig.add_subplot(1,1,1,axisbg=pl.axisbg) # define one subplot (still needed for user-defined font sizes)
952                    sav = pl.xlabel,pl.ylabel,pl.xcoeff,pl.ycoeff,pl.title,pl.swaplab # save titles and labels
953                    # possibility to color lines according to a color map
954                    # ... made so that all plots span the whole color map automatically.
955                    if self.colorbar is not None: 
956                        if self.verbose: print "**** OK. We make a rainbow spaghetti plot with color map ",self.colorbar
957                        ppplot.rainbow(cb=self.colorbar,num=self.howmanyplots)
958                else: 
959                    pl.invert = False ; pl.linestyle = None # don't invert again axis
960                    # set saved titles and labels
961                    if self.plotin is None:
962                        pl.xlabel,pl.ylabel,pl.xcoeff,pl.ycoeff,pl.title,pl.swaplab = sav
963                    else:
964                        prev_plot = self.plotin.p[self.n-1]
965                        pl.xlabel = prev_plot.xlabel
966                        pl.ylabel = prev_plot.ylabel
967                        pl.xcoeff = prev_plot.xcoeff
968                        pl.ycoeff = prev_plot.ycoeff
969                        pl.title = prev_plot.title
970                        pl.swaplab = prev_plot.swaplab
971            else:
972                self.fig.add_subplot(self.subv,self.subh,self.n+1,axisbg=pl.axisbg)
973            if self.verbose: print "**** Done subplot %i / %i " %( self.n+1,self.howmanyplots ) 
974            # finally make the plot
975            pl.make()
976            # possibly print results in a text file
977            if self.savtxt:
978                if self.verbose: print "**** Printing results in a text file"
979                name = pl.var + "%04d" % self.n
980                ppplot.writeascii(field=pl.f,absc=pl.x,name=name) 
981            # increment plot count (and propagate this in plotin)
982            self.n = self.n+1
983            if self.plotin is not None: self.plotin.n = self.n
984        # once completed show the plot (cannot show intermediate plotin)
985        # ... added a fix (customplot=True) for the label problem in basemap
986        if not self.quiet: print "**** PPCLASS. Done step: makeplot"
987        if (self.n == self.howmanyplots):
988            ppplot.save(mode=self.out,filename=self.filename,folder=self.folder,custom=self.customplot,includedate=self.includedate,res=self.res)
989            mpl.close()
990        # SAVE A PICKLE FILE WITH THE self.p ARRAY OF OBJECTS
991        if self.verbose: print "**** Saving session in "+self.filename + ".ppobj"
992        savfile = self.folder + "/" + self.filename + ".ppobj"
993        try: 
994            filehandler = open(savfile, 'w')
995            pickle.dump(self.p, filehandler)
996        except IOError: 
997            if self.verbose: print "!! WARNING !! Saved object file not written. Probably do not have permission to write here."
998        return self
999
1000    ###########################################################
1001    # plot: a shortcut method for the defineplot + plot chain #
1002    ###########################################################
1003    def plot(self,extraplot=0):
1004        self.defineplot(extraplot=extraplot)
1005        self.makeplot()
1006        return self
1007
1008    #######################################################
1009    # getplot: a shortcut method for the get + plot chain #
1010    #######################################################
1011    def getplot(self,extraplot=0):
1012        self.get()
1013        self.plot(extraplot=extraplot)
1014        return self
1015
1016    ###################################################################
1017    # getdefineplot: a shortcut method for the get + defineplot chain #
1018    ###################################################################
1019    def getdefineplot(self,extraplot=0):
1020        self.get()
1021        self.defineplot(extraplot=extraplot)
1022        return self
1023
1024    #################################################################
1025    # func: operation on two pp objects being on status 'definedplot'
1026    # this allows for one field being function of another one
1027    # e.g. u.func(v) means u will be displayed as a function of v
1028    # ... no need to do defineplot after u.func(v), makeplot directly
1029    #################################################################
1030    def func(self,other):
1031        # preamble: for this operation to work, defineplot() must have been done
1032        if self.status != "definedplot":
1033            if self.verbose: print "!! WARNING !! performing defineplot on operand"
1034            self.defineplot()
1035        if other.status != "definedplot":
1036            if self.verbose: print "!! WARNING !! performing defineplot on operand"
1037            other.defineplot()
1038        # check total number of plots
1039        if self.howmanyplots != other.howmanyplots:
1040               print "!! ERROR !! The two operands do not have the same number of subplots."
1041               exit()
1042        # and now operation.
1043        count = 0
1044        while count < self.howmanyplots:
1045           sobj = self.p[count] ; oobj = other.p[count]
1046           if sobj.f.ndim !=1 or oobj.f.ndim !=1:
1047               if self.verbose: print "!! WARNING !! Flattening arrays because more than one-dimensional."
1048               sobj.f = np.ravel(sobj.f)
1049               oobj.f = np.ravel(oobj.f)
1050           sobj.x = oobj.f
1051           sobj.xlabel = oobj.ylabel
1052           if sobj.x.size > sobj.f.size:
1053               if self.verbose:
1054                   print "!! WARNING !! Trying to define y=f(x) with x and y not at the same size.",sobj.x.size,sobj.f.size
1055                   print "!! WARNING !! Modifying x to fit y size but please check." 
1056               sobj.x = sobj.x[0:sobj.f.size]
1057           count = count + 1
1058        return self
1059
1060    ###########################################################
1061    # copyopt: get options from e.g. a parser
1062    # ... allow for simple scripting and user-defined settings
1063    # ... must be called between defineplot and makeplot
1064    # REQUIRED: attributes of opt must be the same as in the pp object
1065    ###########################################################
1066    def getopt(self,opt):
1067        # -- if only one, or less than the number of plots --> we take the first one
1068        # -- if as many as number of plots --> OK, each plot has its own setting
1069        # (except a few cases such as trans)
1070        for iii in range(self.howmanyplots):
1071            if opt.void:
1072                self.p[iii].showcb = False
1073            else:
1074                self.p[iii].showcb = True
1075            ###
1076            try: self.p[iii].trans = opt.trans
1077            except: pass
1078            ###
1079            try: self.p[iii].div = opt.div
1080            except: pass
1081            ###
1082            try: self.p[iii].logy = opt.logy
1083            except: pass
1084            ###
1085            try: self.p[iii].colorbar = opt.colorbar[iii]
1086            except: 
1087                try: self.p[iii].colorbar = opt.colorbar[0] ; self.colorbar = opt.colorbar[0]
1088                except: pass
1089            ###
1090            if opt.void:
1091                self.p[iii].title = ""
1092            else:
1093              try: self.p[iii].title = opt.title[iii]
1094              except: 
1095                try: self.p[iii].title = opt.title[0]
1096                except: pass
1097            ###
1098            if opt.void:
1099                self.p[iii].xlabel = ""
1100            else:
1101              try: self.p[iii].xlabel = opt.xlabel[iii]
1102              except: 
1103                try: self.p[iii].xlabel = opt.xlabel[0]
1104                except: pass
1105            ###
1106            if opt.void:
1107                self.p[iii].ylabel = ""
1108            else:
1109              try: self.p[iii].ylabel = opt.ylabel[iii]
1110              except: 
1111                try: self.p[iii].ylabel = opt.ylabel[0]
1112                except: pass
1113            ###
1114            try: self.p[iii].linestyle = opt.linestyle[iii]
1115            except: 
1116                try: self.p[iii].linestyle = opt.linestyle[0]
1117                except: pass
1118            ###
1119            try: self.p[iii].color = opt.color[iii]
1120            except: 
1121                try: self.p[iii].color = opt.color[0]
1122                except: pass
1123            ###
1124            try: self.p[iii].marker = opt.marker[iii]
1125            except: 
1126                try: self.p[iii].marker = opt.marker[0]
1127                except: pass
1128            ###
1129            try: self.p[iii].legend = opt.legend[iii]
1130            except:
1131                try: self.p[iii].legend = opt.legend[0]
1132                except: pass
1133            ###
1134            try: self.p[iii].proj = opt.proj[iii]
1135            except: 
1136                try: self.p[iii].proj = opt.proj[0]
1137                except: pass
1138            ###
1139            try: self.p[iii].back = opt.back[iii]
1140            except: 
1141                try: self.p[iii].back = opt.back[0]
1142                except: pass
1143            ###
1144            try: self.p[iii].area = opt.area[iii]
1145            except: 
1146                try: self.p[iii].area = opt.area[0]
1147                except: pass
1148            ###
1149            try: self.p[iii].blon = opt.blon[iii]
1150            except: 
1151                try: self.p[iii].blon = opt.blon[0]
1152                except: pass
1153            ###
1154            try: self.p[iii].blat = opt.blat[iii]
1155            except: 
1156                try: self.p[iii].blat = opt.blat[0]
1157                except: pass
1158            ###
1159            try: self.p[iii].vmin = opt.vmin[iii]
1160            except: 
1161                try: self.p[iii].vmin = opt.vmin[0]
1162                except: pass
1163            ###
1164            try: self.p[iii].vmax = opt.vmax[iii]
1165            except: 
1166                try: self.p[iii].vmax = opt.vmax[0]
1167                except: pass
1168            ###
1169            try: self.p[iii].xcoeff = opt.xcoeff[iii]
1170            except:
1171                try: self.p[iii].xcoeff = opt.xcoeff[0]
1172                except: pass
1173            ###
1174            try: self.p[iii].ycoeff = opt.ycoeff[iii]
1175            except:
1176                try: self.p[iii].ycoeff = opt.ycoeff[0]
1177                except: pass
1178            ###
1179            try: self.p[iii].units = opt.units[iii]
1180            except:
1181                try: self.p[iii].units = opt.units[0]
1182                except: pass
1183            ###
1184            try: self.p[iii].wscale = opt.wscale[iii]
1185            except:
1186                try: self.p[iii].wscale = opt.wscale[0]
1187                except: pass
1188            ###
1189            try: self.p[iii].xmin = opt.xmin[iii]
1190            except:
1191                try: self.p[iii].xmin = opt.xmin[0]
1192                except: pass
1193            ###
1194            try: self.p[iii].ymin = opt.ymin[iii]
1195            except:
1196                try: self.p[iii].ymin = opt.ymin[0]
1197                except: pass
1198            ###
1199            try: self.p[iii].xmax = opt.xmax[iii]
1200            except:
1201                try: self.p[iii].xmax = opt.xmax[0]
1202                except: pass
1203            ###
1204            try: self.p[iii].ymax = opt.ymax[iii]
1205            except:
1206                try: self.p[iii].ymax = opt.ymax[0]
1207                except: pass
1208            ###
1209            try: self.p[iii].modx = opt.modx[iii]
1210            except:
1211                try: self.p[iii].modx = opt.modx[0]
1212                except: pass
1213
1214
1215##########################################################
1216### THE ONEREQUEST SUBOBJECT TO PP (ON WHICH IT LOOPS) ###
1217##########################################################
1218class onerequest():
1219
1220    # default settings. mostly initialized to diagnose problem, except dimplot, nplot, verbose, swap_axes, invert_axes
1221    # -------------------------------
1222    def __init__(self):
1223        self.file  = '!! file: I am not set, damned !!'
1224        self.f     = None
1225        self.dim   = None
1226        self.var   = '!! var: I am not set, damned !!'
1227        self.index_x = [] ; self.index_y = [] ; self.index_z = [] ; self.index_t = []
1228        self.index_x2d = [] ; self.index_y2d = []
1229        self.method_x = '!! method_x: I am not set, damned !!'
1230        self.method_y = '!! method_y: I am not set, damned !!'
1231        self.method_z = '!! method_z: I am not set, damned !!'
1232        self.method_t = '!! method_t: I am not set, damned !!'
1233        self.field = None
1234        self.name_x = None ; self.name_y = None ; self.name_z = None ; self.name_t = None
1235        self.dim_x = None ; self.dim_y = None ; self.dim_z = None ; self.dim_t = None
1236        self.field_x = None ; self.field_y = None ; self.field_z = None ; self.field_t = None
1237        self.dimplot = 0
1238        self.nplot = 1
1239        self.absc = None ; self.ordi = None ; self.absclab = None ; self.ordilab = None
1240        self.verbose = True
1241        self.swap_axes = False ; self.invert_axes = False
1242        self.compute = None
1243        self.changetime = None
1244        self.sx = 1 ; self.sy = 1 ; self.sz = 1 ; self.st = 1
1245        self.missing = '!! missing value: I am not set, damned !!'
1246
1247    # open a file. for now it is netcdf. TBD for other formats.
1248    # check that self.var is inside.
1249    # -------------------------------
1250    def openfile(self):
1251        if not os.path.exists(self.file): print '!! ERROR !! I could not find the following file: '+self.file ; exit()
1252        if not os.path.isfile(self.file): print '!! ERROR !! This does not appear to be a file: '+self.file ; exit()
1253        self.f = netCDF4.Dataset(self.file)
1254        if self.verbose: print "**** OK. Opened file "+self.file
1255        if self.var not in self.f.variables.keys(): 
1256            print '!! ERROR !! File '+self.file+' does not contain variable: '+self.var
1257            print '..... try instead with ',self.f.variables.keys() ; exit()
1258
1259    # close a file
1260    # ------------
1261    def closefile(self):
1262        self.f.close()
1263
1264    # copy attributes from another existing object
1265    # --------------------------------------------
1266    def copy(self,source):
1267        for k, v in vars(source).items():
1268            setattr(self,k,v)
1269
1270    # get x,y,z,t dimensions from NETCDF file
1271    # TBD: user could request for a specific altitude dimension
1272    # TBD: staggered variables could request specific dimensions
1273    # -------------------------------
1274    def getdim(self):
1275          # GET SIZES OF EACH DIMENSION
1276          if self.verbose: print "**** OK. Found variable "+self.var
1277          shape = self.f.variables[self.var].shape
1278          self.dim = len(shape)
1279          if self.dim == 1:
1280              if self.verbose: print "**** OK. 1D field. I assume this varies with time."
1281              self.dim_x = 1 ; self.dim_y = 1 ; self.dim_z = 1 ; self.dim_t = shape[0]
1282          elif self.dim == 2:
1283              if self.verbose: print "**** OK. 2D field. I assume this is not-time-varying lat-lon map."
1284              self.dim_x = shape[1] ; self.dim_y = shape[0] ; self.dim_z = 1 ; self.dim_t = 1
1285          elif self.dim == 3:
1286              if self.verbose: print "**** OK. 3D field. I assume this is time-varying lat-lon map."
1287              self.dim_x = shape[2] ; self.dim_y = shape[1] ; self.dim_z = 1 ; self.dim_t = shape[0]
1288          elif self.dim == 4:
1289              if self.verbose: print "**** OK. 4D field."
1290              self.dim_x = shape[3] ; self.dim_y = shape[2] ; self.dim_z = shape[1] ; self.dim_t = shape[0]
1291          # LONGITUDE. Try preset fields. If not present set grid points axis.
1292          self.name_x = "nothing"
1293          for c in glob_listx:
1294            if c in self.f.variables.keys():
1295             self.name_x = c
1296          if self.name_x == "nothing":
1297            self.field_x = np.array(range(self.dim_x))
1298            self.name_x = "x grid points"
1299          else:
1300            self.field_x = self.f.variables[self.name_x]
1301          # LATITUDE. Try preset fields. If not present set grid points axis.
1302          self.name_y = "nothing"
1303          for c in glob_listy:
1304            if c in self.f.variables.keys():
1305             self.name_y = c
1306          if self.name_y == "nothing":
1307            self.field_y = np.array(range(self.dim_y))
1308            self.name_y = "y grid points"
1309          else:
1310            self.field_y = self.f.variables[self.name_y]
1311          # ensure that lon and lat are 2D fields
1312          # 1. simple 1D case (not time-varying)
1313          if len(self.field_x.shape)*len(self.field_y.shape) == 1:
1314               if self.verbose: print "**** OK. recasting lon and lat as 2D fields." 
1315               [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y)
1316          # 2. complex 3D case (time-varying, actually just copied over time axis)
1317          elif len(self.field_x.shape)*len(self.field_y.shape) == 9:
1318               if self.verbose: print "**** OK. reducing lon and lat as 2D fields. get rid of time."
1319               self.field_x = self.field_x[0,:,:]
1320               self.field_y = self.field_y[0,:,:]
1321          # if xy axis are apparently undefined, set 2D grid points axis.
1322          if "grid points" not in self.name_x:
1323            if self.field_x.all() == self.field_x[0,0] \
1324             or self.field_x.min() == self.field_x.max() \
1325             or self.field_y.min() == self.field_y.max():
1326               if self.verbose: print "!! WARNING !! xy axis look undefined. creating non-dummy ones."
1327               self.field_x = np.array(range(self.dim_x)) ; self.name_x = "x grid points"
1328               self.field_y = np.array(range(self.dim_y)) ; self.name_y = "y grid points"
1329               [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y)
1330          if self.dim_x > 1: 
1331               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())
1332          if self.dim_y > 1: 
1333               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())
1334          # ALTITUDE. Try preset fields. If not present set grid points axis.
1335          # WARNING: how do we do if several are available? the last one is chosen.
1336          self.name_z = "nothing"
1337          for c in glob_listz:
1338            if c in self.f.variables.keys():
1339             self.name_z = c
1340          if self.name_z == "nothing":
1341            self.field_z = np.array(range(self.dim_z))
1342            self.name_z = "z grid points"
1343          else:
1344            tabalt = self.f.variables[self.name_z]
1345            # (consider the case where tabtime is not dim 1) TBD: 2D and 3D cases
1346            if tabalt.ndim == 4: 
1347                try:
1348                    self.field_z = tabalt[1,:,0,0] # 4D case. alt is usually third dimension.
1349                                                   # 1 for time to avoid initial 0s
1350                except:
1351                    self.field_z = tabalt[0,:,0,0]
1352                if self.verbose: print "!! WARNING !! "+self.name_z+" is 4D var. We made it 1D."
1353            else: 
1354                self.field_z = self.f.variables[self.name_z][:] # specify dimension
1355            # TBD: problems when self.dim_z != self.field_z.size
1356            if self.field_z.size != self.dim_z:
1357                if self.verbose: print "!! WARNING !! Cannot use this z coordinate. Not enough points. Use simple z axis."
1358                self.field_z = np.array(range(self.dim_z))
1359                self.name_z = "z grid points"
1360          if self.dim_z > 1: 
1361               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())
1362
1363          # TIME. Try preset fields.
1364          self.name_t = "nothing"
1365          for c in glob_listt:
1366            if c in self.f.variables.keys():
1367             self.name_t = c
1368             if self.verbose: print "**** OK. Found time variable: ",c
1369          try:
1370            # speed up: only get first value, last one.
1371            tabtime = self.f.variables[self.name_t]
1372            # (consider the case where tabtime is not dim 1)
1373            # (time is most often the first dimension)
1374            if tabtime.ndim == 2: tabtime = tabtime[:,0]
1375            elif tabtime.ndim == 3: tabtime = tabtime[:,0,0]
1376            elif tabtime.ndim == 4: tabtime = tabtime[:,0,0,0]
1377            # (now proceed) (the +0. just ensures this is a number)
1378            dafirst = tabtime[0] + 0.
1379            if self.dim_t == 1:
1380                self.field_t = np.array([dafirst])
1381            else:
1382                daint = tabtime[1] - dafirst
1383                dalast = dafirst + (self.dim_t-1)*daint
1384                self.field_t = np.linspace(dafirst,dalast,num=self.dim_t)
1385                if self.verbose:
1386                    print "!! WARNING !! WARNING !! Time axis is supposed to be equally spaced !!"
1387                    if dalast != tabtime[self.dim_t-1]:
1388                        print "!! WARNING !! Time axis has been recast to be monotonic",dalast,tabtime[self.dim_t-1]
1389          except:
1390            # ... or if a problem encountered, define a simple time axis
1391            if self.verbose: print "**** OK. There is something weird. Let us go for a simple time axis."
1392            self.field_t = np.array(range(self.dim_t))
1393            self.name_t = "t grid points"
1394          if self.dim_t > 1: 
1395               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())     
1396
1397    # change time axis
1398    # ... add your options here!
1399    # --------------------------
1400    def performtimechange(self):
1401        if self.changetime is not None:
1402            if self.verbose: print "**** OK. Converting time axis:",self.changetime
1403            ### options added by T. Navarro
1404            if self.changetime == "mars_sol2ls":
1405                if "controle" in self.f.variables: 
1406                   self.field_t =  self.field_t \
1407                                 + self.f.variables['controle'][3]%669 \
1408                                 + self.f.variables['controle'][26]
1409                self.field_t = ppcompute.mars_sol2ls(self.field_t)
1410            elif self.changetime == "mars_dayini" and "controle" in self.f.variables:
1411                self.field_t =  self.field_t \
1412                              + self.f.variables['controle'][3]%669 \
1413                              + self.f.variables['controle'][26]
1414            ### options added by A. Spiga
1415            elif self.changetime == "correctls":
1416                # not regularly spaced + handle modulo 360. in files
1417                dafirst = tabtime[0] + 0.
1418                daint = tabtime[1] - dafirst
1419                dalast = dafirst + (self.dim_t-1)*daint
1420                year = 0.
1421                add = np.linspace(dafirst,dalast,num=self.dim_t) ; add[0] = 0.
1422                for iii in range(1,self.dim_t):
1423                  if tabtime[iii] - tabtime[iii-1] < 0: year = year+1.
1424                  add[iii] = year*360.
1425                self.field_t = add + tabtime
1426            elif "mars_meso" in self.changetime:
1427                if 'Times' not in self.f.variables.keys():
1428                    if self.verbose: print "!! WARNING !! Variable Times not in file. Cannot proceed to change of time axis."
1429                else:
1430                    # get the array of strings describing dates
1431                    dates = self.f.variables['Times']
1432                    dates.set_auto_maskandscale(False) # necessary to solve the api Times bug!
1433                    # get ls sol utc from those strings
1434                    ls, sol, utc = ppcompute.mars_date(dates[:])
1435                    # populate self.field_t with the right output from mars_date
1436                    if self.changetime == "mars_meso_ls": 
1437                        self.field_t = ls
1438                        self.name_t = "Ls"
1439                    elif self.changetime == "mars_meso_sol": 
1440                        self.field_t = sol
1441                        self.name_t = "sol"
1442                    elif self.changetime == "mars_meso_utc" \
1443                        and ( self.changetime == "mars_meso_lt" \
1444                              and not hasattr(self.f,'CEN_LON') ): 
1445                        self.field_t = ppcompute.timecorrect(utc)
1446                        self.name_t = "utc"
1447                        if self.method_t == "fixed": 
1448                            self.field_t = self.field_t % 24 # so that the user is not mistaken!
1449                    elif self.changetime == "mars_meso_lt":
1450                        self.field_t = ppcompute.timecorrect(utc) + getattr(self.f,'CEN_LON') / 15.
1451                        self.field_t = ppcompute.timecorrect(self.field_t)
1452                        self.name_t = "local time (center of domain)"
1453                        if self.method_t == "fixed": 
1454                            self.field_t = self.field_t % 24 # so that the user is not mistaken!
1455            else:
1456                if self.verbose: print "!! WARNING !! This time change is not implemented. Nothing is done."
1457            if self.verbose: print "**** OK. new t axis values [%5.1f,%5.1f]" % (self.field_t.min(),self.field_t.max())
1458
1459    # get list of index to be retrieved for time axis
1460    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1461    # -------------------------------
1462    def getindextime(self,dalist=None,ind=None):
1463        if self.method_t == "free": 
1464            self.index_t = np.arange(0,self.dim_t,self.st)
1465            if self.dim_t > 1: 
1466                self.dimplot = self.dimplot + 1 
1467                if self.verbose: print "**** OK. t values. all."
1468            else:               
1469                self.method_t = "fixed"
1470                if self.verbose: print "**** OK. no t dimension."
1471        elif self.method_t == "comp":
1472            start = np.argmin( np.abs( self.field_t - dalist[ind][0] ) )
1473            stop = np.argmin( np.abs( self.field_t - dalist[ind][1] ) )
1474            self.index_t = np.arange(start,stop,self.st)
1475            if self.verbose: print "**** OK. t values. comp over interval ",self.field_t[start],self.field_t[stop]," nvalues=",self.index_t.size
1476        elif self.method_t == "fixed":
1477            self.index_t.append( np.argmin( np.abs( self.field_t - dalist[ind][0] ) ))
1478            if self.verbose: print "**** OK. t values",self.field_t[self.index_t]
1479        else:
1480            print "!! ERROR !! method "+self.method_t+" not supported"
1481        self.index_t = np.array(self.index_t)
1482             
1483    # get list of index to be retrieved for vertical axis
1484    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1485    # -------------------------------
1486    def getindexvert(self,dalist=None,ind=None):
1487        if self.method_z == "free": 
1488            self.index_z = np.arange(0,self.dim_z,self.sz)
1489            if self.dim_z > 1: 
1490                self.dimplot = self.dimplot + 1
1491                if self.verbose: print "**** OK. z values. all."
1492            else:               
1493                self.method_z = "fixed"
1494                if self.verbose: print "**** OK. no z dimension."
1495        elif self.method_z == "comp":
1496            start = np.argmin( np.abs( self.field_z - dalist[ind][0] ) )
1497            stop = np.argmin( np.abs( self.field_z - dalist[ind][1] ) )
1498            self.index_z = np.arange(start,stop,self.sz)
1499            if self.verbose: print "**** OK. z values. comp over interval",self.field_z[start],self.field_z[stop]," nvalues=",self.index_z.size
1500        elif self.method_z == "fixed":
1501            self.index_z.append( np.argmin( np.abs( self.field_z - dalist[ind][0] ) ))
1502            if self.verbose: print "**** OK. z values",self.field_z[self.index_z]
1503        else:
1504            if self.verbose: print "!! ERROR !! method "+self.method_z+" not supported"
1505        self.index_z = np.array(self.index_z)
1506
1507    # get list of index to be retrieved for horizontal grid
1508    # --> index_x and index_y are slices to be retrieved from NETCDF files
1509    # --> index_x2d and index_y2d are the actual (x,y) coordinates corresponding to each relevant point
1510    # [this is slightly more complicated because 2D arrays for lat-lon projection possibly irregular]
1511    # NB: to append index we use lists (the most convenient) then we convert into a numpy.array
1512    ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all)
1513    # -------------------------------
1514    def getindexhori(self,dalistx=None,dalisty=None,indx=None,indy=None):
1515        ## get what is the method over x and y axis
1516        test = self.method_x+self.method_y
1517        ## CASE 0, EASY CASES:
1518        ## - LAT IS FREE (we do here what must be done whatever LON case is)
1519        ## - LON IS FREE (we do here what must be done whatever LAT case is)
1520        ## - LAT IS COMP AND LON IS FREE
1521        ## - LON IS COMP AND LAT IS FREE
1522        if self.method_x == "free" or test in ["compfree","compcomp"]:
1523            self.index_x = range(0,self.dim_x,self.sx)
1524            if self.dim_x > 1: 
1525                if self.method_x == "free": self.dimplot = self.dimplot + 1
1526                if self.verbose: print "**** OK. x values. all."
1527            else:               
1528                self.method_x = "fixed"
1529                if self.verbose: print "**** OK. no x dimension."
1530        if self.method_y == "free" or test in ["freecomp","compcomp"]:
1531            self.index_y = range(0,self.dim_y,self.sy)
1532            if self.dim_y > 1: 
1533                if self.method_y == "free": self.dimplot = self.dimplot + 1
1534                if self.verbose: print "**** OK. y values. all."
1535            else:               
1536                self.method_y = "fixed"
1537                if self.verbose: print "**** OK. no y dimension."
1538        ## CASE 0 above, this is just for continuity for free.
1539        ## ... for comp we have to select bounds.
1540        ## ... TBD: take the bool array strategy for what follows!
1541        if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]:
1542            ### ref1_dirty_hack
1543            ### ... for the moment this is a hack. but actually this is more powerful.
1544            if self.method_x == "comp":
1545                yeah = (self.field_x >= dalistx[indx][0])*(self.field_x <= dalistx[indx][1])
1546                self.index_x = yeah[0,:]
1547            if self.method_y == "comp":
1548                yeah = (self.field_y >= dalisty[indy][0]) * (self.field_y <= dalisty[indy][1])
1549                self.index_y = yeah[:,0]
1550            self.index_x2d = self.index_x
1551            self.index_y2d = self.index_y
1552        ## AND NOW THE LITTLE BIT MORE COMPLICATED CASES
1553        ## CASE 1 LAT AND LON ARE FIXED
1554        elif test == "fixedfixed":
1555            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 ) 
1556                          #TBD: pb with staggered coord
1557            if idx not in self.index_x:  self.index_x.append(idx)
1558            if idy not in self.index_y:  self.index_y.append(idy)
1559            self.index_x2d.append(idx)
1560            self.index_y2d.append(idy)
1561        ## CASE 2 LON IS FIXED BUT NOT LAT
1562        elif test in ["fixedfree","fixedcomp"]:
1563            # find where are requested x values for each y on the free dimension
1564            # NB: this does not work for non-bijective cases e.g. polar stereographic
1565            for iy in range(self.dim_y):
1566              idx = np.argmin( np.abs( self.field_x[iy,:] - dalistx[indx][0] ) )
1567              # if comp is requested we select only indexes which yield values between requested min and max
1568              storeval = (self.method_y == "comp") and (self.field_y[iy,idx] > dalisty[indy][0]) and (self.field_y[iy,idx] < dalisty[indy][1])
1569              storeval = storeval or (self.method_y == "free")
1570              if storeval:
1571                  if idx not in self.index_x:  self.index_x.append(idx)
1572                  if iy not in self.index_y and self.method_y == "comp": self.index_y.append(iy)
1573                  if idx not in self.index_x2d or iy not in self.index_y2d:
1574                    self.index_x2d.append(idx)
1575                    self.index_y2d.append(iy)
1576        ## CASE 3 LAT IS FIXED BUT NOT LON
1577        elif test in ["freefixed","compfixed"]:
1578            # find where are requested y values for each x on the free dimension
1579            # NB: this does not work for non-bijective cases e.g. polar stereographic
1580            for ix in range(self.dim_x):
1581              idy = np.argmin( np.abs( self.field_y[:,ix] - dalisty[indy][0] ) )
1582              # if comp is requested we select only indexes which yield values between requested min and max
1583              storeval = (self.method_x == "comp") and (self.field_x[idy,ix] > dalistx[indx][0]) and (self.field_x[idy,ix] < dalistx[indx][1])
1584              storeval = storeval or (self.method_x == "free")
1585              if storeval:
1586                  if idy not in self.index_y:  self.index_y.append(idy)
1587                  if ix not in self.index_x and self.method_x == "comp": self.index_x.append(ix)
1588                  if ix not in self.index_x2d or idy not in self.index_y2d:
1589                    self.index_x2d.append(ix)
1590                    self.index_y2d.append(idy)
1591        ## check index tab
1592        if len(self.index_x) == 0 or len(self.index_y) == 0:
1593            print "!! ERROR !! no indices found. check prescribed latitudes or longitudes" ; exit()
1594        ## ensure the array is a numpy array for getfield to work
1595        self.index_x = np.array(self.index_x)
1596        self.index_y = np.array(self.index_y)
1597        self.index_x2d = np.array(self.index_x2d)
1598        self.index_y2d = np.array(self.index_y2d)
1599        ### print extrema
1600        printx = self.field_x[np.ix_(self.index_y2d, self.index_x2d)]
1601        printy = self.field_y[np.ix_(self.index_y2d, self.index_x2d)]
1602        if self.verbose: 
1603            print "**** OK. x values (min,max).", printx.min(),printx.max()
1604            print "**** OK. y values (min,max).", printy.min(),printy.max()
1605
1606    # get the field from the NETCDF file and perform averages
1607    # -------------------------------
1608    def getfield(self):
1609        ## first tell what is to be done
1610        if self.verbose:
1611          if self.dimplot > 2:                       print "**** !! WARNING !! "+str(self.dimplot)+"D plots will not be supported!"
1612          elif self.dimplot == 0 and self.verbose:   print "**** OK. 0D value requested."
1613          elif self.dimplot == 1 and self.verbose:   print "**** OK. 1D plot requested."
1614          elif self.verbose:                         print "**** OK. 2D section requested."
1615        # well, now get field from netcdf file
1616        # part below is necessary otherwise there is an index error below
1617        if self.index_x.size == 1: self.index_x = self.index_x[0]
1618        if self.index_y.size == 1: self.index_y = self.index_y[0]
1619        if self.index_z.size == 1: self.index_z = self.index_z[0]
1620        if self.index_t.size == 1: self.index_t = self.index_t[0]
1621        # then retrieve what is requested by user
1622        # each self.dim case corresponds to tests in the beginning of getdim.
1623        time0 = timelib.time()
1624        if self.verbose: print "**** OK. I am getting values from files. Please wait."
1625        if self.dim == 1: 
1626            nt = self.index_t.size ; nz = 1 ; ny = 1 ; nx = 1
1627            self.field = self.f.variables[self.var][self.index_t]
1628        elif self.dim == 2:
1629            nt = 1 ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size
1630            self.field = self.f.variables[self.var][self.index_y,self.index_x]
1631        elif self.dim == 3:
1632            nt = self.index_t.size ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size
1633            self.field = self.f.variables[self.var][self.index_t,self.index_y,self.index_x]
1634            # this is far faster than retrieving each term with a loop
1635        elif self.dim == 4:
1636            nt = self.index_t.size ; nz = self.index_z.size ; ny = self.index_y.size ; nx = self.index_x.size
1637            self.field = self.f.variables[self.var][self.index_t,self.index_z,self.index_y,self.index_x]
1638        else:
1639            print "!! ERROR !! field would have more than four dimensions ?" ; exit()
1640        # dirty hack (AS) ref1_dirty_hack
1641        # waiting for more fundamental modifications. case when self.index_y is a bool array.
1642        # ... be careful if no point...
1643        try:
1644            if type(self.index_x[0]) == np.bool_: nx = np.sum(self.index_x) ## gives the size of the True part!
1645            if type(self.index_y[0]) == np.bool_: ny = np.sum(self.index_y) ## gives the size of the True part!
1646        except:
1647            pass
1648        # NB: ... always 4D array but possibly with "size 1" dimensions
1649        #     ... if one dimension is missing because 1D 2D or 3D requests, make it appear again
1650        self.field = np.reshape(self.field,(nt,nz,ny,nx))
1651        if self.verbose: print "**** OK. I got %7.1e values. This took me %6.4f seconds" % (nx*ny*nz*nt,timelib.time() - time0)
1652        if self.verbose: print "**** OK. I got var "+self.var+" with shape",self.field.shape
1653        # reduce coordinates to useful points
1654        # ... TBD: this should be ordered in the case of non-regular projections
1655        if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]:
1656          # we need 2D coordinates (free) or we get broadcast problem (comp) so we use np.ix
1657          self.field_x = self.field_x[np.ix_(self.index_y2d, self.index_x2d)]
1658          self.field_y = self.field_y[np.ix_(self.index_y2d, self.index_x2d)]
1659        else:
1660          # we are OK with 1D coordinates
1661          self.field_x = self.field_x[self.index_y2d, self.index_x2d]
1662          self.field_y = self.field_y[self.index_y2d, self.index_x2d]
1663          # ... there are special cases with strides
1664          # ... some other fixes will be needed probably TBD
1665          if self.method_x == "free" and self.sx != 1:
1666              self.field_x = self.field_x[self.index_x]
1667          if self.method_y == "free" and self.sy != 1: 
1668              self.field_y = self.field_y[self.index_y]
1669        self.field_z = self.field_z[self.index_z]
1670        self.field_t = self.field_t[self.index_t]
1671        # extract relevant horizontal points
1672        # TBD: is compfree OK with computing on irregular grid?
1673        test = self.method_x + self.method_y
1674        if test in ["fixedfixed","freefree"]:
1675          pass
1676        elif test in ["fixedfree","fixedcomp"] or test in ["freefixed","compfixed"]: 
1677         if self.sx == 1 and self.sy == 1:
1678          time0 = timelib.time()
1679          # now have to obtain the new indexes which correspond to the extracted self.field
1680          # for it to work with unique index, ensure that any index_* is a numpy array
1681          if not isinstance(self.index_x, np.ndarray): self.index_x = np.array([self.index_x])
1682          if not isinstance(self.index_y, np.ndarray): self.index_y = np.array([self.index_y])
1683          if not isinstance(self.index_z, np.ndarray): self.index_z = np.array([self.index_z])
1684          if not isinstance(self.index_t, np.ndarray): self.index_t = np.array([self.index_t])
1685          for val in self.index_x: self.index_x2d[np.where(self.index_x2d == val)] = np.where(self.index_x == val)[0]
1686          for val in self.index_y: self.index_y2d[np.where(self.index_y2d == val)] = np.where(self.index_y == val)[0]
1687               ##### VERY EXPENSIVE
1688               ## recast self.field with 2D horizontal arrays because we might have extracted
1689               ## more than what is to be actually plot or computed, in particular for comps on 2D lat,lon coordinates
1690               #self.field = self.field[np.ix_(self.index_t,self.index_z,self.index_y2d,self.index_x2d)]
1691               #(nt,nz,ny,nx) = self.field.shape       
1692          # prepare the loop on all relevant horizontal points
1693          if self.method_x in ["comp","free"]:   
1694              nnn = self.index_x2d.shape[0] ; what_I_am_supposed_to_do = "keepx"
1695          elif self.method_y in ["comp","free"]: 
1696              nnn = self.index_y2d.shape[0] ; what_I_am_supposed_to_do = "keepy" 
1697          # LOOP to extract only useful values over horizontal dimensions
1698          # only take diagonal terms, do not loop on all self.index_x2d*self.index_y2d
1699          # ... this method is fast enough, perhaps there is a faster way though
1700          # ... (for sure the method with np.diag is much slower)
1701          for iii in range(nnn):
1702           ix = self.index_x2d[iii] ; iy = self.index_y2d[iii]
1703           for iz in range(self.index_z.size):
1704            for it in range(self.index_t.size):
1705              if what_I_am_supposed_to_do == "keepx":    self.field[it,iz,0,ix] = self.field[it,iz,iy,ix]
1706              elif what_I_am_supposed_to_do == "keepy":  self.field[it,iz,iy,0] = self.field[it,iz,iy,ix]
1707          if self.verbose: print "**** OK. I got to pick the right values for your request. This took me %6.4f seconds" % (timelib.time() - time0)
1708          # we only keep the one value that was modified on the dimension which is not free
1709          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))
1710          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))
1711         else:
1712          # there is a problem above if stride != 1. a fix must be found. rewrite might be necessary. TBD
1713          pass
1714        # make a mask in case there are non-NaN missing values. (what about NaN missing values?)
1715        # ... this is important for computations below (see ppcompute)
1716        masked = np.ma.masked_where(np.abs(self.field) > self.missing,self.field)
1717        if masked.mask.any() == True:
1718             if self.verbose: print "!! WARNING !! Values over %5.3e are considered missing values." % self.missing
1719             self.field = masked
1720             self.field.set_fill_value([np.NaN])
1721
1722    # perform computations
1723    # -------------------------------
1724    # available: mean, max, min, meanarea
1725    # TB: integrals? for derivatives, define a function self.dx()
1726    def computations(self): 
1727        nt,nz,ny,nx = self.field.shape
1728        # treat the case of mean on fields normalized with grid mesh area
1729        # ... this is done in the .area() method.
1730        # after that self.field contains field*area/totarea
1731        if "area" in self.compute: 
1732            if "comp" in self.method_x+self.method_y: 
1733                self.area()
1734            else:
1735                if self.verbose: print "!! WARNING !! No area accounted for (computing on t and/or z axis)."
1736        # prepare quadratic mean
1737        if "qmean" in self.compute: self.field = self.field*self.field
1738        # now ready to compute [TBD: we would like to have e.g. mean over x,y and min over t ??]
1739        if self.method_t == "comp":
1740            if self.verbose: print "**** OK. Computing over t axis."
1741            if "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=0)
1742            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=0)
1743            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=0)
1744            else: print "!! ERROR !! operation not supported." ; exit()
1745            nt = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1746        if self.method_z == "comp": 
1747            if self.verbose: print "**** OK. Computing over z axis."
1748            if "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=1)
1749            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=1)
1750            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=1)
1751            else: print "!! ERROR !! operation not supported." ; exit()
1752            nz = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1753        if self.method_y == "comp": 
1754            if self.verbose: print "**** OK. Computing over y axis."
1755            if self.compute == "meanarea": self.field = ppcompute.sum(self.field,axis=2)
1756            elif "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=2)
1757            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=2)
1758            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=2)
1759            else: print "!! ERROR !! operation not supported." ; exit()
1760            ny = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1761            if self.field_x.ndim == 2: self.field_x = self.field_x[0,:] # TBD: this is OK for regular grid but not for irregular
1762        if self.method_x == "comp":
1763            if self.verbose: print "**** OK. Computing over x axis."
1764            if self.compute == "meanarea": self.field = ppcompute.sum(self.field,axis=3)
1765            elif "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=3)
1766            elif self.compute == "min": self.field = ppcompute.min(self.field,axis=3)
1767            elif self.compute == "max": self.field = ppcompute.max(self.field,axis=3)
1768            else: print "!! ERROR !! operation not supported." ; exit()
1769            nx = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx))
1770            if self.field_y.ndim == 2: self.field_y = self.field_y[:,0] # TBD: this is OK for regular grid but not for irregular
1771        # remove all dimensions with size 1 to prepare plot (and check the resulting dimension with dimplot)
1772        self.field = np.squeeze(self.field)
1773        # take root mean square for quadratic mean
1774        if self.compute == "qmean": self.field = np.sqrt(self.field)
1775        # error handling and verbose
1776        if self.field.ndim != self.dimplot: 
1777            print "!! ERROR !! Problem: self.field is different than plot dimensions", self.field.ndim, self.dimplot ; exit()
1778        if self.verbose: 
1779            print "**** OK. Final shape for "+self.var+" after averaging and squeezing",self.field.shape
1780   
1781    # get areas for computations and ponderate self.field by area/totarea
1782    # -------------------------------------------------------------------
1783    def area(self):
1784        if self.verbose: print "**** OK. Get area array for computations."
1785        # create a request object for area
1786        # ... and copy known attributes from self
1787        aire = onerequest()
1788        aire.copy(self)
1789        # get area field name
1790        aire.var = "nothing"
1791        for c in glob_listarea:
1792         if c in aire.f.variables.keys():
1793            aire.var = c
1794        # do not try to calculate areas automatically
1795        if aire.var == "nothing":
1796            print "!! ERROR !! area variable not found... needs to be added in set_ppclass.txt?"
1797            exit()
1798        # define area request dimensions
1799        aire.getdim()
1800        # ensure this is a 2D horizontal request and define indexes
1801        #    ... areas are not supposed to vary with time and height
1802        aire.method_x = "free" ; aire.method_y = "free"
1803        aire.getindexhori() ; aire.dimplot = 2
1804        aire.method_z = "fixed" ; aire.field_z = np.array([0]) ; aire.index_z = np.array([0])
1805        aire.method_t = "fixed" ; aire.field_t = np.array([0]) ; aire.index_t = np.array([0])
1806        # read the 2D area array in netCDF file
1807        aire.getfield()
1808        aire.field = np.squeeze(aire.field)
1809        # reduce with self horizontal indexes
1810        if "fixed" in self.method_x+self.method_y:
1811            aire.field = aire.field[self.index_y,self.index_x]
1812        # calculate total area
1813        # ... 2D comp is easy. 1D comp is a bit less easy but simple array manipulation.
1814        if "free" in self.method_x+self.method_y:
1815            if self.method_x == "free":
1816                totarea = ppcompute.sum(aire.field,axis=0)
1817                totarea = np.reshape(totarea,(1,totarea.size))
1818                totarea = np.tile(totarea,(1,self.index_x))
1819            elif self.method_y == "free":
1820                totarea = ppcompute.sum(aire.field,axis=1)
1821                totarea = np.reshape(totarea,(totarea.size,1))
1822                totarea = np.tile(totarea,(1,self.index_x.size))
1823        elif self.method_x == "comp" and self.method_y == "comp":
1824            aire.field = aire.field[np.ix_(self.index_y, self.index_x)] # reduce to requested indexes only
1825            totarea = ppcompute.sum(ppcompute.sum(aire.field,axis=1),axis=0)
1826        else:
1827            if self.verbose: print "!! WARNING !! Not account for areas. Only averaging over z and/or t axis."
1828        # normalize by total area
1829        print "**** OK. I can now normalize field by areas."
1830        aire.field = aire.field / totarea
1831        # tile area array over self t and z axis so that area field could be multiplied with self.field
1832        aire.field = np.tile(aire.field,(self.index_t.size,self.index_z.size,1,1))
1833        if self.field.shape != aire.field.shape:
1834            print "!! ERROR !! Problem in area(). Check array shapes."
1835            print "Field vs. aire:",self.field.shape,aire.field.shape ; exit()
1836        else:
1837            self.field = self.field*aire.field
1838
1839    # define coordinates for plot
1840    # -------------------------------
1841    def definecoord(self):
1842        I_got_abs = False ; I_got_ord = False
1843        # here is the thing. time is usually taken as an abscissa so we start with time.
1844        if self.method_t ==  "free": 
1845            self.absc = self.field_t ; self.absclab = self.name_t
1846            I_got_abs = True
1847        # then we usually have x as an abscissa.
1848        if self.method_x == "free":
1849            if I_got_abs: 
1850                self.ordi = self.field_x ; self.ordilab = self.name_x
1851                I_got_ord = True
1852            else:         
1853                self.absc = self.field_x ; self.absclab = self.name_x
1854                I_got_abs = True
1855        # ... or we have y
1856        if self.method_y == "free":
1857            if I_got_abs:   
1858                self.ordi = self.field_y ; self.ordilab = self.name_y
1859                I_got_ord = True
1860            else:         
1861                self.absc = self.field_y ; self.absclab = self.name_y
1862                I_got_abs = True
1863        # ... and we end with z because it is usually not an abscissa (profiles).
1864        if self.method_z == "free":
1865            if self.field_z[0] > self.field_z[1]:
1866                self.invert_axes = True # the axis will be turned upside-down
1867            if I_got_abs: 
1868                self.ordi = self.field_z ; self.ordilab = self.name_z
1869                I_got_ord = True
1870            else:
1871                self.absc = self.field_z ; self.absclab = self.name_z
1872                I_got_abs = True
1873                self.swap_axes = True # says that altitude is not supposed to remain as an abscissa
1874        if I_got_abs and self.verbose: print "**** OK. abscissa:",self.absclab, self.absc.shape
1875        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.