In [ ]:

    

import numpy as np
import matplotlib.pyplot as plt

from scipy import integrate
from matplotlib import mlab
from matplotlib import animation

from sympy.abc import x, y
from sympy import sin, atan2, sqrt, lambdify, latex

from JSAnimation import IPython_display

plt.rc('contour', negative_linestyle='solid')


    

In [ ]:

    

def velocity_field(psi):
    u = lambdify((x, y), psi.diff(y), 'numpy')
    v = lambdify((x, y), -psi.diff(x), 'numpy')
    return u, v

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)


    

In [ ]:

    

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

# Initialize velocity field and displace *functions*
radius = 1
def cylinder_stream_function(U=1, R=radius):
    r = sqrt(x**2 + y**2)
    psi = U * (r - R**2 / r) * sin(atan2(y, x))
    return psi

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

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]


    

In [ ]:

    

def cylinder_flow(frames=100, interval=30):
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_aspect('equal')

    # Non-animated.
    stream_function = cylinder_stream_function()
    u, v = velocity_field(stream_function)

    xlim, ylim = (-3, 3), (-3, 3)
    x0, x1 = xlim
    y0, y1 = 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 = u(X, Y)
    V = v(X, Y)
    ax.streamplot(X, Y, U, V, color='0.7')
    c = plt.Circle((0, 0), radius=radius, facecolor='grey')
    ax.add_patch(c)

    # Animated.
    particles, = ax.plot([], [], 'ro')

    def init():
        global pts
        pts = []
        return particles

    dt = 0.05
    displace = displace_func_from_velocity_funcs(u, v)
    def animate(k):
        global pts
        pts = list(pts)
        pts.append((xlim[0], random_y(ylim)))
        pts = displace(pts, dt)
        pts = np.asarray(pts)
        pts = remove_particles(pts, xlim, ylim)
        x, y = np.asarray(pts).transpose()
        particles.set_data(x, y)
        return particles

    return animation.FuncAnimation(fig, animate, init_func=init,
                                frames=frames, interval=interval)


    

In [ ]:

    

cylinder_flow(frames=300)


    

In [ ]: