add more missing examples

Author: rieder

src/amuse/examples/calculate_stellar_isochrone.py

@@ -0,0 +1,100 @@
+"""
+   Evolve a population of N stars.
+   Initial mass function between Mmin and Mmax with stellar evolution
+   for metallicity z.
+"""
+import numpy as np
+from amuse.io import write_set_to_file, read_set_from_file
+from amuse.units import units
+from amuse.datamodel import Particles
+from amuse.community.seba import Seba
+from amuse.community.sse import Sse
+from amuse.community.mesa_r2208.interface import Mesa
+#from amuse.community.mesa import Mesa
+from amuse.community.evtwin import Evtwin
+from amuse.ic.salpeter import new_salpeter_mass_distribution
+
+def calculate_stellar_temperature_and_luminosity(
+        star,
+        evo_code, metallicity=0.02, time_end=100 | units.Myr):
+
+    if evo_code.find("Seba") >= 0:
+        stellar = Seba()
+        stellar.parameters.metallicity = metallicity
+        stellar.particles.add_particle(star)
+        channel_to_framework = stellar.particles.new_channel_to(star)
+        stellar.evolve_model(time_end)
+        channel_to_framework.copy_attributes(
+            ["radius", "temperature", "luminosity"]
+        )
+        stellar.stop()
+    else:
+        if evo_code.find("Mesa") >= 0:
+            stellar = Mesa()
+        elif evo_code.find("Evtwin") >= 0:
+            stellar = Evtwin()
+        stellar.parameters.metallicity = metallicity
+        stellar.particles.add_particle(star)
+        channel_to_framework = stellar.particles.new_channel_to(star)
+        try:
+            if star.mass<=0.8|units.MSun:
+                t_end = 0.1|units.Myr
+            else:
+                t_end = time_end
+            print("Evolved star: m=", star.mass.in_(units.MSun))
+            stellar.evolve_model(t_end)
+            channel_to_framework.copy_attributes(
+                ["radius", "temperature", "luminosity"])
+            print("Successfully evolved star: m=", star.mass.in_(units.MSun))
+        except:
+            print("Failed to evolve star: m=", star.mass.in_(units.MSun))
+        stellar.stop()
+    return star
+
+def main(input_file, time_end, metallicity, evo_code="Seba"):
+    evo_code = evo_code.capitalize()
+    if evo_code.find("Seba") >= 0:
+        filename = "../data/Stellar_Isochrone_Seba.amuse"
+    elif evo_code.find("Mesa") >= 0:
+        filename = "../data/Stellar_Isochrone_Mesa.amuse"
+    elif evo_code.find("Evtwin") >= 0:
+        filename = "../data/Stellar_Isochrone_Evtwin.amuse"
+    else:
+        print("No input code, stop")
+        exit(-1)
+        
+    stars = read_set_from_file(input_file, "amuse", close_file=True)
+    from os.path import exists
+    if exists(filename):
+        stars_done = read_set_from_file(filename, "amuse", close_file=True)
+    else:
+        stars_done = Particles(0)
+    for si in stars:
+        if si not in stars_done:
+            star_evolved = calculate_stellar_temperature_and_luminosity(
+                si.as_set(),
+                evo_code=evo_code, metallicity=metallicity,
+                time_end=time_end)
+            stars_done.add_particle(star_evolved)
+            write_set_to_file(stars_done, filename, overwrite_file=True)
+
+
+def new_option_parser():
+    from amuse.units.optparse import OptionParser
+    result = OptionParser()
+    result.add_option("-C", dest="evo_code", default="Seba",
+                      help="stellar evolution code [SeBa]")
+    result.add_option("--f", dest="input_file",
+                      default="initial_stellar_mass_function.amuse",
+                      help="input data file")
+    result.add_option("-t", dest="time_end", unit=units.Myr,
+                      type="float", default=4500.0 | units.Myr,
+                      help="end time of the simulation [4500] Myr")
+    result.add_option("-z", dest="metallicity", type="float", default=0.02,
+                      help="metalicity [0.02]")
+    return result
+
+
+if __name__ in ('__main__'):
+    o, arguments = new_option_parser().parse_args()
+    main(**o.__dict__)

src/amuse/examples/hr_diagram_cluster.py

@@ -0,0 +1,239 @@
+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.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
+from amuse.ic.salpeter import new_salpeter_mass_distribution
+usage = """\
+usage: %prog [options]
+
+This script will generate HR diagram for an
+evolved cluster of stars with a Salpeter mass
+distribution.
+"""
+
+
+def simulate_stellar_evolution(
+        stellar_evolution=SSE(),
+        number_of_stars=1000,
+        end_time=1000.0 | units.Myr,
+        name_of_the_figure="cluster_HR_diagram.png",
+):
+    """
+    A cluster of stars will be created, with masses following a Salpeter IMF.
+    The stellar evolution will be simulated using any of the legacy codes (SSE,
+    EVtwin, or MESA).
+    Finally, a Hertzsprung-Russell diagram will be produced for the final state
+    of the simulation.
+    """
+    print(
+        "The evolution of",
+        str(number_of_stars),
+        "stars will be ",
+        "simulated until t =",
+        str(end_time),
+        "..."
+    )
+    
+    stellar_evolution.commit_parameters()
+    
+    print(
+        "Deriving a set of",
+        str(number_of_stars),
+        "random masses",
+        "following a Salpeter IMF between 0.1 and 125 MSun (alpha = -2.35)."
+    )
+    
+    salpeter_masses = new_salpeter_mass_distribution(number_of_stars)
+    
+    print("Initializing the particles")
+    stars = datamodel.Particles(number_of_stars)
+    stars.mass = salpeter_masses
+    print("Stars to evolve:")
+    print(stars)
+    stars = stellar_evolution.particles.add_particles(stars)
+    stellar_evolution.commit_particles()
+
+    # print stars.temperature
+
+    print("Start evolving...")
+    stellar_evolution.evolve_model(end_time)
+
+    print("Evolved model successfully.")
+    temperatures = stars.temperature
+    luminosities = stars.luminosity
+
+    stellar_evolution.stop()
+
+    plot_HR_diagram(temperatures, luminosities, name_of_the_figure, end_time)
+            
+    print("All done!")
+
+
+def plot_HR_diagram(temperatures, luminosities, name_of_the_figure, end_time):
+    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...")
+        number_of_stars = len(temperatures)
+        pyplot.figure(figsize=(7, 8))
+        pyplot.title('Hertzsprung-Russell diagram', fontsize=12)
+        pyplot.xlabel(r'T$_{\rm eff}$ (K)')
+        pyplot.ylabel(r'Luminosity (L$_\odot$)')
+        pyplot.loglog(temperatures.value_in(units.K),
+                      luminosities.value_in(units.LSun), "ro")
+
+        xmin, xmax = 20000.0, 2500.0
+        ymin, ymax = 1.e-4, 1.e4
+        pyplot.text(xmin*.75, ymax*0.1, str(number_of_stars) +
+                    " stars\nt="+str(end_time))
+        pyplot.axis([xmin, xmax, ymin, ymax])
+        pyplot.savefig(name_of_the_figure)
+    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 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_short():
+    assert is_mpd_running()
+    code = new_code("sse", 100)
+
+    test_results_path = get_path_to_results()
+    output_file = os.path.join(test_results_path, "cluster_HR_diagram.png")
+
+    simulate_stellar_evolution(
+        code,
+        number_of_stars=100,
+        end_time=2.0 | units.Myr,
+        name_of_the_figure=output_file
+    )
+
+
+def new_commandline_option_parser():
+    result = OptionParser(usage)
+    result.add_option(
+        "-c",
+        "--code",
+        choices=["sse", "evtwin", "mesa"],
+        default="sse",
+        dest="code",
+        metavar="CODE",
+        help="CODE to use for stellar evolution"
+    )
+    result.add_option(
+        "-n",
+        "--number_of_stars",
+        type="int",
+        default=10,
+        dest="number_of_stars",
+        help="Number of stars in the cluster"
+    )
+    result.add_option(
+        "-t",
+        "--end_time",
+        type="float",
+        default=1000.0,
+        dest="end_time",
+        help="Time to evolve to in Myr"
+    )
+    result.add_option(
+        "-C",
+        "--cache",
+        type="string",
+        default=None,
+        dest="cacheDir",
+        help="Use/write cache from directory"
+    )
+    result.add_option(
+        "-S",
+        "--seed",
+        type="int",
+        default=None,
+        dest="salpeterSeed",
+        help="Random number generator seed"
+    )
+    result.add_option(
+        "-p",
+        "--plot_file",
+        type="string",
+        default="cluster_HR_diagram.png",
+        dest="plot_filename",
+        help="Name of the file to plot to"
+    )
+
+    return result
+
+
+if __name__ == '__main__':
+
+    parser = new_commandline_option_parser()
+    (options, arguments) = parser.parse_args()
+    if arguments:
+        parser.error("unknown arguments '{0}'".format(arguments))
+
+    code = new_code(options.code, options.number_of_stars)
+
+    if not (options.cacheDir is None):
+        print("Using cache directory: %s" % (options.cacheDir))
+        code = CachedStellarEvolution(code, options.cacheDir)
+
+    if not (options.salpeterSeed is None):
+        numpy.random.seed(options.salpeterSeed)
+
+    simulate_stellar_evolution(
+        code,
+        number_of_stars=options.number_of_stars,
+        end_time=options.end_time | units.Myr,
+        name_of_the_figure=options.plot_filename
+    )

src/amuse/examples/plot_plummer_stars_and_planets.py

@@ -0,0 +1,28 @@
+from amuse.lab import *
+bodies = read_set_from_file("../data/Pl_N16NO_A1000R.amuse", close_file=True)
+stars = bodies[bodies.name=="star"]
+systems = bodies[bodies.name=="system"]
+planets = bodies[bodies.name=="planet"]
+debris = bodies[bodies.name=="asteroid"]
+
+print(f"N={len(bodies)}, {len(stars)}, {len(systems)}, {len(planets)}, {len(debris)}")
+
+from matplotlib import pyplot as plt
+
+plt.style.use('../lib/matplotlibrc')
+
+figure = plt.figure(figsize=(6, 6))
+ax = plt.gca()
+ax.minorticks_on() # switch on the minor ticks
+ax.locator_params(nbins=3)
+
+plt.xlabel("x [pc]")
+plt.ylabel("y [pc]")
+plt.xlim(-0.01, 0.01)
+plt.ylim(-0.01, 0.01)
+plt.scatter(debris.x.value_in(units.pc), debris.y.value_in(units.pc), c='k', lw=0, s=1, label="debris")
+plt.scatter(planets.x.value_in(units.pc), planets.y.value_in(units.pc), c='y', lw=0, s=20, label="planet")
+plt.scatter(systems.x.value_in(units.pc), systems.y.value_in(units.pc), c='r', lw=0, s=30, label="planetary system")
+plt.scatter(stars.x.value_in(units.pc), stars.y.value_in(units.pc), c='b', lw=0, s=30, label="single star")
+plt.savefig("fig_Pl_N16NO_A1000R.pdf")
+plt.show()