Permalink
Cannot retrieve contributors at this time
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 | |
) |