diff --git a/clusters/__init__.py b/clusters/__init__.py index ad9e8ab..3a096c8 100644 --- a/clusters/__init__.py +++ b/clusters/__init__.py @@ -1 +1,2 @@ -from clusters.find import find_clusters \ No newline at end of file +from clusters.find import find_clusters +from clusters.find import cluster_info_datatype diff --git a/clusters/find.py b/clusters/find.py index 9f17c07..d830479 100644 --- a/clusters/find.py +++ b/clusters/find.py @@ -8,7 +8,6 @@ logger = logging.getLogger('root') - def find_clusters(settings, hits): # Outsource the main cluster finding routine to a Rust compiled library hits_stacked = np.stack((hits['x'], hits['y'], hits['ToA']), axis=-1).astype('int64') @@ -22,8 +21,7 @@ def find_clusters(settings, hits): # Build the maximum size clusters we expect, before filtering cm_chunk = np.zeros((len(clusters), 2, settings.cluster_matrix_size, settings.cluster_matrix_size), dt_clusters) - ci_chunk = np.zeros(len(clusters), dtype=dt_ci) - + ci_chunk = np.zeros(len(clusters), dtype=cluster_info_datatype(settings.cluster_stats)) # Loop over all clusters i = 0 for cluster in clusters: @@ -85,4 +83,15 @@ def build_cluster(c, settings, i, cm, ci): ci[i]['y'] = min_y ci[i]['ToA'] = min_ctoa + if settings.cluster_stats: + ci[i]['sumToT'] = np.sum(c['ToT']) + ci[i]['nHits'] = len(c) + + +def cluster_info_datatype(cluster_stats): + dt = dt_ci_base + + if cluster_stats: + dt = dt + dt_ci_extended + return numpy.dtype(dt) diff --git a/default.cfg b/default.cfg index e810e5f..9bfcc9c 100644 --- a/default.cfg +++ b/default.cfg @@ -30,9 +30,9 @@ cluster_min_sum_tot = 200 cluster_max_sum_tot = 400 # Store cluster_stats -cluster_stats = True +cluster_stats = False ## Event finding algorithm = centroid event_cnn_model = model-200kv-tottoa.h5 -event_cnn_tot_only = False \ No newline at end of file +event_cnn_tot_only = False diff --git a/events/__init__.py b/events/__init__.py index 5f7d86b..0692027 100644 --- a/events/__init__.py +++ b/events/__init__.py @@ -2,4 +2,5 @@ from events.localise import cnn from events.super_resolution import subpixel_event_redistribution from events.chip_edge_correct import chip_edge_correct -from events.predictions import calculate_predictions \ No newline at end of file +from events.predictions import calculate_predictions +from events.localise import event_info_datatype diff --git a/events/localise.py b/events/localise.py index fd75838..70053d2 100644 --- a/events/localise.py +++ b/events/localise.py @@ -3,36 +3,36 @@ from scipy import ndimage import tpx3format -from lib import UserConfigException +from lib.exceptions import UserConfigException from lib.constants import * import numpy as np -from tqdm import tqdm - -# https://github.com/tqdm/tqdm/issues/481 -tqdm.monitor_interval = 0 logger = logging.getLogger('root') -def localise_events(cluster_matrix, cluster_info, method): +def localise_events(cluster_matrix, cluster_info, method, cluster_stats): logger.debug("Started event localization on %d events using method %s" % (len(cluster_info), method)) + events = np.empty(len(cluster_info), event_info_datatype(cluster_stats)) + + # TODO: The events matrix could be seeded by the clusters_info matrix. Like this: + # events = cluster_info[()].astype(dt_event) + if method == "centroid": - events = calculate_centroid(cluster_matrix, cluster_info) + events = calculate_centroid(cluster_matrix, cluster_info, events, cluster_stats) elif method == "random": - events = calculate_random(cluster_matrix, cluster_info) + events = calculate_random(cluster_matrix, cluster_info, events, cluster_stats) elif method == "highest_toa": - events = calculate_toa(cluster_matrix, cluster_info) + events = calculate_toa(cluster_matrix, cluster_info, events, cluster_stats) elif method == "highest_tot": - events = calculate_tot(cluster_matrix, cluster_info) + events = calculate_tot(cluster_matrix, cluster_info, events, cluster_stats) else: raise Exception("Chosen localisation algorithm ('%s') does not exist" % method) return events -def calculate_centroid(cluster_matrix, cluster_info): - events = np.empty(len(cluster_info), dtype=dt_event) +def calculate_centroid(cluster_matrix, cluster_info, events, cluster_stats): # Raise runtime warnings, instead of printing them numpy.seterr(all='raise') @@ -53,12 +53,14 @@ def calculate_centroid(cluster_matrix, cluster_info): events[idx]['ToA'] = cluster_info[idx]['ToA'] - return events + if cluster_stats: + events[idx]['sumToT'] = cluster_info[idx]['sumToT'] + events[idx]['nHits'] = cluster_info[idx]['nHits'] + return events -def calculate_random(cluster_matrix, cluster_info): - events = np.empty(len(cluster_info), dtype=dt_event) +def calculate_random(cluster_matrix, cluster_info, events, cluster_stats): for idx, cluster in enumerate(cluster_matrix): nzy, nzx = np.nonzero(cluster[0]) @@ -80,12 +82,14 @@ def calculate_random(cluster_matrix, cluster_info): events[idx]['ToA'] = cluster_info[idx]['ToA'] - return events + if cluster_stats: + events[idx]['sumToT'] = cluster_info[idx]['sumToT'] + events[idx]['nHits'] = cluster_info[idx]['nHits'] + return events -def calculate_toa(cluster_matrix, cluster_info): - events = np.empty(len(cluster_info), dtype=dt_event) +def calculate_toa(cluster_matrix, cluster_info, events, cluster_stats): for idx, cluster in enumerate(cluster_matrix): if np.max(cluster[1]) == 0: logger.debug("Could not calculate highest_toa: empty ToA cluster Picking random ToT pixel. Cluster_info: %s" % cluster_info[idx]) @@ -115,12 +119,14 @@ def calculate_toa(cluster_matrix, cluster_info): events[idx]['ToA'] = cluster_info[idx]['ToA'] - return events + if cluster_stats: + events[idx]['sumToT'] = cluster_info[idx]['sumToT'] + events[idx]['nHits'] = cluster_info[idx]['nHits'] + return events -def calculate_tot(cluster_matrix, cluster_info): - events = np.empty(len(cluster_info), dtype=dt_event) +def calculate_tot(cluster_matrix, cluster_info, events, cluster_stats): for idx, cluster in enumerate(cluster_matrix): i = np.argmax(cluster[0]) y, x = np.unravel_index(i, cluster[0].shape) @@ -134,10 +140,14 @@ def calculate_tot(cluster_matrix, cluster_info): events[idx]['ToA'] = cluster_info[idx]['ToA'] + if cluster_stats: + events[idx]['sumToT'] = cluster_info[idx]['sumToT'] + events[idx]['nHits'] = cluster_info[idx]['nHits'] + return events -def cnn(cluster_matrix, cluster_info, model, tot_only, hits_cross_extra_offset): +def cnn(cluster_matrix, cluster_info, model, tot_only, hits_cross_extra_offset, cluster_stats): # Delete ToA matrices, required for ToT only CNN if tot_only: cluster_matrix = np.delete(cluster_matrix, 1, 1) @@ -153,7 +163,7 @@ def cnn(cluster_matrix, cluster_info, model, tot_only, hits_cross_extra_offset): predictions = model.predict(cluster_matrix, batch_size=EVENTS_CHUNK_SIZE, verbose=0) # Copy all events from cluster_info as base - events = cluster_info[()].astype(dt_event) + events = cluster_info[()].astype(event_info_datatype(cluster_stats)) # Add prediction offset from cluster origin events['x'] = events['x'] + predictions[:, 1] @@ -171,3 +181,12 @@ def cnn(cluster_matrix, cluster_info, model, tot_only, hits_cross_extra_offset): logger.debug('Removed %d events found outside image matrix shape (%d).' % (deleted, shape)) return events + + +def event_info_datatype(cluster_stats): + dt = dt_event_base + + if cluster_stats: + dt = dt + dt_event_extended + + return numpy.dtype(dt) diff --git a/lib/config.py b/lib/config.py index b8d536a..03b841d 100644 --- a/lib/config.py +++ b/lib/config.py @@ -82,6 +82,7 @@ def parse_config(argv=None): output_group.add_argument("--store_clusters", action='store_true', help="Store /clusters in output file") output_group.add_argument("--store_events", action='store_true', help="Store /events in output file") # output_group.add_argument("--store_predictions", action='store_true', help="Store /predictions in output file") + output_group.add_argument("--cluster_stats", action='store_true', help="Add the cluster sumToT and nHits to cluster_info or events") # Post process options post_process_group = parser.add_argument_group('post processing') diff --git a/lib/constants.py b/lib/constants.py index 4a0b202..f021544 100644 --- a/lib/constants.py +++ b/lib/constants.py @@ -13,23 +13,33 @@ HITS_CHUNK_SIZE = 10_000_000 # Cluster info key indices -dt_ci = numpy.dtype([ +dt_ci_base = [ ('chipId', numpy.uint8), ('x', numpy.uint16), ('y', numpy.uint16), ('ToA', numpy.int64), -]) +] +dt_ci_extended = [ + ('sumToT', numpy.uint16), + ('nHits', numpy.uint8) +] + dt_clusters = 'uint16' CLUSTER_CHUNK_SIZE = 100_000 # Event matrix key indices -dt_event = numpy.dtype([ +dt_event_base = [ ('chipId', numpy.uint8), ('x', numpy.float64), ('y', numpy.float64), - ('ToA', numpy.int64) -]) + ('ToA', numpy.int64), +] +dt_event_extended = [ + ('sumToT', numpy.uint16), + ('nHits', numpy.uint8) +] + EVENTS_CHUNK_SIZE = 100_000 # Control events @@ -39,4 +49,4 @@ CONTROL_OTHER_CHIP_COMMAND = 0x7200 # 29184 # ToT correction matrix shape -TOT_CORRECTION_SHAPE = (1024, 256, 256, 4) \ No newline at end of file +TOT_CORRECTION_SHAPE = (1024, 256, 256, 4) diff --git a/lib/io.py b/lib/io.py index 6d8f117..07a6b98 100644 --- a/lib/io.py +++ b/lib/io.py @@ -7,8 +7,10 @@ import tpx3format from lib import constants -logger = logging.getLogger('root') +import clusters +import events as ev +logger = logging.getLogger('root') class io: write = None @@ -88,9 +90,9 @@ def replace_hits(self, hits): def del_hits(self): del self.write['hits'] - def write_cluster_chunk(self, ci, cm, cms): + def write_cluster_chunk(self, ci, cm, cms, cluster_stats): if 'cluster_info' not in self.write: - self.write.create_dataset('cluster_info', shape=(len(ci),), dtype=constants.dt_ci, maxshape=(None,), + self.write.create_dataset('cluster_info', shape=(len(ci),), dtype=clusters.cluster_info_datatype(cluster_stats), maxshape=(None,), chunks=(constants.CLUSTER_CHUNK_SIZE,), data=ci) self.write.create_dataset('clusters', shape=(len(cm), 2, cms, cms), dtype=constants.dt_clusters, maxshape=(None, 2, cms, cms), @@ -134,10 +136,10 @@ def del_clusters(self): del self.write['cluster_info'] del self.write['clusters'] - def write_event_chunk(self, events): + def write_event_chunk(self, events, cluster_stats): if 'events' not in self.write: shape = (constants.EVENTS_CHUNK_SIZE,) - self.write.create_dataset('events', dtype=constants.dt_event, maxshape=(None,), chunks=shape, data=events) + self.write.create_dataset('events', dtype=ev.event_info_datatype(cluster_stats), maxshape=(None,), chunks=shape, data=events) else: events_f = self.write['events'] old = len(events_f) diff --git a/orchestration/gpu.py b/orchestration/gpu.py index 3ea1863..70f2f9c 100644 --- a/orchestration/gpu.py +++ b/orchestration/gpu.py @@ -5,9 +5,10 @@ import queue import events +import clusters import numpy as np -from lib.constants import EVENTS_CHUNK_SIZE, dt_ci, dt_clusters +from lib.constants import EVENTS_CHUNK_SIZE, dt_clusters class Gpu(Process): @@ -23,7 +24,7 @@ def __init__(self, settings, keep_processing, gq, oq): self.model = None self.clusters = np.zeros((EVENTS_CHUNK_SIZE, 2, self.settings.cluster_matrix_size, self.settings.cluster_matrix_size), dtype=dt_clusters) - self.cluster_info = np.zeros(EVENTS_CHUNK_SIZE, dtype=dt_ci) + self.cluster_info = np.zeros(EVENTS_CHUNK_SIZE, dtype=clusters.cluster_info_datatype(self.settings.cluster_stats)) self.offset = 0 self.chunks = [] @@ -87,8 +88,11 @@ def bunch_clusters(self, cluster_chunk, cluster_info_chunk): # Parse and sent to GPU def parse_clusters_gpu(self): - e = events.cnn(self.clusters[:self.offset], self.cluster_info[:self.offset], self.model, self.settings.event_cnn_tot_only, self.settings.hits_cross_extra_offset) - # e = events.localise_events(self.clusters[:self.offset], self.cluster_info[:self.offset], 'centroid') + e = events.cnn(self.clusters[:self.offset], self.cluster_info[:self.offset], self.model, + self.settings.event_cnn_tot_only, self.settings.hits_cross_extra_offset, + self.settings.cluster_stats + ) + # e = events.localise_events(self.clusters[:self.offset], self.cluster_info[:self.offset], 'centroid', self.settings.cluster_stats) if self.settings.store_events: self.output_queue.put({ diff --git a/orchestration/worker.py b/orchestration/worker.py index d859771..82ea8ec 100644 --- a/orchestration/worker.py +++ b/orchestration/worker.py @@ -99,7 +99,7 @@ def parse_hits(self, hits): return cm_chunk, ci_chunk # From clusters to events - def parse_clusters(self, clusters, cluster_info): - e = events.localise_events(clusters, cluster_info, self.settings.algorithm) + def parse_clusters(self, cl, cluster_info): + e = events.localise_events(cl, cluster_info, self.settings.algorithm, self.settings.cluster_stats) return e diff --git a/orchestration/writer.py b/orchestration/writer.py index 4923ab7..6f2c1e7 100644 --- a/orchestration/writer.py +++ b/orchestration/writer.py @@ -78,8 +78,8 @@ def write_hits(self, hits): self.io.write_hit_chunk(hits) def write_clusters(self, cluster_info, clusters): - self.io.write_cluster_chunk(cluster_info, clusters, self.settings.cluster_matrix_size) + self.io.write_cluster_chunk(cluster_info, clusters, self.settings.cluster_matrix_size, self.settings.cluster_stats) def write_events(self, e): - self.io.write_event_chunk(e) + self.io.write_event_chunk(e, self.settings.cluster_stats)