separate into roll and pitch subplot

viz/viz.py

@@ -6,19 +6,21 @@ import random
 import time
 
 fig = plt.figure()
-ax = plt.subplot(111)
-ax.set_ylim([0, 5])   # set the bounds to be 10, 10
-ax.set_xlim([0, 5])
+
+ax1 = plt.subplot(211)
+ax2 = plt.subplot(212)
 
 class RegrMagic(object):
     """Mock for function Regr_magic()
     """
     def __init__(self):
         self.x = 0.0
+        self.y = 0.0
 
     def __call__(self) -> float:
         self.x += random.randint(-1, 1)
-        return self.x
+        self.y += random.randint(-1, 1)
+        return self.x, self.y
 
 regr_magic = RegrMagic()
 
@@ -26,18 +28,9 @@ def frames():
     while True:
         yield regr_magic()
 
-def plot_line(angle: float):
-     '''
-     angle - Angle you want your end point at in degrees.
-     length - Length of the line you want to plot.
-
-     Will plot the line on a 5 x 5 plot.
-     '''
-
-     # unpack the first point
-     x, y = (2.5, 2.5)
-     length = 20.0
 
+def get_start_end(point: (float, float), length: float, angle: float):
+     x, y = point
      # find the end point
      endx = x + length * math.cos(math.radians(angle))
      endy = y + length * math.sin(math.radians(angle))
@@ -45,10 +38,45 @@ def plot_line(angle: float):
      # find the start point
      startx = x - length * math.cos(math.radians(angle))
      starty = y - length * math.sin(math.radians(angle))
-     ax.clear()
+
+     return (endx, endy, startx, starty)
+
+def plot_line(args):
+     '''
+     angle1 - Angle you want your end point at in degrees.
+     length - Length of the line you want to plot.
+
+     Will plot the line on a 5 x 5 plot.
+     '''
+
+     # unpack the first point
+     x, y = (2., 2.)
+     length = 5.0
+     roll_angle = args[0]
+     pitch_angle = args[1]
+
+     ax1.clear()
+     ax2.clear()
+
+     ax1.set_ylim([0, 4])   # set the bounds to be 10, 10
+     ax2.set_ylim([0, 4])   # set the bounds to be 10, 10
+    
+     ax1.set_xlim([0, 4])
+     ax2.set_xlim([0, 4])
 
      # plot the points
-     return ax.plot([startx, endx], [starty, endy])
+     endx, endy, startx, starty = get_start_end((x,y), length, roll_angle)
+     ax1.plot([startx, endx], [starty, endy], c='black', linewidth='2.0')
+
+     endx, endy, startx, starty = get_start_end((x,y), length, pitch_angle)
+     ax2.plot([startx, endx], [starty, endy], c='black', linewidth='2.0')
 
 ani = animation.FuncAnimation(fig, plot_line, frames=frames, interval=10)
+# save the animation as an mp4.  This requires ffmpeg or mencoder to be
+# installed.  The extra_args ensure that the x264 codec is used, so that
+# the video can be embedded in html5.  You may need to adjust this for
+# your system: for more information, see
+# http://matplotlib.sourceforge.net/api/animation_api.html
+#ani.save('double_pendulum.mp4', fps=30, extra_args=['-vcodec', 'libx264'])
+plt.legend()
 plt.show()