import os.path
from amuse.test.amusetest import get_path_to_results
try:
from matplotlib import pyplot
HAS_MATPLOTLIB = True
from amuse.plot import plot, semilogy, xlabel, ylabel, loglog
except ImportError:
HAS_MATPLOTLIB = False
from amuse.units import units
from amuse.units import constants
from amuse.units.generic_unit_converter import ConvertBetweenGenericAndSiUnits
from amuse.support.exceptions import AmuseException
from amuse.community.mesa.interface import MESA
from amuse.community.gadget2.interface import Gadget2
from amuse.community.fi.interface import Fi
from amuse.ext.star_to_sph import convert_stellar_model_to_SPH, pickle_stellar_model, StellarModel2SPH
from optparse import OptionParser
import numpy
from amuse.datamodel import Particles
from amuse.datamodel import ParticlesSuperset
from amuse.datamodel import Grid
def slowtest1():
head_on_stellar_merger()
def head_on_stellar_merger(
masses = [0.3, 3.0] | units.MSun,
star_age = 310.0 | units.Myr,
maximally_evolved_stars = False,
initial_separation = 4.0 | units.RSun,
angle = numpy.pi / 3,
initial_speed = 3000.0 | units.km / units.s,
initial_speed_perpendicular = 30.0 | units.km / units.s,
number_of_sph_particles = 1000,
t_end = 1.0e4 | units.s,
sph_code = Fi,
steps_per_snapshot = 4,
snapshot_size = 100,
use_stored_stellar_models = True
):
"""
masses: Mass of the two stars
star_age: Initial age of the stars (if maximally_evolved_stars is False)
maximally_evolved_stars: Evolve stars as far as the Stellar Evolution code can get
number_of_sph_particles: Total number of particles of both stars, divided according to their masses
t_end: (Physical, not computational) duration of the hydrodynamics simulation
sph_code: Code to use for the hydrodynamics simulation
steps_per_snapshot: A hydroplot snapshot is generated each time after this many steps (0 or None means no snapshots)
snapshot_size: Size of the snapshot in pixels along one dimension
use_stored_stellar_models: Flag to use previously stored stellar model files (for speed-up).
"""
# Convert some of the input parameters to string, for use in output file names:
n_string = "n" + ("%1.0e"%(number_of_sph_particles)).replace("+0","").replace("+","")
t_end_string = "t" + ("%1.0e"%(t_end.value_in(units.s))).replace("+0","").replace("+","")
masses_string = ("m1_" + ("%0.3e"%(masses[0].value_in(units.MSun))).replace("+0","").replace("+","") +
"_m2_" + ("%0.3e"%(masses[1].value_in(units.MSun))).replace("+0","").replace("+",""))
if maximally_evolved_stars:
star_age_string = "a_max"
else:
star_age_string = "a" + ("%0.3e"%(star_age.value_in(units.Myr))).replace("+0","").replace("+","")
base_output_file_name = os.path.join(get_path_to_results(), "stellar_merger_"+n_string+"_"+t_end_string)
pickle_file_1 = os.path.join(get_path_to_results(), "stellar_merger_"+masses_string+"_"+star_age_string+"_1.pkl")
pickle_file_2 = os.path.join(get_path_to_results(), "stellar_merger_"+masses_string+"_"+star_age_string+"_2.pkl")
if not use_stored_stellar_models or not (os.path.exists(pickle_file_1) and os.path.exists(pickle_file_2)):
stars = Particles(2)
stars.mass = masses
try:
stellar_evolution = MESA()
stellar_evolution.initialize_code()
except:
print "MESA was not built. Returning."
return
stellar_evolution.commit_parameters()
stellar_evolution.particles.add_particles(stars)
stellar_evolution.commit_particles()
if maximally_evolved_stars:
try:
while True:
stellar_evolution.evolve_model()
except AmuseException as exception:
print exception
else:
stellar_evolution.evolve_model(star_age)
if os.path.exists(pickle_file_1):
print "Could not save stellar model 1: file already exists."
else:
pickle_stellar_model(stellar_evolution.particles[0], pickle_file_1)
print "Stellar model 1 saved at:", pickle_file_1
if os.path.exists(pickle_file_2):
print "Could not save stellar model 2: file already exists."
else:
pickle_stellar_model(stellar_evolution.particles[1], pickle_file_2)
print "Stellar model 2 saved at:", pickle_file_2
stellar_evolution.stop()
model_1 = StellarModel2SPH(None, None, pickle_file = pickle_file_1)
model_2 = StellarModel2SPH(None, None, pickle_file = pickle_file_2)
model_1.unpickle_stellar_structure()
model_2.unpickle_stellar_structure()
composition = model_2.composition_profile
midpoints = model_2.midpoints_profile[1:-1]
specific_internal_energy = model_2.specific_internal_energy_profile
number_of_sph_particles_1 = int(round(number_of_sph_particles *
(model_1.mass / (model_1.mass + model_2.mass))))
number_of_sph_particles_2 = number_of_sph_particles - number_of_sph_particles_1
print "Creating initial conditions from a MESA stellar evolution model:"
print model_1.mass, "star consisting of", number_of_sph_particles_1, "particles."
sph_particles_1 = convert_stellar_model_to_SPH(
None,
number_of_sph_particles_1,
seed=12345,
pickle_file = pickle_file_1
).gas_particles
print model_2.mass, "star consisting of", number_of_sph_particles_2, "particles."
sph_particles_2 = convert_stellar_model_to_SPH(
None,
number_of_sph_particles_2,
pickle_file = pickle_file_2
).gas_particles
initial_separation += model_1.radius + model_2.radius
sph_particles_2.x += numpy.cos(angle) * initial_separation
sph_particles_2.y += numpy.sin(angle) * initial_separation
sph_particles_1.vx += numpy.cos(angle) * initial_speed - numpy.sin(angle) * initial_speed_perpendicular
sph_particles_1.vy += numpy.cos(angle) * initial_speed_perpendicular + numpy.sin(angle) * initial_speed
view = [-0.5, 0.5, -0.5, 0.5] * (initial_separation + model_1.radius + model_2.radius)
all_sph_particles = ParticlesSuperset([sph_particles_1, sph_particles_2])
all_sph_particles.move_to_center()
unit_converter = ConvertBetweenGenericAndSiUnits(1.0 | units.RSun, constants.G, t_end)
hydro_legacy_code = sph_code(unit_converter)
n_steps = 100
hydro_legacy_code.parameters.n_smooth = 96
try:
hydro_legacy_code.parameters.timestep = t_end / n_steps
except Exception as exc:
if not "parameter is read-only" in str(exc): raise
hydro_legacy_code.gas_particles.add_particles(all_sph_particles)
times = [] | units.Myr
kinetic_energies = [] | units.J
potential_energies = [] | units.J
thermal_energies = [] | units.J
print "Evolving to:", t_end
for time, i_step in [(i*t_end/n_steps, i) for i in range(1, n_steps+1)]:
hydro_legacy_code.evolve_model(time)
times.append(time)
kinetic_energies.append( hydro_legacy_code.kinetic_energy)
potential_energies.append( hydro_legacy_code.potential_energy)
thermal_energies.append( hydro_legacy_code.thermal_energy)
if steps_per_snapshot and (not i_step % steps_per_snapshot):
hydro_plot(
view,
hydro_legacy_code,
(snapshot_size, snapshot_size),
base_output_file_name + "_hydro_image{0:=03}.png".format(i_step)
)
hydro_legacy_code.gas_particles.new_channel_to(all_sph_particles).copy_attributes(
['mass', 'x','y','z', 'vx','vy','vz', 'u'])
center_of_mass = all_sph_particles.center_of_mass().as_quantity_in(units.RSun)
center_of_mass_velocity = all_sph_particles.center_of_mass_velocity().as_quantity_in(units.km / units.s)
print
print "center_of_mass:", center_of_mass
print "center_of_mass_velocity:", center_of_mass_velocity
all_sph_particles.position -= center_of_mass
sph_midpoints = all_sph_particles.position.lengths()
energy_plot(
times,
kinetic_energies, potential_energies, thermal_energies,
base_output_file_name+"_energy_evolution.png"
)
thermal_energy_plot(
times,
thermal_energies,
base_output_file_name+"_thermal_energy_evolution.png"
)
composition_comparison_plot(
midpoints, composition[0],
sph_midpoints, all_sph_particles.h1,
base_output_file_name+"_composition_h1.png"
)
internal_energy_comparison_plot(
midpoints, specific_internal_energy,
sph_midpoints, all_sph_particles.u,
base_output_file_name+"_new_u.png"
)
hydro_plot(
[-2.0, 2.0, -2.0, 2.0] | units.RSun,
hydro_legacy_code,
(100, 100),
base_output_file_name + "_hydro_image.png"
)
hydro_legacy_code.stop()
print "All done!\n"
def composition_comparison_plot(radii_SE, comp_SE, radii_SPH, comp_SPH, figname):
if not HAS_MATPLOTLIB:
return
pyplot.figure(figsize = (7, 5))
plot(radii_SE.as_quantity_in(units.RSun), comp_SE,
label='stellar evolution model')
plot(radii_SPH, comp_SPH, 'go', label='SPH model')
xlabel('radius')
ylabel('mass fraction')
pyplot.legend()
pyplot.savefig(figname)
print "\nPlot of composition profiles was saved to: ", figname
pyplot.close()
def internal_energy_comparison_plot(radii_SE, u_SE, radii_SPH, u_SPH, figname):
if not HAS_MATPLOTLIB:
return
pyplot.figure(figsize = (7, 5))
semilogy(radii_SE.as_quantity_in(units.RSun), u_SE,
label='stellar evolution model')
semilogy(radii_SPH, u_SPH, 'go', label='SPH model')
xlabel('radius')
ylabel('internal energy')
pyplot.legend()
pyplot.savefig(figname)
print "\nPlot of internal energy profiles was saved to: ", figname
pyplot.close()
def energy_plot(time, E_kin, E_pot, E_therm, figname):
if not HAS_MATPLOTLIB:
return
pyplot.figure(figsize = (5, 5))
plot(time, E_kin.as_quantity_in(units.erg), label='E_kin')
plot(time, E_pot, label='E_pot')
plot(time, E_therm, label='E_therm')
plot(time, E_kin+E_pot+E_therm, label='E_total')
xlabel('Time')
ylabel('Energy')
pyplot.legend(loc=3)
pyplot.savefig(figname)
print "\nPlot of energy evolution was saved to: ", figname
pyplot.close()
def thermal_energy_plot(time, E_therm, figname):
if not HAS_MATPLOTLIB:
return
pyplot.figure(figsize = (5, 5))
plot(time, E_therm.as_quantity_in(units.erg), label='E_therm')
xlabel('Time')
ylabel('Energy')
pyplot.legend(loc=3)
pyplot.savefig(figname)
print "\nPlot of thermal energy evolution was saved to: ", figname
pyplot.close()
def hydro_plot(view, hydro_code, image_size, figname):
"""
view: the (physical) region to plot [xmin, xmax, ymin, ymax]
hydro_code: hydrodynamics code in which the gas to be plotted is defined
image_size: size of the output image in pixels (x, y)
"""
if not HAS_MATPLOTLIB:
return
shape = (image_size[0], image_size[1], 1)
size = image_size[0] * image_size[1]
axis_lengths = [0.0, 0.0, 0.0] | units.m
axis_lengths[0] = view[1] - view[0]
axis_lengths[1] = view[3] - view[2]
grid = Grid.create(shape, axis_lengths)
grid.x += view[0]
grid.y += view[2]
speed = grid.z.reshape(size) * (0 | 1/units.s)
rho, rhovx, rhovy, rhovz, rhoe = hydro_code.get_hydro_state_at_point(grid.x.reshape(size),
grid.y.reshape(size), grid.z.reshape(size), speed, speed, speed)
min_v = 800.0 | units.km / units.s
max_v = 3000.0 | units.km / units.s
min_rho = 3.0e-4 | units.g / units.cm**3
max_rho = 0.1 | units.g / units.cm**3
min_E = 1.0e11 | units.J / units.kg
max_E = 1.0e13 | units.J / units.kg
v_sqr = (rhovx**2 + rhovy**2 + rhovz**2) / rho**2
E = rhoe / rho
log_v = numpy.log((v_sqr / min_v**2)) / numpy.log((max_v**2 / min_v**2))
log_rho = numpy.log((rho / min_rho)) / numpy.log((max_rho / min_rho))
log_E = numpy.log((E / min_E)) / numpy.log((max_E / min_E))
red = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_rho)).reshape(shape)
green = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_v)).reshape(shape)
blue = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_E)).reshape(shape)
alpha = numpy.minimum(numpy.ones_like(log_v), numpy.maximum(numpy.zeros_like(log_v),
numpy.log((rho / (10*min_rho))))).reshape(shape)
rgba = numpy.concatenate((red, green, blue, alpha), axis = 2)
pyplot.figure(figsize = (image_size[0]/100.0, image_size[1]/100.0), dpi = 100)
im = pyplot.figimage(rgba, origin='lower')
pyplot.savefig(figname, transparent=True, dpi = 100)
print "\nHydroplot was saved to: ", figname
pyplot.close()
def new_option_parser():
result = OptionParser()
result.add_option(
"-n",
"--n_sph",
dest="number_of_sph_particles",
type="int",
default = 300,
help="Number of sph particles per star in solar mass"
)
result.add_option(
"-a",
"--mass1",
dest="mass1",
type="float",
default = 1.0,
help="Mass of the first star"
)
result.add_option(
"-b",
"--mass2",
dest="mass2",
type="float",
default = 1.0,
help="Mass of the second star in solar mass"
)
result.add_option(
"-d",
"--distance",
dest="initial_separation",
type="float",
default=4.0,
help="Initial_separation between the two stars in solar radii"
)
result.add_option(
"-v",
"--speed",
dest="initial_speed",
type="float",
default=100.0,
help="Initial relative speed between the two stars in km/s"
)
result.add_option(
"-t",
"--end_time",
dest="end_time",
type="float",
default=2.0e4,
help="Time to evolve to in s"
)
result.add_option(
"-s",
"--t_star",
dest="t_star",
type="float",
default=1.0e4,
help="Age of both stars in Myr"
)
result.add_option(
"-e",
"--max_evolved",
default=False,
action="store_true",
dest="maximally_evolved_stars",
help="Flag to use maximally evolved stars, i.e. until MESA exits"
)
return result
if __name__ == "__main__":
options, arguments = new_option_parser().parse_args()
print "Simulating a head on collision between two stars."
print "Options used:", options
masses = [options.mass1, options.mass2] | units.MSun
head_on_stellar_merger(
masses = masses,
star_age = options.t_star | units.Myr,
initial_separation = options.initial_separation | units.RSun,
initial_speed = options.initial_speed | units.km / units.s,
number_of_sph_particles = options.number_of_sph_particles,
t_end = options.end_time | units.s,
maximally_evolved_stars = options.maximally_evolved_stars
)
Keywords: python, amuse, astrophysics, matplotlib, pylab, example, codex (see how-to-search-examples)