Package qtcm :: Module plot
[hide private]
[frames] | no frames]

Source Code for Module qtcm.plot

  1  #!/usr/bin/python -tt 
  2  #======================================================================= 
  3  #                        General Documentation 
  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  #                       Additional Documentation 
 15  # 
 16  # RCS Revision Code: 
 17  #   $Id: plot.py 4 2008-06-25 01:03:28Z jlin $ 
 18  # 
 19  # Modification History: 
 20  # - 29 May 2008:  Original by Johnny Lin, Physics Department, North 
 21  #   Park University.  Passed passably reasonable tests. 
 22  # 
 23  # Notes: 
 24  # - Written for Python 2.4. 
 25  # - Module docstrings can be tested using the doctest module.  To 
 26  #   test, execute "python plot.py". 
 27  # - See import statements throughout for non-"built-in" packages and 
 28  #   modules required. 
 29  # 
 30  # Copyright (c) 2008 by Johnny Lin.  For licensing, distribution  
 31  # conditions, contact information, and additional documentation see 
 32  # the URL http://www.johnny-lin.com/py_pkgs/qtcm/doc/. 
 33  #======================================================================= 
 34   
 35   
 36   
 37   
 38  #---------------- Module General Import and Declarations --------------- 
 39   
 40  #- Import os and sys.  Also, if you're importing this module in 
 41  #  testing mode, or you're running pydoc on this module via the  
 42  #  command line, import user-specific settings to make sure any  
 43  #  non-standard libraries are found: 
 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  #- Import package version and set module version to package version: 
 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  #- Import numpy/Numeric/numarray as appropriate: 
 61   
 62  import num_settings as num 
 63  from num_settings import N 
 64   
 65   
 66  #- Import matplotlib (must be in this order): 
 67   
 68  from matplotlib import rc as matplotlibrc 
 69  #matplotlibrc('font',**{'family':'sans-serif','sans-serif':['Helvetica']}) 
 70  matplotlibrc('text', usetex=True) 
 71  import pylab 
 72  import matplotlib 
 73  from matplotlib.toolkits.basemap import Basemap 
 74   
 75   
 76  #- Other imports that then are used as module variables: 
 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  #-------------------- Module Function:  nice_levels -------------------- 
 88   
89 -def nice_levels(data, approx_nlev=10, max_nlev=28):
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 #- Default settings and error check: 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 #- See if the cint can be "even". If not, return alternative 151 # contour levels vector: 152 153 #+ Guess a possible cint. Look for an "even" value that is 154 # closest to that: 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 #+ Look around at the "even" values nearby the possible option 178 # for cint. Calculate how many contours each of those cints 179 # would give. Dictionary ncon_count is the number of 180 # contours for a given test cint. test_tcidxs are the indices 181 # in possiblecint to examine in detail; these index values 182 # will be the keys in ncon_count. 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 #+ Select the cint that has the fewest levels if it has at 207 # least nlevels-1. Otherwise, try to find the next cint with 208 # the fewest levels that is below max_nlev. tempcint is what 209 # you get (changed, if warranted) leaving this section: 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 #+ Create levels for case with neg. and pos. contours. There 233 # is the case of all pos., all neg., and mixed pos. and neg. 234 # contours: 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 #+ Sort clevels, reset number of levels, maximum, minimum, 258 # and interval of contour based on the automatic setting: 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 #- Return output: 272 273 return N.arange(tmpcmin, tmpcmax+tmpcint, tmpcint)
274 275 276 277 278 #----------------- Module Function: mpl_latex_script1 ----------------- 279
280 -def mpl_latex_script1(instring):
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 #- Make input string be a list of individual characters: 303 304 instrlist = list(instring) 305 306 307 #- Cycle through each ^ and _, change accordingly, and return: 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 #----------------- Module Function: plot_ncdf_output ------------------ 322
323 -def plot_ncdf_output(id, datafn, **kwds):
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 #- Accomodate other ids. The id that this routine will use 357 # (i.e., iduse corresponds to the name in the netCDF output file) 358 # is called iduse. Set this to id, except for the case where 359 # some aliases of ids are entered in which iduse is the alias: 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 #- Set defined keyword defaults. All are set to None except for 368 # nlatlon which gets an integer: 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 #- Get data and dimensions of iduse to plot: 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 #- Alter data long name to remove any units. The definition 412 # of units as the substring within the [] is the same as used in 413 # defVar in output.F90 of the compiled QTCM model. Remove 414 # underscores and extra whitespace in data_name and data_units, 415 # replacing with a single whitespace character between words: 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 #- Alter dimension long name to remove any units. The definition 428 # of units as the substring within the [] is the same as used in 429 # defVar in output.F90 of the compiled QTCM model. Remove 430 # underscores and extra whitespace in name and units, replacing 431 # with a single whitespace character between words, and 432 # capitalizing like a title: 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 #- Some data checks: 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 #- Choose and describe ranges for lat, lon, and time. The 467 # section cycles through the dictionary of dimensions. idim is 468 # the 1-D array of the values of that dimension. rngs is a 469 # dictionary where each entry corresponds to a dimension, and the 470 # value of the entry is the values of that dimension that are to 471 # be plotted. rngs_idxs are the indices in the original 472 # dimensions array corresponding to the values in rngs. 473 # keys_rngs_sizes_gt_1 is a list of the keys of ranges that have 474 # sizes greater than 1: 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 #- Set plot types (line or contour): 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 #- Set plot axis fields and axis names, depending on what sort 522 # of dimensions will be plotted. If lon is to be plotted, it is 523 # always the x-axis. If lat is to be plotted, it is always the 524 # y-axis. In this section and later on in a few places, I 525 # rely on the count method for a list as a Boolean test: If it 526 # returns 0, consider that False; > 0 is True. The text for the 527 # title and axis labels to be passed to the xlabel, etc. methods, 528 # are called titlename, xname, and yname: 529 530 if plottype == 'line': #+ Choose x-axis vector and 531 x = rngs[keys_rngs_sizes_gt_1[0]] # x/y names for line plot 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': #+ Choose axis vectors and 537 if keys_rngs_sizes_gt_1.count('lon'): # names for contour plot 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 #- Override xname, yname, and titlename with keywords, if they 558 # are not None. titlename receives data_name and data_units 559 # by default: 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 #- Pick data to be plotted and plot: 573 574 pylab.clf() #+ Clear any previous figures 575 pylab.figure(1) #+ Open a pylab figure 576 577 if plottype == 'line': #+ Select data for a line plot 578 y = data[rngs_idxs['time'], # and plot 579 rngs_idxs['lat'], 580 rngs_idxs['lon']] 581 pylab.plot(x, y) 582 583 elif plottype == 'contour': #+ Select data for a contour 584 ritim = rngs_idxs['time'] # plot and plot 585 rilat = rngs_idxs['lat'] 586 rilon = rngs_idxs['lon'] 587 588 589 #* Extract subarrays depending on which two dimensions are 590 # chosen: 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 #* Change zgrid for special case of a lat. vs. time contour 609 # plot. Calculate xgrid and ygrid: 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 #* Set contour levels: 619 620 if plotkwds['levels'] == None: 621 levels = nice_levels(zgrid) 622 else: 623 levels = plotkwds['levels'] 624 625 626 #- Plot (creating continents first if is a lat vs. lon plot) 627 # and write contour levels/color bar as appropriate: 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 #- Add titling. Lat vs. lon plots do not have axis labels because 660 # the map labels already make it clear, and for those plots the 661 # title also includes the time value: 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 #- Output plot to PNG file or screen. The show command seems to 681 # have a problem on my Mac OS X, so save to a temporary file 682 # and use preview to view for fn == None and tmppreview set to 683 # True. Note that the temporary file is not deleted by this 684 # method: 685 686 if plotkwds['fn'] == None: #+ Screen display 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'): #+ Write to file 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 #--------- Addition doctest Examples as Private Module Variable -------- 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 #-------------------------- Main: Test Module ------------------------- 746 747 #- Execute doctest if module is run from command line: 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 # ===== end file ===== 763