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
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
29
31 if res is None:
32 res = []
33
34
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
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
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
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
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
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
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
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):
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
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
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
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
333 fig = None
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
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:
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:
359 plt.figure(fig.number)
360
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()
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
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
505
507 """Used in 'highlight' method of plot_cycles."""
509 self.curr = 0
510 self.axes = paxes.axes
511
512
513 cycles = [thing.cycle[0] for k, thing in paxes.iteritems() if isinstance(thing, pargs) and thing.has_key('cycle')]
514
515
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
537
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
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
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
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()
587 plot._cfl = None
588 plot._cal = None
589
590
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
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