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)