Belle II Software development
test_quantity_extract.py
1#!/usr/bin/env python3
2
3
10
11import unittest
12import ROOT
13
14from quantity_extract import RootQuantityExtract, default_extractor
15
16
17class TestQuantityExtract(unittest.TestCase):
18
19 """
20 Various test cases for the Quantity Extractor class
21 """
22
23 def test_h1d(self):
24 """
25 Test getting the quantities of a TH1d
26 """
27
28 h1 = ROOT.TH1D("ext_test1", "ext_test1", 40, 0, 40.0)
29 h1.Fill(10.0)
30 h1.Fill(20.0)
31 h1.Fill(30.0)
32
33 rext = RootQuantityExtract(default_extractor())
34 res = rext.extract(h1)
35
36 self.assertIn("mean_x", res)
37 self.assertAlmostEqual(res["mean_x"], 20)
38 self.assertIn("mean_y", res)
39 self.assertAlmostEqual(res["mean_y"], 0.075)
40 self.assertIn("mean_y_zero_suppressed", res)
41 self.assertAlmostEqual(res["mean_y_zero_suppressed"], 1.0)
42 self.assertIn("entries", res)
43
44 def test_profile(self):
45 """
46 Test getting the quantities of a TProfile
47 """
48
49 h1 = ROOT.TProfile("ext_test1", "ext_test1", 40, 0, 40.0)
50 h1.Fill(10.0, 4.0)
51 h1.Fill(20.0, 5.0)
52 h1.Fill(30.0, 6.0)
53
54 rext = RootQuantityExtract(default_extractor())
55 res = rext.extract(h1)
56
57 self.assertIn("mean_y", res)
58 self.assertAlmostEqual(res["mean_y"], 0.375)
59 self.assertIn("mean_y_zero_suppressed", res)
60 self.assertAlmostEqual(res["mean_y_zero_suppressed"], 5.0)
61
62 def test_ntuple(self):
63 """
64 Test getting the quantities contained in a TNtuple
65 """
66
67 tn = ROOT.TNtuple("particle_list", "particle_list", "x:y:z:energy")
68 tn.Fill(5, 6, 7, 10)
69
70 rext = RootQuantityExtract(default_extractor())
71 res = rext.extract(tn)
72 self.assertEqual(len(res), 4)
73 self.assertIn("particle_list_z", res)
74 self.assertAlmostEqual(res["particle_list_z"], 7.0)
75
76
77if __name__ == "__main__":
78 unittest.main()