Merge pull request #568 from SamuelBorden/refactor

Added numba defaults for math, fixed some bugs in extended pdfs

src/pygama/math/functions/crystal_ball.py

@@ -9,12 +9,12 @@
 from math import erf
 
 from pygama.math.functions.pygama_continuous import pygama_continuous 
+from pygama.utils import numba_math_defaults_kwargs as nb_kwargs
+from pygama.utils import numba_math_defaults as nb_defaults
 
-kwd = {"parallel": False, "fastmath": True}
-kwd_parallel = {"parallel": True, "fastmath": True}
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_crystal_ball_pdf(x: np.ndarray, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
     r"""
     PDF of a power-law tail plus gaussian. Its range of support is :math:`x\in\mathbb{R}, \beta>0, m>1`. It computes:
@@ -73,7 +73,7 @@ def nb_crystal_ball_pdf(x: np.ndarray, mu: float, sigma: float, beta: float, m:
     return y
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_crystal_ball_cdf(x: np.ndarray, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
     r"""
     CDF for power-law tail plus gaussian. Its range of support is :math:`x\in\mathbb{R}, \beta>0, m>1`. It computes: 
@@ -135,7 +135,7 @@ def nb_crystal_ball_cdf(x: np.ndarray, mu: float, sigma: float, beta: float, m:
     return y
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_crystal_ball_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
     r"""
     Scaled PDF of a power-law tail plus gaussian. 
@@ -161,7 +161,7 @@ def nb_crystal_ball_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: flo
     return area * nb_crystal_ball_pdf(x, mu, sigma, beta, m)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_crystal_ball_scaled_cdf(x: np.ndarray, area: float, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
     r"""
     Scaled CDF for power-law tail plus gaussian. Used for extended binned fits. 
@@ -213,7 +213,7 @@ def cdf_norm(self, x: np.ndarray,  x_lo: float, x_hi: float, mu: float, sigma: f
 
 
     def pdf_ext(self, x: np.ndarray, x_lo: float, x_hi: float, area: float, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
-        return np.diff(nb_crystal_ball_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, beta, m)), nb_crystal_ball_scaled_pdf(x, area, mu, sigma, beta, m)
+        return np.diff(nb_crystal_ball_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, beta, m))[0], nb_crystal_ball_scaled_pdf(x, area, mu, sigma, beta, m)
     def cdf_ext(self, x: np.ndarray, area: float, mu: float, sigma: float, beta: float, m: float) -> np.ndarray:
         return nb_crystal_ball_scaled_cdf(x, area, mu, sigma, beta, m)
 

src/pygama/math/functions/exgauss.py

@@ -13,12 +13,14 @@
 from pygama.math.functions.error_function import nb_erf, nb_erfc
 from pygama.math.functions.gauss import nb_gauss_cdf, nb_gauss_pdf
 
+from pygama.utils import numba_math_defaults_kwargs as nb_kwargs
+from pygama.utils import numba_math_defaults as nb_defaults
+
 limit = np.log(sys.float_info.max)/10
-kwd = {"parallel": False, "fastmath": True}
-kwd_2 = {"parallel": True, "fastmath": True}
 
 
-@nb.njit(**kwd)
+
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_tail_exact(x: float, mu: float, sigma: float, tau: float, tmp: float) -> float:
     r"""
     Exact form of a normalized exponentially modified Gaussian PDF. 
@@ -63,7 +65,7 @@ def nb_gauss_tail_exact(x: float, mu: float, sigma: float, tau: float, tmp: floa
     return tail_f
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_tail_approx(x: np.ndarray, mu: float, sigma: float, tau: float) -> np.ndarray:
     r"""
     Approximate form of a normalized exponentially modified Gaussian PDF
@@ -92,7 +94,7 @@ def nb_gauss_tail_approx(x: np.ndarray, mu: float, sigma: float, tau: float) ->
     return tail_f
 
 
-@nb.njit(**kwd_2)
+@nb.njit(**nb_kwargs)
 def nb_exgauss_pdf(x: np.ndarray, mu: float, sigma: float, tau: float) -> np.ndarray:
     r"""
     Normalized PDF of an exponentially modified Gaussian distribution. Its range of support is :math:`x\in(-\infty,\infty)`, :math:`\tau\in(-\infty,\infty)`
@@ -131,7 +133,7 @@ def nb_exgauss_pdf(x: np.ndarray, mu: float, sigma: float, tau: float) -> np.nda
     return tail_f
 
 
-@nb.njit(**kwd_2)
+@nb.njit(**nb_kwargs)
 def nb_exgauss_cdf(x: np.ndarray, mu: float, sigma: float, tau: float) -> np.ndarray:
     r"""
     The CDF for a normalized exponentially modified Gaussian.  Its range of support is :math:`x\in(-\infty,\infty)`, :math:`\tau\in(-\infty,\infty)`
@@ -170,7 +172,7 @@ def nb_exgauss_cdf(x: np.ndarray, mu: float, sigma: float, tau: float) -> np.nda
     return cdf
 
 
-@nb.njit(**kwd_2)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_exgauss_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float, tau: float) -> np.ndarray:
     r"""
     Scaled PDF of an exponentially modified Gaussian distribution
@@ -195,7 +197,7 @@ def nb_exgauss_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float, t
     return area * nb_exgauss_pdf(x, mu, sigma, tau)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_exgauss_scaled_cdf(x: np.ndarray, area: float, mu: float, sigma: float, tau: float) -> np.ndarray:
     r"""
     Scaled CDF of an exponentially modified Gaussian distribution
@@ -244,7 +246,7 @@ def cdf_norm(self, x: np.ndarray, x_lo: float, x_hi: float, mu: float, sigma: fl
         return self._cdf_norm(x, x_lo, x_hi, mu, sigma, tau)
 
     def pdf_ext(self, x: np.ndarray, x_lo: float, x_hi: float, area: float, mu: float, sigma: float, tau: float) -> np.ndarray:
-        return np.diff(nb_exgauss_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, tau)), nb_exgauss_scaled_pdf(x, area, mu, sigma, tau)
+        return np.diff(nb_exgauss_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, tau))[0], nb_exgauss_scaled_pdf(x, area, mu, sigma, tau)
     def cdf_ext(self, x: np.ndarray, area: float, mu: float, sigma: float, tau: float) -> np.ndarray:
         return nb_exgauss_scaled_cdf(x, area, mu, sigma, tau)
 

src/pygama/math/functions/exponential.py

@@ -8,12 +8,12 @@
 import numpy as np
 
 from pygama.math.functions.pygama_continuous import pygama_continuous
+from pygama.utils import numba_math_defaults_kwargs as nb_kwargs
+from pygama.utils import numba_math_defaults as nb_defaults
 
-kwd = {"parallel": False, "fastmath": True}
-kwd_parallel = {"parallel": True, "fastmath": True}
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_exponential_pdf(x: np.ndarray, mu: float, sigma: float, lamb: float) -> np.ndarray:
     r"""
     Normalised exponential probability density distribution, w/ args: mu, sigma, lamb. Its range of support is :math:`x\in[0,\infty), \lambda>0`. 
@@ -50,7 +50,7 @@ def nb_exponential_pdf(x: np.ndarray, mu: float, sigma: float, lamb: float) -> n
     return y
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_exponential_cdf(x: np.ndarray, mu: float, sigma: float, lamb: float) -> np.ndarray:
     r"""
     Normalised exponential cumulative distribution, w/ args:  mu, sigma, lamb. Its range of support is :math:`x\in[0,\infty), \lambda>0`. 
@@ -87,7 +87,7 @@ def nb_exponential_cdf(x: np.ndarray, mu: float, sigma: float, lamb: float) -> n
     return y
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_exponential_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float, lamb: float) -> np.ndarray:
     r"""
     Scaled exponential probability distribution, w/ args: area, mu, sigma, lambd.
@@ -111,7 +111,7 @@ def nb_exponential_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: floa
     return area * nb_exponential_pdf(x, mu, sigma, lamb)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_exponential_scaled_cdf(x: np.ndarray, area: float, mu: float, sigma: float, lamb: float) -> np.ndarray:
     r"""
     Exponential cdf scaled by the area, used for extended binned fits 
@@ -162,7 +162,7 @@ def cdf_norm(self, x: np.ndarray, x_lo: float, x_hi: float, mu: float, sigma: fl
         return self._cdf_norm(x, x_lo, x_hi, mu, sigma, lamb)
 
     def pdf_ext(self, x: np.ndarray, x_lo: float, x_hi: float, area: float, mu: float, sigma: float, lamb: float) -> np.ndarray:
-        return np.diff(nb_exponential_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, lamb)), nb_exponential_scaled_pdf(x, area, mu, sigma, lamb)
+        return np.diff(nb_exponential_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma, lamb))[0], nb_exponential_scaled_pdf(x, area, mu, sigma, lamb)
     def cdf_ext(self, x: np.ndarray, area: float, mu: float, sigma: float, lamb: float) -> np.ndarray:
         return nb_exponential_scaled_cdf(x, area, mu, sigma, lamb)
 

src/pygama/math/functions/gauss.py

@@ -9,11 +9,10 @@
 
 from pygama.math.functions.error_function import nb_erf
 from pygama.math.functions.pygama_continuous import pygama_continuous 
+from pygama.utils import numba_math_defaults as nb_defaults
 
-kwd = {"parallel": False, "fastmath": True}
 
-
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     """
     Gaussian, unnormalised for use in building PDFs
@@ -39,7 +38,7 @@ def nb_gauss(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     return np.exp(-0.5 * z ** 2)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_amp(x: np.ndarray, mu: float, sigma: float, a: float) -> np.ndarray:
     """
     Gaussian with height as a parameter for FWHM etc.
@@ -67,7 +66,7 @@ def nb_gauss_amp(x: np.ndarray, mu: float, sigma: float, a: float) -> np.ndarray
     return a * np.exp(-0.5 * z ** 2)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_pdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     r"""
     Normalised Gaussian PDF, w/ args: mu, sigma. The support is :math:`(-\infty, \infty)`
@@ -95,7 +94,7 @@ def nb_gauss_pdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     return np.exp(-0.5 * z ** 2) * invnorm
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_cdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     r"""
     Gaussian CDF, w/ args: mu, sigma. The support is :math:`(-\infty, \infty)`
@@ -123,7 +122,7 @@ def nb_gauss_cdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     return 1/2 * (1 + nb_erf(invs*(x - mu)/(np.sqrt(2))))
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
     """
     Gaussian with height as a parameter for fwhm etc.
@@ -145,7 +144,7 @@ def nb_gauss_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float) ->
     return area * nb_gauss_pdf(x, mu, sigma)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_gauss_scaled_cdf(x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
     """
     Gaussian CDF scaled by the number of signal counts for extended binned fits 
@@ -191,7 +190,7 @@ def cdf_norm(self, x: np.ndarray, x_lo: float, x_hi: float, mu: float, sigma: fl
         return self._cdf_norm(x, x_lo, x_hi, mu, sigma)
 
     def pdf_ext(self, x: np.ndarray, x_lo: float, x_hi: float, area: float, mu: float, sigma: float) -> np.ndarray:
-        return np.diff(nb_gauss_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma)), nb_gauss_scaled_pdf(x, area, mu, sigma)
+        return np.diff(nb_gauss_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma))[0], nb_gauss_scaled_pdf(x, area, mu, sigma)
     def cdf_ext(self, x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
         return nb_gauss_scaled_cdf(x, area, mu, sigma)
 

src/pygama/math/functions/linear.py

@@ -6,13 +6,11 @@
 from numba import prange
 
 from pygama.math.functions.pygama_continuous import pygama_continuous 
+from pygama.utils import numba_math_defaults_kwargs as nb_kwargs
+from pygama.utils import numba_math_defaults as nb_defaults
 
 
-kwd = {"parallel": False, "fastmath": True}
-kwd_parallel = {"parallel": True, "fastmath": True}
-
-
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_linear_pdf(x: np.ndarray, x_lo: float, x_hi: float, m: float, b: float) -> np.ndarray:
     r"""
     Normalised linear probability density function, w/ args: m, b. Its range of support is :math:`x\in(x_{lower},x_{upper})`. 
@@ -48,7 +46,7 @@ def nb_linear_pdf(x: np.ndarray, x_lo: float, x_hi: float, m: float, b: float) -
     return result
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_linear_cdf(x: np.ndarray, x_lo: float, x_hi: float, m: float, b: float) -> np.ndarray:
     r"""
     Normalised linear cumulative density function, w/ args: m, b. Its range of support is :math:`x\in(x_{lower},x_{upper})`. 
@@ -84,7 +82,7 @@ def nb_linear_cdf(x: np.ndarray, x_lo: float, x_hi: float, m: float, b: float) -
     return result
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_linear_scaled_pdf(x: np.ndarray, x_lo: float, x_hi: float, area: float, m: float, b: float) -> np.ndarray:
     r"""
     Scaled linear probability distribution, w/ args: m, b.
@@ -110,7 +108,7 @@ def nb_linear_scaled_pdf(x: np.ndarray, x_lo: float, x_hi: float, area: float, m
     return area * nb_linear_pdf(x, x_lo, x_hi, m, b)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_linear_scaled_cdf(x: np.ndarray, x_lo: float, x_hi: float, area: float, m: float, b: float) -> np.ndarray:
     r"""
     Linear cdf scaled by the area for extended binned fits 

src/pygama/math/functions/moyal.py

@@ -9,13 +9,13 @@
 from numba import prange
 from math import erfc
 
-from pygama.math.functions.pygama_continuous import pygama_continuous 
+from pygama.math.functions.pygama_continuous import pygama_continuous
+from pygama.utils import numba_math_defaults_kwargs as nb_kwargs
+from pygama.utils import numba_math_defaults as nb_defaults
 
-kwd = {"parallel": False, "fastmath": True}
-kwd_parallel = {"parallel": True, "fastmath": True}
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_moyal_pdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     r"""
     Normalised Moyal probability distribution function, w/ args: mu, sigma. Its support is :math:`x\in\mathbb{R}`
@@ -47,7 +47,7 @@ def nb_moyal_pdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     return y
 
 
-@nb.njit(**kwd_parallel)
+@nb.njit(**nb_kwargs)
 def nb_moyal_cdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     r"""
     Normalised Moyal cumulative distribution, w/ args: mu, sigma.
@@ -80,7 +80,7 @@ def nb_moyal_cdf(x: np.ndarray, mu: float, sigma: float) -> np.ndarray:
     return y
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_moyal_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
     r"""
     Scaled Moyal probability density function, w/ args: mu, sigma, area.
@@ -102,7 +102,7 @@ def nb_moyal_scaled_pdf(x: np.ndarray, area: float, mu: float, sigma: float) ->
     return area * nb_moyal_pdf(x, mu, sigma)
 
 
-@nb.njit(**kwd)
+@nb.njit(**nb_defaults(parallel=False))
 def nb_moyal_scaled_cdf(x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
     r"""
     Moyal cdf scaled by the area, used for extended binned fits 
@@ -150,7 +150,7 @@ def cdf_norm(self, x: np.ndarray, x_lo: float, x_hi: float, mu: float, sigma: fl
         return self._cdf_norm(x, x_lo, x_hi, mu, sigma)
 
     def pdf_ext(self, x: np.ndarray, x_lo: float, x_hi: float, area: float, mu: float, sigma: float) -> np.ndarray:
-        return np.diff(nb_moyal_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma)), nb_moyal_scaled_pdf(x, area, mu, sigma)
+        return np.diff(nb_moyal_scaled_cdf(np.array([x_lo, x_hi]), area, mu, sigma))[0], nb_moyal_scaled_pdf(x, area, mu, sigma)
     def cdf_ext(self, x: np.ndarray, area: float, mu: float, sigma: float) -> np.ndarray:
         return nb_moyal_scaled_cdf(x, area, mu, sigma)