Stop using Numba to implement 1D functions

[jdbuhler]

Following a discussion on the cosipy Slack, I put together the following quick test to evaluate whether Numba accelerates 1D function evaluation in Astromodels.

In this branch, I replaced all the Numba functions in numba_functions.py with new, pure Numpy implementations in core_functions.py. The Numpy implementations consistently use Numpy rather than Python's math library and remove any explicit loops that were added for Numba acceleration. In a couple of cases (Band and broken powerlaw), I used Numpy masking to replace an unavoidable branch on the argument value for vector inputs.

These functions all pass Astropy's test suite, which checks for agreement with stored values to within single precision, except for SmoothlyBrokenPowerLaw. For that function, the Numba version switches to an asymptotic formula for x < -6 or x > 4, which introduces a relative error of about 1e-6 vs the non-asymptotic case. I used an overflow-safe implementation that avoids the need for the asymptotic case, so I believe my version is actually more accurate. For other functions, agreement is much better than single-precision -- more like 1 part in 10^-14 or better.

To assess the performance impact of this change, I ran each 1D function from powerlaws.py and blackbody.py (the ones that use the Numba implementations) for 1000 iterations with its default parameters on an input x consisting of 1 million random values between 0 and 1. Times were measured after an initial 5-iteration warmup to ensure that I did not measure LLVM or other 1-time overhead for the Numba functions. Here are my results on a 2.3 GHz Intel(R) Xeon(R) Gold 5118 CPU running Linux, using a fresh Conda environment for Astromodels with the latest dependencies as of 8/16/2026:

For all functions except Blackbody and ModifiedBlackbody, the pure Numpy implementation was faster -- occasionally several times faster.

Some thoughts on why we see these results... in the absence of parallelization (which is not done in Astromodels right now), Numba achieves its speedup by (1) eliminating interpreter overhead, (2) implementing basic compiler optimizations like common subexpression elimination, and (3) using looping to fuse multi-step operations and avoid creating intermediate arrays. The third of these happens only if the user codes it that way; some of the Numba functions in Astromodels, such as the cutoff powerlaw implementation, do this, but not all do.

If a Numba implementation does nothing except call Numpy vectorized operations, as in plaw_eval(), it will likely see no benefit and will pay overhead associated with things like calling convention setup. Functions that implement explicit loops over complex operations, like cplaw_eval(), should benefit from eliminating intermediate memory allocations. Unfortunately, on x86, Numba is known to use slower versions of the basic transcendental functions than the Numpy library. Consequently, the expected performance gain is canceled out by these slower functions. (I've seen the same phenomenon with the numexpr package.)

It's also worth noting that I measured performance for a large input x, where the cost of operations likely dominates Numpy allocation and interpreter overhead. Some of these functions might be faster in Numba for much smaller x. But that's probably not the most important case for overall application performance.

Note for future: I suspect the use of Numba in absorption.py could similarly be removed with no ill consequences, but I didn't have a test case for this readily to hand. The uses in extinction.py are better candidates for Numba being faster.

[PreisTo]

Thanks, this looks really good to my ignorant eyes @jdbuhler! Just out of curiosity: did you set OMP_NUM_THREADS and related to 1 for this?

The tests currently seem to be failing only because of our flake8 requirements — a simple run of black should be enough to fix that :)

[jdbuhler]

I did not set OMP_NUM_THREADS because none of the affected Numba code was multithreaded, and I didn't see anything else using matrix multiply. But I can certainly check if it has an impact.

(Update: actually, OMP_NUM_THREADS was already set to 1, but I also set OPENBLAS_NUM_THREADS to 1 and reran everything, since my Numpy is linked to OpenBLAS. The numbers did not change by enough to alter the graph.)

[jdbuhler]

How would I go about updating the expected results for SmoothlyBrokenPowerLaw to get rid of the test failures? They seem to be stored in a file somewhere...

[PreisTo]

@jdbuhler There is a scripts/generate_test_data.py that you can run and as you want to replace I guess you would have to pass 1 as an argument.

However this also means you will have to install (at least) xspec before running so there are no missing entries in the new file.
Be aware xspec==13.0.0 currently breaks astromodels ...