Use Gaussian noise in place of Poisson noise beyond threshold to impr…

…ove performance
nion-software · Aug 8, 2024 · aec8bc5 · aec8bc5
1 parent 56e31cc
commit aec8bc5
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/nionswift_plugin/usim/Noise.py b/nionswift_plugin/usim/Noise.py
@@ -14,11 +14,16 @@ def __init__(self) -> None:
         self.enabled = True
         self.poisson_level: typing.Optional[float] = None
 
-    def apply(self, input: DataAndMetadata.DataAndMetadata) -> DataAndMetadata.DataAndMetadata:
+    def apply(self, input: DataAndMetadata.DataAndMetadata, lambda_thresh: float = 1.0) -> DataAndMetadata.DataAndMetadata:
         if self.enabled and self.poisson_level:
             rs = numpy.random.RandomState()  # use this to avoid blocking other calls to poisson
             input_data = input.data
             assert input_data is not None
-            poisson_data = typing.cast(_NDArray, rs.poisson(self.poisson_level, size=input_data.shape).astype(input_data.dtype))
+            if self.poisson_level > lambda_thresh:
+                # Since it is 'high' lambda, we can approximate it to a normal distribution
+                poisson_data = typing.cast(_NDArray, rs.normal(loc=self.poisson_level, scale=numpy.sqrt(self.poisson_level), size=input_data.shape).astype(input_data.dtype))
+            else:
+                poisson_data = typing.cast(_NDArray, rs.poisson(self.poisson_level, size=input_data.shape).astype(input_data.dtype))
             return input + (poisson_data - self.poisson_level)
         return input
+