Fix spectra coaddition and cosmic masking

[segasai]

This is an a actual PR fixing #2372

[coveralls]

coverage: 30.293% (+0.1%) from 30.164%
when pulling f3bc12b on resol_fix
into 48c1641 on main.

[sbailey]

Thanks. I confirm that the original code is a bug, i.e. we agree on the math of what should be done and that the current code does not do that. Since this implementation is a little different than the #2372 I'd like to study it a bit more and run some test tiles to check for corner cases (e.g. entirely masked spectra).

Ideally I'd like to also add some unit tests that would have caught the original bug, e.g. that a "coaddition" of a single resolution matrix should be the same as its input, and verifying what we expect the impact to be of a masked pixel on 1 out of N exposures, and the same pixel masked on N out of N exposures.

In issue #2372 you documented the impact on chi2; I'd also like to study the impact on final redshifts to decide what to do about Kibo (e.g. a coadds+redshifts-only supplemental run). But that doesn't have to be a blocking factor for wrapping this up.

[segasai]

Thanks for looking into that.

As far as I tested, this implementation should give the same result as the previous one. I just felt it was a bit clearer. but for sure it's worth double checking.
I agree this bug also worth a regression test.
The simplest one is indeed with one masked pixel and one exposure.

FWIW these are my runs https://data.desi.lbl.gov/desi/users/koposov/coadd_test/output/ on random hpx in kibo using this patch. So the redshfits can be directly compared to kibo. (this used earlier version of the patch, but AFAICS it should be exactly equal to what was before).

[segasai]

I've added the test that trips the current desispec but passes after the PR. (coaddition of one spectrum with one masked pixel)

[segasai]

I've added a final (IMO) test that verifies if one has 'true 'spectrum M, Observations with 'random' resolution matrices R_i, random covariance matrices and extracted spectra S_i (obeying S_i= R_i M), then after running coadd the relationship
S_coadded = R_coadded M still holds exactly.
This test also obviously fails with existing desispec.
With this test I am actually very confident that the new code is correct.

[sbailey]

Thanks, this all looks good. In the spirit of not making a major change just before the weekend, I will wait until Monday to merge this, but I don't see any problems. In the meantime, two more tests to request to really drive this home:

• coaddition with one of the exposures entirely masked
• coaddition with a single edge pixel masked (instead of an interior pixel), i.e. test_coadd_single_masked with mpix=0.

I'm not anticipating problems with those.

Let's discuss at the Tuesday data telecon about implications for Kibo. I'm not worried from a quick spot check, but I'd like to quantify that and study some rrdetails chi2 vs. z curves.

[segasai]

I've added the tests. Actually one of the tests did uncover the issue with the code. See b2f71bb

While adding those, I thought I'd add a test of cosmic ray masking, but an intuitive test that I wrote completely fails currently. I.e. here I put a large cosmic in i-th pixel in i-th spectrum.
And I was expecting that after stacking I would get the same result as if I had one less spectrum (with all those 'cosmic-pixels removed), but this fails and the cosmic masking fails to identify the cosmics. I don't quite understand the logic of cosmic masking in coadd(), so I don't know if I'm missing something obvious or not. The cosmic test is below:

    def test_coadd_cosmic(self):
        """                                                                                                                                                               
        Test coadd with cosmic rays                                                                                                                                       
        """
        nspec, nwave = 1000, 30
        s1 = self._random_spectra(nspec, nwave)
        rng = np.random.default_rng(432)
        model0 = rng.uniform(0, 1, size=nwave)
        ivar0 = rng.uniform(size=(nspec, nwave))
        flux0 = model0 + 1./np.sqrt(ivar0) * rng.normal(size=(nspec,nwave))
        resol0 = rng.uniform(0, 1, size=s1.resolution_data['b'].shape)
        s1.flux['b'] = flux0 * 1
        s1.ivar['b'] = ivar0 * 1
        s1.resolution_data['b'] = resol0 * 1
        COSMIC_VAL = 1e9
        xgrid,ygrid =  np.meshgrid(np.arange(nwave), np.arange(nspec))
        mask = ygrid == (xgrid%nwave)
        # make i-th pixel in i-th spectrum a cosmic                                                                                                                       
        s1.flux['b'][mask] = COSMIC_VAL
        s1.fibermap['TARGETID'] = [10] * nspec
        coadd(s1, cosmics_nsig=5)

        # prepare nspec-1 spectra with those cosmic affected pixels excised.
        s2 = self._random_spectra(nspec-1, nwave)
        for i in range(nwave):
            s2.flux['b'][:, i] = flux0[~mask[:,i],i]*1
            s2.ivar['b'][:, i] = ivar0[~mask[:,i],i]*1
            s2.resolution_data['b'][:,:,i] = resol0[~mask[:,i],:,i]*1
        coadd(s2, cosmics_nsig=5)
        print(s1.flux['b'],s2.flux['b'])
        self.assertTrue(np.allclose(s1.flux['b'], s2.flux['b']))
        self.assertTrue(np.allclose(s1.ivar['b'], s2.ivar['b']))
        self.assertTrue(np.allclose(s1.resolution_data['b'],
                                    s2.resolution_data['b']))

[segasai]

Just a few more comments after looking at the code.

• the coadd_cameras has the same issue as coadd() (it can be seen with the existing test, if I replace coadd() by coadd_cameras(). Since the logic there is a bit more complicated, the fix requires more thought. But also coadd_cameras is not really used for much I think, so it's less important.
• Another thing that is a bit strange is the current behaviour when one pixel in all the spectra in the stack is masked, because the current code will then still compute the ivar, ignoring the mask

  desispec/py/desispec/coaddition.py

  Line 571 in 48c1641

  tivar[i][bad] = np.sum(spectra.ivar[b][jj][:,bad],axis=0) # if all masked, keep original ivar

  One situation where that will produce strange results is when the same pixel was masked by cosmic logic in all exposures, then the stack will just ignore that (and it won't be reflected in the mask or ivar). I'm not quite sure it's intended behaviour (but it is also probably very rare occurrence). Also the logic here has implications for what's done with the resolution matrix for these pixels. With the new code the resolution matrix elements that have all zero weights in every exposure will be zero.

[moustakas]

coadd_cameras is used extensively by fastspecfit to write out coadeed spectra and models and also by the GQP working group. @stephjuneau

[weaverba137]

And @adambolton I think did a full coadd across cameras for iron to load into SPARCL. This may need to be reconsidered.

[stephjuneau]

Hi @segasai: confirming what @moustakas and @weaverba137 mentioned above that coadd_across_cameras is very important for GQP and beyond. Any change/streamlining should at least be backward compatible to avoid breaking workflows. Thanks!

[segasai]

@stephjuneau I wasn't talking about breaking any interfaces here. I was just pointing that the function seem to suffer from the same issue that is topic of this PR. But indeed I wasn't aware that this f-n is widely used outside the main pipeline.

[sbailey]

My latest commits include two tests for issues that @segasai identified that we need to fix:

• if a pixel is masked in every exposure, that should be propagated to the coadd even if it is masked for different reasons in each exposure (already fixed; @segasai coded the fix faster than I could code the test)
• if two exposures have unlucky cosmics at the same wavelength, they should both be masked as long as there are N>2 additional exposures that can identify that it really is a cosmic outlier and not a narrow emission line (current code just masks the first one, not the second)

[segasai]

I've committed the masker that can deal with more than one cosmic per pixel. Currently it's doing essentially iterative sigma clipping. That passes the tests, but my only worry are variable objects. With the new code we continue clipping till the chi^2 becomes low enough or we have less than 3 pixels in the stack.
One could potentially stop earlier and not do more than 2 iterations of clipping.
Other than that, I think we dealt with every outstanding issue here.

UPDATE on variable objects -- looking at sv1/dark/5059 RRL with targetid=39628511287183022 with 30 observations:
INFO:coaddition.py:571:_mask_cosmics: masking 47 wavelengths in 30 spectra for targetid=39628511287183022
INFO:coaddition.py:574:_mask_cosmics: masking 6 wavelengths with more than 1 mask per pixel for targetid=39628511287183022
INFO:coaddition.py:571:_mask_cosmics: masking 38 wavelengths in 30 spectra for targetid=39628511287183022
INFO:coaddition.py:574:_mask_cosmics: masking 8 wavelengths with more than 1 mask per pixel for targetid=39628511287183022
INFO:coaddition.py:574:_mask_cosmics: masking 14 wavelengths with more than 1 mask per pixel for targetid=39628511287183022
It doesn't seem too bad, so maybe things can be kept as they are now.

[segasai]

The impact of the cosmic ray fixes is particularly strong on the spectra with sharp gradients.
Here's comparison of old and new stacks. The old stacks tended to mark some of the edges as cosmics:

I assume the same thing is present with narrow emission lines

[segasai]

There was one more bug in the original implementation (and the new one) is that
if there was ever a fully blank spectrum in the stack with ivar=0, then the cosmic masking of any spectrum after that was incorrect, because of this skipping:

desispec/py/desispec/coaddition.py

Line 523 in 4eae355

if np.sum(good)==0 :

When spectra are collected for masking, the ivar=0 ones are skipped, but when doing the masking, that 'off-by-N'

desispec/py/desispec/coaddition.py

Line 561 in 4eae355

k=np.argmax(gradivar[:,bi]*deltagrad[:,bi]**2)

is ignored leading to masking of wrong pixels.
Encountered when running on sv1/dark/5059
The last commit fixes that and adds a test.

[segasai]

Pixel differences for a Lya-forest object, showing impact of the cosmic masking there.
In basically all the pixels with non-zero difference bertween new/old the masking there was happening in the previous iteration of the code and does not happen now.

[sbailey]

Looks good, thanks! So that we don't make major changes going into a weekend, I'll wait until Monday to merge and deploy this.