Cache SSB astrometry calculations

[dlakaplan]

Claude-recommended speedup: #1984

On J0437 data, tried 5 iterations. chi^2 was within 1e-3 with both versions. Average time went from 19.3s to 10.9s.

Have model.clear_ssb_cache() to clear the cache, and also can control with either:
$PINT_DISABLE_CACHE=1 or $PINT_DISABLE_ASTROMETRY_CACHE=1

[jeremy-baier]

amazing !
is the clear_cache method inherited by the timing model class ? or does the clear_cache method just exist as a method of the astrometry delay in particular ?

[dlakaplan]

amazing ! is the clear_cache method inherited by the timing model class ? or does the clear_cache method just exist as a method of the astrometry delay in particular ?

All component level methods are inherited by the timing model. But that brings up a good point which is that there could be other caches elsewhere that we want to deal with. So I changed the name to model.clear_ssb_cache() to make it explicit.

[dlakaplan]

So this mostly works. But oddly not 100%. I added env variables to aid testing. And I find that when I test with/without the cache I usually get the results I expect (without cache is slow, and chi^2 changes very slightly). But not always. I don't yet understand why.

[dlakaplan]

So I am now more confused. Using the master branch or this PR, with or without caching I see bimodal results: sometimes slower, sometimes faster, but when faster I always see a slightly different chi^2 (at the 1e-3 level). I could understand variation in wall-clock time just due to CPU load. But not variation in chi^2. That should be deterministic?

[dlakaplan]

OK, I think I understand more. For some reason, different fits use different numbers of iterations/calls. That leads to pretty different clock times as well as slightly different chi^2. Still not sure why that is, but now I understand the origin of the variation that I see.

Nonetheless, I do see improved performance with this PR in terms of time/call for the ssb methods. With caching enabled:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       43    0.000    0.000    0.697    0.016 astrometry.py:131(ssb_to_psb_xyz_ICRS)
      151    0.001    0.000    0.025    0.000 astrometry.py:1114(ssb_to_psb_xyz_ECL)

or

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       74    0.001    0.000    1.145    0.015 astrometry.py:131(ssb_to_psb_xyz_ICRS)
      272    0.002    0.000    0.037    0.000 astrometry.py:1114(ssb_to_psb_xyz_ECL)

and without:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       43    0.000    0.000    1.649    0.038 astrometry.py:131(ssb_to_psb_xyz_ICRS)
      151    0.003    0.000    0.422    0.003 astrometry.py:1114(ssb_to_psb_xyz_ECL)

and

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       74    0.001    0.000    2.902    0.039 astrometry.py:131(ssb_to_psb_xyz_ICRS)
      272    0.005    0.000    0.795    0.003 astrometry.py:1114(ssb_to_psb_xyz_ECL)

[abhisrkckl]

This looks OK. Shall I merge this?

[dlakaplan]

This looks OK. Shall I merge this?

I think this is OK, and it's something we can turn off, but it should be checked closely.

[codecov]

Codecov Report

❌ Patch coverage is 82.81250% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.33%. Comparing base (036a0bc) to head (685f7a1).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
src/pint/models/astrometry.py	82.81%	6 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1992      +/-   ##
==========================================
- Coverage   70.34%   70.33%   -0.01%     
==========================================
  Files         109      109              
  Lines       25652    25709      +57     
  Branches     4081     4092      +11     
==========================================
+ Hits        18044    18082      +38     
- Misses       6453     6465      +12     
- Partials     1155     1162       +7     

☔ 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.

[dlakaplan]

Look at:
https://nanograv-pint--1992.org.readthedocs.build/en/1992/explanation.html#caching
my goal is to have a section there for every explicitly cached calculation with instructions on how to clear/turn off.

[abhisrkckl]

Is this ready to go?

[dlakaplan]

Is this ready to go?

I would like somebody (you, @scottransom , ...) to check it a little, just to make sure I didn't do anything foolish. Maybe you have already looked.

[abhisrkckl]

Also,

1. Please make sure that methods like adjust_TOAs, apply_clock_corrections, compute_TDBs, compute_posvels etc invalidate the cache.
2. What happens to the astrometry cache when two TOAs objects are merged?

[dlakaplan]

• I changed the name of the property to use_ssb_cache in case we add others later
• I made computation of the key a method, and it will only happen if use_ssb_cache is True
• I added more tests

1. Please make sure that methods like adjust_TOAs, apply_clock_corrections, compute_TDBs, compute_posvels etc invalidate the cache.

I don't know that making them invalidate the cache directly is possible, since those methods only apply to TOAs, not models. But for the methods that change TOAs the cache will invalidate and the values will be recomputed. I added some tests for those (not sure it should apply for compute_posvels: only for those that change tdbld)

2. What happens to the astrometry cache when two TOAs objects are merged?

Works as expected. The cache is invalidated since the times change, and is recomputed.

[abhisrkckl]

Shall I merge this once the tests pass?

[dlakaplan]

Shall I merge this once the tests pass?

Your concerns over the various operations are all OK? If so then OK to merge I think.