예제 #1
0
    def make_heatmap(self,function):
        print("\n")
        print("Creating animated heatmap...")
        plt.figure()
        n=1
        for fname in self.fname_list:
            wholedata = fits.open(fname)
            for i in xrange(1,len(wholedata)):
            	data = wholedata[i].data
                bins = scipy.stats.binned_statistic_2d(data['X'],
                	                	               data['Y'],
                	                	               data['MASS'],
                	                	               statistic=function,
                	                	               bins=50,
                	                	               range=[[-30,30],[-30,30]])
				
                x,y = np.where(~np.isnan(bins[0]))
                plt.scatter(x,y,c=bins[0][np.where(~np.isnan(bins[0]))],
                	                	                	marker='s',
                	                	                	edgecolors='none',
                	                	                	s=200)
                t = n*self.dt
                plt.xlim(0,48)
                plt.ylim(0,48)
                circle1 = plt.Circle((24,24), radius = self.get_event_horizon(t), color = 'r', fill = False)
                fig = plt.gcf()
                fig.gca().add_artist(circle1)
                fig.gca().axes.get_xaxis().set_visible(False)
                fig.gca().axes.get_yaxis().set_visible(False)
                plt.axes().set_aspect('equal')
                plt.title('Nparticles: {}'.format(len(data['X'])))
                cb = plt.colorbar()
                plt.clim(0,50)
                cb.set_label('Number density')
                plt.scatter(24,24,marker="x", color="white")
                fname = "{}/movies/heatmap/frames/{}.png".format(self.dirname,n)
                plt.savefig(fname)
                plt.clf()
                sys.stdout.write('\rFrame {} completed'.format(n))
                sys.stdout.flush()
                n+=1  
            wholedata.close()
        os.system("ffmpeg -framerate 300 -i {}/movies/heatmap/frames/%d.png -c:v libx264 -r 30 -pix_fmt yuv420p {}/movies/heatmap/out.mp4".format(self.dirname,self.dirname))
예제 #2
0
    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)
예제 #3
0
    def SingleParticleNewtonianForce(self, i, soi_radius, collision_radius):        
        forces = np.zeros([self.Nparticles,2])
        
        x1 = self.state[i,0,0]
        y1 = self.state[i,0,1]
        
        m1 = self.masses[i]
        collided1 = 0
        
        if self.interaction=='ClassicalNBody':

            xmin = x1-soi_radius
            xmax = x1+soi_radius
            ymin = y1-soi_radius
            ymax = y1+soi_radius
            sliced_indices_x = np.where(np.logical_and(self.state[:,0,0]>xmin,self.state[:,0,0]<xmax))
            sliced_indices_y = np.where(np.logical_and(self.state[:,0,1]>ymin,self.state[:,0,1]<ymax))
            sliced_indices = np.intersect1d(sliced_indices_x,sliced_indices_y)
            sliced_arr = self.state[sliced_indices]
        
            no_i_indices = np.where(np.logical_and(sliced_arr[:,0,0] == x1,sliced_arr[:,0,1]==y1))
            sliced_arr = np.delete(sliced_arr,no_i_indices,axis=0)
            sliced_masses = self.masses[sliced_indices]
            sliced_masses = np.delete(sliced_masses,no_i_indices,axis=0)
            
            distances2 = np.array(np.transpose(np.matrix((sliced_arr[:,0,0]-x1)**2+(sliced_arr[:,0,1]-y1)**2)))
        
            jforce = np.array(np.transpose(np.matrix(sliced_masses)))/distances2
            jforcedir = np.array(np.transpose(np.matrix([sliced_arr[:,0,0]-x1,sliced_arr[:,0,1]-y1])))/np.sqrt(distances2)
            forces = jforce*jforcedir
        

            for particle_num in xrange(self.Nparticles):
                if particle_num != i:
                	collided2 = 0
                	m2 = self.masses[particle_num]
                	x2 = self.state[particle_num,0,0]
                	y2 = self.state[particle_num,0,1]
                	distance2 = ((x2-x1)**2+(y2-y1)**2)
                	
                	"""
                	if the collision dict is empty, put in the two particles being evaluated.
                
                	if it's not empty, first search to see if the main particle is in there
                	if the main particle is in there, move on to the secondary
                	if the secondary particle is not in there, put it in the entry with the main particle
                	"""
                	if self.Nparticles == 2:
                	    if distance2>500000:
                	        print('\n',distance2)
                	        print(self.state)
                	        raise ValueError("it f****d the duck")
                	elif self.Nparticles == 1:
                	    raise ValueError("it shit the bed")
                	
                	if self.collisions != False and i not in self.cleanup and particle_num not in self.cleanup:
                	    if distance2 < collision_radius:
                	        if self.collision_dict == {}:
                	            self.collision_dict[0] = [(x1,y1),[i,particle_num]]
                	        else:
                	            i_val = farts.dict_check(self.collision_dict, i)
                	            keynum = max(self.collision_dict)+1
                	            if i_val == -1:
                	                self.collision_dict[keynum] = [(x1,y1),[i]]
                	                i_val = keynum
                	            elif farts.dict_check(self.collision_dict, particle_num) == -1:
                	                self.collision_dict[i_val][1].append(particle_num)

            return(np.sum(forces,axis=0))