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)