1
2
3
4
5 """Functions to plot fields from Qtcm model objects and run.
6
7 Module Function:
8 * nice_levels: Compute a vector of "levels" at "nice" increments.
9
10 * plot_ncdf_output: Plot data from QTCM1 netCDF output files.
11 """
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45 import os, sys
46 if (__name__ == "__main__") or \
47 ("pydoc" in os.path.basename(sys.argv[0])):
48 import user
49
50
51
52
53 import package_version as _package_version
54 __version__ = _package_version.version
55 __author__ = _package_version.author
56 __date__ = _package_version.date
57 __credits__ = _package_version.credits
58
59
60
61
62 import num_settings as num
63 from num_settings import N
64
65
66
67
68 from matplotlib import rc as matplotlibrc
69
70 matplotlibrc('text', usetex=True)
71 import pylab
72 import matplotlib
73 from matplotlib.toolkits.basemap import Basemap
74
75
76
77
78 import copy
79 import Scientific.IO.NetCDF as S
80 import shutil
81 import tempfile
82 from where_close import where_close
83
84
85
86
87
88
90 """Compute a vector of "levels" at "nice" increments.
91
92 Returns a 1-D array of "levels" (e.g., contour levels) calculated
93 to give an aesthetically pleasing and human-readable interval,
94 if possible. If not, returns levels for approx_nlev levels
95 between the maximum and minimum of data. In any event, the
96 function will return no more than max_nlev levels.
97
98 Keyword Input Parameter:
99 * data: Array of values to calculate levels for. Can be of any
100 size and shape.
101
102 Keyword Input Parameter:
103 * approx_nlev: Integer referring to approximately how many
104 levels to return. This is the way of adjusting how "coarse"
105 or "fine" to make the vector of levels.
106
107 * max_nlev: The maximum number of levels the function will
108 permit to be returned. The interval of levels will be adjusted
109 to keep the number of levels returned under this value. If
110 approx_nlev is chosen to be greater than or equal to max_nlev,
111 an exception is raised.
112
113 Output:
114 * This function returns a 1-D array of contour levels.
115
116 Function is adaptation of parts of IDL routine contour_plot.pro
117 by Johnny Lin. This is why the capitalization conventions of
118 Python are not strictly followed in this function.
119
120 Examples:
121 >>> z = N.array([-24.5, 50.3, 183.1, 20.])
122 >>> out = nice_levels(z)
123 >>> ['%g' % out[i] for i in range(len(out))]
124 ['-30', '0', '30', '60', '90', '120', '150', '180', '210']
125
126 >>> z = N.array([-24.5, 50.3, 183.1, 20.])
127 >>> out = nice_levels(z, approx_nlev=5)
128 >>> ['%g' % out[i] for i in range(len(out))]
129 ['-50', '0', '50', '100', '150', '200']
130
131 >>> z = N.array([-24.5, 50.3, 183.1, 20.])
132 >>> out = nice_levels(z, approx_nlev=10)
133 >>> ['%g' % out[i] for i in range(len(out))]
134 ['-30', '0', '30', '60', '90', '120', '150', '180', '210']
135 """
136
137
138 if approx_nlev >= max_nlev:
139 raise ValueError, 'max_nlev is too small'
140
141 MAX_zd_ok = N.max(data)
142 MIN_zd_ok = N.min(data)
143
144 nlevels = N.min([approx_nlev, max_nlev])
145 tmpcmax = MAX_zd_ok
146 tmpcmin = MIN_zd_ok
147 tmpcint = N.abs( (tmpcmax-tmpcmin)/float(nlevels) )
148
149
150
151
152
153
154
155
156 guesscint = N.abs( (MAX_zd_ok-MIN_zd_ok)/float(nlevels) )
157
158 if (guesscint > 1e-10) and (guesscint < 1e+10):
159 possiblecint = [1e-10, 1e-9, 1e-8, 1e-7, 1e-6, 1e-5,
160 0.0001, 0.0002, 0.0005,
161 0.001, 0.002, 0.005,
162 0.01, 0.02, 0.05,
163 0.1, 0.2, 0.5,
164 1., 2., 5.,
165 10., 20., 30., 45., 50.,
166 100., 200., 500.,
167 1000., 2000., 5000.,
168 10000., 20000., 50000.,
169 1e+5, 1e+6, 1e+7, 1e+8, 1e+9, 1e+10]
170
171 diffcint = N.abs(possiblecint-guesscint)
172 tempcint = N.compress( diffcint == N.min(diffcint), possiblecint )[0]
173 tcidx = N.compress( where_close( possiblecint, tempcint ),
174 N.arange(N.size(possiblecint)) )[0]
175
176
177
178
179
180
181
182
183
184 if tcidx == 0: tcidx = 1
185 if tcidx == N.size(possiblecint)-1: tcidx = N.size(possiblecint)-2
186
187 ncon_count = {}
188 test_tcidxs = [tcidx-1, tcidx, tcidx+1]
189 for i in test_tcidxs:
190 itcval = possiblecint[i]
191 positivecon = N.arange(max_nlev+2, dtype=float)*itcval
192 negativecon = (N.arange(max_nlev+2, dtype=float)+1.0)*itcval*(-1.0)
193 if (MAX_zd_ok + itcval >= 0) and (MIN_zd_ok - itcval >= 0):
194 ncon_count[i] = N.sum( \
195 N.logical_and(positivecon <= MAX_zd_ok + itcval,
196 positivecon >= MIN_zd_ok - itcval) )
197 elif (MAX_zd_ok + itcval < 0) and (MIN_zd_ok - itcval < 0):
198 ncon_count[i] = N.sum( \
199 N.logical_and(negativecon <= MAX_zd_ok + itcval,
200 negativecon >= MIN_zd_ok - itcval) )
201 else:
202 ncon_count[i] = N.sum(positivecon <= MAX_zd_ok + itcval) \
203 + N.sum(negativecon >= MIN_zd_ok - itcval)
204
205
206
207
208
209
210
211 min_ncon_count = N.min(ncon_count.values())
212 current_best_count = max_nlev
213 for i in test_tcidxs:
214 if (ncon_count[i] == min_ncon_count) and \
215 (ncon_count[i] >= nlevels-1):
216 tempcint = possiblecint[i]
217 current_best_count = ncon_count[i]
218 break
219 elif (ncon_count[i] == min_ncon_count) and \
220 (ncon_count[i] < nlevels-1):
221 continue
222 elif ncon_count[i] > max_nlev:
223 continue
224 else:
225 if N.abs(ncon_count[i]-nlevels) < \
226 N.abs(current_best_count-nlevels):
227 tempcint = possiblecint[i]
228 current_best_count = ncon_count[i]
229 continue
230
231
232
233
234
235
236 positivecon = N.arange(max_nlev+2, dtype=float)*tempcint
237 negativecon = (N.arange(max_nlev+2, dtype=float)+1.0)*tempcint*(-1.0)
238
239 if (MAX_zd_ok + tempcint >= 0) and (MIN_zd_ok - tempcint >= 0):
240 tmpclevels = N.compress( \
241 N.logical_and(positivecon <= MAX_zd_ok + tempcint,
242 positivecon >= MIN_zd_ok - tempcint),
243 positivecon )
244 elif (MAX_zd_ok + tempcint < 0) and (MIN_zd_ok - tempcint < 0):
245 tmpclevels = N.compress( \
246 N.logical_and(negativecon <= MAX_zd_ok + tempcint,
247 negativecon >= MIN_zd_ok - tempcint),
248 negativecon )
249 else:
250 uppercon = N.compress( positivecon <= MAX_zd_ok + tempcint,
251 positivecon )
252 lowercon = N.compress( negativecon >= MIN_zd_ok - tempcint,
253 negativecon )
254 tmpclevels = N.concatenate([lowercon, uppercon])
255
256
257
258
259
260 tmpclevels = N.sort(tmpclevels)
261 if (N.size(tmpclevels) <= max_nlev ) and (N.size(tmpclevels) > 0):
262 nlevels = N.size(tmpclevels)
263 tmpcmax = tmpclevels[-1]
264 tmpcmin = tmpclevels[0]
265 tmpcint = tempcint
266
267 else:
268 pass
269
270
271
272
273 return N.arange(tmpcmin, tmpcmax+tmpcint, tmpcint)
274
275
276
277
278
279
281 """Return instring to handle super/subscripts of one character.
282
283 The string returned expresses instring so that an exponent or
284 subscript with a single character after it (e.g., ^2, _s) can
285 be processed as a LaTeX exponent or superscript in Matplotlib.
286 Any number of these super/subscripts can be present in instring.
287 See http://www.scipy.org/Cookbook/Matplotlib/UsingTex by Darrin
288 Dale for some code snippets incorporated in this function.
289
290 Positional Input Parameter:
291 * instring: String to be formatted.
292
293 Output:
294 * A string, processable by Matplotlib's LaTeX emulation module.
295
296 Examples:
297 >>> mpl_latex_script1('Precipitation [W/m^2]')
298 'Precipitation [W/m$^2$]'
299 >>> mpl_latex_script1('u_0^2 and u_1^2 [m^2/s^2]')
300 'u$_0$$^2$ and u$_1$$^2$ [m$^2$/s$^2$]'
301 """
302
303
304 instrlist = list(instring)
305
306
307
308
309 for ichar in ['^', '_']:
310 num_char = instrlist.count(ichar)
311 for inum in xrange(num_char):
312 cidx = instrlist.index(ichar)
313 instrlist[cidx] = '$' + ichar
314 instrlist[cidx+1] = instrlist[cidx+1] +'$'
315
316 return ''.join(instrlist)
317
318
319
320
321
322
324 """Plot model field id from the data in netCDF file datafn.
325
326 Positional Input Parameter:
327 * id: Name of the id of the field to plot. String.
328
329 * datafn: Filename containing the output data to plot. String.
330
331 Input keyword parameter descriptions are found in the docstring
332 for Qtcm methods ploti, plotm, and other methods that call this
333 private method. In general, those methods take the keyword
334 parameters they receive and pass it along unchanged as keyword
335 parameters to this function. In that sense, this function is
336 seldom used as a stand-alone function, but rather is usually
337 used coupled with a Qtcm instance.
338
339 The data fields read in from the netCDF output file are dimensioned
340 (time, lat, lon). This is different than how the data is stored
341 in the compiled QTCM model fields (lon, lat, time), and at the
342 Python level (lon, lat). The reason this is the case is that
343 f2py automatically makes the arrays passed between the Python
344 and Fortran levels match.
345
346 For a lat vs. lon plot, the contour plot is superimposed onto
347 a cylindrical projection map of the Earth with continents drawn
348 and labeled meridians and parallels. The title also includes
349 the model time, and x- and y-axis labels are not drawn.
350
351 All numerical data used for plotting come from the netCDF output
352 file for consistency (e.g., the dimensions of u1). Currently
353 this method only works for 3-D data arrays (two in space, one
354 in time).
355 """
356
357
358
359
360
361 if id == 'Qc': iduse = 'Prec'
362 elif id == 'FLWut': iduse = 'OLR'
363 elif id == 'STYPE': iduse = 'stype'
364 else: iduse = id
365
366
367
368
369
370 plotkwds_ids = ['lat', 'lon', 'time', 'fn', 'levels', 'title',
371 'xlabel', 'ylabel',
372 'filled', 'nlatlon', 'tmppreview']
373
374 plotkwds = {}
375 for ikey in plotkwds_ids:
376 if kwds.has_key(ikey):
377 plotkwds[ikey] = copy.copy(kwds[ikey])
378 else:
379 plotkwds[ikey] = None
380
381 if not kwds.has_key('nlatlon'):
382 plotkwds['nlatlon'] = 8
383
384
385
386
387 fileobj = S.NetCDFFile(datafn, mode='r')
388 data = N.array(fileobj.variables[iduse].getValue())
389 data_name = fileobj.variables[iduse].long_name
390 data_units = fileobj.variables[iduse].units
391
392 dim = {}
393 dimname = {}
394 dimunits = {}
395
396 dim['lat'] = N.array(fileobj.variables['lat'].getValue())
397 dimname['lat'] = fileobj.variables['lat'].long_name
398 dimunits['lat'] = fileobj.variables['lat'].units
399
400 dim['lon'] = N.array(fileobj.variables['lon'].getValue())
401 dimname['lon'] = fileobj.variables['lon'].long_name
402 dimunits['lon'] = fileobj.variables['lon'].units
403
404 dim['time'] = N.array(fileobj.variables['time'].getValue())
405 dimname['time'] = fileobj.variables['time'].long_name
406 dimunits['time'] = fileobj.variables['time'].units
407
408 fileobj.close()
409
410
411
412
413
414
415
416
417 idx1 = data_name.find('[')
418 idx2 = data_name.find(']')
419 if idx1 != -1 and idx2 != -1:
420 data_name = data_name[:idx1] + data_name[idx2+1:]
421 data_name = data_name.strip()
422
423 data_name = ' '.join(data_name.replace('_',' ').split())
424 data_units = ' '.join(data_units.replace('_',' ').split())
425
426
427
428
429
430
431
432
433
434 for idimkey in dim.keys():
435 idimname = dimname[idimkey]
436 idx1 = idimname.find('[')
437 idx2 = idimname.find(']')
438 if idx1 != -1 and idx2 != -1:
439 idimname = idimname[:idx1] + idimname[idx2+1:]
440 dimname[idimkey] = idimname.strip()
441
442 dimname[idimkey] = \
443 ' '.join(dimname[idimkey].replace('_',' ').split()).title()
444 dimunits[idimkey] = \
445 ' '.join(dimunits[idimkey].replace('_',' ').split()).title()
446
447
448
449
450 if N.rank(data) != 3:
451 raise ValueError, '_plot: can only plot lat, lon, time fields'
452 if not N.allclose(dim['time'], N.sort(dim['time'])):
453 raise ValueError, '_plot: time not monotonically ascending'
454 if not N.allclose(dim['lat'], N.sort(dim['lat'])):
455 raise ValueError, '_plot: lat not monotonically ascending'
456 if not N.allclose(dim['lon'], N.sort(dim['lon'])):
457 raise ValueError, '_plot: lon not monotonically ascending'
458 if N.shape(data)[0] != N.size(dim['time']):
459 raise ValueError, '_plot: data time dim mismatch'
460 if N.shape(data)[1] != N.size(dim['lat']):
461 raise ValueError, '_plot: data lat dim mismatch'
462 if N.shape(data)[2] != N.size(dim['lon']):
463 raise ValueError, '_plot: data lon dim mismatch'
464
465
466
467
468
469
470
471
472
473
474
475
476 rngs = {}
477 rngs_idxs = {}
478 keys_rngs_sizes_gt_1 = []
479 for idimkey in dim.keys():
480 idim = dim[idimkey]
481
482 if plotkwds[idimkey] == None:
483 dim_mask = N.ones( N.size(idim), dtype=int )
484
485 elif N.isscalar(plotkwds[idimkey]):
486 dim_mask = where_close( idim, plotkwds[idimkey] )
487 if N.sum(dim_mask) != 1:
488 raise ValueError, 'no point chosen'
489
490 elif (not N.isscalar(plotkwds[idimkey])) and \
491 N.size(plotkwds[idimkey]) == 1:
492 dim_mask = where_close( idim, plotkwds[idimkey][0] )
493 if N.sum(dim_mask) != 1:
494 raise ValueError, 'no point chosen'
495
496 elif N.size(plotkwds[idimkey]) == 2:
497 dim_mask = N.logical_and( idim >= plotkwds[idimkey][0],
498 idim <= plotkwds[idimkey][-1] )
499
500 else:
501 raise ValueError, 'bad dimension range keyword entry'
502
503 rngs[idimkey] = N.compress( dim_mask, idim )
504 rngs_idxs[idimkey] = N.compress( dim_mask, N.arange(N.size(idim)) )
505 if N.size(rngs[idimkey]) > 1:
506 keys_rngs_sizes_gt_1.append(idimkey)
507
508
509
510
511 if len(keys_rngs_sizes_gt_1) == 0:
512 raise ValueError, 'cannot plot without any fixed dimension'
513 elif len(keys_rngs_sizes_gt_1) == 1:
514 plottype = 'line'
515 elif len(keys_rngs_sizes_gt_1) == 2:
516 plottype = 'contour'
517 else:
518 raise ValueError, 'cannot plot with > 2 varying dimensions'
519
520
521
522
523
524
525
526
527
528
529
530 if plottype == 'line':
531 x = rngs[keys_rngs_sizes_gt_1[0]]
532 xname = dimname[keys_rngs_sizes_gt_1[0]] + ' [' \
533 + dimunits[keys_rngs_sizes_gt_1[0]] + ']'
534 yname = data_name + ' [' + data_units + ']'
535
536 elif plottype == 'contour':
537 if keys_rngs_sizes_gt_1.count('lon'):
538 x = rngs['lon']
539 xname = dimname['lon'] + ' [' + dimunits['lon'] + ']'
540 if keys_rngs_sizes_gt_1.count('lat'):
541 y = rngs['lat']
542 yname = dimname['lat'] + ' [' + dimunits['lat'] + ']'
543 if keys_rngs_sizes_gt_1.count('time'):
544 if keys_rngs_sizes_gt_1.count('lon'):
545 y = rngs['time']
546 yname = dimname['time'] + ' [' + dimunits['time'] + ']'
547 elif keys_rngs_sizes_gt_1.count('lat'):
548 x = rngs['time']
549 xname = dimname['time'] + ' [' + dimunits['time'] + ']'
550 else:
551 raise ValueError, 'bad treatment of time'
552
553 else:
554 raise ValueError, 'unrecognized plottype'
555
556
557
558
559
560
561 if plotkwds['xlabel'] != None:
562 xname = plotkwds['xlabel']
563 if plotkwds['ylabel'] != None:
564 yname = plotkwds['ylabel']
565
566 if plotkwds['title'] != None:
567 titlename = plotkwds['title']
568 else:
569 titlename = data_name + ' [' + data_units + ']'
570
571
572
573
574 pylab.clf()
575 pylab.figure(1)
576
577 if plottype == 'line':
578 y = data[rngs_idxs['time'],
579 rngs_idxs['lat'],
580 rngs_idxs['lon']]
581 pylab.plot(x, y)
582
583 elif plottype == 'contour':
584 ritim = rngs_idxs['time']
585 rilat = rngs_idxs['lat']
586 rilon = rngs_idxs['lon']
587
588
589
590
591
592 if N.size(rngs_idxs['time']) == 1:
593 zgrid = num.MLab.squeeze(data[ ritim[0],
594 rilat[0]:rilat[-1]+1,
595 rilon[0]:rilon[-1]+1 ])
596 elif N.size(rngs_idxs['lat']) == 1:
597 zgrid = num.MLab.squeeze(data[ ritim[0]:ritim[-1]+1,
598 rilat[0],
599 rilon[0]:rilon[-1]+1 ])
600 elif N.size(rngs_idxs['lon']) == 1:
601 zgrid = num.MLab.squeeze(data[ ritim[0]:ritim[-1]+1,
602 rilat[0]:rilat[-1]+1,
603 rilon[0] ])
604 else:
605 raise ValueError, 'unrecognized configuration'
606
607
608
609
610
611 if keys_rngs_sizes_gt_1.count('time') and \
612 keys_rngs_sizes_gt_1.count('lat'):
613 zgrid = N.transpose(zgrid)
614
615 xgrid, ygrid = pylab.meshgrid(x, y)
616
617
618
619
620 if plotkwds['levels'] == None:
621 levels = nice_levels(zgrid)
622 else:
623 levels = plotkwds['levels']
624
625
626
627
628
629 if keys_rngs_sizes_gt_1.count('lon') and \
630 keys_rngs_sizes_gt_1.count('lat'):
631 mapplot = Basemap(projection='cyl', resolution='l',
632 llcrnrlon=N.min(xgrid), llcrnrlat=N.min(ygrid),
633 urcrnrlon=N.max(xgrid), urcrnrlat=N.max(ygrid))
634 mapplot.drawcoastlines()
635 mapplot.drawmeridians(nice_levels(rngs['lon'],
636 approx_nlev=plotkwds['nlatlon']),
637 labels=[1,0,0,1])
638 mapplot.drawparallels(nice_levels(rngs['lat'],
639 approx_nlev=plotkwds['nlatlon']),
640 labels=[1,0,0,1])
641 if plotkwds['filled']:
642 plot = mapplot.contourf(xgrid, ygrid, zgrid, levels)
643 pylab.colorbar(plot, orientation='horizontal', format='%g')
644 else:
645 plot = mapplot.contour(xgrid, ygrid, zgrid, levels)
646 pylab.clabel(plot, inline=1, fontsize=10, fmt='%g')
647 else:
648 if plotkwds['filled']:
649 plot = pylab.contourf(xgrid, ygrid, zgrid, levels)
650 pylab.colorbar(plot, orientation='horizontal', format='%g')
651 else:
652 plot = pylab.contour(xgrid, ygrid, zgrid, levels)
653 pylab.clabel(plot, inline=1, fontsize=10, fmt='%g')
654
655 else:
656 raise ValueError, 'unrecognized plottype'
657
658
659
660
661
662
663 if keys_rngs_sizes_gt_1.count('lon') and \
664 keys_rngs_sizes_gt_1.count('lat'):
665 titlename = titlename + ' at ' \
666 + dimname['time'] + ' ' \
667 + str(rngs['time'][0]) + ' ' \
668 + dimunits['time']
669 titlename = mpl_latex_script1(titlename)
670 pylab.title(titlename)
671 else:
672 titlename = mpl_latex_script1(titlename)
673 xname = mpl_latex_script1(xname)
674 yname = mpl_latex_script1(yname)
675 pylab.xlabel(xname)
676 pylab.ylabel(yname)
677 pylab.title(titlename)
678
679
680
681
682
683
684
685
686 if plotkwds['fn'] == None:
687 if plotkwds['tmppreview'] and sys.platform == 'darwin':
688 outputfn = tempfile.mkstemp('.png','qtcm_')
689 pylab.savefig(outputfn[-1])
690 os.system('open -a /Applications/Preview.app '+outputfn[-1])
691 else:
692 pylab.show()
693
694 elif type(plotkwds['fn']) == type('a'):
695 pylab.savefig(plotkwds['fn'])
696 pylab.close(1)
697
698 else:
699 raise ValueError, 'cannot write to this type of file'
700
701
702
703
704
705
706 __test__ = { 'All positive values data':
707 """
708 >>> z = N.array([[24.5, 50.3],[283.1, 20.]])
709 >>> out = nice_levels(z)
710 >>> ['%g' % out[i] for i in range(len(out))]
711 ['0', '30', '60', '90', '120', '150', '180', '210', '240', '270', '300']
712 """,
713
714 'Exception for max_nlev':
715 """
716 >>> z = N.array([-24.5, 50.3, 183.1, 20.])
717 >>> out = nice_levels(z, approx_nlev=45)
718 Traceback (most recent call last):
719 ...
720 ValueError: max_nlev is too small
721 """,
722
723 'More mpl_latex_script1 examples':
724 """
725 >>> mpl_latex_script1('u_1 variance [m^2/s^2]')
726 'u$_1$ variance [m$^2$/s$^2$]'
727 """,
728
729 'Additional Examples of nice_levels':
730 """
731 >>> z = N.array([-24.5, 50.3, 183.1, 20.])
732 >>> out = nice_levels(z, approx_nlev=45, max_nlev=46)
733 >>> ['%g' % out[i] for i in range(len(out))]
734 ['-25', '-20', '-15', '-10', '-5', '0', '5', '10', '15', '20', '25', '30', '35', '40', '45', '50', '55', '60', '65', '70', '75', '80', '85', '90', '95', '100', '105', '110', '115', '120', '125', '130', '135', '140', '145', '150', '155', '160', '165', '170', '175', '180', '185']
735
736 >>> z = N.array([124.5, 150.3, 183.1, 220.])
737 >>> out = nice_levels(z, approx_nlev=20)
738 >>> ['%g' % out[i] for i in range(len(out))]
739 ['120', '125', '130', '135', '140', '145']
740 """ }
741
742
743
744
745
746
747
748
749 if __name__ == "__main__":
750 """Test the module.
751
752 Note: To help ensure that module testing of this file works, the
753 parent directory to the current directory is added to sys.path.
754 """
755 import doctest, sys, os
756 sys.path.append(os.pardir)
757 doctest.testmod(sys.modules[__name__])
758
759
760
761
762
763