diff --git a/pycs/astro/wl/mass_mapping.py b/pycs/astro/wl/mass_mapping.py
index a879776..8fa05bf 100755
--- a/pycs/astro/wl/mass_mapping.py
+++ b/pycs/astro/wl/mass_mapping.py
@@ -834,8 +834,7 @@ def iks(self, g1, g2, mask, niter=None, dctmax=None):
 
         xg = np.zeros_like(g1, dtype=complex)  # TODO: complex or complex128?
         if dctmax is None:
-            ks = self.gamma_to_cf_kappa(g1 * mask, g2 * mask)
-            lmax = self.get_lmax_dct_inpaint(ks.real, ks.imag)
+            lmax = self.get_lmax_dct_inpaint(g1 * mask, g2 * mask)
         else:
             lmax = dctmax
 
@@ -914,9 +913,9 @@ def _prepare_data(self, InshearData, msg=None, niter=None, Nsigma=None):
             mask = (InshearData.Ncov != 0).astype(int)  # shape = (nx, ny)
         else:
             mask = InshearData.mask
-        InshearData.Ncov[
-            InshearData.Ncov == 0
-        ] = 1e9  # infinite value for no measurement
+        InshearData.Ncov[InshearData.Ncov == 0] = (
+            1e9  # infinite value for no measurement
+        )
         Ncv = InshearData.Ncov / 2.0  # shape = (nx, ny)
 
         # find the minimum noise variance
@@ -930,9 +929,9 @@ def _prepare_data(self, InshearData, msg=None, niter=None, Nsigma=None):
         eta = tau
         # compute signal coefficient
         Esn = eta / Ncv  # shape = (nx, ny)
-        Esn[
-            Esn == np.inf
-        ] = 0  # TODO: useless if we have set Ncv[mask == 0] = 1e9 before
+        Esn[Esn == np.inf] = (
+            0  # TODO: useless if we have set Ncv[mask == 0] = 1e9 before
+        )
 
         return gamma1, gamma2, nx, ny, eta, Esn, mask, ind, tau, niter, Nsigma
 
@@ -1012,9 +1011,10 @@ def prox_wiener_filtering(
 
         if PropagateNoise:
             # Linear operator: uncertainty intervals do not depend on the input images
+            inpshape = gamma1.shape[:-2]  # typically, inpshape = (nimgs,)
             gamma1, gamma2 = self._noise_realizations(
-                InshearData, mask, Nrea=Nrea
-            )  # shape = (Nrea, nx, ny)
+                InshearData, mask, Nrea=Nrea, inpshape=inpshape
+            )  # shape = ([Nimgs], [Nrea], nx, ny)
 
         xg = np.zeros_like(gamma1)
 
@@ -1086,16 +1086,16 @@ def prox_mse(
 
         if PropagateNoise:
             # Linear operator: uncertainty intervals do not depend on the input images
+            inpshape = gamma1.shape[:-2]  # typically, inpshape = (nimgs,)
             gamma1, gamma2 = self._noise_realizations(
-                InshearData, mask, Nrea=Nrea
-            )  # shape = ([Nrea], nx, ny)
+                InshearData, mask, Nrea=Nrea, inpshape=inpshape
+            )  # shape = ([Nimgs], [Nrea], nx, ny)
 
         xg = self.gamma_to_cf_kappa(gamma1, gamma2)  # shape = ([nimgs], nx, ny)
         if Inpaint:
-            # TODO: check code: self.get_resi should be computed on shear maps, not convergence maps.
             # TODO: to be placed before or after "if PropagateNoise"? Inconsistent between methods.
             lmin = 0
-            lmax = self.get_lmax_dct_inpaint(xg.real, xg.imag)
+            lmax = self.get_lmax_dct_inpaint(gamma1, gamma2)
 
         for n in range(niter):
             t1, t2 = self.get_resi(xg, gamma1, gamma2, Esn)  # shape = ([nimgs], nx, ny)
@@ -1200,11 +1200,6 @@ def sparse_wiener_filtering(
         )
 
         RMS_ShearMap = np.sqrt(InshearData.Ncov / 2.0)  # shape = (nx, ny)
-        # shape = ([nimgs], [Nrea], nx, ny)
-        # TODO: complex or complex128? Same question for real-valued arrays
-        xg = np.zeros_like(gamma1, dtype=complex)  # Gaussian + sparse components
-        xs = np.zeros_like(gamma1, dtype=complex)  # sparse component
-        xw = np.zeros_like(gamma1, dtype=complex)  # Gaussian component
         SigmaNoise = np.min(RMS_ShearMap)  # float
         Esn_Sparse = SigmaNoise / RMS_ShearMap  # shape = (nx, ny)
         Esn_Sparse[Esn_Sparse == np.inf] = 0
@@ -1219,9 +1214,8 @@ def sparse_wiener_filtering(
         # the border.
 
         if Inpaint:
-            # TODO: check code: self.get_resi should be computed on shear maps, not convergence maps.
             # TODO: to be placed before or after "if PropagateNoise is True"? Inconsistent between methods.
-            lmin = 0
+            xg = np.zeros_like(gamma1, dtype=complex)
             resi1, resi2 = self.get_resi(xg, gamma1, gamma2, Esn)
             lmax = self.get_lmax_dct_inpaint(resi1, resi2)
 
@@ -1255,7 +1249,11 @@ def sparse_wiener_filtering(
                 InshearData, mask, Nrea=Nrea, inpshape=inpshape
             )  # shape = ([nimgs], [Nrea], nx, ny)
 
-        
+        # shape = ([nimgs], [Nrea], nx, ny)
+        # TODO: complex or complex128? Same question for real-valued arrays
+        xg = np.zeros_like(gamma1, dtype=complex)  # Gaussian + sparse components
+        xs = np.zeros_like(gamma1, dtype=complex)  # sparse component
+        xw = np.zeros_like(gamma1, dtype=complex)  # Gaussian component
 
         for n in range(niter):
             resi1, resi2 = self.get_resi(