This article demonstrates matplotlib's animation module by animating marker particles in a fluid flow around a cylinder. It's a bit long because it ties together a number of different ideas:

- stream functions
- numerical integration
- plotting and animation

Before we really start, let's copy a function from a previous article on stream functions in potential flow.

```
import sympy
from sympy.abc import x, y
def velocity_field(psi):
u = sympy.lambdify((x, y), psi.diff(y), 'numpy')
v = sympy.lambdify((x, y), -psi.diff(x), 'numpy')
return u, v
```

Here, `velocity_field` accepts a symbolic function, `psi`, and returns
functions that calculate the velocity components for coordinates `x` and
`y`. Given a stream function `psi`, we can calculate the velocity at any
point in space. For plotting, however, we need translate these velocities to
displacements; in other words, we need to integrate.

## Calculating particle displacements

To calculate displacements from velocities, we need a numerical integration routine. Below, we wrap three integration functions:

- a simple first-order stepper based on Euler's method
- one using a (simple) Runge-Kutta implementation in matplotlib
- and one based on scipy's odeint function

```
import numpy as np
from matplotlib import mlab
from scipy import integrate
def euler(f, pts, dt):
vel = np.asarray([f(xy) for xy in pts])
return pts + vel * dt
def rk4(f, pts, dt):
new_pts = [mlab.rk4(f, xy, [0, dt])[-1] for xy in pts]
return new_pts
def ode_scipy(f, pts, dt):
new_pts = [integrate.odeint(f, xy, [0, dt])[-1] for xy in pts]
return new_pts
available_integrators = dict(euler=euler, rk4=rk4, scipy=ode_scipy)
```

These integration functions return updated particle coordinates based on a
velocity function, `f`; the initial coordinates, `pts`; and the time step,
`dt`. These functions would be faster if they operated on all coordinates
simultaneously instead of looping; however, that would require you to join the
(N-by-2) coordinates into a single 1D array so that the integration routines
will run properly. After that, you'd have to split them back up for plotting.

Next, we define a factory function (a function that returns a function) that connects our integrator to our velocity-field functions. The returned function will return the next particle position based on the current particle position and the time step.

```
def displace_func_from_velocity_funcs(u_func, v_func, method='rk4'):
"""Return function that calculates particle positions after time step.
Parameters
----------
u_func, v_func : functions
Velocity fields which return velocities at arbitrary coordinates.
method : {'euler' | 'rk4' | 'scipy'}
Integration method to update particle positions at each time step.
"""
def velocity(xy, t=0):
"""Return (u, v) velocities for given (x, y) coordinates."""
# Dummy `t` variable required to work with integrators
# Must return a list (not a tuple) for scipy's integrate functions.
return [u_func(*xy), v_func(*xy)]
odeint = available_integrators[method]
def displace(xy, dt):
return odeint(velocity, xy, dt)
return displace
```

This function looks long, but it's mostly documentation and comments.

## Animations in Matplotlib

Finally, we get to the animation portion. Matplotlib (or more precisely, Ryan May) added the animation module in version 1.1. This module greatly simplifies the process of generating animations. Nevertheless, I wanted a way to reuse animations, which didn't seem terribly easy based on the animation examples. I ended up creating a fairly simple Animation class, which uses (but doesn't subclass) matplotlib's animation class.

The following example uses mpltools.animation to plot particles in a potential flow. There are a few important parts:

`__init__`creates a figure (the name`self.fig`is required) where the animation gets drawn.`init_background`draws any background elements that aren't updated between frames (optional).`update`adds particles on the left side of the frame, updates their positions for a specified time step, and removes particles that leave the right side of the frame.

```
import matplotlib.pyplot as plt
from mpltools.animation import Animation
plt.rc('contour', negative_linestyle='solid')
class StreamFuncAnim(Animation):
def __init__(self, stream_function, dt=0.05, xlim=(-1, 1), ylim=None):
self.dt = dt
# Initialize velocity field and displace *functions*
self.u, self.v = velocity_field(stream_function)
self.displace = displace_func_from_velocity_funcs(self.u, self.v)
# Save bounds of plot
self.xlim = xlim
self.ylim = ylim if ylim is not None else xlim
# Animation objects must create `fig` and `ax` attributes.
self.fig, self.ax = plt.subplots()
self.ax.set_aspect('equal')
def init_background(self):
"""Draw background with streamlines of flow.
Note: artists drawn here aren't removed or updated between frames.
"""
x0, x1 = self.xlim
y0, y1 = self.ylim
# Create 100 x 100 grid of coordinates.
Y, X = np.mgrid[x0:x1:100j, y0:y1:100j]
# Horizontal and vertical velocities corresponding to coordinates.
U = self.u(X, Y)
V = self.v(X, Y)
self.ax.streamplot(X, Y, U, V, color='0.7')
def update(self):
"""Update locations of "particles" in flow on each frame frame."""
pts = []
while True:
pts = list(pts)
pts.append((self.xlim[0], random_y(self.ylim)))
pts = self.displace(pts, self.dt)
pts = np.asarray(pts)
pts = remove_particles(pts, self.xlim, self.ylim)
self.ax.lines = []
x, y = np.asarray(pts).transpose()
lines, = self.ax.plot(x, y, 'ro')
yield lines, # return line so that blit works properly
```

In the `update` method, I've chosen to clear and redraw all the displayed
particles, but it would probably be more efficient to update the particle
positions.

`StreamFuncAnim` calls a couple of small utility functions: First,
`random_y`, which returns a random y-position within the domain so that we
can add particles to the left edge.

```
def random_y(ylim):
yrange = np.diff(ylim)
return yrange * np.random.rand(1)[0] + ylim[0]
```

It also calls `remove_particles`, which removes particles that are outside
the limits of the plot.

```
def remove_particles(pts, xlim, ylim):
if len(pts) == 0:
return []
outside_xlim = (pts[:, 0] < xlim[0]) | (pts[:, 0] > xlim[1])
outside_ylim = (pts[:, 1] < ylim[0]) | (pts[:, 1] > ylim[1])
keep = ~(outside_xlim|outside_ylim)
return pts[keep]
```

Now we have an animation class, `StreamFuncAnim`, that we can use to animate
the flow of particles.

## Particles flowing around a cylinder

For this example, let's copy a function from a previous article on stream functions in potential flow. The following function defines the stream function (in symbolic form) for a cylinder in a uniform flow:

```
import sympy
from sympy.abc import x, y
radius = 1
def cylinder_stream_function(U=1, R=radius):
r = sympy.sqrt(x**2 + y**2)
psi = U * (r - R**2 / r) * sympy.sin(sympy.atan2(y, x))
return psi
```

We could go ahead and use `StreamFuncAnim` to animate particles in this flow,
but instead, let's subclass `StreamFuncAnim` and add a circle to identify
where the cylinder is defined:

```
class CylinderFlow(StreamFuncAnim):
def init_background(self):
StreamFuncAnim.init_background(self)
c = plt.Circle((0, 0), radius=radius, facecolor='none')
self.ax.add_patch(c)
```

Finally, we can show this animation by passing an instance of the stream
function to the animation class and calling its `animate` method:

```
stream_function = cylinder_stream_function()
cylinder_flow = CylinderFlow(stream_function, xlim=(-3, 3))
cylinder_flow.animate(blit=True)
plt.show()
```

Science!

## Final thoughts

It's interesting to note that the Euler method is noticeably less-accurate than
`'rk4'` or scipy's integration routines (as one might expect). Particles
entering and leaving the frame should do so symmetrically (particle paths on
the right of the sphere should mirror those on the left). When using the naive
`'euler'` implementation, however, noticeable errors build up: Particles
entering along the centerline of the cylinder (i.e. the stagnation point) will
drift away from the centerline by the time they reach the opposite side of the
cylinder. I haven't noticed much difference between matplotlib's `'rk4'`
implementation and the more-accurate functions in `scipy.integrate`, but in
real-world use, scipy's functions are preferred.

## Comments

comments powered by Disqus