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.

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?

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.

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

Yes, why not. Thank you.

Do you want me to have a go at swapping in math.lgamma for the script version, @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?

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

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.

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

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 ?

I don't think there's a good reason for the choice of name, I agree with your preference for TSDATE_ENABLE_NUMBA_CACHE=1

@nspope - I'm happy to merge this. Loading time for the library is now down to 2.5s on my laptop with TSDATE_ENABLE_NUMBA_CACHE=1. Does it look OK to you?