plot_stars.py

import sys
from amuse.units import units
from amuse.io import read_set_from_file
import matplotlib.pyplot as plt


def plot_stars(stars, length_unit=units.pc, plot_by_luminosity=True, scale=1/40):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    if plot_by_luminosity:
        if hasattr(stars, "luminosity"):
            s = stars.luminosity.value_in(units.LSun)**(0.5)
        else:
            # approximate the luminosity with mass ^ 3.5
            s = stars.mass.value_in(units.MSun)**(3.5/2)
    else:
        s = 1 / scale
    ax.scatter(
        stars.x.value_in(length_unit),
        stars.y.value_in(length_unit),
        s=s * scale,
        edgecolors="none",
    )
    ax.set_xlim([-5, 5])
    ax.set_ylim([-5, 5])
    ax.set_aspect(1)
    time = stars.get_timestamp()

    if time is not None:
        ax.set_title(f"time = {time.in_(units.Myr)}")
    plt.show()


def main():
    stars = read_set_from_file(sys.argv[1])
    plot_stars(stars, plot_by_luminosity=False)


if __name__ == "__main__":
    main()