added test for std info extraction and neutrino

orcasong/mc_info_extr.py

@@ -130,7 +130,7 @@ def get_std_reco(blob):
 		'std_pos_x':std_pos_x,'std_pos_y':std_pos_y,'std_pos_z':std_pos_z,
 		'std_energy':std_energy,'std_lik_energy':lik_energy,'std_length':std_length,
 		}
-
+		
 	return std_reco_info
 	
 	
@@ -339,7 +339,7 @@ def get_neutrino_mc_info_extr(input_file):
 				 'weight_w1': weight_w1,'weight_w2': weight_w2,'weight_w3': weight_w3,
 				 'n_gen':n_gen,
 				}				 
-
+		print(is_cc)
  		#get all the std reco info		
 		if has_std_reco:
 	

tests/test_core.py

@@ -14,7 +14,8 @@ __author__ = 'Stefan Reck'
 test_dir = os.path.dirname(os.path.realpath(__file__))
 MUPAGE_FILE = os.path.join(test_dir, "data", "mupage.root.h5")
 DET_FILE = os.path.join(test_dir, "data", "KM3NeT_-00000001_20171212.detx")
-
+NEUTRINO_FILE = os.path.join(test_dir, "data", "neutrino_file.h5")
+DET_FILE_NEUTRINO = os.path.join(test_dir, "data", "neutrino_detector_file.detx")
 
 class TestFileBinner(TestCase):
     """ Assert that the filebinner still produces the same output. """
@@ -140,3 +141,171 @@ class TestFileGraph(TestCase):
         }
         for k, v in target.items():
             np.testing.assert_equal(y[k], v)
+
+class TestStdRecoExtractor(TestCase):
+    """ Assert that the neutrino info is extracted correctly File has 18 events. """
+    @classmethod
+    def setUpClass(cls):
+        cls.proc = orcasong.core.FileGraph(
+            max_n_hits=3,
+            time_window=[0, 50],
+            hit_infos=["pos_z", "time", "channel_id"],
+            mc_info_extr=orcasong.mc_info_extr.get_neutrino_mc_info_extr(NEUTRINO_FILE),
+            det_file=DET_FILE_NEUTRINO,
+            add_t0=True,
+        )
+        cls.tmpdir = tempfile.TemporaryDirectory()
+        cls.outfile = os.path.join(cls.tmpdir.name, "binned.h5")
+        cls.proc.run(infile=NEUTRINO_FILE, outfile=cls.outfile)
+        cls.f = h5py.File(cls.outfile, "r")
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.f.close()
+        cls.tmpdir.cleanup()
+
+    def test_keys(self):
+        self.assertSetEqual(set(self.f.keys()), {
+            '_i_event_info', '_i_group_info', '_i_y',
+            'event_info', 'group_info', 'x', 'x_indices', 'y'})
+
+    def test_y(self):
+        y = self.f["y"][()]
+        target = {
+            'weight_w2': np.array([29650.0,
+									297100.0,
+									41450.0,
+									371400.0,
+									1101000000.0,
+									2757000.0,
+									15280000.0,
+									262800000.0,
+									22590.0,
+									24240.0,
+									80030.0,
+									3018000.0,
+									120600.0,
+									872200.0,
+									50440000.0,
+									21540.0,
+									42170.0,
+									25230.0]),
+									
+            'n_gen': np.array([60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0,
+								60000.0]),
+								
+            'dir_z': np.array([-0.896549,
+								-0.835252,
+								0.300461,
+								0.108997,
+								0.128445,
+								-0.543621,
+								-0.23205,
+								-0.297228,
+								0.694932,
+								0.73835,
+								-0.007682,
+								0.437847,
+								-0.126804,
+								0.153432,
+								-0.263229,
+								0.820217,
+								0.452473,
+								0.294217]),
+            
+            'is_cc': np.array([2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0,
+								2.0]),
+								
+            'std_dir_z': np.array([-0.923199825369434,
+									-0.6422689266782661,
+									0.38853917922036363,
+									-0.16690804339142448,
+									-0.01584853496341109,
+									-0.10151549881670698,
+									-0.0409694104272829,
+									-0.32964369874021787,
+									-0.3294926806601529,
+									0.6524241250799204,
+									-0.3899574246450216,
+									0.27872277417339086,
+									0.0019490791409933206,
+									0.20341370281708737,
+									-0.15739475718286297,
+									0.8040250543935723,
+									0.08772622550043882,
+									-0.7766722433951796]),
+									
+            'std_energy': np.array([4.7187625606210775,
+									4.169818842606011,
+									1.0056373761749966,
+									5.908597073055873,
+									12.409377607517195,
+									7.566695371401163,
+									1.3546775620239864,
+									2.659528737837978,
+									1.0056373761749966,
+									2.1968321463948755,
+									1.4821714294894754,
+									10.135831333340658,
+									2.6003934443336765,
+									1.4492149732348223,
+									71.69167874147956,
+									8.094744120333358,
+									3.148088080484504,
+									1.0056373761749966]),
+            
+        }
+        for k, v in target.items():
+            np.testing.assert_equal(y[k], v)
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
\ No newline at end of file