search

danielkoll / PyRADS

A line-by-line longwave radiation code for planetary atmospheres.
MIT License
30 stars 9 forks source link

Option to use numba for ~4x speed-up of OpticalThickness calculation #15

Closed AndrewILWilliams closed 3 years ago

AndrewILWilliams commented 3 years ago

This PR is a first-attempt to close #14 . Turns out it was much easier than I expected to use numba to speed up the compute_Absorption() and compute_Absorption_fixed_cutoff() functions in Absorption_Crosssections_HITRAN2016.py.

To do this, I've made a separate file called Absorption_Crosssections_HITRAN2016_numba.py, and added a use_numba option to the OpticalThickness calculation. Currently it works exactly the same as before, but if you set use_numba=True (for example, in Test01.olr/compute_olr_h2o_numba.py) then the new, numba-accelerated functions are imported from Absorption_Crosssections_HITRAN2016_numba.py.

The code in Absorption_Crosssections_HITRAN2016_numba.py is quite similar to it's non-numba counterpart, but I've made some (minimal) changes to the structure so that it works with numba. For example, I've made separate functions for getGamma() and genAbsgrid(), because numba was having trouble dealing with those functions/lambdas in-loop.

In terms of speed up, on my Macbook Pro Test01.olr/compute_olr_h2o.py takes ~12 seconds to run, whereas Test01.olr/compute_olr_h2o_numba.py takes ~6 seconds to run the first time round (during numba's compilation procedure), but for everytime after that it takes ~3 seconds!

It's not a crazy speed-up, but given how easy it was to implement, I'm happy with a 4x increase in performance.

I think it would be helpful to implement some tests in PyRADS so that we can be sure that numba isn't changing any of the results. Perhaps this could happen next?

danielkoll commented 3 years ago

Crashes with Python2.7, numba 0.47. Works in Python3.7.6, numba 0.48. Will add additional conda environment file, then merge pull request.