def __init__(self,
Nparticles,
M,
a = farts.a0,
dt=0.01,
interaction = "ClassicalNBody",
masses = None,
init_state = None,
save_dir = "/Users/alexdeich/Dropbox/thesis/code/mine/nbody_output"):
self.start_particles = Nparticles
self.Nparticles = Nparticles
self.interaction_type = interaction
self.M = M
self.dt = dt
self.a = a
self.event_horizon = 0
self.cleanup = []
if init_state is not None:
self.use_state = np.copy(init_state)
else:
self.use_state = None
self.save_dir = save_dir
self.init_state = np.copy(init_state)
self.state = np.copy(init_state)
self.collisions = {}
self.fname = ""
if init_state is not None:
if init_state.shape == (Nparticles,2,2):
self.init_state = use_state
else:
raise ValueError("Initial state not the right shape: ",init_state.shape,
"\nShould be ",(Nparticles,2,2))
if masses == None:
self.masses = np.zeros(Nparticles)+0.0001
else:
if len(masses) == Nparticles:
self.masses = masses
else:
raise ValueError("Mass list must be of length ", Nparticles,
"\nGiven length:", len(masses))
def TimeEvolve(self,nsteps,comments):
self.cleanup = []
t=0
##Get init_state
if self.use_state == None:
self.state = np.copy(self.MakeInitialConditions())
self.init_state = np.copy(self.state)
self.Nparticles = len(self.state)
else:
self.state = np.copy(self.use_state)
print("Got initial conditions for {} particles".format(self.start_particles))
primary = self.get_header(nsteps,comments)
frame0 = self.get_hdu()
hdulist = fits.HDUList([primary,frame0])
total_time = 0
for step in xrange(1, nsteps):
stepstart = time.time()
self.state = self.UpdateStateVectorRK4(t)
framen = self.get_hdu()
hdulist.append(framen)
t += self.dt
end = time.time()
steptime = end-stepstart
total_time += steptime
avg = total_time/step
perc = 100*((step+1)/nsteps)
sys.stdout.write('\rFrame {} of {} completed ({}%). Step: {}s, Total: {}s, Estimated time remaining: {}s. Nparticles: {}'.format(step+1,
nsteps,
'%0.1f'%perc,
'%0.4f'%steptime,
'%0.4f'%total_time,
'%i'%(((avg * nsteps)+1)-total_time),
'%i'%self.Nparticles))
sys.stdout.flush()
print("\nWriting to disk...")
filenum = farts.get_filenum(self.save_dir)
self.fname = "{}/nbody_{}_{}.fits".format(self.save_dir,self.start_particles,filenum)
hdulist.writeto(self.fname,clobber=True)
print("Data written at {}".format(self.fname))
def UpdateStateVectorRK4(self,t):
self.event_horizon = self.get_event_horizon(t)
new_state = np.ndarray((self.Nparticles,2,2))
for particle in xrange(self.Nparticles):
kstate = np.copy(self.state)
#do RK4 shit
k1 = self.dt * self.SingleParticleDerivativeVector(kstate,particle, t+self.dt)
kstate[particle] = np.copy(self.state[particle]) + k1/2
k2 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+(self.dt/2))
kstate[particle] = np.copy(self.state[particle]) + k2/2
k3 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+(self.dt/2))
kstate[particle] = np.copy(self.state[particle]) + k3
k4 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+self.dt)
new_state[particle] = np.copy(self.state[particle]) + (1/3)*(k1/2 + k2 + k3 + k4/2)
#Get rid of gobbled or ejected particles
if self.cleanup != []:
for particle in self.cleanup:
print("\n***particle {} shit the bed at step {}***".format(particle,int(t/self.dt)))
self.remove_particle(particle)
new_state = np.delete(new_state,particle,axis=0)
print("***particle {} removed***".format(particle))
self.cleanup.remove(particle)
if self.cleanup == []:
print("***cleanup completed***")
self.cleanup = []
#Cleanup collided particles
if self.collisions != {}:
print("*****")
for key in self.collisions:
p1 = self.collisions[key][1][0]
p2 = self.collisions[key][1][1]
new_particle = self.combine_particles()
print("***{} and {} collided at step{}***".format(p1,p2,int(t/self.dt)))
new_state = np.delete(new_state,[self.collisions[key][0],self.collisions[1]],axis=0)
new_state = np.append(new_state,new_particle,axis = 0)
self.collisions = {}
return(new_state)
def TimeEvolve(self,nsteps,comments,write=True):
self.cleanup = []
self.nsteps = nsteps
t=0
##Get init_state
if self.use_state == None:
self.cleanup = []
self.state = np.copy(farts.make_initial_conditions2(self.start_particles,
self.M,
self.a(0)))
self.init_state = np.copy(self.state)
self.Nparticles = len(self.state)
else:
self.state = np.copy(self.init_state)
print("Got initial conditions for {} particles".format(self.start_particles))
primary = self.get_header(nsteps,comments)
frame0 = self.get_hdu()
filenum = farts.get_filenum(self.save_dir,self.start_particles)
self.dirname = "{}/nbody_{}_{}".format(self.save_dir,self.start_particles,filenum)
os.mkdir(self.dirname)
os.mkdir("{}/data".format(self.dirname))
hdulist = fits.HDUList([primary,frame0])
total_time = 0
savenums = nsteps/1000
print('\n\n')
for step in xrange(1, nsteps):
stepstart = time.time()
self.state = self.UpdateStateVectorRK4(t)
framen = self.get_hdu()
hdulist.append(framen)
t += self.dt
end = time.time()
steptime = end-stepstart
total_time += steptime
avg = total_time/step
perc = 100*((step+1)/nsteps)
sys.stdout.write('\rFrame {} of {} completed ({}%). Step: {}s, Total: {}s, Estimated time remaining: {}s. Nparticles: {}'.format(step+1,
nsteps,
'%0.1f'%perc,
'%0.4f'%steptime,
'%0.4f'%total_time,
'%i'%(((avg * nsteps)+1)-total_time),
'%i'%self.Nparticles))
sys.stdout.flush()
if step%1000 == 0:
if write == True:
print("\nWriting to disk...")
fname = "{}/data/{}.fits".format(self.dirname,step)
hdulist.writeto(fname,clobber=True)
print("Frames {} - {} written at {}".format(step-1000,step,fname))
hdulist = fits.HDUList([primary])
self.fname_list.append(fname)
if len(hdulist)!=1:
print("\nWriting to disk...")
fname = "{}/data/{}.fits".format(self.dirname,step)
hdulist.writeto(fname,clobber=True)
print("Frames {} written at {}".format(step+1,fname))
self.fname_list.append(fname)
print(self.state)
def __init__(self,
Nparticles,
M=1,
a = None,
dt=0.01,
interaction = False,
collisions = 'elastic',
masses = None,
init_state = None,
cr = 0.001,
save_dir = "./output"):
self.start_particles = Nparticles
self.Nparticles = Nparticles
self.interaction = interaction
self.M = M
self.dt = dt
self.a = a
self.cr = cr
self.event_horizon = 0
self.cleanup = []
if init_state is not None:
self.use_state = np.copy(init_state)
else:
self.use_state = None
self.save_dir = save_dir
self.init_state = np.copy(init_state)
self.state = np.copy(init_state)
self.collision_dict = {}
self.collisions=collisions
self.collided_particles = np.array([])
self.fname_list = []
self.dirname = ""
self.plotdata = 0
self.skip_mkdir=False
self.nsteps = 0
self.old_collisions = []
if self.Nparticles == None and self.init_state is not None:
self.Nparticles = len(self.init_state)
if self.a == None:
self.a = farts.a0
if init_state is not None:
if init_state.shape == (Nparticles,2,3):
self.init_state = init_state
else:
raise ValueError("Initial state not the right shape: ",init_state.shape,
"\nShould be ",(Nparticles,2,3))
if masses == None:
self.masses = np.zeros(Nparticles)+0.0001
else:
if len(masses) == Nparticles:
self.masses = masses
else:
raise ValueError("Mass list must be of length ", Nparticles,
"\nGiven length:", len(masses))
self.collision_radius = .05
def UpdateStateVectorRK4(self,t):
self.event_horizon = self.get_event_horizon(t)
new_state = np.ndarray((self.Nparticles,2,3))
for particle in xrange(self.Nparticles):
kstate = np.copy(self.state)
k1 = self.dt * self.SingleParticleDerivativeVector(kstate,particle, t)
kstate[particle] = np.copy(self.state[particle]) + k1/2
k2 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+(self.dt/2))
kstate[particle] = np.copy(self.state[particle]) + k2/2
k3 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+(self.dt/2))
kstate[particle] = np.copy(self.state[particle]) + k3
k4 = self.dt * self.SingleParticleDerivativeVector(kstate,particle,t+self.dt)
new_state[particle] = np.copy(self.state[particle]) + (1/3)*(k1/2 + k2 + k3 + k4/2)
#Get rid of gobbled or ejected particles
if self.cleanup != []:
new_state = np.delete(new_state,self.cleanup,axis=0)
self.remove_particle(particle)
for particle in self.cleanup:
print("\n***particle {} shit the bed at step {}***".format(particle,int(t/self.dt)))
print("***particle {} removed***".format(particle))
self.cleanup.remove(particle)
if self.cleanup == []:
print("***cleanup completed***")
self.cleanup = []
big_particle_list = []
big_vector_list = []
new_masses_list = []
#Cleanup collided particles
if self.collisions == "inelastic":
if self.collision_dict != {}:
for key in self.collision_dict:
particles = self.collision_dict[key][1]
m = 0
x = self.collision_dict[key][0][0]
y = self.collision_dict[key][0][1]
mvx = 0
mvy = 0
for particle in particles:
m += self.masses[particle]
for particle in particles:
mvx += (self.state[particle][1][0]*self.masses[particle])
mvy += (self.state[particle][1][1]*self.masses[particle])
big_particle_list.append(particle)
new_masses_list.append(m)
new_particle = [[x,y],[mvx/m,mvy/m]]
big_vector_list.append(new_particle)
self.masses = np.delete(self.masses,big_particle_list,axis=0)
self.masses = np.append(self.masses,np.array(new_masses_list),axis = 0)
new_state = np.delete(new_state,big_particle_list,axis=0)
new_state = np.append(new_state,np.array(big_vector_list),axis = 0)
self.Nparticles = len(new_state)
self.collision_dict = {}
self.collided_particles = np.array([])
elif self.collisions == 'elastic':
distances = squareform(pdist(self.state[:,0,:]))
ind1, ind2 = np.where(distances < 2 * self.collision_radius)
unique = (ind1 < ind2)
ind1 = ind1[unique]
ind2 = ind2[unique]
new_collisions = zip(ind1,ind2)
for i in new_collisions:
if i not in self.old_collisions:
i1 = i[0]
i2 = i[1]
m1 = self.masses[i1]
m2 = self.masses[i2]
x1 = new_state[i1,0,0]
y1 = new_state[i1,0,1]
z1 = new_state[i2,0,3]
x2 = new_state[i2,0,0]
y2 = new_state[i2,0,1]
z2 = new_state[i2,0,2]
x1d = new_state[i1,1,0]
y1d = new_state[i1,1,1]
z1d = new_state[i1,1,2]
x2d = new_state[i2,1,0]
y2d = new_state[i2,1,1]
z2d = new_state[i2,1,2]
new_state[i1, 1] = [(x1d*(m1-m2) + 2*m2*x2d)/(m1+m2),
(y1d*(m1-m2) + 2*m2*y2d)/(m1+m2),
(z1d*(m1-m2) + 2*m2*z2d)/(m1+m2)]
new_state[i2, 1] = [(x2d*(m2-m1) + 2*m1*x1d)/(m1+m2),
(y2d*(m2-m1) + 2*m1*y1d)/(m1+m2),
(z2d*(m2-m1) + 2*m1*z1d)/(m1+m2)]
self.old_collisions = new_collisions
return(new_state)
