Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JP-3603: reorder tweakreg to reduce iterations through models #8424

Merged
merged 15 commits into from May 18, 2024

Conversation

braingram
Copy link
Collaborator

@braingram braingram commented Apr 16, 2024

This PR reorganizes operations within the tweakreg step to work towards improving image3 performance.

Changes include:

  • Refactor how custom catalogs are resolved (the behavior should be the same although at the the moment there appear to be no tests for this feature, this PR adds tests). This will allow removal of special handling of tweakreg_catalog in ModelContainer (one of the reasons that class opens all models when constructed).
  • use wcs_from_footprints on the in-memory wcs objects to compute the combined wcs for limiting the GAIA query (at the moment I also see no tests for this feature). This required a minor backwards-compatible change in assign_wcs.util.
  • Refactor catalog/corrector and model linking to instead use the catalog/corrector index and model index. These previously were linked by reference and as align_wcs requires that all correctors be in-memory this results in all input image_models also required to be in memory. Breaking this link will (eventually) allow the image_models to be loaded one-by-one instead of having them all in memory at the same time.
  • some minor changes to try and break out portions of TweakRegStep.process into functions to improve readability (and make unit-testing easier)
  • minor change to _write_catalog to return a filename instead of modify the input model

This PR adds a parametrized test for the custom catalog handling in tweakreg (marked as slow as it has many combinations, each taking ~2 seconds on my machine, this will still be run in the CI as tox passes the --slow option to pytest).

Regression tests run:
https://plwishmaster.stsci.edu:8081/blue/organizations/jenkins/RT%2FJWST-Developers-Pull-Requests/detail/JWST-Developers-Pull-Requests/1443/pipeline/195
with one unrelated failure which also appears on main:
https://plwishmaster.stsci.edu:8081/job/RT/job/JWST/2885/testReport/jwst.regtest/test_miri_lrs_slit_spec3/_stable_deps__test_miri_lrs_slit_spec3_user_wcs_shape1_s2d_/

Resolves https://jira.stsci.edu/browse/JP-3603

Closes #8434

Checklist for maintainers

  • added entry in CHANGES.rst within the relevant release section
  • updated or added relevant tests
  • updated relevant documentation
  • added relevant milestone
  • added relevant label(s)
  • ran regression tests, post a link to the Jenkins job below.
    How to run regression tests on a PR
  • Make sure the JIRA ticket is resolved properly

Copy link

codecov bot commented Apr 16, 2024

Codecov Report

Attention: Patch coverage is 66.38655% with 40 lines in your changes are missing coverage. Please review.

Project coverage is 58.01%. Comparing base (781e0e0) to head (80bb840).
Report is 3 commits behind head on master.

Files Patch % Lines
jwst/tweakreg/tweakreg_step.py 63.30% 40 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #8424      +/-   ##
==========================================
+ Coverage   57.93%   58.01%   +0.08%     
==========================================
  Files         387      387              
  Lines       38839    38811      -28     
==========================================
+ Hits        22502    22517      +15     
+ Misses      16337    16294      -43     

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

@braingram braingram force-pushed the tweakreg_model_iter branch 2 times, most recently from e77532e to 3944532 Compare April 16, 2024 20:17
@braingram braingram force-pushed the tweakreg_model_iter branch 2 times, most recently from b6f5959 to fe25dc2 Compare April 18, 2024 17:21
@braingram braingram marked this pull request as ready for review April 18, 2024 17:26
@braingram braingram requested a review from a team as a code owner April 18, 2024 17:26
@braingram braingram requested review from nden and mcara April 18, 2024 17:36
Copy link
Member

@mcara mcara left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are my comments for now.

jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tweakreg_step.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tweakreg_step.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
@braingram
Copy link
Collaborator Author

These are my comments for now.

Thanks for the quick response! I'll leave my responses to the other comments.

@braingram braingram force-pushed the tweakreg_model_iter branch 2 times, most recently from f9a27a9 to f5d49e4 Compare April 19, 2024 11:26
Copy link
Member

@mcara mcara left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To the best of my abilities, this looks fine.

@braingram
Copy link
Collaborator Author

Thanks for the review! Your input was helpful in understanding the code and improving this PR.

@braingram
Copy link
Collaborator Author

braingram commented Apr 19, 2024

To provide a small "teaser" to help illustrate the motivation for these changes the following shows memory used for a tweakreg run on an asn with 33 "members":
Screen Shot 2024-04-19 at 9 41 46 AM
The memory quickly grows and reaches ~4.5 GB heap size (due to ModelContainer opening all models and ModelContainer and tweakreg keeping them open throughout the run).

With this PR (and some major modifications to ModelContainer to allow it to update models on disk) the following graph shows processing the same association:
Screen Shot 2024-04-19 at 9 41 27 AM
with maximum ~1 GB heap size.

Copy link
Collaborator

@nden nden left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a few minor comments about the tests.

jwst/tweakreg/tests/test_tweakreg.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Outdated Show resolved Hide resolved
@braingram braingram force-pushed the tweakreg_model_iter branch 2 times, most recently from 0681bb5 to 0256da3 Compare April 22, 2024 10:56
Copy link
Collaborator

@mairanteodoro mairanteodoro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @braingram! Looks good to me.

I only have a couple of suggestions for improvement:

  • we should add docstring with the purpose of the test;
  • we should break down test_custom_catalog into smaller pieces.

jwst/tweakreg/tests/test_tweakreg.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Show resolved Hide resolved
jwst/tweakreg/tests/test_tweakreg.py Show resolved Hide resolved
@hbushouse hbushouse added this to the Build 11.0 milestone May 6, 2024
CHANGES.rst Show resolved Hide resolved
Copy link
Collaborator

@mairanteodoro mairanteodoro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @braingram! Looks good to me.

@braingram braingram marked this pull request as draft May 9, 2024 18:29
@braingram braingram changed the title reorder tweakreg to reduce iterations through models JP-3603: reorder tweakreg to reduce iterations through models May 15, 2024
@braingram braingram marked this pull request as ready for review May 15, 2024 20:19
@braingram
Copy link
Collaborator Author

@nden thanks for adding the regtest in #8477
I rebased this PR and ran the regtest with the added one passing.

@mcara I believe the changes in this PR should have little to no impact on #8476 but please let me know if this PR impacts that PR (or other planned work on tweakreg).

@hbushouse (or anyone else) any objection to merging this now that it's re-opened after rebasing following the other tweakreg changes?

@hbushouse
Copy link
Collaborator

@mcara @braingram @nden Is everyone agreed that this is ready to merge? No bad side effects with #8476 ?

@mcara
Copy link
Member

mcara commented May 17, 2024

Looks OK to me.

@hbushouse hbushouse merged commit 70bed1d into spacetelescope:master May 18, 2024
27 of 28 checks passed
@braingram braingram deleted the tweakreg_model_iter branch May 18, 2024 12:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Tweakreg memory usage improvements
5 participants