py: Plot tree image example

plot.py

import numpy as np
import h5py as h5
import matplotlib.pyplot as plt
from pathlib import Path

DATA_PATH = Path("tree_images.h5")
ID = "619000000001104"

with h5.File(DATA_PATH, "r") as fp:
    image = fp[f"{ID}/mass"][:]
    first_snap = fp[ID].attrs['first_snap']
    last_snap = fp[ID].attrs['last_snap']

fig, ax = plt.subplots(1, 1, figsize=(12, 12))
ax: plt.Axes

im = ax.pcolormesh(np.arange(image.shape[1]), np.arange(first_snap, last_snap+1), np.log10(image))
plt.colorbar(im, label=r"$\log_{10}(M_{200}/{\rm M_\odot})$", aspect=50)

ax.set_xlim(-2, None)
ax.set_ylabel("snapshot")

ax.set_aspect(0.3)
ax.spines.right.set_visible(False)
ax.spines.top.set_visible(False)
ax.spines.bottom.set_visible(False)
ax.xaxis.set_visible(False)

plt.savefig(f"{ID}.png", bbox_inches="tight")