Simple worm movement plotter. Also saves animation to mp4.

wormMovementPlot.py

import math
import os
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
import h5py

#SPEED_UP = 4
#DT = 0.05


class wormSkeleton:
    skeletonX, skeletonY = None, None
    x_min, y_min = None, None
    n_frames = None
    name = ''   # TODO: add to __init__, grab from wormFile["info"]
    points = None    # this contains our animation data
    
    def __init__(self, worm_file_path):
        """ Load the worm data into wormFile """
        wormFile = h5py.File(worm_file_path, 'r')
        
        self.skeletonX = wormFile["worm"]["posture"]["skeleton"]["x"].value
        self.skeletonY = wormFile["worm"]["posture"]["skeleton"]["y"].value

        # TODO: combine the single-dimension arrays into a two-dimensional array
        #http://stackoverflow.com/questions/16473042/numpy-vstack-vs-column-stack
        #wormPosition = np.column_stack((skeletonX,skeletonY)).T
        
        #print("wormPosition", np.shape(wormPosition))
        #print("skeletonX", np.shape(skeletonX))
        #print("skeletonY", np.shape(skeletonY))
        
        #print("length of array:", len(x[1]))
        #print("length of array:", len(y[1]))


        # NOTE: Dropped frames show up as NaN.  nanmin returns the min ignoring such NaNs.
        self.x_min = np.nanmin(self.skeletonX) # Why is this not a method of the object?
        self.x_max = np.nanmax(self.skeletonX)
        self.y_min = np.nanmin(self.skeletonY)
        self.y_max = np.nanmax(self.skeletonY)

        wormFile.close()
        
        self.n_frames = self.skeletonX.shape[0]   # total frames in the worm file


worm_file_path1 = os.path.join(os.path.dirname(os.path.abspath(__file__)), "unc-8 (rev) on food R_2010_03_19__09_14_57___2___2_features.mat")
  #u'F:/worm_data/segworm_data/features/798 JU258 on food R_2010_11_25__16_34_17___1___9_features.mat'

w = wormSkeleton(worm_file_path1)
sX = w.skeletonX
sY = w.skeletonY

fig = plt.figure()
fig.suptitle('Worm position over time', fontsize=20)
plt.xlabel('x coordinates', fontsize=18)
plt.ylabel('y coordinates', fontsize=16)
# set the axes to the maximum extent of the worm's travels
ax = plt.axes(xlim=(w.x_min, w.x_max), ylim=(w.y_min, w.y_max))


#alternatively: marker='o', linestyle='None'
wormPoints, = ax.plot(w.skeletonX[0,:], w.skeletonY[0,:], color='green', linestyle='point marker', marker='o', markersize=5) 


# ANIMATION OF WORM

# initialization function: plot the background of each frame
def init():
    wormPoints.set_data([], [])
    return wormPoints,

# animation function.  This is called sequentially
def animateFrame(iFrame):
    wormPoints.set_data(sX[iFrame,:], sY[iFrame,:])
    return wormPoints,

# call for the animation to be run
anim = animation.FuncAnimation(fig, func=animateFrame, init_func=init,
                               frames=math.floor(w.n_frames/60), interval=20, blit=True, repeat_delay=100)   #divide total frames by 60 to get about 15 seconds.

# 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
# http://www.wikihow.com/Install-FFmpeg-on-Windows
FFMpegWriter = animation.writers['ffmpeg']
metadata = dict(title='C. elegans movement video', artist='Matplotlib',
        comment='C. elegans movement video from Shafer lab')
writer = FFMpegWriter(fps=15, metadata=metadata)
anim.save('worm_animation.mp4', writer=writer, fps=15, extra_args=['-vcodec', 'libx264'])


print("Done saving worm animation.")    # DEBUG: this never runs.