Python interface.BHTree类代码示例

 def new_gas_code_bhtree(self):
     result = BHTree(self.converter)
     result.parameters.epsilon_squared = self.gas_epsilon ** 2
     result.parameters.timestep = 0.125 * self.interaction_timestep
     result.particles.add_particles(self.new_particles_cluster_as_gas())
     result.commit_particles()
     return result

    def test22(self):
        particles = datamodel.Particles(2)
        particles.x = [0.0,10.0] | nbody_system.length
        particles.y = 0.0 | nbody_system.length
        particles.z = 0.0 | nbody_system.length
        particles.vx =  0.0 | nbody_system.speed
        particles.vy =  0.0 | nbody_system.speed
        particles.vz =  0.0 | nbody_system.speed
        particles.mass = 1.0 | nbody_system.mass

        instance = BHTree()
        instance.particles.add_particles(particles) 
        instance.commit_particles()
        self.assertEquals(instance.particles[0].radius, 0.0 | nbody_system.length)
        p = datamodel.Particle(
            x = 1.0  | nbody_system.length,
            y = 2.0 | nbody_system.length,
            z = 3.0 | nbody_system.length,
            vx = 1.0  | nbody_system.speed,
            vy = 2.0 | nbody_system.speed,
            vz = 3.0 | nbody_system.speed,
            mass = 1.0 | nbody_system.mass,
            radius = 4.0 | nbody_system.length,
        )
        instance.particles.add_particle(p) 
        self.assertEquals(instance.particles[0].radius, 0.0 | nbody_system.length)
        self.assertEquals(instance.particles[1].radius, 0.0 | nbody_system.length)
        self.assertEquals(instance.particles[2].radius, 4.0 | nbody_system.length)
        
        instance.stop()

 def test1(self):
     print "Test collect_required_attributes"
     in_memory = self.new_colliders()
     gravity = BHTree(nbody_system.nbody_to_si(1|units.MSun, 1.0|units.RSun))
     gravity.particles.add_particles(in_memory)
     stellar = SeBa()
     stellar.particles.add_particles(in_memory)
     
     collision = StellarEncounterInHydrodynamics(None, None, verbose=True)
     
     self.assertFalse(hasattr(in_memory, "radius"))
     collision.collect_required_attributes(in_memory, gravity, stellar)
     self.assertTrue(hasattr(in_memory, "radius"))
     self.assertAlmostRelativeEqual(in_memory.radius.sum(), 4.2458 | units.RSun, 3)
     
     from_stellar = stellar.particles.copy()
     for attribute in ["x", "y", "z", "vx", "vy", "vz"]:
         self.assertFalse(hasattr(from_stellar, attribute))
     collision.collect_required_attributes(from_stellar, gravity, stellar)
     
     gravity.stop()
     stellar.stop()
     
     for attribute in ["x", "y", "z", "vx", "vy", "vz"]:
         self.assertTrue(hasattr(from_stellar, attribute))
     self.assertAlmostEqual(from_stellar.position, in_memory.position)
     self.assertAlmostEqual(from_stellar.velocity, in_memory.velocity)

    def test11(self):
       
        convert_nbody = nbody_system.nbody_to_si(5.0 | units.kg, 10.0 | units.m)

        instance = BHTree(convert_nbody)
        
        particles = datamodel.Particles(2)
        self.assertEquals(len(instance.particles), 0)
        
        particles.mass = [15.0, 30.0] | units.kg
        particles.radius =  [10.0, 20.0] | units.m
        particles.position = [[10.0, 20.0, 30.0], [20.0, 40.0, 60.0]] | units.m
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.m / units.s

        
        instance.particles.add_particles(particles)
        
        copyof =  instance.particles.copy()
        
        self.assertAlmostEqual(30 | units.kg, copyof[1].mass, 6)
        
        copyof[1].mass = 35 | units.kg
        
        copyof.copy_values_of_all_attributes_to(instance.particles)
        
        self.assertAlmostEqual(35 | units.kg, instance.particles[1].mass, 6)
        instance.stop()

 def test14(self):
     print "Test14: Testing BHTree parameters (I)"
     convert_nbody = nbody_system.nbody_to_si(1.0 | units.yr, 1.0 | units.AU)
     instance = BHTree(convert_nbody)
     
     value,error = instance.legacy_interface.get_epsilon_squared()
     self.assertEquals(0, error)
     self.assertEquals(0.125, value)
     self.assertAlmostEquals(0.125 | units.AU**2, instance.parameters.epsilon_squared, in_units=units.AU**2)
     for x in [0.01, 0.1, 0.2]:
         instance.parameters.epsilon_squared = x | units.AU**2
         self.assertAlmostEquals(x | units.AU**2, instance.parameters.epsilon_squared, in_units=units.AU**2)
     
     (value, error) = instance.legacy_interface.get_time_step()
     self.assertEquals(0, error)
     self.assertEquals(0.015625, value)
     self.assertAlmostEquals(0.015625 | units.yr, instance.parameters.timestep, in_units=units.yr)
     for x in [0.001, 0.01, 0.1]:
         instance.parameters.timestep = x | units.yr
         self.assertAlmostEquals(x | units.yr, instance.parameters.timestep, in_units=units.yr)
     
     (value, error) = instance.legacy_interface.get_theta_for_tree()
     self.assertEquals(0, error)
     self.assertEquals(0.75, value)
     self.assertEquals(0.75, instance.parameters.opening_angle)
     for x in [0.2, 0.5, 0.7]:
         instance.parameters.opening_angle = x
         self.assertEquals(x, instance.parameters.opening_angle)
     
     (value, error) = instance.legacy_interface.get_use_self_gravity()
     self.assertEquals(0, error)
     self.assertEquals(1, value)
     self.assertEquals(1, instance.parameters.use_self_gravity)
     for x in [0, 1]:
         instance.parameters.use_self_gravity = x
         self.assertEquals(x, instance.parameters.use_self_gravity)
     
     (value, error) = instance.legacy_interface.get_ncrit_for_tree()
     self.assertEquals(0, error)
     self.assertEquals(12, value)
     self.assertEquals(12, instance.parameters.ncrit_for_tree)
     for x in [512, 2048, 4096]:
         instance.parameters.ncrit_for_tree = x
         self.assertEquals(x, instance.parameters.ncrit_for_tree)
     
     (value, error) = instance.legacy_interface.get_dt_dia()
     self.assertEquals(0, error)
     self.assertEquals(1.0, value)
     self.assertAlmostEquals(1.0 | units.yr, instance.parameters.dt_dia, in_units=units.yr)
     for x in [0.1, 10.0, 100.0]:
         instance.parameters.dt_dia = x | units.yr
         self.assertAlmostEquals(x | units.yr, instance.parameters.dt_dia, in_units=units.yr)
     instance.stop()

    def test17(self):
        print "Testing BHTree collision_detection"
        particles = datamodel.Particles(7)
        particles.mass = 0.001 | nbody_system.mass
        particles.radius = 0.01 | nbody_system.length
        particles.x = [-101.0, -100.0, -0.5, 0.5, 100.0, 101.0, 104.0] | nbody_system.length
        particles.y = 0 | nbody_system.length
        particles.z = 0 | nbody_system.length
        particles.velocity = [[2, 0, 0], [-2, 0, 0]]*3 + [[-4, 0, 0]] | nbody_system.speed
        
        instance = BHTree(redirection='none')
        instance.initialize_code()
        instance.parameters.set_defaults()
        
        # Uncommenting any of the following two lines will suppress collision detection
#~        instance.parameters.use_self_gravity = 0
#~        instance.parameters.epsilon_squared = 0.0 | nbody_system.length**2
        
        instance.parameters.opening_angle = 0.1
        instance.particles.add_particles(particles)
        collisions = instance.stopping_conditions.collision_detection
        collisions.enable()
        instance.evolve_model(1.0 | nbody_system.time)
        
        self.assertTrue(collisions.is_set())
        self.assertTrue(instance.model_time < 0.5 | nbody_system.time)
        self.assertEquals(len(collisions.particles(0)), 3)
        self.assertEquals(len(collisions.particles(1)), 3)
        self.assertEquals(len(particles - collisions.particles(0) - collisions.particles(1)), 1)
        self.assertEquals(abs(collisions.particles(0).x - collisions.particles(1).x) < 
                (collisions.particles(0).radius + collisions.particles(1).radius),
                [True, True, True])
        
        sticky_merged = datamodel.Particles(len(collisions.particles(0)))
        sticky_merged.mass = collisions.particles(0).mass + collisions.particles(1).mass
        sticky_merged.radius = collisions.particles(0).radius
        for p1, p2, merged in zip(collisions.particles(0), collisions.particles(1), sticky_merged):
            merged.position = (p1 + p2).center_of_mass()
            merged.velocity = (p1 + p2).center_of_mass_velocity()
        
        print instance.model_time
        print instance.particles
        instance.particles.remove_particles(collisions.particles(0) + collisions.particles(1))
        instance.particles.add_particles(sticky_merged)
        
        instance.evolve_model(1.0 | nbody_system.time)
        print
        print instance.model_time
        print instance.particles
        self.assertTrue(collisions.is_set())
        self.assertTrue(instance.model_time < 1.0 | nbody_system.time)
        self.assertEquals(len(collisions.particles(0)), 1)
        self.assertEquals(len(collisions.particles(1)), 1)
        self.assertEquals(len(instance.particles - collisions.particles(0) - collisions.particles(1)), 2)
        self.assertEquals(abs(collisions.particles(0).x - collisions.particles(1).x) < 
                (collisions.particles(0).radius + collisions.particles(1).radius),
                [True])
        instance.stop()

def evolve_cluster_in_galaxy(N, W0, Rinit, tend, timestep, M, R):

    R_galaxy=0.1 | units.kpc
    M_galaxy=1.6e10 | units.MSun
    converter=nbody_system.nbody_to_si(M_galaxy, R_galaxy)
    galaxy=new_plummer_model(10000,convert_nbody=converter)

    print "com:", galaxy.center_of_mass().in_(units.kpc)
    print "comv:", galaxy.center_of_mass_velocity().in_(units.kms)
   
    print len(galaxy)
    galaxy_code = BHTree(converter, number_of_workers=2)
    galaxy_code.parameters.epsilon_squared = (0.01 | units.kpc)**2
    channe_to_galaxy = galaxy_code.particles.new_channel_to(galaxy)
    channe_to_galaxy.copy()
    galaxy_code.particles.add_particles(galaxy)
    inner_stars = galaxy.select(lambda r: r.length()<Rinit,["position"])
    Minner = inner_stars.mass.sum()
    print "Minner=", Minner.in_(units.MSun)
    print "Ninner=", len(inner_stars)
    vc_inner = (constants.G*Minner/Rinit).sqrt()

    converter=nbody_system.nbody_to_si(Mcluster,Rcluster)
    stars=new_king_model(N,W0,convert_nbody=converter)
    masses = new_powerlaw_mass_distribution(N, 0.1|units.MSun, 100|units.MSun, -2.35)
    stars.mass = masses
    stars.scale_to_standard(converter)
    
    stars.x += Rinit
    stars.vy += 0.8*vc_inner
    cluster_code=ph4(converter, number_of_workers=2)
    cluster_code.particles.add_particles(stars)
    channel_to_stars=cluster_code.particles.new_channel_to(stars)

    system=bridge(verbose=False)
    system.add_system(cluster_code, (galaxy_code,))
    system.add_system(galaxy_code, (cluster_code,))
    system.timestep = 0.1*timestep

    times = quantities.arange(0|units.Myr, tend, timestep)
    for i,t in enumerate(times):
        print "Time=", t.in_(units.Myr)
        channe_to_galaxy.copy()
        channel_to_stars.copy()

        inner_stars =  galaxy.select(lambda r: r.length()<Rinit,["position"])
        print "Minner=", inner_stars.mass.sum().in_(units.MSun)

        system.evolve_model(t,timestep=timestep)
    plot_galaxy_and_stars(galaxy, stars)
    galaxy_code.stop()
    cluster_code.stop()

    def test19(self):
        particles = datamodel.Particles(2)
        particles.x = [0.0,10.0] | nbody_system.length
        particles.y = 0.0 | nbody_system.length
        particles.z = 0.0 | nbody_system.length
        particles.radius = 0.005 | nbody_system.length
        particles.vx =  0.0 | nbody_system.speed
        particles.vy =  0.0 | nbody_system.speed
        particles.vz =  0.0 | nbody_system.speed
        particles.mass = 1.0 | nbody_system.mass

        very_short_time_to_evolve = 1 | units.s
        very_long_time_to_evolve = 1e9 | nbody_system.time
       
        instance = BHTree()
        instance.initialize_code()
        instance.parameters.stopping_conditions_timeout = very_short_time_to_evolve 
        self.assertEquals(instance.parameters.stopping_conditions_timeout, very_short_time_to_evolve)
        instance.parameters.epsilon_squared = (0.01 | nbody_system.length)**2
        instance.particles.add_particles(particles) 
        instance.stopping_conditions.timeout_detection.enable()
        start = time.time()
        instance.evolve_model(very_long_time_to_evolve)
        end = time.time()
        self.assertTrue(instance.stopping_conditions.timeout_detection.is_set())
        self.assertTrue((end-start) < very_short_time_to_evolve.value_in(units.s) + 2)#2 = some overhead compensation
        instance.stop()

    def test20(self):
        particles = datamodel.Particles(2)
        particles.x = [0.0,10.0] | nbody_system.length
        particles.y = 0.0 | nbody_system.length
        particles.z = 0.0 | nbody_system.length
        particles.radius = 0.005 | nbody_system.length
        particles.vx =  0.0 | nbody_system.speed
        particles.vy =  0.0 | nbody_system.speed
        particles.vz =  0.0 | nbody_system.speed
        particles.mass = 1.0 | nbody_system.mass

        very_short_time_to_evolve = 1 | units.s
        very_long_time_to_evolve = 1e9 | nbody_system.time
       
        instance = BHTree()
        instance.initialize_code()
        instance.parameters.stopping_conditions_timeout = very_short_time_to_evolve 
        self.assertEquals(instance.parameters.stopping_conditions_timeout, very_short_time_to_evolve)
        instance.parameters.epsilon_squared = (0.01 | nbody_system.length)**2
        instance.particles.add_particles(particles) 
        codeparticles1 = instance.particles
        instance.particles.add_particle(datamodel.Particle(
            position = [0,1,2] |  nbody_system.length, 
            velocity = [0,0,0] |  nbody_system.speed,
            radius = 0.005 | nbody_system.length,
            mass = 1 | nbody_system.mass
        ))
        codeparticles2 = instance.particles
        self.assertTrue(codeparticles1 is codeparticles2)
        instance.cleanup_code()
        codeparticles3 = instance.particles
        self.assertFalse(codeparticles1 is codeparticles3)
    
        instance.stop()

def nbody_integrator(Ncl, mcl, rcl, t_end, n_steps, escape_velocity_fraction, R):
    converter = nbody_system.nbody_to_si(mcl, rcl)
    bodies = new_plummer_model(Ncl, convert_nbody=converter)

    #estimate of milky way mass by "Mass models of the Milky Way", McMillan
    blackhole_mass = 1.26e12 | units.MSun
    blackhole = Particle(mass=blackhole_mass)
    blackhole.position = [0,0,0] | units.m

    cluster_velocity = [0,0,0] | units.m / units.s
    cluster_position = [0,0,0] | units.parsec
    cluster_position[0] = R
    G_si = converter.to_si(nbody_system.G)
    escape_v = (2*G_si*blackhole_mass/R).sqrt().as_quantity_in(units.m/units.s)
    V = escape_v * escape_velocity_fraction
    cluster_velocity[1] = V
    bodies.move_to_center()
    bodies.velocity += cluster_velocity
    bodies.position += cluster_position
    bodies.add_particle(blackhole)
        

    gravity = BHTree(converter)
    gravity.particles.add_particles(bodies)
    channel_from_gravity_to_framework = gravity.particles.\
        new_channel_to(bodies)

    time = zero
    dt = t_end / float(n_steps)
    x = 0
    base_path = "encounter_plots/"+str(R.value_in(units.parsec))+"_"+\
        str(escape_velocity_fraction)+"_"
    while time < t_end:
        plot_cluster(bodies, base_path+str(x),time, rcl, V)
        time += dt
        gravity.evolve_model(time)
        channel_from_gravity_to_framework.copy()
        x+=1
    plot_cluster(bodies, base_path+str(x),time, rcl, V)

    gravity.stop()
    return V, bodies

    def test9(self):
        instance = BHTree()
        instance.initialize_code()
        instance.parameters.epsilon_squared = 0.00001 | nbody_system.length**2
        
        particles = datamodel.Particles(2)
        particles.mass = [1.0, 1.0] | nbody_system.mass
        particles.radius =  [0.0001, 0.0001] | nbody_system.length
        particles.position = [[0.0,0.0,0.0], [2.0,0.0,0.0]] | nbody_system.length
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | nbody_system.speed
        instance.particles.add_particles(particles)
        
        zero = 0.0 | nbody_system.length
        fx, fy, fz = instance.get_gravity_at_point(zero, 1.0 | nbody_system.length, zero, zero)
        self.assertAlmostEqual(fx, 0.0 | nbody_system.acceleration, 3)
        self.assertAlmostEqual(fy, 0.0 | nbody_system.acceleration, 3)
        self.assertAlmostEqual(fz, 0.0 | nbody_system.acceleration, 3)

        for x in (0.25, 0.5, 0.75):
            x0 = x | nbody_system.length
            x1 = (2.0 - x) | nbody_system.length
            potential0 = instance.get_potential_at_point(zero, x0, zero, zero)
            potential1 = instance.get_potential_at_point(zero, x1, zero, zero)
            fx0, fy0, fz0 = instance.get_gravity_at_point(zero, x0, zero, zero)
            fx1, fy1, fz1 = instance.get_gravity_at_point(zero, x1, zero, zero)
            
            self.assertAlmostEqual(fy0, 0.0 | nbody_system.acceleration, 3)
            self.assertAlmostEqual(fz0, 0.0 | nbody_system.acceleration, 3)
            self.assertAlmostEqual(fy1, 0.0 | nbody_system.acceleration, 3)
            self.assertAlmostEqual(fz1, 0.0 | nbody_system.acceleration, 3)
            
            self.assertAlmostEqual(fx0, -1.0 * fx1, 5)
            fx = (-1.0 / (x0**2) + 1.0 / (x1**2)) * (1.0 | nbody_system.length ** 3 / nbody_system.time ** 2)
            self.assertAlmostEqual(fx, fx0, 2)
            self.assertAlmostEqual(potential0, potential1, 5)
        instance.cleanup_code()
        instance.stop()

    def test8(self):
        convert_nbody = nbody_system.nbody_to_si(5.0 | units.kg, 10.0 | units.m)

        instance = BHTree(convert_nbody)
        instance.commit_parameters()
        
        particles = datamodel.Particles(2)
        self.assertEquals(len(instance.particles), 0)
        
        particles.mass = [15.0, 30.0] | units.kg
        particles.radius =  [10.0, 20.0] | units.m
        particles.position = [[10.0, 20.0, 30.0], [20.0, 40.0, 60.0]] | units.m
        particles.velocity = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.m / units.s

        
        instance.particles.add_particles(particles)
        instance.commit_particles()
        
        instance.particles.mass =  [17.0, 33.0] | units.kg
        
        
        self.assertEquals(instance.get_mass(1), 17.0| units.kg) 
        instance.cleanup_code()
        instance.stop()

    def test15(self):
        print "Test15: Testing effect of BHTree parameter epsilon_squared"
        convert_nbody = nbody_system.nbody_to_si(1.0 | units.MSun, 1.0 | units.AU)
        
        particles = datamodel.Particles(2)
        sun = particles[0]
        sun.mass = 1.0 | units.MSun
        sun.position = [0.0, 0.0, 0.0] | units.AU
        sun.velocity = [0.0, 0.0, 0.0] | units.AU / units.yr
        sun.radius = 1.0 | units.RSun

        earth = particles[1]
        earth.mass = 5.9736e24 | units.kg
        earth.radius = 6371.0 | units.km
        earth.position = [0.0, 1.0, 0.0] | units.AU
        earth.velocity = [2.0*numpy.pi, -0.0001, 0.0] | units.AU / units.yr
        
        initial_direction = math.atan((earth.velocity[0]/earth.velocity[1]))
        final_direction = []
        for log_eps2 in range(-9,10,2):
            instance = BHTree(convert_nbody)
            instance.initialize_code()
            instance.parameters.epsilon_squared = 10.0**log_eps2 | units.AU ** 2
            instance.particles.add_particles(particles)
            instance.commit_particles()
            instance.evolve_model(0.25 | units.yr)
            final_direction.append(math.atan((instance.particles[1].velocity[0]/
                instance.particles[1].velocity[1])))
            instance.stop()
        # Small values of epsilon_squared should result in normal earth-sun dynamics: rotation of 90 degrees
        self.assertAlmostEquals(abs(final_direction[0]), abs(initial_direction+math.pi/2.0), 2)
        # Large values of epsilon_squared should result in ~ no interaction
        self.assertAlmostEquals(final_direction[-1], initial_direction, 2)
        # Outcome is most sensitive to epsilon_squared when epsilon_squared = d(earth, sun)^2
        delta = [abs(final_direction[i+1]-final_direction[i]) for i in range(len(final_direction)-1)]
        self.assertEquals(delta[len(final_direction)//2 -1], max(delta))

    def test23(self):
        particles = datamodel.Particles(2)
        particles.x = [0.0,10.0] | nbody_system.length
        particles.y = 0.0 | nbody_system.length
        particles.z = 0.0 | nbody_system.length
        particles.vx =  1.0 | nbody_system.speed
        particles.vy =  0.0 | nbody_system.speed
        particles.vz =  0.0 | nbody_system.speed
        particles.mass = 0.1 | nbody_system.mass

        instance = BHTree(redirection="none")
        instance.particles.add_particles(particles) 
        instance.commit_particles()
        instance.evolve_model(0.1 | nbody_system.time)
        self.assertFalse(instance.particles[0].vy > 0| nbody_system.speed)
        self.assertAlmostRelativeEquals(instance.particles[0].x , 0.1 | nbody_system.length, 4)
        instance.particles.new_channel_to(particles).copy()
        particles.vy = 1| nbody_system.speed
        particles.new_channel_to(instance.particles).copy()
        
        instance.evolve_model(0.2 | nbody_system.time)
        self.assertTrue(instance.particles[0].vy > 0| nbody_system.speed)
        self.assertAlmostRelativeEquals(instance.particles[0].y , 0.1 | nbody_system.length, 4)
        instance.stop()

 def test16(self):
     numpy.random.seed(0)
     number_of_stars = 2
     stars = plummer.new_plummer_model(number_of_stars)
     stars.radius = 0.00001 | nbody_system.length
     stars.scale_to_standard()
     
     instance = BHTree()
     instance.initialize_code()
     instance.parameters.epsilon_squared = (1.0 / 20.0 / (number_of_stars**0.33333) | nbody_system.length)**2
     instance.parameters.timestep = 0.004 | nbody_system.time
     instance.parameters.timestep = 0.00001 | nbody_system.time
     instance.commit_parameters()
     print instance.parameters.timestep
     instance.particles.add_particles(stars)
     instance.commit_particles()
     energy_total_t0 = instance.potential_energy + instance.kinetic_energy
     request = instance.evolve_model.async(1.0 | nbody_system.time)
     request.result()
     energy_total_t1 = instance.potential_energy + instance.kinetic_energy
     
     self.assertAlmostRelativeEqual(energy_total_t0, energy_total_t1, 3)
     instance.stop()
     numpy.random.seed()