%matplotlib inline
import yt
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid

# fractional=True keyword -- divide the profile data by the sum.
fns = [ "Enzo_64/DD0030/data0030", "Enzo_64/DD0040/data0040"]

fig = plt.figure()
grid = AxesGrid(fig, (0.075,0.075,0.85,0.85),
                nrows_ncols = (1, 1),
                axes_pad = 0.05,
                label_mode = "L",
                share_all = True,
                cbar_location="right",
                cbar_mode="single",
                cbar_size="3%",
                cbar_pad="0%")
fig.suptitle("Gas Phase")

for i, fn in enumerate(fns[:1]):
    ds = yt.load(fn)
    ad = ds.all_data()
    p = yt.PhasePlot(ad, "density", "temperature", "cell_mass", x_bins=256, y_bins=256, weight_field=None, fractional=True)
    #p.set_log('density', False)
    #p.set_log('temperature', False)
    plot = p.plots['cell_mass']
    if i % 2 == 0:
        ptype = "Gas phase ${H_2}$ cooling "
    else:
        ptype = "Gas phase Standard cooling "
    zStr = ptype + "z = %.1f" % ds.current_redshift
    plot.figure = fig
    plot.axes = grid[i].axes
    plot.cax = grid.cbar_axes[i]

    #p['cell_mass'].axes.set_title(zStr, fontsize=20)
    
    p._setup_plots()
    p.plots['cell_mass'].axes.set_aspect(1e-35)

plt.show()