if we are masking a cosmic using gradient we must mask both pixels

leading to large gradient
Add a test of cosmics

py/desispec/coaddition.py

@@ -487,12 +487,15 @@ def _mask_cosmics(wave, flux, ivar, mask, subset, ivarjj_masked,
         ttivar = ivar[j].copy()
         if nbad > 0 :
             ttivar[bad] = np.interp(wave[bad], wave[good], ttivar[good])
-        ttvar = 1./(ttivar + (ttivar == 0))
-        ttflux[1:] = ttflux[1:] - ttflux[:-1]
-        ttvar[1:] = ttvar[1:] + ttvar[:-1]
-        ttflux[0] = 0
-        grad.append(ttflux)
-        gradvar.append(ttvar)
+        # ttivar should not be equal to zero anywhere but just in case
+        # we still protect against it
+        ttvar = 1. / (ttivar + (ttivar == 0))
+
+        # these have one pixel less
+        cur_grad = ttflux[1:] - ttflux[:-1]
+        cur_grad_var = ttvar[1:] + ttvar[:-1]
+        grad.append(cur_grad)
+        gradvar.append(cur_grad_var)
     grad, gradvar = np.array(grad), np.array(gradvar)
     gradivar = (gradvar > 0 ) / np.array(gradvar + (gradvar == 0))
     nspec = grad.shape[0]
@@ -513,9 +516,17 @@ def _mask_cosmics(wave, flux, ivar, mask, subset, ivarjj_masked,
             log.info("masking {} values in {} spectra for targetid={}".format(n_cosmic, nspec,
                                                                               tid))
             badindex = np.where(cosmic_bad)[0]
+            # TODO this is likely not intended
+            # if we have more than cosmic in a the same pixel
+            # this will only mask it once
+            # Also the log is not correct then
             for bi in badindex  :
                 k = np.argmax(gradivar[:, bi] * deltagrad[:, bi]**2)
                 ivarjj_masked[k, bi] = 0.
+                ivarjj_masked[k, bi + 1] = 0
+                # since we are using the maximum value of grad^2
+                # we really cannot say which pixel is responsible for
+                # large gradient hence we must mask two pixels
                 log.debug("masking spec {} wave={}".format(k, wave[bi]))
     return 
 
@@ -630,9 +641,6 @@ def coadd(spectra, cosmics_nsig=None, onetile=False) :
             weights[:, bad] = ivarjj_orig[:, bad]
             # in the case of all masked pixels
             # we still use the variances ignoring masking
-            # TODO One flaw here is if the same pixel was masked
-            # as cosmic in all exposures, then it will be ignored in the result
-            # with no visible flag set
 
             tivar[i] = np.sum(weights, axis=0)
             weights = weights / (tivar[i] + (tivar[i] == 0))

py/desispec/test/test_coadd.py

@@ -10,7 +10,9 @@
 from desispec.spectra import Spectra
 from desispec.io import empty_fibermap
 from desispec.coaddition import (coadd, fast_resample_spectra,
-        spectroperf_resample_spectra, coadd_fibermap, coadd_cameras)
+                                 spectroperf_resample_spectra,
+                                 coadd_fibermap, coadd_cameras,
+                                 _mask_cosmics)
 from desispec.specscore import compute_coadd_scores
 from desispec.resolution import Resolution
 from desispec.maskbits import fibermask
@@ -110,7 +112,28 @@ def _random_spectra(self, ns=3, nw=10, seed=None, with_mask=False,
                 fibermap=fmap,
                 scores=scores,
                 )
-
+    def test_cosmic_masking(self):
+        rng = np.random.default_rng(133)
+        npix, nspec = 10000, 30
+        wave = np.arange(npix)
+        ivar = rng.uniform(0, 1,
+                           size=(nspec, npix)) * (1 + 100 * np.arange(nspec)[:, None])
+        model0 = 0.01 * (wave - (wave)**2 / npix) + rng.uniform(size=npix)
+        flux = model0 + rng.normal(size=(nspec, npix)) / np.sqrt(ivar)
+        COSMIC = 1e6
+        flux[2, 100] = COSMIC
+        flux[0, 130] = COSMIC
+        mask = np.zeros((nspec, npix), dtype=int)
+        ivarjj_masked = ivar * 1
+        cosmics_nsig = 4
+        _mask_cosmics(wave, flux, ivar, mask, np.arange(nspec),
+                      ivarjj_masked, tid=1,
+                      cosmics_nsig=cosmics_nsig)
+        # we mask pixel and 1 neighbor around hence 3 pixel per cosmic
+        self.assertEqual((ivarjj_masked == 0).sum(), 6)
+        self.assertEqual((flux[ivarjj_masked == 0] == COSMIC).sum(),
+                         2)
+
     def test_coadd(self):
         """Test coaddition"""
         nspec, nwave = 3, 10