From 00e8f6cd34814a6eab91993bb94a4fe2e4cfae67 Mon Sep 17 00:00:00 2001
From: Niv Drory <drory@astro.as.utexas.edu>
Date: Tue, 24 Sep 2024 11:39:32 -0500
Subject: [PATCH] vectorize construction of projection matrix in optimal
 extraction.

---
 python/lvmdrp/core/spectrum1d.py       | 37 ++++++++++++++++----------
 python/lvmdrp/functions/imageMethod.py |  6 ++---
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/python/lvmdrp/core/spectrum1d.py b/python/lvmdrp/core/spectrum1d.py
index ffea6454..8af3675c 100644
--- a/python/lvmdrp/core/spectrum1d.py
+++ b/python/lvmdrp/core/spectrum1d.py
@@ -15,7 +15,6 @@
 from lvmdrp.core import plot
 from lvmdrp.core.header import Header
 
-
 def adaptive_smooth(data, start_width, end_width):
     """
     Smooth an array with a filter that adapts in size from start_width to end_width.
@@ -3462,17 +3461,28 @@ def obtainGaussFluxPeaks(self, pos, sigma, replace_error=1e10, plot=False):
         # construct sparse projection matrix
         fact = numpy.sqrt(2.0 * numpy.pi)
         kernel_width = 7 # should exceed 4 sigma
-        vI = []
-        vJ = []
-        vV = []
-        for xx in range(nfibers):
-            for yy in range(int(pos[xx]-kernel_width),int(pos[xx]+kernel_width)):
-                v = numpy.exp(-0.5 * ((yy-pos[xx]) / sigma[xx]) ** 2) / (fact * sigma[xx])
-                if v>0.0001:   # make non-zero and positive definite
-                    vI.append(xx)
-                    vJ.append(yy)
-                    vV.append(v / self._error[yy])
-        B = sparse.csc_matrix((vV, (vJ, vI)), shape=(self._dim, nfibers))
+        # vI = []
+        # vJ = []
+        # vV = []
+        # for xx in range(nfibers):
+        #     for yy in range(int(pos[xx]-kernel_width),int(pos[xx]+kernel_width)+1):
+        #         v = numpy.exp(-0.5 * ((yy-pos[xx]) / sigma[xx]) ** 2) / (fact * sigma[xx])
+        #         if v>=0.0000:   # make non-zero and positive definite
+        #             vI.append(xx)
+        #             vJ.append(yy)
+        #             vV.append(v / self._error[yy])
+        # B = sparse.csc_matrix((vV, (vJ, vI)), shape=(len(self._data), nfibers))
+
+        # nfibers x kernel_size
+        xx = numpy.repeat(numpy.array(range(nfibers)), 2*kernel_width+1)
+        # pixel ranges of fiber images
+        pos_t = numpy.trunc(pos)
+        yyv = numpy.linspace(pos_t-kernel_width, pos_t+kernel_width, 2*kernel_width+1, endpoint=True)
+        # nfibers x kernel_size pixel values
+        v = numpy.exp(-0.5 * ((yyv-pos) / sigma) ** 2) / (fact * sigma)
+        yyv = yyv.T.flatten()
+        v = v.T.flatten() / self._error[yyv.astype(numpy.int32)]
+        B = sparse.csc_matrix((v, (yyv, xx)), shape=(len(self._data), nfibers))
 
         # invert the projection matrix and solve
         ypixels = numpy.arange(self._data.size)
@@ -3485,9 +3495,8 @@ def obtainGaussFluxPeaks(self, pos, sigma, replace_error=1e10, plot=False):
         if bad_pix is not None and bn.nansum(bad_pix) > 0:
             error[bad_pix] = replace_error
 
-        # pyfits.writeto('B1.fits', B1.toarray(), overwrite=True)
+        # pyfits.writeto('B.fits', B.toarray(), overwrite=True)    
         # if plot:
-        #     plt.figure(figsize=(15, 10))
         #     plt.plot(self._data, "ok")
         #     plt.plot(numpy.dot(A * self._error[:, None], out[0]), "-r")
         #     # plt.plot(numpy.dot(A, out[0]), '-r')
diff --git a/python/lvmdrp/functions/imageMethod.py b/python/lvmdrp/functions/imageMethod.py
index 7e8b899a..59f76962 100644
--- a/python/lvmdrp/functions/imageMethod.py
+++ b/python/lvmdrp/functions/imageMethod.py
@@ -47,6 +47,7 @@
 from lvmdrp.core.tracemask import TraceMask
 from lvmdrp.utils.hdrfix import apply_hdrfix
 from lvmdrp.utils.convert import dateobs_to_sjd, correct_sjd
+from lvmdrp.utils.timer import Timer
 
 
 NQUADS = 4
@@ -2672,9 +2673,8 @@ def extract_spectra(
             else:
                 mask = None
         else:
-            (data, error, mask) = img.extractSpecOptimal(
-                trace_mask, trace_fwhm, plot_fig=display_plots
-            )
+            with Timer(name="extract optimal", logger=log.info):
+                (data, error, mask) = img.extractSpecOptimal(trace_mask, trace_fwhm, plot_fig=display_plots)
     elif method == "aperture":
         trace_fwhm = None