Code:Vectors

From Predictive Chemistry
Jump to: navigation, search

This interactive code demp plots three vectors using matplotlib. The "zero" vector, x[0], is colored blue because it is recomputed from the other two vectors, x[1] and x[2].

If you drag them around, you can see that x[0] always shows the sum,

x[1] + x[2]

By changing the "def update_vec" function, you can implement dot products, projections, and reflections.

<source lang="python"> import numpy as np import matplotlib.path as mpath import matplotlib.patches as mpatches import matplotlib.pyplot as plt

Path = mpath.Path

  1. fig, ax = plt.subplots()

ax = plt.gca()

  1. arrows = np.array([[[0.0, 0.0], [0.0, 2.0]]])

arrows = np.array([[0.0, 2.0], [-1,-1], [3,-0.5]])

codes = [Path.MOVETO, Path.LINETO, Path.LINETO,

                     Path.LINETO, Path.LINETO, Path.LINETO]

def path_from_vec(x,y,dx,dy):

   r = (dx*dx+dy*dy)**0.5
   a = 0.5
   c = 0.2*dx/r
   s = 0.2*dy/r
   x1 = x + (r-0.2)*dx/r
   y1 = y + (r-0.2)*dy/r
   pathdata = [(x,y),
               (x1,y1),
               (x1+s-a*c, y1-c-a*s),
               (x1+c, y1+s),
               (x1-s-a*c, y1+c-a*s),
               (x1,y1)]
   return pathdata

def to_path(v, color='green'):

   verts = path_from_vec(0., 0., *v)
   path = mpath.Path(verts, codes)
   patch = mpatches.PathPatch(path, facecolor=color, edgecolor='k', alpha=0.5)
   ax.add_patch(patch)
   return patch

class PathInteractor(object):

   showverts = True
   epsilon = 10  # max pixel distance to count as a vertex hit
   def __init__(self, vecs):
       self.paths = []
       color = 'blue'
       vecs[0] = update_vec(vecs)
       for i,v in enumerate(vecs):
           if i > 0:
               color = 'green'
           path = to_path(v, color)
           self.paths.append(path)
           path.set_animated(True)
       self.ax = self.paths[0].axes
       canvas = self.ax.figure.canvas
       self.vecs = vecs
       self._ind = None  # the active vert
       canvas.mpl_connect('draw_event', self.draw_callback)
       canvas.mpl_connect('button_press_event', self.button_press_callback)
       canvas.mpl_connect('key_press_event', self.key_press_callback)
       canvas.mpl_connect('button_release_event', self.button_release_callback)
       canvas.mpl_connect('motion_notify_event', self.motion_notify_callback)
       self.canvas = canvas
   def draw_callback(self, event):
       self.background = self.canvas.copy_from_bbox(self.ax.bbox)
       for path in self.paths:
           self.ax.draw_artist(path)
       self.canvas.blit(self.ax.bbox)
   def get_ind_under_point(self, event):
       'get the index of the vertex under point if within epsilon tolerance'
       # display coords
       r = np.sqrt(sum(self.vecs*self.vecs,1))
       r += (0.1 - r)*(r < 0.1)
       xy = self.vecs*(r-0.1)/r # move location to arrow head
       xyt = self.paths[0].get_transform().transform(xy)
       xt, yt = xyt[:, 0], xyt[:, 1]
       d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2)
       ind = d.argmin()
       if d[ind] >= self.epsilon:
           ind = None
       return ind
   def button_press_callback(self, event):
       'whenever a mouse button is pressed'
       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 button_release_callback(self, event):
       'whenever a mouse button is released'
       if not self.showverts:
           return
       if event.button != 1:
           return
       self._ind = None
   def key_press_callback(self, event):
       'whenever a key is pressed'
       if not event.inaxes:
           return
       #if event.key == 't':
       #   self.showverts = not self.showverts
       #   self.line.set_visible(self.showverts)
       #   if not self.showverts:
       #       self._ind = None
       #self.canvas.draw()
   def motion_notify_callback(self, event):
       'on mouse movement'
       if not self.showverts:
           return
       if self._ind is None:
           return
       if event.inaxes is None:
           return
       if event.button != 1:
           return
       x, y = event.xdata, event.ydata
       self.vecs[self._ind] = x, y
       self.paths[self._ind].get_path().vertices[:] = \
                   path_from_vec(0., 0., *self.vecs[self._ind])
       # update vec 0 based on callback
       if self._ind != 0:
           self.vecs[0] = update_vec(self.vecs)
           self.paths[0].get_path().vertices[:] = \
                   path_from_vec(0., 0., *self.vecs[0])
       self.canvas.restore_region(self.background)
       for path in self.paths:
           self.ax.draw_artist(path)
       self.canvas.blit(self.ax.bbox)

def update_vec(x):

   return x[1]+x[2]

interactor = PathInteractor(arrows) ax.set_title('drag vertices to update path') ax.set_xlim(-3, 4) ax.set_ylim(-3, 4)

plt.show() </source>