diff --git a/python/aare/ClusterFinder.py b/python/aare/ClusterFinder.py index a2042a4..f678dd1 100644 --- a/python/aare/ClusterFinder.py +++ b/python/aare/ClusterFinder.py @@ -1,5 +1,5 @@ -from ._aare import ClusterFinder_Cluster3x3i +from ._aare import ClusterFinder_Cluster3x3i, ClusterFinder_Cluster2x2i, ClusterFinderMT_Cluster3x3i, ClusterFinderMT_Cluster2x2i, ClusterCollector_Cluster3x3i, ClusterCollector_Cluster2x2i import numpy as np def ClusterFinder(image_size, cluster_size, n_sigma=5, dtype = np.int32, capacity = 1024): @@ -9,6 +9,42 @@ def ClusterFinder(image_size, cluster_size, n_sigma=5, dtype = np.int32, capacit """ if dtype == np.int32 and cluster_size == (3,3): return ClusterFinder_Cluster3x3i(image_size, n_sigma = n_sigma, capacity=capacity) + elif dtype == np.int32 and cluster_size == (2,2): + return ClusterFinder_Cluster2x2i(image_size, n_sigma = n_sigma, capacity=capacity) else: #TODO! add the other formats raise ValueError(f"Unsupported dtype: {dtype}. Only np.int32 is supported.") + + +def ClusterFinderMT(image_size, cluster_size = (3,3), dtype=np.int32, n_sigma=5, capacity = 1024, n_threads = 3): + """ + Factory function to create a ClusterFinderMT object. Provides a cleaner syntax for + the templated ClusterFinderMT in C++. + """ + + if dtype == np.int32 and cluster_size == (3,3): + return ClusterFinderMT_Cluster3x3i(image_size, n_sigma = n_sigma, + capacity = capacity, n_threads = n_threads) + elif dtype == np.int32 and cluster_size == (2,2): + return ClusterFinderMT_Cluster2x2i(image_size, n_sigma = n_sigma, + capacity = capacity, n_threads = n_threads) + else: + #TODO! add the other formats + raise ValueError(f"Unsupported dtype: {dtype}. Only np.int32 is supported.") + + +def ClusterCollector(clusterfindermt, cluster_size = (3,3), dtype=np.int32): + """ + Factory function to create a ClusterCollector object. Provides a cleaner syntax for + the templated ClusterCollector in C++. + """ + + if dtype == np.int32 and cluster_size == (3,3): + return ClusterCollector_Cluster3x3i(clusterfindermt) + elif dtype == np.int32 and cluster_size == (2,2): + return ClusterCollector_Cluster2x2i(clusterfindermt) + + else: + #TODO! add the other formats + raise ValueError(f"Unsupported dtype: {dtype}. Only np.int32 is supported.") + diff --git a/python/aare/__init__.py b/python/aare/__init__.py index b1eb604..096dddf 100644 --- a/python/aare/__init__.py +++ b/python/aare/__init__.py @@ -11,9 +11,10 @@ from ._aare import ROI # from ._aare import ClusterFinderMT, ClusterCollector, ClusterFileSink, ClusterVector_i -from .ClusterFinder import ClusterFinder +from .ClusterFinder import ClusterFinder, ClusterCollector, ClusterFinderMT from .ClusterVector import ClusterVector + from ._aare import fit_gaus, fit_pol1 from ._aare import Interpolator from ._aare import calculate_eta2