def click_event(self, event):

'''Extracts locations from the user interaction

# stash the last event for debugging!

self.ev = event

# if shift is down, clear and bail

if event.key == 'shift':

self.clear()

return

# if no x or y data, bail

if event.xdata is None or event.ydata is None:

return

# if not in our Axes, bail

if event.inaxes is not self.ax:

return

# if left-click, append to points list

if event.button == 1:

self.pt_lst.append((event.xdata, event.ydata))

# if right-click, remove the closest point

elif event.button == 3:

self.remove_pt((event.xdata, event.ydata))

# re-draw (if needed)

self.redraw()

def remove_pt(self, loc):

"""Remove the nearest point.

if len(self.pt_lst) > 0:

self.pt_lst.pop(np.argmin(list(map(lambda x:

np.sqrt((x[0] - loc[0]) ** 2 +

(x[1] - loc[1]) ** 2),

self.pt_lst))))

def redraw(self):

"""Redraw the canvas given the current set of points

# get the current selected points

if len(self.pt_lst) > 0:

x, y = zip(*self.pt_lst)

else:

x, y = [], []

# and update the Line2D with the

self.pt_plot.set_xdata(x)

self.pt_plot.set_ydata(y)

# if we have more than 5 points, create a best-fit closed spline

if len(self.pt_lst) > 5:

SC = SplineCurve.from_pts(self.pt_lst, pix_err=self.pix_err)

new_pts = SC.q_phi_to_xy(0, np.linspace(0, 2 * np.pi, 1000))

center = SC.cntr

self.pt_lst.sort(key=lambda x:

np.arctan2(x[1] - center[1], x[0] - center[0]))

else:

new_pts = ([], [])

# and update the data in the spline Line2D objcet

self.sp_plot.set_xdata(new_pts[0])

self.sp_plot.set_ydata(new_pts[1])

self.canvas.draw_idle()

def __init__(self, ax, pix_err=1):

self.canvas = ax.get_figure().canvas

self.ax = ax

self.cid = None

self.pt_lst = []

self.pt_plot = ax.plot([], [], marker='o',

linestyle='none', zorder=5)[0]

self.sp_plot = ax.plot([], [], lw=3, color='r')[0]

self.pix_err = pix_err

self.connect_sf()

def clear(self):

'''Clears the points'''

self.pt_lst = []

self.redraw()

def connect_sf(self):

if self.cid is None:

self.cid = self.canvas.mpl_connect('button_press_event',

self.click_event)

def disconnect_sf(self):

if self.cid is not None:

self.canvas.mpl_disconnect(self.cid)

self.cid = None

@property

def points(self):

'''Returns the clicked points in the format the rest of the

return np.vstack(self.pt_lst).T

@property

def SplineCurve(self):

curve = SplineCurve.from_pts(self.pt_lst, pix_err=self.pix_err)

'''

@classmethod

def _get_spline(cls, points, pix_err=2, need_sort=True, **kwargs):

'''

if type(points) is np.ndarray:

# make into a list

pt_lst = zip(*points)

# get center

center = np.mean(points, axis=1).reshape(2, 1)

else:

# make a copy of the list

pt_lst = list(points)

# compute center

tmp_fun = lambda x, y: (x[0] + y[0], x[1] + y[1])

center = np.array(reduce(tmp_fun, pt_lst)).reshape(2, 1)

center /= len(pt_lst)

if len(pt_lst) < 5:

raise TooFewPointsException("not enough points")

if need_sort:

# sort the list by angle around center

pt_lst.sort(key=lambda x: np.arctan2(x[1] - center[1],

x[0] - center[0]))

# add first point to end because it is periodic (makes the

# interpolation code happy)

pt_lst.append(pt_lst[0])

# make array for handing in to spline fitting

pt_array = np.vstack(pt_lst).T

# do spline fitting

tck, u = si.splprep(pt_array, s=len(pt_lst) * (pix_err ** 2), per=True)

return tck

@classmethod

def from_pts(cls, new_pts, **kwargs):

tck = cls._get_spline(new_pts, **kwargs)

this = cls(tck)

this.raw_pts = new_pts

return this

def __init__(self, tck):

'''Use `from_pts` class method to construct instance

self.tck = tck

self._cntr = None

self._circ = None

self._th_offset = None

def write_to_hdf(self, parent_group, name=None):

'''

if name is not None:

curve_group = parent_group.create_group(name)

else:

curve_group = parent_group

curve_group.attrs['tck0'] = self.tck[0]

curve_group.attrs['tck1'] = np.vstack(self.tck[1])

curve_group.attrs['tck2'] = self.tck[2]

@property

def circ(self):

'''returns a rough estimate of the circumference'''

if self._circ is None:

new_pts = si.splev(np.linspace(0, 1, 1000), self.tck, ext=2)

self._circ = np.sum(np.sqrt(np.sum(np.diff(new_pts, axis=1) ** 2,

axis=0)))

return self._circ

@property

def cntr(self):

'''returns a rough estimate of the circumference'''

if self._cntr is None:

new_pts = si.splev(np.linspace(0, 1, 1000), self.tck, ext=2)

self._cntr = np.mean(new_pts, 1)

return self._cntr

@property

def th_offset(self):

"""

if self._th_offset is None:

x, y = self.q_phi_to_xy(0, 0) - self.cntr.reshape(2, 1)

self._th_offset = np.arctan2(y, x)

return self._th_offset

@property

def tck0(self):

return self.tck[0]

@property

def tck1(self):

return self.tck[1]

@property

def tck2(self):

return self.tck[2]

def q_phi_to_xy(self, q, phi, cross=None):

'''Converts q, phi pairs -> x, y pairs. All other code that

# make sure data is arrays

q = np.asarray(q)

# convert real units -> interpolation units

phi = np.mod(np.asarray(phi), 2 * np.pi) / (2 * np.pi)

# get the shapes

q_shape, phi_shape = [_.shape if (_.shape != () and

len(_) > 1) else None for

_ in (q, phi)]

# flatten everything

q = q.ravel()

phi = phi.ravel()

# sanity checks on shapes

if cross is False:

if phi_shape != q_shape:

raise ValueError("q and phi must have same" +

" dimensions to broadcast")

if cross is None:

if ((phi_shape is not None) and (q_shape is not None)

and (phi_shape == q_shape)):

cross = False

elif q_shape is None:

cross = False

q = q[0]

else:

cross = True

x, y = si.splev(phi, self.tck, ext=2)

dx, dy = si.splev(phi, self.tck, der=1, ext=2)

norm = np.sqrt(dx ** 2 + dy ** 2)

nx, ny = dy / norm, -dx / norm

# if cross, then

if cross:

data_out = zip(

*map(lambda q_: ((x + q_ * nx).reshape(phi_shape),

(y + q_ * ny).reshape(phi_shape)),

q)

)

else:

data_out = np.vstack([(x + q * nx).reshape(phi_shape),

(y + q * ny).reshape(phi_shape)])