Coverage for src/km3modules/hits.py: 27%

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3# vim:set ts=4 sts=4 sw=4 et: 

4""" 

5Functions and modules to work with hits. 

6""" 

7 

8from numba import njit 

9import numpy as np 

10 

11import km3pipe as kp 

12 

13log = kp.logger.get_logger(__name__) 

14 

15__author__ = "Tamas Gal" 

16 

17 

18@njit 

19def count_multiplicities(times, tmax=20): 

20 """Calculate an array of multiplicities and corresponding coincidence IDs 

21 

22 Note that this algorithm does not take care about DOM IDs, so it has to 

23 be fed with DOM hits. 

24 

25 Parameters 

26 ---------- 

27 times: array[float], shape=(n,) 

28 Hit times for n hits 

29 dt: int [default: 20] 

30 Time window of a coincidence 

31 

32 Returns 

33 ------- 

34 (array[int]), array[int]), shape=(n,) 

35 

36 """ 

37 n = times.shape[0] 

38 mtp = np.ones(n, dtype=np.int64) # multiplicities 

39 cid = np.zeros(n, dtype=np.int64) # coincidence id 

40 idx0 = 0 

41 _mtp = 1 

42 _cid = 0 

43 t0 = times[idx0] 

44 for i in range(1, n): 

45 dt = times[i] - t0 

46 if dt > tmax: 

47 mtp[idx0:i] = _mtp 

48 cid[idx0:i] = _cid 

49 _mtp = 0 

50 _cid += 1 

51 idx0 = i 

52 t0 = times[i] 

53 _mtp += 1 

54 if i == n - 1: 

55 mtp[idx0:] = _mtp 

56 cid[idx0:] = _cid 

57 break 

58 

59 return mtp, cid