added JSAnimation_frametools.py

[rjleveque]

Utility to make animations in notebook or html file. Can probably be improved.

Taken from http://faculty.washington.edu/rjl/classes/am583s2014/notes/labs/lab15.html
where some examples of use can be found.

Used in notebook being developed to illustrate new dtopotools, to appear in clawpack/apps/notebooks/geoclaw.

[ketch]

This will be great to have.

When working in the notebook with PyClaw, we often just store the frames in memory (rather than writing them to disk), and create an animation from that. It would be nice if functions like create_anim didn't assume the files are on disk; actually, I think the step of acquiring the frames could conveniently be abstracted away altogether. This is something I'd love to see across all of VisClaw, but we could start by implementing it here if that sounds reasonable.

[mandli]

I wrote a little while ago, it might contain some added functionality of interest. It unfortunately assumes that the controller claw exists already and contains the data in the frames list but I thought I would share

def add_animation(title, cbar_label, cmap, var, limits, figsize=(8,6)):
    fig = plt.figure(figsize=figsize)
    if isinstance(var, list) or isinstance(var, tuple):
        ax_list = [fig.add_subplot(1,2,1), fig.add_subplot(1,2,2)]
    else:
        ax_list = [fig.add_subplot(1,1,1)]
        var = [var]
        title = [title]
        cbar_label = [cbar_label]

    x, y = claw.frames[0].state.grid.p_centers
    im_list = []
    for (n,ax) in enumerate(ax_list):
        im_list.append( ax.imshow(var[n](0).T, cmap=cmap, extent=[x.min(), x.max(), y.min(), y.max()],
                             vmin=limits[0], vmax=limits[1], interpolation='nearest', origin='lower') )
        ax.set_title(title[n])
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        cbar = fig.colorbar(im_list[n], ax=ax)
        cbar.set_label(cbar_label[n])

    def fplot(frame_number):
        for (n,im) in enumerate(im_list):
            im.set_data(var[n](frame_number).T)
        return im

    return fig, fplot

Note that allows subplots and an arbitrary function for the plotting. A couple of examples from a shallow-water notebook:

Plot velocities:

def extract_velocity(component, frame_num, DRY_TOL=1e-3):
    q = claw.frames[frame_num].q
    u = numpy.zeros(q.shape[1:])
    index = numpy.nonzero((numpy.abs(q[0,:,:]) > DRY_TOL) * (q[0,:,:] != numpy.nan))
    u[index[0],index[1]] = q[component,index[0],index[1]] / q[0,index[0],index[1]]
    return u

u = lambda frame_num:extract_velocity(1, frame_num) 
v = lambda frame_num:extract_velocity(2, frame_num)
fig, fplot = add_animation(['X-Velocity','Y-Velocity'], ['m/s', 'm/s'], 'PiYG', [u, v], [-1.0, 1.0])
animation.FuncAnimation(fig, fplot, frames=len(claw.frames), interval=100)

Plot Vorticity:

def vorticity(frame_num, DRY_TOL=1e-3):
    u = extract_velocity(1, frame_num)
    v = extract_velocity(2, frame_num)

    dx = claw.frames[frame_num].state.grid.delta[0]
    dy = claw.frames[frame_num].state.grid.delta[1]

    u_y = numpy.zeros(u.shape)
    u_y[:,0] = (-3.0*u[:,0] + 4.0 * u[:,1] - u[:,2]) / (2.0 * dy)
    u_y[:,-1] = (u[:,-3] - 4.0 * u[:,-2] + 3.0 * u[:,-1]) / (2.0 * dy)
    u_y[:,1:-1] = (u[:,2:] - u[:,0:-2]) / (2.0 * dy)

    v_x = numpy.zeros(v.shape)
    v_x[0,:] = (-3.0*v[0,:] + 4.0 * v[1,:] - v[2,:]) / (2.0 * dx)
    v_x[-1,:] = (u[-3,:] - 4.0 * u[-2,:] + 3.0 * u[-1,:]) / (2.0 * dx)
    v_x[1:-1,:] = (v[2:,:] - v[0:-2,:]) / (2.0 * dx)

    return v_x - u_y

fig, fplot = add_animation('Vorticity', '1/s', 'PRGn', vorticity, [-0.1, 0.1])
animation.FuncAnimation(fig, fplot, frames=len(claw.frames), interval=100)

[ketch]

I just realized that this PR is only about making an animation from already-generated plots. I guess that what @mandli and I are talking about is different (see #100), though somewhat related.

[rjleveque]

Yes, this code is to create an animation out of a set of png files created by any means, with utilities to help save png files with sequential numbers. At some point we should do more with this and along the lines of what @ketch and @mandli suggest, but for now I'm going to merge this PR since I'm using it in some notebooks to illustrate geoclaw tools.