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 |
---|
10 | import os |
---|
11 | import time as timelib |
---|
12 | import pickle |
---|
13 | # added librairies |
---|
14 | import numpy as np |
---|
15 | import netCDF4 |
---|
16 | import matplotlib.pyplot as mpl |
---|
17 | # personal librairies |
---|
18 | import ppplot |
---|
19 | import ppcompute |
---|
20 | ############################################### |
---|
21 | |
---|
22 | ################################### |
---|
23 | #### HEADER ## |
---|
24 | #### ... executed when imported ## |
---|
25 | ################################### |
---|
26 | # where settings files are located... |
---|
27 | whereset = None |
---|
28 | whereset = ppcompute.findset(whereset) |
---|
29 | # ... we load user-defined automatic settings from set_ppclass.txt |
---|
30 | zefile = "set_ppclass.txt" |
---|
31 | glob_listx = [] ; glob_listy = [] ; glob_listz = [] ; glob_listt = [] |
---|
32 | glob_listarea = [] |
---|
33 | try: |
---|
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) |
---|
39 | for stuff in lines[17].strip().split(';'): glob_listarea.append(stuff) |
---|
40 | except 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 |
---|
48 | def 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). |
---|
60 | def 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 |
---|
82 | def 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) |
---|
89 | def 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 |
---|
102 | |
---|
103 | # look for comas in the input name to separate different names (files, variables,etc ..) |
---|
104 | # (needed by readslices) |
---|
105 | def 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 | ####################### |
---|
120 | class pp(): |
---|
121 | |
---|
122 | # print out a help string when help is invoked on the object |
---|
123 | def __repr__(self): |
---|
124 | whatprint = 'pp object. \"help(pp)\" for more information\n' |
---|
125 | return whatprint |
---|
126 | |
---|
127 | # default settings |
---|
128 | # -- user can define settings by two methods. |
---|
129 | # -- 1. yeah = pp(file="file.nc") |
---|
130 | # -- 2. yeah = pp() ; yeah.file = "file.nc" |
---|
131 | def __init__(self,file=None,var="notset",\ |
---|
132 | filegoal=None,vargoal=None,\ |
---|
133 | x=None,y=None,z=None,t=None,\ |
---|
134 | stridex=1,stridey=1,\ |
---|
135 | stridez=1,stridet=1,\ |
---|
136 | compute="mean",\ |
---|
137 | verbose=False,noproj=False,\ |
---|
138 | superpose=False,\ |
---|
139 | plotin=None,\ |
---|
140 | forcedimplot=-1,\ |
---|
141 | out="gui",\ |
---|
142 | filename="myplot",\ |
---|
143 | folder="./",\ |
---|
144 | includedate=True,\ |
---|
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): |
---|
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 |
---|
165 | self.customplot = False |
---|
166 | self.allfield = None |
---|
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 |
---|
185 | self.includedate = includedate |
---|
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 |
---|
199 | |
---|
200 | # print status |
---|
201 | def printstatus(self): |
---|
202 | if self.filename == "THIS_IS_A_CLONE": |
---|
203 | pass |
---|
204 | else: |
---|
205 | print "**** PPCLASS. Done step: " + self.status |
---|
206 | |
---|
207 | ##################################################### |
---|
208 | # EMULATE OPERATORS + - * / ** << FOR PP() OBJECTS # |
---|
209 | ##################################################### |
---|
210 | |
---|
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 |
---|
244 | self.includedate = other.includedate |
---|
245 | else: |
---|
246 | print "!! ERROR !! argument must be a pp object." ; exit() |
---|
247 | |
---|
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 | |
---|
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 |
---|
286 | the_clone.filename = "THIS_IS_A_CLONE" # trick to avoid additional outputs |
---|
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): |
---|
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(): |
---|
297 | if k != "field": |
---|
298 | setattr(obj,k,v) |
---|
299 | the_clone.status = "retrieved" |
---|
300 | the_clone.filename = self.filename |
---|
301 | return the_clone |
---|
302 | |
---|
303 | # define the operation + on two objects. or with an int/float. |
---|
304 | # ... with selective_copy the self object is not modified. |
---|
305 | def __add__(self,other): |
---|
306 | isnum = self.checktwo(other) |
---|
307 | the_clone = self.selective_copy() |
---|
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): |
---|
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 |
---|
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 |
---|
327 | else: |
---|
328 | obj.field = obj_ref.field + ope |
---|
329 | return the_clone |
---|
330 | |
---|
331 | # define the operation - on two objects. or with an int/float. |
---|
332 | # ... with selective_copy the self object is not modified. |
---|
333 | def __sub__(self,other): |
---|
334 | isnum = self.checktwo(other) |
---|
335 | the_clone = self.selective_copy() |
---|
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): |
---|
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 |
---|
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 |
---|
355 | else: |
---|
356 | obj.field = obj_ref.field - ope |
---|
357 | return the_clone |
---|
358 | |
---|
359 | # define the operation * on two objects. or with an int/float. |
---|
360 | # ... with selective_copy the self object is not modified. |
---|
361 | def __mul__(self,other): |
---|
362 | isnum = self.checktwo(other) |
---|
363 | the_clone = self.selective_copy() |
---|
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): |
---|
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 |
---|
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 |
---|
383 | else: |
---|
384 | obj.field = obj_ref.field * ope |
---|
385 | return the_clone |
---|
386 | |
---|
387 | # define the operation / on two objects. or with an int/float. |
---|
388 | # ... with selective_copy the self object is not modified. |
---|
389 | def __div__(self,other): |
---|
390 | isnum = self.checktwo(other) |
---|
391 | the_clone = self.selective_copy() |
---|
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): |
---|
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 |
---|
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 |
---|
411 | else: |
---|
412 | obj.field = obj_ref.field / ope |
---|
413 | return the_clone |
---|
414 | |
---|
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 | |
---|
439 | # define the operation ** on one object. |
---|
440 | # ... with selective_copy the self object is not modified. |
---|
441 | def __pow__(self,num): |
---|
442 | the_clone = self.selective_copy() |
---|
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): |
---|
450 | obj = the_clone.request[i][j][t][z][y][x] |
---|
451 | obj_ref = self.request[i][j][t][z][y][x] |
---|
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 |
---|
456 | else: |
---|
457 | obj.field = obj_ref.field ** num |
---|
458 | else: |
---|
459 | print "!! ERROR !! To define a power, either an int or a float is needed." ; exit() |
---|
460 | return the_clone |
---|
461 | |
---|
462 | ### TBD: reverse power? for exponentials? |
---|
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 |
---|
510 | if mx == "fixed": self.nplotx = sx.size/2 |
---|
511 | else: self.nplotx = 1 |
---|
512 | if my == "fixed": self.nploty = sy.size/2 |
---|
513 | else: self.nploty = 1 |
---|
514 | if mz == "fixed": self.nplotz = sz.size/2 |
---|
515 | else: self.nplotz = 1 |
---|
516 | if mt == "fixed": self.nplott = st.size/2 |
---|
517 | else: self.nplott = 1 |
---|
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 |
---|
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) |
---|
550 | # change status |
---|
551 | self.status = "defined" |
---|
552 | return self |
---|
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 |
---|
566 | if self.status != "defined": print "!! ERROR !! Please use .define() to define your pp object." ; exit() |
---|
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)] |
---|
574 | ## first get fields |
---|
575 | ## ... only what is needed is extracted from the files |
---|
576 | ## ... and computations are performed |
---|
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() |
---|
585 | obj.computations() |
---|
586 | self.allfield[i][j][t][z][y][x] = obj.field |
---|
587 | # get rid of useless dimensions in allfield |
---|
588 | self.allfield = np.squeeze(self.allfield) |
---|
589 | # change status |
---|
590 | self.status = "retrieved" |
---|
591 | return self |
---|
592 | |
---|
593 | ########################################################## |
---|
594 | # get: a shortcut method for the define + retrieve chain # |
---|
595 | ########################################################## |
---|
596 | def get(self): |
---|
597 | self.define() |
---|
598 | self.retrieve() |
---|
599 | return self |
---|
600 | |
---|
601 | ######################################## |
---|
602 | # smooth: smooth the field in 1D or 2D # |
---|
603 | ######################################## |
---|
604 | ## TBD: smooth not OK with masked array in the end of retrieve() |
---|
605 | def smooth(self,window): |
---|
606 | if self.verbose: |
---|
607 | print "!! WARNING !! Performing a smoothing with a window size",window |
---|
608 | print "!! WARNING !! To come back to unsmoothed file, use .get() again" |
---|
609 | for i in range(self.nfin): |
---|
610 | for j in range(self.nvin): |
---|
611 | for t in range(self.nplott): |
---|
612 | for z in range(self.nplotz): |
---|
613 | for y in range(self.nploty): |
---|
614 | for x in range(self.nplotx): |
---|
615 | obj = self.request[i][j][t][z][y][x] |
---|
616 | if obj.field.ndim == 1: |
---|
617 | print "!! ERROR !! 1D smoothing not supported yet because reduces array sizes." |
---|
618 | exit() |
---|
619 | # TBD TBD TBD |
---|
620 | #obj.field = ppcompute.smooth1d(obj.field,window=window) |
---|
621 | elif obj.field.ndim == 2: |
---|
622 | obj.field = ppcompute.smooth2d(obj.field,window=window) |
---|
623 | |
---|
624 | ############################################################################################## |
---|
625 | # defineplot method |
---|
626 | # --> defineplot first defines arrays of plot objects and set each of them |
---|
627 | # ... simple looping except cases where goal is not main (e.g. contour or vector) |
---|
628 | # --> principle: each onerequest() object gives birth to a subplot |
---|
629 | # --> defineplot vs. makeplot: defining plot and actually plotting it are clearly separated |
---|
630 | # --> THE KEY OUPUT OF defineplot IS AN ARRAY self.p OF PLOT OBJECTS |
---|
631 | # optional arguments |
---|
632 | # --> extraplot: to indicate a number of plots to be added afterwards (use self.plotin) |
---|
633 | # --> loadfile: to use self.p from a previously saved file |
---|
634 | ############################################################################################## |
---|
635 | def defineplot(self,extraplot=0,loadfile=None): |
---|
636 | # ----------------------------------------------------- |
---|
637 | # LOAD MODE: load a self.p object. count plots from it. |
---|
638 | # ----------------------------------------------------- |
---|
639 | if loadfile is not None: |
---|
640 | try: filehandler = open(loadfile, 'r') ; self.p = pickle.load(filehandler) |
---|
641 | except IOError: print "!! ERROR !! Cannot find object file to load." ; exit() |
---|
642 | self.status = "definedplot" ; self.plotin = None |
---|
643 | self.nplot = len(self.p) ; self.howmanyplots = self.nplot |
---|
644 | return |
---|
645 | # ----------------------------------------------------- |
---|
646 | # REGULAR MODE |
---|
647 | # ----------------------------------------------------- |
---|
648 | self.printstatus() |
---|
649 | # check if things were done OK before |
---|
650 | if self.status in ["init","defined"]: |
---|
651 | print "!! ERROR !! Please use .retrieve() to get fields for plots with your pp object." ; exit() |
---|
652 | # check self.plotin (an existing fig on which to add plots afterwards) |
---|
653 | if self.plotin.__class__.__name__ == "pp": |
---|
654 | if self.plotin.fig is None: |
---|
655 | self.plotin = None # this is an additional security in case |
---|
656 | # a pp object is given without figure opened yet. |
---|
657 | elif self.plotin is not None: |
---|
658 | print "!! ERROR !! plotin argument must be a pp object." ; exit() |
---|
659 | # initialize the array of subplot objects |
---|
660 | # either something new or attributes coming from plotin object |
---|
661 | if self.plotin is None: self.p = [ ] |
---|
662 | else: self.p = self.plotin.p |
---|
663 | # create an array of subplot objects |
---|
664 | # ... in theory the order of looping can be changed without any harm |
---|
665 | # ... the only important thing is to keep i,j,t,z,y,x resp. for file,var,t,z,y,x |
---|
666 | count = 0 |
---|
667 | for i in range(self.nfin): |
---|
668 | if self.filegoal[i] == "main": |
---|
669 | for j in range(self.nvin): |
---|
670 | if self.vargoal[j] == "main": |
---|
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 | # look at dimension and append the right plot object |
---|
676 | obj = self.request[i][j][t][z][y][x] |
---|
677 | dp = obj.dimplot |
---|
678 | if dp == 1 or self.forcedimplot == 1: plobj = ppplot.plot1d() |
---|
679 | elif dp == 2 or self.forcedimplot == 2: plobj = ppplot.plot2d() |
---|
680 | elif dp == 0: print "**** OK. VALUES VALUES VALUES",obj.field |
---|
681 | else: print "!! ERROR !! 3D or 4D plots not supported" ; exit() |
---|
682 | # load abscissa and ordinate in obj |
---|
683 | obj.definecoord() |
---|
684 | # we start to define things here before appending |
---|
685 | # (convenient: could be overridden by the user before makeplot) |
---|
686 | # ... the if loop is necessary so that we can loop above on the dp=0 case |
---|
687 | if dp in [1,2]: |
---|
688 | # and define what to do in plobj |
---|
689 | plobj.invert = obj.invert_axes |
---|
690 | plobj.swap = obj.swap_axes |
---|
691 | # axis labels |
---|
692 | plobj.xlabel = obj.absclab ; plobj.ylabel = obj.ordilab |
---|
693 | # superpose or not (this is mostly for saving purpose) |
---|
694 | plobj.superpose = self.superpose |
---|
695 | # get title, colormaps, labels, etc.. from var |
---|
696 | plobj.var = obj.var |
---|
697 | plobj.define_from_var() |
---|
698 | # generic 1D/2D: load field and coord in plot object |
---|
699 | plobj.field = obj.field # field to be plotted |
---|
700 | plobj.absc = obj.absc # abscissa (or longitude) |
---|
701 | plobj.ordi = obj.ordi # ordinate (or latitude) |
---|
702 | # -- useless in 1D but not used anyway |
---|
703 | # specific 1D plot stuff |
---|
704 | if dp == 1: |
---|
705 | # -- a default label |
---|
706 | plobj.label = "" |
---|
707 | if self.nfin > 1: plobj.label = plobj.label + " file #"+str(i+1) |
---|
708 | if self.nvin > 1: plobj.label = plobj.label + " var #"+str(j+1) |
---|
709 | if self.nplott > 1: plobj.label = plobj.label + " t #"+str(t+1) |
---|
710 | if self.nplotz > 1: plobj.label = plobj.label + " z #"+str(z+1) |
---|
711 | if self.nploty > 1: plobj.label = plobj.label + " y #"+str(y+1) |
---|
712 | if self.nplotx > 1: plobj.label = plobj.label + " x #"+str(x+1) |
---|
713 | # specific 2d plot stuff |
---|
714 | if dp == 2: |
---|
715 | # -- light grey background for missing values |
---|
716 | if type(plobj.field).__name__ in 'MaskedArray': plobj.axisbg = '0.75' |
---|
717 | # -- define if it is a map or a plot |
---|
718 | plobj.mapmode = ( obj.method_x+obj.method_y == "freefree" \ |
---|
719 | and "grid points" not in obj.name_x \ |
---|
720 | and not self.noproj ) |
---|
721 | # possible user-defined plot settings shared by all plots |
---|
722 | if self.div is not None: plobj.div = self.div |
---|
723 | if self.xlabel is not None: plobj.xlabel = self.xlabel |
---|
724 | if self.xcoeff is not None: plobj.xcoeff = self.xcoeff |
---|
725 | if self.ylabel is not None: plobj.ylabel = self.ylabel |
---|
726 | if self.ycoeff is not None: plobj.ycoeff = self.ycoeff |
---|
727 | if self.title is not None: plobj.title = self.title |
---|
728 | # -- 1D specific |
---|
729 | if dp == 1: |
---|
730 | if self.lstyle is not None: plobj.lstyle = self.lstyle |
---|
731 | if self.marker is not None: plobj.marker = self.marker |
---|
732 | if self.color is not None: plobj.color = self.color |
---|
733 | if self.label is not None: plobj.label = self.label |
---|
734 | # -- 2D specific |
---|
735 | elif dp == 2: |
---|
736 | if self.proj is not None and not self.noproj: plobj.proj = self.proj |
---|
737 | if self.vmin is not None: plobj.vmin = self.vmin |
---|
738 | if self.vmax is not None: plobj.vmax = self.vmax |
---|
739 | if self.colorb is not None: plobj.colorb = self.colorb |
---|
740 | # finally append plot object |
---|
741 | self.p.append(plobj) |
---|
742 | count = count + 1 |
---|
743 | # self.nplot is number of plot to be defined in this call to defineplot() |
---|
744 | # (because of self.plotin this might less than length of self.p) |
---|
745 | self.nplot = count |
---|
746 | # --- superimposed contours and vectors --- |
---|
747 | # we have to start another loop because we need forward information |
---|
748 | # TBD: there is probably a more flexible way to do that |
---|
749 | count = 0 |
---|
750 | for i in range(self.nfin): |
---|
751 | for j in range(self.nvin): |
---|
752 | for t in range(self.nplott): |
---|
753 | for z in range(self.nplotz): |
---|
754 | for y in range(self.nploty): |
---|
755 | for x in range(self.nplotx): |
---|
756 | goal = self.vargoal[j] + self.filegoal[i] |
---|
757 | obj = self.request[i][j][t][z][y][x] |
---|
758 | if "mainmain" in goal and obj.dimplot == 2: |
---|
759 | # the plot object we consider in the loop |
---|
760 | pl = self.p[count] |
---|
761 | # -- see if there is a contour requested... |
---|
762 | # (we use try because we might be at the end of the list) |
---|
763 | found = 0 |
---|
764 | try: condvar = self.vargoal[j+1] |
---|
765 | except: condvar = "itisok" |
---|
766 | try: condfile = self.filegoal[i+1] |
---|
767 | except: condfile = "itisok" |
---|
768 | # ... get contour |
---|
769 | ########################################## |
---|
770 | # NB: contour is expected to be right after main otherwise it is not displayed |
---|
771 | ########################################## |
---|
772 | if condvar == "contour": |
---|
773 | plobj.addcontour = self.request[i][j+1][t][z][y][x].field ; found += 1 |
---|
774 | if condfile == "contour": |
---|
775 | plobj.addcontour = self.request[i+1][j][t][z][y][x].field ; found += 1 |
---|
776 | # see if there is a vector requested... |
---|
777 | # (we use try because we might be at the end of the list) |
---|
778 | try: condvar = self.vargoal[j+found+1]+self.vargoal[j+found+2] |
---|
779 | except: condvar = "itisok" |
---|
780 | try: condfile = self.filegoal[i+found+1]+self.filegoal[i+found+2] |
---|
781 | except: condfile = "itisok" |
---|
782 | # ... get vector and go directly to the next iteration |
---|
783 | # (in some cases we would do this twice but this is cheap) |
---|
784 | if "vector" in condvar: |
---|
785 | plobj.addvecx = self.request[i][j+found+1][t][z][y][x].field |
---|
786 | plobj.addvecy = self.request[i][j+found+2][t][z][y][x].field |
---|
787 | if "vector" in condfile: |
---|
788 | plobj.addvecx = self.request[i+found+1][j][t][z][y][x].field |
---|
789 | plobj.addvecy = self.request[i+found+2][j][t][z][y][x].field |
---|
790 | count = count + 1 |
---|
791 | # COUNT PLOTS. if 0 just exit. |
---|
792 | # self.howmanyplots is self.nplot + possible extraplots |
---|
793 | self.howmanyplots = self.nplot + extraplot |
---|
794 | if self.howmanyplots > 0: |
---|
795 | if self.verbose: print "**** OK. expect %i plots" % (self.howmanyplots) |
---|
796 | else: |
---|
797 | exit() # because this means that we only had 0D values ! |
---|
798 | # final status |
---|
799 | self.status = "definedplot" |
---|
800 | return self |
---|
801 | |
---|
802 | ############################################################################################## |
---|
803 | # makeplot method |
---|
804 | # --> after defineplot and before makeplot, user-defined plot settings can be easily given |
---|
805 | # simply by modifying the attributes of each elements of self.p |
---|
806 | # --> to change only one plot setting, no need to call defineplot again before makeplot |
---|
807 | # --> in the end, the array self.p of plot objects is saved for easy and convenient replotting |
---|
808 | # --> see practical examples in the folder 'examples' |
---|
809 | ############################################################################################## |
---|
810 | def makeplot(self): |
---|
811 | self.printstatus() |
---|
812 | # a few initial operations |
---|
813 | # ------------------------ |
---|
814 | if "definedplot" not in self.status: |
---|
815 | print "!! ERROR !! Please use .defineplot() before .makeplot() can be used with your pp object." ; exit() |
---|
816 | # open a figure and define subplots |
---|
817 | # --------------------------------- |
---|
818 | if self.plotin is None: |
---|
819 | # start from scratch |
---|
820 | self.fig = mpl.figure(figsize=(16,8)) |
---|
821 | self.subv,self.subh = ppplot.definesubplot(self.howmanyplots,self.fig) |
---|
822 | self.n = 0 |
---|
823 | ## adapted space for labels etc |
---|
824 | ## ... except for ortho because there is no label anyway |
---|
825 | self.customplot = self.p[0].field.ndim == 2 \ |
---|
826 | and self.p[0].mapmode == True \ |
---|
827 | and self.p[0].proj not in ["ortho"] |
---|
828 | if self.customplot: |
---|
829 | margin = 0.07 |
---|
830 | self.fig.subplots_adjust(left=margin,right=1-margin,bottom=margin,top=1-margin) |
---|
831 | else: |
---|
832 | # start from an existing figure. |
---|
833 | # extraplot must have been set in the call to the previous figure. |
---|
834 | self.fig = self.plotin.fig |
---|
835 | self.subv,self.subh = self.plotin.subv,self.plotin.subh |
---|
836 | self.n = self.plotin.n |
---|
837 | self.howmanyplots = self.plotin.howmanyplots |
---|
838 | self.customplot = self.plotin.customplot |
---|
839 | # LOOP on all subplots |
---|
840 | # NB: cannot use 'for pl in self.p' if self.plotin not None |
---|
841 | # -------------------- |
---|
842 | for count in range(self.nplot): |
---|
843 | # the plot object we consider in the loop |
---|
844 | pl = self.p[self.n] |
---|
845 | # before making the plot, create a subplot. the first one is numbered 1 not 0. |
---|
846 | # ... if pl.superpose, we use only one and only figure |
---|
847 | # ... (and we have to be careful with not doing things several times) |
---|
848 | if pl.superpose: |
---|
849 | if self.n == 0: |
---|
850 | self.fig.add_subplot(1,1,1,axisbg=pl.axisbg) # define one subplot (still needed for user-defined font sizes) |
---|
851 | sav = pl.xlabel,pl.ylabel,pl.xcoeff,pl.ycoeff,pl.title,pl.swaplab # save titles and labels |
---|
852 | else: |
---|
853 | pl.invert = False ; pl.lstyle = None # don't invert again axis |
---|
854 | # set saved titles and labels |
---|
855 | if self.plotin is None: |
---|
856 | pl.xlabel,pl.ylabel,pl.xcoeff,pl.ycoeff,pl.title,pl.swaplab = sav |
---|
857 | else: |
---|
858 | prev_plot = self.plotin.p[self.n-1] |
---|
859 | pl.xlabel = prev_plot.xlabel |
---|
860 | pl.ylabel = prev_plot.ylabel |
---|
861 | pl.xcoeff = prev_plot.xcoeff |
---|
862 | pl.ycoeff = prev_plot.ycoeff |
---|
863 | pl.title = prev_plot.title |
---|
864 | pl.swaplab = prev_plot.swaplab |
---|
865 | else: |
---|
866 | self.fig.add_subplot(self.subv,self.subh,self.n+1,axisbg=pl.axisbg) |
---|
867 | if self.verbose: print "**** Done subplot %i / %i " %( self.n+1,self.howmanyplots ) |
---|
868 | # finally make the plot |
---|
869 | pl.make() |
---|
870 | # increment plot count (and propagate this in plotin) |
---|
871 | self.n = self.n+1 |
---|
872 | if self.plotin is not None: self.plotin.n = self.n |
---|
873 | # once completed show the plot (cannot show intermediate plotin) |
---|
874 | # ... added a fix (customplot=True) for the label problem in basemap |
---|
875 | print "**** PPCLASS. Done step: makeplot" |
---|
876 | if (self.n == self.howmanyplots): |
---|
877 | ppplot.save(mode=self.out,filename=self.filename,folder=self.folder,custom=self.customplot,includedate=self.includedate) |
---|
878 | mpl.close() |
---|
879 | # SAVE A PICKLE FILE WITH THE self.p ARRAY OF OBJECTS |
---|
880 | if self.verbose: print "**** Saving session in "+self.filename + ".ppobj" |
---|
881 | savfile = self.folder + "/" + self.filename + ".ppobj" |
---|
882 | try: |
---|
883 | filehandler = open(savfile, 'w') |
---|
884 | pickle.dump(self.p, filehandler) |
---|
885 | except IOError: |
---|
886 | print "!! WARNING !! Saved object file not written. Probably do not have permission to write here." |
---|
887 | return self |
---|
888 | |
---|
889 | ########################################################### |
---|
890 | # plot: a shortcut method for the defineplot + plot chain # |
---|
891 | ########################################################### |
---|
892 | def plot(self,extraplot=0): |
---|
893 | self.defineplot(extraplot=extraplot) |
---|
894 | self.makeplot() |
---|
895 | return self |
---|
896 | |
---|
897 | ####################################################### |
---|
898 | # getplot: a shortcut method for the get + plot chain # |
---|
899 | ####################################################### |
---|
900 | def getplot(self,extraplot=0): |
---|
901 | self.get() |
---|
902 | self.plot(extraplot=extraplot) |
---|
903 | return self |
---|
904 | |
---|
905 | ################################################################### |
---|
906 | # getdefineplot: a shortcut method for the get + defineplot chain # |
---|
907 | ################################################################### |
---|
908 | def getdefineplot(self,extraplot=0): |
---|
909 | self.get() |
---|
910 | self.defineplot(extraplot=extraplot) |
---|
911 | return self |
---|
912 | |
---|
913 | ############################################################## |
---|
914 | # f: operation on two pp objects being on status 'definedplot' |
---|
915 | # this allows for one field being function of another one |
---|
916 | # e.g. u.f(v) means u will be displayed as a function of v |
---|
917 | # ... no need to do defineplot after u.f(v), makeplot directly |
---|
918 | ############################################################## |
---|
919 | def f(self,other): |
---|
920 | # preamble: for this operation to work, defineplot() must have been done |
---|
921 | if self.status != "definedplot": |
---|
922 | if self.verbose: print "!! WARNING !! performing defineplot on operand" |
---|
923 | self.defineplot() |
---|
924 | if other.status != "definedplot": |
---|
925 | if self.verbose: print "!! WARNING !! performing defineplot on operand" |
---|
926 | other.defineplot() |
---|
927 | # check total number of plots |
---|
928 | if self.howmanyplots != other.howmanyplots: |
---|
929 | print "!! ERROR !! The two operands do not have the same number of subplots." |
---|
930 | exit() |
---|
931 | # and now operation. |
---|
932 | count = 0 |
---|
933 | while count < self.howmanyplots: |
---|
934 | sobj = self.p[count] ; oobj = other.p[count] |
---|
935 | if sobj.field.ndim !=1 or oobj.field.ndim !=1: |
---|
936 | if self.verbose: print "!! WARNING !! Flattening arrays because more than one-dimensional." |
---|
937 | sobj.field = np.ravel(sobj.field) |
---|
938 | oobj.field = np.ravel(oobj.field) |
---|
939 | sobj.absc = oobj.field |
---|
940 | sobj.xlabel = oobj.ylabel |
---|
941 | if sobj.absc.size > sobj.field.size: |
---|
942 | if self.verbose: |
---|
943 | print "!! WARNING !! Trying to define y=f(x) with x and y not at the same size.",sobj.absc.size,sobj.field.size |
---|
944 | print "!! WARNING !! Modifying x to fit y size but please check." |
---|
945 | sobj.absc = sobj.absc[0:sobj.field.size] |
---|
946 | count = count + 1 |
---|
947 | return self |
---|
948 | |
---|
949 | ########################################################### |
---|
950 | # copyopt: get options from e.g. a parser |
---|
951 | # ... allow for simple scripting and user-defined settings |
---|
952 | # ... must be called between defineplot and makeplot |
---|
953 | # REQUIRED: attributes of opt must be the same as in the pp object |
---|
954 | ########################################################### |
---|
955 | def getopt(self,opt): |
---|
956 | # -- if only one, or less than the number of plots --> we take the first one |
---|
957 | # -- if as many as number of plots --> OK, each plot has its own setting |
---|
958 | # (except a few cases such as trans) |
---|
959 | for iii in range(self.howmanyplots): |
---|
960 | ### |
---|
961 | try: self.p[iii].trans = opt.trans |
---|
962 | except: pass |
---|
963 | ### |
---|
964 | try: self.p[iii].div = opt.div |
---|
965 | except: pass |
---|
966 | ### |
---|
967 | try: self.p[iii].logy = opt.logy |
---|
968 | except: pass |
---|
969 | ### |
---|
970 | try: self.p[iii].colorb = opt.colorb[iii] |
---|
971 | except: |
---|
972 | try: self.p[iii].colorb = opt.colorb[0] |
---|
973 | except: pass |
---|
974 | ### |
---|
975 | try: self.p[iii].title = opt.title[iii] |
---|
976 | except: |
---|
977 | try: self.p[iii].title = opt.title[0] |
---|
978 | except: pass |
---|
979 | ### |
---|
980 | try: self.p[iii].xlabel = opt.xlabel[iii] |
---|
981 | except: |
---|
982 | try: self.p[iii].xlabel = opt.xlabel[0] |
---|
983 | except: pass |
---|
984 | ### |
---|
985 | try: self.p[iii].ylabel = opt.ylabel[iii] |
---|
986 | except: |
---|
987 | try: self.p[iii].ylabel = opt.ylabel[0] |
---|
988 | except: pass |
---|
989 | ### |
---|
990 | try: self.p[iii].lstyle = opt.lstyle[iii] |
---|
991 | except: |
---|
992 | try: self.p[iii].lstyle = opt.lstyle[0] |
---|
993 | except: pass |
---|
994 | ### |
---|
995 | try: self.p[iii].color = opt.color[iii] |
---|
996 | except: |
---|
997 | try: self.p[iii].color = opt.color[0] |
---|
998 | except: pass |
---|
999 | ### |
---|
1000 | try: self.p[iii].marker = opt.marker[iii] |
---|
1001 | except: |
---|
1002 | try: self.p[iii].marker = opt.marker[0] |
---|
1003 | except: pass |
---|
1004 | ### |
---|
1005 | try: self.p[iii].label = opt.label[iii] |
---|
1006 | except: |
---|
1007 | try: self.p[iii].label = opt.label[0] |
---|
1008 | except: pass |
---|
1009 | ### |
---|
1010 | try: self.p[iii].proj = opt.proj[iii] |
---|
1011 | except: |
---|
1012 | try: self.p[iii].proj = opt.proj[0] |
---|
1013 | except: pass |
---|
1014 | ### |
---|
1015 | try: self.p[iii].back = opt.back[iii] |
---|
1016 | except: |
---|
1017 | try: self.p[iii].back = opt.back[0] |
---|
1018 | except: pass |
---|
1019 | ### |
---|
1020 | try: self.p[iii].area = opt.area[iii] |
---|
1021 | except: |
---|
1022 | try: self.p[iii].area = opt.area[0] |
---|
1023 | except: pass |
---|
1024 | ### |
---|
1025 | try: self.p[iii].blon = opt.blon[iii] |
---|
1026 | except: |
---|
1027 | try: self.p[iii].blon = opt.blon[0] |
---|
1028 | except: pass |
---|
1029 | ### |
---|
1030 | try: self.p[iii].blat = opt.blat[iii] |
---|
1031 | except: |
---|
1032 | try: self.p[iii].blat = opt.blat[0] |
---|
1033 | except: pass |
---|
1034 | ### |
---|
1035 | try: self.p[iii].vmin = opt.vmin[iii] |
---|
1036 | except: |
---|
1037 | try: self.p[iii].vmin = opt.vmin[0] |
---|
1038 | except: pass |
---|
1039 | ### |
---|
1040 | try: self.p[iii].vmax = opt.vmax[iii] |
---|
1041 | except: |
---|
1042 | try: self.p[iii].vmax = opt.vmax[0] |
---|
1043 | except: pass |
---|
1044 | |
---|
1045 | ########################################################## |
---|
1046 | ### THE ONEREQUEST SUBOBJECT TO PP (ON WHICH IT LOOPS) ### |
---|
1047 | ########################################################## |
---|
1048 | class onerequest(): |
---|
1049 | |
---|
1050 | # default settings. mostly initialized to diagnose problem, except dimplot, nplot, verbose, swap_axes, invert_axes |
---|
1051 | # ------------------------------- |
---|
1052 | def __init__(self): |
---|
1053 | self.file = '!! file: I am not set, damned !!' |
---|
1054 | self.f = None |
---|
1055 | self.dim = None |
---|
1056 | self.var = '!! var: I am not set, damned !!' |
---|
1057 | self.index_x = [] ; self.index_y = [] ; self.index_z = [] ; self.index_t = [] |
---|
1058 | self.index_x2d = [] ; self.index_y2d = [] |
---|
1059 | self.method_x = '!! method_x: I am not set, damned !!' |
---|
1060 | self.method_y = '!! method_y: I am not set, damned !!' |
---|
1061 | self.method_z = '!! method_z: I am not set, damned !!' |
---|
1062 | self.method_t = '!! method_t: I am not set, damned !!' |
---|
1063 | self.field = None |
---|
1064 | self.name_x = None ; self.name_y = None ; self.name_z = None ; self.name_t = None |
---|
1065 | self.dim_x = None ; self.dim_y = None ; self.dim_z = None ; self.dim_t = None |
---|
1066 | self.field_x = None ; self.field_y = None ; self.field_z = None ; self.field_t = None |
---|
1067 | self.dimplot = 0 |
---|
1068 | self.nplot = 1 |
---|
1069 | self.absc = None ; self.ordi = None ; self.absclab = None ; self.ordilab = None |
---|
1070 | self.verbose = True |
---|
1071 | self.swap_axes = False ; self.invert_axes = False |
---|
1072 | self.compute = None |
---|
1073 | |
---|
1074 | # open a file. for now it is netcdf. TBD for other formats. |
---|
1075 | # check that self.var is inside. |
---|
1076 | # ------------------------------- |
---|
1077 | def openfile(self): |
---|
1078 | if not os.path.exists(self.file): print '!! ERROR !! I could not find the following file: '+self.file ; exit() |
---|
1079 | if not os.path.isfile(self.file): print '!! ERROR !! This does not appear to be a file: '+self.file ; exit() |
---|
1080 | self.f = netCDF4.Dataset(self.file) |
---|
1081 | if self.verbose: print "**** OK. Opened file "+self.file |
---|
1082 | if self.var not in self.f.variables.keys(): |
---|
1083 | print '!! ERROR !! File '+self.file+' does not contain variable: '+self.var |
---|
1084 | print '..... try instead with ',self.f.variables.keys() ; exit() |
---|
1085 | |
---|
1086 | # copy attributes from another existing object |
---|
1087 | # -------------------------------------------- |
---|
1088 | def copy(self,source): |
---|
1089 | for k, v in vars(source).items(): |
---|
1090 | setattr(self,k,v) |
---|
1091 | |
---|
1092 | # get x,y,z,t dimensions from NETCDF file |
---|
1093 | # TBD: user could request for a specific altitude dimension |
---|
1094 | # TBD: staggered variables could request specific dimensions |
---|
1095 | # ------------------------------- |
---|
1096 | def getdim(self): |
---|
1097 | # GET SIZES OF EACH DIMENSION |
---|
1098 | if self.verbose: print "**** OK. Found variable "+self.var |
---|
1099 | shape = self.f.variables[self.var].shape |
---|
1100 | self.dim = len(shape) |
---|
1101 | if self.dim == 1: |
---|
1102 | if self.verbose: print "**** OK. 1D field. I assume this varies with time." |
---|
1103 | self.dim_x = 1 ; self.dim_y = 1 ; self.dim_z = 1 ; self.dim_t = shape[0] |
---|
1104 | elif self.dim == 2: |
---|
1105 | if self.verbose: print "**** OK. 2D field. I assume this is not-time-varying lat-lon map." |
---|
1106 | self.dim_x = shape[1] ; self.dim_y = shape[0] ; self.dim_z = 1 ; self.dim_t = 1 |
---|
1107 | elif self.dim == 3: |
---|
1108 | if self.verbose: print "**** OK. 3D field. I assume this is time-varying lat-lon map." |
---|
1109 | self.dim_x = shape[2] ; self.dim_y = shape[1] ; self.dim_z = 1 ; self.dim_t = shape[0] |
---|
1110 | elif self.dim == 4: |
---|
1111 | if self.verbose: print "**** OK. 4D field." |
---|
1112 | self.dim_x = shape[3] ; self.dim_y = shape[2] ; self.dim_z = shape[1] ; self.dim_t = shape[0] |
---|
1113 | # LONGITUDE. Try preset fields. If not present set grid points axis. |
---|
1114 | self.name_x = "nothing" |
---|
1115 | for c in glob_listx: |
---|
1116 | if c in self.f.variables.keys(): |
---|
1117 | self.name_x = c |
---|
1118 | if self.name_x == "nothing": |
---|
1119 | self.field_x = np.array(range(self.dim_x)) |
---|
1120 | self.name_x = "x grid points" |
---|
1121 | else: |
---|
1122 | self.field_x = self.f.variables[self.name_x] |
---|
1123 | # LATITUDE. Try preset fields. If not present set grid points axis. |
---|
1124 | self.name_y = "nothing" |
---|
1125 | for c in glob_listy: |
---|
1126 | if c in self.f.variables.keys(): |
---|
1127 | self.name_y = c |
---|
1128 | if self.name_y == "nothing": |
---|
1129 | self.field_y = np.array(range(self.dim_y)) |
---|
1130 | self.name_y = "y grid points" |
---|
1131 | else: |
---|
1132 | self.field_y = self.f.variables[self.name_y] |
---|
1133 | # ensure that lon and lat are 2D fields |
---|
1134 | # 1. simple 1D case (not time-varying) |
---|
1135 | if len(self.field_x.shape)*len(self.field_y.shape) == 1: |
---|
1136 | if self.verbose: print "**** OK. recasting lon and lat as 2D fields." |
---|
1137 | [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y) |
---|
1138 | # 2. complex 3D case (time-varying, actually just copied over time axis) |
---|
1139 | elif len(self.field_x.shape)*len(self.field_y.shape) == 9: |
---|
1140 | if self.verbose: print "**** OK. reducing lon and lat as 2D fields. get rid of time." |
---|
1141 | self.field_x = self.field_x[0,:,:] |
---|
1142 | self.field_y = self.field_y[0,:,:] |
---|
1143 | # if xy axis are apparently undefined, set 2D grid points axis. |
---|
1144 | if "grid points" not in self.name_x: |
---|
1145 | if self.field_x.all() == self.field_x[0,0]: |
---|
1146 | print "!! WARNING !! xy axis look undefined. creating non-dummy ones." |
---|
1147 | self.field_x = np.array(range(self.dim_x)) ; self.name_x = "x grid points" |
---|
1148 | self.field_y = np.array(range(self.dim_y)) ; self.name_y = "y grid points" |
---|
1149 | [self.field_x,self.field_y] = np.meshgrid(self.field_x,self.field_y) |
---|
1150 | if self.dim_x > 1: |
---|
1151 | 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()) |
---|
1152 | if self.dim_y > 1: |
---|
1153 | 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()) |
---|
1154 | # ALTITUDE. Try preset fields. If not present set grid points axis. |
---|
1155 | # WARNING: how do we do if several are available? |
---|
1156 | self.name_z = "nothing" |
---|
1157 | for c in glob_listz: |
---|
1158 | if c in self.f.variables.keys(): |
---|
1159 | self.name_z = c |
---|
1160 | if self.name_z == "nothing": |
---|
1161 | self.field_z = np.array(range(self.dim_z)) |
---|
1162 | self.name_z = "z grid points" |
---|
1163 | else: |
---|
1164 | self.field_z = self.f.variables[self.name_z][:] # specify dimension |
---|
1165 | # TBD: have to check that this is not a 3D field |
---|
1166 | if self.dim_z > 1: |
---|
1167 | 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()) |
---|
1168 | # TIME. Try preset fields. |
---|
1169 | self.name_t = "nothing" |
---|
1170 | for c in glob_listt: |
---|
1171 | if c in self.f.dimensions.keys(): |
---|
1172 | self.name_t = c |
---|
1173 | try: |
---|
1174 | # speed up: only get first value, last one. |
---|
1175 | dafirst = self.f.variables[self.name_t][0] |
---|
1176 | dalast = self.f.variables[self.name_t][self.dim_t-1] |
---|
1177 | self.field_t = np.linspace(dafirst,dalast,num=self.dim_t) |
---|
1178 | if dafirst == dalast: self.field_t = np.array([dafirst]) |
---|
1179 | except: |
---|
1180 | # ... or if a problem encountered, define a simple time axis |
---|
1181 | self.field_t = np.array(range(self.dim_t)) |
---|
1182 | self.name_t = "t grid points" |
---|
1183 | if self.dim_t > 1: |
---|
1184 | 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()) |
---|
1185 | |
---|
1186 | # get list of index to be retrieved for time axis |
---|
1187 | ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all) |
---|
1188 | # ------------------------------- |
---|
1189 | def getindextime(self,dalist=None,ind=None,stride=1): |
---|
1190 | if self.method_t == "free": |
---|
1191 | self.index_t = np.arange(0,self.dim_t,stride) |
---|
1192 | if self.dim_t > 1: |
---|
1193 | self.dimplot = self.dimplot + 1 |
---|
1194 | if self.verbose: print "**** OK. t values. all." |
---|
1195 | else: |
---|
1196 | self.method_t = "fixed" |
---|
1197 | if self.verbose: print "**** OK. no t dimension." |
---|
1198 | elif self.method_t == "comp": |
---|
1199 | start = np.argmin( np.abs( self.field_t - dalist[ind][0] ) ) |
---|
1200 | stop = np.argmin( np.abs( self.field_t - dalist[ind][1] ) ) |
---|
1201 | self.index_t = np.arange(start,stop,stride) |
---|
1202 | if self.verbose: print "**** OK. t values. comp over interval ",self.field_t[start],self.field_t[stop]," nvalues=",self.index_t.size |
---|
1203 | elif self.method_t == "fixed": |
---|
1204 | self.index_t.append( np.argmin( np.abs( self.field_t - dalist[ind][0] ) )) |
---|
1205 | if self.verbose: print "**** OK. t values",self.field_t[self.index_t] |
---|
1206 | else: |
---|
1207 | print "!! ERROR !! method "+self.method_t+" not supported" |
---|
1208 | self.index_t = np.array(self.index_t) |
---|
1209 | |
---|
1210 | # get list of index to be retrieved for vertical axis |
---|
1211 | ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all) |
---|
1212 | # ------------------------------- |
---|
1213 | def getindexvert(self,dalist=None,ind=None,stride=1): |
---|
1214 | if self.method_z == "free": |
---|
1215 | self.index_z = np.arange(0,self.dim_z,stride) |
---|
1216 | if self.dim_z > 1: |
---|
1217 | self.dimplot = self.dimplot + 1 |
---|
1218 | if self.verbose: print "**** OK. z values. all." |
---|
1219 | else: |
---|
1220 | self.method_z = "fixed" |
---|
1221 | if self.verbose: print "**** OK. no z dimension." |
---|
1222 | elif self.method_z == "comp": |
---|
1223 | start = np.argmin( np.abs( self.field_z - dalist[ind][0] ) ) |
---|
1224 | stop = np.argmin( np.abs( self.field_z - dalist[ind][1] ) ) |
---|
1225 | self.index_z = np.arange(start,stop,stride) |
---|
1226 | if self.verbose: print "**** OK. z values. comp over interval",self.field_z[start],self.field_z[stop]," nvalues=",self.index_z.size |
---|
1227 | elif self.method_z == "fixed": |
---|
1228 | self.index_z.append( np.argmin( np.abs( self.field_z - dalist[ind][0] ) )) |
---|
1229 | if self.verbose: print "**** OK. z values",self.field_z[self.index_z] |
---|
1230 | else: |
---|
1231 | if self.verbose: print "!! ERROR !! method "+self.method_z+" not supported" |
---|
1232 | self.index_z = np.array(self.index_z) |
---|
1233 | |
---|
1234 | # get list of index to be retrieved for horizontal grid |
---|
1235 | # --> index_x and index_y are slices to be retrieved from NETCDF files |
---|
1236 | # --> index_x2d and index_y2d are the actual (x,y) coordinates corresponding to each relevant point |
---|
1237 | # [this is slightly more complicated because 2D arrays for lat-lon projection possibly irregular] |
---|
1238 | # NB: to append index we use lists (the most convenient) then we convert into a numpy.array |
---|
1239 | ### TBD: il faudrait ne prendre que les indices qui correspondent a l interieur d un plot (dans all) |
---|
1240 | # ------------------------------- |
---|
1241 | def getindexhori(self,dalistx=None,dalisty=None,indx=None,indy=None,stridex=1,stridey=1): |
---|
1242 | ## get what is the method over x and y axis |
---|
1243 | test = self.method_x+self.method_y |
---|
1244 | ## CASE 0, EASY CASES: |
---|
1245 | ## - LAT IS FREE (we do here what must be done whatever LON case is) |
---|
1246 | ## - LON IS FREE (we do here what must be done whatever LAT case is) |
---|
1247 | ## - LAT IS COMP AND LON IS FREE |
---|
1248 | ## - LON IS COMP AND LAT IS FREE |
---|
1249 | if self.method_x == "free" or test in ["compfree","compcomp"]: |
---|
1250 | self.index_x = range(0,self.dim_x,stridex) |
---|
1251 | if self.dim_x > 1: |
---|
1252 | if self.method_x == "free": self.dimplot = self.dimplot + 1 |
---|
1253 | if self.verbose: print "**** OK. x values. all." |
---|
1254 | else: |
---|
1255 | self.method_x = "fixed" |
---|
1256 | if self.verbose: print "**** OK. no x dimension." |
---|
1257 | if self.method_y == "free" or test in ["freecomp","compcomp"]: |
---|
1258 | self.index_y = range(0,self.dim_y,stridey) |
---|
1259 | if self.dim_y > 1: |
---|
1260 | if self.method_y == "free": self.dimplot = self.dimplot + 1 |
---|
1261 | if self.verbose: print "**** OK. y values. all." |
---|
1262 | else: |
---|
1263 | self.method_y = "fixed" |
---|
1264 | if self.verbose: print "**** OK. no y dimension." |
---|
1265 | ## CASE 0 above, this is just for continuity for free. |
---|
1266 | ## ... for comp we have to select bounds. |
---|
1267 | ## ... TBD: take the bool array strategy for what follows! |
---|
1268 | if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]: |
---|
1269 | ### ref1_dirty_hack |
---|
1270 | ### ... for the moment this is a hack. but actually this is more powerful. |
---|
1271 | if self.method_x == "comp": |
---|
1272 | yeah = (self.field_x >= dalistx[indx][0])*(self.field_x <= dalistx[indx][1]) |
---|
1273 | self.index_x = yeah[0,:] |
---|
1274 | if self.method_y == "comp": |
---|
1275 | yeah = (self.field_y >= dalisty[indy][0]) * (self.field_y <= dalisty[indy][1]) |
---|
1276 | self.index_y = yeah[:,0] |
---|
1277 | self.index_x2d = self.index_x |
---|
1278 | self.index_y2d = self.index_y |
---|
1279 | ## AND NOW THE LITTLE BIT MORE COMPLICATED CASES |
---|
1280 | ## CASE 1 LAT AND LON ARE FIXED |
---|
1281 | elif test == "fixedfixed": |
---|
1282 | 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 ) |
---|
1283 | #TBD: pb with staggered coord |
---|
1284 | if idx not in self.index_x: self.index_x.append(idx) |
---|
1285 | if idy not in self.index_y: self.index_y.append(idy) |
---|
1286 | self.index_x2d.append(idx) |
---|
1287 | self.index_y2d.append(idy) |
---|
1288 | ## CASE 2 LON IS FIXED BUT NOT LAT |
---|
1289 | elif test in ["fixedfree","fixedcomp"]: |
---|
1290 | # find where are requested x values for each y on the free dimension |
---|
1291 | # NB: this does not work for non-bijective cases e.g. polar stereographic |
---|
1292 | for iy in range(self.dim_y): |
---|
1293 | idx = np.argmin( np.abs( self.field_x[iy,:] - dalistx[indx][0] ) ) |
---|
1294 | # if comp is requested we select only indexes which yield values between requested min and max |
---|
1295 | storeval = (self.method_y == "comp") and (self.field_y[iy,idx] > dalisty[indy][0]) and (self.field_y[iy,idx] < dalisty[indy][1]) |
---|
1296 | storeval = storeval or (self.method_y == "free") |
---|
1297 | if storeval: |
---|
1298 | if idx not in self.index_x: self.index_x.append(idx) |
---|
1299 | if iy not in self.index_y and self.method_y == "comp": self.index_y.append(iy) |
---|
1300 | if idx not in self.index_x2d or iy not in self.index_y2d: |
---|
1301 | self.index_x2d.append(idx) |
---|
1302 | self.index_y2d.append(iy) |
---|
1303 | ## CASE 3 LAT IS FIXED BUT NOT LON |
---|
1304 | elif test in ["freefixed","compfixed"]: |
---|
1305 | # find where are requested y values for each x on the free dimension |
---|
1306 | # NB: this does not work for non-bijective cases e.g. polar stereographic |
---|
1307 | for ix in range(self.dim_x): |
---|
1308 | idy = np.argmin( np.abs( self.field_y[:,ix] - dalisty[indy][0] ) ) |
---|
1309 | # if comp is requested we select only indexes which yield values between requested min and max |
---|
1310 | storeval = (self.method_x == "comp") and (self.field_x[idy,ix] > dalistx[indx][0]) and (self.field_x[idy,ix] < dalistx[indx][1]) |
---|
1311 | storeval = storeval or (self.method_x == "free") |
---|
1312 | if storeval: |
---|
1313 | if idy not in self.index_y: self.index_y.append(idy) |
---|
1314 | if ix not in self.index_x and self.method_x == "comp": self.index_x.append(ix) |
---|
1315 | if ix not in self.index_x2d or idy not in self.index_y2d: |
---|
1316 | self.index_x2d.append(ix) |
---|
1317 | self.index_y2d.append(idy) |
---|
1318 | ## check index tab |
---|
1319 | if len(self.index_x) == 0 or len(self.index_y) == 0: |
---|
1320 | print "!! ERROR !! no indices found. check prescribed latitudes or longitudes" ; exit() |
---|
1321 | ## ensure the array is a numpy array for getfield to work |
---|
1322 | self.index_x = np.array(self.index_x) |
---|
1323 | self.index_y = np.array(self.index_y) |
---|
1324 | self.index_x2d = np.array(self.index_x2d) |
---|
1325 | self.index_y2d = np.array(self.index_y2d) |
---|
1326 | ### print extrema |
---|
1327 | printx = self.field_x[np.ix_(self.index_y2d, self.index_x2d)] |
---|
1328 | printy = self.field_y[np.ix_(self.index_y2d, self.index_x2d)] |
---|
1329 | if self.verbose: |
---|
1330 | print "**** OK. x values (min,max).", printx.min(),printx.max() |
---|
1331 | print "**** OK. y values (min,max).", printy.min(),printy.max() |
---|
1332 | |
---|
1333 | # get the field from the NETCDF file and perform averages |
---|
1334 | # ------------------------------- |
---|
1335 | def getfield(self): |
---|
1336 | ## first tell what is to be done |
---|
1337 | if self.dimplot > 2: print "**** !! ERROR !! "+str(self.dimplot)+"D plots not supported!" ; exit() |
---|
1338 | elif self.dimplot == 0 and self.verbose: print "**** OK. 0D value requested." |
---|
1339 | elif self.dimplot == 1 and self.verbose: print "**** OK. 1D plot requested." |
---|
1340 | elif self.verbose: print "**** OK. 2D section requested." |
---|
1341 | # well, now get field from netcdf file |
---|
1342 | # part below is necessary otherwise there is an index error below |
---|
1343 | if self.index_x.size == 1: self.index_x = self.index_x[0] |
---|
1344 | if self.index_y.size == 1: self.index_y = self.index_y[0] |
---|
1345 | if self.index_z.size == 1: self.index_z = self.index_z[0] |
---|
1346 | if self.index_t.size == 1: self.index_t = self.index_t[0] |
---|
1347 | # then retrieve what is requested by user |
---|
1348 | # each self.dim case corresponds to tests in the beginning of getdim. |
---|
1349 | time0 = timelib.time() |
---|
1350 | if self.verbose: print "**** OK. I am getting values from files. Please wait." |
---|
1351 | if self.dim == 1: |
---|
1352 | nt = self.index_t.size ; nz = 1 ; ny = 1 ; nx = 1 |
---|
1353 | self.field = self.f.variables[self.var][self.index_t] |
---|
1354 | elif self.dim == 2: |
---|
1355 | nt = 1 ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size |
---|
1356 | self.field = self.f.variables[self.var][self.index_y,self.index_x] |
---|
1357 | elif self.dim == 3: |
---|
1358 | nt = self.index_t.size ; nz = 1 ; ny = self.index_y.size ; nx = self.index_x.size |
---|
1359 | self.field = self.f.variables[self.var][self.index_t,self.index_y,self.index_x] |
---|
1360 | # this is far faster than retrieving each term with a loop |
---|
1361 | elif self.dim == 4: |
---|
1362 | nt = self.index_t.size ; nz = self.index_z.size ; ny = self.index_y.size ; nx = self.index_x.size |
---|
1363 | self.field = self.f.variables[self.var][self.index_t,self.index_z,self.index_y,self.index_x] |
---|
1364 | else: |
---|
1365 | print "!! ERROR !! field would have more than four dimensions ?" ; exit() |
---|
1366 | # NB: ... always 4D array but possibly with "size 1" dimensions |
---|
1367 | # ... if one dimension is missing because 1D 2D or 3D requests, make it appear again |
---|
1368 | try: |
---|
1369 | self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1370 | except: |
---|
1371 | # dirty hack (AS) ref1_dirty_hack |
---|
1372 | # waiting for more fundamental modifications. case when self.index_y is a bool array. |
---|
1373 | # ... be careful if no point... |
---|
1374 | nx = np.sum(self.index_x) ## gives the size of the True part! |
---|
1375 | ny = np.sum(self.index_y) |
---|
1376 | self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1377 | if self.verbose: print "**** OK. I got %7.1e values. This took me %6.4f seconds" % (nx*ny*nz*nt,timelib.time() - time0) |
---|
1378 | if self.verbose: print "**** OK. I got var "+self.var+" with shape",self.field.shape |
---|
1379 | # reduce coordinates to useful points |
---|
1380 | # ... TBD: this should be ordered in the case of non-regular projections |
---|
1381 | if self.method_x in ["free","comp"] and self.method_y in ["free","comp"]: |
---|
1382 | # we need 2D coordinates (free) or we get broadcast problem (comp) so we use np.ix |
---|
1383 | self.field_x = self.field_x[np.ix_(self.index_y2d, self.index_x2d)] |
---|
1384 | self.field_y = self.field_y[np.ix_(self.index_y2d, self.index_x2d)] |
---|
1385 | else: |
---|
1386 | # we are OK with 1D coordinates |
---|
1387 | self.field_x = self.field_x[self.index_y2d, self.index_x2d] |
---|
1388 | self.field_y = self.field_y[self.index_y2d, self.index_x2d] |
---|
1389 | self.field_z = self.field_z[self.index_z] |
---|
1390 | self.field_t = self.field_t[self.index_t] |
---|
1391 | # extract relevant horizontal points |
---|
1392 | # TBD: is compfree OK with computing on irregular grid? |
---|
1393 | test = self.method_x + self.method_y |
---|
1394 | if test in ["fixedfixed","freefree"]: |
---|
1395 | pass |
---|
1396 | elif test in ["fixedfree","fixedcomp"] or test in ["freefixed","compfixed"]: |
---|
1397 | time0 = timelib.time() |
---|
1398 | # now have to obtain the new indexes which correspond to the extracted self.field |
---|
1399 | # for it to work with unique index, ensure that any index_* is a numpy array |
---|
1400 | if not isinstance(self.index_x, np.ndarray): self.index_x = np.array([self.index_x]) |
---|
1401 | if not isinstance(self.index_y, np.ndarray): self.index_y = np.array([self.index_y]) |
---|
1402 | if not isinstance(self.index_z, np.ndarray): self.index_z = np.array([self.index_z]) |
---|
1403 | if not isinstance(self.index_t, np.ndarray): self.index_t = np.array([self.index_t]) |
---|
1404 | for val in self.index_x: self.index_x2d[np.where(self.index_x2d == val)] = np.where(self.index_x == val)[0] |
---|
1405 | for val in self.index_y: self.index_y2d[np.where(self.index_y2d == val)] = np.where(self.index_y == val)[0] |
---|
1406 | ##### VERY EXPENSIVE |
---|
1407 | ## recast self.field with 2D horizontal arrays because we might have extracted |
---|
1408 | ## more than what is to be actually plot or computed, in particular for comps on 2D lat,lon coordinates |
---|
1409 | #self.field = self.field[np.ix_(self.index_t,self.index_z,self.index_y2d,self.index_x2d)] |
---|
1410 | #(nt,nz,ny,nx) = self.field.shape |
---|
1411 | # prepare the loop on all relevant horizontal points |
---|
1412 | if self.method_x in ["comp","free"]: |
---|
1413 | nnn = self.index_x2d.shape[0] ; what_I_am_supposed_to_do = "keepx" |
---|
1414 | elif self.method_y in ["comp","free"]: |
---|
1415 | nnn = self.index_y2d.shape[0] ; what_I_am_supposed_to_do = "keepy" |
---|
1416 | # LOOP to extract only useful values over horizontal dimensions |
---|
1417 | # only take diagonal terms, do not loop on all self.index_x2d*self.index_y2d |
---|
1418 | # ... this method is fast enough, perhaps there is a faster way though |
---|
1419 | # ... (for sure the method with np.diag is much slower) |
---|
1420 | for iii in range(nnn): |
---|
1421 | ix = self.index_x2d[iii] ; iy = self.index_y2d[iii] |
---|
1422 | for iz in range(self.index_z.size): |
---|
1423 | for it in range(self.index_t.size): |
---|
1424 | if what_I_am_supposed_to_do == "keepx": self.field[it,iz,0,ix] = self.field[it,iz,iy,ix] |
---|
1425 | elif what_I_am_supposed_to_do == "keepy": self.field[it,iz,iy,0] = self.field[it,iz,iy,ix] |
---|
1426 | if self.verbose: print "**** OK. I got to pick the right values for your request. This took me %6.4f seconds" % (timelib.time() - time0) |
---|
1427 | # we only keep the one value that was modified on the dimension which is not free |
---|
1428 | 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)) |
---|
1429 | 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)) |
---|
1430 | # make a mask in case there are non-NaN missing values. (what about NaN missing values?) |
---|
1431 | # ... this is important for computations below (see ppcompute) |
---|
1432 | masked = np.ma.masked_where(np.abs(self.field) > 1e25,self.field) |
---|
1433 | if masked.mask.any() == True: |
---|
1434 | if self.verbose: print "!! WARNING !! Values over +-1e25 are considered missing values." |
---|
1435 | self.field = masked |
---|
1436 | self.field.set_fill_value([np.NaN]) |
---|
1437 | |
---|
1438 | # perform computations |
---|
1439 | # ------------------------------- |
---|
1440 | # available: mean, max, min, meanarea |
---|
1441 | # TB: integrals? for derivatives, define a function self.dx() |
---|
1442 | def computations(self): |
---|
1443 | nt,nz,ny,nx = self.field.shape |
---|
1444 | # treat the case of mean on fields normalized with grid mesh area |
---|
1445 | # ... this is done in the .area() method. |
---|
1446 | # after that self.field contains field*area/totarea |
---|
1447 | if "area" in self.compute: |
---|
1448 | if "comp" in self.method_x+self.method_y: |
---|
1449 | self.area() |
---|
1450 | else: |
---|
1451 | if self.verbose: print "!! WARNING !! No area accounted for (computing on t and/or z axis)." |
---|
1452 | # now ready to compute [TBD: we would like to have e.g. mean over x,y and min over t ??] |
---|
1453 | if self.method_t == "comp": |
---|
1454 | if self.verbose: print "**** OK. Computing over t axis." |
---|
1455 | if "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=0) |
---|
1456 | elif self.compute == "min": self.field = ppcompute.min(self.field,axis=0) |
---|
1457 | elif self.compute == "max": self.field = ppcompute.max(self.field,axis=0) |
---|
1458 | else: print "!! ERROR !! operation not supported." ; exit() |
---|
1459 | nt = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1460 | if self.method_z == "comp": |
---|
1461 | if self.verbose: print "**** OK. Computing over z axis." |
---|
1462 | if "mean" in self.compute: self.field = ppcompute.mean(self.field,axis=1) |
---|
1463 | elif self.compute == "min": self.field = ppcompute.min(self.field,axis=1) |
---|
1464 | elif self.compute == "max": self.field = ppcompute.max(self.field,axis=1) |
---|
1465 | else: print "!! ERROR !! operation not supported." ; exit() |
---|
1466 | nz = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1467 | if self.method_y == "comp": |
---|
1468 | if self.verbose: print "**** OK. Computing over y axis." |
---|
1469 | if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=2) |
---|
1470 | elif self.compute == "min": self.field = ppcompute.min(self.field,axis=2) |
---|
1471 | elif self.compute == "max": self.field = ppcompute.max(self.field,axis=2) |
---|
1472 | elif self.compute == "meanarea": self.field = ppcompute.sum(self.field,axis=2) |
---|
1473 | else: print "!! ERROR !! operation not supported." ; exit() |
---|
1474 | ny = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1475 | if self.field_x.ndim == 2: self.field_x = self.field_x[0,:] # TBD: this is OK for regular grid but not for irregular |
---|
1476 | if self.method_x == "comp": |
---|
1477 | if self.verbose: print "**** OK. Computing over x axis." |
---|
1478 | if self.compute == "mean": self.field = ppcompute.mean(self.field,axis=3) |
---|
1479 | elif self.compute == "min": self.field = ppcompute.min(self.field,axis=3) |
---|
1480 | elif self.compute == "max": self.field = ppcompute.max(self.field,axis=3) |
---|
1481 | elif self.compute == "meanarea": self.field = ppcompute.sum(self.field,axis=3) |
---|
1482 | else: print "!! ERROR !! operation not supported." ; exit() |
---|
1483 | nx = 1 ; self.field = np.reshape(self.field,(nt,nz,ny,nx)) |
---|
1484 | if self.field_y.ndim == 2: self.field_y = self.field_y[:,0] # TBD: this is OK for regular grid but not for irregular |
---|
1485 | # remove all dimensions with size 1 to prepare plot (and check the resulting dimension with dimplot) |
---|
1486 | self.field = np.squeeze(self.field) |
---|
1487 | if self.field.ndim != self.dimplot: |
---|
1488 | print "!! ERROR !! Problem: self.field is different than plot dimensions", self.field.ndim, self.dimplot ; exit() |
---|
1489 | if self.verbose: |
---|
1490 | print "**** OK. Final shape for "+self.var+" after averaging and squeezing",self.field.shape |
---|
1491 | |
---|
1492 | # get areas for computations and ponderate self.field by area/totarea |
---|
1493 | # ------------------------------------------------------------------- |
---|
1494 | def area(self): |
---|
1495 | if self.verbose: print "**** OK. Get area array for computations." |
---|
1496 | # create a request object for area |
---|
1497 | # ... and copy known attributes from self |
---|
1498 | aire = onerequest() |
---|
1499 | aire.copy(self) |
---|
1500 | # get area field name |
---|
1501 | aire.var = "nothing" |
---|
1502 | for c in glob_listarea: |
---|
1503 | if c in aire.f.variables.keys(): |
---|
1504 | aire.var = c |
---|
1505 | # do not try to calculate areas automatically |
---|
1506 | if aire.var == "nothing": |
---|
1507 | print "!! ERROR !! area variable not found... needs to be added in set_ppclass.txt?" |
---|
1508 | exit() |
---|
1509 | # define area request dimensions |
---|
1510 | aire.getdim() |
---|
1511 | # ensure this is a 2D horizontal request and define indexes |
---|
1512 | # ... areas are not supposed to vary with time and height |
---|
1513 | aire.method_x = "free" ; aire.method_y = "free" |
---|
1514 | aire.getindexhori() ; aire.dimplot = 2 |
---|
1515 | aire.method_z = "fixed" ; aire.field_z = np.array([0]) ; aire.index_z = np.array([0]) |
---|
1516 | aire.method_t = "fixed" ; aire.field_t = np.array([0]) ; aire.index_t = np.array([0]) |
---|
1517 | # read the 2D area array in netCDF file |
---|
1518 | aire.getfield() |
---|
1519 | aire.field = np.squeeze(aire.field) |
---|
1520 | # reduce with self horizontal indexes |
---|
1521 | if "fixed" in self.method_x+self.method_y: |
---|
1522 | aire.field = aire.field[self.index_y,self.index_x] |
---|
1523 | # calculate total area |
---|
1524 | # ... 2D comp is easy. 1D comp is a bit less easy but simple array manipulation. |
---|
1525 | if "free" in self.method_x+self.method_y: |
---|
1526 | if self.method_x == "free": |
---|
1527 | totarea = ppcompute.sum(aire.field,axis=0) |
---|
1528 | totarea = np.reshape(totarea,(1,totarea.size)) |
---|
1529 | totarea = np.tile(totarea,(1,self.index_x)) |
---|
1530 | elif self.method_y == "free": |
---|
1531 | totarea = ppcompute.sum(aire.field,axis=1) |
---|
1532 | totarea = np.reshape(totarea,(totarea.size,1)) |
---|
1533 | totarea = np.tile(totarea,(1,self.index_x.size)) |
---|
1534 | elif self.method_x == "comp" and self.method_y == "comp": |
---|
1535 | aire.field = aire.field[np.ix_(self.index_y, self.index_x)] # reduce to requested indexes only |
---|
1536 | totarea = ppcompute.sum(ppcompute.sum(aire.field,axis=1),axis=0) |
---|
1537 | else: |
---|
1538 | if self.verbose: print "!! WARNING !! Not account for areas. Only averaging over z and/or t axis." |
---|
1539 | # normalize by total area |
---|
1540 | print "**** OK. I can now normalize field by areas." |
---|
1541 | aire.field = aire.field / totarea |
---|
1542 | # tile area array over self t and z axis so that area field could be multiplied with self.field |
---|
1543 | aire.field = np.tile(aire.field,(self.index_t.size,self.index_z.size,1,1)) |
---|
1544 | if self.field.shape != aire.field.shape: |
---|
1545 | print "!! ERROR !! Problem in area(). Check array shapes." |
---|
1546 | print "Field vs. aire:",self.field.shape,aire.field.shape ; exit() |
---|
1547 | else: |
---|
1548 | self.field = self.field*aire.field |
---|
1549 | |
---|
1550 | # define coordinates for plot |
---|
1551 | # ------------------------------- |
---|
1552 | def definecoord(self): |
---|
1553 | I_got_abs = False ; I_got_ord = False |
---|
1554 | # here is the thing. time is usually taken as an abscissa so we start with time. |
---|
1555 | if self.method_t == "free": |
---|
1556 | self.absc = self.field_t ; self.absclab = self.name_t |
---|
1557 | I_got_abs = True |
---|
1558 | # then we usually have x as an abscissa. |
---|
1559 | if self.method_x == "free": |
---|
1560 | if I_got_abs: |
---|
1561 | self.ordi = self.field_x ; self.ordilab = self.name_x |
---|
1562 | I_got_ord = True |
---|
1563 | else: |
---|
1564 | self.absc = self.field_x ; self.absclab = self.name_x |
---|
1565 | I_got_abs = True |
---|
1566 | # ... or we have y |
---|
1567 | if self.method_y == "free": |
---|
1568 | if I_got_abs: |
---|
1569 | self.ordi = self.field_y ; self.ordilab = self.name_y |
---|
1570 | I_got_ord = True |
---|
1571 | else: |
---|
1572 | self.absc = self.field_y ; self.absclab = self.name_y |
---|
1573 | I_got_abs = True |
---|
1574 | # ... and we end with z because it is usually not an abscissa (profiles). |
---|
1575 | if self.method_z == "free": |
---|
1576 | if self.field_z[0] > self.field_z[1]: |
---|
1577 | self.invert_axes = True # the axis will be turned upside-down |
---|
1578 | if I_got_abs: |
---|
1579 | self.ordi = self.field_z ; self.ordilab = self.name_z |
---|
1580 | I_got_ord = True |
---|
1581 | else: |
---|
1582 | self.absc = self.field_z ; self.absclab = self.name_z |
---|
1583 | I_got_abs = True |
---|
1584 | self.swap_axes = True # says that altitude is not supposed to remain as an abscissa |
---|
1585 | if I_got_abs and self.verbose: print "**** OK. abscissa:",self.absclab, self.absc.shape |
---|
1586 | if I_got_ord and self.verbose: print "**** OK. ordinate:",self.ordilab, self.ordi.shape |
---|