Add kappa-sigma rejected average to stacking possibilities

Author: pnuu

bin/halostack_cli.py

@@ -88,8 +88,10 @@ def halostack_cli(args):
     images = args['fname_in']
 
     stacks = []
-    for stack in args['stacks']:
-        stacks.append(Stack(stack, len(images), nprocs=args['nprocs']))
+    for i in range(len(args['stacks'])):
+        stacks.append(Stack(args['stacks'][i], len(images),
+                            nprocs=args['nprocs'],
+                            kwargs=args['stack_kwargs'][i]))
 
     base_img_fname = images[0]
     base_img = Image(fname=base_img_fname, nprocs=args['nprocs'])
@@ -202,6 +204,12 @@ def main():
     parser.add_argument("-d", "--median", dest="median_stack_file",
                         default=None, metavar="FILE",
                         help="Output filename of the median stack")
+    parser.add_argument("-S", "--sigma", dest="sigma_stack_file",
+                        default=None, metavar="FILE",
+                        help="Output filename of the kappa-sigma stack")
+    parser.add_argument("-k", "--kappa-sigma-params", dest="kappa_sigma_params",
+                        default=None, metavar="KAPPA,ITERATIONS",
+                        help="Kappa and iterations for the kappa-sigma stack")
     parser.add_argument("-t", "--correlation-threshold",
                         dest="correlation_threshold",
                         default=None, metavar="NUM", type=float,
@@ -267,24 +275,40 @@ def main():
     # Check which stacks will be made
     stacks = []
     stack_fnames = []
+    stack_kwargs = []
     if args['min_stack_file']:
         stacks.append('min')
         stack_fnames.append(args['min_stack_file'])
+        stack_kwargs.append(None)
         LOGGER.debug("Added minimum stack")
     if args['max_stack_file']:
         stacks.append('max')
         stack_fnames.append(args['max_stack_file'])
+        stack_kwargs.append(None)
         LOGGER.debug("Added maximum stack")
     if args['avg_stack_file']:
         stacks.append('mean')
         stack_fnames.append(args['avg_stack_file'])
+        stack_kwargs.append(None)
         LOGGER.debug("Added average stack")
     if args['median_stack_file']:
         stacks.append('median')
         stack_fnames.append(args['median_stack_file'])
+        stack_kwargs.append(None)
+        LOGGER.debug("Added median stack")
+    if args['sigma_stack_file']:
+        stacks.append('sigma')
+        stack_fnames.append(args['sigma_stack_file'])
+        if args['kappa_sigma_params']:
+            tmp = args['kappa_sigma_params'].split(',')
+            stack_kwargs.append({'kappa': float(tmp[0]),
+                                 'max_iters': int(tmp[1])})
+        else:
+            stack_kwargs.append(None)
         LOGGER.debug("Added median stack")
     args['stacks'] = stacks
     args['stack_fnames'] = stack_fnames
+    args['stack_kwargs'] = stack_kwargs
 
     # Check if there's anything to do
     if len(args['stacks']) == 0 and args['save_prefix'] is None or \

halostack/stack.py

@@ -44,9 +44,10 @@ class Stack(object):
     'max' - maximum stack
     'mean' - average stack
     'median' - median stack
+    'sigma' - kappa-sigma stack
     '''
 
-    def __init__(self, mode, num, nprocs=1):
+    def __init__(self, mode, num, nprocs=1, kwargs=None):
         LOGGER.debug("Initializing %s stack for %d images.",
                      mode, num)
         self.stack = None
@@ -58,14 +59,15 @@ def __init__(self, mode, num, nprocs=1):
                                         'calc': None},
                                 'mean': {'update': self._update_mean,
                                          'calc': None},
-                                # 'sigma': {'update': self._update_deep,
-                                #          'calc': self._calculate_sigma},
+                                'sigma': {'update': self._update_deep,
+                                          'calc': self._calculate_sigma},
                                 'median': {'update': self._update_deep,
                                            'calc': self._calculate_median}}
         self._update_func = self._mode_functions[mode]['update']
         self._calculate_func = self._mode_functions[mode]['calc']
         self.num = num
         self._num = 0
+        self._kwargs = kwargs
 
     def add_image(self, img):
         '''Add a frame to the stack.
@@ -178,4 +180,48 @@ def _calculate_sigma(self):
         '''Calculate the sigma-reject average of the stack and return
         the result as Image(dtype=uint16).
         '''
-        LOGGER.error("Sigma stack not implemented.")
+        shape = self.stack['R'].shape
+        img = np.zeros((shape[0], shape[1], 3), dtype=self.stack['R'].dtype)
+
+        try:
+            kappa = self._kwargs["kappa"]
+        except (TypeError, KeyError):
+            kappa = 2
+        try:
+            max_iters = self._kwargs["max_iters"]
+        except (TypeError, KeyError):
+            max_iters = self._num/4
+
+        LOGGER.info("Calculating Sigma-Kappa average.")
+
+        img[:, :, 0] = _sigma_worker(self.stack['R'], kappa, max_iters)
+        img[:, :, 1] = _sigma_worker(self.stack['G'], kappa, max_iters)
+        img[:, :, 2] = _sigma_worker(self.stack['B'], kappa, max_iters)
+
+        return Image(img=img, nprocs=self.nprocs)
+
+def _sigma_worker(data, kappa, max_iters):
+    '''Calculate kappa-sigma mean of the data.'''
+
+    shape = data.shape
+    tile_shape = (shape[-1], 1)
+    data_out = np.empty((shape[0], shape[1]), dtype=STACK_DTYPE) # data.dtype)
+
+    for i in range(shape[0]):
+        num_it = 0
+        row = np.ma.masked_equal(data[i, :, :], 0).T
+        while True:
+            avgs = np.mean(row, 0)
+            diffs = np.abs(np.tile(avgs, tile_shape) - row)
+            stds = np.std(row, 0)
+            idxs = diffs > kappa * np.tile(stds, tile_shape)
+            num_it += 1
+            if np.any(idxs):
+                row = np.ma.masked_where(idxs, row)
+            else:
+                break
+            if num_it >= max_iters:
+                break
+        data_out[i, :] = np.mean(row, 0)
+
+    return data_out