applications example code: test_evrard_collapse.pyΒΆ

[source code]

import os.path
from optparse import OptionParser

try:
    from matplotlib import pyplot
    HAS_MATPLOTLIB = True
    from amuse.plot import plot, xlabel, ylabel
except ImportError:
    HAS_MATPLOTLIB = False

from amuse.test.amusetest import get_path_to_results
from amuse.units import units
from amuse.units import constants
from amuse.units import nbody_system
from amuse.units.generic_unit_converter import ConvertBetweenGenericAndSiUnits
from amuse.ext.evrard_test import new_evrard_gas_sphere
from amuse.community.gadget2.interface import Gadget2
from amuse.community.fi.interface import Fi

from amuse.rfi.core import is_mpd_running
usage = """\
usage: amuse.sh %prog [options]

This example script will simulate the spherical collapse of
an initially-cold adiabatic gas cloud, with a radial density
profile proportional to 1/r.
"""

labels = [['K', 'V', 'U', 'E'], ['K', 'V', 'U', 'E']]
fi_gadget_labels = [['K_fi', 'V_fi', 'U_fi', 'E_fi'], ['K_gadget', 'V_gadget', 'U_gadget', 'E_gadget']]

convert_generic_units = ConvertBetweenGenericAndSiUnits(1.0 | units.kpc, 1.0e10 | units.MSun, constants.G)
convert_nbody_units   = nbody_system.nbody_to_si(       1.0 | units.kpc, 1.0e10 | units.MSun)

def run_evrard(
        hydro_legacy_codes,
        number_of_particles,
        random_seed = None,
        name_of_the_figure = "evrard_collapse_test.png"):

    print "\nThe spherical collapse of an initially-cold adiabatic gas cloud,\n", \
        "consisting of ", str(number_of_particles), "particles will be simulated...\n"

    t_end = 3.0 * convert_nbody_units.to_si(1.0 | nbody_system.time) # (3 natural timescales)
    print "Evolving to (3 natural timescales): ", t_end.as_quantity_in(units.Myr)
    n_steps = 100

    gas = new_evrard_gas_sphere(number_of_particles, convert_nbody_units, do_scale = True, seed = random_seed)
    gas.h_smooth = 0.01 | units.kpc

    times = [] | units.Myr
    kinetic_energies =   []
    potential_energies = []
    thermal_energies =   []

    for hydro_legacy_code in hydro_legacy_codes:
        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(gas)

        kinetic_energies.append([] | units.J)
        potential_energies.append([] | units.J)
        thermal_energies.append([] | units.J)

    for time in [i*t_end/n_steps for i in range(1, n_steps+1)]:
        for i, hydro_legacy_code in enumerate(hydro_legacy_codes):
            hydro_legacy_code.evolve_model(time)
            kinetic_energies[i].append(   hydro_legacy_code.kinetic_energy)
            potential_energies[i].append( hydro_legacy_code.potential_energy)
            thermal_energies[i].append(   hydro_legacy_code.thermal_energy)
        times.append(time)
    for hydro_legacy_code in hydro_legacy_codes:
        hydro_legacy_code.stop()
    energy_plot(times, kinetic_energies, potential_energies, thermal_energies, name_of_the_figure)
    print "All done!\n"

def energy_plot(time, E_kin_list, E_pot_list, E_therm_list, figname):
    if not HAS_MATPLOTLIB:
        return
    pyplot.figure(figsize = (5, 5))
    for i, (E_kin, E_pot, E_therm) in enumerate(zip(E_kin_list, E_pot_list, E_therm_list)):
        plot(time, E_kin.as_quantity_in(units.erg), label=labels[i][0])
        plot(time, E_pot, label=labels[i][1])
        plot(time, E_therm, label=labels[i][2])
        plot(time, E_kin+E_pot+E_therm, label=labels[i][3])
    xlabel('Time')
    ylabel('Energy')
    pyplot.legend(prop={'size':"x-small"}, loc=3)
    pyplot.savefig(figname)
    print "\nPlot of energy evolution was saved to: ", figname
    pyplot.close()

class InstantiateCode(object):

    def fi(self):
        return Fi(convert_nbody_units)

    def gadget(self):
        return Gadget2(convert_generic_units)

    def new_code(self, name_of_the_code):
        if hasattr(self, name_of_the_code):
            return getattr(self, name_of_the_code)()
        else:
            raise Exception("Cannot instantiate code with name '{0}'".format(name_of_the_code))

def new_code(name_of_the_code):
    return InstantiateCode().new_code(name_of_the_code)

def test_evrard_fi_short():
    assert is_mpd_running()
    run_evrard(
        [new_code("fi")],
        100,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_fi_100.png")
    )

def test_evrard_gadget_short():
    assert is_mpd_running()
    run_evrard(
        [new_code("gadget")],
        100,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_gadget_100.png")
    )

def test_evrard_compare_short():
    assert is_mpd_running()
    global labels
    labels = fi_gadget_labels
    run_evrard(
        [new_code("fi"), new_code("gadget")],
        100,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_compare_100.png")
    )

def slowtest_evrard_fi():
    assert is_mpd_running()
    run_evrard(
        [new_code("fi")],
        10000,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_fi_10000.png")
    )

def slowtest_evrard_gadget():
    assert is_mpd_running()
    run_evrard(
        [new_code("gadget")],
        10000,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_gadget_10000.png")
    )

def slowtest_evrard_compare():
    assert is_mpd_running()
    global labels
    labels = fi_gadget_labels
    run_evrard(
        [new_code("fi"), new_code("gadget")],
        10000,
        random_seed = 12345,
        name_of_the_figure = os.path.join(get_path_to_results(), "evrard_test_compare_10000.png")
    )

def new_commandline_option_parser():
    result = OptionParser(usage)
    result.add_option(
        "-c",
        "--code",
        choices=["fi", "gadget"],
        default="fi",
        dest="code",
        metavar="CODE",
        help="CODE to use for hydrodynamics"
    )
    result.add_option(
        "-a",
        "--all_codes",
        default=False,
        action="store_true",
        dest="all_codes",
        metavar="ALLCODE",
        help="ALLCODES flag to use all hydrodynamics solvers"
    )
    result.add_option(
        "-n",
        "--number_of_particles",
        type="int",
        default=100,
        dest="number_of_particles",
        help="Number of particles in the gas cloud"
    )
    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="evrard_collapse_test.png",
        dest="plot_filename",
        help="Name of the file to plot to"
    )
    return result

if __name__ == '__main__':
    if not is_mpd_running():
        print "There is no mpd server running. Please do 'mpd &' first."
        sys.exit()
    parser = new_commandline_option_parser()
    (options, arguments) = parser.parse_args()
    if arguments:
        parser.error("unknown arguments '{0}'".format(arguments))

    if options.all_codes:
        code = [new_code("fi"), new_code("gadget")]
        labels = fi_gadget_labels
    else:
        code = [new_code(options.code)]
    run_evrard(
        code,
        options.number_of_particles,
        random_seed = options.salpeterSeed,
        name_of_the_figure = options.plot_filename
    )

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

This Page