diff --git a/zap/zap.py b/zap/zap.py index 363efd5..767e1af 100644 --- a/zap/zap.py +++ b/zap/zap.py @@ -81,7 +81,7 @@ def process(cubefits, outcubefits='DATACUBE_ZAP.fits', clean=True, zlevel='median', cftype='median', cfwidthSVD=300, cfwidthSP=300, nevals=[], extSVD=None, skycubefits=None, mask=None, interactive=False, ncpu=None, pca_class=None, n_components=None, - overwrite=False, varcurvefits=None): + overwrite=False, varcurvefits=None, zero_nan_slices=True): """ Performs the entire ZAP sky subtraction algorithm. This is the main ZAP function. It works on an input FITS file and @@ -141,7 +141,13 @@ def process(cubefits, outcubefits='DATACUBE_ZAP.fits', clean=True, to False. varcurvefits : str Path for the optional output of the explained variance curves. - + zero_nan_slices : bool + In cases where we use a fixed OUTWCS it is very easy that the data + end up having slices with only NaNs. These should be zero'd in the + same way as the AO gaps. If this is set to True, this will be done. It + goes through the whole cube - under normal conditions the NaN-only + layers should be at the start and end, but this way we can deal with + corner cases. """ logger.info('Running ZAP %s !', __version__) t0 = time() @@ -174,7 +180,8 @@ def _check_file_exists(filename): extSVD = SVDoutput(cubefits, clean=clean, zlevel=zlevel, cftype=cftype, cfwidth=cfwidthSVD, mask=mask) - zobj = Zap(cubefits, pca_class=pca_class, n_components=n_components) + zobj = Zap(cubefits, pca_class=pca_class, n_components=n_components, + zero_nan_slices=zero_nan_slices) zobj._run(clean=clean, zlevel=zlevel, cfwidth=cfwidthSP, cftype=cftype, nevals=nevals, extSVD=extSVD) @@ -336,10 +343,12 @@ class Zap(object): deconstructed stack zlsky : numpy.ndarray A 1d array containing the result of the zero level subtraction - + zeroed_nan_slices : list + If zero_nan_slices is True, this list records the slices where + NaNs have been set to zero """ - def __init__(self, cubefits, pca_class=None, n_components=None): + def __init__(self, cubefits, pca_class=None, n_components=None, zero_nan_slices=True): self.cubefits = cubefits self.ins_mode = None @@ -383,6 +392,20 @@ def __init__(self, cubefits, pca_class=None, n_components=None): else: self.notch_limits = None + # Zero other NaN slices if requested. + self.zeroed_nan_slices = [] + if zero_nan_slices: + logger.info('Zeroing slices with only NaNs') + nl, ny, nx = self.cube.shape + for i in range(nl): + # Check whether there are any finite values. If not, + # set the whole slice to zero + is_finite = np.where(np.isfinite(self.cube[i, :, :])) + if len(is_finite[0]) == 0: + self.cube[i, :, :] = 0.0 + self.zeroed_nan_slices.append(i) + + # NaN Cleaning self.run_clean = False self.nancube = None @@ -721,6 +744,9 @@ def make_cube_from_stack(self, stack, with_nans=False): if self.ins_mode in NOTCH_FILTER_RANGES: lmin, lmax = self.notch_limits cube[lmin:lmax + 1] = np.nan + if len(self.zeroed_nan_slices) > 0: + for i in self.zeroed_nan_slices: + cube[i, :, :] = np.nan return cube def remold(self):