evolve_stars.py

import sys
from amuse.units import units, nbody_system
from amuse.io import read_set_from_file
from amuse.io import write_set_to_file
from amuse.community.ph4 import Ph4


def main():
    time_step = 0.01 | units.Myr
    time_end = 10 | units.Myr
    time_diagnostic = 0.1 | units.Myr
    diagnostic_step = 0

    stars = read_set_from_file(sys.argv[1])
    converter = nbody_system.nbody_to_si(
        stars.mass.sum() / len(stars),
        time_step,
    )
    gravity = Ph4(converter, redirection="none")
    gravity.particles.add_particles(stars)

    while gravity.model_time < time_end:
        gravity.evolve_model(gravity.model_time + time_step)
        if int(gravity.model_time / time_diagnostic) > diagnostic_step:
            print(f"# WRITING STARS TO DISK (TIME = {gravity.model_time.in_(units.Myr)} #")
            write_set_to_file(
                gravity.particles,
                f'stars-{diagnostic_step:04d}.amuse',
                timestamp=gravity.model_time,
            )
            diagnostic_step += 1
    write_set_to_file(
        gravity.particles,
        'stars-end.amuse',
        overwrite_file=True,
        timestamp=gravity.model_time,
    )
    gravity.stop()


if __name__ == "__main__":
    main()