perfect_rubber.py

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import random

#parameter
n = 10
mu = 10.0
p_0 = 0.01

h = 0.5
step = 100
#########
cnt=0
##########
X_now = np.zeros([n+1,2])
X_next = np.zeros([n+1,2])
##########

def diff_theta(X):
    res = np.zeros([n+1])
    for j in np.arange(0,n+1):
        for i in np.arange(j,n+1):
            for k in np.arange(1,i+1):
                mu_i = 1.0
                if i==n:
                    mu_i = mu
                res[j] += mu_i*X[k][0]*np.cos(X[k][1]-X[j][1])
    return res

def diff_p(X):
    res = np.zeros([n+1])
    for j in np.arange(0,n+1):
        for i in np.arange(j,n+1):
            for k in np.arange(0,i+1):
                mu_i = 1.0
                if i==n:
                    mu_i = mu
                res[j] += mu_i*X[k][0]*X[j][0]*np.sin(X[k][1]-X[j][1])
    return res

def f(X):
    res = np.zeros([n+1,2])
    d_p = diff_p(X)
    d_theta = diff_theta(X)

    for i in np.arange(0,n+1):
        res[i][0] = d_p[i]
        res[i][1] = d_theta[i]

    return res

def develop(X):
    k1 = f(X)
    k2 = f(X+h/2*k1)
    k3 = f(X+h/2*k2)
    k4 = f(X+h*k3)

    res = X+h/6*(k1+2*k2+2*k3+k4)
    for i in np.arange(0,n+1):
        res[i][1] = res[i][1]-(int)(res[i][1]/2/np.pi)*2*np.pi
    return res

def init(im):
    global X_now
    X_0 = np.zeros([n+1,2])
    for i in np.arange(0,n+1):
        X_0[i][0] = random.uniform(-p_0,p_0)
        X_0[i][1] = random.randint(0,1)*np.pi+random.uniform(-np.pi*0.2,np.pi*0.2)
        #X_0[i][1] = random.uniform(0,2*np.pi)
    X_0[n][0] = 0.0
    X_now = X_0

    x = np.zeros([n+2])
    y = np.zeros([n+2])
    for i in np.arange(1,n+1):
        x[i] += x[i-1]+np.cos(X_now[i-1][1])
        y[i] += y[i-1]+np.sin(X_now[i-1][1])
    im.append(plt.scatter(x,y))

def animate(i,fig,im):
    global X_now,X_next
    X_next = develop(X_now)
    X_now = X_next
    
    x = np.zeros([n+2])
    y = np.zeros([n+2])
    for i in np.arange(1,n+1):
        x[i] += x[i-1]+np.cos(X_now[i-1][1])
        y[i] += y[i-1]+np.sin(X_now[i-1][1])

    if len(im)>0:
        im[0].remove()
        im.pop()
    im.append(plt.scatter(x,y))
    
    global cnt
    cnt+=1
    print cnt

def main():
    im = []
    
    #Writer = animation.writers['ffmpeg']
    #writer = Writer(fps=15,bitrate=1800)

    fig = plt.figure(figsize=(5,5))
    plt.xlim(-n,n)
    plt.ylim(-n,n)
    init(im)
    ani = animation.FuncAnimation(fig,animate,fargs=(fig,im),frames=step,interval=1)

    ani.save('im5.gif',writer='imagemagick',fps=10)

    #plt.show()

if __name__=='__main__':
    main()