Python dataclasses.Table类代码示例

 def test_template(self):
     n = 10
     channel_ids = np.arange(n)
     dom_ids = np.arange(n)
     times = np.arange(n)
     tots = np.arange(n)
     triggereds = np.ones(n)
     d_hits = {
         'channel_id': channel_ids,
         'dom_id': dom_ids,
         'time': times,
         'tot': tots,
         'triggered': triggereds,
         'group_id': 0,    # event_id
     }
     tab = Table.from_template(d_hits, 'Hits')
     assert tab.name == 'Hits'
     assert tab.split_h5 is True
     assert isinstance(tab, Table)
     ar_hits = {
         'channel_id': np.ones(n, dtype=int),
         'dom_id': np.ones(n, dtype=int),
         'time': np.ones(n, dtype=float),
         'tot': np.ones(n, dtype=float),
         'triggered': np.ones(n, dtype=bool),
         'group_id': np.ones(n, dtype=int),
     }
     tab = Table.from_template(ar_hits, 'Hits')
     assert tab.name == 'Hits'
     assert tab.split_h5 is True
     assert isinstance(tab, Table)

 def test_fromcolumns(self):
     n = 5
     dlist = [
         np.ones(n, dtype=int),
         np.zeros(n, dtype=float),
         0,
     ]
     dt = np.dtype([('a', float), ('b', float), ('c', float)])
     with pytest.raises(ValueError):
         tab = Table(dlist, dtype=dt)
     tab = Table.from_columns(dlist, dtype=dt)
     print(tab.dtype)
     print(tab.shape)
     print(tab)
     assert tab.h5loc == DEFAULT_H5LOC
     assert isinstance(tab, Table)
     tab = Table.from_columns(dlist, dtype=dt, h5loc='/foo')
     print(tab.dtype)
     print(tab.shape)
     print(tab)
     assert tab.h5loc == '/foo'
     assert isinstance(tab, Table)
     bad_dt = [('a', float), ('b', float), ('c', float), ('d', int)]
     with pytest.raises(ValueError):
         tab = Table.from_columns(dlist, dtype=bad_dt)
         print(tab.dtype)
         print(tab.shape)
         print(tab)

 def test_incomplete_template(self):
     n = 10
     channel_ids = np.arange(n)
     dom_ids = np.arange(n)
     # times = np.arange(n)
     tots = np.arange(n)
     triggereds = np.ones(n)
     d_hits = {
         'channel_id': channel_ids,
         'dom_id': dom_ids,
     # 'time': times,
         'tot': tots,
         'triggered': triggereds,
         'group_id': 0,    # event_id
     }
     with pytest.raises(KeyError):
         tab = Table.from_template(d_hits, 'Hits')
         assert tab is not None
     ar_hits = {
         'channel_id': np.ones(n, dtype=int),
         'dom_id': np.ones(n, dtype=int),
     # 'time': np.ones(n, dtype=float),
         'tot': np.ones(n, dtype=float),
         'triggered': np.ones(n, dtype=bool),
         'group_id': np.ones(n, dtype=int),
     }
     with pytest.raises(KeyError):
         tab = Table.from_template(ar_hits, 'Hits')
         assert tab is not None

 def test_fromrows(self):
     dlist = [
         [1, 2, 3],
         [4, 5, 6],
     ]
     dt = np.dtype([('a', float), ('b', float), ('c', float)])
     with pytest.raises(ValueError):
         tab = Table(dlist, dtype=dt)
     tab = Table.from_rows(dlist, dtype=dt)
     print(tab.dtype)
     print(tab.shape)
     print(tab)
     assert tab.h5loc == DEFAULT_H5LOC
     assert isinstance(tab, Table)
     tab = Table.from_rows(dlist, dtype=dt, h5loc='/foo')
     print(tab.dtype)
     print(tab.shape)
     print(tab)
     assert tab.h5loc == '/foo'
     assert isinstance(tab, Table)
     bad_dt = [('a', float), ('b', float), ('c', float), ('d', int)]
     with pytest.raises(ValueError):
         tab = Table.from_rows(dlist, dtype=bad_dt)
         print(tab.dtype)
         print(tab.shape)
         print(tab)

    def test_apply_without_affecting_primary_hit_table(self):
        calib = Calibration(filename=DETX_FILENAME)
        hits = Table({'pmt_id': [1, 2, 1], 'time': [10.1, 11.2, 12.3]})
        hits_compare = hits.copy()
        calib.apply(hits)

        for t_primary, t_calib in zip(hits_compare, hits):
            self.assertAlmostEqual(t_primary, t_calib)

 def test_crash_repr(self):
     a = np.array('', dtype=[('a', '<U1')])
     with pytest.raises(TypeError):
         print(len(a))
     tab = Table(a)
     s = tab.__str__()
     assert s is not None
     r = tab.__repr__()
     assert r is not None

 def test_pos_setter(self):
     tab = Table({
         'pos_x': [1, 0, 0],
         'pos_y': [0, 1, 0],
         'pos_z': [0, 0, 1]
     })
     new_pos = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
     tab.pos = new_pos
     assert np.allclose(new_pos, tab.pos)

 def test_dir_setter(self):
     tab = Table({
         'dir_x': [1, 0, 0],
         'dir_y': [0, 1, 0],
         'dir_z': [0, 0, 1]
     })
     new_dir = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
     tab.dir = new_dir
     assert np.allclose(new_dir, tab.dir)

    def process_event(self, data, blob):
        data_io = BytesIO(data)
        preamble = DAQPreamble(file_obj=data_io)    # noqa
        event = DAQEvent(file_obj=data_io)
        header = event.header

        hits = event.snapshot_hits
        n_hits = event.n_snapshot_hits
        if n_hits == 0:
            return
        dom_ids, channel_ids, times, tots = zip(*hits)
        triggereds = np.zeros(n_hits)
        triggered_map = {}
        for triggered_hit in event.triggered_hits:
            dom_id, pmt_id, time, tot, _ = triggered_hit
            triggered_map[(dom_id, pmt_id, time, tot)] = True
        for idx, hit in enumerate(hits):
            triggereds[idx] = hit in triggered_map

        hit_series = Table.from_template({
            'channel_id': channel_ids,
            'dom_id': dom_ids,
            'time': times,
            'tot': tots,
            'triggered': triggereds,
            'group_id': self.event_id,
        }, 'Hits')

        blob['Hits'] = hit_series

        event_info = Table.from_template(
            {
                'det_id': header.det_id,
        # 'frame_index': self.index,  # header.time_slice,
                'frame_index': header.time_slice,
                'livetime_sec': 0,
                'mc_id': 0,
                'mc_t': 0,
                'n_events_gen': 0,
                'n_files_gen': 0,
                'overlays': event.overlays,
                'trigger_counter': event.trigger_counter,
                'trigger_mask': event.trigger_mask,
                'utc_nanoseconds': header.ticks * 16,
                'utc_seconds': header.time_stamp,
                'weight_w1': 0,
                'weight_w2': 0,
                'weight_w3': 0,    # MC weights
                'run_id': header.run,    # run id
                'group_id': self.event_id,
            },
            'EventInfo'
        )
        blob['EventInfo'] = event_info

        self.event_id += 1
        self.index += 1

 def test_sort(self):
     dt = np.dtype([('a', int), ('b', float), ('c', int)])
     arr = np.array([
         (0, 1.0, 2),
         (3, 7.0, 5),
         (6, 4.0, 8),
     ], dtype=dt)
     tab = Table(arr)
     tab_sort = tab.sorted('b')
     assert_array_equal(tab_sort['a'], np.array([0, 6, 3]))

 def test_drop_columns(self):
     tab = Table({'a': 1, 'b': 2, 'c': 3})
     print(tab.dtype)
     tab = tab.drop_columns(['a', 'b'])
     print(tab.dtype)
     with pytest.raises(AttributeError):
         print(tab.a)
     with pytest.raises(AttributeError):
         print(tab.b)
     print(tab.c)

 def test_add_tables_with_same_colnames_but_different_dtype_order(self):
     cols1 = ('b', 'a')
     tab1 = Table.from_columns([[100, 200], [1, 2]], colnames=cols1)
     self.assertTupleEqual(cols1, tab1.dtype.names)
     cols2 = ('a', 'b')
     tab2 = Table.from_columns([[3, 4, 5], [300, 400, 500]], colnames=cols2)
     added_tab = tab1 + tab2
     self.assertListEqual([1, 2, 3, 4, 5], list(added_tab.a))
     self.assertListEqual([100, 200, 300, 400, 500], list(added_tab.b))
     self.assertListEqual(
         list(added_tab.dtype.names), list(tab1.dtype.names)
     )

 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 extract_event(self):
        blob = self._current_blob
        r = self.event_reader
        r.retrieve_next_event()    # do it at the beginning!

        n = r.number_of_snapshot_hits

        if n > self.buf_size:
            self._resize_buffers(int(n * 3 / 2))

        r.get_hits(
            self._channel_ids, self._dom_ids, self._times, self._tots,
            self._triggereds
        )

        hit_series = Table.from_template({
            'channel_id': self._channel_ids[:n],
            'dom_id': self._dom_ids[:n],
            'time': self._times[:n],
            'tot': self._tots[:n],
            'triggered': self._triggereds[:n],
            'group_id': self.event_index,
        }, 'Hits')

        event_info = Table.from_template({
            'det_id': r.det_id,
            'frame_index': r.frame_index,
            'livetime_sec': 0,
            'mc_id': 0,
            'mc_t': 0,
            'n_events_gen': 0,
            'n_files_gen': 0,
            'overlays': r.overlays,
            'trigger_counter': r.trigger_counter,
            'trigger_mask': r.trigger_mask,
            'utc_nanoseconds': r.utc_nanoseconds,
            'utc_seconds': r.utc_seconds,
            '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['Hits'] = hit_series
        return blob

 def test_from_dict_with_unordered_columns_wrt_to_dtype_fields(self):
     data = {'b': [1, 2], 'c': [3, 4], 'a': [5, 6]}
     dt = [('a', float), ('b', float), ('c', float)]
     tab = Table.from_dict(data, dtype=dt)
     assert np.allclose([1, 2], tab.b)
     assert np.allclose([3, 4], tab.c)
     assert np.allclose([5, 6], tab.a)