Probe sparse dict learning

src/ptychi/api/options/base.py

@@ -545,7 +545,9 @@ class ProbeOrthogonalizeIncoherentModesOptions(FeatureOptions):
 
     method: enums.OrthogonalizationMethods = enums.OrthogonalizationMethods.SVD
     """The method to use for incoherent_mode orthogonalization."""
-
+
+    sort_by_occupancy: bool = False
+    """Keep the probes sorted so that mode with highest occupancy is the 0th shared mode"""
 
 @dataclasses.dataclass
 class ProbeOrthogonalizeOPRModesOptions(FeatureOptions):
@@ -651,25 +653,26 @@ def get_non_data_fields(self) -> dict:
 
 @dataclasses.dataclass
 class SynthesisDictLearnProbeOptions(Options):
-
-    d_mat: Union[ndarray, Tensor] = None
+
+    enabled: bool = False
+    enabled_shared: bool = False
+
+    thresholding_type_shared: str = 'hard'
+    """Choose between 'hard' or 'soft' thresholding."""
+
+    dictionary_matrix: Union[ndarray, Tensor] = None
     """The synthesis sparse dictionary matrix; contains the basis functions 
     that will be used to represent the probe via the sparse code weights."""
 
-    d_mat_conj_transpose: Union[ndarray, Tensor] = None
-    """Conjugate transpose of the synthesis sparse dictionary matrix."""
-
-    d_mat_pinv: Union[ndarray, Tensor] = None
+    dictionary_matrix_pinv: Union[ndarray, Tensor] = None
     """Moore-Penrose pseudoinverse of the synthesis sparse dictionary matrix."""
 
-    probe_sparse_code: Union[ndarray, Tensor] = None
-    """Sparse code weights vector."""
+    sparse_code_probe_shared: Union[ndarray, Tensor] = None
+    """Sparse code weights vector for the shared modes."""
 
-    probe_sparse_code_nnz: float = None
+    sparse_code_probe_shared_nnz: float = None
     """Number of non-zeros we will keep when enforcing sparsity constraint on
-    the sparse code weights vector probe_sparse_code."""
-
-    enabled: bool = False
+    the SHARED sparse code weights vector sparse_code_probe_shared."""
 
 @dataclasses.dataclass
 class PositionCorrectionOptions(Options):

src/ptychi/api/options/lsqml.py

@@ -100,15 +100,20 @@ class LSQMLObjectOptions(base.ObjectOptions):
     propagation always uses all probe modes regardless of this option.
     """
 
+@dataclasses.dataclass
+class LSQMLProbeExperimentalOptions(base.Options):
+    sdl_probe_options: base.SynthesisDictLearnProbeOptions = dataclasses.field(default_factory=base.SynthesisDictLearnProbeOptions)
+
 
 @dataclasses.dataclass
 class LSQMLProbeOptions(base.ProbeOptions):
     optimal_step_size_scaler: float = 0.9
     """
     A scaler for the solved optimal step size (beta_LSQ in PtychoShelves).
     """
-
-
+    experimental: LSQMLProbeExperimentalOptions = dataclasses.field(default_factory=LSQMLProbeExperimentalOptions)
+
+
 @dataclasses.dataclass
 class LSQMLProbePositionOptions(base.ProbePositionOptions):
     pass

src/ptychi/api/task.py

@@ -192,9 +192,11 @@ def build_probe(self):
             self.probe_options.experimental.deep_image_prior_options.enabled
         ):
             self.probe = probe.DIPProbe(**kwargs)
-        elif (
-            isinstance(self.probe_options, api.options.PIEProbeOptions)
-        ) and (
+        elif ( 
+              isinstance(self.probe_options, api.options.PIEProbeOptions)
+              or
+              isinstance(self.probe_options, api.options.LSQMLProbeOptions)
+              ) and (
             self.probe_options.experimental.sdl_probe_options.enabled
         ):
             self.probe = probe.SynthesisDictLearnProbe(**kwargs)