foscat 3.6.1__py3-none-any.whl → 3.7.0__py3-none-any.whl

foscat/scat_cov.py

@@ -33,9 +33,7 @@ testwarn = 0
 
 
 class scat_cov:
-    def __init__(
-        self, s0, s2, s3, s4, s1=None, s3p=None, backend=None, use_1D=False
-    ):
+    def __init__(self, s0, s2, s3, s4, s1=None, s3p=None, backend=None, use_1D=False):
         self.S0 = s0
         self.S2 = s2
         self.S3 = s3
@@ -54,17 +52,17 @@ class scat_cov:
         if self.S1 is None:
             s1 = None
         else:
-            s1 = self.S1.numpy()
+            s1 = self.backend.to_numpy(self.S1)
         if self.S3P is None:
             s3p = None
         else:
-            s3p = self.S3P.numpy()
+            s3p = self.backend.to_numpy(self.S3P)
 
         return scat_cov(
-            (self.S0.numpy()),
-            (self.S2.numpy()),
-            (self.S3.numpy()),
-            (self.S4.numpy()),
+            self.backend.to_numpy(self.S0),
+            self.backend.to_numpy(self.S2),
+            self.backend.to_numpy(self.S3),
+            self.backend.to_numpy(self.S4),
             s1=s1,
             s3p=s3p,
             backend=self.backend,
@@ -94,7 +92,7 @@ class scat_cov:
         )
 
     def conv2complex(self, val):
-        if val.dtype == "complex64" or val.dtype == "complex128":
+        if val.dtype == "complex64" in val.dtype or "complex128" or val.dtype == "torch.complex64" or val.dtype == "torch.complex128" :
             return val
         else:
             return self.backend.bk_complex(val, 0 * val)
@@ -1531,7 +1529,7 @@ class scat_cov:
             if isinstance(x, np.ndarray):
                 return x
             else:
-                return x.numpy()
+                return self.backend.to_numpy(x)
         else:
             return None
 
@@ -1974,7 +1972,7 @@ class scat_cov:
         if self.BACKEND == "numpy":
             s2[:, :, noff:, :] = self.S2
         else:
-            s2[:, :, noff:, :] = self.S2.numpy()
+            s2[:, :, noff:, :] = self.backend.to_numpy(self.S2)
         for i in range(self.S2.shape[0]):
             for j in range(self.S2.shape[1]):
                 for k in range(self.S2.shape[3]):
@@ -1986,7 +1984,7 @@ class scat_cov:
         if self.BACKEND == "numpy":
             s1[:, :, noff:, :] = self.S1
         else:
-            s1[:, :, noff:, :] = self.S1.numpy()
+            s1[:, :, noff:, :] = self.backend.to_numpy(self.S1)
         for i in range(self.S1.shape[0]):
             for j in range(self.S1.shape[1]):
                 for k in range(self.S1.shape[3]):
@@ -2045,12 +2043,12 @@ class scat_cov:
                                     )
                                 )
                             else:
-                                s3[i, j, idx[noff:], k, l_orient] = self.S3.numpy()[
+                                s3[i, j, idx[noff:], k, l_orient] = self.backend.to_numpy(self.S3)[
                                     i, j, j2 == ij - noff, k, l_orient
                                 ]
                                 s3[i, j, idx[:noff], k, l_orient] = (
                                     self.add_data_from_slope(
-                                        self.S3.numpy()[
+                                        self.backend.to_numpy(self.S3)[
                                             i, j, j2 == ij - noff, k, l_orient
                                         ],
                                         noff,
@@ -2358,13 +2356,13 @@ class funct(FOC.FoCUS):
             tmp = self.up_grade(tmp, l_nside * 2, axis=1)
             if image2 is not None:
                 tmpi2 = self.up_grade(tmpi2, l_nside * 2, axis=1)
-
         l_nside = int(np.sqrt(tmp.shape[1] // 12))
         nscale = int(np.log(l_nside) / np.log(2))
         cmat = {}
         cmat2 = {}
+
+        # Loop over scales
         for k in range(nscale):
-            sim = self.backend.bk_abs(self.convol(tmp, axis=1))
             if image2 is not None:
                 sim = self.backend.bk_real(
                     self.backend.bk_L1(
@@ -2375,17 +2373,31 @@ class funct(FOC.FoCUS):
             else:
                 sim = self.backend.bk_abs(self.convol(tmp, axis=1))
 
-            cc = self.backend.bk_reduce_mean(sim[:, :, 0] - sim[:, :, 2], 0)
-            ss = self.backend.bk_reduce_mean(sim[:, :, 1] - sim[:, :, 3], 0)
-            for m in range(smooth_scale):
-                if cc.shape[0] > 12:
-                    cc = self.ud_grade_2(self.smooth(cc))
-                    ss = self.ud_grade_2(self.smooth(ss))
+            # instead of difference between "opposite" channels use weighted average
+            # of cosine and sine contributions using all channels
+            angles = (
+                2 * np.pi * np.arange(self.NORIENT) / self.NORIENT
+            )  # shape: (NORIENT,)
+            angles = angles.reshape(1, 1, self.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)
+            # For simplicity, take first element of the batch
+            cc = weighted_cos[0]
+            ss = weighted_sin[0]
+
+            if smooth_scale > 0:
+                for m in range(smooth_scale):
+                    if cc.shape[0] > 12:
+                        cc = self.ud_grade_2(self.smooth(cc))
+                        ss = self.ud_grade_2(self.smooth(ss))
             if cc.shape[0] != tmp.shape[0]:
                 ll_nside = int(np.sqrt(tmp.shape[1] // 12))
                 cc = self.up_grade(cc, ll_nside)
                 ss = self.up_grade(ss, ll_nside)
 
+            # compute local phase from weighted cos and sin (same as before)
             if self.BACKEND == "numpy":
                 phase = np.fmod(np.arctan2(ss, cc) + 2 * np.pi, 2 * np.pi)
             else:
@@ -2393,87 +2405,87 @@ class funct(FOC.FoCUS):
                     np.arctan2(ss.numpy(), cc.numpy()) + 2 * np.pi, 2 * np.pi
                 )
 
-            iph = (4 * phase / (2 * np.pi)).astype("int")
-            alpha = 4 * phase / (2 * np.pi) - iph
-            mat = np.zeros([sim.shape[1], 4 * 4])
+            # instead of linear interpolation cosine‐based interpolation
+            phase_scaled = self.NORIENT * phase / (2 * np.pi)
+            iph = np.floor(phase_scaled).astype("int")  # lower bin index
+            delta = phase_scaled - iph  # fractional part in [0,1)
+            # interpolation weights
+            w0 = np.cos(delta * np.pi / 2) ** 2
+            w1 = np.sin(delta * np.pi / 2) ** 2
+
+            # build rotation matrix
+            mat = np.zeros([sim.shape[1], self.NORIENT * self.NORIENT])
             lidx = np.arange(sim.shape[1])
-            for l_orient in range(4):
-                mat[lidx, 4 * ((l_orient + iph) % 4) + l_orient] = 1.0 - alpha
-                mat[lidx, 4 * ((l_orient + iph + 1) % 4) + l_orient] = alpha
+            for l in range(self.NORIENT):
+                # Instead of simple linear weights, we use the cosine weights w0 and w1.
+                col0 = self.NORIENT * ((l + iph) % self.NORIENT) + l
+                col1 = self.NORIENT * ((l + iph + 1) % self.NORIENT) + l
+                mat[lidx, col0] = w0
+                mat[lidx, col1] = w1
 
             cmat[k] = self.backend.bk_cast(mat.astype("complex64"))
 
-            mat2 = np.zeros([k + 1, sim.shape[1], 4, 4 * 4])
+            # do same modifications for mat2
+            mat2 = np.zeros(
+                [k + 1, sim.shape[1], self.NORIENT, self.NORIENT * self.NORIENT]
+            )
 
             for k2 in range(k + 1):
-                tmp2 = self.backend.bk_repeat(sim, 4, axis=-1)
+                tmp2 = self.backend.bk_repeat(sim, self.NORIENT, axis=-1)
+
                 sim2 = self.backend.bk_reduce_sum(
                     self.backend.bk_reshape(
-                        mat.reshape(1, mat.shape[0], 16) * tmp2,
-                        [sim.shape[0], cmat[k].shape[0], 4, 4],
+                        mat.reshape(1, mat.shape[0], self.NORIENT * self.NORIENT)
+                        * tmp2,
+                        [sim.shape[0], cmat[k].shape[0], self.NORIENT, self.NORIENT],
                     ),
                     2,
                 )
+
                 sim2 = self.backend.bk_abs(self.convol(sim2, axis=1))
 
-                cc = self.smooth(
-                    self.backend.bk_reduce_mean(sim2[:, :, 0] - sim2[:, :, 2], 0)
+                weighted_cos2 = self.backend.bk_reduce_mean(
+                    sim2 * np.cos(angles), axis=-1
                 )
-                ss = self.smooth(
-                    self.backend.bk_reduce_mean(sim2[:, :, 1] - sim2[:, :, 3], 0)
+                weighted_sin2 = self.backend.bk_reduce_mean(
+                    sim2 * np.sin(angles), axis=-1
                 )
-                for m in range(smooth_scale):
-                    if cc.shape[0] > 12:
-                        cc = self.ud_grade_2(self.smooth(cc))
-                        ss = self.ud_grade_2(self.smooth(ss))
-                if cc.shape[0] != sim.shape[1]:
+
+                cc2 = weighted_cos2[0]
+                ss2 = weighted_sin2[0]
+
+                if smooth_scale > 0:
+                    for m in range(smooth_scale):
+                        if cc2.shape[0] > 12:
+                            cc2 = self.ud_grade_2(self.smooth(cc2))
+                            ss2 = self.ud_grade_2(self.smooth(ss2))
+
+                if cc2.shape[0] != sim.shape[1]:
                     ll_nside = int(np.sqrt(sim.shape[1] // 12))
-                    cc = self.up_grade(cc, ll_nside)
-                    ss = self.up_grade(ss, ll_nside)
+                    cc2 = self.up_grade(cc2, ll_nside)
+                    ss2 = self.up_grade(ss2, ll_nside)
 
                 if self.BACKEND == "numpy":
-                    phase = np.fmod(np.arctan2(ss, cc) + 2 * np.pi, 2 * np.pi)
+                    phase2 = np.fmod(np.arctan2(ss2, cc2) + 2 * np.pi, 2 * np.pi)
                 else:
-                    phase = np.fmod(
-                        np.arctan2(ss.numpy(), cc.numpy()) + 2 * np.pi, 2 * np.pi
+                    phase2 = np.fmod(
+                        np.arctan2(ss2.numpy(), cc2.numpy()) + 2 * np.pi, 2 * np.pi
                     )
-                """
-                for k in range(4):
-                    hp.mollview(np.fmod(phase+np.pi,2*np.pi),cmap='jet',nest=True,hold=False,sub=(2,2,1+k))
-                plt.show()
-                return None
-                """
-                iph = (4 * phase / (2 * np.pi)).astype("int")
-                alpha = 4 * phase / (2 * np.pi) - iph
+
+                phase2_scaled = self.NORIENT * phase2 / (2 * np.pi)
+                iph2 = np.floor(phase2_scaled).astype("int")
+                delta2 = phase2_scaled - iph2
+                w0_2 = np.cos(delta2 * np.pi / 2) ** 2
+                w1_2 = np.sin(delta2 * np.pi / 2) ** 2
                 lidx = np.arange(sim.shape[1])
-                for m in range(4):
-                    for l_orient in range(4):
-                        mat2[
-                            k2, lidx, m, 4 * ((l_orient + iph[:, m]) % 4) + l_orient
-                        ] = (1.0 - alpha[:, m])
-                        mat2[
-                            k2, lidx, m, 4 * ((l_orient + iph[:, m] + 1) % 4) + l_orient
-                        ] = alpha[:, m]
-
-            cmat2[k] = self.backend.bk_cast(mat2.astype("complex64"))
-            """
-            tmp=self.backend.bk_repeat(sim[0],4,axis=1)
-            sim2=self.backend.bk_reduce_sum(self.backend.bk_reshape(mat*tmp,[12*nside**2,4,4]),1)
-
-            cc2=(sim2[:,0]-sim2[:,2])
-            ss2=(sim2[:,1]-sim2[:,3])
-            phase2=np.fmod(np.arctan2(ss2.numpy(),cc2.numpy())+2*np.pi,2*np.pi)
-
-            plt.figure()
-            hp.mollview(phase,cmap='jet',nest=True,hold=False,sub=(2,2,1))
-            hp.mollview(np.fmod(phase2+np.pi,2*np.pi),cmap='jet',nest=True,hold=False,sub=(2,2,2))
-            plt.figure()
-            for k in range(4):
-                hp.mollview((sim[0,:,k]).numpy().real,cmap='jet',nest=True,hold=False,sub=(2,4,1+k),min=-10,max=10)
-                hp.mollview((sim2[:,k]).numpy().real,cmap='jet',nest=True,hold=False,sub=(2,4,5+k),min=-10,max=10)
-
-            plt.show()
-            """
+
+                for m in range(self.NORIENT):
+                    for l in range(self.NORIENT):
+                        col0 = self.NORIENT * ((l + iph2[:, m]) % self.NORIENT) + l
+                        col1 = self.NORIENT * ((l + iph2[:, m] + 1) % self.NORIENT) + l
+                        mat2[k2, lidx, m, col0] = w0_2[:, m]
+                        mat2[k2, lidx, m, col1] = w1_2[:, m]
+                cmat2[k] = self.backend.bk_cast(mat2.astype("complex64"))
 
             if k < l_nside - 1:
                 tmp = self.ud_grade_2(tmp, axis=1)
@@ -2493,11 +2505,12 @@ class funct(FOC.FoCUS):
             image2=None,
             mask=None,
             norm=None,
-            Auto=True,
             calc_var=False,
             cmat=None,
             cmat2=None,
-            out_nside=None
+            Jmax=None,
+            out_nside=None,
+            edge=True
     ):
         """
         Calculates the scattering correlations for a batch of images. Mean are done over pixels.
@@ -2523,9 +2536,6 @@ class funct(FOC.FoCUS):
         norm: None or str
             If None no normalization is applied, if 'auto' normalize by the reference S2,
             if 'self' normalize by the current S2.
-        all_cross: False or True
-            If False compute all the coefficient even the Imaginary part,
-            If True return only the terms computable in the auto case.
         Returns
         -------
         S1, S2, S3, S4 normalized
@@ -2539,7 +2549,7 @@ class funct(FOC.FoCUS):
                 )
                 return None
         if mask is not None:
-            if list(image1.shape) != list(mask.shape)[1:]:
+            if image1.shape[-2] != mask.shape[1] or image1.shape[-1] != mask.shape[2]:
                 print(
                     "The LAST COLUMN of the mask should have the same size ",
                     mask.shape,
@@ -2558,9 +2568,6 @@ class funct(FOC.FoCUS):
         cross = False
         if image2 is not None:
             cross = True
-            all_cross = Auto
-        else:
-            all_cross = False
 
         ### PARAMETERS
         axis = 1
@@ -2596,8 +2603,16 @@ class funct(FOC.FoCUS):
             nside = int(np.sqrt(npix // 12))
 
             J = int(np.log(nside) / np.log(2))  # Number of j scales
-
-        Jmax = J - self.OSTEP  # Number of steps for the loop on scales
+            
+        if (self.use_2D or self.use_1D) and self.KERNELSZ>3:
+            J-=1
+        if Jmax is None:
+            Jmax = J  # Number of steps for the loop on scales
+        if Jmax>J:
+            print('==========\n\n')
+            print('The Jmax you requested is larger than the data size, which may cause problems while computing the scattering transform.')
+            print('\n\n==========')
+        
 
         ### LOCAL VARIABLES (IMAGES and MASK)
         if len(image1.shape) == 1 or (len(image1.shape) == 2 and self.use_2D):
@@ -2621,7 +2636,7 @@ class funct(FOC.FoCUS):
         else:
             vmask = self.backend.bk_cast(mask)  # [Nmask, Npix]
 
-        if self.KERNELSZ > 3:
+        if self.KERNELSZ > 3 and not self.use_2D:
             # if the kernel size is bigger than 3 increase the binning before smoothing
             if self.use_2D:
                 vmask = self.up_grade(
@@ -2645,7 +2660,7 @@ class funct(FOC.FoCUS):
                 if cross:
                     I2 = self.up_grade(I2, nside * 2, axis=axis)
 
-            if self.KERNELSZ > 5:
+            if self.KERNELSZ > 5 and not self.use_2D:
                 # if the kernel size is bigger than 3 increase the binning before smoothing
                 if self.use_2D:
                     vmask = self.up_grade(
@@ -2677,7 +2692,23 @@ class funct(FOC.FoCUS):
 
         ### INITIALIZATION
         # Coefficients
-        S1, S2, S3, S4, S3P = None, None, None, None, None
+        if return_data:
+            S1 = {}
+            S2 = {}
+            S3 = {}
+            S3P = {}
+            S4 = {}
+        else:
+            S1 = []
+            S2 = []
+            S3 = []
+            S4 = []
+            S3P = []
+            VS1 = []
+            VS2 = []
+            VS3 = []
+            VS3P = []
+            VS4 = []
 
         off_S2 = -2
         off_S3 = -3
@@ -2687,11 +2718,6 @@ class funct(FOC.FoCUS):
             off_S3 = -1
             off_S4 = -1
 
-        # Dictionaries for S3 computation
-        M1_dic = {}  # M stands for Module M1 = |I1 * Psi|
-        if cross:
-            M2_dic = {}
-
         # S2 for normalization
         cond_init_P1_dic = (norm == "self") or (
             (norm == "auto") and (self.P1_dic is None)
@@ -2710,7 +2736,12 @@ class funct(FOC.FoCUS):
         if return_data:
             s0 = I1
             if out_nside is not None:
-                s0 = self.backend.bk_reduce_mean(self.backend.bk_reshape(s0,[s0.shape[0],12*out_nside**2,(nside//out_nside)**2]),2)
+                s0 = self.backend.bk_reduce_mean(
+                    self.backend.bk_reshape(
+                        s0, [s0.shape[0], 12 * out_nside**2, (nside // out_nside) ** 2]
+                    ),
+                    2,
+                )
         else:
             if not cross:
                 s0, l_vs0 = self.masked_mean(I1, vmask, axis=1, calc_var=True)
@@ -2720,17 +2751,37 @@ class funct(FOC.FoCUS):
                 )
             vs0 = self.backend.bk_concat([l_vs0, l_vs0], 1)
             s0 = self.backend.bk_concat([s0, l_vs0], 1)
-
         #### COMPUTE S1, S2, S3 and S4
         nside_j3 = nside  # NSIDE start (nside_j3 = nside / 2^j3)
+
+        # a remettre comme avant
+        M1_dic={}
+        
         for j3 in range(Jmax):
+            
+            if edge:
+                if self.mask_mask is None:
+                    self.mask_mask={}
+                if self.use_2D:
+                    if (vmask.shape[1],vmask.shape[2]) not in self.mask_mask:
+                        mask_mask=np.zeros([1,vmask.shape[1],vmask.shape[2]])
+                        mask_mask[0,
+                                  self.KERNELSZ//2:-self.KERNELSZ//2+1,
+                                  self.KERNELSZ//2:-self.KERNELSZ//2+1]=1.0
+                        self.mask_mask[(vmask.shape[1],vmask.shape[2])]=self.backend.bk_cast(mask_mask)
+                    vmask=vmask*self.mask_mask[(vmask.shape[1],vmask.shape[2])]
+                    #print(self.KERNELSZ//2,vmask,mask_mask)
+                    
+                if self.use_1D:
+                    if (vmask.shape[1]) not in self.mask_mask:
+                        mask_mask=np.zeros([1,vmask.shape[1]])
+                        mask_mask[0,
+                                  self.KERNELSZ//2:-self.KERNELSZ//2+1]=1.0
+                        self.mask_mask[(vmask.shape[1])]=self.backend.bk_cast(mask_mask)
+                    vmask=vmask*self.mask_mask[(vmask.shape[1])]
+                    
             if return_data:
-                if S3 is None:
-                    S3 = {}
                 S3[j3] = None
-
-                if S3P is None:
-                    S3P = {}
                 S3P[j3] = None
 
                 if S4 is None:
@@ -2740,12 +2791,12 @@ class funct(FOC.FoCUS):
             ####### S1 and S2
             ### Make the convolution I1 * Psi_j3
             conv1 = self.convol(I1, axis=1)  # [Nbatch, Npix_j3, Norient3]
-
             if cmat is not None:
-                tmp2 = self.backend.bk_repeat(conv1, 4, axis=-1)
+                tmp2 = self.backend.bk_repeat(conv1, self.NORIENT, axis=-1)
                 conv1 = self.backend.bk_reduce_sum(
                     self.backend.bk_reshape(
-                        cmat[j3] * tmp2, [tmp2.shape[0], cmat[j3].shape[0], 4, 4]
+                        cmat[j3] * tmp2,
+                        [tmp2.shape[0], cmat[j3].shape[0], self.NORIENT, self.NORIENT],
                     ),
                     2,
                 )
@@ -2781,33 +2832,31 @@ class funct(FOC.FoCUS):
                 if return_data:
                     if S2 is None:
                         S2 = {}
-                    if out_nside is not None and out_nside<nside_j3:
+                    if out_nside is not None and out_nside < nside_j3:
                         s2 = self.backend.bk_reduce_mean(
-                            self.backend.bk_reshape(s2,[s2.shape[0],
-                                                        12*out_nside**2,
-                                                        (nside_j3//out_nside)**2,
-                                                        s2.shape[2]]),2)
+                            self.backend.bk_reshape(
+                                s2,
+                                [
+                                    s2.shape[0],
+                                    12 * out_nside**2,
+                                    (nside_j3 // out_nside) ** 2,
+                                    s2.shape[2],
+                                ],
+                            ),
+                            2,
+                        )
                     S2[j3] = s2
                 else:
                     if norm == "auto":  # Normalize S2
                         s2 /= P1_dic[j3]
-                    if S2 is None:
-                        S2 = self.backend.bk_expand_dims(
-                            s2, off_S2
+
+                    S2.append(
+                        self.backend.bk_expand_dims(s2, off_S2)
+                    )  # Add a dimension for NS2
+                    if calc_var:
+                        VS2.append(
+                            self.backend.bk_expand_dims(vs2, off_S2)
                         )  # Add a dimension for NS2
-                        if calc_var:
-                            VS2 = self.backend.bk_expand_dims(
-                                vs2, off_S2
-                            )  # Add a dimension for NS2
-                    else:
-                        S2 = self.backend.bk_concat(
-                            [S2, self.backend.bk_expand_dims(s2, off_S2)], axis=2
-                        )
-                        if calc_var:
-                            VS2 = self.backend.bk_concat(
-                                [VS2, self.backend.bk_expand_dims(vs2, off_S2)],
-                                axis=2,
-                            )
 
                 #### S1_auto computation
                 ### Image 1 : S1 = < M1 >_pix
@@ -2825,45 +2874,47 @@ class funct(FOC.FoCUS):
                         )  # [Nbatch, Nmask, Norient3]
 
                 if return_data:
-                    if S1 is None:
-                        S1 = {}
-                    if out_nside is not None and out_nside<nside_j3:
+                    if out_nside is not None and out_nside < nside_j3:
                         s1 = self.backend.bk_reduce_mean(
-                            self.backend.bk_reshape(s1,[s1.shape[0],
-                                                        12*out_nside**2,
-                                                        (nside_j3//out_nside)**2,
-                                                        s1.shape[2]]),2)
+                            self.backend.bk_reshape(
+                                s1,
+                                [
+                                    s1.shape[0],
+                                    12 * out_nside**2,
+                                    (nside_j3 // out_nside) ** 2,
+                                    s1.shape[2],
+                                ],
+                            ),
+                            2,
+                        )
                     S1[j3] = s1
                 else:
                     ### Normalize S1
                     if norm is not None:
                         self.div_norm(s1, (P1_dic[j3]) ** 0.5)
                     ### We store S1 for image1  [Nbatch, Nmask, NS1, Norient3]
-                    if S1 is None:
-                        S1 = self.backend.bk_expand_dims(
-                            s1, off_S2
+                    S1.append(
+                        self.backend.bk_expand_dims(s1, off_S2)
+                    )  # Add a dimension for NS1
+                    if calc_var:
+                        VS1.append(
+                            self.backend.bk_expand_dims(vs1, off_S2)
                         )  # Add a dimension for NS1
-                        if calc_var:
-                            VS1 = self.backend.bk_expand_dims(
-                                vs1, off_S2
-                            )  # Add a dimension for NS1
-                    else:
-                        S1 = self.backend.bk_concat(
-                            [S1, self.backend.bk_expand_dims(s1, off_S2)], axis=2
-                        )
-                        if calc_var:
-                            VS1 = self.backend.bk_concat(
-                                [VS1, self.backend.bk_expand_dims(vs1, off_S2)], axis=2
-                            )
 
             else:  # Cross
                 ### Make the convolution I2 * Psi_j3
                 conv2 = self.convol(I2, axis=1)  # [Nbatch, Npix_j3, Norient3]
                 if cmat is not None:
-                    tmp2 = self.backend.bk_repeat(conv2, 4, axis=-1)
+                    tmp2 = self.backend.bk_repeat(conv2, self.NORIENT, axis=-1)
                     conv2 = self.backend.bk_reduce_sum(
                         self.backend.bk_reshape(
-                            cmat[j3] * tmp2, [tmp2.shape[0], cmat[j3].shape[0], 4, 4]
+                            cmat[j3] * tmp2,
+                            [
+                                tmp2.shape[0],
+                                cmat[j3].shape[0],
+                                self.NORIENT,
+                                self.NORIENT,
+                            ],
                         ),
                         2,
                     )
@@ -2917,14 +2968,19 @@ class funct(FOC.FoCUS):
                         s2 = self.masked_mean(s2, vmask, axis=1, rank=j3)
 
                 if return_data:
-                    if S2 is None:
-                        S2 = {}
-                    if out_nside is not None and out_nside<nside_j3:
+                    if out_nside is not None and out_nside < nside_j3:
                         s2 = self.backend.bk_reduce_mean(
-                            self.backend.bk_reshape(s2,[s2.shape[0],
-                                                        12*out_nside**2,
-                                                        (nside_j3//out_nside)**2,
-                                                        s2.shape[2]]),2)
+                            self.backend.bk_reshape(
+                                s2,
+                                [
+                                    s2.shape[0],
+                                    12 * out_nside**2,
+                                    (nside_j3 // out_nside) ** 2,
+                                    s2.shape[2],
+                                ],
+                            ),
+                            2,
+                        )
                     S2[j3] = s2
                 else:
                     ### Normalize S2_cross
@@ -2932,26 +2988,15 @@ class funct(FOC.FoCUS):
                         s2 /= (P1_dic[j3] * P2_dic[j3]) ** 0.5
 
                     ### Store S2_cross as complex [Nbatch, Nmask, NS2, Norient3]
-                    if not all_cross:
-                        s2 = self.backend.bk_real(s2)
+                    s2 = self.backend.bk_real(s2)
 
-                    if S2 is None:
-                        S2 = self.backend.bk_expand_dims(
-                            s2, off_S2
+                    S2.append(
+                        self.backend.bk_expand_dims(s2, off_S2)
+                    )  # Add a dimension for NS2
+                    if calc_var:
+                        VS2.append(
+                            self.backend.bk_expand_dims(vs2, off_S2)
                         )  # Add a dimension for NS2
-                        if calc_var:
-                            VS2 = self.backend.bk_expand_dims(
-                                vs2, off_S2
-                            )  # Add a dimension for NS2
-                    else:
-                        S2 = self.backend.bk_concat(
-                            [S2, self.backend.bk_expand_dims(s2, off_S2)], axis=2
-                        )
-                        if calc_var:
-                            VS2 = self.backend.bk_concat(
-                                [VS2, self.backend.bk_expand_dims(vs2, off_S2)],
-                                axis=2,
-                            )
 
                 #### S1_auto computation
                 ### Image 1 : S1 = < M1 >_pix
@@ -2968,36 +3013,32 @@ class funct(FOC.FoCUS):
                             MX, vmask, axis=1, rank=j3
                         )  # [Nbatch, Nmask, Norient3]
                 if return_data:
-                    if S1 is None:
-                        S1 = {}
-                    if out_nside is not None and out_nside<nside_j3:
+                    if out_nside is not None and out_nside < nside_j3:
                         s1 = self.backend.bk_reduce_mean(
-                            self.backend.bk_reshape(s1,[s1.shape[0],
-                                                        12*out_nside**2,
-                                                        (nside_j3//out_nside)**2,
-                                                        s1.shape[2]]),2)
+                            self.backend.bk_reshape(
+                                s1,
+                                [
+                                    s1.shape[0],
+                                    12 * out_nside**2,
+                                    (nside_j3 // out_nside) ** 2,
+                                    s1.shape[2],
+                                ],
+                            ),
+                            2,
+                        )
                     S1[j3] = s1
                 else:
                     ### Normalize S1
                     if norm is not None:
                         self.div_norm(s1, (P1_dic[j3]) ** 0.5)
                     ### We store S1 for image1  [Nbatch, Nmask, NS1, Norient3]
-                    if S1 is None:
-                        S1 = self.backend.bk_expand_dims(
-                            s1, off_S2
+                    S1.append(
+                        self.backend.bk_expand_dims(s1, off_S2)
+                    )  # Add a dimension for NS1
+                    if calc_var:
+                        VS1.append(
+                            self.backend.bk_expand_dims(vs1, off_S2)
                         )  # Add a dimension for NS1
-                        if calc_var:
-                            VS1 = self.backend.bk_expand_dims(
-                                vs1, off_S2
-                            )  # Add a dimension for NS1
-                    else:
-                        S1 = self.backend.bk_concat(
-                            [S1, self.backend.bk_expand_dims(s1, off_S2)], axis=2
-                        )
-                        if calc_var:
-                            VS1 = self.backend.bk_concat(
-                                [VS1, self.backend.bk_expand_dims(vs1, off_S2)], axis=2
-                            )
 
             # Initialize dictionaries for |I1*Psi_j| * Psi_j3
             M1convPsi_dic = {}
@@ -3006,7 +3047,7 @@ class funct(FOC.FoCUS):
                 M2convPsi_dic = {}
 
             ###### S3
-            nside_j2=nside_j3
+            nside_j2 = nside_j3
             for j2 in range(0, j3 + 1):  # j2 <= j3
                 if return_data:
                     if S4[j3] is None:
@@ -3041,13 +3082,20 @@ class funct(FOC.FoCUS):
                     if return_data:
                         if S3[j3] is None:
                             S3[j3] = {}
-                        if out_nside is not None and out_nside<nside_j2:
+                        if out_nside is not None and out_nside < nside_j2:
                             s3 = self.backend.bk_reduce_mean(
-                                self.backend.bk_reshape(s3,[s3.shape[0],
-                                                            12*out_nside**2,
-                                                            (nside_j2//out_nside)**2,
-                                                            s3.shape[2],
-                                                            s3.shape[3]]),2)
+                                self.backend.bk_reshape(
+                                    s3,
+                                    [
+                                        s3.shape[0],
+                                        12 * out_nside**2,
+                                        (nside_j2 // out_nside) ** 2,
+                                        s3.shape[2],
+                                        s3.shape[3],
+                                    ],
+                                ),
+                                2,
+                            )
                         S3[j3][j2] = s3
                     else:
                         ### Normalize S3 with S2_j [Nbatch, Nmask, Norient_j]
@@ -3062,23 +3110,18 @@ class funct(FOC.FoCUS):
                             )  # [Nbatch, Nmask, Norient3, Norient2]
 
                         ### Store S3 as a complex [Nbatch, Nmask, NS3, Norient3, Norient2]
-                        if S3 is None:
-                            S3 = self.backend.bk_expand_dims(
-                                s3, off_S3
-                            )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3 = self.backend.bk_expand_dims(
-                                    vs3, off_S3
-                                )  # Add a dimension for NS3
-                        else:
-                            S3 = self.backend.bk_concat(
-                                [S3, self.backend.bk_expand_dims(s3, off_S3)], axis=2
+
+                        # S3.append(self.backend.bk_reshape(s3,[s3.shape[0],s3.shape[1],
+                        #                                      s3.shape[2]*s3.shape[3]]))
+                        S3.append(
+                            self.backend.bk_expand_dims(s3, off_S3)
+                        )  # Add a dimension for NS3
+                        if calc_var:
+                            VS3.append(
+                                self.backend.bk_expand_dims(vs3, off_S3)
                             )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3 = self.backend.bk_concat(
-                                    [VS3, self.backend.bk_expand_dims(vs3, off_S3)],
-                                    axis=2,
-                                )  # Add a dimension for NS3
+                            # VS3.append(self.backend.bk_reshape(vs3,[s3.shape[0],s3.shape[1],
+                            #                                  s3.shape[2]*s3.shape[3]]))
 
                 ### S3_cross = < (I1 * Psi)_j3 x (|I2 * Psi_j2| * Psi_j3)^* >_pix
                 ### S3P_cross = < (I2 * Psi)_j3 x (|I1 * Psi_j2| * Psi_j3)^* >_pix
@@ -3130,19 +3173,33 @@ class funct(FOC.FoCUS):
                         if S3[j3] is None:
                             S3[j3] = {}
                             S3P[j3] = {}
-                        if out_nside is not None and out_nside<nside_j2:
+                        if out_nside is not None and out_nside < nside_j2:
                             s3 = self.backend.bk_reduce_mean(
-                                self.backend.bk_reshape(s3,[s3.shape[0],
-                                                            12*out_nside**2,
-                                                            (nside_j2//out_nside)**2,
-                                                            s3.shape[2],
-                                                            s3.shape[3]]),2)
+                                self.backend.bk_reshape(
+                                    s3,
+                                    [
+                                        s3.shape[0],
+                                        12 * out_nside**2,
+                                        (nside_j2 // out_nside) ** 2,
+                                        s3.shape[2],
+                                        s3.shape[3],
+                                    ],
+                                ),
+                                2,
+                            )
                             s3p = self.backend.bk_reduce_mean(
-                                self.backend.bk_reshape(s3p,[s3.shape[0],
-                                                             12*out_nside**2,
-                                                             (nside_j2//out_nside)**2,
-                                                             s3.shape[2],
-                                                             s3.shape[3]]),2)
+                                self.backend.bk_reshape(
+                                    s3p,
+                                    [
+                                        s3.shape[0],
+                                        12 * out_nside**2,
+                                        (nside_j2 // out_nside) ** 2,
+                                        s3.shape[2],
+                                        s3.shape[3],
+                                    ],
+                                ),
+                                2,
+                            )
                         S3[j3][j2] = s3
                         S3P[j3][j2] = s3p
                     else:
@@ -3166,43 +3223,34 @@ class funct(FOC.FoCUS):
                             )  # [Nbatch, Nmask, Norient3, Norient2]
 
                         ### Store S3 and S3P as a complex [Nbatch, Nmask, NS3, Norient3, Norient2]
-                        if S3 is None:
-                            S3 = self.backend.bk_expand_dims(
-                                s3, off_S3
-                            )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3 = self.backend.bk_expand_dims(
-                                    vs3, off_S3
-                                )  # Add a dimension for NS3
-                        else:
-                            S3 = self.backend.bk_concat(
-                                [S3, self.backend.bk_expand_dims(s3, off_S3)], axis=2
-                            )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3 = self.backend.bk_concat(
-                                    [VS3, self.backend.bk_expand_dims(vs3, off_S3)],
-                                    axis=2,
-                                )  # Add a dimension for NS3
-                        if S3P is None:
-                            S3P = self.backend.bk_expand_dims(
-                                s3p, off_S3
+
+                        # S3.append(self.backend.bk_reshape(s3,[s3.shape[0],s3.shape[1],
+                        #                                      s3.shape[2]*s3.shape[3]]))
+                        S3.append(
+                            self.backend.bk_expand_dims(s3, off_S3)
+                        )  # Add a dimension for NS3
+                        if calc_var:
+                            VS3.append(
+                                self.backend.bk_expand_dims(vs3, off_S3)
                             )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3P = self.backend.bk_expand_dims(
-                                    vs3p, off_S3
-                                )  # Add a dimension for NS3
-                        else:
-                            S3P = self.backend.bk_concat(
-                                [S3P, self.backend.bk_expand_dims(s3p, off_S3)], axis=2
+
+                            # VS3.append(self.backend.bk_reshape(vs3,[s3.shape[0],s3.shape[1],
+                            #                                  s3.shape[2]*s3.shape[3]]))
+
+                        # S3P.append(self.backend.bk_reshape(s3p,[s3.shape[0],s3.shape[1],
+                        #                                      s3.shape[2]*s3.shape[3]]))
+                        S3P.append(
+                            self.backend.bk_expand_dims(s3p, off_S3)
+                        )  # Add a dimension for NS3
+                        if calc_var:
+                            VS3P.append(
+                                self.backend.bk_expand_dims(vs3p, off_S3)
                             )  # Add a dimension for NS3
-                            if calc_var:
-                                VS3P = self.backend.bk_concat(
-                                    [VS3P, self.backend.bk_expand_dims(vs3p, off_S3)],
-                                    axis=2,
-                                )  # Add a dimension for NS3
+                            # VS3P.append(self.backend.bk_reshape(vs3p,[s3.shape[0],s3.shape[1],
+                            #                                  s3.shape[2]*s3.shape[3]]))
 
                 ##### S4
-                nside_j1=nside_j2
+                nside_j1 = nside_j2
                 for j1 in range(0, j2 + 1):  # j1 <= j2
                     ### S4_auto = <(|I1 * psi1| * psi3)(|I1 * psi2| * psi3)^*>
                     if not cross:
@@ -3228,14 +3276,21 @@ class funct(FOC.FoCUS):
                         if return_data:
                             if S4[j3][j2] is None:
                                 S4[j3][j2] = {}
-                            if out_nside is not None and out_nside<nside_j1:
+                            if out_nside is not None and out_nside < nside_j1:
                                 s4 = self.backend.bk_reduce_mean(
-                                    self.backend.bk_reshape(s4,[s4.shape[0],
-                                                                12*out_nside**2,
-                                                                (nside_j1//out_nside)**2,
-                                                                s4.shape[2],
-                                                                s4.shape[3],
-                                                                s4.shape[4]]),2)
+                                    self.backend.bk_reshape(
+                                        s4,
+                                        [
+                                            s4.shape[0],
+                                            12 * out_nside**2,
+                                            (nside_j1 // out_nside) ** 2,
+                                            s4.shape[2],
+                                            s4.shape[3],
+                                            s4.shape[4],
+                                        ],
+                                    ),
+                                    2,
+                                )
                             S4[j3][j2][j1] = s4
                         else:
                             ### Normalize S4 with S2_j [Nbatch, Nmask, Norient_j]
@@ -3259,27 +3314,18 @@ class funct(FOC.FoCUS):
                                     ** 0.5,
                                 )  # [Nbatch, Nmask, Norient3, Norient2, Norient1]
                             ### Store S4 as a complex [Nbatch, Nmask, NS4, Norient3, Norient2, Norient1]
-                            if S4 is None:
-                                S4 = self.backend.bk_expand_dims(
-                                    s4, off_S4
-                                )  # Add a dimension for NS4
-                                if calc_var:
-                                    VS4 = self.backend.bk_expand_dims(
-                                        vs4, off_S4
-                                    )  # Add a dimension for NS4
-                            else:
-                                S4 = self.backend.bk_concat(
-                                    [S4, self.backend.bk_expand_dims(s4, off_S4)],
-                                    axis=2,
+
+                            # S4.append(self.backend.bk_reshape(s4,[s4.shape[0],s4.shape[1],
+                            #                                  s4.shape[2]*s4.shape[3]*s4.shape[4]]))
+                            S4.append(
+                                self.backend.bk_expand_dims(s4, off_S4)
+                            )  # Add a dimension for NS4
+                            if calc_var:
+                                # VS4.append(self.backend.bk_reshape(vs4,[s4.shape[0],s4.shape[1],
+                                #                              s4.shape[2]*s4.shape[3]*s4.shape[4]]))
+                                VS4.append(
+                                    self.backend.bk_expand_dims(vs4, off_S4)
                                 )  # Add a dimension for NS4
-                                if calc_var:
-                                    VS4 = self.backend.bk_concat(
-                                        [
-                                            VS4,
-                                            self.backend.bk_expand_dims(vs4, off_S4),
-                                        ],
-                                        axis=2,
-                                    )  # Add a dimension for NS4
 
                         ### S4_cross = <(|I1 * psi1| * psi3)(|I2 * psi2| * psi3)^*>
                     else:
@@ -3305,14 +3351,21 @@ class funct(FOC.FoCUS):
                         if return_data:
                             if S4[j3][j2] is None:
                                 S4[j3][j2] = {}
-                            if out_nside is not None and out_nside<nside_j1:
+                            if out_nside is not None and out_nside < nside_j1:
                                 s4 = self.backend.bk_reduce_mean(
-                                    self.backend.bk_reshape(s4,[s4.shape[0],
-                                                                12*out_nside**2,
-                                                                (nside_j1//out_nside)**2,
-                                                                s4.shape[2],
-                                                                s4.shape[3],
-                                                                s4.shape[4]]),2)
+                                    self.backend.bk_reshape(
+                                        s4,
+                                        [
+                                            s4.shape[0],
+                                            12 * out_nside**2,
+                                            (nside_j1 // out_nside) ** 2,
+                                            s4.shape[2],
+                                            s4.shape[3],
+                                            s4.shape[4],
+                                        ],
+                                    ),
+                                    2,
+                                )
                             S4[j3][j2][j1] = s4
                         else:
                             ### Normalize S4 with S2_j [Nbatch, Nmask, Norient_j]
@@ -3336,41 +3389,33 @@ class funct(FOC.FoCUS):
                                     ** 0.5,
                                 )  # [Nbatch, Nmask, Norient3, Norient2, Norient1]
                             ### Store S4 as a complex [Nbatch, Nmask, NS4, Norient3, Norient2, Norient1]
-                            if S4 is None:
-                                S4 = self.backend.bk_expand_dims(
-                                    s4, off_S4
-                                )  # Add a dimension for NS4
-                                if calc_var:
-                                    VS4 = self.backend.bk_expand_dims(
-                                        vs4, off_S4
-                                    )  # Add a dimension for NS4
-                            else:
-                                S4 = self.backend.bk_concat(
-                                    [S4, self.backend.bk_expand_dims(s4, off_S4)],
-                                    axis=2,
+                            # S4.append(self.backend.bk_reshape(s4,[s4.shape[0],s4.shape[1],
+                            #                                  s4.shape[2]*s4.shape[3]*s4.shape[4]]))
+                            S4.append(
+                                self.backend.bk_expand_dims(s4, off_S4)
+                            )  # Add a dimension for NS4
+                            if calc_var:
+
+                                # VS4.append(self.backend.bk_reshape(vs4,[s4.shape[0],s4.shape[1],
+                                #                              s4.shape[2]*s4.shape[3]*s4.shape[4]]))
+                                VS4.append(
+                                    self.backend.bk_expand_dims(vs4, off_S4)
                                 )  # Add a dimension for NS4
-                                if calc_var:
-                                    VS4 = self.backend.bk_concat(
-                                        [
-                                            VS4,
-                                            self.backend.bk_expand_dims(vs4, off_S4),
-                                        ],
-                                        axis=2,
-                                    )  # Add a dimension for NS4
-                            nside_j1=nside_j1 // 2
-                        nside_j2=nside_j2 // 2
-                        
+
+                            nside_j1 = nside_j1 // 2
+                        nside_j2 = nside_j2 // 2
+
             ###### Reshape for next iteration on j3
             ### Image I1,
             # downscale the I1 [Nbatch, Npix_j3]
             if j3 != Jmax - 1:
-                I1_smooth = self.smooth(I1, axis=1)
-                I1 = self.ud_grade_2(I1_smooth, axis=1)
+                I1 = self.smooth(I1, axis=1)
+                I1 = self.ud_grade_2(I1, axis=1)
 
                 ### Image I2
                 if cross:
-                    I2_smooth = self.smooth(I2, axis=1)
-                    I2 = self.ud_grade_2(I2_smooth, axis=1)
+                    I2 = self.smooth(I2, axis=1)
+                    I2 = self.ud_grade_2(I2, axis=1)
 
                 ### Modules
                 for j2 in range(0, j3 + 1):  # j2 =< j3
@@ -3388,8 +3433,9 @@ class funct(FOC.FoCUS):
                             M2_dic[j2], axis=1
                         )  # [Nbatch, Npix_j3, Norient3]
                         M2_dic[j2] = self.ud_grade_2(
-                            M2_smooth, axis=1
+                            M2, axis=1
                         )  # [Nbatch, Npix_j3, Norient3]
+            
                 ### Mask
                 vmask = self.ud_grade_2(vmask, axis=1)
 
@@ -3404,7 +3450,33 @@ class funct(FOC.FoCUS):
             self.P1_dic = P1_dic
             if cross:
                 self.P2_dic = P2_dic
+        """
+        Sout=[s0]+S1+S2+S3+S4
 
+        if cross:
+            Sout=Sout+S3P
+        if calc_var:
+            SVout=[vs0]+VS1+VS2+VS3+VS4
+            if cross:
+                VSout=VSout+VS3P
+            return self.backend.bk_concat(Sout, 2),self.backend.bk_concat(VSout, 2)
+
+        return self.backend.bk_concat(Sout, 2)
+        """
+        if not return_data:
+            S1 = self.backend.bk_concat(S1, 2)
+            S2 = self.backend.bk_concat(S2, 2)
+            S3 = self.backend.bk_concat(S3, 2)
+            S4 = self.backend.bk_concat(S4, 2)
+            if cross:
+                S3P = self.backend.bk_concat(S3P, 2)
+            if calc_var:
+                VS1 = self.backend.bk_concat(VS1, 2)
+                VS2 = self.backend.bk_concat(VS2, 2)
+                VS3 = self.backend.bk_concat(VS3, 2)
+                VS4 = self.backend.bk_concat(VS4, 2)
+                if cross:
+                    VS3P = self.backend.bk_concat(VS3P, 2)
         if calc_var:
             if not cross:
                 return scat_cov(
@@ -3487,10 +3559,17 @@ class funct(FOC.FoCUS):
             M_dic[j2], axis=1
         )  # [Nbatch, Npix_j3, Norient3, Norient2]
         if cmat2 is not None:
-            tmp2 = self.backend.bk_repeat(MconvPsi, 4, axis=-1)
+            tmp2 = self.backend.bk_repeat(MconvPsi, self.NORIENT, axis=-1)
             MconvPsi = self.backend.bk_reduce_sum(
                 self.backend.bk_reshape(
-                    cmat2[j3][j2] * tmp2, [tmp2.shape[0], cmat2[j3].shape[1], 4, 4, 4]
+                    cmat2[j3][j2] * tmp2,
+                    [
+                        tmp2.shape[0],
+                        cmat2[j3].shape[1],
+                        self.NORIENT,
+                        self.NORIENT,
+                        self.NORIENT,
+                    ],
                 ),
                 3,
             )
@@ -3566,6 +3645,27 @@ class funct(FOC.FoCUS):
                 )  # [Nbatch, Nmask, Norient3, Norient2, Norient1]
                 return s4
 
+    def to_gaussian(self,x):
+        from scipy.stats import norm
+        from scipy.interpolate import interp1d
+
+        idx=np.argsort(x.flatten())
+        p = (np.arange(1, idx.shape[0] + 1) - 0.5) / idx.shape[0]
+        im_target=x.flatten()
+        im_target[idx] = norm.ppf(p)
+        
+        # Interpolation cubique
+        self.f_gaussian = interp1d(im_target[idx], x.flatten()[idx], kind='cubic')
+        self.val_min=im_target[idx[0]]
+        self.val_max=im_target[idx[-1]]
+        return im_target.reshape(x.shape)
+
+
+    def from_gaussian(self,x):
+        
+        x=self.backend.bk_clip_by_value(x,self.val_min,self.val_max)
+        return self.f_gaussian(self.backend.to_numpy(x))
+    
     def square(self, x):
         if isinstance(x, scat_cov):
             if x.S1 is None:
@@ -3615,42 +3715,47 @@ class funct(FOC.FoCUS):
             return self.backend.bk_abs(self.backend.bk_sqrt(x))
 
     def reduce_mean(self, x):
-        
+
         if isinstance(x, scat_cov):
-            result = self.backend.bk_reduce_sum(self.backend.bk_abs(x.S0)) + \
-                     self.backend.bk_reduce_sum(self.backend.bk_abs(x.S2)) + \
-                     self.backend.bk_reduce_sum(self.backend.bk_abs(x.S3)) + \
-                     self.backend.bk_reduce_sum(self.backend.bk_abs(x.S4))
-                
-            N = self.backend.bk_size(x.S0)+self.backend.bk_size(x.S2)+ \
-                self.backend.bk_size(x.S3)+self.backend.bk_size(x.S4)
-                
+            result = (
+                self.backend.bk_reduce_sum(self.backend.bk_abs(x.S0))
+                + self.backend.bk_reduce_sum(self.backend.bk_abs(x.S2))
+                + self.backend.bk_reduce_sum(self.backend.bk_abs(x.S3))
+                + self.backend.bk_reduce_sum(self.backend.bk_abs(x.S4))
+            )
+
+            N = (
+                self.backend.bk_size(x.S0)
+                + self.backend.bk_size(x.S2)
+                + self.backend.bk_size(x.S3)
+                + self.backend.bk_size(x.S4)
+            )
+
             if x.S1 is not None:
-                result = result+self.backend.bk_reduce_sum(self.backend.bk_abs(x.S1))
+                result = result + self.backend.bk_reduce_sum(self.backend.bk_abs(x.S1))
                 N = N + self.backend.bk_size(x.S1)
             if x.S3P is not None:
-                result = result+self.backend.bk_reduce_sum(self.backend.bk_abs(x.S3P))
+                result = result + self.backend.bk_reduce_sum(self.backend.bk_abs(x.S3P))
                 N = N + self.backend.bk_size(x.S3P)
-            return result/self.backend.bk_cast(N)
+            return result / self.backend.bk_cast(N)
         else:
             return self.backend.bk_reduce_mean(x, axis=0)
-                
 
     def reduce_mean_batch(self, x):
-        
+
         if isinstance(x, scat_cov):
-            
-            sS0=self.backend.bk_reduce_mean(x.S0, axis=0)
-            sS2=self.backend.bk_reduce_mean(x.S2, axis=0)
-            sS3=self.backend.bk_reduce_mean(x.S3, axis=0)
-            sS4=self.backend.bk_reduce_mean(x.S4, axis=0)
-            sS1=None
-            sS3P=None
+
+            sS0 = self.backend.bk_reduce_mean(x.S0, axis=0)
+            sS2 = self.backend.bk_reduce_mean(x.S2, axis=0)
+            sS3 = self.backend.bk_reduce_mean(x.S3, axis=0)
+            sS4 = self.backend.bk_reduce_mean(x.S4, axis=0)
+            sS1 = None
+            sS3P = None
             if x.S1 is not None:
                 sS1 = self.backend.bk_reduce_mean(x.S1, axis=0)
             if x.S3P is not None:
                 sS3P = self.backend.bk_reduce_mean(x.S3P, axis=0)
-                
+
             result = scat_cov(
                 sS0,
                 sS2,
@@ -3664,22 +3769,22 @@ class funct(FOC.FoCUS):
             return result
         else:
             return self.backend.bk_reduce_mean(x, axis=0)
-    
+
     def reduce_sum_batch(self, x):
-        
+
         if isinstance(x, scat_cov):
-            
-            sS0=self.backend.bk_reduce_sum(x.S0, axis=0)
-            sS2=self.backend.bk_reduce_sum(x.S2, axis=0)
-            sS3=self.backend.bk_reduce_sum(x.S3, axis=0)
-            sS4=self.backend.bk_reduce_sum(x.S4, axis=0)
-            sS1=None
-            sS3P=None
+
+            sS0 = self.backend.bk_reduce_sum(x.S0, axis=0)
+            sS2 = self.backend.bk_reduce_sum(x.S2, axis=0)
+            sS3 = self.backend.bk_reduce_sum(x.S3, axis=0)
+            sS4 = self.backend.bk_reduce_sum(x.S4, axis=0)
+            sS1 = None
+            sS3P = None
             if x.S1 is not None:
                 sS1 = self.backend.bk_reduce_sum(x.S1, axis=0)
             if x.S3P is not None:
                 sS3P = self.backend.bk_reduce_sum(x.S3P, axis=0)
-                
+
             result = scat_cov(
                 sS0,
                 sS2,
@@ -3693,7 +3798,7 @@ class funct(FOC.FoCUS):
             return result
         else:
             return self.backend.bk_reduce_mean(x, axis=0)
-    
+
     def reduce_distance(self, x, y, sigma=None):
 
         if isinstance(x, scat_cov):
@@ -3729,11 +3834,13 @@ class funct(FOC.FoCUS):
             return result
         else:
             if sigma is None:
-                tmp=x-y
+                tmp = x - y
             else:
-                tmp=(x-y)/sigma
+                tmp = (x - y) / sigma
             # do abs in case of complex values
-            return self.backend.bk_abs(self.backend.bk_reduce_mean(self.backend.bk_square(tmp)))
+            return self.backend.bk_abs(
+                self.backend.bk_reduce_mean(self.backend.bk_square(tmp))
+            )
 
     def reduce_sum(self, x):
 
@@ -3909,3 +4016,122 @@ class funct(FOC.FoCUS):
         return scat_cov(
             s0, s2, s3, s4, s1=s1, s3p=s3p, backend=self.backend, use_1D=self.use_1D
         )
+
+    def synthesis(self,
+                  image_target,
+                  nstep=4,
+                  seed=1234,
+                  edge=True,
+                  to_gaussian=True,
+                  EVAL_FREQUENCY=100,
+                  NUM_EPOCHS = 300):
+                      
+        import foscat.Synthesis as synthe
+        import time
+
+        def The_loss(u,scat_operator,args):
+            ref  = args[0]
+            sref = args[1]
+            
+            # compute scattering covariance of the current synthetised map called u
+            learn=scat_operator.reduce_mean_batch(scat_operator.eval(u,edge=edge))
+            
+            # make the difference withe the reference coordinates
+            loss=scat_operator.reduce_distance(learn,ref,sigma=sref)
+    
+            return loss
+
+        if to_gaussian:
+            # Change the data histogram to gaussian distribution
+            im_target=self.to_gaussian(image_target)
+        else:
+            im_target=image_target
+            
+        axis=len(im_target.shape)-1
+        if self.use_2D:
+            axis-=1
+        if axis==0:
+            im_target=self.backend.bk_expand_dims(im_target,0)
+
+        # compute the number of possible steps
+        if self.use_2D:
+            jmax=int(np.min([np.log(im_target.shape[1]),np.log(im_target.shape[2])])/np.log(2))
+        elif self.use_1D:
+            jmax=int(np.log(im_target.shape[1])/np.log(2))
+        else:
+            jmax=int((np.log(im_target.shape[1]//12)/np.log(2))/2)
+            nside=2**jmax
+
+        if nstep>jmax-1:
+            nstep=jmax-1
+
+        t1=time.time()
+        tmp={}
+        tmp[nstep-1]=im_target
+        for l in range(nstep-2,-1,-1):
+            tmp[l]=self.ud_grade_2(tmp[l+1],axis=1)
+            
+        if not self.use_2D and not self.use_1D:
+            l_nside=nside//(2**(nstep-1))
+        
+        for k in range(nstep):
+            if k==0:
+                np.random.seed(seed)
+                if self.use_2D:
+                    imap=np.random.randn(tmp[k].shape[0],
+                                         tmp[k].shape[1],
+                                        tmp[k].shape[2])
+                else:
+                    imap=np.random.randn(tmp[k].shape[0],
+                                         tmp[k].shape[1])
+            else:
+                axis=1
+                # if the kernel size is bigger than 3 increase the binning before smoothing
+                if self.use_2D:
+                    imap = self.up_grade(
+                        omap, imap.shape[axis] * 2, axis=1, nouty=imap.shape[axis + 1] * 2
+                    )
+                elif self.use_1D:
+                    imap = self.up_grade(omap, imap.shape[axis] * 2, axis=1)
+                else:
+                    imap = self.up_grade(omap, l_nside, axis=axis)
+                    
+            # compute the coefficients for the target image
+            ref,sref=self.eval(tmp[k],calc_var=True,edge=edge)
+
+            # compute the mean of the population does nothing if only one map is given
+            ref=self.reduce_mean_batch(ref)
+            sref=self.reduce_mean_batch(sref)
+
+            # define a loss to minimize
+            loss=synthe.Loss(The_loss,self,ref,sref)
+        
+            sy = synthe.Synthesis([loss])
+
+            # initialize the synthesised map
+            if self.use_2D:
+                print('Synthesis scale [ %d x %d ]'%(imap.shape[1],imap.shape[2]))
+            elif self.use_1D:
+                print('Synthesis scale [ %d ]'%(imap.shape[1]))
+            else:
+                print('Synthesis scale nside=%d'%(l_nside))
+                l_nside*=2
+            
+            # do the minimization
+            omap=sy.run(imap,
+                        EVAL_FREQUENCY=EVAL_FREQUENCY,
+                        NUM_EPOCHS = NUM_EPOCHS)
+
+        t2=time.time()
+        print('Total computation %.2fs'%(t2-t1))
+
+        if to_gaussian:
+            omap=self.from_gaussian(omap)
+
+        if axis==0:
+            return omap[0]
+        else:
+            return omap
+            
+    def to_numpy(self,x):
+        return self.backend.to_numpy(x)