Package PyDSTool :: Package PyCont :: Module Plotting
[hide private]
[frames] | no frames]

Source Code for Module PyDSTool.PyCont.Plotting

  1  """ Plotting class and function 
  2   
  3      Drew LaMar, May 2006 
  4  """ 
  5   
  6  from PyDSTool.common import args 
  7  from PyDSTool.matplotlib_import import * 
  8   
  9  # THESE ARE REPEATS FROM CONTINUATION!  MAKE SURE AND UPDATE!!! 
 10  all_point_types = ['P', 'RG', 'LP', 'BP', 'H', 'BT', 'ZH', 'CP', 'GH', 'DH', 'LPC', 'PD', 
 11                     'NS', 'MX', 'UZ', 'B'] 
 12  all_curve_types = ['EP', 'LP', 'H', 'FP', 'LC'] 
 13   
 14  ##### 
 15  _classes = ['pargs', 'KeyEvent'] 
 16   
 17  _functions = ['initializeDisplay'] 
 18   
 19  __all__ = _classes + _functions 
 20  ##### 
 21   
22 -class pargs(args):
23 - def __getitem__(self, key):
24 if len(key.split(':')) > 1: 25 (curve, pt) = key.split(':') 26 return args.__getitem__(args.__getitem__(self, curve), pt) 27 else: 28 return args.__getitem__(self, key)
29
30 - def __repr__(self, ct=0, res=None):
31 if res is None: 32 res = [] 33 34 # This is so I can make things alphabetical 35 all_pargs = [] 36 all_args = [] 37 deleted = self.has_key('deleted') and self.deleted 38 if not deleted: 39 for k, v in self.iteritems(): 40 if k in ['point', 'curve', 'cycle']: 41 if v[0].get_label() != '_nolegend_': 42 all_args.append(('Legend', v[0].get_label())) 43 elif k == 'text': 44 if v is not None: 45 all_args.append(('Label', v.get_text().lstrip(' '))) 46 elif k == 'type': 47 all_args.append(('Type', self['type'])) 48 elif isinstance(v, pargs): 49 all_pargs.append((k, v)) 50 51 if len(all_args) == 0 and len(all_pargs) == 0: 52 if self.has_key('deleted'): 53 res.append(ct*4*' ' + 'Deleted\n') 54 else: 55 res.append(ct*4*' ' + 'Empty\n') 56 else: 57 all_args.sort() 58 if len(all_args) > 0: 59 for n in all_args: 60 res.append(ct*4*' ' + n[0] + ': ' + n[1] + '\n') 61 62 all_pargs.sort() 63 if len(all_pargs) > 0: 64 ct += 1 65 for n in all_pargs: 66 res.append((ct-1)*4*' ' + n[0] + ':' + '\n') 67 n[1].__repr__(ct=ct, res=res) 68 69 if ct <= 1: 70 return reduce(lambda x, y: x+y, res)
71 72 __str__ = __repr__ 73
74 - def fromLabel(self, label, key=None):
75 point = None 76 if self.has_key('text') and self.text.get_text().lstrip(' ') == label or \ 77 self.has_key('point') and self.point[0].get_label() == label or \ 78 self.has_key('curve') and self.curve[0].get_label() == label: 79 return (key, self) 80 for k, v in self.iteritems(): 81 if isinstance(v, pargs): 82 point = v.fromLabel(label, key=k) 83 if point is not None: 84 break 85 return point
86
87 - def get(self, objtype, bylabel=None, byname=None, bytype=None, bylegend=None, obj=None, ct=0):
88 if objtype not in ['point', 'text', 'cycle', 'curve']: 89 raise 'Object type must be point, text, cycle, or curve' 90 91 ct += 1 92 if isinstance(bylabel, str): 93 bylabel = [bylabel] 94 if isinstance(byname, str): 95 byname = [byname] 96 if isinstance(bytype, str): 97 bytype = [bytype] 98 if isinstance(bylegend, str): 99 bylegend = [bylegend] 100 101 if obj is None: 102 obj = [] 103 if bylabel is None and byname is None and bytype is None and bylegend is None: 104 if objtype in ['point', 'text', 'cycle']: 105 bytype = all_point_types 106 elif objtype == 'curve': 107 bytype = all_curve_types 108 for k, v in self.iteritems(): 109 if isinstance(v, pargs): 110 if v.has_key(objtype): 111 if bylabel is not None and v.has_key('text') and v.text.get_text().lstrip(' ') in bylabel or \ 112 byname is not None and k in byname or \ 113 bytype is not None and k.strip('0123456789') in bytype or \ 114 bylegend is not None and ((v.has_key('curve') and v.curve[0].get_label() in bylegend) or (v.has_key('cycle') and v.cycle[0].get_label() in bylegend)): 115 obj.append((k,v[objtype])) 116 v.get(objtype, bylabel=bylabel, byname=byname, bytype=bytype, bylegend=bylegend, obj=obj, ct=ct) 117 118 if ct == 1: 119 return obj
120
121 - def toggleLabel(self, visible='on', refresh=True):
122 if self.has_key('text'): 123 self.text.set_visible(visible == 'on') 124 125 if refresh: 126 self.refresh()
127
128 - def toggleLabels(self, visible='on', bylabel=None, byname=None, bytype=None, ct=0):
129 ct += 1 130 if isinstance(bylabel, str): 131 bylabel = [bylabel] 132 if isinstance(byname, str): 133 byname = [byname] 134 if isinstance(bytype, str): 135 bytype = [bytype] 136 137 if bylabel is None and byname is None and bytype is None: 138 bytype = all_point_types # Turn off all points 139 for k, v in self.iteritems(): 140 if bytype is not None and k.strip('0123456789') in bytype or \ 141 byname is not None and k in byname or \ 142 bylabel is not None and isinstance(v, pargs) and v.has_key('text') and v.text.get_text().lstrip(' ') in bylabel: 143 v.toggleLabel(visible=visible, refresh=False) 144 elif isinstance(v, pargs): 145 v.toggleLabels(visible=visible, bylabel=bylabel, byname=byname, bytype=bytype, ct=ct) 146 147 if ct == 1: 148 self.refresh()
149
150 - def togglePoint(self, visible='on', refresh=True):
151 if self.has_key('point'): 152 self.point[0].set_visible(visible == 'on') 153 154 if refresh: 155 self.refresh()
156
157 - def togglePoints(self, visible='on', bylabel=None, byname=None, bytype=None, ct=0):
158 ct += 1 159 if isinstance(bylabel, str): 160 bylabel = [bylabel] 161 if isinstance(byname, str): 162 byname = [byname] 163 if isinstance(bytype, str): 164 bytype = [bytype] 165 166 if bylabel is None and byname is None and bytype is None: 167 bytype = all_point_types # Turn off all points 168 for k, v in self.iteritems(): 169 if bytype is not None and k.strip('0123456789') in bytype or \ 170 byname is not None and k in byname or \ 171 bylabel is not None and isinstance(v, pargs) and v.has_key('text') and v.text.get_text().lstrip(' ') in bylabel: 172 v.togglePoint(visible=visible, refresh=False) 173 elif isinstance(v, pargs): 174 v.togglePoints(visible=visible, bylabel=bylabel, byname=byname, bytype=bytype, ct=ct) 175 176 if ct == 1: 177 self.refresh()
178
179 - def toggleCurve(self, visible='on', refresh=True):
180 if self.has_key('curve'): 181 for k, v in self.iteritems(): 182 if isinstance(v, pargs): 183 v.toggleLabel(visible=visible, refresh=False) 184 v.togglePoint(visible=visible, refresh=False) 185 for line in self.curve: 186 line.set_visible(visible == 'on') 187 188 if refresh: 189 self.refresh()
190
191 - def toggleCurves(self, visible='on', bylegend=None, byname=None, bytype=None, ct=0):
192 ct += 1 193 if isinstance(bylegend, str): 194 bylegend = [bylegend] 195 if isinstance(byname, str): 196 byname = [byname] 197 if isinstance(bytype, str): 198 bytype = [bytype] 199 200 if bylegend is None and byname is None and bytype is None: 201 bytype = all_curve_types 202 for k, v in self.iteritems(): 203 if bytype is not None and isinstance(v, pargs) and v.has_key('type') and v.type in bytype or \ 204 byname is not None and k in byname or \ 205 bylegend is not None and isinstance(v, pargs) and v.has_key('curve') and v.curve[0].get_label() in bylegend: 206 v.toggleCurve(visible=visible, refresh=False) 207 elif isinstance(v, pargs): 208 v.toggleCurves(visible=visible, bylegend=bylegend, byname=byname, bytype=bytype, ct=ct) 209 210 if ct == 1: 211 self.refresh()
212
213 - def toggleCycle(self, visible='on', refresh=True):
214 if self.has_key('cycle'): 215 for line in self.cycle: 216 line.set_visible(visible == 'on') 217 218 if refresh: 219 self.refresh()
220
221 - def toggleCycles(self, visible='on', bylegend=None, byname=None, bytype=None, ct=0):
222 ct += 1 223 if isinstance(bylegend, str): 224 bylegend = [bylegend] 225 if isinstance(byname, str): 226 byname = [byname] 227 if isinstance(bytype, str): 228 bytype = [bytype] 229 230 if bylegend is None and byname is None and bytype is None: 231 bytype = all_point_types 232 for k, v in self.iteritems(): 233 if bytype is not None and k.strip('0123456789') in bytype or \ 234 byname is not None and k in byname or \ 235 bylegend is not None and isinstance(v, pargs) and v.has_key('cycle') and v.cycle[0].get_label() in bylegend: 236 v.toggleCycle(visible=visible, refresh=False) 237 elif isinstance(v, pargs): 238 v.toggleCycles(visible=visible, bylegend=bylegend, byname=byname, bytype=bytype, ct=ct) 239 240 if ct == 1: 241 self.refresh()
242
243 - def toggleAll(self, visible='on', bylabel=None, byname=None, bytype=None, ct=0):
244 ct += 1 245 if isinstance(bylabel, str): 246 bylabel = [bylabel] 247 if isinstance(byname, str): 248 byname = [byname] 249 if isinstance(bytype, str): 250 bytype = [bytype] 251 252 if bylabel is None and byname is None and bytype is None: 253 bytype = all_point_types 254 for k, v in self.iteritems(): 255 if bytype is not None and k.strip('0123456789') in bytype or \ 256 byname is not None and k in byname or \ 257 bylabel is not None and isinstance(v, pargs) and ((v.has_key('text') and v.text.get_text().lstrip(' ') in bylabel) or (v.has_key('curve') and v.curve[0].get_label() in bylabel) or (v.has_key('cycle') and v.cycle[0].get_label() in bylabel)): 258 v.toggleLabel(visible=visible, refresh=False) 259 v.togglePoint(visible=visible, refresh=False) 260 v.toggleCycle(visible=visible, refresh=False) 261 elif isinstance(v, pargs): 262 v.toggleAll(visible=visible, bylabel=bylabel, byname=byname, bytype=bytype, ct=ct) 263 264 if ct == 1: 265 self.refresh()
266
267 - def setLabel(self, label, refresh=True):
268 if self.has_key('text'): 269 self.text.set_text(' '+label) 270 271 if refresh: 272 self.refresh()
273
274 - def setLabels(self, label, bylabel=None, byname=None, bytype=None, ct=0):
275 ct += 1 276 if isinstance(bylabel, str): 277 bylabel = [bylabel] 278 if isinstance(byname, str): 279 byname = [byname] 280 if isinstance(bytype, str): 281 bytype = [bytype] 282 283 if bylabel is None and byname is None and bytype is None: 284 bytype = all_point_types # Turn off all points 285 for k, v in self.iteritems(): 286 if bytype is not None and k.strip('0123456789') in bytype or \ 287 byname is not None and k in byname or \ 288 bylabel is not None and isinstance(v, pargs) and v.has_key('text') and v.text.get_text().lstrip(' ') in bylabel: 289 v.setLabel(label, refresh=False) 290 elif isinstance(v, pargs): 291 v.setLabels(label, bylabel=bylabel, byname=byname, bytype=bytype, ct=ct) 292 293 if ct == 1: 294 self.refresh()
295
296 - def setLegend(self, legend, refresh=True):
297 if self.has_key('curve'): 298 for piece in self.curve: 299 piece.set_label(legend) 300 elif self.has_key('cycle'): 301 for piece in self.cycle: 302 piece.set_label(legend) 303 304 if refresh: 305 self.refresh()
306
307 - def setLegends(self, legend, bylegend=None, byname=None, bytype=None, ct=0):
308 ct += 1 309 if isinstance(bylegend, str): 310 bylabel = [bylegend] 311 if isinstance(byname, str): 312 byname = [byname] 313 if isinstance(bytype, str): 314 bytype = [bytype] 315 316 if bylegend is None and byname is None and bytype is None: 317 if v.has_key('curve'): 318 bytype = all_curve_types 319 else: 320 bytype = all_point_types # Turn off all points 321 for k, v in self.iteritems(): 322 if bytype is not None and isinstance(v, pargs) and ((v.has_key('curve') and v.type in bytype) or (v.has_key('cycle') and k.strip('0123456789') in bytype)) or \ 323 byname is not None and k in byname or \ 324 bylegend is not None and isinstance(v, pargs) and ((v.has_key('curve') and v.curve[0].get_label() in bylegend) or (v.has_key('cycle') and v.cycle[0].get_label() in bylegend)): 325 v.setLegend(legend, refresh=False) 326 elif isinstance(v, pargs): 327 v.setLegends(legend, bylegend=bylegend, byname=byname, bytype=bytype, ct=ct) 328 329 if ct == 1: 330 self.refresh()
331
332 - def refresh(self):
333 fig = None # Used to save current figure (reset when done) 334 335 if not self.has_key('deleted'): 336 if 'fig' in self.keys(): 337 plt.figure(self.fig.number) 338 plt.draw() 339 elif 'axes' in self.keys(): 340 plt.figure(self.axes.figure.number) 341 #Would like to say: plt.axes(self.axes) Doesn't work, though. 342 plt.draw() 343 elif 'curve' in self.keys(): 344 plt.figure(self.curve[0].figure.number) 345 plt.draw() 346 elif 'point' in self.keys(): 347 plt.figure(self.point[0].figure.number) 348 plt.draw() 349 elif 'cycle' in self.keys(): 350 plt.figure(self.cycle[0].figure.number) 351 plt.draw() 352 else: # Only here in plot structure 353 fig = plt.gcf() 354 for k, v in self.iteritems(): 355 if isinstance(v, pargs): 356 v.refresh() 357 358 if fig is not None: # Reset to current figure 359 plt.figure(fig.number)
360
361 - def clean(self):
362 """Cleans plotting structure (e.g. deleted) Also gets rid of nonexisting figures 363 if they've been deleted by an outside source. This method is a little clumsy since 364 it has to spawn a dummy plot if a plot has been deleted, but who cares, it works.""" 365 deleted = [] 366 for k, v in self.iteritems(): 367 recursive_clean = True 368 if isinstance(v, pargs): 369 if v.has_key('deleted'): 370 deleted.append(k) 371 recursive_clean = False 372 elif v.has_key('fig'): 373 fig_check = plt.figure(v.fig.number) 374 if fig_check != v.fig: 375 plt.close(fig_check) 376 deleted.append(k) 377 recursive_clean = False 378 elif v.has_key('axes'): 379 try: 380 fig_axes = plt.axes(v.axes) 381 except: 382 fig_axes = None 383 384 if fig_axes is None: 385 deleted.append(k) 386 recursive_clean = False 387 elif fig_axes not in v.axes.figure.axes: 388 print 'Warning: Axes were deleted without using plt.delaxes().' 389 v.axes.figure.axes.append(fig_axes) 390 plt.draw() 391 elif v.has_key('curve'): 392 for piece in v.curve: 393 if piece not in piece.axes.lines: 394 deleted.append(k) 395 v.delete() # Remove remaining pieces of curve in pyplot 396 recursive_clean = False 397 break 398 elif v.has_key('cycle'): 399 for piece in v.cycle: 400 if piece not in piece.axes.lines: 401 deleted.append(k) 402 v.delete() 403 recursive_clean = False 404 break 405 elif v.has_key('point'): 406 if v.point[0] not in v.point[0].axes.lines or v.text not in v.text.axes.texts: 407 deleted.append(k) 408 if v.text in v.text.axes.texts: 409 v.text.axes.texts.remove(v.text) 410 if v.point[0] in v.point[0].axes.lines: 411 v.point[0].axes.lines.remove(v.point[0]) 412 recursive_clean = False 413 414 if recursive_clean: 415 v.clean() 416 417 for k in deleted: 418 self.pop(k)
419
420 - def clear(self, refresh=True):
421 if not self.has_key('deleted'): 422 if self.has_key('fig'): 423 plt.figure(self.fig.number) 424 plt.clf() 425 remove = [k for k in self.keys() if k != 'fig'] 426 elif self.has_key('axes'): 427 title = self.axes.title.get_text() 428 self.axes.clear() 429 plt.axes(self.axes) 430 plt.title(title) 431 remove = [k for k in self.keys() if k != 'axes'] 432 if refresh: 433 self.refresh() 434 elif self.has_key('curve'): 435 remove = [k for k in self.keys() if k != 'curve'] 436 for k in remove: 437 if isinstance(self[k], pargs): 438 self[k].clear(refresh=False) 439 if refresh: 440 self.refresh() 441 elif self.has_key('point'): 442 if self.point[0] in self.point[0].axes.lines: 443 self.point[0].axes.lines.remove(self.point[0]) 444 if self.text in self.point[0].axes.texts: 445 self.point[0].axes.texts.remove(self.text) 446 remove = [k for k in self.keys() if k != 'point'] 447 if refresh: 448 self.refresh() 449 self.deleted = True 450 else: 451 remove = [] 452 453 if remove != []: 454 for k in remove: 455 self.pop(k) 456 else: 457 print 'Object is deleted.'
458
459 - def clearall(self):
460 for v in self.itervalues(): 461 if isinstance(v, pargs): 462 v.clear(refresh=False) 463 self.refresh()
464
465 - def delete(self, refresh=True):
466 if not self.has_key('deleted'): 467 if self.has_key('fig'): 468 self.clear(refresh=False) 469 plt.close() 470 self.deleted = True 471 elif self.has_key('axes'): 472 self.clear(refresh=False) 473 plt.delaxes() 474 self.deleted = True 475 elif self.has_key('curve'): 476 self.clear(refresh=False) 477 for curve in self.curve: 478 if curve in self.curve[0].axes.lines: 479 self.curve[0].axes.lines.remove(curve) 480 if refresh: 481 self.refresh() 482 self.deleted = True 483 elif self.has_key('cycle'): 484 self.clear(refresh=False) 485 for cycle in self.cycle: 486 if cycle in self.cycle[0].axes.lines: 487 self.cycle[0].axes.lines.remove(cycle) 488 if refresh: 489 self.refresh() 490 self.deleted = True 491 elif self.has_key('point'): 492 self.clear(refresh=refresh) 493 else: 494 for v in self.itervalues(): 495 if isinstance(v, pargs): 496 v.delete(refresh=refresh) 497 else: 498 print 'Object is already deleted.'
499
500 - def deleteall(self):
501 for v in self.itervalues(): 502 if isinstance(v, pargs): 503 v.delete(refresh=False) 504 self.refresh()
505
506 -class KeyEvent(object):
507 """Used in 'highlight' method of plot_cycles."""
508 - def __init__(self, paxes):
509 self.curr = 0 510 self.axes = paxes.axes 511 512 # Get list of cycles from pargs class axes 513 cycles = [thing.cycle[0] for k, thing in paxes.iteritems() if isinstance(thing, pargs) and thing.has_key('cycle')] 514 # Put in order as in axes.lines (this should be same as order along curve, or user 515 # specified in cycles when plot_cycles was called) 516 self.cycles = [] 517 for line in self.axes.lines: 518 if line in cycles: 519 self.cycles.append(line) 520 self.axes.set_title(self.cycles[0].get_label()) 521 522 self.bgd_lw = 0.2 523 self.fgd_lw = 1.5 524 for line in self.cycles: 525 line.set(linewidth=self.bgd_lw) 526 self.cycles[0].set(linewidth=self.fgd_lw) 527 plt.draw() 528 plt.connect('key_press_event', self.__call__)
529
530 - def __call__(self, event):
531 if event.key == 'right': self.change_curr(1) 532 elif event.key == 'left': self.change_curr(-1) 533 elif event.key == 'up': self.change_bgd(1) 534 elif event.key == 'down': self.change_bgd(-1) 535 elif event.key in ('+', '='): self.change_fgd(1) 536 elif event.key in ('-', '_'): self.change_fgd(-1)
537
538 - def change_bgd(self, up):
539 if up == 1: 540 if self.bgd_lw+0.1 > self.fgd_lw: 541 self.bgd_lw = self.fgd_lw 542 else: 543 self.bgd_lw += 0.1 544 else: 545 if self.bgd_lw-0.1 < 0: 546 self.bgd_lw = 0 547 else: 548 self.bgd_lw -= 0.1 549 550 for ct, line in enumerate(self.cycles): 551 if ct != self.curr: 552 line.set(linewidth=self.bgd_lw) 553 plt.draw() 554 555 print 'Background linewidth = %f' % self.bgd_lw
556
557 - def change_fgd(self, up):
558 if up == 1: 559 self.fgd_lw += 0.1 560 else: 561 if self.fgd_lw-0.1 < self.bgd_lw: 562 self.fgd_lw = self.bgd_lw 563 else: 564 self.fgd_lw -= 0.1 565 566 self.cycles[self.curr].set(linewidth=self.fgd_lw) 567 plt.draw() 568 569 print 'Foreground linewidth = %f' % self.fgd_lw
570
571 - def change_curr(self, up):
572 self.cycles[self.curr].set(linewidth=self.bgd_lw) 573 if up == 1: 574 self.curr = (self.curr+1) % len(self.cycles) 575 else: 576 self.curr = (self.curr-1) % len(self.cycles) 577 self.cycles[self.curr].set(linewidth=self.fgd_lw) 578 self.axes.set_title(self.cycles[self.curr].get_label()) 579 plt.draw()
580
581 -def initializeDisplay(plot, figure=None, axes=None):
582 """If figure = 'new', then it will create a new figure with name fig# 583 If figure is an integer, that figure # will be selected 584 """ 585 586 plot.clean() # Clean up plot structure 587 plot._cfl = None 588 plot._cal = None 589 590 # Handle figure 591 if figure is None: 592 if len(plot) <= 3: 593 figure = plt.gcf() 594 else: 595 raise ValueError('Please specify a figure.') 596 597 cfl = None 598 if isinstance(figure, plt.Figure): 599 for k, v in plot.iteritems(): 600 if isinstance(v, pargs) and v.fig == figure: 601 cfl = k 602 break 603 elif isinstance(figure, str): 604 if figure == 'new': 605 figure = plt.figure() 606 else: 607 cfl = figure 608 if cfl not in plot.keys(): 609 plot[cfl] = pargs() 610 plot[cfl].fig = plt.figure() 611 elif isinstance(figure, int): 612 fighandle = plt.figure(figure) 613 cfl = 'fig' + str(figure) 614 if cfl not in plot.keys(): 615 plot[cfl] = pargs() 616 plot[cfl].fig = fighandle 617 else: 618 raise TypeError("Invalid type for figure argument") 619 620 if cfl is None: 621 cfl = 'fig1' 622 ct = 1 623 while cfl in plot.keys(): 624 ct += 1 625 cfl = 'fig'+repr(ct) 626 plot[cfl] = pargs() 627 plot[cfl].fig = figure 628 629 plt.figure(plot[cfl].fig.number) 630 631 # Handle axes 632 if not axes: 633 if len(plot[cfl]) <= 2: 634 axes = plt.gca() 635 else: 636 raise 'Please specify axes.' 637 elif isinstance(axes, tuple): 638 if len(axes) == 3: 639 axes = plt.subplot(axes[0],axes[1],axes[2]) 640 else: 641 raise 'Tuple must be of length 3' 642 643 cal = None 644 if isinstance(axes, plt.Axes): 645 for k, v in plot[cfl].iteritems(): 646 if isinstance(v, pargs) and v.axes == axes: 647 cal = k 648 break 649 650 if cal is None: 651 cal = 'axes1' 652 ct = 1 653 while cal in plot[cfl].keys(): 654 ct += 1 655 cal = 'axes'+repr(ct) 656 plot[cfl][cal] = pargs() 657 plot[cfl][cal].axes = axes 658 plot[cfl][cal].axes.set_title(cal) 659 elif isinstance(axes, str): 660 cal = axes 661 if cal not in plot[cfl].keys(): 662 plot[cfl][cal] = pargs() 663 plot[cfl][cal].axes = plt.axes() 664 plot[cfl][cal].axes.set_title(cal) 665 666 plt.axes(plot[cfl][cal].axes) 667 668 plot._cfl = cfl 669 plot._cal = cal
670