[ÇóÖú] Çë½ÌÒ»¸ö¹ØÓÚÖØÁ¦ÏµÍ³ÏÂÖʵãÇòƽÅ×ÂäµØ·´µ¯µÄÔ˶¯¹ì¼£ÊýѧģÐͽ¨Á¢µÄÎÊÌâ

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

def fun(w, t, p):
    """Define the  equation for throw
    Arguments:
        w: vector of the state variables:
            w = [x,y, vx, vy]
        t: time
        p: vector of the parameters:
            p = [m, mu, g, e]
    :return: f
    """
    x, y, vx, vy = w
    m, mu, g, k = p

    # create f = (x',y', vx', vy'):

    f = [vx,                          # x' = dx/dt = vx
         -mu * vx / m,                 # vx' = dvx/dt = -mu * vx / m
         vy,                           # y' = dy/dt = vy
         -mu * vy / m - g,             # vy' = dvy/dt = -mu * vy /m - g
         ]
    return f

def main():
    # Parameter values
    # Masses
    m = 1
    # Friction coefficients
    mu = 0.1
    # Acceleration
    g = 9.8
    # Restitution coefficients
    k = 0.3

    # Initial conditions
    x = 0
    y = 0.5
    vx = 1
    vy = 0


    # ODE solver parameters
    abserr = 1.0e-8
    relerr = 1.0e-6
    # stop time for movement simulation
    stoptime = 0.7
    # delta time
    numpoints = 30

    # Create the time samples for the output of the ODE solver.
    t = [stoptime * float(i) / (numpoints - 1) for i in range(numpoints)]

    # Pack up the parameters and initial conditions:
    p = [m, mu, g, k]
    w0 = [x, y, vx, vy]

    # Call the ODE solver
    wsol = odeint(fun, w0, t, args=(p,), atol=abserr, rtol=relerr)

    with open('ball_movement.dat', 'w') as fi:
        # print and save the solution.
        for t1, w1 in zip(t, wsol):
            if w1[2] >= 0:
                print >> fi, t1, w1[0], w1[2]
            else:
                # print >> fi, "The ball is landing to floor!"

                x = w1[0]
                y = w1[2]
                vx = w1[1]
                vy = w1[3]
    # wsol = odeint(fun, w0, )

    # Plot the solution that was generated

    from numpy import loadtxt
    from pylab import figure, plot, xlabel, grid, hold, legend, title, savefig
    from matplotlib.font_manager import FontProperties

    t, x, y = loadtxt('ball_movement.dat', unpack=True)

    figure(1, figsize=(6, 4.5))

    xlabel('t')
    grid(True)
    hold(True)
    lw = 1

    plot(x, y, 'b', linewidth=lw)
    #plot(t, y, 'g', linewidth=lw)

    legend((r'$x$', r'$y$'), prop=FontProperties(size=16))
    title('Ball Mass System')
    savefig('ball_movement.png', dpi=100)


if __name__ == "__main__":
    main()