thomson_problem.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import random
import math
######################################
#h^5 order
h=0.05

n=5

step=5000

eps=1e-3

######################################

def rho_ij(theta_i,phi_i,theta_j,phi_j):
    tmp=1 \
        -np.sin(theta_i)*np.cos(phi_i)*np.sin(theta_j)*np.cos(phi_j) \
        -np.sin(theta_i)*np.sin(phi_i)*np.sin(theta_j)*np.sin(phi_j) \
        -np.cos(theta_i)*np.cos(theta_j)
    if tmp==0:
        tmp+=eps
    return tmp**0.5

def rho_ij_inverse_theta_i(theta_i,phi_i,theta_j,phi_j):
    tmp=-np.cos(theta_i)*np.cos(phi_i)*np.sin(theta_j)*np.cos(phi_j) \
        -np.cos(theta_i)*np.sin(phi_i)*np.sin(theta_j)*np.sin(theta_j) \
        +np.sin(theta_i)*np.cos(theta_j)
    return -tmp/(rho_ij(theta_i,phi_i,theta_j,phi_j)**3)

def rho_ij_inverse_phi_i(theta_i,phi_i,theta_j,phi_j):
    tmp=np.sin(theta_i)*np.sin(phi_i)*np.sin(theta_j)*np.cos(phi_j) \
        -np.sin(theta_i)*np.cos(phi_i)*np.sin(theta_j)*np.sin(phi_j)
    return -tmp/(rho_ij(theta_i,phi_i,theta_j,phi_j)**3)

#############################################################

def ftheta_dd(X):
    res=np.empty(n)
    for i in range(n):
        res[i]=np.sin(X[2][i])*np.cos(X[2][i])*X[1][i]**2
    for i in range(n):
        for j in range(n):
            if i==j:
                continue
            res[i]-=rho_ij_inverse_theta_i(X[2][i],X[3][i],X[2][j],X[3][j])
    return res

def fphi_dd(X):
    res=np.empty(n)
    for i in range(n):
        if np.sin(X[2][i]==0):
            res[i]=0
            continue
        res[i]=-2*X[0][i]*X[1][i]*np.cos(X[2][i])/np.sin(X[2][i])
        for j in range(n):
            if i==j:
                continue
            res-=rho_ij_inverse_phi_i(X[2][i],X[3][i],X[2][j],X[3][j])/(np.sin(X[2][i])**2)
    return res

def potential(X):
    res=0
    for i in range(n):
        for j in range(i+1,n):
            res+=1/rho_ij(X[2][i],X[3][i],X[2][j],X[3][j])
    return res

############################################################3

def function(X):
    res=np.empty((4,n))
    res[0]=ftheta_dd(X)
    res[1]=fphi_dd(X)
    res[2]=X[0]
    res[3]=X[1]
    return res
############################################################

def rotation(ans):
    theta=ans[2][0]
    phi=ans[3][0]

    Y=np.empty((3,n))
    Y[0]=[np.sin(ans[2][i])*np.cos(ans[3][i]) for i in range(n)]
    Y[1]=[np.sin(ans[2][i])*np.sin(ans[3][i]) for i in range(n)]
    Y[2]=[np.cos(ans[2][i]) for i in range(n)]
    
    Rx=np.matrix([[1,0,0],
                    [0,np.cos(theta),-np.sin(theta)],
                    [0,np.sin(theta),np.cos(theta)]
                ])
    Ry=np.matrix([[np.cos(-theta),0,np.sin(-theta)],
                    [0,1,0],
                    [-np.sin(-theta),0,np.cos(-theta)]
                ])
    Rz=np.matrix([[np.cos(-phi),-np.sin(-phi),0],
                    [np.sin(-phi),np.cos(-phi),0],
                    [0,0,1]
                ])
    res=Ry*Rz*Y
    return res.tolist()

def projection(Y):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    plt.hold(True)

    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)

    x = np.outer(np.cos(u), np.sin(v))
    y = np.outer(np.sin(u), np.sin(v))
    z = np.outer(np.ones(np.size(u)), np.cos(v))
    ax.plot_wireframe(x, y, z,rstride=7,cstride=7)
    
    Y=list(Y)
    ax.scatter(Y[0],Y[1],Y[2],color='#ff0000',s=100)

    plt.show()


def main():
    X=np.zeros((step+1,4,n))
    for i in range(n):
        X[0][2][i]=random.uniform(0,np.pi) 
        X[0][3][i]=random.uniform(0,2*np.pi) 
    
    ans=X[0]
    potential_min=potential(X[0])
    for s in np.arange(step):
        k1=function(X[s])
        k2=function(X[s]+h/2*k1)
        k3=function(X[s]+h/2*k2)
        k4=function(X[s]+h*k3)
        X[s+1]=X[s]+h/6*(k1+2*k2+2*k3+k4)
        X[s+1][2]=X[s+1][2]-np.pi*np.floor(X[s+1][2]/np.pi)
        X[s+1][3]=X[s+1][3]-2*np.pi*np.floor(X[s+1][3]/2/np.pi)
        
        tmp=potential(X[s+1])
        #print tmp
        if tmp<potential_min:
            potential_min=tmp
            ans=X[s+1]
        """
        ok=True
        for i in range(n):
            if np.abs(X[s+1][2][i]-X[s][2][i])>h**5 or np.abs(X[s+1][3][i]-X[s][3][i])>h**5:
                ok=False
        if ok:
            break
        """
        if s%(step/10)==0:
            print '[',
            for i in range(s/(step/10)+1):
                print '*',
            for i in range(9-s/(step/10)):
                print ' ',
            print ']'
    
    Y=rotation(ans)
    projection(Y)

if __name__=='__main__':
    main()