Fluxcalib variance fix

py/desispec/fluxcalibration.py

@@ -1459,23 +1459,27 @@ def apply_flux_calibration(frame, fluxcalib):
 
     """
     F'=F/C
-    Var(F') = Var(F)/C**2 + F**2*(  d(1/C)/dC )**2*Var(C)
-    = 1/(ivar(F)*C**2) + F**2*(1/C**2)**2*Var(C)
-    = 1/(ivar(F)*C**2) + F**2*Var(C)/C**4
-    = 1/(ivar(F)*C**2) + F**2/(ivar(C)*C**4)
+    Var(F') = Var(F)/C^2 + F^2*(  d(1/C)/dC )^2*Var(C)
+    = 1/(ivar(F)*C^2) + F^2*(1/C^2)^2*Var(C)
+    = 1/(ivar(F)*C^2) + F^2*Var(C)/C^4
+    = 1/(ivar(F)*C^2) + F^2/(ivar(C)*C^4)
+    ivar(F') = C^4 * ivar(F) * ivar(C)/ (C^2 * ivar(C) + F^2 * ivar(F))
     """
-    # for fiber in range(nfibers) :
-    #     C = fluxcalib.calib[fiber]
-    #     flux[fiber]=frame.flux[fiber]*(C>0)/(C+(C==0))
-    #     ivar[fiber]=(ivar[fiber]>0)*(civar[fiber]>0)*(C>0)/(   1./((ivar[fiber]+(ivar[fiber]==0))*(C**2+(C==0))) + flux[fiber]**2/(civar[fiber]*C**4+(civar[fiber]*(C==0)))   )
 
     C = fluxcalib.calib
-    frame.flux = frame.flux * (C>0) / (C+(C==0))
-    frame.ivar *= (fluxcalib.ivar>0) * (C>0)
+    good = (fluxcalib.ivar > 0) & (C > 0) & (frame.ivar > 0)
     for i in range(nfibers) :
-        ok=np.where(frame.ivar[i]>0)[0]
-        if ok.size>0 :
-            frame.ivar[i,ok] = 1./( 1./(frame.ivar[i,ok]*C[i,ok]**2)+frame.flux[i,ok]**2/(fluxcalib.ivar[i,ok]*C[i,ok]**4)  )
+        ok = good[i]
+        if ok.any():
+            frame.ivar[i, ok] = (C[i, ok]**4 * frame.ivar[i, ok] *
+                                 fluxcalib.ivar[i, ok] /
+                                 (C[i, ok]**2 * fluxcalib.ivar[i, ok] +
+                                  frame.flux[i, ok]**2 * frame.ivar[i, ok]
+                                  ))
+        frame.ivar[i, ~ok] = 0
+    # It is important we update flux *after*
+    # updating variance
+    frame.flux = frame.flux * (C > 0) / (C + (C == 0))
 
     if fluxcalib.fibercorr is not None and frame.fibermap is not None :
         if "PSF_TO_FIBER_FLUX" in fluxcalib.fibercorr.dtype.names :

py/desispec/test/test_flux_calibration.py

@@ -136,21 +136,31 @@ def test_apply_fluxcalibration(self):
         nwave = len(wave)
         nspec = 3
         flux = np.random.uniform(0.9, 1.0, size=(nspec, nwave))
-        ivar = np.ones_like(flux)
-        origframe = Frame(wave, flux, ivar, spectrograph=0)
+        ivar = np.random.uniform(0.9, 1.1, size=flux.shape)
+        origframe = Frame(wave, flux.copy(), ivar.copy(), spectrograph=0)
 
-        #define fluxcalib object
-        calib = np.ones_like(origframe.flux)
+        # efine fluxcalib object
+        calib = np.random.uniform(.5, 1.5, size=origframe.flux.shape)
         mask = np.zeros(origframe.flux.shape, dtype=np.uint32)
-        calib[0] *= 0.5
-        calib[1] *= 1.5
+
+        ivar_big = 1e20 * np.ones_like(origframe.flux)
 
         # fc with essentially no error
-        fcivar = 1e20 * np.ones_like(origframe.flux)
-        fc = FluxCalib(origframe.wave, calib, fcivar,mask)
+        fc = FluxCalib(origframe.wave, calib, ivar_big, mask)
+        frame = copy.deepcopy(origframe)
+        apply_flux_calibration(frame, fc)
+        self.assertTrue(np.allclose(frame.ivar, calib**2 * ivar))
+
+        # spectrum with essentially no error
+        # but large calibration error
+        # in this case the S/N should be the same as of the
+        # calibration vector
+        fc = FluxCalib(origframe.wave, calib, ivar, mask)
         frame = copy.deepcopy(origframe)
+        frame.ivar = ivar_big
         apply_flux_calibration(frame, fc)
-        self.assertTrue(np.allclose(frame.ivar, calib**2))
+        self.assertTrue(np.allclose(frame.flux**2 * frame.ivar,
+                                    calib**2 * ivar))
 
         # origframe.flux=0 should result in frame.flux=0
         fcivar = np.ones_like(origframe.flux)