An animated 3D matplotlib figure that rotates incrementally as the plot is drawn.

animated_rotated_3dtrajectory.py

import argparse
import matplotlib.pyplot as plot
from matplotlib import animation
from mpl_toolkits.mplot3d.axes3d import Axes3D
import numpy as np
import scipy.io
import ar.ar as AR

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = 'Cilia AR Subspace Plotting',
        epilog = 'lol moar cilia', add_help = 'How to use',
        prog = 'python -m cilia.pipelines.ar.subspace')
    parser.add_argument('-i', '--input', required = True,
        help = 'Path to patches created by running AR cross-validation script.')
    parser.add_argument('-q', '--degree', type = int, required = True,
        help = 'Number of principal components to use.')
    parser.add_argument('-r', '--rois', required = True,
        help = 'Path to dictionary containing the names of the ROIs used.')
    parser.add_argument('-f', '--files', required = True,
        help = 'Path to text file containing ROIs to build a figure from.')
    parser.add_argument('-o', '--output', required = True,
        help = 'Output directory for any plots.')

    args = vars(parser.parse_args())

    # Load the matfile.
    m = scipy.io.loadmat(args['input'])
    rois = scipy.io.loadmat(args['rois'])['rois']
    D = np.array(m['D'])
    frames = D.shape[1] / len(rois)
    examples = np.loadtxt(args['files'], dtype = np.str)
    N_trajectories = np.size(examples)

    # Compute the subspace.
    X, C = AR.state_space(D, args['degree'])
    x_t = np.zeros(shape = (N_trajectories, frames, 3))
    for i, roi in list(enumerate(rois)):
        roi = roi.strip()
        if roi in examples:
            # Which one?
            index = np.where(examples == roi)[0][0]

            # Extract the state-space for this particular ROI.
            x_t[index] = X[:3, (i * frames):(i + 1) * frames].T


    # Locate the ROIs of interest in D.
    # Set up figure & 3D axis for animation
    fig = plot.figure()
    ax = fig.add_axes([0, 0, 1, 1], projection='3d')
    #ax.axis('off')

    # choose a different color for each trajectory
    colors = ['b', 'r']

    # set up lines and points
    lines = sum([ax.plot([], [], [], '-', c=c)
                 for c in colors], [])
    pts = sum([ax.plot([], [], [], 'o', c=c)
               for c in colors], [])

    # prepare the axes limits
    ax.set_xlim((-2, 2))
    ax.set_ylim((-2, 2))
    ax.set_zlim((-2, 2))
    ax.set_ylabel("$x_2$", fontsize = 24)
    ax.set_xlabel("$x_1$", fontsize = 24)
    ax.set_zlabel("$x_3$", fontsize = 24)

    # set point-of-view: specified by (altitude degrees, azimuth degrees)
    ax.view_init(2, 0)


    # initialization function: plot the background of each frame
    def init():
        for line, pt in zip(lines, pts):
            line.set_data([], [])
            line.set_3d_properties([])

            pt.set_data([], [])
            pt.set_3d_properties([])
        return lines + pts

    # animation function.  This will be called sequentially with the frame number
    def animate(i):
        # we'll step two time-steps per frame.  This leads to nice results.
        i = i % x_t.shape[1]

        for line, pt, xi in zip(lines, pts, x_t):
            x, y, z = xi[:i].T
            line.set_data(x, y)
            line.set_3d_properties(z)

            pt.set_data(x[-1:], y[-1:])
            pt.set_3d_properties(z[-1:])

        ax.view_init(30, 0.3 * i * 5)
        fig.canvas.draw()
        return lines + pts

    # instantiate the animator.
    anim = animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=500, interval=30, blit=True)

    # Save as mp4. This requires mplayer or ffmpeg to be installed
    anim.save('ar_space.mp4', fps=15, extra_args=['-vcodec', 'libx264'])