NBabel

 1 #!/usr/bin/python
 2 
 3 from __future__ import division
 4 
 5 import operator
 6 import pickle
 7 import numpy as Npy
 8 from numpy import *
 9 from numpy.random import *
10 from random import seed, random
11 from math import sqrt, pow, sin, cos, atan, fabs
12 from sys import exit, stdout
13 from copy import deepcopy
14 
15 PI = 4*atan(1)
16 
17 class Cluster() :
18     def __init__(self, N) :
19         self.m = zeros((N,1))
20         self.r = zeros((N,3))
21         self.v = zeros((N,3))
22     def __iter__(self):
23         return self
24     def __repr__(self) :
25         tt = 'mass:\n%s\n' % (self.m)
26         tt += 'position:\n%s\n' % (self.r)
27         tt += 'velocity:\n%s' % (self.v)
28         return tt
29     def sort(self) :
30         cl_w_key = [ (dot(self.r[i], self.r[i]),
31                      copy(self.m[i]), copy(self.r[i]), copy(self.v[i]))
32                          for i in range(len(self.m))]
33         cl_w_key.sort(key=operator.itemgetter(0))
34         for i in range(len(self.m)) :
35             self.m[i] = cl_w_key[i][1]
36             self.r[i] = cl_w_key[i][2]
37             self.v[i] = cl_w_key[i][3]
38 
39     def get_acc(self) :
40         acc = zeros(shape(self.r))
41         for i in range(len(self.r)) :
42             Ri = self.r[i]
43             mi = self.m[i][0]
44             for j in range(i+1, len(self.r)) :
45                 Rj = self.r[j]
46                 mj = self.m[j][0]
47                 Rij = Ri-Rj
48                 apre = pow(dot(Rij,Rij), -3.0/2)
49                 acc[i] -= mj*apre*Rij
50                 acc[j] += mi*apre*Rij
51         return acc
52 
53     def energy(self) :
54         KE = 0.5*sum(self.v*self.v*self.m)
55         PE = 0.0
56         for i in range(len(self.r)) :
57             Ri = self.r[i]
58             mi = self.m[i][0]
59             for j in range(i+1, len(self.r)) :
60                 Rj = self.r[j]
61                 mj = self.m[j][0]
62                 Rij = Ri-Rj
63                 PE -= mi*mj/sqrt(dot(Rij,Rij))
64         return PE+KE,KE,PE
65 
66     def LagrangianRadii(self) :
67         LR_perc = array([0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.99]
68                   )
69         clu_cop = deepcopy(self)
70         clu_cop.sort()
71         LR_vals = LagRad(clu_cop, LR_perc)
72         lr = zeros(len(LR_perc))
73         for i in range(len(LR_perc)) :
74             lr[i] = LR_vals[i]
75         ll = column_stack((LR_perc[:,newaxis],lr[:,newaxis]))
76         return ll
77 
78     def scale(self) :
79         # this function scales positions, such that PE = -1/2
80         # and scales velocities, such that KE = 1/4
81         En,KE,PE = self.energy()
82         self.r *= -PE*2.0
83         self.v *= 1.0/sqrt(4.0*KE)
84 
85     def set_CM_to_zero(self) :
86         Mtot = self.m.sum()
87         CMr = (self.r*self.m).sum(axis=0) / Mtot
88         CMv = (self.v*self.m).sum(axis=0) / Mtot
89         self.r = self.r - CMr
90         self.v = self.v - CMv
91         CMr = (self.r*self.m).sum(axis=0) / Mtot
92         CMv = (self.v*self.m).sum(axis=0) / Mtot
93 
94     def print_star(t) :
95         EN,KE,PE = energy(clu)
96         print t,
97         print EN+0.25, KE-0.25, PE+0.5,
98         for i in range(len(self.m)) :
99             for k in range(3) :
100                 print self.r[i][k],
101                 for k in range(3) :
102                     print self.v[i][k],
103         print
104         stdout.flush()
105 
106 def lin_interp(x_y, x) :
107     # given a list of values x_y (both monotonically
108     # increasing), find the corresponding y value for x
109     if x<x_y[0][0] or x>=x_y[-1][0] :
110         print "out of bounds"
111         exit(-1)
112     b = 0
113     e = len(x_y)-1
114     while 1: # loop invariant x_y[b][1]<= y < x_y[e][1]
115         m = int((e+b)/2)
116         if x_y[m][0] > x : e=m
117         else : b=m
118         if e-b==1 : break
119     # XXX diagnostics only
120     if e-b!=1 or x>=x_y[e][0] or x<x_y[b][0] :
121         print 'problem with binary search!'
122         exit(-2)
123     # XXX end of diagnostics
124     return (x_y[b][1] +
125         (x_y[e][1]-x_y[b][1])*(x-x_y[b][0])/(x_y[e][0]-x_y[b][0]))
126 
127 def LagRad(clu, list=[0.01, 0.1, 0.2, 0.3, 0.4,
128               0.5, 0.6, 0.7, 0.8, 0.9, 0.99]) :
129     # calculates the lagrange radii of the cluster,
130     # for the values given in the list
131     # this routine assumes that the cluster is sorted
132     # it does not assume that the total mass is unity
133     m_r = [[0.0, 0.0]]
134     mtot = 0.0
135     for i in range(len(clu.m)) :
136         r = sqrt(dot(clu.r[i], clu.r[i]))
137         mtot += clu.m[i][0]
138         m_r.append([mtot,r])
139     m_r = [[x[0]/mtot, x[1]] for x in m_r]
140     # XXX diagnostics only
141 #    for x in m_r:
142 #        print x
143     # end of diagnostics
144     radii = []
145     for LR in list :
146         radii.append(lin_interp(m_r, LR))
147     return radii
148 
149 if __name__=="__main__":
150 
151     fd = open('Plummer.in', 'r')
152     cl = pickle.load(fd)
153     ET,KE,PE = cl.energy()
154     print ET, KE, PE
155     print cl.LagrangianRadii()
156 
157 

 1 #!/usr/bin/python
 2 
 3 from __future__ import division
 4 
 5 import operator
 6 from numpy import *
 7 from math import pow
 8 from sys import exit, stdout
 9 import pickle
10 import Cluster as sc
11 
12 def drift_cluster(clu,dt) :
13     clu.r += clu.v*dt
14 
15 def kick_cluster(cl,dt) :
16     clu.v += cl.get_acc(clu)*dt
17 
18 def evolve_cluster_DKD_leapfrog(clu, dt) :
19     # Drift-Kick-Drift leapfrog
20     drift_cluster(clu,dt/2.0)
21     kick_cluster(clu,dt)
22     drift_cluster(clu,dt/2.0)
23 
24 if __name__=="__main__":
25     fd = open('Plummer.in', 'r')
26     cl = pickle.load(fd)
27     ET0,KE0,PE0 = cl.energy()
28 
29     t = 0.0
30     tend = 1.0
31     dt = 1e-3
32     k = 0
33     # Predictor-Corrector Leapfrog
34     acc_0 = cl.get_acc()
35     while t<tend :
36         cl.r += dt*cl.v + 0.5*dt*dt*acc_0
37         acc_1 = cl.get_acc()
38         cl.v += 0.5*dt*(acc_0+acc_1)
39         acc_0 = acc_1
40         t += dt
41         k += 1
42         if k%10==0 :
43             ET,KE,PE = cl.energy()
44             print "t= ", t, " E= ", ET, KE, PE, \
45                   " dE= ", (ET-ET0)/ET0, (KE-KE0)/KE0, (PE-PE0)/PE0
46             print cl.LagrangianRadii()