Example #1
0
def get_galaxies_in_orbit(m_a=10.e11|units.MSun,
                          m_b=10.e11|units.MSun,
                          ecc=0.5,
                          r_min=25.|units.kpc,
                          t_start=None):
  """
  binary galaxy with orbit of given parameters
  -- if ecc=>1, start at t_start (p625, t_start=-10=-10*100Myr)
  -- if ecc<1, start at apocenter (p644)
  """
  
  converter=nbody_system.nbody_to_si(m_a+m_b,1|units.kpc)
  semi = r_min/(1.-ecc)
  
  # relative position and velocity vectors at the pericenter using kepler
  kepler = Kepler_twobody(converter)
  kepler.initialize_code()
  kepler.initialize_from_elements(mass=(m_a+m_b), semi=semi, ecc=ecc, periastron=r_min) # at periastron
  
  # evolve back till initial position
  if ( ecc<1. ):
    kepler.return_to_apastron()
  else:
    kepler.transform_to_time(t_start)
  
  # get time of the orbit
  t_orbit = kepler.get_time()
  
  rl = kepler.get_separation_vector()
  r = [rl[0].value_in(units.AU), rl[1].value_in(units.AU), rl[2].value_in(units.AU)] | units.AU
  vl = kepler.get_velocity_vector()
  v = [vl[0].value_in(units.kms), vl[1].value_in(units.kms), vl[2].value_in(units.kms)] | units.kms
  
  kepler.stop()
  
  # assign particle atributes
  galaxies = Particles(2)
  galaxies[0].mass = m_a
  galaxies[0].position = (0,0,0) | units.AU
  galaxies[0].velocity = (0,0,0) | units.kms
  galaxies[1].mass = m_b
  galaxies[1].position = r
  galaxies[1].velocity = v
  
  # identification
  galaxies[0].id = 'a0'
  galaxies[1].id = 'b0'
  
  galaxies.move_to_center()
  
  return galaxies, t_orbit
def relative_position_and_velocity_from_orbital_elements(
        mass1, mass2, semimajor_axis, eccentricity, mean_anomaly, seed=None):
    """
  Function that returns relative positions and velocity vectors or orbiters with masses 
  mass2 of the central body with mass mass1 in Cartesian coordinates;
  for vectors of orbital elements -- semi-major axes, eccentricities, mean anomalies.
  3D orientation of orbits (inclination, longitude of ascending node and argument of periapsis) are random.
  (cos(incl) is uniform -1--1, longitude of ascending node and argument of periapsis are uniform 0--2pi)
  Assuming mass1 is static in the center [0,0,0] m, [0,0,0] km/s (that is mass2<<mass1)
  """
    position_vectors = []
    velocity_vectors = []
    converter = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.AU)
    kepler = Kepler(converter)
    kepler.initialize_code()
    r_vec = (0., 0., 0.) | units.AU
    v_vec = (0., 0., 0.) | units.kms
    # to change seed for each particle
    if seed is not None:
        i = 0
    for m2_i, a_i, ecc_i, ma_i in zip(mass2, semimajor_axis, eccentricity,
                                      mean_anomaly):
        #print m2_i, a_i, ecc_i, ma_i
        if seed is not None:
            kepler.set_random(seed + i)
            i = i + 1
        kepler.initialize_from_elements(mass=(mass1 + m2_i),
                                        semi=a_i,
                                        ecc=ecc_i,
                                        mean_anomaly=ma_i,
                                        random_orientation=-1)
        ri = kepler.get_separation_vector()
        vi = kepler.get_velocity_vector()
        # this is to get ~half of the orbits retrograde (that is with inclination
        # of 90--180 degrees) --> velocity = -velocity
        vel_vec_dir = numpy.random.random()
        if (vel_vec_dir <= 0.5):
            vel_orientation = 1.
        else:
            vel_orientation = -1.
        position_vectors.append([ri[0], ri[1], ri[2]])
        velocity_vectors.append([
            vel_orientation * vi[0], vel_orientation * vi[1],
            vel_orientation * vi[2]
        ])
    kepler.stop()
    return position_vectors, velocity_vectors
Example #3
0
def relative_position_and_velocity_from_orbital_elements(mass1,
                                                         mass2,
                                                         semimajor_axis,
                                                         eccentricity,
                                                         mean_anomaly,
                                                         seed=None):
  """
  Function that returns relative positions and velocity vectors or orbiters with masses 
  mass2 of the central body with mass mass1 in Cartesian coordinates;
  for vectors of orbital elements -- semi-major axes, eccentricities, mean anomalies.
  3D orientation of orbits (inclination, longitude of ascending node and argument of periapsis) are random.
  (cos(incl) is uniform -1--1, longitude of ascending node and argument of periapsis are uniform 0--2pi)
  Assuming mass1 is static in the center [0,0,0] m, [0,0,0] km/s (that is mass2<<mass1)
  """
  position_vectors = []
  velocity_vectors = []
  converter = nbody_system.nbody_to_si(1|units.MSun,1|units.AU)
  kepler = Kepler(converter)
  kepler.initialize_code()
  r_vec = (0.,0.,0.) | units.AU
  v_vec = (0.,0.,0.) | units.kms
  # to change seed for each particle
  if seed is not None:
    i=0
  for m2_i, a_i, ecc_i, ma_i in zip(mass2, semimajor_axis, eccentricity, mean_anomaly):
    #print m2_i, a_i, ecc_i, ma_i
    if seed is not None:
      kepler.set_random(seed+i)
      i=i+1
    kepler.initialize_from_elements(mass=(mass1+m2_i),semi=a_i,ecc=ecc_i,mean_anomaly=ma_i,random_orientation=-1)
    ri = kepler.get_separation_vector()
    vi = kepler.get_velocity_vector()
    # this is to get ~half of the orbits retrograde (that is with inclination
    # of 90--180 degrees) --> velocity = -velocity
    vel_vec_dir = numpy.random.random()
    if (vel_vec_dir<=0.5):
      vel_orientation = 1.
    else:
      vel_orientation = -1.
    position_vectors.append([ri[0], ri[1], ri[2]])
    velocity_vectors.append([vel_orientation*vi[0], vel_orientation*vi[1], vel_orientation*vi[2]])
  kepler.stop()
  return position_vectors, velocity_vectors
Example #4
0
def run_kepler(mass, semi, ecc, time):

    kep = Kepler(redirection='none')
    kep.initialize_code()

    kep.set_longitudinal_unit_vector(1.0, 1.0,
                                     0.0)
    kep.initialize_from_elements(mass, semi, ecc)
    a,e = kep.get_elements()
    p = kep.get_periastron()
    print "elements:", a, e, p
    kep.transform_to_time(time)
    x,y,z = kep.get_separation_vector()
    print "separation:", x,y,z
    x,y,z = kep.get_longitudinal_unit_vector()
    print "longitudinal:", x,y,z

    pos = [1, 0, 0] | nbody_system.length
    vel = [0, 0.5, 0] | nbody_system.speed
    kep.initialize_from_dyn(mass, pos[0], pos[1], pos[2],
                            vel[0], vel[1], vel[2])
    a,e = kep.get_elements()
    p = kep.get_periastron()
    print "elements:", a, e, p
    kep.transform_to_time(time)
    x,y,z = kep.get_separation_vector()
    print "separation:", x,y,z
    x,y,z = kep.get_velocity_vector()
    print "velocity:", x,y,z
    x,y,z = kep.get_longitudinal_unit_vector()
    print "longitudinal:", x,y,z

    kep.set_random(42)
    kep.make_binary_scattering(0.5 | nbody_system.mass,
                               0.5,
                               0.5 | nbody_system.mass,
                               0.0 | nbody_system.speed,
                               0.0 | nbody_system.length,
                               1.e-6,
                               0)
    kep.stop()
    def test11(self):
        """
        testing orbital_elements_for_rel_posvel_arrays for unbound orbits
        """
        
        from amuse.community.kepler.interface import Kepler
        
        numpy.random.seed(66)
        N = 10
        
        mass_sun = 1. | units.MSun
        mass1 = numpy.ones(N) * mass_sun
        mass2 = numpy.zeros(N) | units.MSun
        semi_major_axis=-1000.*(random.random(N)) | units.AU 
        eccentricity = (1.+random.random(N))*10.-9.
        inclination = numpy.pi*random.random(N)
        longitude_of_the_ascending_node = 2.*numpy.pi*random.random(N)-numpy.pi
        argument_of_periapsis = 2.*numpy.pi*random.random(N)-numpy.pi      
        
        # kepler.initialize_from_elements initializes orbits with mean_anomaly=0 and true_anomaly=0
        true_anomaly = 0.*(360.*random.random(N)-180.)
        
        comets = datamodel.Particles(N)
        
        converter = nbody_system.nbody_to_si(1|units.MSun,1|units.AU)
        kepler = Kepler(converter)
        kepler.initialize_code()
        for i,arg in enumerate(zip(mass1,mass2,semi_major_axis,eccentricity,true_anomaly,inclination, 
                                   longitude_of_the_ascending_node,argument_of_periapsis)):
            kepler.initialize_from_elements(mass=(mass1[i]+mass2[i]),
                                            semi=semi_major_axis[i],
                                            ecc=eccentricity[i])
            ri = kepler.get_separation_vector()
            vi = kepler.get_velocity_vector()
            
            om = longitude_of_the_ascending_node[i]
            w = argument_of_periapsis[i]
            incl = inclination[i]
            a1 = ([numpy.cos(om), -numpy.sin(om), 0.0], [numpy.sin(om), numpy.cos(om), 0.0], [0.0, 0.0, 1.0])
            a2 = ([1.0, 0.0, 0.0], [0.0, numpy.cos(incl), -numpy.sin(incl)], [0.0, numpy.sin(incl), numpy.cos(incl)])
            a3 = ([numpy.cos(w), -numpy.sin(w), 0.0], [numpy.sin(w), numpy.cos(w), 0.0], [0.0, 0.0, 1.0])
            A = numpy.dot(numpy.dot(a1,a2),a3)

            r_vec = numpy.dot(A,numpy.reshape(ri,3,1))
            v_vec = numpy.dot(A,numpy.reshape(vi,3,1))
          
            r = (0.0, 0.0, 0.0) | units.AU
            v = (0.0, 0.0, 0.0) | (units.AU / units.day)
            r[0] = r_vec[0]
            r[1] = r_vec[1]
            r[2] = r_vec[2]
            v[0] = v_vec[0]
            v[1] = v_vec[1]
            v[2] = v_vec[2]
  
            comets[i].mass = mass2[i]
            comets[i].position = r_vec
            comets[i].velocity = v_vec
        
        kepler.stop()
        
        semi_major_axis_ext, eccentricity_ext, ta_ext, inclination_ext, \
        longitude_of_the_ascending_node_ext, argument_of_periapsis_ext = \
        orbital_elements(comets.position,
                                               comets.velocity,
                                               comets.mass + mass_sun,
                                               G=constants.G)
        
        self.assertAlmostEqual(semi_major_axis,semi_major_axis_ext.in_(units.AU))
        self.assertAlmostEqual(eccentricity,eccentricity_ext)
        self.assertAlmostEqual(inclination,inclination_ext)
        self.assertAlmostEqual(longitude_of_the_ascending_node,longitude_of_the_ascending_node_ext)
        self.assertAlmostEqual(argument_of_periapsis,argument_of_periapsis_ext)
        self.assertAlmostEqual(true_anomaly,ta_ext)
Example #6
0
def run_ph4(infile = None, outfile = None,
             number_of_stars = 100, number_of_binaries = 0,
             end_time = 10 | nbody_system.time,
             delta_t = 1 | nbody_system.time,
             n_workers = 1, use_gpu = 1, gpu_worker = 1,
             salpeter = 0,
             accuracy_parameter = 0.1,
             softening_length = 0.0 | nbody_system.length,
             manage_encounters = 1, random_seed = 1234):

    if random_seed <= 0:
        numpy.random.seed()
        random_seed = numpy.random.randint(1, pow(2,31)-1)
    numpy.random.seed(random_seed)
    print "random seed =", random_seed

    if infile != None: print "input file =", infile
    print "end_time =", end_time.number
    print "delta_t =", delta_t.number
    print "n_workers =", n_workers
    print "use_gpu =", use_gpu
    print "manage_encounters =", manage_encounters
    print "\ninitializing the gravity module"
    sys.stdout.flush()

    init_smalln()
    
    # Note that there are actually three GPU options:
    #
    #	1. use the GPU code and allow GPU use (default)
    #	2. use the GPU code but disable GPU use (-g)
    #	3. use the non-GPU code (-G)

    if gpu_worker == 1:
        try:
            gravity = grav(number_of_workers = n_workers,
                           redirection = "none", mode = "gpu")
        except Exception as ex:
            gravity = grav(number_of_workers = n_workers,
                           redirection = "none")
    else:
        gravity = grav(number_of_workers = n_workers,
                       redirection = "none")

    gravity.initialize_code()
    gravity.parameters.set_defaults()

    #-----------------------------------------------------------------

    if infile == None:

        print "making a Plummer model"
        stars = new_plummer_model(number_of_stars)

        id = numpy.arange(number_of_stars)
        stars.id = id+1

        print "setting particle masses and radii"
        if salpeter == 0:
            print 'equal masses'
            total_mass = 1.0 | nbody_system.mass
            scaled_mass = total_mass / number_of_stars
        else:
            print 'salpeter mass function'
            scaled_mass = new_salpeter_mass_distribution_nbody(number_of_stars) 
        stars.mass = scaled_mass

        print "centering stars"
        stars.move_to_center()
        print "scaling stars to virial equilibrium"
        stars.scale_to_standard(smoothing_length_squared
                                    = gravity.parameters.epsilon_squared)

    else:

        # Read the input data.  Units are dynamical (sorry).
        # Format:  id  mass  pos[3]  vel[3]

        print "reading file", infile

        id = []
        mass = []
        pos = []
        vel = []

        f = open(infile, 'r')
        count = 0
        for line in f:
            if len(line) > 0:
                count += 1
                cols = line.split()
                if count == 1: snap = int(cols[0])
                elif count == 2: number_of_stars = int(cols[0])
                elif count == 3: time = float(cols[0]) | nbody_system.time
                else:
                    if len(cols) >= 8:
                        id.append(int(cols[0]))
                        mass.append(float(cols[1]))
                        pos.append((float(cols[2]),
                                    float(cols[3]), float(cols[4])))
                        vel.append((float(cols[5]),
                                    float(cols[6]), float(cols[7])))
        f.close()

        stars = datamodel.Particles(number_of_stars)
        stars.id = id
        stars.mass = mass | nbody_system.mass
        stars.position = pos | nbody_system.length
        stars.velocity = vel | nbody_system.speed
        #stars.radius = 0. | nbody_system.length

    total_mass = stars.mass.sum()
    ke = pa.kinetic_energy(stars)
    kT = ke/(1.5*number_of_stars)

    if number_of_binaries > 0:

        # Turn selected stars into binary components.
        # Only tested for equal-mass case.

        kep = Kepler(redirection = "none")
        kep.initialize_code()

        added_mass = 0.0 | nbody_system.mass

        # Work with energies rather than semimajor axes.

        Emin = 10*kT
        Emax = 20*kT
        ecc = 0.1

        id_count = number_of_stars
        nbin = 0
        for i in range(0, number_of_stars,
                       number_of_stars/number_of_binaries):

            # Star i is CM, becomes component, add other star at end.

            nbin += 1

            mass = stars[i].mass
            new_mass = numpy.random.uniform()*mass	# uniform q?
            mbin = mass + new_mass
            fac = new_mass/mbin
            E = Emin + numpy.random.uniform()*(Emax-Emin)
            a = 0.5*nbody_system.G*mass*new_mass/E

            kep.initialize_from_elements(mbin, a, ecc)
            dr = quantities.AdaptingVectorQuantity()
            dr.extend(kep.get_separation_vector())
            dv = quantities.AdaptingVectorQuantity()
            dv.extend(kep.get_velocity_vector())

            newstar = datamodel.Particles(1)
            newstar.mass = new_mass
            newstar.position = stars[i].position + (1-fac)*dr
            newstar.velocity = stars[i].velocity + (1-fac)*dv
            # stars[i].mass = mass
            stars[i].position = stars[i].position - fac*dr
            stars[i].velocity = stars[i].velocity - fac*dv

            id_count += 1
            newstar.id = id_count
            stars.add_particles(newstar)
            added_mass += new_mass

            if nbin >= number_of_binaries: break

        kep.stop()

        print 'created', nbin, 'binaries'
        sys.stdout.flush()

        stars.mass = stars.mass * total_mass/(total_mass+added_mass)
        number_of_stars += nbin

    # Set dynamical radii (assuming virial equilibrium and standard
    # units).  Note that this choice should be refined, and updated
    # as the system evolves.  Probably the choice of radius should be
    # made entirely in the multiples module.  TODO.  In these units,
    # M = 1 and <v^2> = 0.5, so the mean 90-degree turnaround impact
    # parameter is
    #
    #		b_90 = G (m_1+m_2) / vrel^2
    #		     = 2 <m> / 2<v^2>
    #		     = 2 / N			for equal masses
    #
    # Taking r_i = m_i / 2<v^2> = m_i in virial equilibrium means
    # that, approximately, "contact" means a 90-degree deflection (r_1
    # + r_2 = b_90).  A more conservative choice with r_i less than
    # this value will isolates encounters better, but also place more
    # load on the large-N dynamical module.

    stars.radius = stars.mass.number | nbody_system.length

    time = 0.0 | nbody_system.time
    # print "IDs:", stars.id.number

    print "recentering stars"
    stars.move_to_center()
    sys.stdout.flush()

    #-----------------------------------------------------------------

    if softening_length < 0.0 | nbody_system.length:

        # Use ~interparticle spacing.  Assuming standard units here.  TODO

        eps2 = 0.25*(float(number_of_stars))**(-0.666667) \
			| nbody_system.length**2
    else:
        eps2 = softening_length*softening_length
    print 'softening length =', eps2.sqrt()

    gravity.parameters.timestep_parameter = accuracy_parameter
    gravity.parameters.epsilon_squared = eps2
    gravity.parameters.use_gpu = use_gpu
    # gravity.parameters.manage_encounters = manage_encounters

    print ''
    print "adding particles"
    # print stars
    sys.stdout.flush()
    gravity.particles.add_particles(stars)
    gravity.commit_particles()

    print ''
    print "number_of_stars =", number_of_stars
    print "evolving to time =", end_time.number, \
          "in steps of", delta_t.number
    sys.stdout.flush()

    # Channel to copy values from the code to the set in memory.
    channel = gravity.particles.new_channel_to(stars)

    stopping_condition = gravity.stopping_conditions.collision_detection
    stopping_condition.enable()

    # Debugging: prevent the multiples code from being called.
    if 0:
        stopping_condition.disable()
        print 'stopping condition disabled'
        sys.stdout.flush()

    # -----------------------------------------------------------------
    # Create the coupled code and integrate the system to the desired
    # time, managing interactions internally.

    kep = init_kepler(stars[0], stars[1])
    multiples_code = multiples.Multiples(gravity, new_smalln, kep)

    multiples_code.neighbor_perturbation_limit = 0.1
    #multiples_code.neighbor_distance_factor = 1.0
    #multiples_code.neighbor_veto = False
    #multiples_code.neighbor_distance_factor = 2.0
    multiples_code.neighbor_veto = True

    print ''
    print 'multiples_code.initial_scale_factor =', \
        multiples_code.initial_scale_factor
    print 'multiples_code.neighbor_perturbation_limit =', \
        multiples_code.neighbor_perturbation_limit
    print 'multiples_code.neighbor_veto =', \
        multiples_code.neighbor_veto
    print 'multiples_code.final_scale_factor =', \
        multiples_code.final_scale_factor
    print 'multiples_code.initial_scatter_factor =', \
        multiples_code.initial_scatter_factor
    print 'multiples_code.final_scatter_factor =', \
        multiples_code.final_scatter_factor
    print 'multiples_code.retain_binary_apocenter =', \
        multiples_code.retain_binary_apocenter
    print 'multiples_code.wide_perturbation_limit =', \
        multiples_code.wide_perturbation_limit

    pre = "%%% "
    E0,cpu0 = print_log(pre, time, multiples_code)

    while time < end_time:

        time += delta_t
        multiples_code.evolve_model(time)

        # Copy values from the module to the set in memory.

        channel.copy()
    
        # Copy the index (ID) as used in the module to the id field in
        # memory.  The index is not copied by default, as different
        # codes may have different indices for the same particle and
        # we don't want to overwrite silently.

        channel.copy_attribute("index_in_code", "id")

        print_log(pre, time, multiples_code, E0, cpu0)
        sys.stdout.flush()

    #-----------------------------------------------------------------

    if not outfile == None:

        # Write data to a file.

        f = open(outfile, 'w')

        #--------------------------------------------------
        # Need to save top-level stellar data and parameters.
        # Need to save multiple data and parameters.

        f.write('%.15g\n'%(time.number))
        for s in multiples_code.stars: write_star(s, f)

        #--------------------------------------------------

        f.close()
        print 'wrote file', outfile

    print ''
    gravity.stop()
Example #7
0
def run_ph4(infile = None, outfile = None,
            number_of_stars = 100, number_of_binaries = 0,
            end_time = 10 | nbody_system.time,
            delta_t = 1 | nbody_system.time,
            n_workers = 1, use_gpu = 1, gpu_worker = 1,
            salpeter = 0,
            accuracy_parameter = 0.1,
            softening_length = 0.0 | nbody_system.length,
            manage_encounters = 1, random_seed = 1234):

    if random_seed <= 0:
        numpy.random.seed()
        random_seed = numpy.random.randint(1, pow(2,31)-1)
    numpy.random.seed(random_seed)
    print "random seed =", random_seed

    if infile != None: print "input file =", infile
    print "end_time =", end_time.number
    print "delta_t =", delta_t.number
    print "n_workers =", n_workers
    print "use_gpu =", use_gpu
    print "manage_encounters =", manage_encounters
    print "\ninitializing the gravity module"
    sys.stdout.flush()

    init_smalln()
    
    # Note that there are actually three GPU options:
    #
    #	1. use the GPU code and allow GPU use (default)
    #	2. use the GPU code but disable GPU use (-g)
    #	3. use the non-GPU code (-G)

    if gpu_worker == 1:
        try:
            gravity = grav(number_of_workers = n_workers,
                           redirection = "none", mode = "gpu")
        except Exception as ex:
            gravity = grav(number_of_workers = n_workers,
                           redirection = "none")
    else:
        gravity = grav(number_of_workers = n_workers,
                       redirection = "none")

    gravity.initialize_code()
    gravity.parameters.set_defaults()

    #-----------------------------------------------------------------

    if infile == None:

        print "making a Plummer model"
        stars = new_plummer_model(number_of_stars)

        id = numpy.arange(number_of_stars)
        stars.id = id+1

        print "setting particle masses and radii"
        if salpeter == 0:
            print 'equal masses'
            total_mass = 1.0 | nbody_system.mass
            scaled_mass = total_mass / number_of_stars
        else:
            print 'salpeter mass function'
            scaled_mass = new_salpeter_mass_distribution_nbody(number_of_stars) 
        stars.mass = scaled_mass

        print "centering stars"
        stars.move_to_center()
        print "scaling stars to virial equilibrium"
        stars.scale_to_standard(smoothing_length_squared
                                    = gravity.parameters.epsilon_squared)

    else:

        # Read the input data.  Units are dynamical (sorry).
        # Format:  id  mass  pos[3]  vel[3]

        print "reading file", infile

        id = []
        mass = []
        pos = []
        vel = []

        f = open(infile, 'r')
        count = 0
        for line in f:
            if len(line) > 0:
                count += 1
                cols = line.split()
                if count == 1: snap = int(cols[0])
                elif count == 2: number_of_stars = int(cols[0])
                elif count == 3: time = float(cols[0]) | nbody_system.time
                else:
                    if len(cols) >= 8:
                        id.append(int(cols[0]))
                        mass.append(float(cols[1]))
                        pos.append((float(cols[2]),
                                    float(cols[3]), float(cols[4])))
                        vel.append((float(cols[5]),
                                    float(cols[6]), float(cols[7])))
        f.close()

        stars = datamodel.Particles(number_of_stars)
        stars.id = id
        stars.mass = mass | nbody_system.mass
        stars.position = pos | nbody_system.length
        stars.velocity = vel | nbody_system.speed
        #stars.radius = 0. | nbody_system.length

    total_mass = stars.mass.sum()
    ke = pa.kinetic_energy(stars)
    kT = ke/(1.5*number_of_stars)

    if number_of_binaries > 0:

        # Turn selected stars into binary components.
        # Only tested for equal-mass case.

        kep = Kepler(redirection = "none")
        kep.initialize_code()

        added_mass = 0.0 | nbody_system.mass

        # Work with energies rather than semimajor axes.

        Emin = 10*kT
        Emax = 20*kT
        ecc = 0.1

        id_count = number_of_stars
        nbin = 0
        for i in range(0, number_of_stars,
                       number_of_stars/number_of_binaries):

            # Star i is CM, becomes component, add other star at end.

            nbin += 1

            mass = stars[i].mass
            new_mass = numpy.random.uniform()*mass	# uniform q?
            mbin = mass + new_mass
            fac = new_mass/mbin
            E = Emin + numpy.random.uniform()*(Emax-Emin)
            a = 0.5*nbody_system.G*mass*new_mass/E

            kep.initialize_from_elements(mbin, a, ecc)
            dr = quantities.AdaptingVectorQuantity()
            dr.extend(kep.get_separation_vector())
            dv = quantities.AdaptingVectorQuantity()
            dv.extend(kep.get_velocity_vector())

            newstar = datamodel.Particles(1)
            newstar.mass = new_mass
            newstar.position = stars[i].position + (1-fac)*dr
            newstar.velocity = stars[i].velocity + (1-fac)*dv
            # stars[i].mass = mass
            stars[i].position = stars[i].position - fac*dr
            stars[i].velocity = stars[i].velocity - fac*dv

            id_count += 1
            newstar.id = id_count
            stars.add_particles(newstar)
            added_mass += new_mass

            if nbin >= number_of_binaries: break

        kep.stop()

        print 'created', nbin, 'binaries'
        sys.stdout.flush()

        stars.mass = stars.mass * total_mass/(total_mass+added_mass)
        number_of_stars += nbin

    # Set dynamical radii (assuming virial equilibrium and standard
    # units).  Note that this choice should be refined, and updated
    # as the system evolves.  Probably the choice of radius should be
    # made entirely in the multiples module.  TODO.  In these units,
    # M = 1 and <v^2> = 0.5, so the mean 90-degree turnaround impact
    # parameter is
    #
    #		b_90 = G (m_1+m_2) / vrel^2
    #		     = 2 <m> / 2<v^2>
    #		     = 2 / N			for equal masses
    #
    # Taking r_i = m_i / 2<v^2> = m_i in virial equilibrium means
    # that, approximately, "contact" means a 90-degree deflection (r_1
    # + r_2 = b_90).  A more conservative choice with r_i less than
    # this value will isolates encounters better, but also place more
    # load on the large-N dynamical module.

    stars.radius = stars.mass.number | nbody_system.length

    time = 0.0 | nbody_system.time
    # print "IDs:", stars.id.number

    print "recentering stars"
    stars.move_to_center()
    sys.stdout.flush()

    #-----------------------------------------------------------------

    if softening_length < 0.0 | nbody_system.length:

        # Use ~interparticle spacing.  Assuming standard units here.  TODO

        eps2 = 0.25*(float(number_of_stars))**(-0.666667) \
			| nbody_system.length**2
    else:
        eps2 = softening_length*softening_length
    print 'softening length =', eps2.sqrt()

    gravity.parameters.timestep_parameter = accuracy_parameter
    gravity.parameters.epsilon_squared = eps2
    gravity.parameters.use_gpu = use_gpu
    # gravity.parameters.manage_encounters = manage_encounters

    print ''
    print "adding particles"
    # print stars
    sys.stdout.flush()
    gravity.particles.add_particles(stars)
    gravity.commit_particles()

    print ''
    print "number_of_stars =", number_of_stars
    print "evolving to time =", end_time.number, \
          "in steps of", delta_t.number
    sys.stdout.flush()

    # Channel to copy values from the code to the set in memory.
    channel = gravity.particles.new_channel_to(stars)

    stopping_condition = gravity.stopping_conditions.collision_detection
    stopping_condition.enable()

    # Debugging: prevent the multiples code from being called.
    if 0:
        stopping_condition.disable()
        print 'stopping condition disabled'
        sys.stdout.flush()

    # -----------------------------------------------------------------
    # Create the coupled code and integrate the system to the desired
    # time, managing interactions internally.

    kep = init_kepler(stars[0], stars[1])
    multiples_code = multiples.Multiples(gravity, new_smalln, kep)

    multiples_code.neighbor_perturbation_limit = 0.1
    #multiples_code.neighbor_distance_factor = 1.0
    #multiples_code.neighbor_veto = False
    #multiples_code.neighbor_distance_factor = 2.0
    multiples_code.neighbor_veto = True

    print ''
    print 'multiples_code.initial_scale_factor =', \
        multiples_code.initial_scale_factor
    print 'multiples_code.neighbor_perturbation_limit =', \
        multiples_code.neighbor_perturbation_limit
    print 'multiples_code.neighbor_veto =', \
        multiples_code.neighbor_veto
    print 'multiples_code.final_scale_factor =', \
        multiples_code.final_scale_factor
    print 'multiples_code.initial_scatter_factor =', \
        multiples_code.initial_scatter_factor
    print 'multiples_code.final_scatter_factor =', \
        multiples_code.final_scatter_factor
    print 'multiples_code.retain_binary_apocenter =', \
        multiples_code.retain_binary_apocenter
    print 'multiples_code.wide_perturbation_limit =', \
        multiples_code.wide_perturbation_limit

    pre = "%%% "
    E0,cpu0 = print_log(pre, time, multiples_code)

    while time < end_time:

        time += delta_t
        multiples_code.evolve_model(time)

        # Copy values from the module to the set in memory.

        channel.copy()
    
        # Copy the index (ID) as used in the module to the id field in
        # memory.  The index is not copied by default, as different
        # codes may have different indices for the same particle and
        # we don't want to overwrite silently.

        channel.copy_attribute("index_in_code", "id")

        print_log(pre, time, multiples_code, E0, cpu0)
        sys.stdout.flush()

    #-----------------------------------------------------------------

    if not outfile == None:

        # Write data to a file.

        f = open(outfile, 'w')

        #--------------------------------------------------
        # Need to save top-level stellar data and parameters.
        # Need to save multiple data and parameters.

        f.write('%.15g\n'%(time.number))
        for s in multiples_code.stars: write_star(s, f)

        #--------------------------------------------------

        f.close()
        print 'wrote file', outfile

    print ''
    gravity.stop()
Example #8
0
def compress_binary_components(comp1, comp2, scale):

    # Compress the two-body system consisting of comp1 and comp2 to
    # lie within distance scale of one another.

    pos1 = comp1.position
    pos2 = comp2.position
    sep12 = ((pos2-pos1)**2).sum()

    if sep12 > scale*scale:
        print '\ncompressing components', int(comp1.id.number), \
              'and', int(comp2.id.number), 'to separation', scale.number
        sys.stdout.flush()
        mass1 = comp1.mass
        mass2 = comp2.mass
        total_mass = mass1 + mass2
        vel1 = comp1.velocity
        vel2 = comp2.velocity
        cmpos = (mass1*pos1+mass2*pos2)/total_mass
        cmvel = (mass1*vel1+mass2*vel2)/total_mass

        # For now, create and delete a temporary kepler
        # process to handle the transformation.  Obviously
        # more efficient to define a single kepler at the
        # start of the calculation and reuse it.

        kep = Kepler(redirection = "none")
        kep.initialize_code()
        mass = comp1.mass + comp2.mass
        rel_pos = pos2 - pos1
        rel_vel = vel2 - vel1
        kep.initialize_from_dyn(mass,
                                rel_pos[0], rel_pos[1], rel_pos[2],
                                rel_vel[0], rel_vel[1], rel_vel[2])
        M,th = kep.get_angles()
        a,e = kep.get_elements()
        if e < 1:
            peri = a*(1-e)
            apo = a*(1+e)
        else:
            peri = a*(e-1)
            apo = 2*a		# OK - used ony to reset scale
        limit = peri + 0.01*(apo-peri)
        if scale < limit: scale = limit

        if M < 0:
            # print 'approaching'
            kep.advance_to_periastron()
            kep.advance_to_radius(limit)
        else:
            # print 'receding'
            if kep.get_separation() < scale:
                kep.advance_to_radius(limit)
            else:
                kep.return_to_radius(scale)

        # a,e = kep.get_elements()
        # r = kep.get_separation()
        # E,J = kep.get_integrals()
        # print 'kepler: a,e,r =', a.number, e.number, r.number
        # print 'E, J =', E, J

        # Note: if periastron > scale, we are now just past periastron.

        new_rel_pos = kep.get_separation_vector()
        new_rel_vel = kep.get_velocity_vector()
        kep.stop()

        # Enew = 0
        # r2 = 0
        # for k in range(3):
        #     Enew += 0.5*(new_rel_vel[k].number)**2
        #     r2 += (new_rel_pos[k].number)**2
        # rnew = math.sqrt(r2)
        # Enew -= mass.number/r1
        # print 'E, Enew, rnew =', E.number, E1, r1

        # Problem: the vectors returned by kepler are lists,
        # not numpy arrays, and it looks as though we can say
        # comp1.position = pos, but not comp1.position[k] =
        # xxx, as we'd like...  Also, we don't know how to
        # copy a numpy array with units...  TODO

        newpos1 = pos1 - pos1	# stupid trick to create zero vectors
        newpos2 = pos2 - pos2	# with the proper form and units...
        newvel1 = vel1 - vel1
        newvel2 = vel2 - vel2

        frac2 = mass2/total_mass
        for k in range(3):
            dxk = new_rel_pos[k]
            dvk = new_rel_vel[k]
            newpos1[k] = cmpos[k] - frac2*dxk
            newpos2[k] = cmpos[k] + (1-frac2)*dxk
            newvel1[k] = cmvel[k] - frac2*dvk
            newvel2[k] = cmvel[k] + (1-frac2)*dvk

        # Perform the changes to comp1 and comp2, and recursively
        # transmit them to the (currently absolute) coordinates of
        # all lower components.

        offset_particle_tree(comp1, newpos1-pos1, newvel1-vel1)
        offset_particle_tree(comp2, newpos2-pos2, newvel2-vel2)
Example #9
0
def CutOrAdvance(enc_bodies, primary_sysID, converter=None):
    bodies = enc_bodies.copy()
    if converter==None:
        converter = nbody_system.nbody_to_si(bodies.mass.sum(), 2 * np.max(bodies.radius.number) | bodies.radius.unit)
    systems = stellar_systems.get_heirarchical_systems_from_set(bodies, converter=converter, RelativePosition=False)
    # Deal with Possible Key Issues with Encounters with 3+ Star Particles Being Run More than Other Systems ...
    if int(primary_sysID) not in systems.keys():
        print "...: Error: Previously run binary system has been found! Not running this system ..."
        print primary_sysID
        print systems.keys()
        print "---------------------------------"
        return None
    # As this function is pulling from Multiples, there should never be more or less than 2 "Root" Particles ...
    if len(systems) != 2:
        print "...: Error: Encounter has more roots than expected! Total Root Particles:", len(systems)
        print bodies
        print "---------------------------------"
        return None
    # Assign the Primary System to #1 and Perturbing System to #2
    sys_1 = systems[int(primary_sysID)]
    secondary_sysID = [key for key in systems.keys() if key!=int(primary_sysID)][0]
    sys_2 = systems[secondary_sysID]
    print 'All System Keys:', systems.keys()
    print 'Primary System Key:', primary_sysID
    print 'System 1 IDs:', sys_1.id
    print 'System 2 IDs:', sys_2.id
    # Calculate Useful Quantities
    mass_ratio = sys_2.mass.sum()/sys_1.mass.sum()
    total_mass = sys_1.mass.sum() + sys_2.mass.sum()
    rel_pos = sys_1.center_of_mass() - sys_2.center_of_mass()
    rel_vel = sys_1.center_of_mass_velocity() - sys_2.center_of_mass_velocity()
    # Initialize Kepler Worker
    kep = Kepler(unit_converter = converter, redirection = 'none')
    kep.initialize_code()
    kep.initialize_from_dyn(total_mass, rel_pos[0], rel_pos[1], rel_pos[2], rel_vel[0], rel_vel[1], rel_vel[2])
    # Check to See if the Periastron is within the Ignore Distance for 10^3 Perturbation
    p = kep.get_periastron()
    ignore_distance = mass_ratio**(1./3.) * 600 | units.AU
    if p > ignore_distance:
        print "Encounter Ignored due to Periastron of", p.in_(units.AU), "and an IgnoreDistance of",ignore_distance
        kep.stop()
        print "---------------------------------"
        return None
    # Move the Particles to be Relative to their Respective Center of Mass
    cm_sys_1, cm_sys_2 = sys_1.center_of_mass(), sys_2.center_of_mass()
    cmv_sys_1, cmv_sys_2 = sys_1.center_of_mass_velocity(), sys_2.center_of_mass_velocity()
    for particle in sys_1:
        particle.position -= cm_sys_1
        particle.velocity -= cmv_sys_1
    for particle in sys_2:
        particle.position -= cm_sys_2
        particle.velocity -= cmv_sys_2
    # Check to See if the Planets are Closer than the Ignore Distance
    # Note: This shouldn't happen in the main code, but this prevents overshooting the periastron in debug mode.
    if kep.get_separation() > ignore_distance:
        kep.advance_to_radius(ignore_distance)
    # Advance the Center of Masses to the Desired Distance in Reduced Mass Coordinates
    x, y, z = kep.get_separation_vector()
    rel_pos_f = rel_pos.copy()
    rel_pos_f[0], rel_pos_f[1], rel_pos_f[2] = x, y, z
    vx, vy, vz = kep.get_velocity_vector()
    rel_vel_f = rel_vel.copy()
    rel_vel_f[0], rel_vel_f[1], rel_vel_f[2] = vx, vy, vz
    # Transform to Absolute Coordinates from Kepler Reduced Mass Coordinates
    cm_pos_1, cm_pos_2 = sys_2.mass.sum() * rel_pos_f / total_mass, -sys_1.mass.sum() * rel_pos_f / total_mass
    cm_vel_1, cm_vel_2 = sys_2.mass.sum() * rel_vel_f / total_mass, -sys_1.mass.sum() * rel_vel_f / total_mass
    # Move the Particles to the New Postions of their Respective Center of Mass
    for particle in sys_1:
        particle.position += cm_pos_1
        particle.velocity += cm_vel_1
    for particle in sys_2:
        particle.position += cm_pos_2
        particle.velocity += cm_vel_2
    # Stop Kepler and Return the Systems as a Particle Set
    kep.stop()
    # Collect the Collective Particle Set to be Returned Back
    final_set = Particles()
    final_set.add_particles(sys_1)
    final_set.add_particles(sys_2)
    print "---------------------------------"
    return final_set
def run_ph4(options,
            time=None,
            stars=None,
            mc_root_to_tree=None,
            randomize=True):
    infile = options.infile
    outfile = options.outfile
    restart_file = options.restart_file
    number_of_stars = options.N
    number_of_binaries = options.Nbin
    end_time = options.t_end | nbody_system.time
    delta_t = options.delta_t | nbody_system.time
    n_workers = options.n_workers
    use_gpu = options.use_gpu
    gpu_worker = options.gpu_worker
    salpeter = options.salpeter
    accuracy_parameter = options.accuracy_parameter
    softening_length = options.softening_length | nbody_system.length
    manage_encounters = options.manage_encounters
    random_seed = options.random_seed

    if randomize:
        if random_seed <= 0:
            numpy.random.seed()
            random_seed = numpy.random.randint(1, pow(2, 31) - 1)
        numpy.random.seed(random_seed)
        print "random seed =", random_seed

    if infile is not None: print "input file =", infile
    if restart_file is not None: print "restart file =", restart_file
    if restart_file is not None and infile is not None:
        print "restart file overrides input file"
    print "end_time =", end_time.number
    print "delta_t =", delta_t.number
    print "n_workers =", n_workers
    print "use_gpu =", use_gpu
    print "manage_encounters =", manage_encounters
    print "n =", number_of_stars
    print "nbin=", number_of_binaries
    print "\ninitializing the gravity module"
    sys.stdout.flush()

    init_smalln()

    # Note that there are actually three GPU options:
    #
    #	1. use the GPU code and allow GPU use (default)
    #	2. use the GPU code but disable GPU use (-g)
    #	3. use the non-GPU code (-G)

    if gpu_worker == 1:
        try:
            #gravity = GravityModule(number_of_workers = n_workers,
            #               redirection = "xterm")
            gravity = GravityModule(number_of_workers=n_workers,
                                    redirection="none",
                                    mode="gpu")
        except Exception as ex:
            gravity = GravityModule(number_of_workers=n_workers,
                                    redirection="none")
    else:
        gravity = GravityModule(number_of_workers=n_workers,
                                redirection="none")

    gravity.initialize_code()
    gravity.parameters.set_defaults()

    if softening_length < 0.0 | nbody_system.length:

        # Use ~interparticle spacing.  Assuming standard units here.  TODO

        eps2 = 0.25*(float(number_of_stars))**(-0.666667) \
            | nbody_system.length**2
    else:
        eps2 = softening_length * softening_length
    print 'softening length =', eps2.sqrt()

    gravity.parameters.timestep_parameter = accuracy_parameter
    gravity.parameters.epsilon_squared = eps2
    gravity.parameters.use_gpu = use_gpu

    kep = Kepler(redirection="none")
    kep.initialize_code()

    multiples_code = None
    Xtra = numpy.zeros(2)

    #-----------------------------------------------------------------

    if (restart_file is None
            or not os.path.exists(restart_file + ".stars.hdf5")
        ) and infile is None and stars is None:

        print "making a Plummer model"
        stars = new_plummer_model(number_of_stars)

        id = numpy.arange(number_of_stars)
        stars.id = id + 1

        print "setting particle masses and radii"
        if salpeter == 0:
            print 'equal masses'
            total_mass = 1.0 | nbody_system.mass
            scaled_mass = total_mass / number_of_stars
        else:
            print 'salpeter mass function'
            scaled_mass = new_salpeter_mass_distribution_nbody(number_of_stars)
        stars.mass = scaled_mass

        print "centering stars"
        stars.move_to_center()
        print "scaling stars to virial equilibrium"
        stars.scale_to_standard(
            smoothing_length_squared=gravity.parameters.epsilon_squared)
        time = 0.0 | nbody_system.time

        total_mass = stars.mass.sum()
        ke = pa.kinetic_energy(stars)
        kT = ke / (1.5 * number_of_stars)

        # Set dynamical radii (assuming virial equilibrium and standard
        # units).  Note that this choice should be refined, and updated
        # as the system evolves.  Probably the choice of radius should be
        # made entirely in the multiples module.  TODO.  In these units,
        # M = 1 and <v^2> = 0.5, so the mean 90-degree turnaround impact
        # parameter is
        #
        #		b_90 = G (m_1+m_2) / vrel^2
        #		     = 2 <m> / 2<v^2>
        #		     = 2 / N			for equal masses
        #
        # Taking r_i = m_i / 2<v^2> = m_i in virial equilibrium means
        # that, approximately, "contact" means a 90-degree deflection (r_1
        # + r_2 = b_90).  A more conservative choice with r_i less than
        # this value will isolates encounters better, but also place more
        # load on the large-N dynamical module.

        stars.radius = stars.mass.number | nbody_system.length

        if number_of_binaries > 0:

            # Turn selected stars into binary components.
            # Only tested for equal-mass case.

            added_mass = 0.0 | nbody_system.mass

            # Work with energies rather than semimajor axes.

            Emin = 10 * kT
            Emax = 20 * kT
            ecc = 0.1

            id_count = number_of_stars
            nbin = 0
            for i in range(0, number_of_stars,
                           number_of_stars / number_of_binaries):

                # Star i is CM, becomes component, add other star at end.

                nbin += 1

                mass = stars[i].mass
                #new_mass = numpy.random.uniform()*mass	# uniform q?
                new_mass = mass  # uniform q?
                mbin = mass + new_mass
                fac = new_mass / mbin
                E = Emin + numpy.random.uniform() * (Emax - Emin)
                a = 0.5 * nbody_system.G * mass * new_mass / E

                kep.initialize_from_elements(mbin, a, ecc)
                dr = quantities.AdaptingVectorQuantity()
                dr.extend(kep.get_separation_vector())
                dv = quantities.AdaptingVectorQuantity()
                dv.extend(kep.get_velocity_vector())

                newstar = datamodel.Particles(1)
                newstar.mass = new_mass
                newstar.position = stars[i].position + (1 - fac) * dr
                newstar.velocity = stars[i].velocity + (1 - fac) * dv
                newstar.radius = newstar.mass.number | nbody_system.length
                #newstar.radius = 3.0*stars[i].radius    # HACK: try to force collision
                # stars[i].mass = mass
                stars[i].position = stars[i].position - fac * dr
                stars[i].velocity = stars[i].velocity - fac * dv

                id_count += 1
                newstar.id = id_count
                stars.add_particles(newstar)
                added_mass += new_mass

                if nbin >= number_of_binaries: break

            kep.stop()

            print 'created', nbin, 'binaries'
            sys.stdout.flush()

            stars.mass = stars.mass * total_mass / (total_mass + added_mass)
            number_of_stars += nbin
            Xtra = numpy.zeros(2)
        print "recentering stars"
        stars.move_to_center()
        sys.stdout.flush()

        stars.savepoint(time)

        print ''
        print "adding particles"
        # print stars
        sys.stdout.flush()
        gravity.particles.add_particles(stars)
        gravity.commit_particles()
    else:
        print "Restart detected.  Loading parameters from restart."
        new_end = options.t_end
        stars, time, multiples_code, Xtra = MRest.read_state_from_file(
            restart_file, gravity, new_smalln, kep)
        options.t_end = new_end

    total_mass = stars.mass.sum()
    ke = pa.kinetic_energy(stars)
    kT = ke / (1.5 * number_of_stars)

    # print "IDs:", stars.id.number

    print ''
    print "number_of_stars =", number_of_stars
    print "evolving to time =", end_time.number, \
          "in steps of", delta_t.number
    sys.stdout.flush()

    # Channel to copy values from the code to the set in memory.
    channel = gravity.particles.new_channel_to(stars)

    stopping_condition = gravity.stopping_conditions.collision_detection
    stopping_condition.enable()

    # -----------------------------------------------------------------
    # Create the coupled code and integrate the system to the desired
    # time, managing interactions internally.

    kep = init_kepler(stars[0], stars[1])
    if not multiples_code:
        multiples_code = multiples.Multiples(gravity, new_smalln, kep)

        multiples_code.neighbor_distance_factor = 1.0
        multiples_code.neighbor_veto = True
        #multiples_code.neighbor_distance_factor = 2.0
        #multiples_code.neighbor_veto = True
        multiples_code.retain_binary_apocenter = False

    print ''
    print 'multiples_code.initial_scale_factor =', \
        multiples_code.initial_scale_factor
    print 'multiples_code.neighbor_distance_factor =', \
        multiples_code.neighbor_distance_factor
    print 'multiples_code.neighbor_veto =', \
        multiples_code.neighbor_veto
    print 'multiples_code.final_scale_factor =', \
        multiples_code.final_scale_factor
    print 'multiples_code.initial_scatter_factor =', \
        multiples_code.initial_scatter_factor
    print 'multiples_code.final_scatter_factor =', \
        multiples_code.final_scatter_factor
    print 'multiples_code.retain_binary_apocenter =', \
        multiples_code.retain_binary_apocenter

    #    if mc_root_to_tree is not None:
    #        multiples_code.root_to_tree = mc_root_to_tree
    #        print 'multiples code re-loaded with binary trees snapshot'

    pre = "%%% "
    E0, cpu0 = print_log(pre, time, multiples_code)

    while time < end_time:

        time += delta_t
        multiples_code.evolve_model(time)

        # Copy values from the module to the set in memory.

        channel.copy()

        # Copy the index (ID) as used in the module to the id field in
        # memory.  The index is not copied by default, as different
        # codes may have different indices for the same particle and
        # we don't want to overwrite silently.

        channel.copy_attribute("index_in_code", "id")

        print_log(pre, time, multiples_code, E0, cpu0)
        stars.savepoint(time)
        MRest.write_state_to_file(time,
                                  stars,
                                  gravity,
                                  multiples_code,
                                  options.restart_file,
                                  Xtra,
                                  backup=1)
        sys.stdout.flush()

    #-----------------------------------------------------------------

    if not outfile is None:

        # Write data to a file.

        f = open(outfile, 'w')

        #--------------------------------------------------
        # Need to save top-level stellar data and parameters.
        # Need to save multiple data and parameters.

        f.write('%.15g\n' % time.number)
        for s in multiples_code.stars:
            write_star(s, f)

        #--------------------------------------------------

        f.close()
        print 'wrote file', outfile

    print ''
    gravity.stop()
Example #11
0
def compress_binary_components(comp1, comp2, scale):

    # Compress the two-body system consisting of comp1 and comp2 to
    # lie within distance scale of one another.

    pos1 = comp1.position
    pos2 = comp2.position
    sep12 = ((pos2 - pos1)**2).sum()

    if sep12 > scale * scale:
        print('\ncompressing components', int(comp1.id.number), \
              'and', int(comp2.id.number), 'to separation', scale.number)
        sys.stdout.flush()
        mass1 = comp1.mass
        mass2 = comp2.mass
        total_mass = mass1 + mass2
        vel1 = comp1.velocity
        vel2 = comp2.velocity
        cmpos = (mass1 * pos1 + mass2 * pos2) / total_mass
        cmvel = (mass1 * vel1 + mass2 * vel2) / total_mass

        # For now, create and delete a temporary kepler
        # process to handle the transformation.  Obviously
        # more efficient to define a single kepler at the
        # start of the calculation and reuse it.

        kep = Kepler(redirection="none")
        kep.initialize_code()
        mass = comp1.mass + comp2.mass
        rel_pos = pos2 - pos1
        rel_vel = vel2 - vel1
        kep.initialize_from_dyn(mass, rel_pos[0], rel_pos[1], rel_pos[2],
                                rel_vel[0], rel_vel[1], rel_vel[2])
        M, th = kep.get_angles()
        a, e = kep.get_elements()
        if e < 1:
            peri = a * (1 - e)
            apo = a * (1 + e)
        else:
            peri = a * (e - 1)
            apo = 2 * a  # OK - used ony to reset scale
        limit = peri + 0.01 * (apo - peri)
        if scale < limit: scale = limit

        if M < 0:
            # print 'approaching'
            kep.advance_to_periastron()
            kep.advance_to_radius(limit)
        else:
            # print 'receding'
            if kep.get_separation() < scale:
                kep.advance_to_radius(limit)
            else:
                kep.return_to_radius(scale)

        # a,e = kep.get_elements()
        # r = kep.get_separation()
        # E,J = kep.get_integrals()
        # print 'kepler: a,e,r =', a.number, e.number, r.number
        # print 'E, J =', E, J

        # Note: if periastron > scale, we are now just past periastron.

        new_rel_pos = kep.get_separation_vector()
        new_rel_vel = kep.get_velocity_vector()
        kep.stop()

        # Enew = 0
        # r2 = 0
        # for k in range(3):
        #     Enew += 0.5*(new_rel_vel[k].number)**2
        #     r2 += (new_rel_pos[k].number)**2
        # rnew = math.sqrt(r2)
        # Enew -= mass.number/r1
        # print 'E, Enew, rnew =', E.number, E1, r1

        # Problem: the vectors returned by kepler are lists,
        # not numpy arrays, and it looks as though we can say
        # comp1.position = pos, but not comp1.position[k] =
        # xxx, as we'd like...  Also, we don't know how to
        # copy a numpy array with units...  TODO

        newpos1 = pos1 - pos1  # stupid trick to create zero vectors
        newpos2 = pos2 - pos2  # with the proper form and units...
        newvel1 = vel1 - vel1
        newvel2 = vel2 - vel2

        frac2 = mass2 / total_mass
        for k in range(3):
            dxk = new_rel_pos[k]
            dvk = new_rel_vel[k]
            newpos1[k] = cmpos[k] - frac2 * dxk
            newpos2[k] = cmpos[k] + (1 - frac2) * dxk
            newvel1[k] = cmvel[k] - frac2 * dvk
            newvel2[k] = cmvel[k] + (1 - frac2) * dvk

        # Perform the changes to comp1 and comp2, and recursively
        # transmit them to the (currently absolute) coordinates of
        # all lower components.

        offset_particle_tree(comp1, newpos1 - pos1, newvel1 - vel1)
        offset_particle_tree(comp2, newpos2 - pos2, newvel2 - vel2)
def get_orbit_ini(m0, m1, peri, ecc, incl, omega, 
                  rel_force=0.01, r_disk=50|units.AU):
  converter=nbody_system.nbody_to_si(1|units.MSun,1|units.AU)
  
  # semi-major axis
  if ecc!=1.0:
    semi = peri/(1.0-ecc)
  else:
    semi = 1.0e10 | units.AU
  
  # relative position and velocity vectors at the pericenter using kepler
  kepler = Kepler_twobody(converter)
  kepler.initialize_code()
  kepler.initialize_from_elements(mass=(m0+m1), semi=semi, ecc=ecc, periastron=peri) # at pericenter
  
  # moving particle backwards to radius r where: F_m1(r) = rel_force*F_m0(r_disk)
  #r_disk = peri
  r_ini = r_disk*(1.0 + numpy.sqrt(m1/m0)/rel_force)
  kepler.return_to_radius(radius=r_ini)
  
  rl = kepler.get_separation_vector()
  r = [rl[0].value_in(units.AU), rl[1].value_in(units.AU), rl[2].value_in(units.AU)] | units.AU
  vl = kepler.get_velocity_vector()
  v = [vl[0].value_in(units.kms), vl[1].value_in(units.kms), vl[2].value_in(units.kms)] | units.kms
  period_kepler = kepler.get_period()
  time_peri = kepler.get_time()
  
  kepler.stop()
  
  # rotation of the orbital plane by inclination and argument of periapsis
  a1 = ([1.0, 0.0, 0.0], [0.0, numpy.cos(incl), -numpy.sin(incl)], [0.0, numpy.sin(incl), numpy.cos(incl)])
  a2 = ([numpy.cos(omega), -numpy.sin(omega), 0.0], [numpy.sin(omega), numpy.cos(omega), 0.0], [0.0, 0.0, 1.0])
  rot = numpy.dot(a1,a2)
  r_au = numpy.reshape(r.value_in(units.AU), 3, 1)
  v_kms = numpy.reshape(v.value_in(units.kms), 3, 1)
  r_rot = numpy.dot(rot, r_au) | units.AU
  v_rot = numpy.dot(rot, v_kms) | units.kms
  
  bodies = Particles(2)
  bodies[0].mass = m0
  bodies[0].radius = 1.0|units.RSun
  bodies[0].position = (0,0,0) | units.AU
  bodies[0].velocity = (0,0,0) | units.kms
  bodies[1].mass = m1
  bodies[1].radius = 1.0|units.RSun
  bodies[1].x = r_rot[0]
  bodies[1].y = r_rot[1]
  bodies[1].z = r_rot[2]
  bodies[1].vx = v_rot[0]
  bodies[1].vy = v_rot[1]
  bodies[1].vz = v_rot[2]
  
  bodies.age = 0.0 | units.yr
  bodies.move_to_center()
  
  print "\t r_rel_ini  = ", r_rot.in_(units.AU)
  print "\t v_rel_ini  = ", v_rot.in_(units.kms)
  print "\t time since peri = ", time_peri.in_(units.yr)
  a_orbit, e_orbit, p_orbit = orbital_parameters(r_rot, v_rot, (m0+m1))
  print "\t a = ", a_orbit.in_(units.AU), "\t e = ", e_orbit, "\t period = ", p_orbit.in_(units.yr)
  
  return bodies, time_peri
def get_orbit_ini(m0,
                  m1,
                  peri,
                  ecc,
                  incl,
                  omega,
                  rel_force=0.01,
                  r_disk=50 | units.AU):
    converter = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.AU)

    # semi-major axis
    if ecc != 1.0:
        semi = peri / (1.0 - ecc)
    else:
        semi = 1.0e10 | units.AU

    # relative position and velocity vectors at the pericenter using kepler
    kepler = Kepler_twobody(converter)
    kepler.initialize_code()
    kepler.initialize_from_elements(mass=(m0 + m1),
                                    semi=semi,
                                    ecc=ecc,
                                    periastron=peri)  # at pericenter

    # moving particle backwards to radius r where: F_m1(r) = rel_force*F_m0(r_disk)
    #r_disk = peri
    r_ini = r_disk * (1.0 + numpy.sqrt(m1 / m0) / rel_force)
    kepler.return_to_radius(radius=r_ini)

    rl = kepler.get_separation_vector()
    r = [
        rl[0].value_in(units.AU), rl[1].value_in(units.AU), rl[2].value_in(
            units.AU)
    ] | units.AU
    vl = kepler.get_velocity_vector()
    v = [
        vl[0].value_in(units.kms), vl[1].value_in(units.kms), vl[2].value_in(
            units.kms)
    ] | units.kms
    period_kepler = kepler.get_period()
    time_peri = kepler.get_time()

    kepler.stop()

    # rotation of the orbital plane by inclination and argument of periapsis
    a1 = ([1.0, 0.0,
           0.0], [0.0, numpy.cos(incl),
                  -numpy.sin(incl)], [0.0,
                                      numpy.sin(incl),
                                      numpy.cos(incl)])
    a2 = ([numpy.cos(omega), -numpy.sin(omega),
           0.0], [numpy.sin(omega), numpy.cos(omega), 0.0], [0.0, 0.0, 1.0])
    rot = numpy.dot(a1, a2)
    r_au = numpy.reshape(r.value_in(units.AU), 3, 1)
    v_kms = numpy.reshape(v.value_in(units.kms), 3, 1)
    r_rot = numpy.dot(rot, r_au) | units.AU
    v_rot = numpy.dot(rot, v_kms) | units.kms

    bodies = Particles(2)
    bodies[0].mass = m0
    bodies[0].radius = 1.0 | units.RSun
    bodies[0].position = (0, 0, 0) | units.AU
    bodies[0].velocity = (0, 0, 0) | units.kms
    bodies[1].mass = m1
    bodies[1].radius = 1.0 | units.RSun
    bodies[1].x = r_rot[0]
    bodies[1].y = r_rot[1]
    bodies[1].z = r_rot[2]
    bodies[1].vx = v_rot[0]
    bodies[1].vy = v_rot[1]
    bodies[1].vz = v_rot[2]

    bodies.age = 0.0 | units.yr
    bodies.move_to_center()

    print "\t r_rel_ini  = ", r_rot.in_(units.AU)
    print "\t v_rel_ini  = ", v_rot.in_(units.kms)
    print "\t time since peri = ", time_peri.in_(units.yr)
    a_orbit, e_orbit, p_orbit = orbital_parameters(r_rot, v_rot, (m0 + m1))
    print "\t a = ", a_orbit.in_(
        units.AU), "\t e = ", e_orbit, "\t period = ", p_orbit.in_(units.yr)

    return bodies, time_peri