Merge d5dcccf into 2d4c000

csxtools/image/stack.py

@@ -1,6 +1,9 @@
 import numpy as np
 from ..ext import image as extimage
 
+import logging
+logger = logging.getLogger(__name__)
+
 
 def stackmean(array):
     """Cacluate the mean of a stack
@@ -24,25 +27,44 @@ def stackmean(array):
     return X
 
 
-def stacksum(array):
+def stacksum(array, norm=False):
     """Cacluate the sum of a stack
 
     This function calculates the sum of a stack of images (or any array).
     It ignores values that are np.NAN and does not include them in the sum
     calculation. It assumes an array of shape (.. i, j, x, y) where x and y
     are the size of the returned array (x, y).
 
+    If norm is True then the output sum is corrected for elements where NaNs
+    are encountered and renormalized to the value expected based on the stack
+    size.
+
     Parameters
     ----------
     array : array_like
         Input array of at least 3 dimensions.
 
+    norm : bool
+        If true then normalize output to account for NaNs.
+
     Returns
     -------
     tuple
         tuple of 2 arrays of the sum and number of points in the sum
     """
     X, Y = extimage.stackprocess(array, 0)
+
+    total_elements = array.size / (array.shape[-1] * array.shape[-2])
+
+    if not norm and np.sum(Y != total_elements):
+        logger.warning("stacksum encountered NaN values and excluded these "
+                       "values from the sum. Consider using the number of "
+                       "points, to renormalize the image. Hint: use "
+                       "norm=True")
+
+    if norm:
+        X = X * (total_elements / Y)
+
     return X, Y