Improve Python leapfrog integration performance

[jamesgrimmett]

See discussions in #365

This change aligns the Python implementation of the leapfrog integrator with the C code by combining the drift calculations to reduce computations in the inner integration loop.

Also, the checks for non-axisymmetric and dissipative Potentials are moved from the internal (e.g., _evaluateRforces) to public (e.g., evaluateRforces) methods in order to avoid calling the checks at every loop iteration.

Provides a ~50% reduction in computation time.

Profiling main using this script

         55743016 function calls (49982710 primitive calls) in 27.938 seconds

   Ordered by: cumulative time
   List reduced from 99 to 50 due to restriction <50>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   27.938   27.938 Orbits.py:1338(integrate)
        1    0.000    0.000   27.938   27.938 integrateFullOrbit.py:714(integrateFullOrbit)
        1    0.001    0.001   27.938   27.938 integrateFullOrbit.py:766(integrate_for_map)
        1    0.415    0.415   27.937   27.937 symplecticode.py:38(leapfrog)
   320015    1.003    0.000   26.539    0.000 integrateFullOrbit.py:1236(_rectForce)
   320015    0.467    0.000   10.233    0.000 Potential.py:2257(_evaluatezforces)
   320015    0.477    0.000   10.107    0.000 Potential.py:2123(_evaluateRforces)

Profiling this branch;

         15121126 function calls (15121090 primitive calls) in 13.951 seconds

   Ordered by: cumulative time
   List reduced from 98 to 50 due to restriction <50>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   13.951   13.951 Orbits.py:1338(integrate)
        1    0.000    0.000   13.951   13.951 integrateFullOrbit.py:714(integrateFullOrbit)
        1    0.001    0.001   13.951   13.951 integrateFullOrbit.py:766(integrate_for_map)
        1    0.296    0.296   13.950   13.950 symplecticode.py:38(leapfrog)
   320015    0.937    0.000   12.978    0.000 integrateFullOrbit.py:1236(_rectForce)
   640030    6.469    0.000    7.289    0.000 PowerSphericalPotentialwCutoff.py:205(_mass)
   320015    0.313    0.000    5.725    0.000 Potential.py:2267(_evaluatezforces)
   320015    0.304    0.000    5.458    0.000 Potential.py:2133(_evaluateRforces)

[jamesgrimmett]

Hi @jobovy, here are the changes discussed in #365. As far as I can tell there is just one small change needed for the tests; passing pot=None to Orbit.from_fit now raises a PotentialError instead of an AttributeError, which I hope is an improvement. Otherwise, I think test_dissipative_noVelocityError and test_nonaxierror_function still provide the necessary coverage. I tested with pytest -vxs tests/test_potential.py tests/test_orbit.py, I hope that is sufficient.

Though not necessarily related to the leapfrog integration, I also move the non-axi check from evaluatePotentials to _evaluatePotentials for completeness. Just let me know if you'd prefer that change not to be included.

On a related note; I also looked at adding a isDissipative attribute to potentials, similar to isNonAxi. It looks like it would require a few more changes than anticipated, mainly due to the way that RZToverticalPotential is implemented/tested (the non-Potential objects used for testing currently need to pass through _isDissipative without issue, to be caught by later checks). Anyway, I thought that might be better not included in this change.

Let me know if there's anything I've missed.

EDIT: Oops, a couple of extra commits for some things that I missed. I have edited this comment to match.

[codecov]

Codecov Report

Attention: Patch coverage is 96.29630% with 2 lines in your changes missing coverage. Please review.

Project coverage is 93.29%. Comparing base (e593ab6) to head (ac1785f).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
galpy/potential/Potential.py	87.50%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (e593ab6) and HEAD (ac1785f). Click for more details.

HEAD has 1 upload less than BASE

Flag	BASE (e593ab6)	HEAD (ac1785f)
	13	12

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #690      +/-   ##
==========================================
- Coverage   99.91%   93.29%   -6.62%     
==========================================
  Files         200      200              
  Lines       29400    29413      +13     
  Branches      564      564              
==========================================
- Hits        29374    27441    -1933     
- Misses         26     1972    +1946     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[jobovy]

Thanks @jamesgrimmett, this looks great! A few comments:

• There's a test failure in the actionAngle tests, because it seems like the estimateDeltaStaeckel function was relying on the non-axi check to error when called with non-axisymmetric potentials (there are two test failures, but they are both the same underlying issue). I think the best thing to do is to just put a non-axi check in estimateDeltaStaeckel and just raise a PotentialError("Calling estimateDeltaStaeckel with non-axisymmetric potentials is not supported") there to keep the old behavior (and then the tests will pass). You can run the relevant tests with pytest -vxs tests/test_actionAngle.py -k interface_staeckel_PotentialErrors.
• For consistency, we should probably make the same change to the evaluateplanarPotentials/Rforces/phitorques in planarPotential.py and in evaluatelinearpotentials/forces in linearPotential.py.
• Could you add an entry in the HISTORY.txt about these changes (probably two entries, one for the leapfrog change, one for the dissipative/non-axi changes).

[jamesgrimmett]

Hi @jobovy, I've made the change that you suggested to resolve the actionAngle test failures - that works, thanks!

For consistency, we should probably make the same change to the evaluateplanarPotentials/Rforces/phitorques in planarPotential.py and in evaluatelinearpotentials/forces in linearPotential.py.

I made this change, but in the process I have now added an isDissipative attribute to the force/potential classes (contrary to my earlier comment). I think this was needed for removing the non-axi/dissipative checks from the internal functions in planarPotentials.py, because the _isDissipative result is still used later in the _evaluateplanarRforces/phitorques logic.

I have also adjusted the _isNonAxi and _isDissipative functions to catch the AttributeError and raise a PotentialError instead, if the input Pot is missing an isNonAxi or isDissipative attribute. This is so that incorrect usage such as plp = potential.RZToverticalPotential("something else", 1.2) still behaves similarly (raises a PotentialError), and the tests don't need to be adjusted after the changes to _isDissipative.

Hope that's all OK, it does increase the scope of this change a bit, so just let me know if you had something else in mind.

[jobovy]

Thanks for these updates! I like the addition of the isDissipative attribute, but I've left a few minor comments on the implementation that I think should be easy to address. There is also now a new slight issue that there are five lines in verticalPotential.py that are not covered by the tests:

https://app.codecov.io/gh/jobovy/galpy/pull/690/blob/galpy/potential/verticalPotential.py

These used to be covered, but aren't anymore because the new code in _isDissipative now raises the PotentialError. Just in case future changes to _isDissipative could let these bad inputs through again, I think it's good to just keep the checks in verticalPotential.py, so perhaps we should just label them when # pragma: no cover so they are excluded from the coverage calculation (because they should never be hit). You can just put the # pragma: no cover after the relevant colon (e.g.. after except: or after else:.

Once this is done and the coverage is back to what it was, I think this will be good to be merged!

I don't know how familiar you are with rebasing, but if you are, it could be good to clean up the commit history a bit, by combining the commits that update the leapfrog implementation, the commits to HISTORY.txt and perhaps combining some other commits as well. I try to keep the commit history relatively clean these days and normally I would just squash a PR into a single commit to do that, but in this case it would be good to separate the PR into a few commits. But if you're not familiar with rebasing, I can also clean the commit history up a bit before merging.

[jamesgrimmett]

Thanks for your comments - I've resolved most, but just had a couple of questions for the remaining two.

I'm happy to do a rebase once I finish up the last couple of changes. I'll try to group everything into a few commits, but of course feel free to re-rebase if you can see a cleaner way to group the changes.

[jobovy]

Hi @jamesgrimmett,

Thanks for these final tweaks and for the clean rebase. I made some final very minor tweaks to verticalPotential.py, swapping the conversion.get_physical(Pot) and _isDissipative checks in two places so they are better covered by the tests (because there were still two lines not covered). So I'll merge this once the CI clears. Thanks again!