python overload

Author: casperdcl

python/petsird/helpers/__init__.py

@@ -83,11 +83,37 @@ def make_detection_bin(
                                [expanded_detection_bin])[0]
 
 
+@typing.overload
 def get_detection_efficiency(scanner: petsird.ScannerInformation,
                              type_of_module_pair: petsird.TypeOfModulePair,
                              detection_bin_1: petsird.DetectionBin,
                              detection_bin_2: petsird.DetectionBin) -> float:
     """Compute the detection efficiency for a pair of detectors"""
+    ...
+
+
+@typing.overload
+def get_detection_efficiency(scanner: petsird.ScannerInformation,
+                             type_of_module_pair: petsird.TypeOfModulePair,
+                             event: petsird.CoincidenceEvent) -> float:
+    """Compute the detection efficiency for a coincidence event"""
+    ...
+
+
+def get_detection_efficiency(
+        scanner: petsird.ScannerInformation,
+        type_of_module_pair: petsird.TypeOfModulePair,
+        event_or_detection_bin_1: typing.Union[petsird.CoincidenceEvent,
+                                               petsird.DetectionBin],
+        detection_bin_2: petsird.DetectionBin = None) -> float:
+    """Compute the detection efficiency"""
+    if isinstance(event_or_detection_bin_1, petsird.DetectionBin):
+        detection_bin_1 = event_or_detection_bin_1
+        assert detection_bin_2 is not None, "Second detection bin must be provided"
+    else:
+        detection_bin_1, detection_bin_2 = event_or_detection_bin_1.detection_bins[:
+                                                                                   2]
+
     if scanner.detection_efficiencies is None:
         # should never happen really, but this way, we don't crash.
         return 1.
@@ -139,13 +165,3 @@ def get_detection_efficiency(scanner: petsird.ScannerInformation,
             expanded_det_bin1.energy_index]
 
     return eff
-
-
-def get_event_detection_efficiency(
-        scanner: petsird.ScannerInformation,
-        type_of_module_pair: petsird.TypeOfModulePair,
-        event: petsird.CoincidenceEvent) -> float:
-    """Compute the detection efficiency for a coincidence event"""
-    return get_detection_efficiency(scanner, type_of_module_pair,
-                                    event.detection_bins[0],
-                                    event.detection_bins[1])

python/petsird/helpers/analysis.py

@@ -8,8 +8,8 @@
 
 import petsird
 import petsird.helpers.geometry
-from petsird.helpers import (expand_detection_bin,
-                             get_event_detection_efficiency, get_num_det_els)
+from petsird.helpers import (expand_detection_bin, get_detection_efficiency,
+                             get_num_det_els)
 
 
 def parserCreator():
@@ -123,7 +123,7 @@ def parserCreator():
                                     ", ",
                                     expanded_det_bin1,
                                 )
-                                eff = get_event_detection_efficiency(
+                                eff = get_detection_efficiency(
                                     scanner, mtype_pair, event)
                                 print("    efficiency:", eff)
 

python/petsird/helpers/generator.py

@@ -9,8 +9,7 @@
 
 import numpy
 import petsird
-from petsird.helpers import (get_event_detection_efficiency,
-                             get_num_detection_bins)
+from petsird.helpers import get_detection_efficiency, get_num_detection_bins
 
 # these are constants for now
 NUMBER_OF_EVENT_ENERGY_BINS = 3
@@ -268,8 +267,8 @@ def get_events(header: petsird.Header,
             # short-cut to directly generate a random detection bin
             event.detection_bins[1] = get_random_uint(count1)
 
-            if get_event_detection_efficiency(header.scanner,
-                                              type_of_module_pair, event) > 0:
+            if get_detection_efficiency(header.scanner, type_of_module_pair,
+                                        event) > 0:
                 # in coincidence, we can get out of the loop
                 break