test_root.py

#!/usr/bin/env python

# BSD 3-Clause License; see https://github.com/scikit-hep/aghast/blob/master/LICENSE

import pytest
import numpy

import aghast
from aghast import *

ROOT = pytest.importorskip("ROOT")
connect = pytest.importorskip("aghast._connect._root")

data = [
    -0.319,
    -1.949,
    -1.511,
    1.176,
    0.695,
    -0.286,
    0.392,
    -0.094,
    -0.714,
    1.223,
    1.811,
    -2.57,
    -0.014,
    0.328,
    -0.084,
    -0.13,
    -0.751,
    0.47,
    1.558,
    0.714,
    -0.999,
    0.336,
    -2.51,
    -0.532,
    2.495,
    -0.085,
    0.165,
    0.764,
    0.865,
    -0.604,
    -1.083,
    2.501,
    -1.074,
    -0.853,
    -0.241,
    -0.346,
    0.472,
    -0.017,
    1.013,
    0.887,
    2.154,
    -1.354,
    0.966,
    -1.118,
    -1.374,
    -0.928,
    -0.664,
    2.803,
    -0.133,
    1.814,
    0.844,
    -0.719,
    1.239,
    0.49,
    0.333,
    0.518,
    0.079,
    -1.33,
    -0.048,
    1.335,
    -2.108,
    -1.772,
    -0.68,
    0.151,
    -0.479,
    0.749,
    0.589,
    -1.048,
    -0.491,
    1.125,
    1.064,
    0.4,
    0.349,
    0.193,
    -0.645,
    0.038,
    0.536,
    -0.675,
    0.732,
    -1.442,
    -0.889,
    -0.976,
    0.889,
    1.296,
    1.231,
    0.934,
    -1.359,
    1.602,
    0.186,
    -0.622,
    -0.08,
    -0.887,
    0.109,
    0.418,
    0.945,
    -0.081,
    -0.32,
    -1.309,
    -0.497,
    0.346,
]


def check1d(before, after):
    assert before.GetNbinsX() == after.GetNbinsX()
    for i in range(before.GetNbinsX() + 2):
        assert before.GetBinContent(i) == after.GetBinContent(i)
        assert before.GetBinError(i) == after.GetBinError(i)
    assert before.GetEntries() == after.GetEntries()
    assert before.GetMean() == after.GetMean()
    assert before.GetStdDev() == after.GetStdDev()
    assert before.GetTitle() == after.GetTitle()
    assert before.GetXaxis().GetTitle() == after.GetXaxis().GetTitle()
    assert bool(before.GetXaxis().GetLabels()) == bool(after.GetXaxis().GetLabels())
    assert before.GetXaxis().IsVariableBinSize() == after.GetXaxis().IsVariableBinSize()
    if before.GetXaxis().GetLabels():
        assert list(before.GetXaxis().GetLabels()) == list(after.GetXaxis().GetLabels())
    elif before.GetXaxis().IsVariableBinSize():
        beforeedges = numpy.full(before.GetNbinsX() + 1, 999, dtype=numpy.float64)
        before.GetXaxis().GetLowEdge(beforeedges)
        beforeedges[-1] = before.GetXaxis().GetBinUpEdge(before.GetNbinsX())
        afteredges = numpy.full(after.GetNbinsX() + 1, 123, dtype=numpy.float64)
        after.GetXaxis().GetLowEdge(afteredges)
        afteredges[-1] = after.GetXaxis().GetBinUpEdge(after.GetNbinsX())
        assert numpy.array_equal(before, after)
    else:
        assert before.GetXaxis().GetBinLowEdge(1) == after.GetXaxis().GetBinLowEdge(1)
        assert before.GetXaxis().GetBinUpEdge(
            before.GetNbinsX()
        ) == after.GetXaxis().GetBinUpEdge(after.GetNbinsX())


num = 0


def rootname():
    global num
    num += 1
    return "name-{0}".format(num)


@pytest.mark.parametrize("cls", [ROOT.TH1C, ROOT.TH1S, ROOT.TH1I, ROOT.TH1F, ROOT.TH1D])
def test_root_oned(cls):
    before = cls(rootname(), "title", 5, -2.0, 2.0)
    before.GetXaxis().SetTitle("title2")
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)

    before = cls(rootname(), "title", 5, -2.0, 2.0)
    before.GetXaxis().SetTitle("title2")
    for x in data:
        before.Fill(x)
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)

    before = cls(rootname(), "title", 5, -2.0, 2.0)
    before.GetXaxis().SetTitle("title2")
    for i, x in enumerate(["one", "two", "three", "four", "five"]):
        before.GetXaxis().SetBinLabel(i + 1, x)
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)

    before = cls(rootname(), "title", 5, -5.0, 5.0)
    before.GetXaxis().SetTitle("title2")
    for i, x in enumerate(["one", "two", "three", "four", "five"]):
        before.GetXaxis().SetBinLabel(i + 1, x)
    for x in data:
        before.Fill(x)
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)

    edges = numpy.array([-5.0, -3.0, 0.0, 5.0, 10.0, 100.0], dtype=numpy.float64)
    before = cls(rootname(), "title", 5, edges)
    before.GetXaxis().SetTitle("title2")
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)

    edges = numpy.array([-5.0, -3.0, 0.0, 5.0, 10.0, 100.0], dtype=numpy.float64)
    before = cls(rootname(), "title", 5, edges)
    before.GetXaxis().SetTitle("title2")
    for x in data:
        before.Fill(x)
    after = aghast.to_root(aghast.from_root(before), rootname())
    check1d(before, after)