Fix a remaining bug in orthogonalize_probes and also fix some lingering issues with the liveplotting for unified mode fancy ptycho

src/cdtools/models/fancy_ptycho.py

@@ -635,7 +635,7 @@ class FancyPtycho(CDIModel):
         def get_probes(idx):
             basis_prs = self.probe * self.probe_support[..., :, :]
             prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
-                        axis=-4)
+                        axis=-3)
             ortho_probes = analysis.orthogonalize_probes_t(prs)
 
             if mode.lower() == 'amplitude':

src/cdtools/tools/analysis/analysis.py

@@ -164,11 +164,10 @@ def orthogonalize_probes_t(
         orthogonalized_probes = orthogonalized_probes.numpy()
         reexpressed_weight_matrix = reexpressed_weight_matrix.numpy()
 
-    to_return = (orthogonalized_probes,)
     if return_reexpressed_weights:
-        to_return += (reexpressed_weight_matrix,)
-    
-    return to_return
+        return orthogonalized_probes, reexpressed_weight_matrix
+    else:
+        return orthogonalized_probes
 
         
 
@@ -795,12 +794,12 @@ def calc_mode_power_fractions(
     """
 
     if not assume_preorthogonalized:
-        ortho_probes, reexpressed_weights = \
-            orthogonalize_probes_t(
-                probes,
-                weight_matrix=weight_matrix,
-                n_probe_dims=n_probe_dims,
-            )
+        ortho_probes = orthogonalize_probes_t(
+            probes,
+            weight_matrix=weight_matrix,
+            n_probe_dims=n_probe_dims,
+            return_reexpressed_weights=False
+        )
     else:
         weight_slice = np.s_[...,] + np.s_[None,] * n_probe_dims
         if weight_matrix is None:

src/cdtools/tools/plotting/plotting.py

@@ -758,7 +758,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
 
         ax_im = axes[1].images[-1]
         ax_im.set_data(im)
+        ax_im.norecurse=False
         update_colorbar(ax_im)
+        #plt.draw()
 
 
     #

tests/tools/test_analysis.py

@@ -80,7 +80,8 @@ def test_orthogonalize_probes_t():
     ]
     
     for weight_matrix in weight_matrices:
-        ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix)
+        ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix,
+                                     return_reexpressed_weights=True)
         assert isinstance(ortho_probes_np, np.ndarray)
         assert isinstance(rwm_np, np.ndarray)
         
@@ -88,7 +89,8 @@ def test_orthogonalize_probes_t():
         wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
                 else weight_matrix)
 
-        ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t)
+        ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t,
+                                   return_reexpressed_weights=True)
         assert t.is_tensor(ortho_probes_t)
         assert t.is_tensor(rwm_t)