textbook example code: explode_evolved_star.pyΒΆ

[source code]

import numpy
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 generic_unit_system
from amuse.units import nbody_system
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 *

from amuse.datamodel import Particles
from amuse.datamodel import Grid

def inject_supernova_energy(gas_particles):
    Rinner = 10|units.RSun
    inner = gas_particles.select(lambda pos : pos.length_squared() < Rinner**2, ["position"])
    print , "innermost particles selected."
    print "Adding", (1.0e51 | units.erg) / inner.total_mass(), "to each of",\
        len(inner), "inner most particles of the expliding statr."
    inner.u += (1.0e51 | units.erg) / inner.total_mass()

def setup_stellar_evolution_model():
    out_pickle_file = os.path.join(get_path_to_results(), "super_giant_stellar_structure.pkl")
    if os.path.exists(out_pickle_file):
        return out_pickle_file

    stellar_evolution = MESA(redirection = "none")
    stars =  Particles(1)
    stars.mass = 10.0 | units.MSun
    stellar_evolution.particles.add_particles(stars)

    """
    print "Evolving a MESA star with mass:", stellar_evolution.particles[0].mass
    try:
        while True:
            stellar_evolution.evolve_model()
    except AmuseException as ex:
        print "Evolved star to t=", stellar_evolution.particles[0].age,
        print "Mass:", stellar_evolution.particles[0].mass
        print "Radius:", stellar_evolution.particles[0].radius
        print "core-mass:", stellar_evolution.particles[0].core_mass
    """
    while stellar_evolution.particles[0].stellar_type <= 12|units.stellar_type:
        stellar_evolution.evolve_model()
#    while stellar_evolution.particles[0].stellar_type <= 3|units.stellar_type:
#        stellar_evolution.evolve_model()

    pickle_stellar_model(stellar_evolution.particles[0], out_pickle_file)
    stellar_evolution.stop()
    return out_pickle_file

def run_supernova():
    # options:
    use_hydro_code = Gadget2 # Fi -or- Gadget2
    hydro_code_options = dict(number_of_workers=3) # e.g. dict(use_gl = True) or dict(redirection = "none")
#    number_of_sph_particles = 30000
    t_end = 1.0e5 | units.s
    number_of_sph_particles = 3000
    t_end = 1.0e4 | units.s

    pickle_file = setup_stellar_evolution_model()

    print "Creating initial conditions from a MESA stellar evolution model..."
    model = convert_stellar_model_to_SPH(
        None,
        number_of_sph_particles,
        seed = 12345,
        pickle_file = pickle_file,
#        base_grid_options = dict(type = "glass", target_rms = 0.01),
        with_core_particle = True,
        target_core_mass = 1.4|units.MSun
    )
    print "model=", model.core_particle
    core, gas_without_core, core_radius = model.core_particle, model.gas_particles, model.core_radius
    """
    if len(core):
        print "Created", len(gas_without_core), "SPH particles and one 'core-particle':\n", core
        print "Setting gravitational smoothing to:", core_radius
    else:
        print "Warning: Only SPH particles created."
    """

    inject_supernova_energy(gas_without_core)

    print "\nEvolving (SPH) to:", t_end
    n_steps = 100

    unit_converter = ConvertBetweenGenericAndSiUnits(1.0 | units.RSun, constants.G, t_end)
    hydro_code = use_hydro_code(unit_converter, **hydro_code_options)

    try:
        hydro_code.parameters.timestep = t_end / n_steps
    except Exception as exc:
        if not "parameter is read-only" in str(exc): raise

    hydro_code.parameters.epsilon_squared = core_radius**2
    hydro_code.parameters.n_smooth_tol = 0.01
    hydro_code.gas_particles.add_particles(gas_without_core)
    hydro_code.dm_particles.add_particle(core)

    times = [] | units.s
    potential_energies = [] | units.J
    kinetic_energies =   [] | units.J
    thermal_energies =   [] | units.J
    for time, i_step in [(i*t_end/n_steps, i) for i in range(0, n_steps+1)]:
        hydro_code.evolve_model(time)
        times.append(time)
        potential_energies.append( hydro_code.potential_energy)
        kinetic_energies.append(   hydro_code.kinetic_energy)
        thermal_energies.append(   hydro_code.thermal_energy)
        hydro_plot(
            [-1.0, 1.0, -1.0, 1.0] * (350 | units.RSun),
            hydro_code,
            (100, 100),
            os.path.join(get_path_to_results(), "supernova_hydro_image{0:=03}.png".format(i_step))
        )

    energy_plot(times, kinetic_energies, potential_energies, thermal_energies,
        os.path.join(get_path_to_results(), "supernova_energy_evolution.png"))

    hydro_code.stop()
    print "All done!\n"

from prepare_figure import single_frame, figure_frame, set_tickmarks
from distinct_colours import get_distinct

def energy_plot(time, E_kin, E_pot, E_therm, figname):
    if not HAS_MATPLOTLIB:
        return
    x_label = 'Time [hour]'
    y_label = 'Energy [foe]'
    single_frame(x_label, y_label, logx=False, logy=False, xsize=14, ysize=10, ymin=-1, ymax=-1)
    cols = get_distinct(4)

    FOE = 1.e+51 | units.erg
    hour = 1|units.hour
    pyplot.plot(time/hour, E_kin/FOE, label='E_kin', c=cols[0])
    pyplot.plot(time/hour, E_pot/FOE, label='E_pot', c=cols[1])
    pyplot.plot(time/hour, E_therm/FOE, label='E_therm', c=cols[2])
    pyplot.plot(time/hour, (E_kin+E_pot+E_therm)/FOE, label='E_total',
                c=cols[3])

    pyplot.legend(loc=4)
    pyplot.savefig(figname)
    print "\nPlot of 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-9 | units.g / units.cm**3
    max_rho = 1.0e-5 | 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, facecolor='k', edgecolor='k')
    print "\nHydroplot was saved to: ", figname
    pyplot.close()


if __name__ == "__main__":
    print "Test run to mimic a supernova in SPH"
    print
    print "Details:"
    print "First a high-mass star is evolved up to the super giant phase using MESA. " \
        "Then it is converted to SPH particles with the convert_stellar_model_to_SPH " \
        "procedure (with a non-SPH 'core' particle). Finally the internal energies of " \
        "the inner particles are increased, such that the star gains the 10^51 ergs " \
        "released in supernova explosions."
    print
    run_supernova()

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