Context Navigation

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

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