posts/plotpick/plotpick.py

   1 #!/usr/bin/env python
   2
   3 """Pick data-points with a cursor.
   4
   5 Reads in ASCII data from datafile, or, if datafile is not given, from
   6 stdin.  Data files define an array of (x,y) points and consist of two
   7 columns (x&y) seperated by TABs, with ENDLINEs between the points:
   8
   9 x0      y0
  10 x1      y1
  11 ...
  12
  13 The points are plotted in a window, and mousing over the plot displays
  14 a cursor that snaps to the nearest point.  Right-clicking (button 3)
  15 will print the TAB seperated coordinates (x, y, point-index) to stdout.
  16 """
  17
  18 import matplotlib.pyplot as _pyplot
  19 import numpy as _numpy
  20
  21
  22 __version__ = '0.3'
  23
  24
  25 class Picker (object):
  26     """Pick points from a plot
  27     """
  28     def __init__(self, axes, x=None, y=None, selected=None, highlight=False):
  29         self.axes = axes
  30         if highlight:
  31             if selected is None:
  32                 selected = []
  33             sx = [x[i] for i in selected]
  34             sy = [y[i] for i in selected]
  35             hold = self.axes.ishold()
  36             axes.hold(True)
  37             self.highlight_line, = axes.plot(sx, sy, 'r.', zorder=5)
  38             self.axes.hold(hold)
  39         self.selected = selected
  40         self.highlight = highlight
  41         self.txt = self.axes.title
  42
  43     def onpick(self, event):
  44         thisline = event.artist
  45         xdata = thisline.get_xdata()
  46         ydata = thisline.get_ydata()
  47         for i in event.ind:
  48             if self.highlight:
  49                 if i in self.selected:
  50                     self.selected.remove(i)
  51                 else:
  52                     self.selected.append(i)
  53                 sx = [xdata[i] for i in self.selected]
  54                 sy = [ydata[i] for i in self.selected]
  55                 self.highlight_line.set_data(sx, sy)
  56                 _pyplot.draw()
  57             else:
  58                 print('{}\t{}\t{}'.format(xdata[i], ydata[i], i))
  59             break  # only deal with the first picked point
  60
  61
  62 if __name__ == "__main__" :
  63     from argparse import ArgumentParser as _ArgumentParser
  64     import sys as _sys
  65
  66     parser = _ArgumentParser(
  67         description=__doc__, version=__version__)
  68     parser.add_argument(
  69         '--highlight', action='store_const', const=True, default=False,
  70         help=('Highlight selected points and print them afterwards, instead '
  71               'of printing them as they are clicked'))
  72     parser.add_argument(
  73         'datafile', nargs='?', default=None,
  74         help='Path to the datafile.  Defaults to stdin')
  75
  76     args = parser.parse_args()
  77
  78     if args.datafile:
  79         data = _numpy.genfromtxt(args.datafile)
  80     else:
  81         data = _numpy.genfromtxt(_sys.stdin)
  82     x = data[:,0].squeeze()
  83     y = data[:,1].squeeze()
  84
  85     figure = _pyplot.figure()
  86     axes = figure.add_subplot(1, 1, 1)
  87     axes.plot(x, y, '.', picker=2)
  88     axes.autoscale(tight=True)
  89     picker = Picker(axes=axes, highlight=args.highlight)
  90     figure.canvas.mpl_connect('pick_event', picker.onpick)
  91
  92     print('#x\ty\tindex')
  93     _pyplot.show()
  94     if args.highlight:
  95         for i in picker.selected:
  96             point = (x[i], y[i], i)
  97             print('\t'.join(str(x) for x in point))