Made frame_by_frame mode of cube_fix_badpix_isolated function compati…

…ble with correct_only, when provided a bad pixel map

vip_hci/preproc/badpixremoval.py

@@ -321,6 +321,7 @@ def cube_fix_badpix_isolated(array, bpm_mask=None, correct_only=False,
                 else:
                     excl_mask_tmp = None
                 res = frame_fix_badpix_isolated(array[i], bpm_mask=bpm_mask_tmp,
+                                                correct_only=correct_only,
                                                 sigma_clip=sigma_clip,
                                                 num_neig=num_neig, size=size,
                                                 protect_mask=protect_mask,
@@ -334,7 +335,8 @@ def cube_fix_badpix_isolated(array, bpm_mask=None, correct_only=False,
         else:
             if verbose:
                 print("Cleaning frames using ADACS' multiprocessing approach")
-            # dummy calling the function to create cached version of the code prior to forking
+            # dummy calling the function to create cached version of the code
+            # prior to forking
             if bpm_mask is not None:
                 bpm_mask_dum = bpm_mask[0]
             else:
@@ -345,6 +347,7 @@ def cube_fix_badpix_isolated(array, bpm_mask=None, correct_only=False,
                 excl_mask_dum = None
             # point of dummy call
             frame_fix_badpix_isolated(array[0], bpm_mask=bpm_mask_dum,
+                                      correct_only=correct_only,
                                       sigma_clip=sigma_clip, num_neig=num_neig,
                                       size=size, protect_mask=protect_mask,
                                       excl_mask=excl_mask_dum, verbose=False,
@@ -373,6 +376,7 @@ def mp_clean_isolated(j, frame, bpm_mask=None, sigma_clip=3,
                                   excl_mask=None, verbose=False, cxy=None,
                                   ignore_nan=True, full_output=True):
                 sh_res = frame_fix_badpix_isolated(frame, bpm_mask,
+                                                   correct_only=correct_only,
                                                    sigma_clip=sigma_clip,
                                                    num_neig=num_neig, size=size,
                                                    protect_mask=protect_mask,
@@ -406,7 +410,8 @@ def _mp_clean_isolated(args):
                 dict_kwargs = {'bpm_mask': bpm_mask_tmp,
                                'sigma_clip': sigma_clip, 'num_neig': num_neig,
                                'size': size, 'protect_mask': protect_mask,
-                               'excl_mask': excl_mask_tmp, 'cxy': (cx[j], cy[j]),
+                               'excl_mask': excl_mask_tmp,
+                               'cxy': (cx[j], cy[j]),
                                'ignore_nan': ignore_nan}
                 args.append([j, array[j], dict_kwargs])
 
@@ -428,7 +433,7 @@ def _mp_clean_isolated(args):
             excl_mask = np.zeros(array.shape[-2:], dtype=bool)
         elif excl_mask.ndim == 3:
             excl_mask = np.median(excl_mask, axis=0)
-        elif excl_mask is None:
+        else:
             msg = "Input exclusion mask should have same last 2 dims as array\n"
             assert excl_mask.shape == array.shape[-2:], msg
         ind_excl = np.where(excl_mask)
@@ -437,13 +442,13 @@ def _mp_clean_isolated(args):
                 bpm_mask = np.zeros(array.shape[-2:], dtype=bool)
             elif bpm_mask.ndim == 3:
                 bpm_mask = np.median(bpm_mask, axis=0)
-            bpm_mask = bpm_mask+excl_mask
+            all_excl_mask = bpm_mask+excl_mask
             ori_nan_mask = np.where(np.isnan(np.nanmean(array, axis=0)))
             ind = clip_array(np.nanmean(array, axis=0), sigma_clip, sigma_clip,
-                             bpm_mask, neighbor=neigh, num_neighbor=num_neig,
-                             mad=mad)
-            final_bpm = np.zeros_like(array[0], dtype=bool)
-            final_bpm[ind] = 1
+                             all_excl_mask, neighbor=neigh,
+                             num_neighbor=num_neig, mad=mad)
+            final_bpm = bpm_mask.copy()
+            final_bpm[ind] += 1
             if ignore_nan:
                 final_bpm[ori_nan_mask] = 0
             if protect_mask:
@@ -485,8 +490,9 @@ def cube_fix_badpix_annuli(array, fwhm, cy=None, cx=None, sig=5., bpm_mask=None,
                            min_thr=None, max_thr=None, min_thr_np=None,
                            bad_values=None, full_output=False):
     """
-    Function to correct the bad pixels annulus per annulus (centered on the
-    provided location of the star), in an input frame or cube.
+    Correct bad pixels in concentric annuli centered on the provided location\
+    of the star, in an input frame or cube.
+
     This function is faster than ``cube_fix_badpix_clump``; hence to be
     preferred in all cases where there is only one bright source with circularly
     symmetric PSF. The bad pixel values are replaced by:
@@ -564,8 +570,8 @@ def cube_fix_badpix_annuli(array, fwhm, cy=None, cx=None, sig=5., bpm_mask=None,
         [full_output=True] The bad pixel map or the cube of bpix maps
     ann_frame_cumul: 2 or 3d array
         [full_output=True] The cube of defined annuli
-    """
 
+    """
     ndims = array.ndim
     assert ndims == 2 or ndims == 3, "Object is not two or three dimensional.\n"