Resample speedup

[t-brandt]

Partially resolves JP-3997
Resolves JP-4216
Resolves RCAL-1311

Closes spacetelescope/jwst#10115

This PR supersedes #485 and #367

The reproject functions and calc_pixmap and calc_gwcs_pixmap functions in alignment and outlier_detection are all removed in favor of a copy of the calc_pixmap function from drizzle which now lives in alignment. The rewritten calc_pixmap accepts two new inputs: stepsize (duplicating an argument in drizzlepac), which performs the full WCS pixel mapping only every stepsize pixels in each direction, and order, either 1 or 3, to perform either bilinear or bicubic interpolation in between. These parameters should be exposed at the top-level calls to outlier detection and resample.

Tasks

• update or add relevant tests
• update relevant docstrings and / or docs/ page
• Does this PR change any API used downstream? (if not, label with no-changelog-entry-needed)
  • write news fragment(s) in changes/: echo "changed something" > changes/<PR#>.<changetype>.rst (see changelog readme for instructions)
  • run regression tests with this branch installed ("git+https://github.com/<fork>/stcal@<branch>")
    • jwst regression test
    • romancal regression test

[codecov]

Codecov Report

❌ Patch coverage is 95.53073% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.08%. Comparing base (75d09de) to head (8924d92).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/stcal/resample/utils.py	91.17%	6 Missing ⚠️
src/stcal/resample/resample.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #488      +/-   ##
==========================================
- Coverage   91.54%   89.08%   -2.47%     
==========================================
  Files          63       63              
  Lines        8715     8792      +77     
==========================================
- Hits         7978     7832     -146     
- Misses        737      960     +223     

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

[emolter]

The jwst repository calls gwcs_blot and the Resample class. So in order for this to be accessible to jwst, I think the two new parameters stepsize and order need to be propagated one more step and become optional inputs to those two.

For changelog fragments, you'll need a few different ones, something like the following:

• 488.alignment.rst: Rename calc_pixmap to calc_gwcs_pixmap, and add interpolation options to calc_gwcs_pixmap to speed up resample computations. (something about the amount of speedup here).
• 488.alignment.rst: Remove util.reproject method
• 488.outlier_detection.rst: Remove utils.reproject method. Remove utils.calc_gwcs_pixmap; use alignment.resample_utils.calc_gwcs_pixmap instead.
• 488.breaking.rst: Remove alignment.util.reproject, outlier_detection.utils.reproject, and outlier_detection.utils.calc_gwcs_pixmap. Rename alignment.util.calc_pixmap to alignment.util.calc_gwcs_pixmap

[t-brandt]

I added two changelog fragments attempting to match the current filenames and the spirit of your suggestions, I propagated stepsize and order up as needed, and I fixed the formatting. Tests are passing. This PR still needs unit tests of the interpolation aspect of calc_pixmap.

[emolter]

Thanks for doing so! It looks like I have what I need to ensure this looks good on multiple types of JWST data.

RE the changelog, what you have looks good. We are about to change some of the categories - this is part of what we were arguing about at hack day - so they unfortunatley may need to be renamed and split by step if #486 gets in before this does.

[emolter]

@t-brandt I started testing this by hooking it into the jwst repository (see this early PR draft). I think I found a small bug, which is that the output pixel map seems to be the wrong shape for non-square data.

For example, for the unit test jwst.outlier_detection.tests.test_outlier_step_no_outliers, the input data shape is (21,20), and this gets passed into calc_pixmap properly (data and shape parameters are consistent) but the output has shape (20,21,2) when it should be (21,20,2).

Changing the last line of calc_pixmap to have some transposes, i.e. pixmap = np.dstack([x.T, y.T]), allows the test to pass. (I didn't check if this was the mathematically correct fix.)

[t-brandt]

Thanks, @emolter -- the x and y order were reversed in one call. It didn't matter if they were the same as for a square array. I think it should be fixed now.

[emolter]

Fix looks like it works, thanks. I'm still testing this on different jwst datasets, will update when that's done.

[emolter]

I've now tested this on all jwst imaging regression test data. Things are looking excellent- the differences in output data are almost all smaller than our tolerances, and those that aren't are still small enough that I'm not worried. I've also confirmed that the runtime improvements can be substantial for large JWST mosaics. See comments on spacetelescope/jwst#10137 for details.

I'm ready to approve this pending two things:

• Regression tests against jwst main. Although I don't anticipate differences, we did consolidate calc_pixmap into one function, and the now-removed stcal version did have some small differences with the current version. I've kicked those off here: https://github.com/spacetelescope/RegressionTests/actions/runs/21004645391
• Additional unit test coverage, which I think is already on Tim's radar.

[emolter]

There are three failures in the regression tests (https://github.com/spacetelescope/RegressionTests/actions/runs/21004645391) pointing to jwst/main and this PR branch.

They look small so I'm not super worried, but we didn't expect any changes at all, so I think it'd be worth at least understanding where they are coming from. I presume it is indeed related to some difference between the old stcal calc_gwcs_pixmap and the new calc_pixmap.

test_nirspec_mos_spec3[v000000056-crf] had just one pixel in one of several sci extensions flip from a real value to NaN.
test_nirspec_mos_spec3[v000000056-s2d] had just four pixels change value, but some of those changes look large.
test_nirspec_fs_spec3_moving_target[crf] had 21 pixels flip from real values to NaN (also in just one extension), which is still only around 0.1% of pixels

Does anyone have insight here?

edit with more details:
For the failing SlitModel in test_nirspec_mos_spec3[v000000056-s2d], there's some noise at the numerical precision level. The plot below shows the diff between the science data in the PR branch vs our truth file. The green/yellow highlighted points are the 10 largest differences. These are:

Difference statistics:
  Shape: (31, 1291) (spatial x wavelength)
  Max absolute difference: 1.032e+00
  RMS difference: 6.677e-03

  Top 10 largest differences:
    1. [y= 29, x= 200]: -1.032e+00
    2. [y= 29, x= 199]: -5.989e-01
    3. [y= 28, x= 200]: -2.680e-01
    4. [y= 28, x= 199]: -1.142e-01
    5. [y=  5, x= 552]: +6.104e-05
    6. [y=  9, x= 553]: +3.052e-05
    7. [y= 10, x= 552]: -1.526e-05
    8. [y=  8, x= 553]: +7.629e-06
    9. [y=  9, x= 554]: +3.815e-06
    10. [y=  5, x= 554]: -3.815e-06

So the four points the regtest flagged (the four in the upper left) are indeed substantially larger differences than the rest. Still unclear to me what caused this.

For the failing second exposure in the MultiExposureModel in test_nirspec_fs_spec3_moving_target[crf], it looks like the newly-NaN pixels are all along one edge; see below, where I'm plotting the diff between the PR branch and the truth file and coloring the newly-NaN pixels in magenta.

[emolter]

Looking again at the analysis I did above, I think it's fine to chalk it up to just small numerical differences between the old and new calc_pixmap. I got an opinion from Melanie and she agrees. So I'm fine ok-ifying these whenever the PR goes in.

[stscieisenhamer]

As per discussion in rcalpr#2136, rcal is good-to-go with this PR.

[mcara]

With regard to "wiggling" of 3rd order splines... I still think this is a potential issue albeit, as of now, it is only theoretical. I performed extensive (oversampling each output pixel x10) search for these wiggling using WCS of several imaging instruments and I was not able to find any. I did not perform any testing on spectral WCS which is not possible at this moment. Attempt to test this on spectral WCS failed because resulting pixmap for spectral WCS is all NaN.

The reason for all NaN in spectral WCS is that often they have bounding boxes smaller than image shape and WCS outside of the bounding box is NaN. When coarse grid used to construct interpolating splines with RectBivariateSpline includes NaN, constructed splines have all coefficients set to NaN resulting in a full pixmap being all NaN.

So, my recommendation is to modify the code for the case of stepsize>1 to use only regions within the bounding box to construct interpolating polynomials and set reconstructed full pixmap values to NaN for pixels outside the bounding box.

I assume this is the reason why stepsize=1 is hardcoded in https://github.com/emolter/jwst/blob/65bec7d2cdcc35f1efa6380c203ca6d9cfd729fd/jwst/resample/resample_spec.py#L214-L215 in spacetelescope/jwst#10137 for spectral data. Otherwise it would simply fail.

[mcara]

Looking again at the analysis I did above, I think it's fine to chalk it up to just small numerical differences between the old and new calc_pixmap. I got an opinion from Melanie and she agrees. So I'm fine ok-ifying these whenever the PR goes in.

Yes, I believe the differences are specifically due to differences in the algorithm of calc_pixmap() from drizzle.utils (same as in this PR for the default stepsize=1) and the algorithm in calc_gwcs_pixmap() used only in gwcs_blot().

To be safe, if desired, I recommend making a separate PR that does only one thing: switch gwcs_blot() from calc_gwcs_pixmap() to drizzle.utils.calc_pixmap(). This is a very straightforward change. You should see the same failures as in this PR. You can then merge that PR, okify regression tests, and then rebase this PR. There should not be any failures.

[mcara]

However, regression tests run on this PR test only the default case of stepsize=1. There are no unit or regression tests to test the case of stepsize>1. I highly recommend adding unit tests for this PR even for the case of stepsize=1 since this PR chose to not reuse drizzle.utils.calc_pixmap() which is tested.

[mcara]

I left separate comments with some recommendations. However, existing unit test failures should be fixed and I recommend that new unit tests be added to test calc_pixmap for both stepsize == 1 (or use drizzle.utils.calc_pixmap() for stepsize=1) and stepsize > 1.

Also, IMO, for stepsize>1 bounding box must be taken into account.

[t-brandt]

@mcara, I fully agree that we need unit tests to verify behavior for stepsize>1. I am stretch rather thin, but I was hoping to put together a unit test with a realistic WCS, e.g. from NIRCAM, and verifying that the transformation is nearly identical with stepsize=10, and better with order=3 than with order=1. Do you have the time and expertise to write such a test? And for spectroscopic data, I see two paths forward. I'll try to take care of the bounding box issue but may not get to it for a week or two. Another choice would be to check whether the bounding box differs from the array shape, and if so, force stepsize=1.

[mcara]

One more question for everyone: Why is calc_pixmap living in alignment instead of resample module? I know that it is used by outlier detection but I do not see any connection to alignment.

CC: @braingram @emolter

[braingram]

@mcara, I fully agree that we need unit tests to verify behavior for stepsize>1. I am stretch rather thin, but I was hoping to put together a unit test with a realistic WCS, e.g. from NIRCAM, and verifying that the transformation is nearly identical with stepsize=10, and better with order=3 than with order=1. Do you have the time and expertise to write such a test? And for spectroscopic data, I see two paths forward. I'll try to take care of the bounding box issue but may not get to it for a week or two. Another choice would be to check whether the bounding box differs from the array shape, and if so, force stepsize=1.

@t-brandt I think it's best if you add any unit tests you'd like to this PR and @mcara can review them.

[braingram]

One more question for everyone: Why is calc_pixmap living in alignment instead of resample module? I know that it is used by outlier detection but I do not see any connection to alignment.

CC: @braingram @emolter

I agree with @mcara that resample seems like the most intuitive home for calc_pixmap.

[mcara]

Regression tests:

• romancal: https://github.com/spacetelescope/RegressionTests/actions/runs/21886414324
• jwst: https://github.com/spacetelescope/RegressionTests/actions/runs/21886394010

[emolter]

LGTM, one comment about the test coverage but it doesn't feel important enough to hold up the PR

[emolter]

@zacharyburnett would you please re-trigger the readthedocs build? It appears to have timed out and is required for merge. It looks to me like you might be the only one with permission in here

[zacharyburnett]

Sure, restarted

[emolter]

https://app.readthedocs.org/projects/stcal/builds/31372985/ shows the most recent build worked, so why hasn't the CI summary updated to reflect that?

Does anyone with permission feel like overriding the branch protections to merge this? @zacharyburnett @tapastro