got it working from a movable polygon example here: https://matplotlib.org/stable/gallery/event_handling/poly_editor.html
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D
from matplotlib.artist import Artist
global new_freqs
def dist(x, y):
"""
Return the distance between two points.
"""
d = x - y
return np.sqrt(np.dot(d, d))
def dist_point_to_segment(p, s0, s1):
"""
Get the distance of a point to a segment.
*p*, *s0*, *s1* are *xy* sequences
This algorithm from
http://www.geomalgorithms.com/algorithms.html
"""
v = s1 - s0
w = p - s0
c1 = np.dot(w, v)
if c1 <= 0:
return dist(p, s0)
c2 = np.dot(v, v)
if c2 <= c1:
return dist(p, s1)
b = c1 / c2
pb = s0 + b * v
return dist(p, pb)
class PolygonInteractor:
"""
A polygon editor.
Key-bindings
't' toggle vertex markers on and off. When vertex markers are on,
you can move them, delete them
'd' delete the vertex under point
'i' insert a vertex at point. You must be within epsilon of the
line connecting two existing vertices
"""
showverts = True
epsilon = 5 # max pixel distance to count as a vertex hit
def __init__(self, ax, poly, start_freqs):
if poly.figure is None:
raise RuntimeError('You must first add the polygon to a figure '
'or canvas before defining the interactor')
self.ax = ax
canvas = poly.figure.canvas
self.poly = poly
x, y = zip(*self.poly.xy)
self.line = Line2D(x, y,
marker='o', markerfacecolor='r',
animated=True)
self.ax.add_line(self.line)
ax.vlines(x,linestyle="--", ymin=0.0, ymax=1, alpha=0.9, color=col, gid="new_lines")
self.original_freqs = start_freqs #variable for storing the starting frequencies
self.orig_cols = col # array or starting colours
self.new_cols = self.orig_cols # array or new colours, same as original to begin with
self.issue_spacing = 1.0e6 # variable to store the value of kids to close together
self.cid = self.poly.add_callback(self.poly_changed)
self._ind = None # the active vert
canvas.mpl_connect('draw_event', self.on_draw)
canvas.mpl_connect('button_press_event', self.on_button_press)
canvas.mpl_connect('key_press_event', self.on_key_press)
canvas.mpl_connect('button_release_event', self.on_button_release)
canvas.mpl_connect('motion_notify_event', self.on_mouse_move)
self.canvas = canvas
def draw_new_positions(self):
for i in range(len(self.new_cols)):
if self.poly.xy[i,0] != self.original_freqs[i]:
if self.poly.xy[i,0] < self.original_freqs[i]:
self.new_cols[i] = "purple" #if the kid has moved backwward show purple
elif (self.poly.xy[i+1,0]-self.poly.xy[i,0]) < self.issue_spacing or (self.poly.xy[i,0]-self.poly.xy[i-1,0]) < self.issue_spacing :
self.new_cols[i] = "black" #if the kid to close the the ones next to it show black
else:
self.new_cols[i] = "orange" #if the kid has moved and is positioned ok show orange
else:
self.new_cols[i] = self.orig_cols[i]
new_lines = self.ax.vlines(self.poly.xy[:-1,0], ymin=-1, ymax=0, linestyle="--", color=self.new_cols, alpha=0.9, gid="new_lines") #new line to move where mouse is
self.ax.draw_artist(new_lines) # drawing the line on moving the mouse
self.canvas.blit(self.ax.bbox) # blitting the canvas to render moving
new_lines.remove()
def on_draw(self, event):
self.background = self.canvas.copy_from_bbox(self.ax.bbox)
self.ax.draw_artist(self.poly)
self.ax.draw_artist(self.line)
# do not need to blit here, this will fire before the screen is
# updated
def poly_changed(self, poly):
"""This method is called whenever the pathpatch object is called."""
# only copy the artist props to the line (except visibility)
vis = self.line.get_visible()
Artist.update_from(self.line, poly)
self.line.set_visible(vis) # don't use the poly visibility state
def get_ind_under_point(self, event):
"""
Return the index of the point closest to the event position or *None*
if no point is within ``self.epsilon`` to the event position.
"""
# display coords
xy = np.asarray(self.poly.xy)
xyt = self.poly.get_transform().transform(xy)
xt, yt = xyt[:, 0], xyt[:, 1]
d = np.hypot(xt - event.x, yt - event.y)
indseq, = np.nonzero(d == d.min())
ind = indseq[0]
if d[ind] >= self.epsilon:
ind = None
return ind
def on_button_press(self, event):
"""Callback for mouse button presses."""
if not self.showverts:
return
if event.inaxes is None:
return
if event.button != 1:
return
self._ind = self.get_ind_under_point(event)
def on_button_release(self, event):
"""Callback for mouse button releases."""
if not self.showverts:
return
if event.button != 1:
return
self._ind = None
def on_key_press(self, event):
"""Callback for key presses."""
if not event.inaxes:
return
if event.key == 't': #toggles the movable points on and off
self.showverts = not self.showverts
self.line.set_visible(self.showverts)
if not self.showverts:
self._ind = None
elif event.key == ' ': #prints the x vals of all polygon points (which are the new frequencies) to the console
new_freqs = self.poly.xy[:,0]
for i in range(len(new_freqs)-1):
print("{:.1f},".format(new_freqs[i]))
# print(len(new_freqs))
elif event.key == 'l': #save new frequencies to csv file and show final plot
new_freqs = self.poly.xy[:-1,0]
new_data = 0
new_data = np.zeros((len(new_freqs), 2))
new_data[:,0] = data["kid_id"]
new_data[:,1] = new_freqs
new_data_df = pd.DataFrame(data=new_data, columns=["kid_id", "f0"]) #makes a new data frame to save to csv with all new positions
new_data_df.to_csv("new_kid_positions.csv", index=False)
plt.close()
plt.figure("new array", dpi=150)
for i in range(len(new_data_df["f0"])):
if self.poly.xy[i,0] == self.original_freqs[i]:
col="green"
else:
col="orange"
plt.axvline(new_data_df["f0"][i]/1e6, color=col, linestyle="--", linewidth=1.5, alpha=0.9)
plt.plot([],[], color="orange", label="Moved")
plt.plot([],[], color="green", label="Not moved")
plt.legend(loc="best")
plt.xlabel("Frequency (MHz)")
plt.ylabel("")
plt.title("Altered array")
plt.grid()
plt.show()
if self.line.stale:
self.canvas.draw_idle()
def on_mouse_move(self, event):
"""Callback for mouse movements."""
if not self.showverts:
return
if self._ind is None:
return
if event.inaxes is None:
return
if event.button != 1:
self.moving_line.remove()
return
x, y = event.xdata, 0
self.poly.xy[self._ind] = x, y
if self._ind == 0:
self.poly.xy[-1] = x, y
elif self._ind == len(self.poly.xy) - 1:
self.poly.xy[0] = x, y
self.line.set_data(zip(*self.poly.xy))
# self.remove_series()
# f = x
# self.add_series(f, "new_lines", len(self.poly.xy[:,0])-1)
# ax.axvline(x, ymin=-1, ymax=1, linestyle="--", color="orange", alpha=0.9, animated=True)
self.canvas.restore_region(self.background)
self.ax.draw_artist(self.poly)
self.ax.draw_artist(self.line)
# self.moving_line = self.ax.axvline(x, ymin=-1, ymax=1, linestyle="--", color="orange", alpha=0.9) #new line to move where mouse is
# self.ax.draw_artist(self.moving_line) # drawing the line on moving the mouse
# self.moving_line.remove()
self.draw_new_positions()
self.canvas.blit(self.ax.bbox) # blitting the canvas to render moving
if __name__ == '__main__':
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
#reading in freequency positions
data = pd.read_csv("sorted_kids.csv")
#getting color for each line
col = [] #empty list
for i in range(len(data["issue"])):
if data["issue"][i] == 1:
col.append("red")
else:
col.append("green")
# col = np.array(col)
xs = data["f0"]
ys = np.zeros_like(xs)
poly = Polygon(np.column_stack([xs, ys]), animated=True)
fig, ax = plt.subplots(dpi=150, figsize=(12,6))
ax.add_patch(poly)
p = PolygonInteractor(ax, poly, xs)
ax.set_title('Click a point to drag. spacebar=print freqs. T=toggle move. L=save and show final solution')
ax.plot([],[], color="black", alpha=0.0, label=r"TOP: ORIGINAL ARRAY")
ax.plot([],[], color="red", label="to close")
ax.plot([],[], color="green", label="spacing ok")
ax.plot([],[], color="black", alpha=0.0, label="\nBOT: ALTEERED ARRAY")
ax.plot([],[], color="red", label="orig position & to close")
ax.plot([],[], color="green", label="orig position & ok")
ax.plot([],[], color="purple", label="moved backward!")
ax.plot([],[], color="black", label="to close still")
ax.plot([],[], color="orange", label="moved & ok")
ax.legend(loc=1, fontsize=6)
plt.xlabel("Frequency (Hz)")
ax.set_xlim((min(data["f0"])-10e6, max(data["f0"])+10e6))
ax.set_ylim((-1.5, 1.5))
plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
ax.axes.get_yaxis().set_visible(False)
plt.show()
