add numba jit cache=True

desihub · Nov 8, 2024 · 59fd3a4 · 59fd3a4
1 parent c3b4047
commit 59fd3a4
Show file tree

Hide file tree

Showing 10 changed files with 34 additions and 34 deletions.
diff --git a/py/fastspecfit/continuum.py b/py/fastspecfit/continuum.py
@@ -480,7 +480,7 @@ def _get_cflux(cwave, linear_fit=False, siglo=2., sighi=2.,
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True, fastmath=True)
+    @jit(nopython=True, nogil=True, fastmath=True, cache=True)
     def attenuate(M, A, zfactors, wave, dustflux):
         """
         Compute attenuated version of a model spectrum,
@@ -519,7 +519,7 @@ def attenuate(M, A, zfactors, wave, dustflux):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True, fastmath=True)
+    @jit(nopython=True, nogil=True, fastmath=True, cache=True)
     def attenuate_nodust(M, A, zfactors):
         """
         Compute attenuated version of a model spectrum M,

diff --git a/py/fastspecfit/emline_fit/interface.py b/py/fastspecfit/emline_fit/interface.py
@@ -320,7 +320,7 @@ def __init__(self,
 
         """
 
-        @jit(nopython=True, nogil=True)
+        @jit(nopython=True, nogil=True, cache=True)
         def _suppress_negative_fluxes(endpts, M):
             """
             suppress negative fluxes arising from resolution matrix
@@ -438,7 +438,7 @@ def find_peak_amplitudes(line_parameters,
 
     """
 
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _update_line_maxima(max_amps, line_models):
         """
         Given an array of line waveforms and an array of maximum
@@ -592,7 +592,7 @@ def _build_multimodel_core(line_parameters,
         consumer_fun((endpts, M))
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _add_patches(obs_bin_centers,
                  model_fluxes,
                  patch_endpts,
@@ -639,7 +639,7 @@ def _add_patches(obs_bin_centers,
 ###################################################################
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def mulWMJ(w, M, Jsp):
     """Compute the sparse matrix product P = WMJ, where
     W is a diagonal weight matrix
@@ -717,7 +717,7 @@ def mulWMJ(w, M, Jsp):
     return (endpts, J)
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _prepare_bins(centers, camerapix):
     """
     Convert bin centers to the info needed by the optimizer,

diff --git a/py/fastspecfit/emline_fit/jacobian.py b/py/fastspecfit/emline_fit/jacobian.py
@@ -12,7 +12,7 @@
 )
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def emline_model_jacobian(line_parameters,
                           log_obs_bin_edges,
                           ibin_widths,
@@ -181,7 +181,7 @@ def emline_model_jacobian(line_parameters,
 
 
 @staticmethod
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def patch_jacobian(obs_bin_centers,
                    obs_weights,
                    patch_endpts,

diff --git a/py/fastspecfit/emline_fit/model.py b/py/fastspecfit/emline_fit/model.py
@@ -11,7 +11,7 @@
 )
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def emline_model(line_wavelengths,
                  line_parameters,
                  log_obs_bin_edges,
@@ -77,7 +77,7 @@ def emline_model(line_wavelengths,
     return model_fluxes
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def emline_perline_models(line_wavelengths,
                           line_parameters,
                           log_obs_bin_edges,
@@ -157,7 +157,7 @@ def emline_perline_models(line_wavelengths,
     return (endpts, line_profiles)
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def emline_model_core(line_wavelength,
                       line_amplitude,
                       line_vshift,

diff --git a/py/fastspecfit/emline_fit/params_mapping.py b/py/fastspecfit/emline_fit/params_mapping.py
@@ -93,7 +93,7 @@ def mapFreeToFull(self, freeParms, out=None, patchDoublets=True):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _mapFreeToFull(freeParms, nParms, sources, factors,
                        doubletPatches, fullParms, patchDoublets):
 
@@ -137,7 +137,7 @@ def getJacobian(self, freeParms):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _matvec(J_S, v, w):
         """
         Multiply sparse parameter Jacobian J_S * v, writing result to w.
@@ -161,7 +161,7 @@ def _matvec(J_S, v, w):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _add_rmatvec(J_S, v, w):
         """
         Multiply v with sparse parameter Jacobian J_S.T,

diff --git a/py/fastspecfit/emline_fit/sparse_rep.py b/py/fastspecfit/emline_fit/sparse_rep.py
@@ -187,7 +187,7 @@ def _rmatvec(self, v):
     # Multiply ideal Jacobian J * v, writing result to w.
     #
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _matvec_J(J, v, w):
         """
         Multiply partial Jacobian J with v,
@@ -209,7 +209,7 @@ def _matvec_J(J, v, w):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _rmatvec_J(J, v, w):
         """
         Multiply v with partial Jacobian J.T,
@@ -236,7 +236,7 @@ def _rmatvec_J(J, v, w):
             w[i] = acc
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _matvec_J_add(J, v, w):
     """
     Multiply partial Jacobian J with v,

diff --git a/py/fastspecfit/emline_fit/utils.py b/py/fastspecfit/emline_fit/utils.py
@@ -11,7 +11,7 @@
 MAX_SDEV = 5.
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def norm_pdf(a):
     """
     PDF of standard normal distribution at a point a
@@ -22,7 +22,7 @@ def norm_pdf(a):
     return 1/SQRT_2PI * np.exp(-0.5 * a**2)
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def norm_cdf(a):
     """
     Approximate the integral of a standard normal PDF from -infty to a.
@@ -51,7 +51,7 @@ def norm_cdf(a):
     return y
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def max_buffer_width(log_obs_bin_edges, line_sigmas, padding=0):
     """
     Compute a safe estimate of the number of nonzero bin fluxes possible

diff --git a/py/fastspecfit/igm.py b/py/fastspecfit/igm.py
@@ -98,7 +98,7 @@ def full_IGM(self, z, lobs):
 
 
     @staticmethod
-    @jit(nopython=True, nogil=True)
+    @jit(nopython=True, nogil=True, cache=True)
     def _full_IGM(z, lobs, scale_tau, igm_params):
 
         lam, cALAF1, cALAF2, cALAF3, cADLA1, cADLA2 = igm_params
@@ -119,7 +119,7 @@ def _full_IGM(z, lobs, scale_tau, igm_params):
         return np.exp(-scale_tau * (tau_LC + tau_LS))
 
 
-@jit(nopython=True, fastmath=True, nogil=True)
+@jit(nopython=True, fastmath=True, nogil=True, cache=True)
 def _tLSLAF(zS, lobs, lam,
             cALAF1, cALAF2, cALAF3):
     """
@@ -146,7 +146,7 @@ def _tLSLAF(zS, lobs, lam,
     return r
 
 
-@jit(nopython=True, fastmath=True, nogil=True)
+@jit(nopython=True, fastmath=True, nogil=True, cache=True)
 def _tLSDLA(zS, lobs, lam,
             cADLA1, cADLA2):
     """
@@ -170,7 +170,7 @@ def _tLSDLA(zS, lobs, lam,
     return r
 
 
-@jit(nopython=True, fastmath=True, nogil=True)
+@jit(nopython=True, fastmath=True, nogil=True, cache=True)
 def _tLCDLA(zS, lobs):
     """
     Lyman continuum, DLA
@@ -207,7 +207,7 @@ def _tLCDLA(zS, lobs):
     return r
 
 
-@jit(nopython=True, fastmath=True, nogil=True)
+@jit(nopython=True, fastmath=True, nogil=True, cache=True)
 def _tLCLAF(zS, lobs):
     """
     Lyman continuum, LAF

diff --git a/py/fastspecfit/resolution.py b/py/fastspecfit/resolution.py
@@ -75,7 +75,7 @@ def dot(self, v, out=None):
 
 
     @staticmethod
-    @jit(nopython=True, fastmath=False, nogil=True)
+    @jit(nopython=True, fastmath=False, nogil=True, cache=True)
     def _matvec(D, v, out):
         """
         Compute matrix-vector product of a Resolution matrix A and a vector v.
@@ -104,7 +104,7 @@ def _matvec(D, v, out):
 
 
     @staticmethod
-    @jit(nopython=True, fastmath=False, nogil=True)
+    @jit(nopython=True, fastmath=False, nogil=True, cache=True)
     def _dia_to_rows(D):
         """
         Convert a diagonally sparse matrix M in the form

diff --git a/py/fastspecfit/util.py b/py/fastspecfit/util.py
@@ -349,7 +349,7 @@ def minfit(x, y, return_coeff=False):
 # array copies and redundant summation
 # on each iteration
 #
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def sigmaclip(c, low=3., high=3.):
 
     n  = len(c)
@@ -386,19 +386,19 @@ def sigmaclip(c, low=3., high=3.):
 
 # Numba's quantile impl is much faster
 # than Numpy's standard version
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def quantile(A, q):
     return np.quantile(A, q)
 
 
 # Numba's median impl is also faster
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def median(A):
     return np.median(A)
 
 
 # Open-coded Numba trapz is much faster than np.traz
-@jit(nopython=True, nogil=True, fastmath=True)
+@jit(nopython=True, nogil=True, fastmath=True, cache=True)
 def trapz(y, x):
     res = 0.
     for i in range(len(x) - 1):
@@ -467,7 +467,7 @@ def trapz_rebin_pre(bin_centers):
     return (edges, ibw)
 
 
-@jit(nopython=True, nogil=True, fastmath=True)
+@jit(nopython=True, nogil=True, fastmath=True, cache=True)
 def _trapz_rebin(x, y, edges, ibw, out):
     """
     Trapezoidal rebinning
@@ -533,7 +533,7 @@ def y_at(edge_x, j): # j: largest j s.t. x[j] < edge_x
     return results
 
 
-@jit(nopython=True, nogil=True)
+@jit(nopython=True, nogil=True, cache=True)
 def centers2edges(centers):
     """
     Convert bin centers to bin edges, guessing at what you probably meant.