/amuse

Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
Raw Blame History
231 lines (188 sloc) 6.76 KB
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
)