from pyhs3.data import BinnedData
from pyhs3.axes import BinnedAxis

# Create 2D binned data (3x4 = 12 bins)
contents = [1.0, 2.0, 3.0, 4.0,
            5.0, 6.0, 7.0, 8.0,
            9.0, 10.0, 11.0, 12.0]

data = BinnedData(
    name="2d_hist",
    type="binned",
    contents=contents,
    axes=[
        BinnedAxis(name="x", min=0.0, max=3.0, nbins=3),
        BinnedAxis(name="y", min=0.0, max=4.0, nbins=4),
    ]
)

# Convert to hist and plot as heatmap
h = data.to_hist()
plt.pcolormesh(*h.axes.edges.T, h.values().T, cmap="viridis")
plt.colorbar(label="Events")
plt.xlabel("x")
plt.ylabel("y")
plt.title("2D Histogram")