plot.py

#! /usr/bin/env python
'''Plot the output of the planet-disk simulation.'''

from __future__ import division

import sys

from soapp.plot.voronoi import voronoi
from soapp.parser.arepo.parser import ArepoParser

from matplotlib.pyplot import *
from matplotlib.colors import Normalize
import numpy as np


def parse_output(fname):
    '''Parses output of particular frame.'''
    parser = ArepoParser()
    return parser.parse(fname)


def plot(dump):
    # Shortcuts
    rho = dump['density']
    pos = dump['position']
    vel = dump['velocity']
    e_int = dump['internal_energy']

    # Plot stuff
    aspect_ratio = 16 / 9
    height = 10.8
    figure(figsize=(aspect_ratio * height, height))
    ax = subplot(111)
    title('t = {}'.format(dump.time))
    xlim(0.0, 1.0)
    ylim(0.15, 0.85)
    voronoi_diagram = voronoi(ax, pos.data['f0'][:len(rho.data)],
        pos.data['f1'][:len(rho.data)],
        rho.data, edgecolors='0.8', linewidths=0.3)
    voronoi_diagram.set_norm(Normalize(0.0, 4.0))
    colorbar(voronoi_diagram)

if __name__ == '__main__':
    # Parse stuff
    dump = parse_output(sys.argv[1])
    # Plot it
    plot(dump)
    show()