Skip to content
Snippets Groups Projects

Resolve "Best way to get number of hit DOMs"

Merged Tamas Gal requested to merge 46-best-way-to-get-number-of-hit-doms into master
3 files
+ 107
1
Compare changes
  • Side-by-side
  • Inline
Files
3
+ 45
0
#!/usr/bin/env python3
#!/usr/bin/env python3
 
import numba as nb
 
import numpy as np
import uproot
import uproot
 
# 110 MB based on the size of the largest basket found so far in km3net
# 110 MB based on the size of the largest basket found so far in km3net
BASKET_CACHE_SIZE = 110 * 1024**2
BASKET_CACHE_SIZE = 110 * 1024**2
BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
BASKET_CACHE = uproot.cache.ThreadSafeArrayCache(BASKET_CACHE_SIZE)
@@ -40,3 +43,45 @@ def to_num(value):
@@ -40,3 +43,45 @@ def to_num(value):
except (ValueError, TypeError):
except (ValueError, TypeError):
pass
pass
return value
return value
 
 
 
@nb.jit(nopython=True)
 
def unique(array, dtype=np.int64):
 
"""Return the unique elements of an array with a given dtype.
 
 
The performance is better for pre-sorted input arrays.
 
 
"""
 
n = len(array)
 
out = np.empty(n, dtype)
 
last = array[0]
 
entry_idx = 0
 
out[entry_idx] = last
 
for i in range(1, n):
 
current = array[i]
 
if current == last: # shortcut for sorted arrays
 
continue
 
already_present = False
 
for j in range(entry_idx + 1):
 
if current == out[j]:
 
already_present = True
 
break
 
if not already_present:
 
entry_idx += 1
 
out[entry_idx] = current
 
last = current
 
return out[:entry_idx+1]
 
 
 
@nb.jit(nopython=True)
 
def uniquecount(array, dtype=np.int64):
 
"""Count the number of unique elements in a jagged Awkward1 array."""
 
n = len(array)
 
out = np.empty(n, dtype)
 
for i in range(n):
 
sub_array = array[i]
 
if len(sub_array) == 0:
 
out[i] = 0
 
else:
 
out[i] = len(unique(sub_array))
 
return out
Loading