diff --git a/msaexp/resample_numba.py b/msaexp/resample_numba.py index 409fcac..179a3ca 100644 --- a/msaexp/resample_numba.py +++ b/msaexp/resample_numba.py @@ -207,7 +207,9 @@ def pixel_integrated_gaussian_numba(x, mu, sigma, dx=None, normalization=1.0): @jit(nopython=True, fastmath=True, error_model="numpy") def compute_igm(z, wobs, scale_tau=1.0): """ - Calculate Inoue+14 IGM transmission + Calculate + [Inoue+ (2014)](https://ui.adsabs.harvard.edu/abs/2014MNRAS.442.1805I) IGM + transmission, reworked from `~eazy.igm.Inoue14` Parameters ---------- @@ -218,12 +220,12 @@ def compute_igm(z, wobs, scale_tau=1.0): Observed-frame wavelengths, Angstroms scale_tau : float - Scale factor multiplied to tau_igm + Scalar multiplied to tau_igm Returns ------- igmz : array-like - IGM transmission + IGM transmission factor """ _LAF = np.array( @@ -322,7 +324,8 @@ def compute_igm(z, wobs, scale_tau=1.0): ADLA1 = _DLA[2] ADLA2 = _DLA[3] - _pow = lambda a, b: a**b + def _pow(a, b): + return a**b #### # Lyman series, Lyman-alpha forest @@ -342,11 +345,13 @@ def compute_igm(z, wobs, scale_tau=1.0): tau = np.zeros_like(wobs) zS = z - + + # Explicit iteration should be fast in JIT for i, wi in enumerate(wobs): if wi > 1300.0 * (1 + zS): continue - + + # Iterate over Lyman series for j, lsj in enumerate(ALAM): # LS LAF if wi < lsj * (1 + zS): @@ -365,8 +370,9 @@ def compute_igm(z, wobs, scale_tau=1.0): else: tau[i] += ADLA2[j] * (wi / lsj) ** 3 - # LC DLA + # Lyman Continuum if wi < lamL * (1 + zS): + # LC DLA if zS < z1DLA: tau[i] += ( 0.2113 * _pow(1 + zS, 2) @@ -413,7 +419,6 @@ def compute_igm(z, wobs, scale_tau=1.0): - 0.2496 * _pow(wi / lamL, 2.1) ) else: - if wi > lamL * (1.0 + z2LAF): tau[i] += 5.221e-4 * ( _pow(1 + zS, 3.4) * _pow(wi / lamL, 2.1) @@ -433,4 +438,5 @@ def compute_igm(z, wobs, scale_tau=1.0): ) igmz = np.exp(-scale_tau * tau) + return igmz diff --git a/msaexp/tests/test_numba.py b/msaexp/tests/test_numba.py new file mode 100644 index 0000000..c6796d7 --- /dev/null +++ b/msaexp/tests/test_numba.py @@ -0,0 +1,45 @@ +import numpy as np + + +def test_igm(): + """ + Test IGM implementation + """ + + try: + from ..resample_numba import compute_igm + except ImportError: + return None + + wobs = np.linspace(0.2, 1.0, 32) * 1.0e4 + + igmz = compute_igm(3.1, wobs, scale_tau=1.0) + _tau = np.log(igmz) + + assert np.allclose( + igmz[:6], + np.array( + [ + 0.09095587, + 0.09023704, + 0.09558526, + 0.1098157, + 0.13983981, + 0.20491541, + ] + ), + rtol=1.0e-3, + ) + + # Scaled tau + igm2 = compute_igm(3.1, wobs, scale_tau=2.0) + _tau2 = np.log(igm2) + assert np.allclose(_tau * 2, _tau2) + + # Low-z, no IGM + igmz = compute_igm(0.1, wobs, scale_tau=1.0) + assert np.allclose(igmz, 1.0) + + # High-z, completely absorbed + igmz = compute_igm(10, wobs, scale_tau=1.0) + assert np.allclose(igmz, 0.0, rtol=1.0e-6) diff --git a/setup.py b/setup.py index b303aa0..beda28e 100755 --- a/setup.py +++ b/setup.py @@ -1,4 +1,5 @@ #!/usr/bin/env python from setuptools import setup -setup() \ No newline at end of file + +setup()