probe power constraint scaling and fft using backward normalization #39
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
fft scaling, the power the probe gets scaled to is P / N^2 (where the array size of the probe is N x N) instead of P. This small fix should account for this; alternatively we can initialize the probe power constraint to be P * N^2 and not use this small fix, but that seems like a clunky way of doing things?
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR fixes the probe power constraint by compensating for the unnormalized FFT in the forward propagator, ensuring the probe’s power remains at the intended value.
- Divide the summed squared magnitude by the number of elements to correct FFT scaling
- Recompute
power_correctionusing the normalized probe power
Comments suppressed due to low confidence (2)
src/ptychi/data_structures/probe.py:369
- Add an inline comment explaining that dividing by the number of elements compensates for the FFT’s lack of normalization, preserving the intended total probe power.
propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) / self.data.size().numel()
src/ptychi/data_structures/probe.py:369
- Add or update unit tests to verify that
constrain_probe_powercorrectly normalizes the probe power after FFT propagation, covering both the old and new behavior.
propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) / self.data.size().numel()
src/ptychi/data_structures/probe.py
Outdated
| propagated_probe = propagator.propagate_forward(probe_composed) | ||
| propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) | ||
| power_correction = torch.sqrt(self.probe_power / propagated_probe_power) | ||
| propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) / self.data.size().numel() |
There was a problem hiding this comment.
Consider using propagated_probe.numel() instead of self.data.size().numel() to clearly reference the propagated tensor’s element count and avoid possible mismatches if shapes differ.
Suggested change
| propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) / self.data.size().numel() | |
| propagated_probe_power = torch.sum(propagated_probe.abs() ** 2) / propagated_probe.numel() |
… to `FourierPropagator.norm`
Collaborator
|
Sorry it slipped out of my sight..merging it now |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
When using the probe power constraint, because of the way PtyChi does fft scaling, the power the probe gets scaled to is P / N^2 (where the array size of the probe is N x N) instead of P.
This small fix should account for this; alternatively we can initialize the probe power constraint to be P * N^2 and not use this small fix, but that seems like a clunky way of doing things?
Another way to account for this is to not use the propagate_forward() part of the code, directly compute the probe norm e.g. torch.sum(probe_composed.abs() ** 2), and then rescale the probe and sample using this power correction?