Create a binary star pair

create_and_follow_binary.py

#!/usr/bin/env python
# coding: utf-8

import sys
import numpy
import argparse
from amuse.units import nbody_system, units, constants
from amuse.ic.plummer import new_plummer_model
from amuse.community.ph4 import Ph4
from amuse.community.smalln import Smalln
from amuse.community.kepler import Kepler
from amuse.couple import multiples

from amuse.support.console import set_printing_strategy

from amuse.ext.orbital_elements import get_orbital_elements_from_binaries, generate_binaries

def new_argument_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-n", "--number_of_stars", type=int, default=100,
        dest="number_of_stars",
        help="number of stars [100]",
    )
    parser.add_argument(
        "-t", "--end_time", type=float, default=50.,
        dest="end_time",
        help="end time [50.0] (Myr)",
    )
    parser.add_argument(
        "-s", "--random_seed", type=int, default=None,
        dest="random_seed",
        help="random seed to use [random]",
    )
    parser.add_argument(
        "-r", "--radius", type=float, default=0.5,
        dest="radius",
        help="cluster radius [0.5] (parsec)",
    )
    parser.add_argument(
        "-o", "--outputfile", default="output.txt",
        dest="outputfile",
        help="name of the output file [output.txt]",
    )
    return parser.parse_args()


class GravityWithMultiples:
    def __init__(self, stars, converter=None, time_step=0.01 | units.Myr):
        self.time_step = time_step
        self.mass_total = stars.total_mass()
        self.number_of_particles = len(stars)
        if converter is None:
            self.converter = nbody_system.nbody_to_si(
                self.mass_total,
                self.time_step
            )
        else:
            self.converter = converter
        self.model_time = 0 * self.time_step

        self.SmallN = None
        self.gravity = Ph4(convert_nbody=self.converter)
        self.gravity.parameters.epsilon_squared = (self.converter.to_si(0.0 | nbody_system.length))**2

        self.gravity.particles.add_particles(stars)

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

        self.init_smalln()
        self.kep = Kepler(unit_converter=self.converter)
        self.kep.initialize_code()
        self.multiples_code = multiples.Multiples(
            self.gravity, self.new_smalln, self.kep, constants.G
        )
        self.multiples_code.neighbor_perturbation_limit = 0.05
        self.multiples_code.global_debug = 1

        #  global_debug = 0: no output from multiples
        #                 1: minimal output
        #                 2: debugging output
        #                 3: even more output

        print('')
        print(f'multiples_code.neighbor_veto = {self.multiples_code.neighbor_veto}')
        print(f'multiples_code.neighbor_perturbation_limit = {self.multiples_code.neighbor_perturbation_limit}')
        print(f'multiples_code.retain_binary_apocenter = {self.multiples_code.retain_binary_apocenter}')
        print(f'multiples_code.wide_perturbation_limit = {self.multiples_code.wide_perturbation_limit}')

        self.energy_0 = self.print_diagnostics(self.multiples_code)

    @property
    def stars(self):
        return self.multiples_code.stars
    
    @property
    def particles(self):
        return self.multiples_code.particles
    
    def init_smalln(self):
        self.SmallN = Smalln(convert_nbody=self.converter)

    def new_smalln(self):
        self.SmallN.reset()
        return self.SmallN

    def stop_smalln(self):
        self.SmallN.stop()

    def print_diagnostics(self, grav, energy_0=None):
        # Simple diagnostics.

        energy_kinetic = grav.kinetic_energy
        energy_potential = grav.potential_energy
        (
            self.number_of_multiples,
            self.number_of_binaries,
            self.energy_in_multiples,
        ) = grav.get_total_multiple_energy()
        energy = energy_kinetic + energy_potential + self.energy_in_multiples
        print('')
        print(f'Time = {grav.get_time().in_(units.Myr)}')
        print(f'    top-level kinetic energy = {energy_kinetic}')
        print(f'    top-level potential energy = {energy_potential}')
        print(f'    total top-level energy = {energy_kinetic + energy_potential}')
        print(f'    {self.number_of_multiples} multiples, total energy = {self.energy_in_multiples}')
        print(f'    uncorrected total energy ={energy}')

        # Apply known corrections.

        energy_tidal = (
            grav.multiples_external_tidal_correction
            + grav.multiples_internal_tidal_correction
        )  # tidal error
        energy_error = grav.multiples_integration_energy_error  # integration error

        energy -= energy_tidal + energy_error
        print(f'    corrected total energy = {energy}')

        if energy_0 is not None: print('    relative energy error=', (energy-energy_0)/energy_0)

        return energy

    def evolve_model(self, t_end):
        self.multiples_code.evolve_model(t_end)
        self.print_diagnostics(self.multiples_code, self.energy_0)
        self.model_time = self.multiples_code.model_time

    def stop(self):
        self.gravity.stop()
        self.kep.stop()
        self.stop_smalln()


def main():
    set_printing_strategy("astro")  # use more practical units

    args = new_argument_parser()
    if args.random_seed is None:
        args.random_seed = numpy.random.randint(1e9)
    # Fix random seed for reproducibility
    numpy.random.seed(args.random_seed)
    number_of_stars = args.number_of_stars
    radius = args.radius | units.pc
    total_mass = number_of_stars | units.MSun
    converter = nbody_system.nbody_to_si(total_mass, radius)
    with open(args.outputfile, 'w') as out:
        out.write(f"# Random seed used: {args.random_seed}\n")
        out.write(f"# Number of stars: {args.number_of_stars}\n")
        out.write(f"# Cluster radius: {args.radius}\n")

    stars = new_plummer_model(number_of_stars, convert_nbody=converter)
    stars.mass = stars.total_mass() / number_of_stars
    stars.scale_to_standard(
        convert_nbody=converter,
        smoothing_length_squared=0 | units.m**2,
    )
    index = numpy.arange(number_of_stars)
    stars.id = index + 1
    stars.radius = (0.5 / number_of_stars) | units.parsec

    model = GravityWithMultiples(stars)
    channel_from_code = model.multiples_code.stars.new_channel_to(stars)

    time_step = converter.to_si(0.01 | nbody_system.time)

    # stars.savepoint(model.model_time)

    time = [] | units.Myr
    semi_major_axis = [] | units.au
    eccentricity = []
    print("Evolving until the first binary forms...")
    while len(model.stars.get_binaries()) == 0:
        model.evolve_model(model.model_time + time_step)
        channel_from_code = model.stars.new_channel_to(stars)
        channel_from_code.copy()

    time_of_first_binary = model.model_time
    while model.model_time < time_of_first_binary + (args.end_time | units.Myr):
        model.evolve_model(model.model_time + time_step)
        channel_from_code = model.stars.new_channel_to(stars)
        channel_from_code.copy()
        binaries = model.stars.get_binaries()
        if len(binaries) > 0:
            binary = binaries[0]
            mass1, mass2, a, e, x1, x2, x3, x4 = get_orbital_elements_from_binaries(
                binary[0], binary[1]
            )
            time.append(model.model_time)
            semi_major_axis.append(a[0])
            eccentricity.append(e[0])

    with open(args.outputfile, "a") as out:
        for i in range(len(time)):
            out.write(f"{time[i].value_in(units.Myr)} {semi_major_axis[i].value_in(units.pc)} {eccentricity[i]}\n")
    return model


if __name__ == "__main__":
    main()