clarify the code

desihub · Sep 25, 2024 · a3abe0c · a3abe0c
1 parent d911f70
commit a3abe0c
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/py/desispec/coaddition.py b/py/desispec/coaddition.py
@@ -567,13 +567,15 @@ def coadd(spectra, cosmics_nsig=None, onetile=False) :
             ww = spectra.resolution_data[b].shape[1]//2
             # resolution kernel width
             npix = spectra.resolution_data[b].shape[2]
-            for r in range(spectra.resolution_data[b].shape[1]) :
-                l1, l2 = max(r - ww, 0) , min(npix + (r - ww), npix)
-                l1x, l2x = l1 - (r - ww), l2 - (r - ww)
-                norm = np.sum(spectra.ivar[b][jj, l1:l2] ,axis=0)
-                trdata[i, r, l1x:l2x] = np.sum(spectra.ivar[b][jj, l1:l2] *
-                                    spectra.resolution_data[b][jj, r, l1x:l2x],
-                                               axis=0) / (norm + (norm == 0))
+            # indices of the corresponding variance point
+            # that needs to be used for ivar weights
+            res_indices = (np.arange(npix)[None, :]
+                           + np.arange(-ww, ww+1)[:, None]) % npix
+            res_whts = np.array([_[res_indices] for _ in spectra.ivar[b][jj]])
+            res_norm = res_whts.sum(axis=0)
+            trdata[i, :, : ] = np.sum(res_whts *
+                                    spectra.resolution_data[b][jj, : , : ],
+                                               axis=0) / (res_norm + (res_norm == 0))
             bad=(tivar[i]==0)
             if np.sum(bad)>0 :
                 tivar[i][bad] = np.sum(spectra.ivar[b][jj][:,bad],axis=0) # if all masked, keep original ivar