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,
)
