Direct Ptycho Holiday Cleanup

[gvarnavi]

This PR introduces a unified API for the following 5 deconvolution_kernels:

• 'single-sideband' / 'ssb' / 'aberration-corrected-bright-field' / 'acbf'

$$ \tilde{\phi}_{\text{SSB}}(\boldsymbol{q}) = \sum_{\boldsymbol{k}} \frac{\Gamma^{*}(\boldsymbol{q}, \boldsymbol{k}) \, G(\boldsymbol{q}, \boldsymbol{k})} {\lvert \Gamma(\boldsymbol{q}, \boldsymbol{k}) \rvert} $$

• 'optimum-bright-field' / 'obf'

$$ \tilde{\phi}_{\text{OBF}}(\boldsymbol{q}) = \frac{\sum_{\boldsymbol{k}} \Gamma^{*}(\boldsymbol{q}, \boldsymbol{k})\, G(\boldsymbol{q}, \boldsymbol{k})} {\sqrt{\sum_{\boldsymbol{k}} \lvert \Gamma(\boldsymbol{q}, \boldsymbol{k}) \rvert^{2}}} $$

• 'matched-filter' / 'mf' (Note this is a k-loop variant of WDD)

$$ \tilde{\phi}_{\text{MF}}(\boldsymbol{q}) = \frac{\sum_{\boldsymbol{k}} \Gamma^{*}(\boldsymbol{q}, \boldsymbol{k})\, G(\boldsymbol{q}, \boldsymbol{k})} {\sum_{\boldsymbol{k}} \lvert \Gamma(\boldsymbol{q}, \boldsymbol{k}) \rvert^{2} + \epsilon(\boldsymbol{q})} $$

• 'parallax' / 'prlx' / 'tilt-corrected-bright-field' / 'tcbf'

$$ \tilde{\phi}_{\text{prlx}}(\boldsymbol{q}) = \sum_{\boldsymbol{k}} G(\boldsymbol{q}, \boldsymbol{k}) \, \mathrm{sgn}[\sin \chi(\boldsymbol{q})] \, \mathrm{e}^{-\mathrm{i}\nabla_{\boldsymbol{k}} \chi(\boldsymbol{k}) \cdot \boldsymbol{q}} $$

• 'center-of-mass' / 'icom'

$$ \tilde{\phi}_{\text{iCOM}}(\boldsymbol{q}) = \sum_{\boldsymbol{k}} \frac{\boldsymbol{q} \cdot \bigl[ G(\boldsymbol{q}, \boldsymbol{k})\, \boldsymbol{k} \bigr]} {\mathrm{i}\, \lvert \boldsymbol{q} \rvert^{2}} $$

In addition, the PR does the following:

• Generalized fitting / optimization to now return a HyperparameterState with knowledge of the optimized / fixed parameters, e.g.

HyperparameterState(
  fixed_aberrations={'defocus': 500},
  optimized_aberrations={'C12': 0.10539382682653327, 'phi12': -0.5827360245016954},
  optimized_rotation_angle=-0.2625944124075851,
)

• Added simple StreamingParallax implementation -- mostly for reference, see writeup here Parallax UDF LiberTEM/LiberTEM#1788

• Added ability to estimate the spectral SNR and Fourier shell correlation (FSC) using two half-set reconstructions by check-boarding the BF mask. It seems to mostly agree with the analytical expressions on white-noise objects

• Updated tutorial notebooks updating notebooks with latest refactor quantem-tutorials#2 ready to be merged once this is merged

[gvarnavi]

Alrighty @smribet, ready for review!

Least-squares aberration fitting for direct ptychography

Finally added a new least-squares–based aberration fitting routine for direct ptychography.

API changes

• The previous method fit_hyperparameters has been renamed to fit_hyperparameters_cross_correlation.
• A new method, fit_hyperparameters_least_squares, is introduced.

Basic usage

direct_ptycho = direct_ptycho.fit_hyperparameters_least_squares(
    aberration_coefs = {"C10":0,"C12":0,"phi12":0},
    # rotation_angle=None,
    cartesian_basis="low_order",
    # fit_method="recursive",
    q_signal_weight=0.0,
    num_q_modes=12,
    # unwrap_method = "reliability-sorting",
    # two_pass_unwrapping = False,
    verbose=True,
)

Argument notes

• aberration_coefs
  Initial aberration guess, which is explicitly deconvolved from the measured data prior to fitting.
  If None, the current values stored in HyperparameterState are used.
  In the example above, zeros are passed explicitly to demonstrate robustness to poor initial guesses.

• rotation_angle
  Fixed value during least-squares fitting.
  If provided, it overrides the stored rotation angle, but is not optimized by this routine.

• cartesian_basis and fit_method
  Control which aberrations are fitted and how the fit is staged.

  • If cartesian_basis is a list of lists of strings, it specifies an explicit fitting schedule.
    Example:

    [["C10"], ["C10", "C30"], ["C50"]]

    Fits C10, then refines with C10 + C30, then refines again with C50.

  • If cartesian_basis is a list of strings:

    • fit_method="global":
      Fit all specified aberrations in a single solve.
    • fit_method="recursive":
      Fit aberrations in cumulative rounds of increasing radial order.
      Example:

      ["C10", "C12_a", "C12_b", "C21_a", "C21_b"]

      → two rounds:
      ["C10", "C12_a", "C12_b"], then
      ["C10", "C12_a", "C12_b", "C21_a", "C21_b"].
    • fit_method="sequential":
      Fit aberrations in exclusive rounds of increasing radial order.
      The same list above yields:
      ["C10", "C12_a", "C12_b"], then
      ["C21_a", "C21_b"].

• q_signal_weight and num_q_modes
  Control which spatial frequencies are used in the fit.

  • Spatial frequencies are ranked using a weighted combination of:
    • proximity to the BF edge, and
    • measured signal power.
  • q_signal_weight → 0 favors proximity to the BF edge.
    q_signal_weight → ∞ favors signal power.
  • num_q_modes selects how many spatial frequencies are retained after ranking.

• unwrap_method
  Torch-native phase unwrapping:

  • "reliability-sorting": port of the scikit-image algorithm (more robust, slower).
  • "poisson": FFT-based least-squares unwrapping (faster, but tends to underestimate phase).

• two_pass_unwrapping
  If True, performs a second masked unwrapping pass to improve robustness in partially connected overlap regions.

Tutorial notebooks have been updated to reflect these changes electronmicroscopy/quantem-tutorials#2

[gvarnavi]

@smribet See new tutorial notebook on hyperparameter optimization!