import unittest
import pandas as pd
import nnfwtbn
[docs]class HistTestBase(unittest.TestCase):
"""
Test the implementation of hist().
The implementation is tested by inspecting the returned uhepp objects.
"""
[docs] def setUp(self):
"""Set up a toy dataframe and processes"""
self.data = pd.DataFrame({
"m": [10, 10, 10, 10, 10],
"w": [ 1, 2, 1, 3, 8],
"p": [ 1, 1, 1, 2, 3]
})
self.process_a = nnfwtbn.Process("a", lambda d: d.p == 1)
self.process_b = nnfwtbn.Process("b", lambda d: d.p == 2)
self.mc_stack = nnfwtbn.McStack(self.process_a, self.process_b)
self.process_x = nnfwtbn.Process("x", lambda d: d.p == 3)
self.data_stack = nnfwtbn.McStack(self.process_x)
[docs] def test_yield_base(self):
"""Check the bin contents"""
hist = nnfwtbn.hist(self.data, variable="m", weight="w", bins=[0, 20],
stacks=[self.mc_stack, self.data_stack],
return_uhepp=True)
self.assertEqual(hist.yields["a__s0_p0"].base, [0, 4, 0])
self.assertEqual(hist.yields["b__s0_p1"].base, [0, 3, 0])
self.assertEqual(hist.yields["x__s1_p0"].base, [0, 8, 0])
[docs] def test_yiele_stat(self):
"""Check the statistical uncertainties"""
hist = nnfwtbn.hist(self.data, variable="m", weight="w", bins=[0, 20],
stacks=[self.mc_stack, self.data_stack],
return_uhepp=True)
self.assertAlmostEqual(hist.yields["a__s0_p0"].stat[1]**2, 6)
self.assertAlmostEqual(hist.yields["b__s0_p1"].stat[1]**2, 9)
self.assertAlmostEqual(hist.yields["x__s1_p0"].stat[1]**2, 64)