jit a_d functions inside thrust.py

[jtegedor]

This implements a recommendation from @astrojuanlu in #1076 to jit some functions on trust.py, which makes a significant speed improvement.

Execution times of test_thrust.py test:

Current (without jit):

============================= slowest 15 durations ==============================
83.19s call     tests/tests_twobody/test_thrust.py::test_leo_geo_numerical[0.49741883681838395]
33.42s call     tests/tests_twobody/test_thrust.py::test_geo_cases_numerical[0.0-0.4]
29.90s call     tests/tests_twobody/test_thrust.py::test_geo_cases_numerical[0.4-0.0]
10.02s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_numerical[0.0-0.1245]
9.41s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_numerical[0.1245-0.0]
4.78s call     tests/tests_twobody/test_thrust.py::test_soyuz_standard_gto_numerical
0.55s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.1-20.0-83.043-1.6789]
0.27s call     tests/tests_twobody/test_thrust.py::test_soyuz_standard_gto_delta_v
0.26s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_time_and_delta_v[0.0-0.1245]
0.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.4-20.0-61.522-1.7592]
============= 17 passed, 1 skipped, 2 warnings in 172.84s (0:02:52) =============

New (with jit):

============================= slowest 15 durations ==============================
78.02s call     tests/tests_twobody/test_thrust.py::test_leo_geo_numerical[0.49741883681838395]
4.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_numerical[0.4-0.0]
3.99s call     tests/tests_twobody/test_thrust.py::test_geo_cases_numerical[0.0-0.4]
3.51s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_numerical[0.0-0.1245]
3.37s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_numerical[0.1245-0.0]
1.47s call     tests/tests_twobody/test_thrust.py::test_soyuz_standard_gto_numerical
0.44s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.1-20.0-83.043-1.6789]
0.31s call     tests/tests_twobody/test_thrust.py::test_soyuz_standard_gto_delta_v
0.24s call     tests/tests_twobody/test_thrust.py::test_sso_disposal_time_and_delta_v[0.0-0.1245]
0.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.8-10.0-16.304-1.9799]
0.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.6-16.0-40.0-1.7241]
0.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.4-20.0-61.522-1.7592]
0.01s call     tests/tests_twobody/test_thrust.py::test_geo_cases_beta_dnd_delta_v[0.2-20.0-76.087-1.689]
============= 17 passed, 1 skipped, 2 warnings in 96.86s (0:01:36) ==============

[codecov]

Codecov Report

Merging #1250 (5528a45) into main (1a25403) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #1250      +/-   ##
==========================================
+ Coverage   91.27%   91.29%   +0.01%     
==========================================
  Files          79       79              
  Lines        4127     4135       +8     
  Branches      360      360              
==========================================
+ Hits         3767     3775       +8     
  Misses        270      270              
  Partials       90       90              

Impacted Files	Coverage Δ
src/poliastro/twobody/thrust/change_a_inc.py	100.00% <100.00%> (ø)
src/poliastro/twobody/thrust/change_argp.py	100.00% <100.00%> (ø)
...oliastro/twobody/thrust/change_ecc_quasioptimal.py	100.00% <100.00%> (ø)
src/poliastro/twobody/thrust/change_inc_ecc.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1a25403...5528a45. Read the comment docs.

[astrojuanlu]

Aw yeah, these are the speed improvements I wanted to see! 😀 Thanks @jtegedor ! Reviewing right now

[astrojuanlu]

Approved with just one question 🚀

[astrojuanlu]

What error were you getting without forceobj=True? I'm open to merge this as is and investigate later, but asking in case it's an easy fix.

[jtegedor]

This is the Numbat warning that appears without forceobj=True:

  /workspaces/poliastro/src/poliastro/twobody/thrust/change_a_inc.py:46: NumbaWarning: 
  Compilation is falling back to object mode WITH looplifting enabled because Function "a_d" failed type inference due to: Internal error at <numba.core.typeinfer.CallConstraint object at 0x7fc12d092280>.
  tuple index out of range
  During: resolving callee type: type(CPUDispatcher(<function beta at 0x7fc12df97550>))
  During: typing of call at /workspaces/poliastro/src/poliastro/twobody/thrust/change_a_inc.py (52)
  
  Enable logging at debug level for details.
  
  File "src/poliastro/twobody/thrust/change_a_inc.py", line 52:
      def a_d(t0, u_, k):
          <source elided>
          # Change sign of beta with the out-of-plane velocity
          beta_ = beta(t0, V_0=V_0, f=f, beta_0=beta_0_) * np.sign(r[0] * (inc_f - inc_0))
          ^
  
    @jit

tests/tests_twobody/test_thrust.py::test_leo_geo_numerical[0.49741883681838395]
  /usr/local/lib/python3.8/site-packages/numba/core/object_mode_passes.py:151: NumbaWarning: Function "a_d" was compiled in object mode without forceobj=True.
  
  File "src/poliastro/twobody/thrust/change_a_inc.py", line 47:
      @jit
      def a_d(t0, u_, k):
      ^
  
    warnings.warn(errors.NumbaWarning(warn_msg,

tests/tests_twobody/test_thrust.py::test_leo_geo_numerical[0.49741883681838395]
  /usr/local/lib/python3.8/site-packages/numba/core/object_mode_passes.py:161: NumbaDeprecationWarning: 
  Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.
  
  For more information visit https://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit
  
  File "src/poliastro/twobody/thrust/change_a_inc.py", line 47:
      @jit
      def a_d(t0, u_, k):
      ^

I did not investigate further, any suggestion?

[astrojuanlu]

Wow, that's super weird.

It will take me some time to get to debugging this, but I don't want to block the pull request because of that. Let's merge it - it will save us a lot of testing time, and we can look at the numba error later.