applications example code: test_HRdiagram_tracks.pyΒΆ

[source code]

import sys
import numpy
import os
import warnings

from optparse import OptionParser

from amuse.units import units
from amuse.community.sse.interface import SSE
from amuse.community.evtwin.interface import EVtwin
from amuse.community.evtwin2sse.interface import EVtwin2SSE
from amuse.community.mesa.interface import MESA

from amuse.community.cachedse.interface import CachedStellarEvolution

from amuse.test.amusetest import get_path_to_results

from amuse import datamodel
from amuse.rfi.core import is_mpd_running
usage = """\
usage: %prog [options]

This script will generate HR diagram tracks for
stars with specified a masses.
"""
def stellar_remnant_state(star):
    return 10 <= star.stellar_type.value_in(units.stellar_type) and \
        star.stellar_type.value_in(units.stellar_type) < 16

def simulate_evolution_tracks(
    stellar_evolution = SSE(),
    masses = [0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0] | units.MSun,
    name_of_the_figure = "HR_evolution_tracks.png"
    ):
    """
    For every mass in the `masses' array, a stellar evolution track across the Hertzsprung-Russell
    diagram will be calculated and plotted. Each star will be created, evolved and removed one by
    one. This is only necessary because the time span of each track is different (a solar mass star
    evolution track takes billions of years, but we don't want to also evolve high mass stars for
    billions of years) In most applications the stars have to be evolved up to a common end time,
    which can be more easily accomplished by creating an array (stars = datamodel.Stars(number_of_stars))
    and using evolve_model(end_time = ...).
    """
    number_of_stars=len(masses)
    all_tracks_luminosity = []
    all_tracks_temperature = []
    all_tracks_stellar_type = []
    stellar_evolution.commit_parameters()

    print "The evolution across the Hertzsprung-Russell diagram of ", str(number_of_stars), \
        " stars with\nvarying masses will be simulated..."

    for j in range(number_of_stars):
        star = datamodel.Particle()
        star.mass = masses[j]
        print "Created new star with mass: ", star.mass

        star = stellar_evolution.particles.add_particle(star)
        stellar_evolution.commit_particles()

        luminosity_at_time     = [] | units.LSun
        temperature_at_time     = [] | units.K
        stellar_type_at_time = [] | units.stellar_type

        stopped_evolving = False
#        Evolve this star until it changes into a compact stellar remnant (white dwarf, neutron star, or black hole)
        while not stellar_remnant_state(star) and not stopped_evolving:
            luminosity_at_time.append(star.luminosity)
            temperature_at_time.append(star.temperature)
            stellar_type_at_time.append(star.stellar_type)
            previous_age = star.age
            try:
                stellar_evolution.evolve_model()
                stopped_evolving = (star.age == previous_age) # Check whether the age has stopped increasing
            except Exception as ex:
                print str(ex)
                stopped_evolving = True
        if stopped_evolving:
            print "Age did not increase during timestep. Aborted evolving..."
        else:
            stellar_type_at_time.append(star.stellar_type)
            # Fudged: final stellar type annotation at previous (Teff, L);
            # BHs and neutron stars would otherwise fall off the chart.
            luminosity_at_time.append(luminosity_at_time[-1])
            temperature_at_time.append(temperature_at_time[-1])
        print " ... evolved model to t = " + str(star.age.as_quantity_in(units.Myr))
        print "Star has now become a: ", star.stellar_type, "(stellar_type: "+str(star.stellar_type.value_in(units.stellar_type))+")"
        print
        all_tracks_luminosity.append(luminosity_at_time)
        all_tracks_temperature.append(temperature_at_time)
        all_tracks_stellar_type.append(stellar_type_at_time)

#        Remove the star before creating the next one. See comments at the top.
        stellar_evolution.particles.remove_particle(star)

    stellar_evolution.stop()

    plot_HR_diagram(masses, all_tracks_luminosity, all_tracks_temperature,
        all_tracks_stellar_type, name_of_the_figure)

    print "All done!"

def plot_HR_diagram(masses, luminosity_tracks, temperature_tracks, stellar_type_tracks, plotfile):
    try:
#        This removes the need for ssh -X to be able to do plotting
        import matplotlib
        matplotlib.use("Agg", warn=False)

        from matplotlib import pyplot
        print "Plotting the data..."
        pyplot.figure(figsize = (7, 8))
        pyplot.title('Hertzsprung-Russell diagram', fontsize=12)
        pyplot.xlabel('Effective Temperature (K)')
        pyplot.ylabel('Luminosity (solar luminosity)')

#        Define some strings for line formatting (colors, symbols, etc.), used recurrently when many stars are simulated
        plot_format_strings_lines=["r-","y-","c-","b-","m-"]
        len_fmt_str_lin=len(plot_format_strings_lines)
        plot_format_strings_symbols=["r^","y^","c^","b^","m^","rs","ys","cs","bs","ms"]
        len_fmt_str_sym=len(plot_format_strings_symbols)

        number_of_stars = len(masses)
        for j in range(number_of_stars):
#            Plot track of the current star j
            x_values = temperature_tracks[j].value_in(units.K)
            y_values = luminosity_tracks[j].value_in(units.LSun)
            pyplot.loglog(x_values, y_values, plot_format_strings_lines[j%len_fmt_str_lin])

#            Plot symbols whenever this star has switched to the next stellar evolution phase
            x_values = []
            y_values = []
            text_values = []
            previous_type = 15 | units.stellar_type
            for i, type in enumerate(stellar_type_tracks[j]):
                if not type == previous_type:
                    x_values.append(temperature_tracks[j][i].value_in(units.K))
                    y_values.append(luminosity_tracks[j][i].value_in(units.LSun))
                    text_values.append(stellar_type_tracks[j][i].value_in(units.stellar_type))
                    previous_type = type

            pyplot.loglog(x_values, y_values, plot_format_strings_symbols[j%len_fmt_str_sym])
            text_offset_factor_x=1.05
            text_offset_factor_y=0.6
            for i, phase in enumerate(text_values):
                pyplot.annotate(str(int(phase)), xy=(x_values[i],y_values[i]), \
                    xytext=(x_values[i]*text_offset_factor_x,y_values[i]*text_offset_factor_y))
            text_offset_factor_x=1.1
            text_offset_factor_y=0.9
            pyplot.annotate(str(masses[j]),xy=(x_values[0],y_values[0]), \
                xytext=(x_values[0]*text_offset_factor_x,y_values[0]*text_offset_factor_y), \
                color='g', horizontalalignment='right')

        pyplot.axis([300000., 2500., 1.e-2, 1.e6])
#        Or use these axes to also view neutron stars and black holes:
#         pyplot.axis([1.e7, 2500., 1.e-11, 1.e6])
        pyplot.savefig(plotfile)
        print "Meaning of the stellar evolution phase markers (black numbers):"
        for i in range(16):
            print str(i)+": ", (i | units.stellar_type)
    except ImportError:
        print "Unable to produce plot: couldn't find matplotlib."

class InstantiateCode(object):

    def sse(self, number_of_stars):
        return SSE()

    def evtwin(self, number_of_stars):
        result = EVtwin()
        result.initialize_code()
        if number_of_stars > result.parameters.maximum_number_of_stars:
            result.parameters.maximum_number_of_stars = number_of_stars
            warnings.warn("You're simulating a large number of stars with EVtwin. This may not be such a good idea...")
        return result

    def mesa(self, number_of_stars):
        result = MESA()
        result.initialize_code()
        if number_of_stars > (10):
            warnings.warn("You're simulating a large number of stars with MESA. This may not be such a good idea...")
        if number_of_stars > (1000):
            raise Exception("You want to simulate with more than 1000 stars using MESA, this is not supported")
        return result

    def evtwin2sse(self, number_of_stars):
        result = EVtwin2SSE()
        result.initialize_code()
        # TODO add maximum_number_of_stars parameter to Evtwin2SSE
        #if number_of_stars > result.parameters.maximum_number_of_stars:
        #     result.parameters.maximum_number_of_stars = number_of_stars
        #     warnings.warn("You're simulating a large number of stars with EVtwin. This may not be such a good idea...")
        return result

    def new_code(self, name_of_the_code, number_of_stars):
        if hasattr(self, name_of_the_code):
            return getattr(self, name_of_the_code)(number_of_stars)
        else:
            raise Exception("Cannot instantiate code with name '{0}'".format(name_of_the_code))

def new_code(name_of_the_code, number_of_stars):
    return InstantiateCode().new_code(name_of_the_code, number_of_stars)

def test_simulate_one_star():
    assert is_mpd_running()
    code = new_code("sse", 1)
    test_results_path = get_path_to_results()
    output_file = os.path.join(test_results_path, "HR_evolution_tracks.png")
    simulate_evolution_tracks(
        code,
        [20.0] | units.MSun,
        name_of_the_figure=output_file,
    )

def new_commandline_option_parser():
    result = OptionParser(usage)
    result.add_option(
        "-c",
        "--code",
        choices=["sse", "evtwin", "mesa", "evtwin2sse"],
        default="sse",
        dest="code",
        metavar="CODE",
        help="CODE to use for stellar evolution"
    )
    result.add_option(
        "-C",
        "--cache",
        type="string",
        default=None,
        dest="cacheDir",
        help="Use/write cache from directory"
    )
    result.add_option(
        "-p",
        "--plot_file",
        type="string",
        default="HR_evolution_tracks.png",
        dest="plot_filename",
        help="Name of the file to plot to"
    )
    return result

if __name__ == '__main__':
    if not is_mpd_running():
        print "There is no mpd server running. Please do 'mpd &' first."
        sys.exit()
    parser = new_commandline_option_parser()
    (options, arguments) = parser.parse_args()
    if arguments:
        parser.error("unknown arguments '{0}'".format(arguments))
    mass_list = [0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0] | units.MSun

    code = new_code(options.code, len(mass_list))
    if not (options.cacheDir is None):
        print "Using cache directory: %s" % (options.cacheDir)
# As a special case, we use caching of the underlying models instead of the model output for EVtwin2SSE
        if (options.code == "evtwin2sse"):
            code.cache_underlying_models(options.cacheDir)
        else:
            code = CachedStellarEvolution(code, options.cacheDir)
    simulate_evolution_tracks(
        code,
        masses = mass_list,
        name_of_the_figure=options.plot_filename
    )

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

This Page