diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 438f8fe7f..011359a9c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -72,7 +72,7 @@ jobs: HEXRD_EXAMPLE_REPO_PATH: ${{ github.workspace }}/examples NUMBA_DISABLE_JIT: 0 run: | - NUMBA_DISABLE_JIT=1 coverage run --source hexrd -m pytest tests/ + NUMBA_JIT_COVERAGE=1 coverage run --source hexrd -m pytest tests/ coverage combine coverage xml -i if: ${{ matrix.config.os == 'ubuntu-latest'}} diff --git a/hexrd/core/constants.py b/hexrd/core/constants.py index 126690451..a7eac458a 100644 --- a/hexrd/core/constants.py +++ b/hexrd/core/constants.py @@ -138,49 +138,6 @@ ] ).astype(np.float64) -c_coeff_exp1exp = np.array( - [ - 0.999999584, - -0.249992399, - 0.055514994, - -0.010315766, - 0.001535370, - -0.000142164, - ] -).astype(np.complex128) - -cnum_exp1exp = np.array( - [ - 1.0, - 99.0, - 3952.0, - 82544.0, - 979524.0, - 6712860.0, - 25815840.0, - 51369120.0, - 44339040.0, - 10628640.0, - 0.0, - ] -).astype(np.complex128) - -cden_exp1exp = np.array( - [ - 1.0, - 100.0, - 4050.0, - 86400.0, - 1058400.0, - 7620480.0, - 31752000.0, - 72576000.0, - 81648000.0, - 36288000.0, - 3628800.0, - ] -).astype(np.complex128) - """ >> @AUTHOR: Saransh Singh, Lawrence Livermore National Lab, diff --git a/hexrd/core/fitting/peakfunctions.py b/hexrd/core/fitting/peakfunctions.py index a70789507..ca2f131a0 100644 --- a/hexrd/core/fitting/peakfunctions.py +++ b/hexrd/core/fitting/peakfunctions.py @@ -29,12 +29,9 @@ from numba import njit import copy from hexrd.core import constants -from hexrd.core.constants import ( - c_erf, - cnum_exp1exp, - cden_exp1exp, - c_coeff_exp1exp, -) +from hexrd.core.fitting.special import erfc, exp1_complex_numba, wofz + +exp1exp = exp1_complex_numba gauss_width_fact = constants.sigma_to_fwhm lorentz_width_fact = 2.0 @@ -49,88 +46,6 @@ 'pink_beam_dcs': 8, } -""" -cutom function to compute the complementary error function -based on rational approximation of the convergent Taylor -series. coefficients found in -Formula 7.1.26 -Handbook of Mathematical Functions, -Abramowitz and Stegun -Error is < 1.5e-7 for all x -""" - - -@njit(cache=True, nogil=True) -def erfc(x): - # save the sign of x - sign = np.sign(x) - x = np.abs(x) - - # constants - a1, a2, a3, a4, a5, p = c_erf - - # A&S formula 7.1.26 - t = 1.0 / (1.0 + p * x) - y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) - erf = sign * y # erf(-x) = -erf(x) - return 1.0 - erf - - -""" -cutom function to compute the exponential integral -based on Padé approximation of exponential integral -function. coefficients found in pg. 231 Abramowitz -and Stegun, eq. 5.1.53 -""" - - -@njit(cache=True, nogil=True) -def exp1exp_under1(x): - f = np.zeros(x.shape).astype(np.complex128) - for i in range(6): - xx = x ** (i + 1) - f += c_coeff_exp1exp[i] * xx - - return (f - np.log(x) - np.euler_gamma) * np.exp(x) - - -""" -cutom function to compute the exponential integral -based on Padé approximation of exponential integral -function. coefficients found in pg. 415 Y. Luke, The -special functions and their approximations, vol 2 -(1969) Elsevier -""" - - -@njit(cache=True, nogil=True) -def exp1exp_over1(x): - num = np.zeros(x.shape).astype(np.complex128) - den = np.zeros(x.shape).astype(np.complex128) - - for i in range(11): - p = 10 - i - if p != 0: - xx = x**p - num += cnum_exp1exp[i] * xx - den += cden_exp1exp[i] * xx - else: - num += cnum_exp1exp[i] - den += cden_exp1exp[i] - - return (num / den) * (1.0 / x) - - -@njit(cache=True, nogil=True) -def exp1exp(x): - mask = np.sign(x.real) * np.abs(x) > 1.0 - - f = np.zeros(x.shape).astype(np.complex128) - f[mask] = exp1exp_over1(x[mask]) - f[~mask] = exp1exp_under1(x[~mask]) - - return f - # ============================================================================= # 1-D Gaussian Functions diff --git a/hexrd/core/fitting/special.py b/hexrd/core/fitting/special.py new file mode 100644 index 000000000..21d28f03c --- /dev/null +++ b/hexrd/core/fitting/special.py @@ -0,0 +1,444 @@ +import math +import numpy as np +from numba import njit, vectorize, complex128, float64 +from hexrd.core.constants import c_erf + +_EULER_GAMMA = 0.577215664901532860606512090082402431 +_EPS = 2.220446049250313e-16 +_TINY = 1.0e-300 +_MAX_ITER = 1000 + +_SMALL_SERIES_CUTOFF = 4.0 +_LEFT_SERIES_CUTOFF = 100.0 +_MAX_LOG = 709.782712893384 # approximately log(max float64) + + +# Fixed N=11 modified trapezoidal-rule approximation. +# h = sqrt(pi / (N + 1)), H = pi / h +_NTRAP = 11 +_H_STEP = math.sqrt(math.pi / (_NTRAP + 1.0)) +_H_CAP = math.pi / _H_STEP + +_MAX_EXP = 709.782712893384 # approx log(float64 max) +_MIN_EXP = -745.1332191019411 # exp(x) underflows to 0 below this + + +@njit(cache=True, nopython=True) +def erfc(x): + # save the sign of x + sign = np.sign(x) + x = np.abs(x) + + # constants + a1, a2, a3, a4, a5, p = c_erf + + # A&S formula 7.1.26 + t = 1.0 / (1.0 + p * x) + y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) + erf = sign * y # erf(-x) = -erf(x) + return 1.0 - erf + + +@njit(cache=True, nopython=True) +def _complex_log_principal_numba(z): + """ + Principal complex log, written manually so signed zero in imag + is handled through atan2. + """ + return complex( + math.log(math.hypot(z.real, z.imag)), + math.atan2(z.imag, z.real), + ) + + +@njit(cache=True, nopython=True) +def _exp_neg_complex_numba(z): + """ + Compute exp(-z) with simple overflow/underflow guards. + """ + x = z.real + y = z.imag + + if x > 745.0: + return complex(0.0, 0.0) + + if x < -_MAX_LOG: + # The true result is generally outside float64 range. + return complex(math.inf, math.nan) + + scale = math.exp(-x) + return complex(scale * math.cos(y), -scale * math.sin(y)) + + +@njit(cache=True, nopython=True) +def _exp1_series_numba(z): + """ + Logarithmic power series: + + E1(z) = -gamma - log(z) - sum_{k=1}^inf (-z)^k / (k k!) + + Valid on the principal branch, away from z = 0. + """ + term = -z + s = term + + for k in range(2, _MAX_ITER + 1): + term *= -z / k + add = term / k + s += add + + if abs(add) <= _EPS * max(1.0, abs(s)): + break + + return complex(-_EULER_GAMMA, 0.0) - _complex_log_principal_numba(z) - s + + +@njit(cache=True, nopython=True) +def _exp1_contfrac_numba(z): + """ + Continued fraction via modified Lentz iteration. + Good away from the origin and away from the difficult part + of the negative-real branch cut. + """ + b = z + 1.0 + c = complex(1.0 / _TINY, 0.0) + d = 1.0 / b + h = d + + for i in range(1, _MAX_ITER + 1): + a = -float(i * i) + b += 2.0 + + denom = a * d + b + if abs(denom) < _TINY: + denom = complex(_TINY, 0.0) + d = 1.0 / denom + + c = b + a / c + if abs(c) < _TINY: + c = complex(_TINY, 0.0) + + delta = c * d + h *= delta + + if abs(delta - 1.0) <= _EPS: + break + + return h * _exp_neg_complex_numba(z) + + +@njit(cache=True, nopython=True) +def exp1_complex_scalar_numba(z): + """ + Numba-compatible complex exponential integral E1(z). + + This is the complex analogue of scipy.special.exp1 for complex128 input. + For exact negative-real complex inputs, the side of the branch cut is + selected by the sign of the imaginary zero: + complex(-x, +0.0) -> -i*pi side + complex(-x, -0.0) -> +i*pi side + """ + x = z.real + y = z.imag + + if math.isnan(x) or math.isnan(y): + return complex(math.nan, math.nan) + + if x == 0.0 and y == 0.0: + return complex(math.inf, 0.0) + + if math.isinf(x) or math.isinf(y): + if x == math.inf and not math.isinf(y): + return complex(0.0, 0.0) + return complex(math.nan, math.nan) + + az = abs(z) + + near_negative_real_axis = ( + x < 0.0 and az <= _LEFT_SERIES_CUTOFF and abs(y) <= max(2.0, 0.25 * abs(x)) + ) + + if az <= _SMALL_SERIES_CUTOFF or near_negative_real_axis: + return _exp1_series_numba(z) + + out = _exp1_contfrac_numba(z) + + # If the continued fraction path is used exactly on the negative-real + # branch cut, insert the correct signed-zero branch jump. + if x < 0.0 and y == 0.0: + if math.copysign(1.0, y) < 0.0: + out += complex(0.0, math.pi) + else: + out += complex(0.0, -math.pi) + + return out + + +@vectorize([complex128(complex128)], nopython=True) +def exp1_complex_numba(z): + """ + Vectorized complex ufunc version. + """ + return exp1_complex_scalar_numba(z) + + +@njit(cache=True, nopython=True) +def exp1_real_scalar_numba(x): + """ + Real-valued SciPy-like wrapper. + + For real negative x, scipy.special.exp1(x) returns nan. + For complex negative x + 0j, use exp1_complex_scalar_numba instead. + """ + if math.isnan(x): + return math.nan + + if x < 0.0: + return math.nan + + if x == 0.0: + return math.inf + + if math.isinf(x): + return 0.0 + + return exp1_complex_scalar_numba(complex(x, 0.0)).real + + +@vectorize([float64(float64)], nopython=True) +def exp1_real_numba(x): + """ + Vectorized real ufunc version. + """ + return exp1_real_scalar_numba(x) + + +# def exp1exp(x): +# if x.dtype == float: +# return exp1_real_numba(x) +# elif x.dtype == complex: +# return exp1_complex_numba(x) + + +@njit(cache=True, nopython=True) +def _signed_inf_from_trig(v): + if v > 0.0: + return math.inf + if v < 0.0: + return -math.inf + return 0.0 + + +@njit(cache=True, nopython=True) +def _cexp_parts(a, b): + """ + Compute exp(a + i b) with simple overflow/underflow guards. + """ + if a > _MAX_EXP: + c = math.cos(b) + s = math.sin(b) + return complex(_signed_inf_from_trig(c), _signed_inf_from_trig(s)) + + if a < _MIN_EXP: + return complex(0.0, 0.0) + + ea = math.exp(a) + return complex(ea * math.cos(b), ea * math.sin(b)) + + +@njit(cache=True, nopython=True) +def _exp_neg_z2(z): + """ + Compute exp(-z*z), where z = x + i y. + -z^2 = (y^2 - x^2) - 2 i x y + """ + x = z.real + y = z.imag + return _cexp_parts(y * y - x * x, -2.0 * x * y) + + +@njit(cache=True, nopython=True) +def _trap_correction(z, plus): + """ + Correction term: + + 2 exp(-z^2) / (1 + exp(-2 i H z)) if plus=True + 2 exp(-z^2) / (1 - exp(-2 i H z)) if plus=False + + Uses a reciprocal form when exp(-2 i H z) would be very large. + """ + x = z.real + y = z.imag + + log_abs_b = 2.0 * _H_CAP * y + + if log_abs_b > 40.0: + # B = exp(-2 i H z), inv_b = 1/B = exp(2 i H z) + inv_b = _cexp_parts( + -2.0 * _H_CAP * y, + 2.0 * _H_CAP * x, + ) + + # 2 exp(-z^2) / B = 2 exp(-z^2 + 2 i H z) + a_over_b = 2.0 * _cexp_parts( + y * y - x * x - 2.0 * _H_CAP * y, + -2.0 * x * y + 2.0 * _H_CAP * x, + ) + + if plus: + return a_over_b / (1.0 + inv_b) + else: + return a_over_b / (inv_b - 1.0) + + else: + a = 2.0 * _exp_neg_z2(z) + b = _cexp_parts( + log_abs_b, + -2.0 * _H_CAP * x, + ) + + if plus: + return a / (1.0 + b) + else: + return a / (1.0 - b) + + +@njit(cache=True, nopython=True) +def _wofz_midpoint(z): + """ + Midpoint-rule approximation used in part of the first quadrant. + """ + z2 = z * z + az = complex(0.0, 2.0 / _H_CAP) * z + + h0 = 0.5 * _H_STEP + t2 = h0 * h0 + s = math.exp(-t2) / (z2 - t2) + + for k in range(1, _NTRAP + 1): + t = _H_STEP * (k + 0.5) + t2 = t * t + s += math.exp(-t2) / (z2 - t2) + + return az * s + + +@njit(cache=True, nogil=True) +def _wofz_modified_midpoint(z): + """ + Modified midpoint-rule approximation. + """ + return _wofz_midpoint(z) + _trap_correction(z, True) + + +@njit(cache=True, nopython=True) +def _wofz_modified_trapezoid(z): + """ + Modified trapezoidal-rule approximation. + """ + z2 = z * z + az = complex(0.0, 2.0 / _H_CAP) * z + + t = _H_STEP * _NTRAP + t2 = t * t + s = math.exp(-t2) / (z2 - t2) + + for k in range(1, _NTRAP): + t = _H_STEP * k + t2 = t * t + s += math.exp(-t2) / (z2 - t2) + + return complex(0.0, 1.0 / _H_CAP) / z + az * s + _trap_correction(z, False) + + +@njit(cache=True, nopython=True) +def _wofz_first_quadrant(z): + """ + Evaluate w(z) for Re(z) >= 0, Im(z) >= 0. + """ + x = z.real + y = z.imag + + rzh = x / _H_STEP + buff = abs(rzh - math.floor(rzh) - 0.5) + + if y >= max(x, _H_CAP): + return _wofz_midpoint(z) + + elif y < x and buff <= 0.25: + return _wofz_modified_trapezoid(z) + + else: + return _wofz_modified_midpoint(z) + + +@njit(cache=True, nopython=True) +def wofz_complex_scalar_numba(z): + """ + Numba-compatible scalar complex Faddeeva function. + + Approximate replacement for: + + scipy.special.wofz(z) + + for complex128-like finite inputs. + """ + x = z.real + y = z.imag + + if math.isnan(x) or math.isnan(y): + return complex(math.nan, math.nan) + + # Basic SciPy-like limiting behavior for common infinite cases. + # For ordinary finite inputs, the main algorithm below is used. + if math.isinf(x) or math.isinf(y): + if y == -math.inf and x == 0.0: + return complex(math.inf, 0.0) + if y == -math.inf: + return complex(math.nan, math.nan) + return complex(0.0, 0.0) + + # Map arbitrary z to the first quadrant using symmetries. + xneg = x < 0.0 + yneg = y < 0.0 + not_both = xneg != yneg + + zw = z + + if xneg: + zw = -zw + + if not_both: + zw = complex(zw.real, -zw.imag) + + w = _wofz_first_quadrant(zw) + + if not_both: + w = complex(w.real, -w.imag) + + if yneg: + # w(z) = 2 exp(-z^2) - w(-z) + w = 2.0 * _exp_neg_z2(z) - w + + return w + + +@njit(cache=True, nopython=True) +def wofz_real_scalar_numba(x): + """ + Convenience wrapper for real x, returning complex w(x). + """ + return wofz_complex_scalar_numba(complex(x, 0.0)) + + +@vectorize( + [complex128(complex128), complex128(float64)], + nopython=True, +) +def wofz(z): + """ + Vectorized ufunc version. + + Works on complex128 arrays and float64 arrays. + """ + return wofz_complex_scalar_numba(z) diff --git a/hexrd/core/fitting/utils.py b/hexrd/core/fitting/utils.py index 5d0ce8b9f..ed93fe1fd 100644 --- a/hexrd/core/fitting/utils.py +++ b/hexrd/core/fitting/utils.py @@ -2,15 +2,12 @@ import numpy as np from numba import njit +from hexrd.core.fitting.special import erfc, exp1_complex_numba -from hexrd.core.constants import ( - c_erf, - cnum_exp1exp, - cden_exp1exp, - c_coeff_exp1exp, -) from hexrd.core.matrixutil import uniqueVectors +# alias for exp1 function +exp1exp = exp1_complex_numba # ============================================================================= # LMFIT Parameter munging utilities @@ -125,93 +122,6 @@ def _set_peak_center_bounds(params, window_range, min_sep=0.01): prev_peak = curr_peak -# ============================================================================= -# DCS-related utilities -# ============================================================================= - -""" -cutom function to compute the complementary error function -based on rational approximation of the convergent Taylor -series. coefficients found in -Formula 7.1.26 -Handbook of Mathematical Functions, -Abramowitz and Stegun -Error is < 1.5e-7 for all x -""" - - -@njit(cache=True, nogil=True) -def erfc(x): - # save the sign of x - sign = np.sign(x) - x = np.abs(x) - - # constants - a1, a2, a3, a4, a5, p = c_erf - - # A&S formula 7.1.26 - t = 1.0 / (1.0 + p * x) - y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) - erf = sign * y # erf(-x) = -erf(x) - return 1.0 - erf - - -""" -cutom function to compute the exponential integral -based on Padé approximation of exponential integral -function. coefficients found in pg. 231 Abramowitz -and Stegun, eq. 5.1.53 -""" - - -@njit(cache=True, nogil=True) -def exp1exp_under1(x): - f = np.zeros(x.shape).astype(np.complex128) - for i in range(6): - xx = x ** (i + 1) - f += c_coeff_exp1exp[i] * xx - - return (f - np.log(x) - np.euler_gamma) * np.exp(x) - - -""" -cutom function to compute the exponential integral -based on Padé approximation of exponential integral -function. coefficients found in pg. 415 Y. Luke, The -special functions and their approximations, vol 2 -(1969) Elsevier -""" - - -@njit(cache=True, nogil=True) -def exp1exp_over1(x): - num = np.zeros(x.shape).astype(np.complex128) - den = np.zeros(x.shape).astype(np.complex128) - - for i in range(11): - p = 10 - i - if p != 0: - xx = x**p - num += cnum_exp1exp[i] * xx - den += cden_exp1exp[i] * xx - else: - num += cnum_exp1exp[i] - den += cden_exp1exp[i] - - return (num / den) * (1.0 / x) - - -@njit(cache=True, nogil=True) -def exp1exp(x): - mask = np.sign(x.real) * np.abs(x) > 1.0 - - f = np.zeros(x.shape).astype(np.complex128) - f[mask] = exp1exp_over1(x[mask]) - f[~mask] = exp1exp_under1(x[~mask]) - - return f - - @njit(cache=True, nogil=True) def _calc_alpha(alpha, x0): a0, a1 = alpha diff --git a/hexrd/powder/wppf/peakfunctions.py b/hexrd/powder/wppf/peakfunctions.py index fcb283ad3..8c14faa34 100644 --- a/hexrd/powder/wppf/peakfunctions.py +++ b/hexrd/powder/wppf/peakfunctions.py @@ -29,8 +29,9 @@ import copy from hexrd.core import constants from numba import vectorize, float64, njit, prange -from hexrd.core.fitting.peakfunctions import erfc, exp1exp +from hexrd.core.fitting.special import erfc, exp1_complex_numba +exp1exp = exp1_complex_numba # from scipy.special import erfc, exp1 # addr = get_cython_function_address("scipy.special.cython_special", "exp1") @@ -521,7 +522,7 @@ def _gaussian_pink_beam(alpha, beta, fwhm_g, tth, tth_list): t1 = erfc(y) t2 = erfc(z) g = np.zeros(tth_list.shape) - zmask = np.abs(del_tth) > 5.0 + zmask = np.abs(del_tth) > 100.0 g[~zmask] = (0.5 * (alpha * beta) / (alpha + beta)) * np.exp(u[~zmask]) * t1[ ~zmask ] + np.exp(v[~zmask]) * t2[~zmask] @@ -550,7 +551,12 @@ def _lorentzian_pink_beam(alpha, beta, fwhm_l, tth, tth_list): # f1 = exp1(p) # f2 = exp1(q) - y = -(alpha * beta) / (np.pi * (alpha + beta)) * (f1 + f2).imag + y = np.zeros(tth_list.shape) + zmask = np.abs(del_tth) > 100.0 + + y[~zmask] = ( + -(alpha * beta) / (np.pi * (alpha + beta)) * (f1[~zmask] + f2[~zmask]).imag + ) mask = np.isnan(y) y[mask] = 0.0 diff --git a/tests/core/fitting/test_fitting_utils.py b/tests/core/fitting/test_fitting_utils.py index d168b1540..ea28c6634 100644 --- a/tests/core/fitting/test_fitting_utils.py +++ b/tests/core/fitting/test_fitting_utils.py @@ -83,15 +83,6 @@ def test_peak_center_bounds(mock_uniq, params): def test_math_funcs(): - assert np.isclose(utils.erfc(np.array([0.0, 100.0]))[0], 1.0) - assert np.isclose(utils.erfc(np.array([-100.0]))[0], 2.0) - - assert not np.isnan(utils.exp1exp_under1(np.array([0.1 + 0j]))[0]) - assert not np.isnan(utils.exp1exp_over1(np.array([2.0 + 0j]))[0]) - assert not np.any( - np.isnan(utils.exp1exp(np.array([0.1, 2.0], dtype=np.complex128))) - ) - assert utils._calc_alpha(np.array([1.0, 2.0]), 0.0) == 1.0 assert utils._calc_beta(np.array([3.0, 4.0]), 0.0) == 3.0 diff --git a/tests/core/fitting/test_peakfunctions.py b/tests/core/fitting/test_peakfunctions.py index c44bd7961..39e41e44f 100644 --- a/tests/core/fitting/test_peakfunctions.py +++ b/tests/core/fitting/test_peakfunctions.py @@ -1,8 +1,9 @@ # tests/test_peakfunctions_simplified.py import numpy as np import pytest -from scipy.special import erfc as scipy_erfc +from scipy.special import erfc as scipy_erfc, wofz as scipy_wofz, exp1 as scipy_exp1 import hexrd.core.fitting.peakfunctions as pf +import hexrd.core.fitting.special as sp # ----------------- small utilities ----------------- @@ -21,18 +22,41 @@ def finite_diff_derivative(func, p, x, idx, rel_step=1e-6): # ----------------- special functions ----------------- -def test_erfc_and_exp1exp_smoke(): - vals = np.array([-5.0, -1.0, 0.0, 0.5, 1.0, 3.2], dtype=float) - got = pf.erfc(vals) +def test_erfc(): + # test real values + vals = np.random.uniform(-5, 5, 5) + got = sp.erfc(vals) assert got.shape == vals.shape assert np.all((got >= 0.0) & (got <= 2.0)) assert np.allclose(got, scipy_erfc(vals), rtol=1e-6, atol=1e-6) - x_pos = np.array([0.1, 0.5, 1.0, 2.0], dtype=float) - for fn in (pf.exp1exp_under1, pf.exp1exp_over1, pf.exp1exp): - arr = fn(x_pos) - assert arr.shape == x_pos.shape - assert np.all(np.isfinite(arr)) + +def test_wofz(): + # test complex values + vals = np.random.uniform(-5, 5, 5) + 1j * np.random.uniform(-5, 5, 5) + got = sp.wofz(vals) + assert got.shape == vals.shape + assert np.allclose(got, scipy_wofz(vals), rtol=1e-6, atol=1e-6) + + +def test_exp1(): + # test real values + vals = np.random.uniform(-5, 5, 5) + got = sp.exp1_real_numba(vals) + assert got.shape == vals.shape + # filter out nans. these are infs in scipy + mask = ~np.isnan(got) + assert np.allclose(got[mask], scipy_exp1(vals)[mask], rtol=1e-6, atol=1e-6) + + # test complex values + vals = np.random.uniform(-5, 5, 5) + 1j * np.random.uniform(-5, 5, 5) + got = sp.exp1_complex_numba(vals) + scipy_got = scipy_exp1(vals) + assert got.shape == vals.shape + + # filter out nans. these are infs in scipy + mask = ~np.isnan(got) + assert np.allclose(got[mask], scipy_got[mask], rtol=1e-6, atol=1e-6) # ----------------- 1D Gaussian ----------------- @@ -80,12 +104,8 @@ def test_gaussian1d_full_and_derivatives(): np.testing.assert_allclose(dmat[0:3, :], analytic_no_bg) for i in range(3): - fd = finite_diff_derivative( - pf._gaussian1d_no_bg, p[:3], x, i, rel_step=1e-6 - ) - np.testing.assert_allclose( - analytic_no_bg[i, :], fd, rtol=1e-5, atol=1e-8 - ) + fd = finite_diff_derivative(pf._gaussian1d_no_bg, p[:3], x, i, rel_step=1e-6) + np.testing.assert_allclose(analytic_no_bg[i, :], fd, rtol=1e-5, atol=1e-8) # ----------------- Lorentzian ----------------- @@ -113,9 +133,7 @@ def test_lorentzian_unit_and_full_and_deriv(): np.testing.assert_allclose(no_bg_out, A * expected_unit, rtol=1e-12) full_out = pf.lorentzian1d(np.array([A, x0, FWHM, c0, c1]), x) - np.testing.assert_allclose( - full_out, A * expected_unit + c0 + c1 * x, rtol=1e-12 - ) + np.testing.assert_allclose(full_out, A * expected_unit + c0 + c1 * x, rtol=1e-12) assert int(np.argmax(no_bg_out)) == int(np.argmin(np.abs(x - x0))) p = np.array([2.5, -0.3, 0.9], dtype=float) @@ -191,25 +209,19 @@ def test_pvoigt_and_split_variants(): FWHM_plus = 1.6 n_minus = 0.2 n_plus = 0.7 - p_no_bg = np.array( - [A, x0, FWHM_minus, FWHM_plus, n_minus, n_plus], dtype=float - ) + p_no_bg = np.array([A, x0, FWHM_minus, FWHM_plus, n_minus, n_plus], dtype=float) x_small = np.linspace(-2.0, 2.0, 401) out_no_bg = pf._split_pvoigt1d_no_bg(p_no_bg, x_small) expected = np.zeros_like(x_small) xr = x_small >= x0 xl = x_small < x0 - expected[xr] = A * pf._unit_pvoigt1d( - np.array([x0, FWHM_plus, n_plus]), x_small[xr] - ) + expected[xr] = A * pf._unit_pvoigt1d(np.array([x0, FWHM_plus, n_plus]), x_small[xr]) expected[xl] = A * pf._unit_pvoigt1d( np.array([x0, FWHM_minus, n_minus]), x_small[xl] ) np.testing.assert_allclose(out_no_bg, expected) np.testing.assert_allclose( - pf.split_pvoigt1d( - np.hstack([p_no_bg, np.array([0.05, -0.01])]), x_small - ), + pf.split_pvoigt1d(np.hstack([p_no_bg, np.array([0.05, -0.01])]), x_small), expected + 0.05 - 0.01 * x_small, ) @@ -217,8 +229,7 @@ def test_pvoigt_and_split_variants(): x_sym = np.linspace(-2.0, 2.0, 201) np.testing.assert_allclose( pf._split_pvoigt1d_no_bg(p_sym, x_sym), - p_sym[0] - * pf._unit_pvoigt1d(np.array([p_sym[1], p_sym[2], p_sym[4]]), x_sym), + p_sym[0] * pf._unit_pvoigt1d(np.array([p_sym[1], p_sym[2], p_sym[4]]), x_sym), ) p_zeroA = np.array( @@ -245,9 +256,7 @@ def test_calc_alpha_beta_and_mixing_factor_behavior_and_wrappers(): assert np.allclose(got_beta, expected_beta, atol=1e-12) for fwhm_g, fwhm_l in [(0.1, 0.2), (0.5, 0.5), (1.0, 2.0), (2.5, 0.3)]: - eta, fwhm = pf._mixing_factor_pv( - np.float64(fwhm_g), np.float64(fwhm_l) - ) + eta, fwhm = pf._mixing_factor_pv(np.float64(fwhm_g), np.float64(fwhm_l)) assert np.isfinite(eta) and np.isfinite(fwhm) assert 0.0 <= eta <= 1.0 poly = ( @@ -283,9 +292,7 @@ def test_calc_alpha_beta_and_mixing_factor_behavior_and_wrappers(): L = pf._lorentzian_pink_beam(p_l, x) expected_full = eta * L + (1.0 - eta) * G + b0 + b1 * x out_wrapper = pf.pink_beam_dcs(p_full, x) - np.testing.assert_allclose( - out_wrapper, expected_full, rtol=1e-8, atol=1e-12 - ) + np.testing.assert_allclose(out_wrapper, expected_full, rtol=1e-8, atol=1e-12) x_lm = np.linspace(-0.8, 0.8, 101) out_lm = pf.pink_beam_dcs_lmfit( @@ -293,21 +300,15 @@ def test_calc_alpha_beta_and_mixing_factor_behavior_and_wrappers(): ) alpha_lm = pf._calc_alpha(np.array([alpha0, alpha1], dtype=float), x0) beta_lm = pf._calc_beta(np.array([beta0, beta1], dtype=float), x0) - p_g_lm = np.hstack( - ([A, x0], np.array([alpha_lm, beta_lm, fwhm_g], dtype=float)) - ) - p_l_lm = np.hstack( - ([A, x0], np.array([alpha_lm, beta_lm, fwhm_l], dtype=float)) - ) + p_g_lm = np.hstack(([A, x0], np.array([alpha_lm, beta_lm, fwhm_g], dtype=float))) + p_l_lm = np.hstack(([A, x0], np.array([alpha_lm, beta_lm, fwhm_l], dtype=float))) eta_lm, _ = pf._mixing_factor_pv(fwhm_g, fwhm_l) expected_lm = eta_lm * pf._lorentzian_pink_beam(p_l_lm, x_lm) + ( 1.0 - eta_lm ) * pf._gaussian_pink_beam(p_g_lm, x_lm) np.testing.assert_allclose(out_lm, expected_lm, rtol=1e-8, atol=1e-12) - p = np.array( - [10.0, 0.5, 0.01, 0.02, 0.03, 0.01, 0.15, 0.20, 1.0, 0.1], dtype=float - ) + p = np.array([10.0, 0.5, 0.01, 0.02, 0.03, 0.01, 0.15, 0.20, 1.0, 0.1], dtype=float) x_small = np.linspace(-1.0, 2.0, 7) y_no_bg = pf._pink_beam_dcs_no_bg(p[:-2], x_small) y_full = pf.pink_beam_dcs(p, x_small) @@ -346,9 +347,7 @@ def test_tanh_and_2d_coord_transform_and_gauss2d_rot_and_3d(): assert np.allclose(out2, out2[0, 0]) p_rot = np.array([1.0, 0.0, 0.0, 1.0, 1.0, np.pi / 3]) - out_rot = pf._gaussian2d_rot_no_bg( - p_rot, np.zeros((3, 3)), np.zeros((3, 3)) - ) + out_rot = pf._gaussian2d_rot_no_bg(p_rot, np.zeros((3, 3)), np.zeros((3, 3))) assert out_rot.shape == (3, 3) and np.isfinite(out_rot).all() p_full = np.array([1.0, 0.0, 0.0, 1.0, 1.0, 2.0, 0.5, -0.5]) @@ -359,20 +358,14 @@ def test_tanh_and_2d_coord_transform_and_gauss2d_rot_and_3d(): p_gr = np.array([1.0, 0.0, 0.0, 1.0, 1.0, np.pi / 4, 1.0, 0.3, 0.2]) out_gr = pf.gaussian2d_rot(p_gr, np.zeros((3, 3)), np.zeros((3, 3))) - base_r = pf._gaussian2d_rot_no_bg( - p_gr[:6], np.zeros((3, 3)), np.zeros((3, 3)) - ) + base_r = pf._gaussian2d_rot_no_bg(p_gr[:6], np.zeros((3, 3)), np.zeros((3, 3))) assert np.allclose(out_gr, base_r + 1.0 + 0.3 * 0 - 0.2 * 0) xq = np.array([[-1.0, 1.0], [-1.0, 1.0]]) yq = np.array([[-1.0, -1.0], [1.0, 1.0]]) p2d = np.array([1.0, 0.0, 0.0, 1.0, 2.0, 1.5, 2.5, 0.1, 0.2, 0.3, 0.4]) out2d = pf._split_pvoigt2d_no_bg(p2d, xq, yq) - assert ( - out2d.shape == (2, 2) - and np.isfinite(out2d).all() - and (out2d >= 0).all() - ) + assert out2d.shape == (2, 2) and np.isfinite(out2d).all() and (out2d >= 0).all() out2drot = pf._split_pvoigt2d_rot_no_bg( np.hstack((p2d, np.pi / 6)), np.zeros((3, 3)), np.zeros((3, 3)) @@ -380,9 +373,7 @@ def test_tanh_and_2d_coord_transform_and_gauss2d_rot_and_3d(): assert out2drot.shape == (3, 3) and np.isfinite(out2drot).all() p_split_bg = np.hstack((p2d, np.pi / 7, 1.0, 0.2, -0.1)) - out_sp_bg = pf.split_pvoigt2d_rot( - p_split_bg, np.zeros((4, 4)), np.zeros((4, 4)) - ) + out_sp_bg = pf.split_pvoigt2d_rot(p_split_bg, np.zeros((4, 4)), np.zeros((4, 4))) base_sp = pf._split_pvoigt2d_rot_no_bg( p_split_bg[:12], np.zeros((4, 4)), np.zeros((4, 4)) ) @@ -417,9 +408,7 @@ def _build_peak_row(pktype): if pktype == "split_pvoigt": return np.array([0.9, 0.0, 0.7, 1.1, 0.25, 0.75], dtype=float) if pktype == "pink_beam_dcs": - return np.array( - [1.3, 0.1, 0.01, 0.002, 0.02, -0.0003, 0.2, 0.4], dtype=float - ) + return np.array([1.3, 0.1, 0.01, 0.002, 0.02, -0.0003, 0.2, 0.4], dtype=float) raise ValueError