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

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

UTIL PYTHON planetoplot_v2. added support for wind vectors in 2D section (not only maps, thanks Benjamin for the request). fixed stride problems (not all though). now to stride vectors simply use the stride command, this is more flexible. added stride to pp.py. also corrected a bug for logx axis (thanks Alizee)

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