Implement caching of the mass for the PowerSphericalPotentialwCutoff potential

[jamesgrimmett]

Hi @jobovy, here are the second set of changes to close off #701.

I've taken two different approaches in two commits, happy to rebase this branch to contain whichever you prefer.

In 020eb77 I've implemented the caching of the mass in the Rforce and zforce functions. I left the Planar functions as-is (including the nargs/potentialArgs) because I don't think there would be any performance benefit from caching in the Planar case, but just let me know if you'd prefer the caching to be included for the Planar functions as well for consistency (or if I'm mistaken about the lack of benefit).

In 84da14b I instead implemented the caching inside of the mass function itself. This requires that the potentialArgs be passed to the mass function, but reduces the duplicate code in each of the calling functions. In this case the caching needs to be set up for any situation that the mass function going to be called, i.e. including the Planar case.

I guess the tradeoff is that first option is potentially more flexible, whereas the second option is a bit tidier and reduces some repetition in the code.

In terms of performance both commits are the same; compared to the main branch I see a 25% performance improvement. E.g., using this simple test case;

    ts = np.linspace(0,1000,100001)*u.Gyr
    sun = Orbit()
    t = time.time()
    sun.integrate(ts, MWPotential2014, method="symplec4_c")
    print(f"Duration: {time.time() - t}")

With the current main branch I get

Duration: 7.169827222824097

and with this change I get

Duration: 5.4198899269104

Finally I haven't implemented any caching for the Python integrator, but am happy to include that here if you like (on that front, functools.lru_cache is an interesting option).

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.90%. Comparing base (f254193) to head (eb7d081).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #720   +/-   ##
=======================================
  Coverage   99.90%   99.90%           
=======================================
  Files         202      202           
  Lines       29743    29754   +11     
  Branches      606      608    +2     
=======================================
+ Hits        29716    29727   +11     
  Misses         27       27           

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

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jamesgrimmett]

Noted the tests failing for tests/test_orbit.py -k 'not test_energy_jacobi_conservation', I've converted back to draft and will take a look at the cause

[jobovy]

Hi @jamesgrimmett , thanks for this PR and implementing the two options! While I think doing the caching in the mass function is more elegant, I think it might be best to go with the caching in the force functions themselves, because (a) it's more consistent with how caching is doing for other potentials, (b) it avoids unnecessary caching when evaluating the potential or using planarRforce, and (c) you could then actually cache the entire - amp * m / pow(r2,1.5) part of the force, just needing to multiply by R or Z to get the R/Zforce. It probably won't make a real difference for the performance, but if we're going to be caching, we might as well cache as much as possible. Thanks!

P.S.: I wouldn't worry about caching in the Python code since that is not very performant anyway.

[jamesgrimmett]

Thanks @jobovy, all good points! Agreed that is the more practical change, and consistency is good.
I've reverted back to the first option, and made the change to cache that entire force magnitude(?) term. As you guessed, still the same ~25% performance improvement, but there isn't any sense in not caching it and should save a few numbers from being crunched. Just let me know if there should be any other adjustments.

[jobovy]

Looks great, thanks so much!