File:Mplwp ballistic trajectory drag nodrag.svg

#!/usr/bin/python
# -*- coding: utf8 -*-

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from math import *

code_website = 'http://commons.wikimedia.org/wiki/User:Geek3/mplwp'
try:
    import mplwp
except ImportError, er:
    print 'ImportError:', er
    print 'You need to download mplwp.py from', code_website
    exit(1)

name = 'mplwp_ballistic_trajectory_drag_nodrag.svg'
fig = mplwp.fig_standard(mpl)

xlim = 0,1; fig.gca().set_xlim(xlim)
ylim = -0.35*355/515.,0.65*355/515.; fig.gca().set_ylim(ylim)
mplwp.mark_axeszero(fig.gca())

from scipy.integrate import odeint
from scipy.optimize import brentq
def ballistic(g, k, xy0, v0, alpha0, tt):
    # use a four-dimensional vector function vec = [x, y, vx, vy]
    def dif(vec, t):
        v = sqrt(vec[2]**2 + vec[3]**2)
        return [vec[2], vec[3], -k*v*vec[2], -g -k*v*vec[3]]
    
    # solve the differential equation numerically
    vec = odeint(dif, [xy0[0], xy0[1], v0*cos(alpha0), v0*sin(alpha0)],
        np.concatenate(([0], tt)))[1:,:]
    return vec[:,0], vec[:,1] # return x(tt) and y(tt)

g = 1.0
v0 = 1.0
alpha0 = radians(45)
for k, label in [(0, 'without drag'), (1, 'with air drag')]:
    t1 = brentq(lambda t: ballistic(g,k,[0,0],v0,alpha0,[t])[1][0],0.1,5)
    t2 = brentq(lambda t:
        ballistic(g,k,[0,0],v0,alpha0,[t])[1][0] - ylim[0], 0.1, 5)
    t = np.linspace(0, t1, 5001)
    x, y = ballistic(g, k, [0, 0], v0, alpha0, t)
    plt.plot(x, y, label=label)
    if k > 0:
        t = np.linspace(t1, t2, 5001)
        x, y = ballistic(g, k, [0, 0], v0, alpha0, t)
        plt.plot(x, y, '--', color=fig.gca().lines[-1].get_color())
        

mpl.rc('legend', borderaxespad=1.0)
plt.legend(loc='upper center').get_frame().set_alpha(0.9)
plt.savefig(name)
mplwp.postprocess(name)