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

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

UTIL PYTHON planetoplot_v2. units: set an additional column set_var.txt and put units over colorbar. 1D bounding and formatting. possible mode nothing. corrected bug coast/sea.

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