BUG: Fix bugs from merge

AdvancedPhotonSource · Jul 9, 2024 · a7a40ec · a7a40ec
1 parent 14338f2
commit a7a40ec
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/src/tike/ptycho/solvers/_preconditioner.py b/src/tike/ptycho/solvers/_preconditioner.py
@@ -46,7 +46,7 @@ def _psi_preconditioner(
 ) -> npt.NDArray:
 
     psi_update_denominator = cp.zeros(
-        shape=parameters.psi.shape[-2:],
+        shape=parameters.psi.shape,
         dtype=parameters.psi.dtype,
     )
 
@@ -58,25 +58,25 @@ def make_certain_args_constant(
         nonlocal psi_update_denominator
 
         probe_amp = _probe_amp_sum(parameters.probe)[:, 0]
-        psi_update_denominator[...] = operator.diffraction.patch.adj(
+        psi_update_denominator[0] = operator.diffraction.patch.adj(
             patches=probe_amp,
-            images=psi_update_denominator,
+            images=psi_update_denominator[0],
             positions=parameters.scan[lo:hi],
         )
 
-        probe1 = probe[:, 0]
-        for i in range(1, len(psi)):
+        probe1 = parameters.probe[:, 0]
+        for i in range(1, len(parameters.psi)):
             probe1 = operator.diffraction.diffraction.fwd(
                 probe=probe1,
-                scan=scan[lo:hi],
-                psi=psi[i-1],
+                scan=parameters.scan[lo:hi],
+                psi=parameters.psi[i-1],
             )
             probe1 = operator.diffraction.propagation.fwd(probe1)
             probe_amp = _probe_amp_sum(probe1)
             psi_update_denominator[i] = operator.diffraction.patch.adj(
                 patches=probe_amp,
                 images=psi_update_denominator[i],
-                positions=scan[lo:hi],
+                positions=parameters.scan[lo:hi],
             )
 
     tike.communicators.stream.stream_and_modify2(
@@ -87,7 +87,7 @@ def make_certain_args_constant(
         hi=len(parameters.scan),
     )
 
-    return psi_update_denominator[None, ...]
+    return psi_update_denominator
 
 
 @cp.fuse()

diff --git a/src/tike/ptycho/solvers/lstsq.py b/src/tike/ptycho/solvers/lstsq.py
@@ -287,7 +287,7 @@ def lstsq_grad(
                 v=object_options.v,
                 m=object_options.m,
                 mdecay=object_options.mdecay,
-                errors=list(float(cp.mean(x)) for x in algorithm_options.costs[-3:]),
+                errors=list(float(np.mean(x)) for x in algorithm_options.costs[-3:]),
                 beta=beta_object,
                 memory_length=3,
             )