plot_amuse.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Plot hydro and/or stars
"""
import os
import logging
import numpy
import copy
import argparse

import matplotlib
from matplotlib import pyplot
# import matplotlib.cm as cm

from amuse.datamodel import Particles
from amuse.io import read_set_from_file
from amuse.units import units, nbody_system
try:
    from amuse.ext.mapper import MapHydro
except ImportError:
    from mapper import MapHydro


logger = logging.getLogger(__name__)


def velocity_divergence(vx_field, vy_field, velocity_unit=units.kms):
    # Return divergence of velocity fields

    div = numpy.ufunc.reduce(
        numpy.add,
        [
            -numpy.gradient(vx_field.value_in(velocity_unit), axis=1),
            -numpy.gradient(vy_field.value_in(velocity_unit), axis=0),
        ]
    ) | velocity_unit

    return div


def plot_velocity(
    ax,
    maps,
    x_axis="vx",
    y_axis="vy",
    velocity_unit=units.kms,
    color='black',
    normalise_by_counts=True,
    style="streamplot",  # or "quiver"
):
    if normalise_by_counts:
        vxmap = getattr(maps, x_axis) / maps.counts
        vymap = getattr(maps, y_axis) / maps.counts
        # vzmap = getattr(maps, z_axis) / maps.counts

    x_range = numpy.arange(
        maps.xmin, maps.xmax, maps.width/maps.bins[0]
    )
    y_range = numpy.arange(
        maps.ymin, maps.ymax, maps.height/maps.bins[1]
    )
    vx_values = vxmap.value_in(velocity_unit)
    vy_values = vymap.value_in(velocity_unit)
    x_grid, y_grid = numpy.meshgrid(x_range, y_range)

    if style == "streamplot":
        return ax.streamplot(
            x_grid, y_grid,
            vx_values, vy_values,
            color=color,
        )
    return ax.quiver(
        x_grid, y_grid,
        vx_values, vy_values,
        color=color,
    )


def plot_divergence(
    ax,
    maps,
    x_axis="vx",
    y_axis="vy",
    contours=False,
    normalise_by_counts=True,
    vmin=None,
    vmax=None,
):
    x_range = numpy.arange(
        maps.xmin.value_in(units.pc), maps.xmax.value_in(units.pc),
        (maps.width / maps.bins[0]).value_in(units.pc),
    )
    y_range = numpy.arange(
        maps.ymin.value_in(units.pc), maps.ymax.value_in(units.pc),
        (maps.height / maps.bins[1]).value_in(units.pc),
    )
    x_grid, y_grid = numpy.meshgrid(x_range, y_range)
    div = velocity_divergence(
        getattr(maps, x_axis) / maps.counts,
        getattr(maps, y_axis) / maps.counts,
    )
    if not normalise_by_counts:
        div = div * maps.counts

    if contours:
        return ax.contour(
            x_grid, y_grid,
            div.value_in(units.kms),
            levels=(-12, -4, 4, 12),
            # colors=["blue", "red"],
            # cmap="coolwarm",
            cmap='bwr',
        )
    minmax = max(-div.min(), div.max()).value_in(units.kms)
    if not vmin:
        vmin = -minmax
    if not vmax:
        vmax = minmax
    # minmax = 500
    # print(minmax)
    return ax.imshow(
        div.value_in(units.kms),
        extent=maps.extent,
        vmin=vmin,
        vmax=vmax,
        cmap='bwr',
        origin='lower',
    )


def plot_column_density(
    ax,
    maps,
    unit_col_density=units.MSun * units.pc**-2,
    vmin=-1,
    vmax=4,
    cmap='viridis',
):
    cmap = copy.copy(matplotlib.cm.get_cmap(cmap))
    cmap.set_bad('k', alpha=1.0)

    column_density = maps.column_density
    logscale = numpy.log10(column_density.value_in(unit_col_density))
    return ax.imshow(
        logscale,
        extent=maps.extent,
        vmin=vmin,
        vmax=vmax,
        origin='lower',
        cmap=cmap,
    )


def plot_temperature(
    ax,
    maps,
    vmin=0,
    vmax=5,
    cmap="inferno",
):
    cmap = copy.copy(matplotlib.cm.get_cmap(cmap))
    cmap.set_bad('k', alpha=1.0)
    temperature_map = maps.temperature
    logscale_temperature_map = numpy.log10(
        temperature_map.value_in(units.K)
    )

    return ax.imshow(
        logscale_temperature_map,
        extent=maps.extent,
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        origin="lower",
    )


def plot_hydro_and_stars(
    time,
    maps,
    title="",
    plot="column density",
    length_unit=units.parsec,
    use_fresco=False,
    image_size_scale=1,
    starscale=1,
    dpi=150,
    vmin=None,
    vmax=None,
    sinks_color='red',
    stars_color='white',
    plot_velocities=False,
):
    "Plot gas and stars"

    left = 0.2
    bottom = 0.1
    right = 1
    top = 0.9
    # fig = pyplot.figure(figsize=(6, 5))
    image_size = image_size_scale * maps.bins
    figwidth = image_size[0] / dpi / (right - left)
    figheight = image_size[1] / dpi / (top - bottom)
    figsize = (figwidth, figheight)
    fig = pyplot.figure(figsize=figsize, dpi=dpi)
    wspace = 0.5
    fig.subplots_adjust(
        left=left, right=right, top=top, bottom=bottom, wspace=wspace)

    ax = fig.add_subplot(1, 1, 1)
    ax.set_title(plot)
    if plot == "column density":
        plot_column_density(ax, maps)

    # if plot == "temperature":
    #     plot_temperature(ax, maps)

    if plot == "divergence":
        plot_divergence(
            ax,
            maps,
            normalise_by_counts=True,
            vmin=vmin,
            vmax=vmax,
        )

    if plot == "mass flux":
        plot_divergence(
            ax,
            maps,
            normalise_by_counts=False,
            vmin=vmin,
            vmax=vmax,
        )

    if plot in ["stars", "stars colour"]:
        ax.set_facecolor('black')

    cmap = copy.copy(matplotlib.cm.get_cmap("coolwarm"))
    if plot == "stars colour":
        stars_color = maps.stars.mass.value_in(units.MSun)

    if maps.sinks is not None:
        if not maps.sinks.is_empty():
            # Scale sinks the same way as stars
            s = starscale * 0.1 * (
                    (maps.sinks.mass / (7 | units.MSun))**(3.5 / 2)
                )
            x = getattr(maps.sinks, maps.axes[0]).value_in(length_unit)
            y = getattr(maps.sinks, maps.axes[1]).value_in(length_unit)
            c = (
                "black" if plot in [
                    "temperature", "divergence", "mass flux",
                ]
                else sinks_color
            )
            ax.scatter(x, y, s=s, c=c, cmap=cmap, lw=0)

    if maps.stars is not None:
        if not maps.stars.is_empty():
            s = starscale * 0.1 * (
                (maps.stars.mass / (7 | units.MSun))**(3.5 / 2)
            )
            x = getattr(maps.stars, maps.axes[0]).value_in(length_unit)
            y = getattr(maps.stars, maps.axes[1]).value_in(length_unit)
            c = (
                "black" if plot in [
                    "temperature", "divergence", "mass flux",
                ]
                else stars_color
            )
            if plot == "stars colour":
                s = starscale * maps.stars.mass.value_in(units.MSun)
            if not use_fresco:
                use_fresco = 0
            ax.scatter(x, y, s=s, c=c, cmap=cmap, lw=0, alpha=1-use_fresco)

    if maps.stars is not None and use_fresco:
        from amuse.ext.fresco.fresco import make_image
        converter = nbody_system.nbody_to_si(
            maps.stars.total_mass(),
            maps.width,
        )
        stars = maps.stars.copy()
        stars.x -= maps.offset[0]
        stars.y -= maps.offset[1]
        stars.z -= maps.offset[2]
        gas = maps.gas.copy()
        if gas is not None:
            gas.x -= maps.offset[0]
            gas.y -= maps.offset[1]
            gas.z -= maps.offset[2]

        fresco_image, vmax = make_image(
            stars=stars,
            gas=gas,
            converter=converter,
            image_width=maps.width,
            image_size=maps.bins,
            percentile=0.9995,
            calc_temperature=True,
            age=0 | units.Myr,
            vmax=None,
            sourcebands='ubvri',
            zoom_factor=maps.bins[0]/2048,
            psf_type='hubble',
            psf_sigma=1.0,
            return_vmax=True,
            extinction=True,
        )
        ax.imshow(
            fresco_image,
            extent=maps.extent,
            alpha=use_fresco,
            origin='lower',
        )

    if plot_velocities:
        plot_velocity(
            ax, maps,
            velocity_unit=units.kms,
            color=(
                "black" if plot in ["temperature", "divergence"]
                else stars_color
            ),
            # smoothing_length=100 | units.pc,
            normalise_by_counts=True,
        )

    ax.set_xlim(maps.extent[0], maps.extent[1])
    ax.set_ylim(maps.extent[2], maps.extent[3])
    ax.set_xlabel("%s [%s]" % (maps.axes[0], length_unit))
    ax.set_ylabel("%s [%s]" % (maps.axes[1], length_unit))

    # pyplot.title(title)
    fig.suptitle(title)

    return fig, ax


def new_argument_parser():
    "Parse command line arguments"
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-s',
        dest='starsfilename',
        default='',
        help='file containing stars (optional) []',
    )
    parser.add_argument(
        '-i',
        dest='sinksfilename',
        default='',
        help='file containing sinks (optional) []',
    )
    parser.add_argument(
        '-g',
        dest='gasfilename',
        default='',
        help='file containing gas (optional) []',
    )
    parser.add_argument(
        '-o',
        dest='imagefilename',
        default='test',
        help='write image to this file [test]',
    )
    parser.add_argument(
        '-n',
        dest='bins',
        default=800,
        type=int,
        help='number of bins (800)',
    )
    parser.add_argument(
        '-x',
        dest='x',
        default=0,
        type=float,
        help='Central X coordinate (0 [pc])',
    )
    parser.add_argument(
        '-y',
        dest='y',
        default=0,
        type=float,
        help='Central Y coordinate (0 [pc])',
    )
    parser.add_argument(
        '-z',
        dest='z',
        default=0,
        type=float,
        help='Central Z coordinate (0 [pc])',
    )
    parser.add_argument(
        '-w',
        dest='w',
        default=10,
        type=float,
        help='Width in pc (10)',
    )
    parser.add_argument(
        '-d',
        dest='thickness',
        default=0,
        type=float,
        help='Thickness in pc (0)',
    )
    parser.add_argument(
        '--com',
        dest='use_com',
        action='store_true',
        default=False,
        help='Center on center of mass [False]',
    )
    parser.add_argument(
        '--fresco',
        dest='fresco_fraction',
        default=0,
        type=float,
        help='Fresco mix-in fraction (0-1)',
    )
    parser.add_argument(
        '-X',
        dest='x_axis',
        default='x',
        help='Horizontal axis ["x"]',
    )
    parser.add_argument(
        '-Y',
        dest='y_axis',
        default='y',
        help='Vertical axis ["y"]',
    )
    parser.add_argument(
        '-Z',
        dest='z_axis',
        default='z',
        help='Line-of-sight axis ["z"]',
    )
    parser.add_argument(
        '--starscale',
        dest='starscale',
        type=float,
        default=1,
        help='Scale of stars and sinks (1)',
    )

    return parser.parse_args()


def main():
    o = new_argument_parser()
    gasfilename = o.gasfilename
    starsfilename = o.starsfilename
    sinksfilename = o.sinksfilename
    imagefilename = o.imagefilename
    bins = o.bins
    offset_x = o.x | units.pc
    offset_y = o.y | units.pc
    offset_z = o.z | units.pc
    width = o.w | units.pc
    x_axis = o.x_axis
    y_axis = o.y_axis
    z_axis = o.z_axis
    starscale = o.starscale
    dpi = 150
    fresco_fraction = 0  # fresco mix-in rate; 0 = not at all, 1 = only fresco
    plots = []

    if os.path.isfile(starsfilename):
        stars = read_set_from_file(
            starsfilename,
            "amuse",
        ) if starsfilename != "" else Particles()
    else:
        stars = Particles()
    if os.path.isfile(sinksfilename):
        sinks = read_set_from_file(
            sinksfilename,
            "amuse",
        ) if sinksfilename != "" else Particles()
    else:
        sinks = Particles()
    if gasfilename:
        gas = read_set_from_file(
            gasfilename,
            "amuse",
        )
        if hasattr(gas, "itype"):
            gas = gas[gas.itype == 1]
        # gas.h_smooth = gas.h
        if hasattr(gas, "u"):
            plots.append("temperature")
    else:
        gas = Particles()

    if not gas.is_empty():
        plots += ["column density", "divergence", "mass flux"]
        if hasattr(gas, 'u'):
            plots.append("temperature")
    if not stars.is_empty():
        plots += ["stars", "stars colour"]

    if not hasattr(gas, "h_smooth"):
        gas.h_smooth = 1 | units.pc

    if not gas.is_empty():
        mtot = gas.total_mass()
        com = mtot * gas.center_of_mass()
        time = gas.get_timestamp()
    else:
        mtot = 0 | units.MSun
        com = [0, 0, 0] | units.pc * units.MSun
        time = False
    if not sinks.is_empty():
        mtot += sinks.total_mass()
        com += sinks.total_mass() * sinks.center_of_mass()
        if not time:
            time = sinks.get_timestamp()
    if not stars.is_empty():
        mtot += stars.total_mass()
        com += stars.total_mass() * stars.center_of_mass()
        if not time:
            time = stars.get_timestamp()
    com = com / mtot
    if o.use_com:
        offset_x = com[0]
        offset_y = com[1]
        offset_z = com[2]

    if time is None:
        print('Unable to get timestamp, set time to 0 Myr.')
        time = 0.0 | units.Myr

    converter = nbody_system.nbody_to_si(
        1 | units.pc, 1 | units.MSun,
    )

    width = o.w | units.pc
    bins = o.bins
    thickness = o.thickness | units.pc

    if thickness > 0 | units.pc:
        if z_axis == "x":
            gas_slab = gas[(gas.x - offset_x)**2 < (thickness/2)**2]
        elif z_axis == "y":
            gas_slab = gas[(gas.y - offset_y)**2 < (thickness/2)**2]
        else:  # if z_axis == "z":
            gas_slab = gas[(gas.z - offset_z)**2 < (thickness/2)**2]
    else:
        gas_slab = gas

    maps = MapHydro(gas_slab, stars, sinks)
    maps.width = width
    maps.bins = bins
    maps.offset = (offset_x, offset_y, offset_z)
    maps.axes = (x_axis, y_axis, z_axis)

    for plot in plots:
        if plot == "density":
            use_fresco = fresco_fraction
        else:
            use_fresco = 0
        figure, ax = plot_hydro_and_stars(
            time,
            maps,
            title="time = %06.2f %s" % (
                time.value_in(units.Myr),
                units.Myr,
            ),
            plot=plot,
            # colorbar=True,  # causes weird interpolation
            # alpha_sfe=0.02,
            # stars_are_sinks=False,
            use_fresco=use_fresco,
            length_unit=units.parsec,
            dpi=dpi,
            starscale=starscale,
        )
        plot_cluster_locations = False
        if plot_cluster_locations:
            from find_clusters import find_clusters
            clusters = find_clusters(stars, convert_nbody=converter,)
            for cluster in clusters:
                cluster_com = cluster.center_of_mass()
                cluster_x = cluster_com[0].value_in(units.parsec)
                cluster_y = cluster_com[1].value_in(units.parsec)
                lagrangian = cluster.LagrangianRadii(converter)
                lr90 = lagrangian[0][-2]
                size = lr90.value_in(units.parsec)
                # print("Circle with x, y, z: ", x, y, s)
                circle = pyplot.Circle(
                    (cluster_x, cluster_y), size, color='r', fill=False,
                )
                ax.add_artist(circle)
        pyplot.savefig(
            plot.replace(" ", "_")
            + "-" + imagefilename + ".png",
            dpi=dpi,
        )
    print(
        "This plotting script is (c) Steven Rieder. Please cite \"Rieder et al. (in prep.)\" if you use it."
    )


if __name__ == "__main__":
    main()