Allow the numba cache to be used, for development

[hyanwong]

Copied from SGKIT. Fixes #438

[nspope]

I ran into the same issues as you here-- the bulk of stuff is still being compiled on import. I think pulling all the EP stuff into a low-level extension is the way to go.

[hyanwong]

I think pulling all the EP stuff into a low-level extension is the way to go.

Can we do this while keeping it in python (I presume not)? Or are you thinking of writing the low-level e.g. in rust?

[nspope]

It may be possible to build the extension with numba, but I think it'd be stabler and easier to maintain a C or rust extension.

[nspope]

I could make the edits to the scipy special functions on this branch if you'd like @hyanwong ?

[hyanwong]

Yes, why not. Thank you.

[hyanwong]

Do you want me to have a go at swapping in math.lgamma for the script version, @nspope ?

[nspope]

Sure, sorry to drop the ball on this. Where I'm thinking is, https://github.com/tskit-dev/tsdate/blob/240593d01a1a77561dadebee82cfb5fae67f83bc/tsdate/hypergeo.py#L72C17-L74C26

the _gammaln_f8 should be changed to math.lgamma -- hopefully that's enough to make more functions cache-able?

[hyanwong]

No problem,. I'll give it a whirl. Thanks for the pointer.

[hyanwong]

OK, with the new commits and the math.lgamma function this now seems to work, and cuts import down from 29 seconds (with TSDATE_DISABLE_NUMBA_CACHE=1) to 9 seconds (with TSDATE_DISABLE_NUMBA_CACHE=0) on my laptop.

I put _hyp2f1_unity as an inner function within _hyp2f1_laplace, as I couldn't figure out how to pass things neatly otherwise. I guess you could rework that if it grates or is expected to make things slow @nspope. I haven't figured out how to make _gammainc_inv (and the related approximate_gamma_iqr) cacheable either, so that would make loading faster too, if you can figure it out.

[hyanwong]

Aha. I spotted some more functions I could mark to cache. Now down to 5 secs for import.

[hyanwong]

I guess this is good to go, at least for the time being.

To enable the caching (e.g. for development), you need to explicitly set TSDATE_DISABLE_NUMBA_CACHE=0 which seems a bit of a weird way to phrase it (I would have chosen TSDATE_ENABLE_NUMBA_CACHE=1 I think). I just copied that from sgkit: is there any reason for the choice of name / polarisation in sgkit, @jeromekelleher ?