def test_append_columns(self):
tab = Table(self.arr)
print(tab)
with pytest.raises(ValueError):
tab = tab.append_columns('new', [1, 2, 3, 4])
tab = tab.append_columns('new', [1, 2, 3])
print(tab)
assert tab.new[0] == 1
assert tab.new[-1] == 3
tab = tab.append_columns('bar', 0)
print(tab)
assert tab.bar[0] == 0
assert tab.bar[-1] == 0
tab = tab.append_columns('lala', [1])
print(tab)
assert tab.lala[0] == 1
assert tab.lala[-1] == 1
with pytest.raises(ValueError):
tab = tab.append_columns(['m', 'n'], [1, 2])
with pytest.raises(ValueError):
tab = tab.append_columns(['m', 'n'], [[1], [2]])
tab = tab.append_columns(['m', 'n'], [[1, 1, 2], [2, 4, 5]])
print(tab)
assert tab.m[0] == 1
assert tab.m[-1] == 2
assert tab.n[0] == 2
assert tab.n[-1] == 5
def test_append_columns(self):
tab = Table(self.arr)
print(tab)
with pytest.raises(ValueError):
tab = tab.append_columns("new", [1, 2, 3, 4])
tab = tab.append_columns("new", [1, 2, 3])
print(tab)
assert tab.new[0] == 1
assert tab.new[-1] == 3
tab = tab.append_columns("bar", 0)
print(tab)
assert tab.bar[0] == 0
assert tab.bar[-1] == 0
tab = tab.append_columns("lala", [1])
print(tab)
assert tab.lala[0] == 1
assert tab.lala[-1] == 1
with pytest.raises(ValueError):
tab = tab.append_columns(["m", "n"], [1, 2])
with pytest.raises(ValueError):
tab = tab.append_columns(["m", "n"], [[1], [2]])
tab = tab.append_columns(["m", "n"], [[1, 1, 2], [2, 4, 5]])
print(tab)
assert tab.m[0] == 1
assert tab.m[-1] == 2
assert tab.n[0] == 2
assert tab.n[-1] == 5
def test_append_columns_which_is_too_long(self):
tab = Table({'a': [1, 2, 3]})
with pytest.raises(ValueError):
tab.append_columns('b', values=[4, 5, 6, 7])
def test_append_columns_duplicate(self):
tab = Table({'a': 1})
with pytest.raises(ValueError):
tab = tab.append_columns(['a'], np.array([[2]]))
def test_append_columns_with_mismatching_lengths_raises(self):
tab = Table({'a': [1, 2, 3]})
with pytest.raises(ValueError):
tab.append_columns(colnames=['b', 'c'], values=[[4, 5, 6], [7, 8]])
def test_append_columns_modifies_dtype(self):
tab = Table({'a': [1, 2]})
tab = tab.append_columns('group_id', [0, 1])
assert 'group_id' in tab.dtype.names
def test_append_column_which_is_too_short_raises(self):
tab = Table({'a': [1, 2, 3]})
with pytest.raises(ValueError):
tab = tab.append_columns('b', [4, 5])
def test_append_columns_with_single_value(self):
tab = Table({'a': 1})
tab = tab.append_columns('group_id', 0)
assert 0 == tab.group_id[0]
def test_append_columns_with_multiple_values(self):
tab = Table({'a': [1, 2]})
tab = tab.append_columns('group_id', [0, 1])
assert 0 == tab.group_id[0]
assert 1 == tab.group_id[1]
def extract_event(self):
blob = Blob()
r = self.event_reader
r.retrieve_next_event() # do it at the beginning!
n = r.n_fits
if n > self.buf_size:
self._resize_buffers(int(n * 3 / 2))
r.get_fits(
self._pos_xs,
self._pos_ys,
self._pos_zs,
self._dir_xs,
self._dir_ys,
self._dir_zs,
self._ndfs,
self._times,
self._qualities,
self._energies,
)
fit_collection = Table({
'pos_x': self._pos_xs[:n],
'pos_y': self._pos_ys[:n],
'pos_z': self._pos_zs[:n],
'dir_x': self._dir_xs[:n],
'dir_y': self._dir_ys[:n],
'dir_z': self._dir_zs[:n],
'ndf': self._ndfs[:n],
'time': self._times[:n],
'quality': self._qualities[:n],
'energy': self._energies[:n],
},
h5loc='/jfit')
fit_collection = fit_collection.append_columns(['event_id'],
[self.event_index])
# TODO make this into a datastructure
event_info = Table.from_template({
'det_id': 0,
'frame_index': 0,
'livetime_sec': 0,
'MC ID': 0,
'MC time': 0,
'n_events_gen': 0,
'n_files_gen': 0,
'overlays': 0,
'trigger_counter': 0,
'trigger_mask': 0,
'utc_nanoseconds': 0,
'utc_seconds': 0,
'weight_w1': np.nan,
'weight_w2': np.nan,
'weight_w3': np.nan,
'run_id': 0,
'group_id': self.event_index,
}, 'EventInfo')
self.event_index += 1
blob['EventInfo'] = event_info
blob['JFit'] = fit_collection
return blob
def test_append__single_column(self):
tab = Table({"a": 1})
print(tab.dtype)
tab = tab.append_columns(["b"], np.array([[2]]))
print(tab.dtype)
print(tab.b)
def test_append_columns_with_mismatching_lengths_raises(self):
tab = Table({"a": [1, 2, 3]})
with pytest.raises(ValueError):
tab.append_columns(colnames=["b", "c"], values=[[4, 5, 6], [7, 8]])
def test_append_columns_which_is_too_long(self):
tab = Table({"a": [1, 2, 3]})
with pytest.raises(ValueError):
tab.append_columns("b", values=[4, 5, 6, 7])
def test_append_columns_duplicate(self):
tab = Table({"a": 1})
with pytest.raises(ValueError):
tab = tab.append_columns(["a"], np.array([[2]]))
def test_append_column_which_is_too_short_raises(self):
tab = Table({"a": [1, 2, 3]})
with pytest.raises(ValueError):
tab = tab.append_columns("b", [4, 5])
def test_append_columns_modifies_dtype(self):
tab = Table({"a": [1, 2]})
tab = tab.append_columns("group_id", [0, 1])
assert "group_id" in tab.dtype.names
def test_append_columns_with_multiple_values(self):
tab = Table({"a": [1, 2]})
tab = tab.append_columns("group_id", [0, 1])
assert 0 == tab.group_id[0]
assert 1 == tab.group_id[1]
def test_append_columns_with_single_value(self):
tab = Table({"a": 1})
tab = tab.append_columns("group_id", 0)
assert 0 == tab.group_id[0]
def apply(self,
hits,
no_copy=False,
correct_slewing=True,
slewing_variant=3):
"""Add x, y, z, t0 (and du, floor if DataFrame) columns to the hits."""
if not no_copy:
try:
hits = hits.copy()
except AttributeError: # probably a km3io object
pass
if isinstance(hits, (ak.Array, ak.Record, km3io.rootio.Branch)):
if hasattr(hits, "dom_id"):
hits = Table(
dict(
dom_id=hits.dom_id,
channel_id=hits.channel_id,
time=hits.t,
tot=hits.tot,
triggered=hits.trig,
))
else: # mc_hits in km3io
hits = Table(
dict(
pmt_id=hits.pmt_id,
time=hits.t,
a=hits.a,
# TODO: Not all MC files have these two fields
# pure_a=hits.pure_a,
# pure_t=hits.pure_t,
origin=hits.origin,
))
if istype(hits, "DataFrame"):
# do we ever see McHits here?
hits = Table.from_template(hits, "Hits")
is_mc = None
if hasattr(hits, "dom_id") and hasattr(hits, "channel_id"):
try:
(
dir_x,
dir_y,
dir_z,
du,
floor,
pos_x,
pos_y,
pos_z,
t0,
pmt_id,
) = _get_calibration_for_hits(hits,
self._calib_by_dom_and_channel)
except KeyError as e:
self.log.critical("Wrong calibration (DETX) data provided.")
raise
is_mc = False
elif hasattr(hits, "pmt_id"):
try:
(
dir_x,
dir_y,
dir_z,
du,
floor,
pos_x,
pos_y,
pos_z,
t0,
dom_id,
channel_id,
) = _get_calibration_for_mchits(hits, self._calib_by_pmt_id)
except KeyError as e:
self.log.critical("Wrong calibration (DETX) data provided.")
raise
is_mc = True
else:
raise TypeError("Don't know how to apply calibration to '{0}'. "
"We need at least 'dom_id' and 'channel_id', or "
"'pmt_id'.".format(hits.name))
if hasattr(hits, "time") and not is_mc:
if hits.time.dtype != t0.dtype:
time = hits.time.astype("f4") + t0.astype("f4")
hits = hits.drop_columns(["time"])
hits = hits.append_columns(["time"], [time])
else:
hits.time += t0
hits_data = {}
for colname in hits.dtype.names:
hits_data[colname] = hits[colname]
calib = {
"dir_x": dir_x,
"dir_y": dir_y,
"dir_z": dir_z,
"du": du.astype(np.uint8),
"floor": floor.astype(np.uint8),
"pos_x": pos_x,
"pos_y": pos_y,
"pos_z": pos_z,
"t0": t0,
}
if is_mc:
calib["dom_id"] = dom_id.astype(np.int32)
calib["channel_id"] = channel_id.astype(np.int32)
else:
calib["pmt_id"] = pmt_id.astype(np.int32)
hits_data.update(calib)
if correct_slewing and not is_mc:
hits_data["time"] -= slew(hits_data["tot"],
variant=slewing_variant)
return Table(hits_data,
h5loc=hits.h5loc,
split_h5=hits.split_h5,
name=hits.name)
def test_append__single_column(self):
tab = Table({'a': 1})
print(tab.dtype)
tab = tab.append_columns(['b'], np.array([[2]]))
print(tab.dtype)
print(tab.b)
