Separate noise estimation from sample weight calculation

[ChaitanyaChawak]

Summary

• Separates noise estimation from sample-weight computation in train_utils so the per-observation noise sigma can be reused by the SHE validation workflow.
• Closes Separate noise estimation from sample weight computation #225

What’s changed

• Added train_utils.estimate_noise_sigma method which returns the per-observation noise standard deviation.
• Refactored calculate_sample_weights to call estimate_noise_sigma internally.
• Added private _is_masked_loss helper to help simplify calculate_sample_weights method.
• Added tests covering the new methods

How to test / verify

• Check if CI passes

Scope

Indicate the type of PR:

• Feature
• Bug fix
• Hotfix
• Documentation / process change
• Internal / refactor
• Release

This PR is part of the "Compatibility with Euclid SHE Validation" milestone.

Changelog

Did this PR introduce user-visible changes?
If yes, a Scriv changelog fragment must be added and committed.

• Changelog fragment added (if applicable)

Reviewer Checklist

Reviewers should confirm the following before approving and merging:

• The PR targets the correct base branch (develop, or main for release PRs)
• The PR is assigned to the developer
• Appropriate labels are applied
• The PR is included in relevant projects and/or milestones
• Description clearly explains what has changed
• Issue references included, if applicable
• Code and documentation adhere to current standards (ruff)
• Documentation updates included, if relevant
• CI tests are passing
• All reviewer comments have been addressed

Next Steps / Notes (if applicable)

Will use this abstracted out estimate_noise_sigma method for in the she_validation_adaptor required by the SHE Validation pipeline

[jeipollack]

Thanks, Chai, for your work on this PR. You implemented all the right ingredients to arrive at the correct outcome. There is just a question of the reusability of estimate_noise_sigma() as a general image-analysis tool, since it is currently coupled to the training logic.

One general point I'd like to ask you to keep in mind for this and future tasks: when the implementation exposes something that doesn't fit the conceptual responsibility of the code you're adding, stop and consider whether the design should change. Then make an intentional decision about whether that change belongs in the current scope.

If the improvement is small and directly related to the task, it can be addressed in the current PR. If it's larger or more loosely related, it can be captured as a follow-up issue. If it's necessary to implement the task cleanly, then the refactoring should be done as part of the work, even if that means separating it into a prerequisite PR.

This is the strategy I've been applying in the Quality Control Pipeline Milestone, and I think it has been working well.

[ChaitanyaChawak]

Thanks a lot @jeipollack for this thorough review! As suggested, I have moved all the NoiseEstimator utilities into a noise.py file, and added a separate estimate_noise_batch method. This method is now no longer depends on the loss or any training-specific elements. Thanks a lot for pointing this shortcoming in the original PR. I have also updated the tests for the updated parts, and the PR is now ready for a second review pass 🙏

[jeipollack]

Thanks @ChaitanyaChawak for the updates. Moving NoiseEstimator into noise.py and adding the batch estimation method is a good improvement. I think there are still a few places where we can simplify the separation of responsibilities further. I'm also getting better at spotting these things now than when I first began working on the project. I admit that I am partially responsible for these code bits, which contained mixed logic. 🙈

To start, I think the training-specific resolution of outputs should happen in general_train_cycle(), immediately before calling calculate_sample_weights(). The loss is only needed here to determine how the training outputs representation should be interpreted (i.e. whether to extract an image and mask). Once that has been resolved, calculate_sample_weights() doesn't need to know about the loss function at all.

This would give calculate_sample_weights() a much simpler contract along the lines of:

calculate_sample_weights(
    images,
    masks=None,
    apply_sigmoid=False,
    sigmoid_max_val=5.0,
    sigmoid_power_k=1.0,
)

I also think use_sample_weights is better handled by the caller. Rather than having calculate_sample_weights() decide whether it should calculate anything, general_train_cycle() could conditionally call it and otherwise leave sample_weight = None. That would give the function a single responsibility: calculating the weights from the supplied images/masks.

Finally, I wonder whether the construction of NoiseEstimator could encapsulate more of the noise-estimation configuration. At the moment the caller needs to know both how to extract the image dimensions and how the exclusion-window radius is calculated:

img_dim = (outputs.shape[1], outputs.shape[2])
win_rad = np.ceil(outputs.shape[1] / 3.33)
std_est = NoiseEstimator(img_dim=img_dim, win_rad=win_rad)

The window-radius calculation in particular feels like knowledge that belongs with NoiseEstimator, rather than with the training code. It may be worth considering whether the estimator can derive this from the image dimensions itself, while retaining the ability to configure the radius if needed.

Putting this all together, I think this would leave us with a clearer boundary: general_train_cycle() understands the training representation and loss; calculate_sample_weights() understands how to turn image noise into training weights; and NoiseEstimator owns the actual noise-estimation logic and its configuration. With this tidy-up, we may not even need estimate_noise_sigma method because the functionality is already handled by estimate_noise_per_batch, which I think you can re-use in your SHE Validation adapter.

Hope that makes sense! And, apologies for the back and forth. Just know I have such exchanges with Rory, too! ;-)

[ChaitanyaChawak]

yess, I think that makes a lot of sense! I've updated the calculate_sample_weights() function to make it much more simple and the training-specific resolution of the output is now done in general_train_cycle just before the sample weights are calculated.

I have also moved the window radius calculation in NoiseEstimator by introducing an optional win_rad parameter during init. The default value of this parameter is calculated in a newly introduced default_win_rad method using the provided img_dim.

Finally, as you correctly predicted, all these changes rendered estimate_noise_sigma obsolete and it has now been removed.

I've also updated the relevant sections of the train_utils_test.py to correctly test for the changes made in this PR and fix the CI pytest failure ✌️

[jeipollack]

Hi @ChaitanyaChawak, thanks for your positivity in our series of iterations. I spotted only 1-2 small code items that I think could be cleaned up. And noted now that I don't see a changelog fragment for this PR. we're also adding batch noise-estimation functionality to NoiseEstimator, so I think this is worth capturing in the changelog.

One these are completed I would be happy to approve this PR and merge.

[ChaitanyaChawak]

@jeipollack : Done!

[jeipollack]

Hi @ChaitanyaChawak, I looked over your latest commits, and left a couple of comments concerning the CHANGELOG and a unit test.

[ChaitanyaChawak]

@jeipollack thanks for the speedy review! I've made the changes acc to your suggestions 🙌