Pastebin

from yt.mods import *
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
import os.path
import numpy as np


Zoom = np.array([2, 16])
Zoom = 15360/Zoom
Time = np.array(range(0,2))/1.

if len(sys.argv) == 2:
    last = int(sys.argv[1])
else:
    # Gets the last data file
    for i in range(10000,0,-1):
        last = i;
        if not os.path.exists('./data.%04d.3d.hdf5' %i): continue
        else: break

Time = Time*last
Time[0] += 1.


fns = []
for i in Time:
     fns.append('data.%04d.3d.hdf5' %i)

fig = plt.figure()

# See http://matplotlib.org/mpl_toolkits/axes_grid/api/axes_grid_api.html
# These choices of keyword arguments produce a four panel plot with a single
# shared narrow colorbar on the right hand side of the multipanel plot. Axes
# labels are drawn for all plots since we're slicing along different directions
# for each plot.
grid = AxesGrid(fig, (0.075,0.075,0.85,0.85),
                nrows_ncols = (np.size(Time), np.size(Zoom)),
                axes_pad = 0.05,
                label_mode = "L",
                share_all = False,
                cbar_location="right",
                cbar_mode="single",
                cbar_size="3%",
                cbar_pad="0%")

for row, fn in enumerate(fns):
    # Load the data and create a single plot
    pf = load(fn) # load data
    dd = pf.h.all_data()
    t = pf.current_time*pf['years']/1000;
    masses = dd['particle_mass']
    if masses.size == 0:
        c = [0,0,0]
    else:
        index = masses.argmax()
        c = [dd['particle_position_' + direction][index] for direction in ['x', 'y', 'z']]
    sp = pf.h.sphere(c, 50./pf['au'])
    L = sp.quantities['AngularMomentumVector']()

    for col, z in enumerate(Zoom):
        p = ProjectionPlot(pf, 'z', 'Density', c, width=(z, 'au'))

        # Ensure the colorbar limits match for all plots
        p.set_zlim('Density', 5e-2, 1e3)

        # This forces the ProjectionPlot to redraw itself on the AxesGrid axes.
        plot = p.plots['Density']
        plot.figure = fig
        plot.axes = grid[row*np.size(Zoom) + col].axes
        plot.cax = grid.cbar_axes[row*np.size(Zoom) + col]

        # Finally, this actually redraws the plot.
        p._setup_plots()

plt.savefig('multiplot6x4TimeSeries.png')