hist2dtest.py

import ROOT
import random
import sys
from array import array


def make_th2d(name, title, data, xmark=lambda x: x, ymark=lambda y: str(y), zrange=(0, 1), normal_y=False):
    """
    Build TH2D from a 2D array of values; use dict keys as axis labels

    :param normal_y: sort Y axis labels in asceding order (default: false)
    :type normal_y: bool
    :param name: histogram name (must be unique)
    :type name: str
    :param title: histogram title
    :type title: str
    :type data: dict(dict(int)) | dict(dict(float))
    :param data:
    :param xmark: function to format X labels (e.g. add mark)
    :type xmark: lambda | function
    :param ymark: function to format Y labels (e.g. add mark)
    :type ymark: lambda | function
    :param zrange: range on Z axis (default: (0,1))
    :type zrange: tuple | list

    :rtype: ROOT.TH2D
    """

    x_labels = set()
    y_labels = set()
    for k in data.keys():
        x_labels.add(k)
        y_labels.update(data[k].keys())

    x_labels = sorted(x_labels)
    y_labels = sorted(y_labels)

    h2 = ROOT.TH2D(name, title, len(x_labels), 0, len(x_labels), len(y_labels), 0, len(y_labels))
    for ix, x in enumerate(x_labels):
        for iy, y in enumerate(y_labels):
            try:
                z = data[x][y]
                h2.Fill(ix, iy, z)
            except KeyError:
                pass

    for ix, x in enumerate(x_labels):
        h2.GetXaxis().SetBinLabel(ix + 1, xmark(x))

    for iy, y in enumerate(y_labels):
        h2.GetYaxis().SetBinLabel(iy + 1, ymark(y))

    if zrange[0] is not None:
        h2.SetMinimum(zrange[0])

    if zrange[1] is not None:
        h2.SetMaximum(zrange[1])

    h2.GetXaxis().SetLabelFont(83)
    h2.GetXaxis().SetLabelSize(14)

    h2.GetYaxis().SetLabelFont(83)
    h2.GetYaxis().SetLabelSize(14)

    h2.GetZaxis().SetLabelFont(83)
    h2.GetZaxis().SetLabelSize(14)

    h2.GetXaxis().LabelsOption("v")

    ROOT.SetOwnership(h2, False)
    return h2


def make_data(nx, ny):
    data = {}
    for x in range(nx):
        data[x] = {}
        for y in range(ny):
            data[x][y] = random.random()

    return data


def main():
    if len(sys.argv) < 3:
        print("Usage: hist2dtest nx ny")
        exit(0)

    nx = int(sys.argv[1])
    ny = int(sys.argv[2])
    ROOT.gStyle.SetOptStat(0)
    ROOT.gROOT.SetBatch(True)

    colors_r = array('d', [1.0, 1.0])
    colors_g = array('d', [1.0, 0.0])
    colors_b = array('d', [0.62, 0.0])
    colors_s = array('d', [0.0, 1.0])
    ROOT.TColor.CreateGradientColorTable(2, colors_s, colors_r, colors_g, colors_b, 100)

    data = make_data(nx, ny)
    width = max(800, 15 * nx)
    c = ROOT.TCanvas("c0", "c0", width, max(600, 15 * ny))
    pad = ROOT.TPad("pad", "my pad", 0.05, 0.05, 0.95, 0.99)
    pad.Draw()
    pad.cd()
    pad.SetGrid()

    h2 = make_th2d("TEST", "THIS IS A TEST", data, str, str)
    h2.SetContour(20)
    h2.Draw("COLZ")

    c.Update()
    c.SaveAs("plot_{0}x{1}.png".format(nx, ny))


if __name__ == "__main__":
    main()