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

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

UTIL PYTHON planetoplot_v2. oops previous commit. removed a dirty print.

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