simple example code: grav_stellar_simple.pyΒΆ

(Source code)

"""
    simple coupling of a stellar evolution and a gravitational code to simulate a cluster of stars.
"""
from __future__ import print_function

import os

from amuse.units.optparse import OptionParser
from amuse.units import units, nbody_system
from amuse.datamodel.particles import Channels

from amuse.community.hermite0.interface import Hermite
from amuse.community.seba.interface import SeBa
from amuse.couple.bridge import Bridge

from amuse.ic.plummer import new_plummer_model
from amuse.ic.salpeter import new_salpeter_mass_distribution

from matplotlib import pyplot
from amuse import plot as aplot

from amuse.support.console import set_printing_strategy


def create_stars(number_of_stars, size):
    masses = new_salpeter_mass_distribution(number_of_stars, mass_min = 2|units.MSun)
    converter = nbody_system.nbody_to_si(masses.sum(), size)
    stars = new_plummer_model(number_of_stars, convert_nbody=converter)
    stars.mass = masses
    stars.zams_mass = masses

    return stars, converter

def plot_results(stars, time):
    mass_loss = stars.zams_mass - stars.mass

    x = stars.x.in_(units.parsec)
    y = stars.y.in_(units.parsec)

    pyplot.figure(figsize=(8,8))
    aplot.plot(x, y, "*")

    for x, y, mass_loss in zip(x.number, y.number, mass_loss):
        pyplot.annotate("%0.2f"%abs(mass_loss.number), xy=(x,y+2),
            horizontalalignment='center', verticalalignment='bottom')

    pyplot.axis('equal')
    pyplot.xlim([-60, 60])
    pyplot.ylim([-60, 60])
    aplot.xlabel("x")
    aplot.ylabel("y")
    pyplot.title("time = " + str(time))

    if not os.path.exists("plots"):
        os.mkdir("plots")

    name = "plots/plot_{0:=05}.png".format(int(time.value_in(units.Myr)))
    print("creating", name)
    pyplot.savefig(name)
    pyplot.close()

def gravity_and_stellar_evolution(number_of_stars, size, end_time, sync_timestep=1|units.Myr, plot_timestep=10|units.Myr):

    stars, converter = create_stars(number_of_stars, size)

    gravity = Hermite(converter)
    stellar = SeBa()
    bridge = Bridge()

    gravity.particles.add_particles(stars)
    stellar.particles.add_particles(stars)

    bridge.add_system(gravity)
    bridge.add_system(stellar)
    bridge.channels.add_channel(stellar.particles.new_channel_to(gravity.particles, attributes=["mass", "radius"]))

    bridge.timestep = sync_timestep

    plot_channels = Channels()
    plot_channels.add_channel(stellar.particles.new_channel_to(stars))
    plot_channels.add_channel(gravity.particles.new_channel_to(stars))

    time = 0 | units.Myr
    while time <= end_time:

        bridge.evolve_model(time)

        plot_channels.copy()
        plot_results(stars, time)

        time += plot_timestep

def parse_arguments():
    parser = OptionParser()
    parser.add_option("-N", dest="number_of_stars", type="int", default=100,
        help="The number of stars in the cluster [%default].")
    parser.add_option("-s", dest="size", type="float", unit=units.parsec, default=10,
        help="The total size of the cluster [%default %unit].")
    parser.add_option("-t", dest="end_time", type="float", unit=units.Gyr, default=0.1,
        help="The end time of the simulation [%default %unit].")

    options, args = parser.parse_args()
    return options.__dict__

if __name__ == "__main__":
    options = parse_arguments()
    set_printing_strategy("custom", preferred_units = [units.MSun, units.parsec, units.Myr], precision=3)
    gravity_and_stellar_evolution(**options)

Keywords: python, amuse, astrophysics, matplotlib, pylab, example, codex (see how-to-search-examples)

This Page