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