#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy
import matplotlib
from scipy.integrate import odeint
from scipy.optimize import fsolve
import dynamics
import json
# function to convert string tab description json to numpy.ndarray
def json2arr(str):
"""
Convert a string of json structure to a numpy.array
Args:
str: string
describing a list (from a json file)
:rtype: numpy.array
:return: the numpy.array of the json structure
:Example:
>>> string = "[1,2,3]"
>>> Array = json2arr(string)
"""
return eval("numpy.array("+str+",dtype=float)")
class Nbody(object):
tol = 1e-10
[docs] def __init__(self, N=2, mass=None,init=None,dim=4,rotating=None,jsonFile=None,zdzInit=None):
"""
Constructor
:param N: int The number of bodies. default 2.
:param mass: numpy.array of masses of the bodies.
default None that implies each body has a mass equal to one.
:param init: numpy.array of initial conditions (ex: [x1, y1, dx1, dy1,x2,y2,dx2,dy2...] in dimension 4 (2D)).
default None, the bodies are on a circle.
:param dim: number of dimension (space and velocity). default 4 (planar motion)
:param rotating: rotation matrix (dim/2 times dim/2) for the frame. Default None (no rotation)
:param jsonFile: name of the json file to construct the object. default none.
"""
if(isinstance(jsonFile,str)):
json_data = open(jsonFile).read()
data = json.loads(json_data)
# size of the system
self._N = int(data["N"])
# size of the system dynamics for each body
self._dim = int(data["dim"])
# the mass initialized to one
if(not data.has_key("mass")):
self.mass = numpy.ones(self._N)
else:
self.mass = json2arr(data["mass"])
# the position velocity initialization
if(not data.has_key("init")): # initialization for N-gon
self.z = numpy.ones(self._N*self._dim)
for i in range(self._N):
#print i*self._dim
deg = 2.0*numpy.pi/self._N
#print i*deg
self.z[(i*self._dim)+0] = numpy.cos(i*deg)
self.z[(i*self._dim)+1] = numpy.sin(i*deg)
self.z[(i*self._dim)+2] = -numpy.sin(i*deg)
self.z[(i*self._dim)+3] = numpy.cos(i*deg)
else:
self.z = json2arr(data["init"])
# the state and the linearized state around the solution
if(not data.has_key("zdzInit")): #
self.zdz = numpy.ones(2*self._N*self._dim)
else:
self.zdz = json2arr(data["zdzInit"])
# matrix of rotating coordinates
if(not data.has_key("rotating")):
self.rotating = None
else:
self.rotating = json2arr(data["rotating"])
else: # we consider the other argumets
# size of the system
self._N = N
# size of the system dynamics for each body
self._dim = dim
# the mass initialized to one
if(not isinstance(mass,numpy.ndarray)):
self.mass = numpy.ones(N)
else:
self.mass = mass.copy()
# the position velocity initialization
if(not isinstance(init, numpy.ndarray)): # initialization for N-gon
self.z = numpy.ones(self._N*self._dim)
for i in range(self._N):
#print i*self._dim
deg = 2.0*numpy.pi/self._N
#print i*deg
self.z[(i*self._dim)+0] = numpy.cos(i*deg)
self.z[(i*self._dim)+1] = numpy.sin(i*deg)
self.z[(i*self._dim)+2] = -numpy.sin(i*deg)
self.z[(i*self._dim)+3] = numpy.cos(i*deg)
else:
self.z = init.copy()
# zdz initialization
if(not isinstance(zdzInit, numpy.ndarray)):
self.zdz = numpy.ones(2*self._N*self._dim)
else:
self.zdz = zdzInit.copy()
# matrix of rotating coordinates
if(not isinstance(rotating,numpy.ndarray)):
self.rotating = None
else:
self.rotating = rotating.copy()
[docs] def get_N(self):
"""
Get the number of bodies
:return: (int) Return the size of the system
"""
return self._N
[docs] def get_dim(self):
"""
Get the dimensions (position+velocity of one body)
:return: (int) Return the size of the system dynamics for each body
"""
return self._dim
[docs] def get_z(self):
"""
Get the position of one body
:return: Return the current state (numpy.array of dimension N*dim)
"""
return self.z
[docs] def get_r(self,I):
"""
Get the position of one body
:param I: the number of the considered body (starting at 0)
:return: Return the position of the Ith Body
"""
return self.z[I*self._dim:(I*self._dim+self._dim/2)]
[docs] def get_v(self,I):
"""
Get the velocity of one body
:param I: the number of the considered body (starting at 0)
:return: Return the velocity of the Ith Body
"""
return self.z[(I*self._dim+self._dim/2):(I*self._dim+self._dim)]
[docs] def integrate(self,t,fileName=None):
"""
Method to integrate a solution for a given N-body problem
:param t: (numpy array) t of all the times for which we want to get a point
:param fileName: (string) optional parameter to save the solution in a external file.
The first column of the file contains the time steps. The others contain
the value of the phase state of each body.
:return: the integrate solution at each times in a numpy array of size of t
"""
z0 = self.z.copy()
sol = odeint(dynamics.system, z0, t, rtol = self.tol, atol = self.tol, args=(self.mass, self._N,self._dim,self.rotating))
if(isinstance(fileName, str)): # if a name of file is given, save the sol into the file
print numpy.array([t]).shape, sol.shape
dat = numpy.hstack((numpy.array([t]).T,sol))
print dat.shape
numpy.savetxt(fileName, dat)
return sol
[docs] def Dintegrate(self,t,fileName=None):
"""
Method to integrate a solution for a given N-body problem and the linearized system around the solution
:param t: (numpy array) t of all the times for which we want to get a point
:param fileName: (string) optional parameter to save the solution in a external file.
The first column of the file contains the time steps. The others contain
the value of the phase state of each body.
:return: the integrate solution at each times in a numpy array of size of t
"""
zdz0 = self.zdz.copy()
sol = odeint(dynamics.Dsystem, zdz0, t, rtol = self.tol, atol = self.tol, args=(self.mass, self._N,self._dim,self.rotating))
if(isinstance(fileName, str)): # if a name of file is given, save the sol into the file
print numpy.array([t]).shape, sol.shape
dat = numpy.hstack((numpy.array([t]).T,sol))
print dat.shape
numpy.savetxt(fileName, dat)
return sol
[docs] def move(self,t):
"""
Move the state integrating self.z during the time t
:param t: time (float) for integration
:return: move the self.z to the point after integration during t
"""
sol = odeint(system, self.z, [0.0,t], rtol = self.tol, atol = self.tol, args=(self.mass, self._N,self._dim,self.rotating))
self.z = sol[-1,:].copy()
[docs] def monodromy(self,t):
"""
Compute the monodromy matrix starting à self.z position pour a periodic solution of period t. The computation is done integrating the linearized equation around the periodic solution.
:param t: (float) time for integration. The period of the periodic solution
:return: (numpy.array) the monodromy matrix. Dim (N*dim, N*dim)
"""
# initialization
M = numpy.zeros((self._N*self._dim,self._N*self._dim))
# save the zdz value
zdzSave = self.zdz.copy()
T=[0.0,t/2.0,t]
for i in range(self._N*self._dim):
ei = numpy.zeros(self._N*self._dim)
ei[i] = 1.0
Init = numpy.hstack((self.z,ei))
self.zdz = Init.copy()
sol = self.Dintegrate(T)
M[i,:] = sol[-1,self._N*self._dim:].copy()
# restore the zdz value
self.zdz = zdzSave.copy()
return M
[docs] def monodromyDF(self,t,dx=1e-09):
"""
Compute the monodromy matrix starting à self.z position pour a periodic solution of period t. The computation is done by finite differences.
:param t: (float) time for integration. The period of the periodic solution
:return: (numpy.array) the monodromy matrix. Dim (N*dim, N*dim)
"""
# initialization
M = numpy.zeros((self._N*self._dim,self._N*self._dim))
# save the z value
zSave = self.z.copy()
T=[0.0,t/2.0,t]
Tm=[0.0,-t/2,-t]
for i in range(self._N*self._dim):
ei = numpy.zeros(self._N*self._dim)
ei[i] = dx
self.z = self.z+ei
sol = self.integrate(T)
self.z = zSave.copy()
self.z = self.z-ei
sol2 = self.integrate(Tm)
self.z = zSave.copy()
M[i,:] = (sol[-1,:] - sol2[-1,:])/(2.0*dx)
# restore the zdz value
return M
def shoot(self,t,z0):
T = [0.0, t]
sol = odeint(system, z0, T, rtol = 1e-10, atol = 1e-10,args=(self.mass, self._N,self._dim))
return numpy.hstack(sol[1,:])
def shootFChoregraphy(self,z0):
T = 2.0*numpy.pi/self._N
sol = self.shoot(T,z0)
#print "init = ", self.z
#print "sol = ", sol
zeros = numpy.empty(self._N*self._dim)
zeros[0:(self._N-1)*self._dim] = sol[0:(self._N-1)*self._dim]-z0[self._dim:(self._N)*self._dim]
zeros[(self._N-1)*self._dim:(self._N)*self._dim] = sol[(self._N-1)*self._dim:(self._N)*self._dim]-z0[0:self._dim]
return zeros
def shootChoregraphy(self):
return fsolve(self.shootFChoregraphy, self.z,xtol=1.0e-06,full_output=True)
def massHomotopy(self, Mass,epsinit = 1e-03, epsmin = 1e-12, epsmax = 0.1, itmax=100):
print "-------------- Mass Homotopy --------------"
# save the initial mass
massinit = self.mass.copy()
# save the initial z
zinit = self.z.copy()
# homotopy parameter
lbd = 0.0
# number of iteration
it = 0
# step
eps = epsinit
zold = self.z.copy()
while(lbd < 1.0 and it < itmax):
lbdold = lbd
lbd = min(lbd + eps,1.0)
it = it +1
self.z = self.z + lbd*(self.z-zold)
self.mass = (1.0-lbd)*massinit+lbd*Mass
(z, infodict, ier, mesg) = self.shootChoregraphy()
if(ier==1): # if the newton did converge
zold = self.z.copy()
self.z = z.copy()
eps = min(1.1*eps,epsmax)
print "Success Iteration:",it," output:", ier, mesg
print "eps = ", eps, ", lambda = ", lbd
else:
lbd = lbdold
eps = max(eps*0.75,epsmin)
print "Fail Iteration:",it," output:", ier,mesg
print "eps = ", eps, ", lambda = ", lbd
print "---------------------------------------"
if(it==itmax):
print "Maximal number of iteration"
#self.z = zinit.copy()
#self.mass = massinit.copy()
else:
print "Mass Homotopy success"