Modify NSClean behaviour for subarray data

[stscijgbot-jp]

Issue JP-3654 was created on JIRA by David Law:

At present, the NSClean step does not work in subarray mode as it tries to handle too much data at the same time and runs out of memory (as opposed to the full array implementation, which goes line by line).  This ticket is to consider a rework of the implementation to allow for future application to subarray data.

See #8547 for some related discussion.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

To clarify: it does work for all subarrays except for ALLSLITS, which is significantly larger than the others.  NSClean is explicitly disabled when subarray = ALLSLITS, but it should work for all other NIRSpec modes.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Melanie Clarke That is true, but the computational cost can be prohibitive in a configuration like BOTS because of the implementation.  On, say, a 2048x30 pixel subarray, the matrices can be ~60,000 square as opposed to 512 square in a row-by-row mode on a full array.  This makes the NSClean step too expensive to run on these configurations when there are many integrations.  If this is refactored into something like 60 1,000 x 1,000 matrices, the speedup should be at least a factor of 60 (maybe even 60^2), making NSClean straightforward to run on spectra with many integrations.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Agreed, it is very slow for BOTS currently.  Better performance would be very welcome!

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

After consulting with Bernie Rauscher, the algorithm currently uses frequencies below 1 kHz to correct 1/f and frequencies of 50 kHz to correct the odd/even (or alternating columns) effect.  He is in agreement that operating in chunks is a viable way forward.  The chunks should be significantly larger than 100 pixels (corresponding to 1 kHz frequencies).  My suggestion is to use 1000 pixel chunks.  Bernie also mentioned that, if the chunks are disjoint, there can be edge effects that are apparent.  We agreed that this could be overcome using overlapping chunks with smoothly varying weights.  This would come at the cost of double the computational work, but because of the scaling of matrix operations with array size, it would still be far faster than the current implementation.  If all watchers are in agreement with this approach (~1000 pixel overlapping chunks for Fourier fitting) I will go ahead and implement it.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Tagging David Law and Melanie Clarke as a reminder–if the fix that Bernie and I discussed sounds good, I will proceed with implementation. 

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

To the extent that I understand it, that sounds fine to me.

Eddie Bergeron and I have discussed elsewhere (JP-3437, JP-3611) better ways to address alternating column noise, including eliminating a similar 50 kHz correction from IRS2 reference pixel handling in favor of a true column-based correction.  Tagging him here, in case he has opinions on keeping that correction in NSClean or not.

[stscijgbot-jp]

Comment by David Law on JIRA:

Timothy Brandt Sounds reasonable to me as well.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

I have a draft implementation that I would like to refine and test.  Do any of you, Christian Hayes, David Law, Melanie Clarke, have a suggestion for a small amount of data where I expect NSClean on subarray data to do something useful so that I can run comparisons?  Thanks!

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Sure - you can try jw02757002001_03104_00001_nrs1_rate.fits, which is one of the example data sets in the fixed slit NSClean notebooks.

[stscijgbot-jp]

Comment by Christian Hayes on JIRA:

Timothy Brandt, here is one in another subarray if its useful in addition to Melanie's example: jw02122004001_04102_00002_nrs1_rate.fits also has a clear 1/f pattern.  If you wanted some more cases I'm happy to pull a few others out.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Ok David Law, Christian Hayes, and Melanie Clarke, I have something that seems to work as well as the current implementation (results are generally close to identical, sometimes very slightly better, sometimes very slightly worse).  It runs much faster than the current implementation, taking just under 0.1 second for a BOTS exposure on my laptop (depending on the choice of a tunable parameter); this should scale to perhaps 5-10 seconds on a full frame treated as a subarray.  The corresponding numbers for the current implementation are about 3 seconds for a single BOTS exposure and unmeasurable for the full frame due to out of memory errors.  There should be no memory issues now.  

I have two questions before I start putting together a pull request.

The current implementation of NSClean is essentially absent from the unit tests.  I guess I can make a test for consistency between the new and old subarray approaches; I imagine that I should take this route and forget about the lack of unit tests in the rest of the module?

Part of my checks for this involved making series of plots using some extra code that I added to the routines; I have no intention of leaving this in for the PR version.  What documentation should I provide for the algorithmic change to NSClean and its effects on the data?

Thanks!

[stscijgbot-jp]

Comment by David Law on JIRA:

Timothy Brandt

1) If it's easy to add in a consistency test between old and new approaches please do.  I wouldn't dig too deeply into this at the moment though if we're going to potentially be looking at moving 1/f correction to detector1, which would presumably require new tests anyway.

2) Readthedocs is probably the best place for the algorithm change description, and can be included in the same PR (i.e., https://github.com/spacetelescope/jwst/blob/master/docs/jwst/nsclean/main.rst).  Looks like this is pretty bare bones at the moment so any information that you add will help.  Any plots and the like comparing performance should be attached to this ticket for future reference.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

I have attached comparison plots for BOTS data (from the NSClean notebook): 500 pixels at either end of the spectrum.  It looks good to me; the new approach arguably performs a tiny bit better than the current implementation.  I will try to write up a test and proceed with the PR.  Run time is about a factor of around 30 faster with the new implementation on 32x2048 subarrays and a factor of around 100 faster on 64x2048 subarrays (0.1-0.2 seconds in both cases with the new implementation vs. 3-5 seconds and 20-40 seconds with the current one).  The improvement in runtime will be a larger factor for larger subarrays.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Looks good to me too!  This performance fix will be definitely be necessary for running NSClean in detector1 on group data.

Do you have a branch I can test on?  I ran into an issue with some ringing at the edges of subarrays when running the existing NSClean with changes in JP-3639 - I'd like to see if your new implementation fixes that or if I need to work harder on how the masking is working for that mode.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Melanie Clarke Not yet; I will get on that shortly.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Okay, thanks.  No rush!

[stscijgbot-jp]

Comment by David Law on JIRA:

Thanks- those comparisons look good to me as well.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Melanie Clarke try it here:

https://github.com/t-brandt/jwst/tree/nsclean_subarray_speedup

The size of the chunks to use is a tunable parameter, and the current behavior can be recovered by setting this to infinity (or any number larger than the number of available pixels for fitting 1/f noise).  I have defaulted to 512.  I expect that making this number larger will have little impact on results and will slightly degrade the run time.  

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Thank you!  I'll check it out.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

I temporarily merged your work with mine for JP-3639 to check it out on some test data I've been working with.

A preliminary check with FS subarray file jw02757002001_03104_00001_nrs1_rate.fits looks good – the update naturally fixes the edge effects I was seeing with the older version.  

BOTS file jw01366003001_04101_00001-seg002_nrs1_rateints.fits also looks good - output is comparable to the previous version, and the new version of the step took 22s instead of 16m!

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Ok, I am almost done with a revision of the full frame version of NSClean to operate fairly differently.  I will expose some additional free parameters via optional arguments to the function.  At this point, I should I add these modifications to the present branch/fork, or should I make a new branch?  The revisions to the full frame version of NSClean are important to get acceptable behavior in general in full frame images (c.f. https://jira.stsci.edu/browse/JP-3639).  It builds on the changes to the subarray version of NSClean and uses that as its base, ultimately discarding the current full frame version, which is why I put it here (it is as much a generalization of the subarray version as anything).  

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

It seems fine to me to make the modifications for full frame on this branch.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Cross-posting from JP-3639:

I have redone the full frame NSClean algorithm here:

https://github.com/t-brandt/jwst/tree/nsclean_subarray_speedup

It isn't really ready for a PR yet–there are some things that really should go into the mask calculation that I should refactor.  This is a legacy from the current implementation, which also mixes the mask creation between various places.  I would be curious for people to try it though and tell me if it seems to work.  I can write up a more detailed document outlining the changes and motivations if desired.  Briefly, only the very lowest 1/f frequencies can be fit channel-by-channel without overfitting, while combining the four output channels gives you access to a wider range of frequencies.  We can access both sets of frequencies using the appropriate set(s) of pixels in one or in all readout channels.  Proper use of the reference pixels from something like SIRS reduces read noise by about 30% over not using them at all, while a 1/f correction, assuming ~half of the array is available for the correction, can get you another 20% or so, with a much larger reduction is the visually obvious low-frequency noise.  Computational cost should be a few seconds per full frame image, I am guessing 3s-10s depending on computer; this would apply to each group.  This is a bit slower than the current implementation of the full frame correction, but still much faster than the runtime before JP-3653.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Timothy Brandt - what is the status on the various changes for NSClean?  At this point, perhaps we should get the performance improvements for subarray mode in first, then we can finish testing and integration for full-frame changes on another ticket.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

The subarray speedup should be ready to go, I think.  The full frame version should, I think, also be ready to go (though there are a number of decisions I made in that implementation, and I have an "aggressiveness" parameter I intend to expose to the user).  The subarray version should not be used as-is for a 1/f correction.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Sorry, to clarify - which subarray version should not be used as-is?

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

Sorry–the subarray version should not be used to correct full-frame data in the 1/f correction step.  The subarray version should be good to go for subarray data.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Ah, okay, thanks!  Since there are separate decisions that were made for the full frame correction, rather than just unifying it with the subarray implementation, can you please send in a separate PR with just the subarray speed up for this ticket?  I think the science team might want to test the full frame differences more, but the subarray speed up should be more straightforward.

[stscijgbot-jp]

Comment by Timothy Brandt on JIRA:

I just submitted a PR for only the changes to subarray mode: #8745

I explicitly removed reference pixels from the fitting; this removed a lot of the remaining ringing.  Aside from that the code is as before.

[stscijgbot-jp]

Comment by Melanie Clarke on JIRA:

Fixed by #8745

[stscijgbot-jp]

Comment by Christian Hayes on JIRA:

I've tested this and it works as expected:  subarray data process through the nsclean step significantly faster and now allslits data can be processed (and reasonably quickly) through nsclean.  There are some very minor differences in the resulting data as Tim mentioned, but nothing concerning.