import binascii
from collections import namedtuple
import os
import uproot
import uproot3
import numpy as np
import numba as nb
TIMESLICE_FRAME_BASKET_CACHE_SIZE = 523 * 1024 ** 2 # [byte]
SUMMARYSLICE_FRAME_BASKET_CACHE_SIZE = 523 * 1024 ** 2 # [byte]
BASKET_CACHE_SIZE = 110 * 1024 ** 2
BASKET_CACHE = uproot3.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
# Parameters for PMT rate conversions, since the rates in summary slices are
# stored as a single byte to save space. The values from 0-255 can be decoded
# using the `get_rate(value)` function, which will yield the actual rate
# in Hz.
MINIMAL_RATE_HZ = 2.0e3
MAXIMAL_RATE_HZ = 2.0e6
RATE_FACTOR = np.log(MAXIMAL_RATE_HZ / MINIMAL_RATE_HZ) / 255
CHANNEL_BITS_TEMPLATE = np.zeros(31, dtype=bool)
BranchConfiguration = namedtuple(
field_names=["branch_address", "interpretation"], typename="BranchConfiguration"
)
class SummarysliceReader:
"""
A reader for summaryslices which are loaded as chunks given by step_size.
To be used as an iterator (`for chunks in SummarysliceReader(...): ...`)
"""
TREE_ADDR = "KM3NET_SUMMARYSLICE/KM3NET_SUMMARYSLICE"
_subbranches = [
BranchConfiguration(
"KM3NETDAQ::JDAQSummarysliceHeader",
uproot.interpretation.numerical.AsDtype(
[
(" cnt", "u4"),
(" vers", "u2"),
(" cnt2", "u4"),
(" vers2", "u2"),
(" cnt3", "u4"),
(" vers3", "u2"),
("detector_id", ">i4"),
("run", ">i4"),
("frame_index", ">i4"),
(" cnt4", "u4"),
(" vers4", "u2"),
("UTC_seconds", ">u4"),
("UTC_16nanosecondcycles", ">u4"),
]
),
),
BranchConfiguration(
"vector<KM3NETDAQ::JDAQSummaryFrame>",
uproot.interpretation.jagged.AsJagged(
uproot.interpretation.numerical.AsDtype(
[
("dom_id", ">i4"),
("dq_status", ">u4"),
("hrv", ">u4"),
("fifo", ">u4"),
("status3", ">u4"),
("status4", ">u4"),
]
+ [(f"ch{c}", "u1") for c in range(31)]
),
header_bytes=10,
),
),
]
def __init__(self, fobj, step_size=1000):
if isinstance(fobj, str):
self._fobj = uproot.open(fobj)
else:
self._fobj = fobj
self._step_size = step_size
self._branch = self._fobj[self.TREE_ADDR]
self.ChunksConstructor = namedtuple(
field_names=["headers", "slices"], typename="SummarysliceChunk"
)
def _chunks_generator(self):
for chunk in self._branch.iterate(
dict(self._subbranches), step_size=self._step_size
):
yield self.ChunksConstructor(
*[getattr(chunk, bc.branch_address) for bc in self._subbranches]
)
def __iter__(self):
self._chunks = self._chunks_generator()
return self
def __next__(self):
return next(self._chunks)
def __len__(self):
return int(np.ceil(self._branch.num_entries / self._step_size))
def __repr__(self):
return f"<{self.__class__.__name__} {self._branch.num_entries} items, step_size={self._step_size} ({len(self)} chunks)>"
@nb.vectorize(
[nb.int32(nb.int8), nb.int32(nb.int16), nb.int32(nb.int32), nb.int32(nb.int64)]
)
def get_rate(value): # pragma: no cover
"""Return the rate in Hz from the short int value"""
if value == 0:
return 0
else:
return MINIMAL_RATE_HZ * np.exp(value * RATE_FACTOR)
@nb.guvectorize(
"void(i8, b1[:], b1[:])", "(), (n) -> (n)", target="parallel", nopython=True
)
def unpack_bits(value, bits_template, out): # pragma: no cover
"""Return a boolean array for a value's bit representation.
This function also accepts arrays as input, the output shape will be
NxM where N is the number of input values and M the length of the
``bits_template`` array, which is just a dummy array, due to the weird
signature system of numba.
Parameters
----------
value: int or np.array(int) with shape (N,)
The binary value of containing the bit information
bits_template: np.array() with shape (M,)
The template for the output array, the only important is its shape
Returns
-------
np.array(bool) either with shape (M,) or (N, M)
"""
for i in range(bits_template.shape[0]):
out[30 - i] = value & (1 << i) > 0
def get_channel_flags(value):
"""Returns the hrv/fifo flags for the PMT channels (hrv/fifo)
Parameters
----------
value : int32
The integer value to be parsed.
"""
channel_bits = np.bitwise_and(value, 0x7FFFFFFF)
flags = unpack_bits(channel_bits, CHANNEL_BITS_TEMPLATE)
return np.flip(flags, axis=-1)
def get_number_udp_packets(value):
"""Returns the number of received UDP packets (dq_status)
Parameters
----------
value : int32
The integer value to be parsed.
"""
return np.bitwise_and(value, 0x7FFF)
def get_udp_max_sequence_number(value):
"""Returns the maximum sequence number of the received UDP packets (dq_status)
Parameters
----------
value : int32
The integer value to be parsed.
"""
return np.right_shift(value, 16)
def has_udp_trailer(value):
"""Returns the UDP Trailer flag (fifo)
Parameters
----------
value : int32
The integer value to be parsed.
"""
return np.any(np.bitwise_and(value, np.left_shift(1, 31)))
class OnlineReader:
"""Reader for online ROOT files"""
def __init__(self, filename):
self._fobj = uproot3.open(filename)
self._filename = filename
self._events = None
self._timeslices = None
self._summaryslices = None
self._uuid = binascii.hexlify(self._fobj._context.uuid).decode("ascii")
@property
def uuid(self):
return self._uuid
def close(self):
self._fobj.close()
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
@property
def events(self):
if self._events is None:
tree = self._fobj["KM3NET_EVENT"]
headers = tree["KM3NETDAQ::JDAQEventHeader"].array(
uproot3.interpret(tree["KM3NETDAQ::JDAQEventHeader"], cntvers=True)
)
snapshot_hits = tree["snapshotHits"].array(
uproot3.asjagged(
uproot3.astable(
uproot3.asdtype(
[
("dom_id", ">i4"),
("channel_id", "u1"),
("time", "<u4"),
("tot", "u1"),
]
)
),
skipbytes=10,
)
)
triggered_hits = tree["triggeredHits"].array(
uproot3.asjagged(
uproot3.astable(
uproot3.asdtype(
[
("dom_id", ">i4"),
("channel_id", "u1"),
("time", "<u4"),
("tot", "u1"),
(" cnt", "u4"),
(" vers", "u2"),
("trigger_mask", ">u8"),
]
)
),
skipbytes=10,
)
)
self._events = OnlineEvents(headers, snapshot_hits, triggered_hits)
return self._events
@property
def timeslices(self):
if self._timeslices is None:
self._timeslices = Timeslices(self._fobj)
return self._timeslices
@property
def summaryslices(self):
if self._summaryslices is None:
self._summaryslices = SummarysliceReader(
uproot.open(self._filename)
) # TODO: remove when using uproot4
return self._summaryslices
class Timeslices:
"""A simple wrapper for timeslices"""
def __init__(self, fobj):
self._fobj = fobj
self._timeslices = {}
self._read_streams()
def _read_streams(self):
"""Read the L0, L1, L2 and SN streams if available"""
streams = set(
s.split(b"KM3NET_TIMESLICE_")[1].split(b";")[0]
for s in self._fobj.keys()
if b"KM3NET_TIMESLICE_" in s
)
for stream in streams:
tree = self._fobj[b"KM3NET_TIMESLICE_" + stream][
b"KM3NETDAQ::JDAQTimeslice"
]
headers = tree[b"KM3NETDAQ::JDAQTimesliceHeader"][b"KM3NETDAQ::JDAQHeader"][
b"KM3NETDAQ::JDAQChronometer"
]
if len(headers) == 0:
continue
superframes = tree[b"vector<KM3NETDAQ::JDAQSuperFrame>"]
hits_dtype = np.dtype([("pmt", "u1"), ("tdc", "<u4"), ("tot", "u1")])
hits_buffer = superframes[
b"vector<KM3NETDAQ::JDAQSuperFrame>.buffer"
].lazyarray(
uproot3.asjagged(
uproot3.astable(uproot3.asdtype(hits_dtype)), skipbytes=6
),
basketcache=uproot3.cache.ThreadSafeArrayCache(
TIMESLICE_FRAME_BASKET_CACHE_SIZE
),
)
self._timeslices[stream.decode("ascii")] = (
headers,
superframes,
hits_buffer,
)
setattr(
self,
stream.decode("ascii"),
TimesliceStream(headers, superframes, hits_buffer),
)
def stream(self, stream, idx):
ts = self._timeslices[stream]
return Timeslice(*ts, idx, stream)
def __str__(self):
return "Available timeslice streams: {}".format(
", ".join(s for s in self._timeslices.keys())
)
def __repr__(self):
return str(self)
class TimesliceStream:
def __init__(self, headers, superframes, hits_buffer):
# self.headers = headers.lazyarray(
# uproot3.asjagged(uproot3.astable(
# uproot3.asdtype(
# np.dtype([('a', 'i4'), ('b', 'i4'), ('c', 'i4'),
# ('d', 'i4'), ('e', 'i4')]))),
# skipbytes=6),
# basketcache=uproot3.cache.ThreadSafeArrayCache(
# TIMESLICE_FRAME_BASKET_CACHE_SIZE))
self.headers = headers
self.superframes = superframes
self._hits_buffer = hits_buffer
# def frames(self):
# n_hits = self._superframe[
# b'vector<KM3NETDAQ::JDAQSuperFrame>.numberOfHits'].lazyarray(
# basketcache=BASKET_CACHE)[self._idx]
# module_ids = self._superframe[
# b'vector<KM3NETDAQ::JDAQSuperFrame>.id'].lazyarray(basketcache=BASKET_CACHE)[self._idx]
# idx = 0
# for module_id, n_hits in zip(module_ids, n_hits):
# self._frames[module_id] = hits_buffer[idx:idx + n_hits]
# idx += n_hits
class Timeslice:
"""A wrapper for a timeslice"""
def __init__(self, header, superframe, hits_buffer, idx, stream):
self.header = header
self._frames = {}
self._superframe = superframe
self._hits_buffer = hits_buffer
self._idx = idx
self._stream = stream
self._n_frames = None
@property
def frames(self):
if not self._frames:
self._read_frames()
return self._frames
def _read_frames(self):
"""Populate a dictionary of frames with the module ID as key"""
hits_buffer = self._hits_buffer[self._idx]
n_hits = self._superframe[
b"vector<KM3NETDAQ::JDAQSuperFrame>.numberOfHits"
].lazyarray(basketcache=BASKET_CACHE)[self._idx]
try:
module_ids = self._superframe[
b"vector<KM3NETDAQ::JDAQSuperFrame>.id"
].lazyarray(basketcache=BASKET_CACHE)[self._idx]
except KeyError:
module_ids = (
self._superframe[
b"vector<KM3NETDAQ::JDAQSuperFrame>.KM3NETDAQ::JDAQModuleIdentifier"
]
.lazyarray(
uproot3.asjagged(
uproot3.astable(uproot3.asdtype([("dom_id", ">i4")]))
),
basketcache=BASKET_CACHE,
)[self._idx]
.dom_id
)
idx = 0
for module_id, n_hits in zip(module_ids, n_hits):
self._frames[module_id] = hits_buffer[idx : idx + n_hits]
idx += n_hits
def __len__(self):
if self._n_frames is None:
self._n_frames = len(
self._superframe[b"vector<KM3NETDAQ::JDAQSuperFrame>.id"].lazyarray(
basketcache=BASKET_CACHE
)[self._idx]
)
return self._n_frames
def __str__(self):
return "{} timeslice with {} frames.".format(self._stream, len(self))
def __repr__(self):
return "<{}: {} entries>".format(self.__class__.__name__, len(self.header))
class OnlineEvents:
"""A simple wrapper for online events"""
def __init__(self, headers, snapshot_hits, triggered_hits):
self.headers = headers
self.snapshot_hits = snapshot_hits
self.triggered_hits = triggered_hits
def __getitem__(self, item):
return OnlineEvent(
self.headers[item], self.snapshot_hits[item], self.triggered_hits[item]
)
def __len__(self):
return len(self.headers)
def __str__(self):
return "Number of events: {}".format(len(self.headers))
def __repr__(self):
return str(self)
class OnlineEvent:
"""A wrapper for a online event"""
def __init__(self, header, snapshot_hits, triggered_hits):
self.header = header
self.snapshot_hits = snapshot_hits
self.triggered_hits = triggered_hits
def __str__(self):
return "Online event with {} snapshot hits and {} triggered hits".format(
len(self.snapshot_hits), len(self.triggered_hits)
)
def __repr__(self):
return str(self)