Merge pull request #180 from desihub/template_fft_caching

Template fft caching

.github/workflows/python-package.yml

@@ -71,7 +71,10 @@ jobs:
             - name: Coveralls
               env:
                 GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+                COVERALLS_PARALLEL: true
+                COVERALLS_FLAG_NAME: "${{ matrix.os }},${{ matrix.python-version }},${{ matrix.package_level }}"
                 COVERALLS_SERVICE_NAME: github
+                COVERALLS_SERVICE_JOB_ID: "${{ github.run_id }}"
                 COVERALLS_SERVICE_NUMBER: "${{ github.workflow }}-${{ github.run_number }}"
               run: coveralls
 

doc/changes.rst

@@ -5,13 +5,15 @@ Change Log
 3.0.0 (not released yet)
 ------------------------
 
+* Cache Fourier transforms of the templates for speed [`PR #180`_].
 * Update line-list (drop broad HeI lines) and rest wavelengths [`PR #179`_].
 * Near-total rewrite of both the continuum and emission-line fitting engines and
   major updates to the organizational infrastructure of the code, all with the
   goal of maximizing speed and accuracy [`PR #177`_].
 
 .. _`PR #177`: https://github.com/desihub/fastspecfit/pull/177
 .. _`PR #179`: https://github.com/desihub/fastspecfit/pull/179
+.. _`PR #180`: https://github.com/desihub/fastspecfit/pull/180
 
 2.5.2 (2024-04-28)
 ------------------

py/fastspecfit/continuum.py

@@ -46,10 +46,12 @@ def __init__(self, data, templates, phot, igm, ebv_guess=0.05,
         # Optionally ignore templates which are older than the age of the
         # universe at the redshift of the object.
         if constrain_age:
-            agekeep = _younger_than_universe(templates.info['age'].value, data['tuniv'])
+            self.agekeep = self._younger_than_universe(templates.info['age'].value, data['tuniv'])
+            self.nage = len(self.agekeep)
         else:
-            agekeep = np.arange(templates.ntemplates)
-        self.agekeep = agekeep
+            # use default slice instead of arange to avoid copying templates
+            self.agekeep = slice(None, None)
+            self.nage = templates.ntemplates
 
         # Get preprocessing data to accelerate continuum_to_photometry()
         # but ONLY when it is called with the default filters=None
@@ -390,9 +392,7 @@ def templates2data(self, templateflux, templatewave, redshift=0., dluminosity=No
 
         # broaden for velocity dispersion
         if vdisp is not None:
-            vd_templateflux = Templates.convolve_vdisp(templateflux, vdisp,
-                                                       pixsize_kms=Templates.PIXKMS,
-                                                       pixkms_bounds=self.pixkms_bounds)
+            vd_templateflux = self.templates.convolve_vdisp(templateflux, vdisp)
         else:
             vd_templateflux = templateflux
 
@@ -556,7 +556,7 @@ def call_nnls(modelflux, flux, ivar, xparam=None, debug=False,
         if debug:
             if xivar > 0:
                 leg = r'${:.1f}\pm{:.1f}$'.format(xbest, 1. / np.sqrt(xivar))
-                #leg = r'${:.3f}\pm{:.3f}\ (\chi^2_{{min}}={:.2f})$'.format(xbest, 1/np.sqrt(xivar), chi2min)
+                #leg = r'${:.3f}\pm{:.3f}\ (\chi^2_{{min}}={:.2f})$'.format(xbest, 1./np.sqrt(xivar), chi2min)
             else:
                 leg = r'${:.3f}$'.format(xbest)
 
@@ -591,11 +591,14 @@ def call_nnls(modelflux, flux, ivar, xparam=None, debug=False,
         return chi2min, xbest, xivar, bestcoeff
 
 
-    # compute attenuated version of a model spectrum,
-    # including contribution of dust emission
     @staticmethod
     @numba.jit(nopython=True, fastmath=True, nogil=True)
     def attenuate(M, A, zfactors, wave, dustflux):
+        """
+        Compute attenuated version of a model spectrum,
+        including contribution of dust emission.
+
+        """
 
         # Concurrently replace M by M * (atten ** ebv) and
         # compute (by trapezoidal integration) integral of
@@ -627,19 +630,24 @@ def attenuate(M, A, zfactors, wave, dustflux):
         for i in range(len(M)):
             M[i] = (M[i] + lbol_diff * dustflux[i]) * zfactors[i]
 
-    # compute attenuated version of a model spectrum M,
-    # without dust emission
     @staticmethod
     @numba.jit(nopython=True, fastmath=True, nogil=True)
     def attenuate_nodust(M, A, zfactors):
+        """
+        Compute attenuated version of a model spectrum M,
+        without dust emission.
+
+        """
+
         # final result is
         # M * (atten ** ebv) * zfactors
         for i in range(len(M)):
             M[i] *= A[i] * zfactors[i]
 
 
     def build_stellar_continuum(self, templateflux, templatecoeff,
-                                ebv, vdisp=None, dust_emission=True):
+                                ebv, vdisp=None, conv_pre=None,
+                                dust_emission=True):
 
         """Build a stellar continuum model.
 
@@ -657,6 +665,9 @@ def build_stellar_continuum(self, templateflux, templatecoeff,
             Velocity dispersion in km/s. If `None`, do not convolve to the
             specified velocity dispersion (usually because `templateflux` has
             already been smoothed to some nominal value).
+        conv_pre: :class:`tuple` or None
+            Optional preprocessing data to accelerate template convolution with vdisp
+            (may be present only if vdisp is not None).
         dust_emission : :class:`bool`
             Model impact of infrared dust emission spectrum. Energy-balance is used
             to compute the normalization of this spectrum.
@@ -667,14 +678,34 @@ def build_stellar_continuum(self, templateflux, templatecoeff,
             Full-wavelength, native-resolution, observed-frame model spectrum.
 
         """
-        # [1] - Compute the weighted sum of the templates.
-        contmodel = templateflux.dot(templatecoeff)
+        if conv_pre is None or vdisp > Templates.MAX_PRE_VDISP:
+            # [1] - Compute the weighted sum of the templates.
+            contmodel = templateflux.dot(templatecoeff)
 
-        # [2] - Optionally convolve to the desired velocity dispersion.
-        if vdisp is not None:
-            contmodel = Templates.convolve_vdisp(contmodel, vdisp,
-                                                 pixsize_kms=Templates.PIXKMS,
-                                                 pixkms_bounds=self.pixkms_bounds)
+            # [2] - Optionally convolve to the desired velocity dispersion.
+            if vdisp is not None:
+                contmodel = self.templates.convolve_vdisp(contmodel, vdisp)
+        else:
+            # if conv_pre is present, it contains flux values for non-convolved
+            # regions of template fluxes, plus FTs of tempaltes for convolved
+            # region.  Both must be combined using template coefficients.
+            flux_lohi, ft_flux_mid, fft_len = conv_pre
+
+            # [1] - Compute the weighted sum of the templates.
+            cont_lohi   = flux_lohi.dot(templatecoeff)
+            ft_cont_mid = ft_flux_mid.dot(templatecoeff)
+
+            # [2] - convolve to the desired velocity dispersion.
+            # Use the vdisp convolution that takes precomputed FT
+            # of flux for convolved region
+            flux_len = templateflux.shape[0]
+            contmodel = self.templates.convolve_vdisp_from_pre(
+                cont_lohi, ft_cont_mid, flux_len, fft_len, vdisp)
+
+            # sanity check for debugging
+            #contmodel0 = templateflux.dot(templatecoeff)
+            #contmodel0 = self.templates.convolve_vdisp(contmodel0, vdisp)
+            #print("DIFF ", np.max(np.abs(contmodel - contmodel0)))
 
         # [3] - Apply dust attenuation; ToDo: allow age-dependent
         # attenuation. Also compute the bolometric luminosity before and after
@@ -687,9 +718,8 @@ def build_stellar_continuum(self, templateflux, templatecoeff,
 
         # [5] - Redshift factors.
 
-        # do this part in Numpy because it is very slow
-        # in Numba unless accelerated transcendentals are
-        # available via, e.g., Intel SVML.
+        # Do this part in Numpy because it is very slow in Numba unless
+        # accelerated transcendentals are available via, e.g., Intel SVML.
         A = self.lg_atten * ebv
         np.exp(A, out=A)
 
@@ -787,7 +817,7 @@ def continuum_to_photometry(self, contmodel,
 
 
     def _stellar_objective(self, params, templateflux,
-                           dust_emission, fit_vdisp,
+                           dust_emission, fit_vdisp, conv_pre,
                            objflam, objflamistd,
                            specflux, specistd,
                            synthphot, synthspec):
@@ -808,7 +838,8 @@ def _stellar_objective(self, params, templateflux,
         # allocating new
         fullmodel = self.build_stellar_continuum(
             templateflux, templatecoeff,
-            ebv=ebv, vdisp=vdisp, dust_emission=dust_emission)
+            ebv=ebv, vdisp=vdisp, conv_pre=conv_pre,
+            dust_emission=dust_emission)
 
         # save the full model each time we compute the objective;
         # after optimization, the final full model will be
@@ -843,7 +874,7 @@ def _stellar_objective(self, params, templateflux,
         return resid
 
 
-    def fit_stellar_continuum(self, templateflux, fit_vdisp,
+    def fit_stellar_continuum(self, templateflux, fit_vdisp, conv_pre=None,
                               vdisp_guess=250., ebv_guess=0.05,
                               coeff_guess=None,
                               vdisp_bounds=(75., 500.), ebv_bounds=(0., 3.),
@@ -860,6 +891,9 @@ def fit_stellar_continuum(self, templateflux, fit_vdisp,
         fit_vdisp : :class:`bool`
             If true, solve for the velocity dispersion;
             if false, use a nominal dispersion.
+        conv_pre : :class:`tuple` of None
+            If not None, preprocessing data for convolving templateflux
+            with vdisp values.  (Occurs only if fit_vdisp is True.)
         vdisp_guess : :class:`float`
             Guess for scalar value of the velocity dispersion if fitting.
         ebv_guess : :class:`float`
@@ -922,6 +956,7 @@ def fit_stellar_continuum(self, templateflux, fit_vdisp,
             'templateflux':  templateflux,
             'dust_emission': dust_emission,
             'fit_vdisp':     fit_vdisp,
+            'conv_pre':      conv_pre,
             'objflam':       objflam,
             'objflamistd':   objflamistd,
             'specflux':      specflux,
@@ -1057,7 +1092,7 @@ def continuum_fastphot(redshift, objflam, objflamivar, CTools,
     data = CTools.data
     templates = CTools.templates
     agekeep = CTools.agekeep
-    nage = len(agekeep)
+    nage = CTools.nage
 
     ebv = 0.
     ebvivar = 0.
@@ -1154,7 +1189,7 @@ def _continuum_fastspec_legacy(redshift, specwave, specflux, specivar,
     templates = CTools.templates
     phot = CTools.phot
     agekeep = CTools.agekeep
-    nage = len(agekeep)
+    nage = CTools.nage
 
     vdisp_nominal = templates.vdisp_nominal
     ndof_cont = np.sum(specivar > 0.)
@@ -1333,7 +1368,7 @@ def continuum_fastspec(redshift, objflam, objflamivar, CTools,
     phot = CTools.phot
     templates = CTools.templates
     agekeep = CTools.agekeep
-    nage = len(agekeep)
+    nage = CTools.nage
 
     # Combine all three cameras; we will unpack them to build the
     # best-fitting model (per-camera) below.
@@ -1423,17 +1458,21 @@ def continuum_fastspec(redshift, objflam, objflamivar, CTools,
 
         if compute_vdisp:
             input_templateflux = templates.flux[:, agekeep]
+            input_conv_pre = templates.conv_pre_select(templates.conv_pre, agekeep)
             input_templateflux_nolines = templates.flux_nolines[:, agekeep]
+            input_conv_pre_nolines = templates.conv_pre_select(templates.conv_pre_nolines, agekeep)
         else:
             # Use the cached templates with nominal velocity dispersion
             input_templateflux = templates.flux_nomvdisp[:, agekeep]
+            input_conv_pre = None
             input_templateflux_nolines = templates.flux_nolines_nomvdisp[:, agekeep]
+            input_conv_pre_nolines = None
             log.info('Insufficient wavelength coverage to compute velocity dispersion.')
 
         t0 = time.time()
         ebv, vdisp, coeff, resid = CTools.fit_stellar_continuum(
             input_templateflux, # [npix,nage]
-            fit_vdisp=compute_vdisp,
+            fit_vdisp=compute_vdisp, conv_pre=input_conv_pre,
             vdisp_guess=templates.vdisp_nominal,
             #ebv_guess=ebv, coeff_guess=coeff_guess, # don't bias the answer...?
             objflam=objflam, objflamistd=objflamistd,
@@ -1480,7 +1519,8 @@ def continuum_fastspec(redshift, objflam, objflamivar, CTools,
             sedmodel = CTools.optimizer_saved_contmodel
             sedmodel_nolines = CTools.build_stellar_continuum(
                 input_templateflux_nolines, coeff, ebv=ebv,
-                vdisp=(vdisp if compute_vdisp else None), dust_emission=False)
+                vdisp=(vdisp if compute_vdisp else None),
+                conv_pre=input_conv_pre_nolines, dust_emission=False)
 
     desimodel_nolines = CTools.continuum_to_spectroscopy(sedmodel_nolines)
 
@@ -1653,8 +1693,8 @@ def continuum_specfit(data, result, templates, igm, phot,
         logmstar = np.log10(CTools.massnorm * masstot)
         zzsun = np.log10(coeff.dot(mstars * 10.**tinfo['zzsun']) / masstot) # mass-weighted
         age = coeff.dot(tinfo['age']) / coefftot / 1e9           # luminosity-weighted [Gyr]
-        #age = coeff.dot(mstars * tinfo['age']) / masstot / 1e9        # mass-weighted [Gyr]
-        sfr = CTools.massnorm * coeff.dot(tinfo['sfr'])                          # [Msun/yr]
+        #age = coeff.dot(mstars * tinfo['age']) / masstot / 1e9  # mass-weighted [Gyr]
+        sfr = CTools.massnorm * coeff.dot(tinfo['sfr'])          # [Msun/yr]
         if templates.use_legacy_fitting:
             AV = coeff.dot(tinfo['av']) / coefftot # luminosity-weighted [mag]
         else:

py/fastspecfit/qa.py

@@ -398,7 +398,7 @@ def major_formatter(x, pos):
         else:
             contmodel = CTools.build_stellar_continuum(
                 templates.flux_nolines, fastspec['COEFF'],
-                vdisp=fastspec['VDISP'], 
+                vdisp=fastspec['VDISP'], conv_pre=templates.conv_pre_nolines, 
                 ebv=fastspec['AV'] / Templates.klambda(5500.)
             )