Context Navigation

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

Last change on this file since 915 was 915, checked in by aslmd, 12 years ago
UTIL PYTHON planetoplot_v2. error corrected in previous commit.
File size: 68.4 KB

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

Download in other formats: