def setUp(self): self.input_blob_1 = { "Hits": kp.Table({ "x": [4, 5, 6], "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], }), "EventInfo": kp.Table({"pad": 1.0}), }
def setUp(self): self.input_blob_1 = { "Hits": kp.Table({ "x": [4, 5, 6], 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], }), "EventInfo": kp.Table({"pad": 1.}) }
def process(self, blob): if self.kind == "offline": n = blob["event"].n_hits if n == 0: return blob hits = blob["event"].hits hits_data = { "channel_id": hits.channel_id, "dom_id": hits.dom_id, "time": hits.t, "tot": hits.tot, "triggered": hits.trig, } if self.with_calibration: hits_data["pos_x"] = hits.pos_x hits_data["pos_y"] = hits.pos_y hits_data["pos_z"] = hits.pos_z hits_data["dir_x"] = hits.dir_x hits_data["dir_y"] = hits.dir_y hits_data["dir_z"] = hits.dir_z hits_data["tdc"] = hits.tdc blob["Hits"] = kp.Table( hits_data, h5loc="/hits", split_h5=self.split, name="Hits", ) if self.kind == "mc": n = blob["event"].n_mc_hits if n == 0: return blob mc_hits = blob["event"].mc_hits blob["McHits"] = kp.Table( { "a": mc_hits.a, "origin": mc_hits.origin, "pmt_id": mc_hits.pmt_id, "time": mc_hits.t, }, h5loc="/mc_hits", split_h5=self.split, name="McHits", ) if self.kind == "online": raise NotImplementedError( "The extraction of online (DAQ) hits is not implemented yet." ) return blob
def _parse_eventinfo(self, event, sim_program): wgt1, wgt2, wgt3, wgt4 = self._parse_wgts(event.w) tab_data = { "event_id": event.id, "run_id": event.run_id, "weight_w1": wgt1, "weight_w2": wgt2, "weight_w3": wgt3, "weight_w4": wgt4, "timestamp": event.t_sec, "nanoseconds": event.t_ns, "mc_time": event.mc_t, "trigger_mask": event.trigger_mask, "trigger_counter": event.trigger_counter, "overlays": event.overlays, "det_id": event.det_id, "frame_index": event.frame_index, "mc_run_id": event.mc_run_id, } if sim_program != None: # unfold the info in the w2list w2list_dict = self._unfold_w2list(event.w2list, sim_program) tab_data.update(w2list_dict) info = kp.Table(tab_data, h5loc="/event_info", name="EventInfo") return info
def test_mc_info_maker_dtype(self): """ Test the mcinfo maker on some dummy data. """ def extractor(blob): hits = blob["Hits"] return {"dom_id_0": hits.dom_id[0], "time_2": hits.time[2]} in_blob = { "Hits": kp.Table({ 'dom_id': np.array([2, 3, 3], dtype="int8"), 'time': np.array([10.1, 11.2, 12.3], dtype="float32"), }) } module = modules.McInfoMaker(extractor=extractor, store_as="test", to_float64=False) out_blob = module.process(in_blob) np.testing.assert_array_equal(out_blob["test"]["dom_id_0"], np.array([ 2, ], dtype="int8")) np.testing.assert_array_equal(out_blob["test"]["time_2"], np.array([ 12.3, ], dtype="float32"))
def test_mc_info_maker(self): """ Test the mcinfo maker on some dummy data. """ def extractor(blob): hits = blob["Hits"] return {"dom_id_0": hits.dom_id[0], "time_2": hits.time[2]} in_blob = { "Hits": kp.Table({ 'dom_id': [2, 3, 3], 'channel_id': [0, 1, 2], 'time': [10.1, 11.2, 12.3] }) } module = modules.McInfoMaker(extractor=extractor, store_as="test") out_blob = module.process(in_blob) self.assertSequenceEqual(list(out_blob.keys()), ["Hits", "test"]) self.assertSequenceEqual(list(out_blob["test"].dtype.names), ('dom_id_0', 'time_2')) np.testing.assert_array_equal(out_blob["test"]["dom_id_0"], np.array([ 2, ], dtype="float64")) np.testing.assert_array_equal(out_blob["test"]["time_2"], np.array([ 12.3, ], dtype="float64"))
def process(self, blob): n = np.random.randint(2, 20) tab = kp.Table( { "lambda": np.random.rand(n), "energy": np.random.rand(n) * 1000 }, h5loc="/reco", name="Quality Parameter", ) event_info = kp.Table({"run_id": np.random.randint(10000, 20000)}, h5loc="/event_info") blob["Reco"] = tab blob["EventInfo"] = event_info blob["Header"] = self.header return blob
def process(self, blob): tracks = blob["TrackIns"] muons = tracks[tracks.type == 5] muon = kp.Table(muons[np.argmax(muons.energy)]) hits = blob["CalibHits"] dist = kp.math.pld3(hits.pos, muon.pos, muon.dir) self.distances.append(dist) return blob
def test_2d_xt_binning(self): # (3 x 2) x-t binning bin_edges_list = [["x", [3.5, 4.5, 5.5, 6.5]], ["time", [0.5, 2, 3.5]]] module = modules.ImageMaker(bin_edges_list=bin_edges_list) in_blob = { "Hits": kp.Table({ "x": [4, 5, 6], "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } target = { "Hits": kp.Table({ "x": [4, 5, 6], "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), "samples": np.array([[ [1, 0], [0, 1], [0, 1], ]]), } out_blob = module.process(in_blob) self.assertSetEqual(set(out_blob.keys()), set(target.keys())) np.testing.assert_array_almost_equal( np.array(out_blob["Hits"].view("<f8")), np.array(target["Hits"].view("<f8"))) np.testing.assert_array_almost_equal(np.array(out_blob["samples"]), np.array(target["samples"]))
def test_time_preproc_mchits_t0_and_center(self): module = modules.TimePreproc(add_t0=True, center_time=True) target = { "Hits": kp.Table({ 'time': [-1.1, 0., 1.1], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), "McHits": kp.Table({ 'time': [-1.2, -0.2, 0.8], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), } out_blob = module.process(self.in_blob_mc) self.assertSetEqual(set(out_blob.keys()), set(target.keys())) np.testing.assert_array_almost_equal( np.array(out_blob["McHits"].view("<f8")), np.array(target["McHits"].view("<f8")))
def setUp(self): self.in_blob = { "Hits": kp.Table({ "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } self.in_blob_mc = { "Hits": kp.Table({ "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), "McHits": kp.Table({ "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), }
def setUp(self): self.in_blob = { "Hits": kp.Table({ 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } self.in_blob_mc = { "Hits": kp.Table({ 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), "McHits": kp.Table({ 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) }
def test_1d_binning_no_hits(self): # (1, ) t binning bin_edges_list = [["time", [3.5, 4.5]]] module = modules.ImageMaker(bin_edges_list=bin_edges_list) in_blob = { "Hits": kp.Table({ 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } target = { "Hits": kp.Table({ 'time': [1., 2., 3.], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }), "samples": np.array([ [ 0, ], ]) } out_blob = module.process(in_blob) self.assertSetEqual(set(out_blob.keys()), set(target.keys())) np.testing.assert_array_almost_equal( np.array(out_blob["Hits"].view("<f8")), np.array(target["Hits"].view("<f8"))) np.testing.assert_array_almost_equal(np.array(out_blob["samples"]), np.array(target["samples"]))
def test_input_blob_1_max_n_hits(self): input_blob_long = { "Hits": kp.Table({ "x": np.random.rand(1000).astype("float32"), }) } result = modules.PointMaker( max_n_hits=10, hit_infos=("x", ), time_window=None, dset_n_hits=None, ).process(input_blob_long)["samples"] self.assertSequenceEqual(result.shape, (1, 10, 2)) self.assertTrue( all(np.isin(result[0, :, 0], input_blob_long["Hits"]["x"])))
def _parse_mc_tracks(self, mc_tracks): dct = { "dir_x": mc_tracks.dir_x, "dir_y": mc_tracks.dir_y, "dir_z": mc_tracks.dir_z, "pos_x": mc_tracks.pos_x, "pos_y": mc_tracks.pos_y, "pos_z": mc_tracks.pos_z, "energy": mc_tracks.E, "time": mc_tracks.t, "pdgid": mc_tracks.pdgid, "id": mc_tracks.id, "length": mc_tracks.len, } if self._read_usr_data: dct.update(self._parse_usr_to_dct(mc_tracks)) return kp.Table(dct, name="McTracks", h5loc="/mc_tracks", split_h5=self.split)
def test_it(self): # (3 x 2) x-t binning bin_edges_list = [ ["x", [3.5, 4.5, 5.5, 6.5]], ["time", [0.5, 2, 3.5]], ["z", [1, 4]], ] in_blob = { "Hits": kp.Table({ "x": [4, 5, 6, 6], "time": [1.0, 2.0, 3.0, 50], "z": [0, 3, 4, 5], "t0": [0.1, 0.2, 0.3, 0.4], "triggered": [0, 1, 1, 1], }) } target = { "x": { "hist": np.array([0.0, 0.0, 0.0, 1.0, 0.0, 1.0]), "hist_bin_edges": np.array([3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5]), "bin_edges": [3.5, 4.5, 5.5, 6.5], "cut_off": np.array([0.0, 0.0]), }, "time": { "hist": np.array([0.0, 2.0]), "hist_bin_edges": [0.5, 2, 3.5], "bin_edges": [0.5, 2, 3.5], "cut_off": np.array([0.0, 1.0]), }, "z": { "hist": np.array([0.0, 2.0]), "hist_bin_edges": np.array([1.0, 2.5, 4.0]), "bin_edges": [1, 4], "cut_off": np.array([1.0, 1.0]), }, } module = modules.BinningStatsMaker(bin_edges_list=bin_edges_list, res_increase=2) module.process(in_blob) output = module.finish() check_dicts_n_ray(output, target)
def test_time_preproc_center(self): module = modules.TimePreproc(add_t0=False, center_time=True) target = { "Hits": kp.Table({ "time": [-1.0, 0.0, 1.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } out_blob = module.process(self.in_blob) self.assertSetEqual(set(out_blob.keys()), set(target.keys())) np.testing.assert_array_equal(np.array(out_blob["Hits"]), np.array(target["Hits"]))
def test_unknown_field(self): # (3 x 2) x-t binning bin_edges_list = [["aggg", [3.5, 4.5, 5.5, 6.5]], ["time", [0.5, 2, 3.5]]] module = modules.ImageMaker(bin_edges_list=bin_edges_list) in_blob = { "Hits": kp.Table({ "x": [4, 5, 6], "time": [1.0, 2.0, 3.0], "t0": [0.1, 0.2, 0.3], "triggered": [0, 1, 1], }) } with self.assertRaises(ValueError): module.process(in_blob)
def test_it(self): # (3 x 2) x-t binning bin_edges_list = [["x", [3.5, 4.5, 5.5, 6.5]], ["time", [0.5, 2, 3.5]], ["z", [1, 4]]] in_blob = { "Hits": kp.Table({ "x": [4, 5, 6, 6], 'time': [1., 2., 3., 50], "z": [0, 3, 4, 5], "t0": [0.1, 0.2, 0.3, 0.4], "triggered": [0, 1, 1, 1], }) } target = { 'x': { 'hist': np.array([0., 0., 0., 1., 0., 1.]), 'hist_bin_edges': np.array([3.5, 4., 4.5, 5., 5.5, 6., 6.5]), 'bin_edges': [3.5, 4.5, 5.5, 6.5], 'cut_off': np.array([0., 0.]) }, 'time': { 'hist': np.array([0., 2.]), 'hist_bin_edges': [0.5, 2, 3.5], 'bin_edges': [0.5, 2, 3.5], 'cut_off': np.array([0., 1.]) }, 'z': { 'hist': np.array([0., 2.]), 'hist_bin_edges': np.array([1., 2.5, 4.]), 'bin_edges': [1, 4], 'cut_off': np.array([1., 1.]) } } module = modules.BinningStatsMaker(bin_edges_list=bin_edges_list, res_increase=2) module.process(in_blob) output = module.finish() check_dicts_n_ray(output, target)
def process(self, blob): table = kp.Table({"x": np.random.randn(self.n)}, h5loc=self.h5loc) blob["RandomNumbers"] = table return blob
def configure(self): self.header = kp.Table({"is_mc": True}, h5singleton=True, h5loc="/header")
def _put_tracks_into_blob(self, blob, tracks, reco_identifier, n_tracks): """ Put a certain type of "tracks" in the blob and give specific name. Parameters ---------- tracks : awkward array The tracks object to be put in the blob eventually. Can be only best tracks. identifier : string A string to name the kp table. n_tracks : int The number of tracks from before. Use to distinguish between best and all tracks. """ reco_tracks = dict( pos_x=tracks.pos_x, pos_y=tracks.pos_y, pos_z=tracks.pos_z, dir_x=tracks.dir_x, dir_y=tracks.dir_y, dir_z=tracks.dir_z, E=tracks.E, rec_type=tracks.rec_type, t=tracks.t, likelihood=tracks.lik, length=tracks.len, # do all recos have this? ) if n_tracks != 1: reco_tracks.update( id=tracks.id, idx=np.arange(n_tracks), ) n_columns = max(km3io.definitions.fitparameters.values()) + 1 fitinf_array = np.ma.filled( ak.to_numpy(ak.pad_none(tracks.fitinf, target=n_columns, axis=-1)), fill_value=np.nan, ).astype("float32") fitinf_split = np.split(fitinf_array, fitinf_array.shape[-1], axis=-1) if n_tracks == 1: for fitparam, idx in km3io.definitions.fitparameters.items(): reco_tracks[fitparam] = fitinf_split[idx][0] else: for fitparam, idx in km3io.definitions.fitparameters.items(): reco_tracks[fitparam] = fitinf_split[idx][:, 0] blob["Reco_" + reco_identifier] = kp.Table( reco_tracks, h5loc=f"/reco/" + reco_identifier, name="Reco " + reco_identifier, split_h5=self.split, ) # write out the rec stages only once with all tracks if n_tracks != 1: _rec_stage = np.array(ak.flatten(tracks.rec_stages)._layout) _counts = ak.count(tracks.rec_stages, axis=1) _idx = np.repeat(np.arange(n_tracks), _counts) blob["RecStages"] = kp.Table( dict(rec_stage=_rec_stage, idx=_idx), # Just to save space, we specify smaller dtypes. # We assume there will be never more than 32767 # reco tracks for a single reconstruction type. dtypes=[("rec_stage", np.int16), ("idx", np.uint16)], h5loc=f"/reco/rec_stages", name="Reconstruction Stages", split_h5=self.split, )