Use Numba, take 2

[astrojuanlu]

In gh-17 there was an attempt to accelerate the code using Numba, but as described in gh-46 after the change was released, it ended up making the code slower.

However, I am quite stubborn and was disappointed to see Numba go, so I gave it a try. I am happy to see that some basic tests pass, and that the code is 6.5x faster :tada:

In [1]: line1, line2 = (
   ...:   "1 00005U 58002B   00179.78495062  .00000023  00000-0  28098-4 0  4753",
   ...:   "2 00005  34.2682 348.7242 1859667 331.7664  19.3264 10.82419157413667",
   ...: )

In [2]: from sgp4.api import Satrec

In [3]: sat = Satrec.twoline2rv(line1, line2)

In [4]: sat.sgp4(0, 0)
Out[4]: 
(0,
 (2561.0165140549384, -10709.951057126229, -727.886667131307),
 (4.149113751894791, 1.3483989470230675, -2.927579443622721))

In [5]: sat.sgp4(sat.jdsatepoch + 10, 0)
Out[5]: 
(0,
 (-3316.08121585171, -6182.251823725894, -4684.562483409189),
 (6.290093860106019, -2.6388564186743766, 1.3415255224216052))

In [6]: %timeit sat.sgp4(sat.jdsatepoch + 10, 0)
12 µs ± 448 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [7]: from sgp4.fast.model import Satrec as FastSatrec, twoline2rv as fast_twoline2rv

In [8]: fast_sat = fast_twoline2rv(FastSatrec(), line1, line2)

In [9]: fast_sat.sgp4(0, 0)
Out[9]: 
(0,
 (2561.0165140549384, -10709.951057126229, -727.886667131307),
 (4.149113751894791, 1.3483989470230675, -2.927579443622721))

In [10]: fast_sat.sgp4(sat.jdsatepoch + 10, 0)
Out[10]: 
(0,
 (-3316.08121585171, -6182.251823725894, -4684.562483409189),
 (6.290093860106019, -2.6388564186743766, 1.3415255224216052))

In [11]: %timeit fast_sat.sgp4(sat.jdsatepoch + 10, 0)
1.78 µs ± 15.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

This departs from gh-17 in three ways:

• The key to success is to never use Numba in "object mode", and instead force always "nopython mode" (@jit(nopython=True) or @njit). In turn, this requires jitting more functions and making compilation times slower, but the resulting code should be way faster.
  • This also means that a new Satrec class had to be created using the experimental @jitclass functionality.
• I didn't touch almost anything from the original codebase, and instead I copy-pasted and tweaked what I needed, adding a new sgp4.fast package. This of course is open to discussion, but made testing the code easier. The advantage is that, as compilation times are so slow (~5 seconds on my laptop) it's only worth it if the propagation is called many times.
• As numba does not support string formatting, I had to change the functions in propagation.py to simplify the error messages, and make them static. This is the only real code modification I had to make, everything else works.

Some other comments:

• Interestingly, to make this work I had to make ephtype and revnum strings, but in the C++ code, they are int and long respectively. I wonder if this difference is intentional :warning:
• Numba does support passing np.datetime64 objects around, but not creating them in nopython mode https://github.com/numba/numba/issues/6555. Besides, datetime.datetime objects are not supported at all. As a result, I didn't include an .epoch attribute.

Hope that this serves at least as a conversation starter on what would be the best way to add this functionality to python-sgp4.

[brandon-rhodes]

Interesting, thanks for sharing! I have a backlog of other issues at the moment, so it might be a few weeks before I can check this branch out locally to compare the fast files to the normal ones and see all that has changed. In the meantime I'll think about whether this is something that python-sgp4 itself can host and support indefinitely going forward, or whether I should encourage you to publish a separate Python package that would let you support directly the users that need this specially crafted version of the code.

To clarify: this helps the case where the C++ SGP4 fails to compile on a machine (or isn't available as a pre-compiled Python wheel), and the library falls back to doing pure Python?

[astrojuanlu]

Interesting, thanks for sharing! I have a backlog of other issues at the moment, so it might be a few weeks before I can check this branch out locally to compare the fast files to the normal ones and see all that has changed. In the meantime I'll think about whether this is something that python-sgp4 itself can host and support indefinitely going forward, or whether I should encourage you to publish a separate Python package that would let you support directly the users that need this specially crafted version of the code.

Excellent, thanks a lot!

To clarify: this helps the case where the C++ SGP4 fails to compile on a machine (or isn't available as a pre-compiled Python wheel), and the library falls back to doing pure Python?

It would be a way to have "native performance" without needing to compile C++. The trade-off here would be the installation of numba & llvmlite itself - they support a lot of platforms and produce binary wheels for several architectures, but I am not sure what's the intersection between "platforms in which C++ cannot be compiled or isn't available as pre-compiled wheel" and "platforms in which numba is available as pre-compiled wheel and doesn't require compiling LLVM".

(And about "native performance": I still didn't benchmark your C++ wrapper against my Numba version, but I will do it very soon)

[brandon-rhodes]

(And about "native performance": I still didn't benchmark your C++ wrapper against my Numba version, but I will do it very soon)

Hey, @astrojuanlu! Looking at open issues and pull requests on my dashboard, I thought I'd circle back to this PR and start the conversation up again. Did you ever get a chance to check the performance against the C++ wrapper?

Looking back over the code, I suspect that this would work best as a separate project. It could be named “NumbaSGP4” or somesuch, and I could add a mention of it to the docs for folks who need (a) a faster-than-Python SGP4 module, but (b) can't install or compile the C++ version bundled here in this library.

The big reason I suspect this code doesn't belong here in this package is compatibility: sgp4 works to remain compatible with Python 2.7, whereas the syntax used by @jitclass — which, wow, by the way, is really cool, I've never seen a class decorated with C types so succinctly in Python before! — weren't available until Python 3.5. As a separate Python package, your code would be able to state its Python 3.5 compatibility right in its metadata. But as a mere subdirectory fast/ in this package, it wouldn't carry its own metadata.

But maybe I'm being overly finicky? Since several months have passed since I was last able to take a look at it, let me know where things stand from your end, and we can figure out what to do with this PR. Thanks!

[astrojuanlu]

Hi @brandon-rhodes ! I ended up benchmarking the C++ wrapper as well, yes :) You can see that in https://github.com/astrojuanlu/sgp4-benchmarks/blob/master/benchmarks/test_cpp_wrapper.py, and the report is in https://opensatcom.org/2020/12/28/omm-assessment-sgp4-benchmarks/ (second link)

I understand your concerns and I'm fine with making this its own package. Besides, if it breaks for whatever reason people will complain here, but I fear I might not have enough free time available to maintain it (since the project started I have left the space industry... for a job at Read the Docs!).

I will close this PR without removing my branch, and will publish these changes as a standalone package with a different name.

Thanks a lot again @brandon-rhodes ! 🙌🏽

[brandon-rhodes]

I understand your concerns and I'm fine with making this its own package.

Thanks for understanding why I don't want this package to take on a third SGP4 implementation! It’s always awkward to say “no” to a pull request that involved hard work that really accomplished something, so thanks for your gracious answer.

I have left the space industry... for a job at Read the Docs

WOW! Congratulations! I enjoyed the one Write the Docs conference I attended (in Portland), it was interesting to see how wide a community the effort has gathered, even outside of traditional programming documentation projects.

And, again, if you'd like to comment back here later with the package name (numba_sgp4 might be better than the name I suggested above, by the way), I'll be happy to add a mention in this package's README.