JP-3653: Nsclean speedup

[t-brandt]

Resolves JP-3653

Closes #8548

This PR improves the run time of the NSClean algorithm (not counting the time to construct a pixel mask in NIRSpec) by a factor between 10 and 20. This is achieved by representing a diagonal weight matrix as a vector and using vector multiplication and broadcasting to achieve equivalent results without extra multiplications by zero and additions of zero. The changes are entirely within NSClean.fit(). Nothing outside of this routine appears any different, and I have verified that the results of the calculation are the same.

Checklist for PR authors (skip items if you don't have permissions or they are not applicable)

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

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.

Project coverage is 58.52%. Comparing base (b7e0b10) to head (52ebaf5). Report is 337 commits behind head on master.

Files with missing lines	Patch %	Lines
jwst/nsclean/lib.py	0.00%	3 Missing :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## master #8547 +/- ## ========================================== + Coverage 58.02% 58.52% +0.50% ========================================== Files 388 388 Lines 38977 38958 -19 ========================================== + Hits 22617 22802 +185 + Misses 16360 16156 -204 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[drlaw1558]

@t-brandt Looks good to me; can you please add a change log entry? @hbushouse Can you please kick off a regression test run for this?

[t-brandt]

Ok, I updated CHANGES.rst.

[drlaw1558]

Thanks- minor comment that the number of dashes underneath the title 'nsclean' should be the same number of characters as the step name. I'm not honestly sure why, but @hbushouse has frequently had to remind me about it.

[hbushouse]

Thanks- minor comment that the number of dashes underneath the title 'nsclean' should be the same number of characters as the step name. I'm not honestly sure why, but @hbushouse has frequently had to remind me about it.

The number of underline characters needs to be at least as long as the title above it, but can be longer (without doing harm). It just can't be shorter or the sphinx parser will barf.

[t-brandt]

Ok, I adjusted the number of dashes.

[braingram]

I tried to run regtests for this PR: https://plwishmaster.stsci.edu:8081/blue/organizations/jenkins/RT%2FJWST-Developers-Pull-Requests/detail/JWST-Developers-Pull-Requests/1513/tests Unfortunately there are many errors that look unrelated (mostly in fgs related tests) but I'm not sure what's causing them. Here's a snippet from the traceback:

        # Make the query
        params = {'pattern': pattern}
        with requests.get(search_url, params=params, headers=headers) as r:
>           url_paths = r.json()['files']
E           KeyError: 'files'
/data1/jenkins/workspace/RT/JWST-Developers-Pull-Requests/clone/jwst/regtest/regtestdata.py:500: KeyError

[melanieclarke]

It looks like these changes are only in the NSClean class. Is it possible to make similar changes in the NSCleanSubarray class too? I'm wondering if this would enable support for the ALLSLITS subarray - right now, it uses so much memory that it crashes the kernel if attempted, so support for this subarray is explicitly turned off.

[t-brandt]

@melanieclarke The NSCleanSubarray already does something close to the changes I made. Unfortunately, I think that the reason it fails is https://github.com/spacetelescope/jwst/blob/41f24b7a4ade69a77d294bc759dd07a80c28d09f/jwst/nsclean/lib.py#L471 This results in a matrix of size npix x npix (see https://github.com/spacetelescope/jwst/blob/41f24b7a4ade69a77d294bc759dd07a80c28d09f/jwst/nsclean/lib.py#L465) If a subarray is 1000x100 pixels, for example, then Line 471 referenced above produces a matrix with 1e10 elements, which may not fit in RAM.

This problem does not affect the full array version of NSClean because that operates row-by-row, so that there are never more than 2048 pixels to be fit simultaneously, meaning that the equivalent matrix is at most 2048^2: https://github.com/spacetelescope/jwst/blob/41f24b7a4ade69a77d294bc759dd07a80c28d09f/jwst/nsclean/lib.py#L163

This array size/memory issue in NSCleanSubarray can be fixed, but not quite so simply as with the pull request here. If it is a priority for the mission I can spend a bit more time to propose a fix.

[drlaw1558]

Maybe I'm missing something, but can we simply swap subarray mode to run row-by-row as well?

[t-brandt]

@drlaw1558 Yes, we can do that. The best thing is probably to do something intermediate: to process as many rows simultaneously as we think we can profitably fit the Fourier modes. Assuming I am interpreting Bernie's code right, line https://github.com/spacetelescope/jwst/blob/41f24b7a4ade69a77d294bc759dd07a80c28d09f/jwst/nsclean/lib.py#L336 defines the cutoff and cuton frequencies as ~50 kHz and ~1 kHz, which would correspond to 2 pixels and 100 pixels. So that would imply that doing a few rows at a time would be best, I think, with the number depending on subarray size. I don't think this would be a ton of work to implement, but I think it would be good to consult with Bernie to confirm my intuition. We might then compute backgrounds over overlapping regions and then use weighted averages to construct a smooth combined background.

[drlaw1558]

@t-brandt Sounds like that's sufficiently different that it makes sense to consider separately rather than tying it to the work currently here. I'll file a separate ticket to track it; https://jira.stsci.edu/browse/JP-3654

[hbushouse]

Another regtest run started at https://plwishmaster.stsci.edu:8081/job/RT/job/JWST-Developers-Pull-Requests/1524

[hbushouse]

Regtest run https://plwishmaster.stsci.edu:8081/job/RT/job/JWST-Developers-Pull-Requests/1525/ had no failures or differences, so this looks good.