Added dictionary learning (DL) functionality to LSQML, cleaned up tensor operations in RPIE + DL #42

a4894z · 2025-08-05T19:18:29Z

Features/fixes

Added dictionary learning (DL) functionality to LSQML, and cleaned up the rPIE + DL calculations for tensor products using torch.einsum. This pull request is only for independent step length calculation for the sample and sparse code update directions; the joint solution for step length calculation for the sample and sparse code update directions is coming soon.

Related issues (optional)

n/s

Mentions

Checklist

Have you...

Formatted your code properly adhering to PEP-8? Considering using RUFF to format your code automatically.
Nope, need to finally learn how to do this
Resolved all merge conflicts with the main branch?
[ ]
Checked the diffs between your branch and the main branch to ensure that your changes are not introducing any regressions or unnecessary/irrelevant changes?

optimal step for sparse code update using uncoupled object and probe step calculation. after that, do the coupled step length calc

changes to rPIE and synthesissparseprobe class with cleaner tensor products using torch.einsum

Copilot

Pull Request Overview

This pull request adds dictionary learning (DL) functionality to the LSQML reconstructor and refactors tensor operations in both PIE and LSQML reconstructors using torch.einsum for improved clarity and efficiency. The implementation focuses on independent step length calculation for sample and sparse code update directions.

Added complete dictionary learning support to LSQML with optimal step size calculation for sparse code updates
Refactored tensor operations in PIE reconstructor using torch.einsum for better readability and performance
Added configuration option for averaging sparse codes over scan positions

Reviewed Changes

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

Show a summary per file

File	Description
src/ptychi/reconstructors/pie.py	Refactored sparse code update calculations using `torch.einsum` and improved tensor reshaping logic
src/ptychi/reconstructors/lsqml.py	Added complete dictionary learning functionality with optimal step length calculation for sparse code updates
src/ptychi/data_structures/probe.py	Updated sparse code weight computation and probe generation to use `torch.einsum` operations
src/ptychi/api/task.py	Extended probe building logic to support LSQML with dictionary learning options
src/ptychi/api/options/lsqml.py	Added experimental options structure for LSQML probe configuration
src/ptychi/api/options/base.py	Added `use_avg_spos_sparse_code` configuration option for sparse code averaging behavior

Copilot · 2025-08-05T19:19:32Z

src/ptychi/reconstructors/lsqml.py

+                # if self.parameter_group.probe.use_avg_spos_sparse_code:
+                #     delta_p_i = torch.tile( delta_p_i, ( n_spos, 1, 1 ) )


This commented-out code should either be implemented or removed. Leaving commented code in production reduces maintainability and clarity.

Suggested change

# if self.parameter_group.probe.use_avg_spos_sparse_code:

# delta_p_i = torch.tile( delta_p_i, ( n_spos, 1, 1 ) )

Copilot · 2025-08-05T19:19:33Z

src/ptychi/reconstructors/pie.py

-                    sel = sparse_code_sorted[0][self.parameter_group.probe.probe_sparse_code_nnz, :]
+                    sel = sparse_code_sorted[0][..., self.parameter_group.probe.probe_sparse_code_nnz]

+                    #(TODO: soft thresholding option as default?)


TODO comments should include more specific information about implementation timeline, requirements, or be converted to proper issue tracking.

Suggested change

#(TODO: soft thresholding option as default?)

# TODO: Consider implementing soft thresholding as the default option for enforcing sparsity.

# See issue #123 on GitHub (https://github.com/AdvancedPhotonSource/pty-chi/issues/123) for requirements and discussion.

Copilot · 2025-08-05T19:19:33Z

src/ptychi/reconstructors/pie.py

                if (self.parameter_group.probe.representation == "sparse_code"):
-                    # TODO: move this into SynthesisDictLearnProbe class
+
+                    # TODO: move these into SynthesisDictLearnProbe class           


TODO comments should include more specific information about implementation timeline, requirements, or be converted to proper issue tracking.

Suggested change

# TODO: move these into SynthesisDictLearnProbe class

# TODO: Move these calculations into SynthesisDictLearnProbe class for better modularity.

# See issue tracker: https://github.com/AdvancedPhotonSource/pty-chi/issues/XXX

# Target: Refactor by Q3 2025. Requirements: Move rc, n_scpm, n_spos, obj_patches_conj, conjT_i_delta_exwv_i, and related logic.

Copilot · 2025-08-05T19:19:33Z

src/ptychi/reconstructors/lsqml.py

+                                     torch.conj( obj_patches_vec ), 
+                                     numer)          
+
+                # real is used to throw away small imag part due to numerical precision errors                    


The comment about using .real to handle numerical precision errors should be more specific about when this occurs and potential alternatives.

Suggested change

# real is used to throw away small imag part due to numerical precision errors

# In theory, numer/denom should be real, but small imaginary parts can arise due to floating-point

# precision errors in complex arithmetic. We use .real to discard these. Alternatively, one could use

# torch.real_if_close or check that the imaginary part is negligible before discarding it.

define a separate representation = 'sparse_code' to the OPR options since I'm using the shared probe update representation currently.

mdw771 · 2025-08-14T17:44:55Z

src/ptychi/data_structures/opr_mode_weights.py

        chi: Tensor,
        delta_p_i: Tensor,
        delta_p_hat: Tensor,
+        probe_current_slice: Tensor,


Make this argument optional since it is only used by SDL. Default to None:

probe_current_slice: Optional[Tensor] = None

mdw771 · 2025-08-14T17:46:48Z

src/ptychi/data_structures/opr_mode_weights.py

        chi: Tensor,
        delta_p_i: Tensor,
        delta_p_hat: Tensor,
+        probe_current_slice: Tensor,


This should be an optional argument too

mdw771 · 2025-08-14T19:15:55Z

src/ptychi/data_structures/opr_mode_weights.py

        if batch_size == 1:
            return

+        update_eigenmode = probe.optimization_enabled(current_epoch)


Let's not put this check here. The parent function update_variable_probe already checks it. update_opr_probe_modes_and_weights is only called when probe.optimization_enabled() returns True.

mdw771 · 2025-08-14T19:21:20Z

src/ptychi/data_structures/opr_mode_weights.py

+
+        if (probe.representation == "sparse_code") and update_eigenmode:
+
+            sz = delta_p_i.shape


Let's make things inside the if block a separate method, something like update_opr_mode_sparse_code, and move it to SynthesisDictLearnProbe, then just call probe.update_opr_mode_sparse_code(...) here.

mdw771 · 2025-08-14T19:23:51Z

src/ptychi/data_structures/opr_mode_weights.py

+
+            probe_current_slice_vec = torch.reshape( probe_current_slice[:,0,...], (Nspos, rc) ).T
+
+            #==================================================================================


Please keep comments succinct; avoid using ASCII bars to divide segments. Major segments in a code block should be separated into different child functions/methods instead of by bars.

mdw771 · 2025-08-14T19:39:53Z

src/ptychi/data_structures/opr_mode_weights.py


-        if probe.optimization_enabled(current_epoch):
-            probe.set_data(probe_data)
+        if probe.optimization_enabled(current_epoch):   # and not (probe.representation == "sparse_code")


You don't need to set it again, so you can uncomment the and not

mdw771 · 2025-08-14T19:49:06Z

src/ptychi/data_structures/probe.py

        sparse_code_probe = self.get_sparse_code_weights()
        self.register_parameter("sparse_code_probe", torch.nn.Parameter(sparse_code_probe))

+        use_avg_spos_sparse_code = self.options.experimental.sdl_probe_options.use_avg_spos_sparse_code


Since it's just a boolean you don't have to register it as buffer. Also, instead of setting it as an attribute in __init__, let's just reference it from self.options on the fly in the method where it is used. This way if the user changes the value in the options object in the middle of a reconstruction, the new value can take effect dynamically.

mdw771 · 2025-08-14T19:52:33Z

src/ptychi/reconstructors/lsqml.py

-            )
+            if (self.parameter_group.probe.representation == "sparse_code"):
+
+                rc = chi.shape[-1] * chi.shape[-2]


I'd also make everything side the if a separate method and call it here. Also avoid using ASCII dividers.
If the routines for sparse code update is the same across all reconstructors or if there is at least something common between them, please put the common parts in SynthesisDictLearnProbe.

mdw771 · 2025-08-14T19:53:51Z

src/ptychi/reconstructors/lsqml.py

+                #======================================================================                              
+                # sparse code update directions vs scan position and shared probe modes  
+
+                obj_patches_slice_i_conj = torch.conj( obj_patches[:, i_slice, ...] )


Please remove extra space in the brackets. Make code more compact by removing unnecessary blank spaces.

mdw771 · 2025-08-14T20:06:53Z

src/ptychi/reconstructors/pie.py

+
                if (self.parameter_group.probe.representation == "sparse_code"):
-                    # TODO: move this into SynthesisDictLearnProbe class
+


Separate method for all everything inside the if. If the routines for sparse code update is the same across all reconstructors or if there is at least something common between them, please put the common parts in SynthesisDictLearnProbe.

implementing the other suggestions from Ming

mdw771 · 2025-08-28T15:56:48Z

src/ptychi/data_structures/opr_mode_weights.py

+            # Start from the second OPR mode which is the first after the main mode - i.e., the first eigenmode.
+            for i_opr_mode in range(1, probe.n_opr_modes):
+                # Just take the first incoherent mode.
+                eigenmode_i = probe.get_mode_and_opr_mode(mode=0, opr_mode=i_opr_mode)


Here you can regenerate the eigenmodes using update SDL coefficients (probe.generate() to update the probe.data attribute) so that you can reuse the existing code

mdw771 · 2025-08-28T15:57:34Z

src/ptychi/data_structures/opr_mode_weights.py

+        update_eigenmode = probe.optimization_enabled(current_epoch)        # why is this needed again? To even get into this function, we need this to already be true?
+        update_eigenmode_weights = self.eigenmode_weight_optimization_enabled(current_epoch)
+
+        if self.options.use_optimal_update:


Split SDL and optimal OPR weight and eigenmode updates into 2 pull requests

mdw771 · 2025-08-28T15:57:50Z

src/ptychi/data_structures/opr_mode_weights.py

+            Ws = (weights_data[ indices, 1:]).to(torch.complex64)
+
+            Tsconj_chi = (obj_patches[:,0,...].conj() * chi[:,0,...])
+            Tsconj_chi = adjoint_shift_probe_update_direction( indices, Tsconj_chi[:,None,...], first_mode_only=True)


This one should be replaceable by delta_p_i

mdw771 · 2025-08-28T15:58:27Z

src/ptychi/data_structures/opr_mode_weights.py

+            Tsconj_chi = (obj_patches[:,0,...].conj() * chi[:,0,...])
+            Tsconj_chi = adjoint_shift_probe_update_direction( indices, Tsconj_chi[:,None,...], first_mode_only=True)
+
+            chi = adjoint_shift_probe_update_direction( indices, chi, first_mode_only=True)


Probably ignorable because we have been mixing shifted and unshifted variables before and it works. In that case the adjoint_shift_probe_update argument can be removed

Ashish Tripathi added 3 commits July 31, 2025 13:37

infrastructure for LSQML + sDL

3c1aeb7

calculation of sparse code update direction, next step is computing

cb2f641

optimal step for sparse code update using uncoupled object and probe step calculation. after that, do the coupled step length calc

independent optimal step length calc for LSQML sample and sparse code;

ec900f7

changes to rPIE and synthesissparseprobe class with cleaner tensor products using torch.einsum

a4894z requested review from Copilot and mdw771 August 5, 2025 19:18

Copilot AI reviewed Aug 5, 2025

View reviewed changes

Ashish Tripathi added 2 commits August 6, 2025 11:11

lsqml correct step length

f0ed9dd

added dictionary learning to the OPR mode updates, I still need to

519b961

define a separate representation = 'sparse_code' to the OPR options since I'm using the shared probe update representation currently.

AdvancedPhotonSource deleted a comment from Copilot AI Aug 14, 2025

mdw771 requested changes Aug 14, 2025

View reviewed changes

Ashish Tripathi added 3 commits August 19, 2025 17:02

Work in progress reflecting some of the changes Ming requested, still

91773fe

implementing the other suggestions from Ming

style changes and more control logic tweaking

5979d66

refactor of sparse code + dictionary learning shared and OPR mode update

be8ea9a

a4894z requested a review from mdw771 August 27, 2025 20:53

updated a doc string

42b5187

mdw771 requested changes Aug 28, 2025

View reviewed changes

a4894z mentioned this pull request Sep 8, 2025

Probe sparse dict learning #45

Open

3 tasks

		# if self.parameter_group.probe.use_avg_spos_sparse_code:
		# delta_p_i = torch.tile( delta_p_i, ( n_spos, 1, 1 ) )

	#(TODO: soft thresholding option as default?)
	# TODO: Consider implementing soft thresholding as the default option for enforcing sparsity.
	# See issue #123 on GitHub (https://github.com/AdvancedPhotonSource/pty-chi/issues/123) for requirements and discussion.

-                    # TODO: move these into SynthesisDictLearnProbe class
+                    # TODO: Move these calculations into SynthesisDictLearnProbe class for better modularity.
+                    # See issue tracker: https://github.com/AdvancedPhotonSource/pty-chi/issues/XXX
+                    # Target: Refactor by Q3 2025. Requirements: Move rc, n_scpm, n_spos, obj_patches_conj, conjT_i_delta_exwv_i, and related logic.

-                # real is used to throw away small imag part due to numerical precision errors
+                # In theory, numer/denom should be real, but small imaginary parts can arise due to floating-point
+                # precision errors in complex arithmetic. We use .real to discard these. Alternatively, one could use
+                # torch.real_if_close or check that the imaginary part is negligible before discarding it.


		if (probe.representation == "sparse_code") and update_eigenmode:

		sz = delta_p_i.shape


		probe_current_slice_vec = torch.reshape( probe_current_slice[:,0,...], (Nspos, rc) ).T

		#==================================================================================


		if (self.parameter_group.probe.representation == "sparse_code"):
		# TODO: move this into SynthesisDictLearnProbe class

Added dictionary learning (DL) functionality to LSQML, cleaned up tensor operations in RPIE + DL #42

Are you sure you want to change the base?

Added dictionary learning (DL) functionality to LSQML, cleaned up tensor operations in RPIE + DL #42

Uh oh!

Conversation

a4894z commented Aug 5, 2025

Features/fixes

Related issues (optional)

Mentions

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants