foscat 3.7.3__py3-none-any.whl → 3.8.2__py3-none-any.whl

foscat/scat_cov.py

@@ -2277,14 +2277,14 @@ class funct(FOC.FoCUS):
                     if tmp.S1 is not None:
                         nS1 = np.expand_dims(tmp.S1, 0)
                 else:
-                    nS0 = np.expand_dims(tmp.S0.numpy(), 0)
-                    nS2 = np.expand_dims(tmp.S2.numpy(), 0)
-                    nS3 = np.expand_dims(tmp.S3.numpy(), 0)
-                    nS4 = np.expand_dims(tmp.S4.numpy(), 0)
+                    nS0 = np.expand_dims(self.backend.to_numpy(tmp.S0), 0)
+                    nS2 = np.expand_dims(self.backend.to_numpy(tmp.S2), 0)
+                    nS3 = np.expand_dims(self.backend.to_numpy(tmp.S3), 0)
+                    nS4 = np.expand_dims(self.backend.to_numpy(tmp.S4), 0)
                     if tmp.S3P is not None:
-                        nS3P = np.expand_dims(tmp.S3P.numpy(), 0)
+                        nS3P = np.expand_dims(self.backend.to_numpy(tmp.S3P), 0)
                     if tmp.S1 is not None:
-                        nS1 = np.expand_dims(tmp.S1.numpy(), 0)
+                        nS1 = np.expand_dims(self.backend.to_numpy(tmp.S1), 0)
 
                 if S0 is None:
                     S0 = nS0
@@ -2304,24 +2304,24 @@ class funct(FOC.FoCUS):
                         S3P = np.concatenate([S3P, nS3P], 0)
                     if tmp.S1 is not None:
                         S1 = np.concatenate([S1, nS1], 0)
-            sS0 = np.std(S0, 0)
-            sS2 = np.std(S2, 0)
-            sS3 = np.std(S3, 0)
-            sS4 = np.std(S4, 0)
-            mS0 = np.mean(S0, 0)
-            mS2 = np.mean(S2, 0)
-            mS3 = np.mean(S3, 0)
-            mS4 = np.mean(S4, 0)
+            sS0 = self.backend.bk_cast(np.std(S0, 0))
+            sS2 = self.backend.bk_cast(np.std(S2, 0))
+            sS3 = self.backend.bk_cast(np.std(S3, 0))
+            sS4 = self.backend.bk_cast(np.std(S4, 0))
+            mS0 = self.backend.bk_cast(np.mean(S0, 0))
+            mS2 = self.backend.bk_cast(np.mean(S2, 0))
+            mS3 = self.backend.bk_cast(np.mean(S3, 0))
+            mS4 = self.backend.bk_cast(np.mean(S4, 0))
             if tmp.S3P is not None:
-                sS3P = np.std(S3P, 0)
-                mS3P = np.mean(S3P, 0)
+                sS3P = self.backend.bk_cast(np.std(S3P, 0))
+                mS3P = self.backend.bk_cast(np.mean(S3P, 0))
             else:
                 sS3P = None
                 mS3P = None
 
             if tmp.S1 is not None:
-                sS1 = np.std(S1, 0)
-                mS1 = np.mean(S1, 0)
+                sS1 = self.backend.bk_cast(np.std(S1, 0))
+                mS1 = self.backend.bk_cast(np.mean(S1, 0))
             else:
                 sS1 = None
                 mS1 = None
@@ -2375,14 +2375,13 @@ class funct(FOC.FoCUS):
 
             # instead of difference between "opposite" channels use weighted average
             # of cosine and sine contributions using all channels
-            angles = (
+            angles = self.backend.bk_cast((
                 2 * np.pi * np.arange(self.NORIENT) / self.NORIENT
-            )  # shape: (NORIENT,)
-            angles = angles.reshape(1, 1, self.NORIENT)
+            ).reshape(1, 1, self.NORIENT))  # shape: (NORIENT,)
 
             # we use cosines and sines as weights for sim
-            weighted_cos = self.backend.bk_reduce_mean(sim * np.cos(angles), axis=-1)
-            weighted_sin = self.backend.bk_reduce_mean(sim * np.sin(angles), axis=-1)
+            weighted_cos = self.backend.bk_reduce_mean(sim * self.backend.bk_cos(angles), axis=-1)
+            weighted_sin = self.backend.bk_reduce_mean(sim * self.backend.bk_sin(angles), axis=-1)
             # For simplicity, take first element of the batch
             cc = weighted_cos[0]
             ss = weighted_sin[0]
@@ -2402,7 +2401,8 @@ class funct(FOC.FoCUS):
                 phase = np.fmod(np.arctan2(ss, cc) + 2 * np.pi, 2 * np.pi)
             else:
                 phase = np.fmod(
-                    np.arctan2(ss.numpy(), cc.numpy()) + 2 * np.pi, 2 * np.pi
+                    np.arctan2(self.backend.to_numpy(ss), 
+                                self.backend.to_numpy(cc)) + 2 * np.pi, 2 * np.pi
                 )
 
             # instead of linear interpolation cosine‐based interpolation
@@ -2435,7 +2435,8 @@ class funct(FOC.FoCUS):
 
                 sim2 = self.backend.bk_reduce_sum(
                     self.backend.bk_reshape(
-                        mat.reshape(1, mat.shape[0], self.NORIENT * self.NORIENT)
+                        self.backend.bk_cast(
+                        mat.reshape(1, mat.shape[0], self.NORIENT * self.NORIENT))
                         * tmp2,
                         [sim.shape[0], cmat[k].shape[0], self.NORIENT, self.NORIENT],
                     ),
@@ -2445,10 +2446,10 @@ class funct(FOC.FoCUS):
                 sim2 = self.backend.bk_abs(self.convol(sim2, axis=1))
 
                 weighted_cos2 = self.backend.bk_reduce_mean(
-                    sim2 * np.cos(angles), axis=-1
+                    sim2 * self.backend.bk_cos(angles), axis=-1
                 )
                 weighted_sin2 = self.backend.bk_reduce_mean(
-                    sim2 * np.sin(angles), axis=-1
+                    sim2 * self.backend.bk_sin(angles), axis=-1
                 )
 
                 cc2 = weighted_cos2[0]
@@ -2469,7 +2470,8 @@ class funct(FOC.FoCUS):
                     phase2 = np.fmod(np.arctan2(ss2, cc2) + 2 * np.pi, 2 * np.pi)
                 else:
                     phase2 = np.fmod(
-                        np.arctan2(ss2.numpy(), cc2.numpy()) + 2 * np.pi, 2 * np.pi
+                        np.arctan2(self.backend.to_numpy(ss2), 
+                                self.backend.to_numpy(cc2)) + 2 * np.pi, 2 * np.pi
                     )
 
                 phase2_scaled = self.NORIENT * phase2 / (2 * np.pi)
@@ -4791,17 +4793,37 @@ class funct(FOC.FoCUS):
         This function is strongly inspire by the package https://github.com/SihaoCheng/scattering_transform
         Done by Sihao Cheng and Rudy Morel.
         '''
-        if_xodd = (data_f.shape[-2]%2==1)
-        if_yodd = (data_f.shape[-1]%2==1)
-        result = self.backend.backend.cat(
-            (self.backend.backend.cat(
-                ( data_f[...,:dx+if_xodd, :dy+if_yodd] , data_f[...,-dx:, :dy+if_yodd]
-                ), -2),
-              self.backend.backend.cat(
-                ( data_f[...,:dx+if_xodd, -dy:] , data_f[...,-dx:, -dy:]
-                ), -2)
-            ),-1)
-        return result
+        
+        if self.backend.BACKEND=='torch':
+            if_xodd = (data_f.shape[-2]%2==1)
+            if_yodd = (data_f.shape[-1]%2==1)
+            result = self.backend.backend.cat(
+                (self.backend.backend.cat(
+                    ( data_f[...,:dx+if_xodd, :dy+if_yodd] , data_f[...,-dx:, :dy+if_yodd]
+                    ), -2),
+                self.backend.backend.cat(
+                    ( data_f[...,:dx+if_xodd, -dy:] , data_f[...,-dx:, -dy:]
+                    ), -2)
+                ),-1)
+            return result
+        else:
+            # Check if the last two dimensions are odd
+            if_xodd = self.backend.backend.cast(self.backend.backend.shape(data_f)[-2] % 2 == 1, self.backend.backend.int32)
+            if_yodd = self.backend.backend.cast(self.backend.backend.shape(data_f)[-1] % 2 == 1, self.backend.backend.int32)
+
+            # Extract four regions
+            top_left = data_f[..., :dx+if_xodd, :dy+if_yodd]
+            top_right = data_f[..., -dx:, :dy+if_yodd]
+            bottom_left = data_f[..., :dx+if_xodd, -dy:]
+            bottom_right = data_f[..., -dx:, -dy:]
+
+            # Concatenate along the last two dimensions
+            top = self.backend.backend.concat([top_left, top_right], axis=-2)
+            bottom = self.backend.backend.concat([bottom_left, bottom_right], axis=-2)
+            result = self.backend.backend.concat([top, bottom], axis=-1)
+
+            return result
+
     # ---------------------------------------------------------------------------
     #
     # utility functions for computing scattering coef and covariance
@@ -4824,13 +4846,15 @@ class funct(FOC.FoCUS):
         This function is strongly inspire by the package https://github.com/SihaoCheng/scattering_transform
         Done by Sihao Cheng and Rudy Morel.
         '''
-        edge_masks = self.backend.backend.empty((J, M, N))
-        X, Y = self.backend.backend.meshgrid(self.backend.backend.arange(M), self.backend.backend.arange(N), indexing='ij')
+        edge_masks = np.empty((J, M, N))
+        X, Y = np.meshgrid(np.arange(M), np.arange(N), indexing='ij')
         for j in range(J):
             edge_dx = min(M//4, 2**j*d0)
             edge_dy = min(N//4, 2**j*d0)
             edge_masks[j] = (X>=edge_dx) * (X<=M-edge_dx) * (Y>=edge_dy) * (Y<=N-edge_dy)
-        return edge_masks.to(self.backend.torch_device)
+        edge_masks = edge_masks[:,None,:,:]
+        edge_masks = edge_masks / edge_masks.mean((-2,-1))[:,:,None,None]
+        return self.backend.bk_cast(edge_masks)
     
     # ---------------------------------------------------------------------------
     #
@@ -4914,6 +4938,11 @@ class funct(FOC.FoCUS):
         if S4_criteria is None:
             S4_criteria = 'j2>=j1'
         
+        if self.all_bk_type == "float32":
+            C_ONE=np.complex64(1.0)
+        else:
+            C_ONE=np.complex128(1.0)
+            
         # determine jmax and nside corresponding to the input map
         im_shape = data.shape
         if self.use_2D:
@@ -4921,18 +4950,24 @@ class funct(FOC.FoCUS):
                 nside = np.min([im_shape[0], im_shape[1]])
                 M,N = im_shape[0],im_shape[1]
                 N_image =  1
+                N_image2 = 1
             else:
                 nside = np.min([im_shape[1], im_shape[2]])
                 M,N = im_shape[1],im_shape[2]
                 N_image = data.shape[0]
+                if data2 is not None:
+                    N_image2 = data2.shape[0]
             J = int(np.log(nside) / np.log(2))-1  # Number of j scales
         elif self.use_1D:
             if len(data.shape) == 2:
                 npix = int(im_shape[1])  # Number of pixels
                 N_image =  1
+                N_image2 = 1
             else:
                 npix = int(im_shape[0])  # Number of pixels
                 N_image = data.shape[0]
+                if data2 is not None:
+                    N_image2 = data2.shape[0]
 
             nside = int(npix)
 
@@ -4941,9 +4976,12 @@ class funct(FOC.FoCUS):
             if len(data.shape) == 2:
                 npix = int(im_shape[1])  # Number of pixels
                 N_image =  1
+                N_image2 =  1
             else:
                 npix = int(im_shape[0])  # Number of pixels
                 N_image = data.shape[0]
+                if data2 is not None:
+                    N_image2 = data2.shape[0]
 
             nside = int(np.sqrt(npix // 12))
 
@@ -4957,7 +4995,475 @@ class funct(FOC.FoCUS):
             print('\n\n==========')
         
         L=self.NORIENT
+        norm_factor_S3=1.0
         
+        if self.backend.BACKEND=='torch':
+            if (M,N,J,L) not in self.filters_set:
+                self.filters_set[(M,N,J,L)] = self.computer_filter(M,N,J,L) #self.computer_filter(M,N,J,L)
+                
+            filters_set = self.filters_set[(M,N,J,L)]
+            
+            #weight = self.weight
+            if use_ref:
+                if normalization=='S2': 
+                    ref_S2 = self.ref_scattering_cov_S2
+                else: 
+                    ref_P11 = self.ref_scattering_cov['P11']
+
+            # convert numpy array input into self.backend.bk_ tensors
+            data = self.backend.bk_cast(data)
+            data_f = self.backend.bk_fftn(data, dim=(-2,-1))
+            if data2 is not None:
+                data2 = self.backend.bk_cast(data2)
+                data2_f = self.backend.bk_fftn(data2, dim=(-2,-1))
+            
+            # initialize tensors for scattering coefficients
+            S2 = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
+            S1 = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
+            
+            Ndata_S3 = J*(J+1)//2
+            Ndata_S4 = J*(J+1)*(J+2)//6
+            J_S4={}
+            
+            S3 = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+            if data2 is not None:
+                S3p = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+            S4_pre_norm = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
+            S4 = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
+            
+            # variance
+            if get_variance:
+                S2_sigma = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
+                S1_sigma = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
+                S3_sigma = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+                if data2 is not None:
+                    S3p_sigma = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+                S4_sigma = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
+                
+            if iso_ang:
+                S3_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+                S4_iso = self.backend.bk_zeros((N_image,Ndata_S4,L,L), dtype=data_f.dtype) 
+                if get_variance:
+                    S3_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+                    S4_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S4,L,L), dtype=data_f.dtype) 
+                if data2 is not None:
+                    S3p_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+                    if get_variance:
+                        S3p_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+            
+            #
+            if edge: 
+                if (M,N,J) not in self.edge_masks:
+                    self.edge_masks[(M,N,J)] = self.get_edge_masks(M,N,J)
+                edge_mask=self.edge_masks[(M,N,J)]
+            else: 
+                edge_mask = 1
+                
+            # calculate scattering fields
+            if data2 is None:
+                if self.use_2D:
+                    if len(data.shape) == 2:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[None,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        ).abs()
+                    else:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[:,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        ).abs()
+                elif self.use_1D:
+                    if len(data.shape) == 1:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[None,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        ).abs()
+                    else:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[:,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        ).abs()
+                else:
+                    print('todo')   
+                
+                S2 = (I1**2 * edge_mask).mean((-2,-1))
+                S1  = (I1 * edge_mask).mean((-2,-1))
+
+                if get_variance:
+                    S2_sigma = (I1**2 * edge_mask).std((-2,-1))
+                    S1_sigma  = (I1 * edge_mask).std((-2,-1))
+                    
+            else:
+                if self.use_2D:
+                    if len(data.shape) == 2:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[None,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        )
+                        I2 = self.backend.bk_ifftn(
+                            data2_f[None,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        )
+                    else:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[:,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        )
+                        I2 = self.backend.bk_ifftn(
+                            data2_f[:,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
+                        )
+                elif self.use_1D:
+                    if len(data.shape) == 1:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[None,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        )
+                        I2 = self.backend.bk_ifftn(
+                            data2_f[None,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        )
+                    else:
+                        I1 = self.backend.bk_ifftn(
+                            data_f[:,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        )
+                        I2 = self.backend.bk_ifftn(
+                            data2_f[:,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
+                        )
+                else:
+                    print('todo')   
+        
+                I1=self.backend.bk_real(I1*self.backend.bk_conjugate(I2))
+            
+                S2 = self.backend.bk_reduce_mean((I1 * edge_mask),axis=(-2,-1))
+                if get_variance:
+                    S2_sigma = self.backend.bk_reduce_std((I1 * edge_mask),axis=(-2,-1))
+            
+                I1=self.backend.bk_L1(I1)
+            
+                S1  = self.backend.bk_reduce_mean((I1 * edge_mask),axis=(-2,-1))
+
+                if get_variance:
+                    S1_sigma  = self.backend.bk_reduce_std((I1 * edge_mask),axis=(-2,-1))
+                
+            I1_f= self.backend.bk_fftn(I1, dim=(-2,-1))
+                
+            if pseudo_coef != 1:
+                I1 = I1**pseudo_coef
+            
+            Ndata_S3=0
+            Ndata_S4=0
+            
+            # calculate the covariance and correlations of the scattering fields
+            # only use the low-k Fourier coefs when calculating large-j scattering coefs.
+            for j3 in range(0,J):
+                J_S4[j3]=Ndata_S4
+                
+                dx3, dy3 = self.get_dxdy(j3,M,N)
+                I1_f_small = self.cut_high_k_off(I1_f[:,:j3+1], dx3, dy3) # Nimage, J, L, x, y
+                data_f_small = self.cut_high_k_off(data_f, dx3, dy3)
+                if data2 is not None:
+                    data2_f_small = self.cut_high_k_off(data2_f, dx3, dy3)
+                if edge:
+                    I1_small = self.backend.bk_ifftn(I1_f_small, dim=(-2,-1), norm='ortho')
+                    data_small = self.backend.bk_ifftn(data_f_small, dim=(-2,-1), norm='ortho')
+                    if data2 is not None:
+                        data2_small = self.backend.bk_ifftn(data2_f_small, dim=(-2,-1), norm='ortho')
+                        
+                wavelet_f3 = self.cut_high_k_off(filters_set[j3], dx3, dy3) # L,x,y
+                _, M3, N3 = wavelet_f3.shape
+                wavelet_f3_squared = wavelet_f3**2
+                edge_dx = min(4, int(2**j3*dx3*2/M))
+                edge_dy = min(4, int(2**j3*dy3*2/N))
+                
+                # a normalization change due to the cutoff of frequency space
+                fft_factor = 1 /(M3*N3) * (M3*N3/M/N)**2
+                for j2 in range(0,j3+1):
+                    I1_f2_wf3_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3.view(1,1,L,M3,N3)
+                    I1_f2_wf3_2_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3_squared.view(1,1,L,M3,N3)
+                    if edge:
+                        I12_w3_small = self.backend.bk_ifftn(I1_f2_wf3_small, dim=(-2,-1), norm='ortho')
+                        I12_w3_2_small = self.backend.bk_ifftn(I1_f2_wf3_2_small, dim=(-2,-1), norm='ortho')
+                    if use_ref:
+                        if normalization=='P11':
+                            norm_factor_S3 = (ref_S2[:,None,j3,:] * ref_P11[:,j2,j3,:,:]**pseudo_coef)**0.5
+                        if normalization=='S2':
+                            norm_factor_S3 = (ref_S2[:,None,j3,:] * ref_S2[:,j2,:,None]**pseudo_coef)**0.5
+                    else:
+                        if normalization=='P11':
+                            # [N_image,l2,l3,x,y]
+                            P11_temp = (I1_f2_wf3_small.abs()**2).mean((-2,-1)) * fft_factor
+                            norm_factor_S3 = (S2[:,None,j3,:] * P11_temp**pseudo_coef)**0.5
+                        if normalization=='S2':
+                            norm_factor_S3 = (S2[:,None,j3,:] * S2[:,j2,:,None]**pseudo_coef)**0.5
+
+                    if not edge:
+                        S3[:,Ndata_S3,:,:] = (
+                            data_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                        ).mean((-2,-1)) * fft_factor / norm_factor_S3
+                        
+                        if get_variance:
+                            S3_sigma[:,Ndata_S3,:,:] = (
+                                data_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                            ).std((-2,-1)) * fft_factor / norm_factor_S3
+                    else:
+                        S3[:,Ndata_S3,:,:] = (
+                            data_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
+                        )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].mean((-2,-1)) * fft_factor / norm_factor_S3
+                        if get_variance:
+                            S3_sigma[:,Ndata_S3,:,:] = (
+                                data_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
+                                )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].std((-2,-1)) * fft_factor / norm_factor_S3
+                    if data2 is not None:
+                        if not edge:
+                            S3p[:,Ndata_S3,:,:] = (
+                                data2_f_small.view(N_image2,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                            ).mean((-2,-1)) * fft_factor / norm_factor_S3
+                        
+                            if get_variance:
+                                S3p_sigma[:,Ndata_S3,:,:] = (
+                                    data2_f_small.view(N_image2,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                                ).std((-2,-1)) * fft_factor / norm_factor_S3
+                        else:
+                            S3p[:,Ndata_S3,:,:] = (
+                                data2_small.view(N_image2,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
+                            )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].mean((-2,-1)) * fft_factor / norm_factor_S3
+                            if get_variance:
+                                S3p_sigma[:,Ndata_S3,:,:] = (
+                                    data2_small.view(N_image2,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
+                                    )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].std((-2,-1)) * fft_factor / norm_factor_S3
+                    Ndata_S3+=1
+                    if j2 <= j3:
+                        beg_n=Ndata_S4
+                        for j1 in range(0, j2+1):
+                            if eval(S4_criteria):
+                                if not edge:
+                                    if not if_large_batch:
+                                        # [N_image,l1,l2,l3,x,y]
+                                        S4_pre_norm[:,Ndata_S4,:,:,:] = (
+                                            I1_f_small[:,j1].view(N_image,L,1,1,M3,N3) * 
+                                            self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,1,L,L,M3,N3))
+                                        ).mean((-2,-1)) * fft_factor
+                                        if get_variance:
+                                            S4_sigma[:,Ndata_S4,:,:,:] = (
+                                                I1_f_small[:,j1].view(N_image,L,1,1,M3,N3) * 
+                                                self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,1,L,L,M3,N3))
+                                            ).std((-2,-1)) * fft_factor
+                                    else:
+                                        for l1 in range(L):
+                                            # [N_image,l2,l3,x,y]
+                                            S4_pre_norm[:,Ndata_S4,l1,:,:] = (
+                                                I1_f_small[:,j1,l1].view(N_image,1,1,M3,N3) * 
+                                                self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,L,L,M3,N3))
+                                            ).mean((-2,-1)) * fft_factor
+                                            if get_variance:
+                                                S4_sigma[:,Ndata_S4,l1,:,:] = (
+                                                    I1_f_small[:,j1,l1].view(N_image,1,1,M3,N3) * 
+                                                    self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,L,L,M3,N3))
+                                                ).std((-2,-1)) * fft_factor
+                                else:
+                                    if not if_large_batch:
+                                        # [N_image,l1,l2,l3,x,y]
+                                        S4_pre_norm[:,Ndata_S4,:,:,:] = (
+                                            I1_small[:,j1].view(N_image,L,1,1,M3,N3) * self.backend.bk_conjugate(
+                                                I12_w3_2_small.view(N_image,1,L,L,M3,N3)
+                                            )
+                                        )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                        if get_variance:
+                                            S4_sigma[:,Ndata_S4,:,:,:] = (
+                                                I1_small[:,j1].view(N_image,L,1,1,M3,N3) * self.backend.bk_conjugate(
+                                                    I12_w3_2_small.view(N_image,1,L,L,M3,N3)
+                                                )
+                                            )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].std((-2,-1)) * fft_factor
+                                    else:
+                                        for l1 in range(L):
+                                        # [N_image,l2,l3,x,y]
+                                            S4_pre_norm[:,Ndata_S4,l1,:,:] = (
+                                                I1_small[:,j1].view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(
+                                                    I12_w3_2_small.view(N_image,L,L,M3,N3)
+                                                )
+                                            )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                            if get_variance:
+                                                S4_sigma[:,Ndata_S4,l1,:,:] = (
+                                                    I1_small[:,j1].view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(
+                                                        I12_w3_2_small.view(N_image,L,L,M3,N3)
+                                                    )
+                                                )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                    
+                                Ndata_S4+=1
+                                
+                        if normalization=='S2':
+                            if use_ref: 
+                                P = ((ref_S2[:,j3:j3+1,:,None,None] * ref_S2[:,j2:j2+1,None,:,None] )**(0.5*pseudo_coef))
+                            else: 
+                                P = ((S2[:,j3:j3+1,:,None,None] * S2[:,j2:j2+1,None,:,None] )**(0.5*pseudo_coef))
+                                                        
+                            S4[:,beg_n:Ndata_S4,:,:,:]=S4_pre_norm[:,beg_n:Ndata_S4,:,:,:].clone()/P
+                                
+                            if get_variance:
+                                S4_sigma[:,beg_n:Ndata_S4,:,:,:] = S4_sigma[:,beg_n:Ndata_S4,:,:,:]/P
+                        else:
+                            S4=S4_pre_norm
+                            
+            # average over l1 to obtain simple isotropic statistics
+            if iso_ang:
+                S2_iso = S2.mean(-1)
+                S1_iso = S1.mean(-1)
+                for l1 in range(L):
+                    for l2 in range(L):
+                        S3_iso[...,(l2-l1)%L] += S3[...,l1,l2]
+                        if data2 is not None:
+                            S3p_iso[...,(l2-l1)%L] += S3p[...,l1,l2]
+                        for l3 in range(L):
+                            S4_iso[...,(l2-l1)%L,(l3-l1)%L] += S4[...,l1,l2,l3]
+                S3_iso /= L; S4_iso /= L
+                if data2 is not None:
+                    S3p_iso /= L
+                
+                if get_variance:
+                    S2_sigma_iso = S2_sigma.mean(-1)
+                    S1_sigma_iso = S1_sigma.mean(-1)
+                    for l1 in range(L):
+                        for l2 in range(L):
+                            S3_sigma_iso[...,(l2-l1)%L] += S3_sigma[...,l1,l2]
+                            if data2 is not None:
+                                S3p_sigma_iso[...,(l2-l1)%L] += S3p_sigma[...,l1,l2]
+                            for l3 in range(L):
+                                S4_sigma_iso[...,(l2-l1)%L,(l3-l1)%L] += S4_sigma[...,l1,l2,l3]
+                    S3_sigma_iso /= L; S4_sigma_iso /= L
+                    if data2 is not None:
+                        S3p_sigma_iso /= L
+            
+            mean_data=self.backend.bk_zeros((N_image,1), dtype=data.dtype) 
+            std_data=self.backend.bk_zeros((N_image,1), dtype=data.dtype)
+            
+            if data2 is None:
+                mean_data[:,0]=data.mean((-2,-1))
+                std_data[:,0]=data.std((-2,-1))
+            else:
+                mean_data[:,0]=(data2*data).mean((-2,-1))
+                std_data[:,0]=(data2*data).std((-2,-1))
+            
+            if get_variance:
+                ref_sigma={}
+                if iso_ang:
+                    ref_sigma['std_data']=std_data
+                    ref_sigma['S1_sigma']=S1_sigma_iso
+                    ref_sigma['S2_sigma']=S2_sigma_iso
+                    ref_sigma['S3_sigma']=S3_sigma_iso
+                    if data2 is not None:
+                        ref_sigma['S3p_sigma']=S3p_sigma_iso
+                    ref_sigma['S4_sigma']=S4_sigma_iso
+                else:
+                    ref_sigma['std_data']=std_data
+                    ref_sigma['S1_sigma']=S1_sigma
+                    ref_sigma['S2_sigma']=S2_sigma
+                    ref_sigma['S3_sigma']=S3_sigma
+                    if data2 is not None:
+                        ref_sigma['S3p_sigma']=S3p_sigma
+                    ref_sigma['S4_sigma']=S4_sigma
+            
+            if data2 is None:
+                if iso_ang:
+                    if ref_sigma is not None:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/ref_sigma['std_data'],
+                            std_data/ref_sigma['std_data'],
+                            (S2_iso/ref_sigma['S2_sigma']).reshape((N_image, -1)).log(), 
+                            (S1_iso/ref_sigma['S1_sigma']).reshape((N_image, -1)).log(),
+                            (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
+                            (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
+                            (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
+                            (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                    else:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/std_data,
+                            std_data,
+                            S2_iso.reshape((N_image, -1)).log(), 
+                            S1_iso.reshape((N_image, -1)).log(),
+                            S3_iso.reshape((N_image, -1)).real, 
+                            S3_iso.reshape((N_image, -1)).imag,
+                            S4_iso.reshape((N_image, -1)).real, 
+                            S4_iso.reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                else:
+                    if ref_sigma is not None:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/ref_sigma['std_data'],
+                            std_data/ref_sigma['std_data'],
+                            (S2/ref_sigma['S2_sigma']).reshape((N_image, -1)).log(), 
+                            (S1/ref_sigma['S1_sigma']).reshape((N_image, -1)).log(),
+                            (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
+                            (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
+                            (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
+                            (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                    else:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/std_data,
+                            std_data,
+                            S2.reshape((N_image, -1)).log(), 
+                            S1.reshape((N_image, -1)).log(),
+                            S3.reshape((N_image, -1)).real, 
+                            S3.reshape((N_image, -1)).imag,
+                            S4.reshape((N_image, -1)).real, 
+                            S4.reshape((N_image, -1)).imag,
+                            ),dim=-1)
+            else:
+                if iso_ang:
+                    if ref_sigma is not None:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/ref_sigma['std_data'],
+                            std_data/ref_sigma['std_data'],
+                            (S2_iso/ref_sigma['S2_sigma']).reshape((N_image, -1)), 
+                            (S1_iso/ref_sigma['S1_sigma']).reshape((N_image, -1)),
+                            (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
+                            (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
+                            (S3p_iso/ref_sigma['S3p_sigma']).reshape((N_image, -1)).real, 
+                            (S3p_iso/ref_sigma['S3p_sigma']).reshape((N_image, -1)).imag,
+                            (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
+                            (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                    else:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/std_data,
+                            std_data,
+                            S2_iso.reshape((N_image, -1)), 
+                            S1_iso.reshape((N_image, -1)),
+                            S3_iso.reshape((N_image, -1)).real, 
+                            S3_iso.reshape((N_image, -1)).imag,
+                            S3p_iso.reshape((N_image, -1)).real, 
+                            S3p_iso.reshape((N_image, -1)).imag,
+                            S4_iso.reshape((N_image, -1)).real, 
+                            S4_iso.reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                else:
+                    if ref_sigma is not None:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/ref_sigma['std_data'],
+                            std_data/ref_sigma['std_data'],
+                            (S2/ref_sigma['S2_sigma']).reshape((N_image, -1)), 
+                            (S1/ref_sigma['S1_sigma']).reshape((N_image, -1)),
+                            (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
+                            (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
+                            (S3p/ref_sigma['S3p_sigma']).reshape((N_image, -1)).real, 
+                            (S3p/ref_sigma['S3p_sigma']).reshape((N_image, -1)).imag,
+                            (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
+                            (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                    else:
+                        for_synthesis = self.backend.backend.cat((
+                            mean_data/std_data,
+                            std_data,
+                            S2.reshape((N_image, -1)), 
+                            S1.reshape((N_image, -1)),
+                            S3.reshape((N_image, -1)).real, 
+                            S3.reshape((N_image, -1)).imag,
+                            S3p.reshape((N_image, -1)).real, 
+                            S3p.reshape((N_image, -1)).imag,
+                            S4.reshape((N_image, -1)).real, 
+                            S4.reshape((N_image, -1)).imag,
+                            ),dim=-1)
+                    
+            if not use_ref: 
+                self.ref_scattering_cov_S2=S2
+            
+            if get_variance:
+                return for_synthesis,ref_sigma
+                
+            return for_synthesis
+            
         if (M,N,J,L) not in self.filters_set:
             self.filters_set[(M,N,J,L)] = self.computer_filter(M,N,J,L) #self.computer_filter(M,N,J,L)
             
@@ -4978,46 +5484,41 @@ class funct(FOC.FoCUS):
             data2_f = self.backend.bk_fftn(data2, dim=(-2,-1))
         
         # initialize tensors for scattering coefficients
-        S2 = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
-        S1 = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
         
         Ndata_S3 = J*(J+1)//2
         Ndata_S4 = J*(J+1)*(J+2)//6
         J_S4={}
         
-        S3 = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+        S3 = []
         if data2 is not None:
-            S3p = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype)
-        S4_pre_norm = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
-        S4 = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
+            S3p = []
+        S4_pre_norm = [] 
+        S4 = [] 
         
         # variance
         if get_variance:
-            S2_sigma = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
-            S1_sigma = self.backend.bk_zeros((N_image,J,L), dtype=data.dtype)
-            S3_sigma = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype) 
+            S3_sigma = [] 
             if data2 is not None:
-                S3p_sigma = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype)
-            S4_sigma = self.backend.bk_zeros((N_image,Ndata_S4,L,L,L), dtype=data_f.dtype) 
+                S3p_sigma = []
+            S4_sigma = [] 
             
         if iso_ang:
-            S3_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+            S3_iso = []
             if data2 is not None:
-                S3p_iso = self.backend.bk_zeros((N_image,Ndata_S3,L,L), dtype=data_f.dtype)
+                S3p_iso = []
                 
-            S4_iso = self.backend.bk_zeros((N_image,Ndata_S4,L,L), dtype=data_f.dtype) 
+            S4_iso = []
             if get_variance:
-                S3_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
+                S3_sigma_iso = []
                 if data2 is not None:
-                    S3p_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S3,L), dtype=data_f.dtype) 
-                S4_sigma_iso = self.backend.bk_zeros((N_image,Ndata_S4,L,L), dtype=data_f.dtype) 
+                    S3p_sigma_iso = [] 
+                S4_sigma_iso = [] 
             
         #
         if edge: 
             if (M,N,J) not in self.edge_masks:
                 self.edge_masks[(M,N,J)] = self.get_edge_masks(M,N,J)
-            edge_mask = self.edge_masks[(M,N,J)][:,None,:,:]
-            edge_mask = edge_mask / edge_mask.mean((-2,-1))[:,:,None,None]
+            edge_mask = self.edge_masks[(M,N,J)]
         else: 
             edge_mask = 1
             
@@ -5025,31 +5526,31 @@ class funct(FOC.FoCUS):
         if data2 is None:
             if self.use_2D:
                 if len(data.shape) == 2:
-                    I1 = self.backend.bk_ifftn(
+                    I1 = self.backend.bk_abs(self.backend.bk_ifftn(
                         data_f[None,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
-                    ).abs()
+                    ))
                 else:
-                    I1 = self.backend.bk_ifftn(
+                    I1 = self.backend.bk_abs(self.backend.bk_ifftn(
                         data_f[:,None,None,:,:] * filters_set[None,:J,:,:,:], dim=(-2,-1)
-                    ).abs()
+                    ))
             elif self.use_1D:
                 if len(data.shape) == 1:
-                    I1 = self.backend.bk_ifftn(
+                    I1 = self.backend.bk_abs(self.backend.bk_ifftn(
                         data_f[None,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
-                    ).abs()
+                    ))
                 else:
-                    I1 = self.backend.bk_ifftn(
+                    I1 = self.backend.bk_abs(self.backend.bk_ifftn(
                         data_f[:,None,None,:] * filters_set[None,:J,:,:], dim=(-1)
-                    ).abs()
+                    ))
             else:
                 print('todo')
             
-            S2 = (I1**2 * edge_mask).mean((-2,-1))
-            S1  = (I1 * edge_mask).mean((-2,-1))
+            S2 = self.backend.bk_reduce_mean((I1**2 * edge_mask),axis=(-2,-1))
+            S1 = self.backend.bk_reduce_mean(I1 * edge_mask,axis=(-2,-1))
 
             if get_variance:
-                S2_sigma = (I1**2 * edge_mask).std((-2,-1))
-                S1_sigma  = (I1 * edge_mask).std((-2,-1))
+                S2_sigma = self.backend.bk_reduce_std((I1**2 * edge_mask),axis=(-2,-1))
+                S1_sigma  = self.backend.bk_reduce_std((I1 * edge_mask),axis=(-2,-1))
                 
             I1_f= self.backend.bk_fftn(I1, dim=(-2,-1))
             
@@ -5089,16 +5590,16 @@ class funct(FOC.FoCUS):
             
             I1=self.backend.bk_real(I1*self.backend.bk_conjugate(I2))
             
-            S2 = (I1 * edge_mask).mean((-2,-1))
+            S2 = self.backend.bk_reduce_mean((I1 * edge_mask),axis=(-2,-1))
             if get_variance:
-                S2_sigma = (I1 * edge_mask).std((-2,-1))
+                S2_sigma = self.backend.bk_reduce_std((I1 * edge_mask),axis=(-2,-1))
             
             I1=self.backend.bk_L1(I1)
             
-            S1  = (I1 * edge_mask).mean((-2,-1))
+            S1  = self.backend.bk_reduce_mean((I1 * edge_mask),axis=(-2,-1))
 
             if get_variance:
-                S1_sigma  = (I1 * edge_mask).std((-2,-1))
+                S1_sigma  = self.backend.bk_reduce_std((I1 * edge_mask),axis=(-2,-1))
                 
             I1_f= self.backend.bk_fftn(I1, dim=(-2,-1))
             
@@ -5129,170 +5630,156 @@ class funct(FOC.FoCUS):
             edge_dx = min(4, int(2**j3*dx3*2/M))
             edge_dy = min(4, int(2**j3*dy3*2/N))
             # a normalization change due to the cutoff of frequency space
-            fft_factor = 1 /(M3*N3) * (M3*N3/M/N)**2
+            if self.all_bk_type == "float32":
+                fft_factor = np.complex64(1 /(M3*N3) * (M3*N3/M/N)**2)
+            else:
+                fft_factor = np.complex128(1 /(M3*N3) * (M3*N3/M/N)**2)
             for j2 in range(0,j3+1):
-                I1_f2_wf3_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3.view(1,1,L,M3,N3)
-                I1_f2_wf3_2_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3_squared.view(1,1,L,M3,N3)
+                #I1_f2_wf3_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3.view(1,1,L,M3,N3)
+                #I1_f2_wf3_2_small = I1_f_small[:,j2].view(N_image,L,1,M3,N3) * wavelet_f3_squared.view(1,1,L,M3,N3)
+                I1_f2_wf3_small = self.backend.bk_reshape(I1_f_small[:,j2],[N_image,1,L,1,M3,N3]) * self.backend.bk_reshape(wavelet_f3,[1,1,1,L,M3,N3])
+                I1_f2_wf3_2_small = self.backend.bk_reshape(I1_f_small[:,j2],[N_image,1,L,1,M3,N3]) * self.backend.bk_reshape(wavelet_f3_squared,[1,1,1,L,M3,N3])
                 if edge:
                     I12_w3_small = self.backend.bk_ifftn(I1_f2_wf3_small, dim=(-2,-1), norm='ortho')
                     I12_w3_2_small = self.backend.bk_ifftn(I1_f2_wf3_2_small, dim=(-2,-1), norm='ortho')
                 if use_ref:
                     if normalization=='P11':
                         norm_factor_S3 = (ref_S2[:,None,j3,:] * ref_P11[:,j2,j3,:,:]**pseudo_coef)**0.5
+                        norm_factor_S3 = self.backend.bk_complex(norm_factor_S3,0*norm_factor_S3)
                     elif normalization=='S2':
                         norm_factor_S3 = (ref_S2[:,None,j3,:] * ref_S2[:,j2,:,None]**pseudo_coef)**0.5
-                    norm_factor_S3 = 1.0
+                        norm_factor_S3 = self.backend.bk_complex(norm_factor_S3,0*norm_factor_S3)
+                    else:
+                        norm_factor_S3 = C_ONE
                 else:
                     if normalization=='P11':
                         # [N_image,l2,l3,x,y]
-                        P11_temp = (I1_f2_wf3_small.abs()**2).mean((-2,-1)) * fft_factor
+                        P11_temp = self.backend.bk_reduce_mean((I1_f2_wf3_small.abs()**2),axis=(-2,-1)) * fft_factor
                         norm_factor_S3 = (S2[:,None,j3,:] * P11_temp**pseudo_coef)**0.5
+                        norm_factor_S3 = self.backend.bk_complex(norm_factor_S3,0*norm_factor_S3)
                     elif normalization=='S2':
-                        norm_factor_S3 = (S2[:,None,j3,:] * S2[:,j2,:,None]**pseudo_coef)**0.5
-                    norm_factor_S3 = 1.0
+                        norm_factor_S3 = (S2[:,None,j3,None,:] * S2[:,None,j2,:,None]**pseudo_coef)**0.5
+                        norm_factor_S3 = self.backend.bk_complex(norm_factor_S3,0*norm_factor_S3)
+                    else:
+                        norm_factor_S3 = C_ONE
+                        
 
                 if not edge:
-                    S3[:,Ndata_S3,:,:] = (
-                        data_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
-                    ).mean((-2,-1)) * fft_factor / norm_factor_S3
-                    
+                    S3.append(self.backend.bk_reduce_mean(
+                        self.backend.bk_reshape(data_f_small,[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                    ,axis=(-2,-1)) * fft_factor / norm_factor_S3)
                     if get_variance:
-                        S3_sigma[:,Ndata_S3,:,:] = (
-                            data_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
-                        ).std((-2,-1)) * fft_factor / norm_factor_S3
+                        S3_sigma.append(self.backend.bk_reduce_std(
+                            self.backend.bk_reshape(data_f_small,[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                        ,axis=(-2,-1)) * fft_factor / norm_factor_S3)
                 else:
-                    
-                    S3[:,Ndata_S3,:,:] = (
-                        data_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
-                    )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].mean((-2,-1)) * fft_factor / norm_factor_S3
+                    S3.append(self.backend.bk_reduce_mean(
+                        (self.backend.bk_reshape(data_small,[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(I12_w3_small)
+                    )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy],axis=(-2,-1)) * fft_factor / norm_factor_S3)
                     if get_variance:
-                        S3_sigma[:,Ndata_S3,:,:] = (
-                            data_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
-                            )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].std((-2,-1)) * fft_factor / norm_factor_S3
+                        S3_sigma.apend(self.backend.bk_reduce_std(
+                            (self.backend.bk_reshape(data_small,[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(I12_w3_small)
+                            )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy],axis=(-2,-1)) * fft_factor / norm_factor_S3)
                 if data2 is not None:
                     if not edge:
-                        S3p[:,Ndata_S3,:,:] = (
-                            data2_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
-                        ).mean((-2,-1)) * fft_factor / norm_factor_S3
+                        S3p.append(self.backend.bk_reduce_mean(
+                            (self.backend.bk_reshape(data2_f_small,[N_image2,1,1,1,M3,N3]) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                        ),axis=(-2,-1)) * fft_factor / norm_factor_S3)
                     
                         if get_variance:
-                            S3p_sigma[:,Ndata_S3,:,:] = (
-                                data2_f_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I1_f2_wf3_small)
-                            ).std((-2,-1)) * fft_factor / norm_factor_S3
+                            S3p_sigma.append(self.backend.bk_reduce_std(
+                                (self.backend.bk_reshape(data2_f_small,[N_image2,1,1,1,M3,N3]) * self.backend.bk_conjugate(I1_f2_wf3_small)
+                            ),axis=(-2,-1)) * fft_factor / norm_factor_S3)
                     else:
                     
-                        S3p[:,Ndata_S3,:,:] = (
-                            data2_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
-                        )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].mean((-2,-1)) * fft_factor / norm_factor_S3
+                        S3p.append(self.backend.bk_reduce_mean(
+                            (self.backend.bk_reshape(data2_small,[N_image2,1,1,1,M3,N3]) * self.backend.bk_conjugate(I12_w3_small)
+                        )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy],axis=(-2,-1)) * fft_factor / norm_factor_S3)
                         if get_variance:
-                            S3p_sigma[:,Ndata_S3,:,:] = (
-                                data2_small.view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(I12_w3_small)
-                                )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy].std((-2,-1)) * fft_factor / norm_factor_S3
-                                
-                Ndata_S3+=1
+                            S3p_sigma.append(self.backend.bk_reduce_std(
+                                (self.backend.bk_reshape(data2_small,[N_image2,1,1,1,M3,N3]) * self.backend.bk_conjugate(I12_w3_small)
+                                )[...,edge_dx:M3-edge_dx, edge_dy:N3-edge_dy],axis=(-2,-1)) * fft_factor / norm_factor_S3)
+                              
                 if j2 <= j3:
-                    beg_n=Ndata_S4
+                    if normalization=='S2':
+                        if use_ref: 
+                            P = 1/((ref_S2[:,j3:j3+1,:,None,None] * ref_S2[:,j2:j2+1,None,:,None] )**(0.5*pseudo_coef))
+                        else: 
+                            P = 1/(((S2[:,j3:j3+1,:,None,None] * S2[:,j2:j2+1,None,:,None] )**(0.5*pseudo_coef)))
+                        P=self.backend.bk_complex(P,0.0*P)
+                    else:
+                        P=C_ONE
+                        
                     for j1 in range(0, j2+1):
-                        if eval(S4_criteria):
                             if not edge:
                                 if not if_large_batch:
                                     # [N_image,l1,l2,l3,x,y]
-                                    S4_pre_norm[:,Ndata_S4,:,:,:] = (
-                                        I1_f_small[:,j1].view(N_image,L,1,1,M3,N3) * 
-                                        self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,1,L,L,M3,N3))
-                                    ).mean((-2,-1)) * fft_factor
+                                    S4.append(self.backend.bk_reduce_mean(
+                                        (self.backend.bk_reshape(I1_f_small[:,j1],[N_image,1,L,1,1,M3,N3]) * 
+                                        self.backend.bk_conjugate(self.backend.bk_reshape(I1_f2_wf3_2_small,[N_image,1,1,L,L,M3,N3]))
+                                    ),axis=(-2,-1)) * fft_factor*P)
                                     if get_variance:
-                                        S4_sigma[:,Ndata_S4,:,:,:] = (
-                                            I1_f_small[:,j1].view(N_image,L,1,1,M3,N3) * 
-                                            self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,1,L,L,M3,N3))
-                                        ).std((-2,-1)) * fft_factor
+                                        S4_sigma.append(self.backend.bk_reduce_std(
+                                            (self.backend.bk_reshape(I1_f_small[:,j1],[N_image,1,L,1,1,M3,N3]) * 
+                                            self.backend.bk_conjugate(self.backend.bk_reshape(I1_f2_wf3_2_small,[N_image,1,1,L,L,M3,N3]))
+                                        ),axis=(-2,-1)) * fft_factor*P)
                                 else:
                                     for l1 in range(L):
                                         # [N_image,l2,l3,x,y]
-                                        S4_pre_norm[:,Ndata_S4,l1,:,:] = (
-                                            I1_f_small[:,j1,l1].view(N_image,1,1,M3,N3) * 
-                                            self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,L,L,M3,N3))
-                                        ).mean((-2,-1)) * fft_factor
+                                        S4.append(self.backend.bk_reduce_mean(
+                                            (self.backend.bk_reshape(I1_f_small[:,j1,l1],[N_image,1,1,1,M3,N3]) * 
+                                            self.backend.bk_conjugate(self.backend.bk_reshape(I1_f2_wf3_2_small,[N_image,1,L,L,M3,N3]))
+                                        ),axis=(-2,-1)) * fft_factor*P)
                                         if get_variance:
-                                            S4_sigma[:,Ndata_S4,l1,:,:] = (
-                                                I1_f_small[:,j1,l1].view(N_image,1,1,M3,N3) * 
-                                                self.backend.bk_conjugate(I1_f2_wf3_2_small.view(N_image,L,L,M3,N3))
-                                            ).std((-2,-1)) * fft_factor
+                                            S4_sigma.append(self.backend.bk_reduce_std(
+                                                (self.backend.bk_reshape(I1_f_small[:,j1,l1],[N_image,1,1,1,M3,N3]) * 
+                                                self.backend.bk_conjugate(self.backend.bk_reshape(I1_f2_wf3_2_small,[N_image,1,L,L,M3,N3]))
+                                            ),axis=(-2,-1)) * fft_factor*P)
                             else:
                                 if not if_large_batch:
                                     # [N_image,l1,l2,l3,x,y]
-                                    S4_pre_norm[:,Ndata_S4,:,:,:] = (
-                                        I1_small[:,j1].view(N_image,L,1,1,M3,N3) * self.backend.bk_conjugate(
-                                            I12_w3_2_small.view(N_image,1,L,L,M3,N3)
+                                    S4.append(self.backend.bk_reduce_mean(
+                                        (self.backend.bk_reshape(I1_small[:,j1],[N_image,1,L,1,1,M3,N3]) * self.backend.bk_conjugate(
+                                            self.backend.bk_reshape(I12_w3_2_small,[N_image,1,1,L,L,M3,N3])
                                         )
-                                    )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                    )[...,edge_dx:-edge_dx, edge_dy:-edge_dy],axis=(-2,-1)) * fft_factor*P)
                                     if get_variance:
-                                        S4_sigma[:,Ndata_S4,:,:,:] = (
-                                            I1_small[:,j1].view(N_image,L,1,1,M3,N3) * self.backend.bk_conjugate(
-                                                I12_w3_2_small.view(N_image,1,L,L,M3,N3)
+                                        S4_sigma.append(self.backend.bk_reduce_std(
+                                            (self.backend.bk_reshape(I1_small[:,j1],[N_image,1,L,1,1,M3,N3]) * self.backend.bk_conjugate(
+                                                self.backend.bk_reshape(I12_w3_2_small,[N_image,1,1,L,L,M3,N3])
                                             )
-                                        )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].std((-2,-1)) * fft_factor
+                                        )[...,edge_dx:-edge_dx, edge_dy:-edge_dy],axis=(-2,-1)) * fft_factor*P)
                                 else:
                                     for l1 in range(L):
                                     # [N_image,l2,l3,x,y]
-                                        S4_pre_norm[:,Ndata_S4,l1,:,:] = (
-                                            I1_small[:,j1].view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(
-                                                I12_w3_2_small.view(N_image,L,L,M3,N3)
+                                        S4.append(self.backend.bk_reduce_mean(
+                                            (self.backend.bk_reshape(I1_small[:,j1],[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(
+                                                self.backend.bk_reshape(I12_w3_2_small,[N_image,1,L,L,M3,N3])
                                             )
-                                        )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                        )[...,edge_dx:-edge_dx, edge_dy:-edge_dy],axis=(-2,-1)) * fft_factor*P)
                                         if get_variance:
-                                            S4_sigma[:,Ndata_S4,l1,:,:] = (
-                                                I1_small[:,j1].view(N_image,1,1,M3,N3) * self.backend.bk_conjugate(
-                                                    I12_w3_2_small.view(N_image,L,L,M3,N3)
+                                            S4_sigma.append(self.backend.bk_reduce_std(
+                                                (self.backend.bk_reshape(I1_small[:,j1],[N_image,1,1,1,M3,N3]) * self.backend.bk_conjugate(
+                                                    self.backend.bk_reshape(I12_w3_2_small,[N_image,1,L,L,M3,N3])
                                                 )
-                                            )[...,edge_dx:-edge_dx, edge_dy:-edge_dy].mean((-2,-1)) * fft_factor
+                                            )[...,edge_dx:-edge_dx, edge_dy:-edge_dy],axis=(-2,-1)) * fft_factor*P)
                                 
-                            Ndata_S4+=1
-                            
-                    if normalization=='S2':
-                        if use_ref: 
-                            P = (ref_S2[:,j3,:,None,None] * ref_S2[:,j2,None,:,None] )**(0.5*pseudo_coef)
-                        else: 
-                            P = ((S2[:,j3,:,None,None] * S2[:,j2,None,:,None] )**(0.5*pseudo_coef))
-                        S4[:,beg_n:Ndata_S4,:,:,:]=S4_pre_norm[:,beg_n:Ndata_S4,:,:,:].clone()/(P.clone())
-                            
-                        if get_variance:
-                            S4_sigma[:,beg_n:Ndata_S4,:,:,:] = S4_sigma[:,beg_n:Ndata_S4,:,:,:] / (P)
-                    else:
-                        S4[:,beg_n:Ndata_S4,:,:,:]=S4_pre_norm[:,beg_n:Ndata_S4,:,:,:].clone()
-                            
-                        if get_variance:
-                            S4_sigma[:,beg_n:Ndata_S4,:,:,:] = S4_sigma[:,beg_n:Ndata_S4,:,:,:]
-                            
-        """
-        # define P11 from diagonals of S4
-        for j1 in range(J):
-            for l1 in range(L):
-                P11[:,j1,:,l1,:] = S4_pre_norm[:,j1,j1,:,l1,l1,:].real
-       
-                
-        if normalization=='S4':
-            if use_ref: 
-                P = ref_P11
-            else: 
-                P = P11
-            #.view(N_image,J,1,J,L,1,L) * .view(N_image,1,J,J,1,L,L)
-            S4 = S4_pre_norm / (
-                P[:,:,None,:,:,None,:] * P[:,None,:,:,None,:,:]
-            )**(0.5*pseudo_coef)
-        
-        
+        S3=self.backend.bk_concat(S3,axis=1)
+        S4=self.backend.bk_concat(S4,axis=1)
         
+        if get_variance:
+            S3_sigma=self.backend.bk_concat(S3_sigma,axis=1)
+            S4_sigma=self.backend.bk_concat(S4_sigma,axis=1)
+            
+        if data2 is not None:
+            S3p=self.backend.bk_concat(S3p,axis=1)
+            if get_variance:
+                S3p_sigma=self.backend.bk_concat(S3p_sigma,axis=1)
         
-        # get a single, flattened data vector for_synthesis
-        select_and_index        = self.get_scattering_index(J, L, normalization, S4_criteria)
-        index_for_synthesis     = select_and_index['index_for_synthesis']
-        index_for_synthesis_iso = select_and_index['index_for_synthesis_iso']
-        """
         # average over l1 to obtain simple isotropic statistics
         if iso_ang:
-            S2_iso = S2.mean(-1)
-            S1_iso = S1.mean(-1)
+            S2_iso = self.backend.bk_reduce_mean(S2,axis=(-1))
+            S1_iso = self.backend.bk_reduce_mean(S1,axis=(-1))
             for l1 in range(L):
                 for l2 in range(L):
                     S3_iso[...,(l2-l1)%L] += S3[...,l1,l2]
@@ -5305,8 +5792,8 @@ class funct(FOC.FoCUS):
                 S3p_iso /= L
             
             if get_variance:
-                S2_sigma_iso = S2_sigma.mean(-1)
-                S1_sigma_iso = S1_sigma.mean(-1)
+                S2_sigma_iso = self.backend.bk_reduce_mean(S2_sigma,axis=(-1))
+                S1_sigma_iso = self.backend.bk_reduce_mean(S1_sigma,axis=(-1))
                 for l1 in range(L):
                     for l2 in range(L):
                         S3_sigma_iso[...,(l2-l1)%L] += S3_sigma[...,l1,l2]
@@ -5318,10 +5805,12 @@ class funct(FOC.FoCUS):
                 if data2 is not None:
                     S3p_sigma_iso /= L
         
-        mean_data=self.backend.bk_zeros((N_image,1), dtype=data.dtype) 
-        std_data=self.backend.bk_zeros((N_image,1), dtype=data.dtype) 
-        mean_data[:,0]=data.mean((-2,-1))
-        std_data[:,0]=data.std((-2,-1))
+        if data2 is None:
+            mean_data=self.backend.bk_reshape(self.backend.bk_reduce_mean(data,axis=(-2,-1)),[N_image,1])
+            std_data=self.backend.bk_reshape(self.backend.bk_reduce_std(data,axis=(-2,-1)),[N_image,1])
+        else:
+            mean_data=self.backend.bk_reshape(self.backend.bk_reduce_mean(data*data2,axis=(-2,-1)),[N_image,1])
+            std_data=self.backend.bk_reshape(self.backend.bk_reduce_std(data*data2,axis=(-2,-1)),[N_image,1])
         
         if get_variance:
             ref_sigma={}
@@ -5345,105 +5834,105 @@ class funct(FOC.FoCUS):
         if data2 is None:
             if iso_ang:
                 if ref_sigma is not None:
-                    for_synthesis = self.backend.backend.cat((
+                    for_synthesis = self.backend.bk_concat((
                         mean_data/ref_sigma['std_data'],
                         std_data/ref_sigma['std_data'],
-                        (S2_iso/ref_sigma['S2_sigma']).reshape((N_image, -1)).log(), 
-                        (S1_iso/ref_sigma['S1_sigma']).reshape((N_image, -1)).log(),
-                        (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
-                        (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
-                        (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
-                        (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_log(S2_iso/ref_sigma['S2_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_log(S1_iso/ref_sigma['S1_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3_iso/ref_sigma['S3_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3_iso/ref_sigma['S3_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4_iso/ref_sigma['S4_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4_iso/ref_sigma['S4_sigma']),[N_image, -1]),
+                        ),axis=-1)
                 else:
-                    for_synthesis = self.backend.backend.cat((
+                    for_synthesis = self.backend.bk_concat((
                         mean_data/std_data,
                         std_data,
-                        S2_iso.reshape((N_image, -1)).log(), 
-                        S1_iso.reshape((N_image, -1)).log(),
-                        S3_iso.reshape((N_image, -1)).real, 
-                        S3_iso.reshape((N_image, -1)).imag,
-                        S4_iso.reshape((N_image, -1)).real, 
-                        S4_iso.reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_log(S2_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_log(S1_iso),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3_iso),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4_iso),[N_image, -1]),
+                        ),axis=-1)
             else:
                 if ref_sigma is not None:
-                    for_synthesis = self.backend.backend.cat((
+                    for_synthesis = self.backend.bk_concat((
                         mean_data/ref_sigma['std_data'],
                         std_data/ref_sigma['std_data'],
-                        (S2/ref_sigma['S2_sigma']).reshape((N_image, -1)).log(), 
-                        (S1/ref_sigma['S1_sigma']).reshape((N_image, -1)).log(),
-                        (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
-                        (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
-                        (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
-                        (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_log(S2/ref_sigma['S2_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_log(S1/ref_sigma['S1_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3/ref_sigma['S3_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3/ref_sigma['S3_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4/ref_sigma['S4_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4/ref_sigma['S4_sigma']),[N_image, -1]),
+                        ),axis=-1)
                 else:
-                    for_synthesis = self.backend.backend.cat((
-                        mean_data/std_data,
-                        std_data,
-                        S2.reshape((N_image, -1)).log(), 
-                        S1.reshape((N_image, -1)).log(),
-                        S3.reshape((N_image, -1)).real, 
-                        S3.reshape((N_image, -1)).imag,
-                        S4.reshape((N_image, -1)).real, 
-                        S4.reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                    for_synthesis = self.backend.bk_concat((
+                            mean_data/std_data,
+                            std_data,
+                            self.backend.bk_reshape(self.backend.bk_log(S2),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_log(S1),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_real(S3),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_imag(S3),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_real(S4),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_imag(S4),[N_image, -1])
+                        ),axis=-1)
         else:
             if iso_ang:
                 if ref_sigma is not None:
                     for_synthesis = self.backend.backend.cat((
                         mean_data/ref_sigma['std_data'],
                         std_data/ref_sigma['std_data'],
-                        (S2_iso/ref_sigma['S2_sigma']).reshape((N_image, -1)), 
-                        (S1_iso/ref_sigma['S1_sigma']).reshape((N_image, -1)),
-                        (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
-                        (S3_iso/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
-                        (S3p_iso/ref_sigma['S3p_sigma']).reshape((N_image, -1)).real, 
-                        (S3p_iso/ref_sigma['S3p_sigma']).reshape((N_image, -1)).imag,
-                        (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
-                        (S4_iso/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_real(S2_iso/ref_sigma['S2_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_real(S1_iso/ref_sigma['S1_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3_iso/ref_sigma['S3_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3_iso/ref_sigma['S3_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3p_iso/ref_sigma['S3p_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3p_iso/ref_sigma['S3p_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4_iso/ref_sigma['S4_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4_iso/ref_sigma['S4_sigma']),[N_image, -1]),
+                        ),axis=-1)
                 else:
                     for_synthesis = self.backend.backend.cat((
                         mean_data/std_data,
                         std_data,
-                        S2_iso.reshape((N_image, -1)), 
-                        S1_iso.reshape((N_image, -1)),
-                        S3_iso.reshape((N_image, -1)).real, 
-                        S3_iso.reshape((N_image, -1)).imag,
-                        S3p_iso.reshape((N_image, -1)).real, 
-                        S3p_iso.reshape((N_image, -1)).imag,
-                        S4_iso.reshape((N_image, -1)).real, 
-                        S4_iso.reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_real(S2_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_real(S1_iso),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3_iso),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3p_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3p_iso),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4_iso),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4_iso),[N_image, -1]),
+                        ),axis=-1)
             else:
                 if ref_sigma is not None:
                     for_synthesis = self.backend.backend.cat((
                         mean_data/ref_sigma['std_data'],
                         std_data/ref_sigma['std_data'],
-                        (S2/ref_sigma['S2_sigma']).reshape((N_image, -1)), 
-                        (S1/ref_sigma['S1_sigma']).reshape((N_image, -1)),
-                        (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).real, 
-                        (S3/ref_sigma['S3_sigma']).reshape((N_image, -1)).imag,
-                        (S3p/ref_sigma['S3p_sigma']).reshape((N_image, -1)).real, 
-                        (S3p/ref_sigma['S3p_sigma']).reshape((N_image, -1)).imag,
-                        (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).real, 
-                        (S4/ref_sigma['S4_sigma']).reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                        self.backend.bk_reshape(self.backend.bk_real(S2/ref_sigma['S2_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_real(S1/ref_sigma['S1_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3/ref_sigma['S3_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3/ref_sigma['S3_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S3p/ref_sigma['S3p_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S3p/ref_sigma['S3p_sigma']),[N_image, -1]),
+                        self.backend.bk_reshape(self.backend.bk_real(S4/ref_sigma['S4_sigma']),[N_image, -1]), 
+                        self.backend.bk_reshape(self.backend.bk_imag(S4/ref_sigma['S4_sigma']),[N_image, -1]),
+                        ),axis=-1)
                 else:
-                    for_synthesis = self.backend.backend.cat((
-                        mean_data/std_data,
-                        std_data,
-                        S2.reshape((N_image, -1)), 
-                        S1.reshape((N_image, -1)),
-                        S3.reshape((N_image, -1)).real, 
-                        S3.reshape((N_image, -1)).imag,
-                        S3p.reshape((N_image, -1)).real, 
-                        S3p.reshape((N_image, -1)).imag,
-                        S4.reshape((N_image, -1)).real, 
-                        S4.reshape((N_image, -1)).imag,
-                        ),dim=-1)
+                    for_synthesis = self.backend.bk_concat((
+                            mean_data/std_data,
+                            std_data,
+                            self.backend.bk_reshape(self.backend.bk_real(S2),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_real(S1),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_real(S3),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_imag(S3),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_real(S3p),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_imag(S3p),[N_image, -1]),
+                            self.backend.bk_reshape(self.backend.bk_real(S4),[N_image, -1]), 
+                            self.backend.bk_reshape(self.backend.bk_imag(S4),[N_image, -1])
+                        ),axis=-1)
                 
         if not use_ref: 
             self.ref_scattering_cov_S2=S2
@@ -5835,6 +6324,7 @@ class funct(FOC.FoCUS):
                   to_gaussian=True,
                   use_variance=False,
                   synthesised_N=1,
+                  input_image=None,
                   iso_ang=False,
                   EVAL_FREQUENCY=100,
                   NUM_EPOCHS = 300):
@@ -5892,14 +6382,22 @@ class funct(FOC.FoCUS):
         
         for k in range(nstep):
             if k==0:
-                np.random.seed(seed)
-                if self.use_2D:
-                    imap=np.random.randn(synthesised_N,
-                                         tmp[k].shape[1],
-                                        tmp[k].shape[2])
+                if input_image is None:
+                    np.random.seed(seed)
+                    if self.use_2D:
+                        imap=np.random.randn(synthesised_N,
+                                             tmp[k].shape[1],
+                                            tmp[k].shape[2])
+                    else:
+                        imap=np.random.randn(synthesised_N,
+                                             tmp[k].shape[1])
                 else:
-                    imap=np.random.randn(synthesised_N,
-                                         tmp[k].shape[1])
+                    if self.use_2D:
+                        imap=self.backend.bk_reshape(self.backend.bk_tile(self.backend.bk_cast(input_image.flatten()),synthesised_N),
+                                                        [synthesised_N,tmp[k].shape[1],tmp[k].shape[2]])
+                    else:
+                        imap=self.backend.bk_reshape(self.backend.bk_tile(self.backend.bk_cast(input_image.flatten()),synthesised_N),
+                                                        [synthesised_N,tmp[k].shape[1]])
             else:
                 # Increase the resolution between each step
                 if self.use_2D: