#!/usr/bin/python
# -*- coding: utf8 -*-
importmatplotlib.pyplotasplt
importmatplotlibasmpl
importnumpyasnp
frommathimport *
code_website = 'http://commons.wikimedia.org/wiki/User:Geek3/mplwp'
try:
 importmplwp
except ImportError, er:
 print 'ImportError:', er
 print 'You need to download mplwp.py from', code_website
 exit(1)
name = 'mplwp_ballistic_trajectories_velocities.svg'
fig = mplwp.fig_standard(mpl)
xlim = 0,2.6; fig.gca().set_xlim(xlim)
ylim = 0,2.6*355/515.; fig.gca().set_ylim(ylim)
fig.gca().xaxis.set_major_locator(mpl.ticker.MultipleLocator(0.4))
fig.gca().yaxis.set_major_locator(mpl.ticker.MultipleLocator(0.4))
fromscipy.integrateimport odeint
fromscipy.optimizeimport brentq
defballistic(g, k, xy0, v0, alpha0, tt):
 # use a four-dimensional vector function vec = [x, y, vx, vy]
 defdif(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)], tt)
 return vec[:,0], vec[:,1] # return x(tt) and y(tt)
g = 1.0
k = 1.0
alpha0 = pi/4
for v0 in np.linspace(0, 10, 6)[1:]:
 t1 = brentq(lambda t: ballistic(g,k,[0,0],v0,alpha0,[0,t])[1][1],0.1,5)
 t = np.linspace(0, t1, 5001)
 x, y = ballistic(g, k, [0, 0], v0, alpha0, t)
 while len(y) > 1 and y[-2] <= 0.0: x = x[:-1]; y = y[:-1]
 plt.plot(x, y,
 label=ur'$v_0=,円{:.0f}$'.format(v0))
mpl.rc('legend', borderaxespad=1.0)
plt.legend(loc='upper left').get_frame().set_alpha(0.9)
plt.savefig(name)
mplwp.postprocess(name)