shared probe mode power fraction in descending order

[a4894z]

sometimes orthogonalization of the shared probe modes causes the 0th mode to not have the most power content; this fix ensures that the shared modes will always be sorted in descending order.

If the shared probe modes are not sorted in this way, updating the OPRs (which are defined only for the 0th probe mode) can have strange discontinuous behavior.

[Copilot]

Pull request overview

This PR addresses discontinuities in OPR updates by ensuring the dominant shared probe mode remains the 0th incoherent mode (highest power/occupancy), even after incoherent-mode orthogonalization.

Changes:

• Add an (opt-in) step to sort shared incoherent probe modes by occupancy/power in constrain_incoherent_modes_orthogonality().
• Introduce a new option flag ProbeOrthogonalizeIncoherentModesOptions.sort_by_occupancy to control this behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 6 comments.

File	Description
src/ptychi/data_structures/probe.py	Adds occupancy-based sorting of shared incoherent modes prior to orthogonalization.
src/ptychi/api/options/base.py	Adds the sort_by_occupancy configuration option for incoherent-mode orthogonalization.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

shared_occupancy is normalized by torch.sum(abs(probe[0])**2); if the probe is all zeros (or extremely small), this will produce NaNs/Infs and make the sort order unstable. Since you only need a descending ordering, sort by unnormalized per-mode power (and handle the zero-total-power case) to avoid divide-by-zero.

Suggested change

	shared_occupancy = torch.sum(torch.abs(probe[0, ...]) ** 2, (-2, -1)) / torch.sum(
	torch.abs(probe[0, ...]) ** 2
	)
	shared_occupancy = torch.sort(shared_occupancy, dim=0, descending=True)
	probe[0, ...] = probe[0, shared_occupancy[1], ...]
	# Sort incoherent modes by their (unnormalized) power to avoid
	# potential divide-by-zero when total power is zero or extremely small.
	per_mode_power = torch.sum(torch.abs(probe[0, ...]) ** 2, dim=(-2, -1))
	# If all modes have zero power, skip reordering as they are equivalent.
	if torch.any(per_mode_power != 0):
	sorted_power, sort_indices = torch.sort(per_mode_power, dim=0, descending=True)
	probe[0, ...] = probe[0, sort_indices, ...]

[Copilot]

This change introduces a new behavior flag (sort_by_occupancy) but there’s no unit test ensuring that enabling it actually keeps the 0th incoherent mode as the highest-power mode (and that it remains true after orthogonalization). Please add a small Probe unit test that constructs two incoherent modes with known power ordering, enables orthogonalize_incoherent_modes.sort_by_occupancy, runs constrain_incoherent_modes_orthogonality(), and asserts the ordering is descending.

[Copilot]

PR description says this fix “ensures that the shared modes will always be sorted in descending order”, but sort_by_occupancy defaults to False, so the behavior is unchanged unless callers explicitly enable it. Either make this enabled by default / apply the sorting unconditionally when incoherent-mode orthogonalization is enabled, or update the PR description/docs to reflect that it’s opt-in.

[mdw771]

Change this docstring to "If True, keep the probes sorted so that mode with highest occupancy is the 0th shared mode."