foscat 2025.5.2__py3-none-any.whl → 2025.6.3__py3-none-any.whl

foscat/scat_cov.py

@@ -92,12 +92,7 @@ class scat_cov:
         )
 
     def conv2complex(self, val):
-        if (
-            val.dtype == "complex64"
-            or val.dtype == "complex128"
-            or val.dtype == "torch.complex64"
-            or val.dtype == "torch.complex128"
-        ):
+        if self.backend.bk_is_complex(val):
             return val
         else:
             return self.backend.bk_complex(val, 0 * val)
@@ -107,7 +102,7 @@ class scat_cov:
     def flatten(self):
         tmp = [
             self.conv2complex(
-                self.backend.bk_reshape(self.S0, [self.S1.shape[0], self.S0.shape[1]])
+                self.backend.bk_reshape(self.S0, [self.S1.shape[0], self.S0.shape[1]*self.S0.shape[2]])
             )
         ]
         if self.use_1D:
@@ -177,21 +172,9 @@ class scat_cov:
                     ],
                 )
             ),
-            self.backend.bk_reshape(
-                self.S3,
-                [
-                    self.S3.shape[0],
-                    self.S3.shape[1]
-                    * self.S3.shape[2]
-                    * self.S3.shape[3]
-                    * self.S3.shape[4],
-                ],
-            ),
-        ]
-        if self.S3P is not None:
-            tmp = tmp + [
+            self.conv2complex(
                 self.backend.bk_reshape(
-                    self.S3P,
+                    self.S3,
                     [
                         self.S3.shape[0],
                         self.S3.shape[1]
@@ -200,22 +183,39 @@ class scat_cov:
                         * self.S3.shape[4],
                     ],
                 )
+            ),
+        ]
+        if self.S3P is not None:
+            tmp = tmp + [
+                self.conv2complex(
+                    self.backend.bk_reshape(
+                        self.S3P,
+                        [
+                            self.S3.shape[0],
+                            self.S3.shape[1]
+                            * self.S3.shape[2]
+                            * self.S3.shape[3]
+                            * self.S3.shape[4],
+                        ],
+                    )
+                )
             ]
 
         tmp = tmp + [
-            self.backend.bk_reshape(
-                self.S4,
-                [
-                    self.S3.shape[0],
-                    self.S4.shape[1]
-                    * self.S4.shape[2]
-                    * self.S4.shape[3]
-                    * self.S4.shape[4]
-                    * self.S4.shape[5],
-                ],
+            self.conv2complex(
+                self.backend.bk_reshape(
+                    self.S4,
+                    [
+                        self.S4.shape[0],
+                        self.S4.shape[1]
+                        * self.S4.shape[2]
+                        * self.S4.shape[3]
+                        * self.S4.shape[4]
+                        * self.S4.shape[5],
+                    ],
+                )
             )
         ]
-
         return self.backend.bk_concat(tmp, 1)
 
     # ---------------------------------------------−---------
@@ -2354,9 +2354,9 @@ class funct(FOC.FoCUS):
             tmpi2 = image2
         if upscale:
             l_nside = int(np.sqrt(tmp.shape[1] // 12))
-            tmp = self.up_grade(tmp, l_nside * 2, axis=1)
+            tmp = self.up_grade(tmp, l_nside * 2)
             if image2 is not None:
-                tmpi2 = self.up_grade(tmpi2, l_nside * 2, axis=1)
+                tmpi2 = self.up_grade(tmpi2, l_nside * 2)
         l_nside = int(np.sqrt(tmp.shape[1] // 12))
         nscale = int(np.log(l_nside) / np.log(2))
         cmat = {}
@@ -2367,12 +2367,12 @@ class funct(FOC.FoCUS):
             if image2 is not None:
                 sim = self.backend.bk_real(
                     self.backend.bk_L1(
-                        self.convol(tmp, axis=1)
-                        * self.backend.bk_conjugate(self.convol(tmpi2, axis=1))
+                        self.convol(tmp)
+                        * self.backend.bk_conjugate(self.convol(tmpi2))
                     )
                 )
             else:
-                sim = self.backend.bk_abs(self.convol(tmp, axis=1))
+                sim = self.backend.bk_abs(self.convol(tmp))
 
             # instead of difference between "opposite" channels use weighted average
             # of cosine and sine contributions using all channels
@@ -2442,12 +2442,12 @@ class funct(FOC.FoCUS):
 
             for k2 in range(k + 1):
 
-                tmp2 = self.backend.bk_repeat(sim, self.NORIENT, axis=1)
+                tmp2 = self.backend.bk_expand_dims(sim,-2)
 
                 sim2 = self.backend.bk_reduce_sum(
                     self.backend.bk_reshape(
                         self.backend.bk_cast(
-                            mat.reshape(1, self.NORIENT * self.NORIENT, mat.shape[1])
+                            mat.reshape(1, self.NORIENT, self.NORIENT, mat.shape[1])
                         )
                         * tmp2,
                         [sim.shape[0], self.NORIENT, self.NORIENT, mat.shape[1]],
@@ -2455,14 +2455,14 @@ class funct(FOC.FoCUS):
                     1,
                 )
 
-                sim2 = self.backend.bk_abs(self.convol(sim2, axis=-1))
+                sim2 = self.backend.bk_abs(self.convol(sim2))
 
                 angles = self.backend.bk_reshape(angles, [1, self.NORIENT, 1, 1])
                 weighted_cos2 = self.backend.bk_reduce_mean(
-                    sim2 * self.backend.bk_cos(angles), axis=1
+                    sim2 * self.backend.bk_cos(angles), axis=-3
                 )
                 weighted_sin2 = self.backend.bk_reduce_mean(
-                    sim2 * self.backend.bk_sin(angles), axis=1
+                    sim2 * self.backend.bk_sin(angles), axis=-3
                 )
 
                 cc2 = weighted_cos2[0]
@@ -2471,13 +2471,13 @@ class funct(FOC.FoCUS):
                 if smooth_scale > 0:
                     for m in range(smooth_scale):
                         if cc2.shape[1] > 12:
-                            cc2, _ = self.ud_grade_2(self.smooth(cc2, axis=1), axis=1)
-                            ss2, _ = self.ud_grade_2(self.smooth(ss2, axis=1), axis=1)
+                            cc2, _ = self.ud_grade_2(self.smooth(cc2))
+                            ss2, _ = self.ud_grade_2(self.smooth(ss2))
 
                 if cc2.shape[1] != sim.shape[2]:
                     ll_nside = int(np.sqrt(sim.shape[2] // 12))
-                    cc2 = self.up_grade(cc2, ll_nside, axis=1)
-                    ss2 = self.up_grade(ss2, ll_nside, axis=1)
+                    cc2 = self.up_grade(cc2, ll_nside)
+                    ss2 = self.up_grade(ss2, ll_nside)
 
                 if self.BACKEND == "numpy":
                     phase2 = np.fmod(np.arctan2(ss2, cc2) + 2 * np.pi, 2 * np.pi)
@@ -2509,9 +2509,9 @@ class funct(FOC.FoCUS):
                 )
 
             if k < l_nside - 1:
-                tmp, _ = self.ud_grade_2(tmp, axis=1)
+                tmp, _ = self.ud_grade_2(tmp)
                 if image2 is not None:
-                    tmpi2, _ = self.ud_grade_2(tmpi2, axis=1)
+                    tmpi2, _ = self.ud_grade_2(tmpi)
         return cmat, cmat2
 
     def div_norm(self, complex_value, float_value):
@@ -2534,6 +2534,7 @@ class funct(FOC.FoCUS):
         edge=True,
         nside=None,
         cell_ids=None,
+        spin=0
     ):
         """
         Calculates the scattering correlations for a batch of images. Mean are done over pixels.
@@ -2559,11 +2560,14 @@ 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.
+        spin : Integer
+            If different from 0 compute spinned data (U,V to Divergence/Rotational spin==1) or (Q,U to E,B spin=2).
+            This implies that the input data is 2*12*nside^2.
         Returns
         -------
         S1, S2, S3, S4 normalized
         """
-
+        
         return_data = self.return_data
 
         # Check input consistency
@@ -2576,8 +2580,8 @@ class funct(FOC.FoCUS):
         if mask is not None:
             if self.use_2D:
                 if (
-                    image1.shape[-2] != mask.shape[1]
-                    or image1.shape[-1] != mask.shape[2]
+                    image1.shape[-2] != mask.shape[-1]
+                    or image1.shape[-1] != mask.shape[-1]
                 ):
                     print(
                         "The LAST 2 COLUMNs of the mask should have the same size ",
@@ -2588,7 +2592,7 @@ class funct(FOC.FoCUS):
                     )
                     return None
             else:
-                if image1.shape[-1] != mask.shape[1]:
+                if image1.shape[-1] != mask.shape[-1]:
                     print(
                         "The LAST COLUMN of the mask should have the same size ",
                         mask.shape,
@@ -2609,7 +2613,6 @@ class funct(FOC.FoCUS):
             cross = True
         l_nside = 2**32  # not initialize if 1D or 2D
         ### PARAMETERS
-        axis = 1
         # determine jmax and nside corresponding to the input map
         im_shape = image1.shape
         if self.use_2D:
@@ -2637,10 +2640,7 @@ class funct(FOC.FoCUS):
 
             J = int(np.log(nside) / np.log(2))  # Number of j scales
         else:
-            if len(image1.shape) == 2:
-                npix = int(im_shape[1])  # Number of pixels
-            else:
-                npix = int(im_shape[0])  # Number of pixels
+            npix=int(im_shape[-1])
 
             if nside is None:
                 nside = int(np.sqrt(npix // 12))
@@ -2659,7 +2659,7 @@ class funct(FOC.FoCUS):
             print("\n\n==========")
 
         ### LOCAL VARIABLES (IMAGES and MASK)
-        if len(image1.shape) == 1 or (len(image1.shape) == 2 and self.use_2D):
+        if len(image1.shape) == 1 or (len(image1.shape) == 2 and self.use_2D) or (len(image1.shape) == 2 and spin>0):
             I1 = self.backend.bk_cast(
                 self.backend.bk_expand_dims(image1, 0)
             )  # Local image1 [Nbatch, Npix]
@@ -2684,26 +2684,26 @@ class funct(FOC.FoCUS):
             # if the kernel size is bigger than 3 increase the binning before smoothing
             if self.use_2D:
                 vmask = self.up_grade(
-                    vmask, I1.shape[axis] * 2, axis=1, nouty=I1.shape[axis + 1] * 2
+                    vmask, I1.shape[-2] * 2, nouty=I1.shape[-1] * 2
                 )
                 I1 = self.up_grade(
-                    I1, I1.shape[axis] * 2, axis=axis, nouty=I1.shape[axis + 1] * 2
+                    I1, I1.shape[-2] * 2, nouty=I1.shape[-1] * 2
                 )
                 if cross:
                     I2 = self.up_grade(
-                        I2, I2.shape[axis] * 2, axis=axis, nouty=I2.shape[axis + 1] * 2
+                        I2, I2.shape[-2] * 2, nouty=I2.shape[-1] * 2
                     )
             elif self.use_1D:
-                vmask = self.up_grade(vmask, I1.shape[axis] * 2, axis=1)
-                I1 = self.up_grade(I1, I1.shape[axis] * 2, axis=axis)
+                vmask = self.up_grade(vmask, I1.shape[-1] * 2)
+                I1 = self.up_grade(I1, I1.shape[-1] * 2)
                 if cross:
-                    I2 = self.up_grade(I2, I2.shape[axis] * 2, axis=axis)
+                    I2 = self.up_grade(I2, I2.shape[-1] * 2)
                 nside = nside * 2
             else:
-                I1 = self.up_grade(I1, nside * 2, axis=axis)
-                vmask = self.up_grade(vmask, nside * 2, axis=1)
+                I1 = self.up_grade(I1, nside * 2)
+                vmask = self.up_grade(vmask, nside * 2)
                 if cross:
-                    I2 = self.up_grade(I2, nside * 2, axis=axis)
+                    I2 = self.up_grade(I2, nside * 2)
 
                 nside = nside * 2
 
@@ -2711,29 +2711,28 @@ class funct(FOC.FoCUS):
                 # if the kernel size is bigger than 3 increase the binning before smoothing
                 if self.use_2D:
                     vmask = self.up_grade(
-                        vmask, I1.shape[axis] * 2, axis=1, nouty=I1.shape[axis + 1] * 2
+                        vmask, I1.shape[-2] * 2, nouty=I1.shape[-1] * 2
                     )
                     I1 = self.up_grade(
-                        I1, I1.shape[axis] * 2, axis=axis, nouty=I1.shape[axis + 1] * 2
+                        I1, I1.shape[-2] * 2, nouty=I1.shape[-1] * 2
                     )
                     if cross:
                         I2 = self.up_grade(
                             I2,
-                            I2.shape[axis] * 2,
-                            axis=axis,
-                            nouty=I2.shape[axis + 1] * 2,
+                            I2.shape[-2] * 2,
+                            nouty=I2.shape[-1] * 2,
                         )
                 elif self.use_1D:
-                    vmask = self.up_grade(vmask, I1.shape[axis] * 2, axis=1)
-                    I1 = self.up_grade(I1, I1.shape[axis] * 2, axis=axis)
+                    vmask = self.up_grade(vmask, I1.shape[-1] * 2)
+                    I1 = self.up_grade(I1, I1.shape[-1] * 2)
                     if cross:
-                        I2 = self.up_grade(I2, I2.shape[axis] * 2, axis=axis)
+                        I2 = self.up_grade(I2, I2.shape[-1] * 2)
                     nside = nside * 2
                 else:
-                    I1 = self.up_grade(I1, nside * 2, axis=axis)
-                    vmask = self.up_grade(vmask, nside * 2, axis=1)
+                    I1 = self.up_grade(I1, nside * 2)
+                    vmask = self.up_grade(vmask, nside * 2)
                     if cross:
-                        I2 = self.up_grade(I2, nside * 2, axis=axis)
+                        I2 = self.up_grade(I2, nside * 2)
                     nside = nside * 2
 
         # Normalize the masks because they have different pixel numbers
@@ -2762,6 +2761,7 @@ class funct(FOC.FoCUS):
         off_S2 = -2
         off_S3 = -3
         off_S4 = -4
+            
         if self.use_1D:
             off_S2 = -1
             off_S3 = -1
@@ -2793,13 +2793,20 @@ class funct(FOC.FoCUS):
                 )
         else:
             if not cross:
-                s0, l_vs0 = self.masked_mean(I1, vmask, axis=1, calc_var=True)
+                s0, l_vs0 = self.masked_mean(I1,
+                                             vmask,
+                                             calc_var=True)
             else:
                 s0, l_vs0 = self.masked_mean(
-                    self.backend.bk_L1(I1 * I2), vmask, axis=1, calc_var=True
-                )
-            vs0 = self.backend.bk_concat([l_vs0, l_vs0], 1)
-            s0 = self.backend.bk_concat([s0, l_vs0], 1)
+                    self.backend.bk_L1(I1 * I2),
+                    vmask,
+                    calc_var=True)
+                
+            vs0 = self.backend.bk_concat([l_vs0, l_vs0], -1)
+            s0 = self.backend.bk_concat([s0, l_vs0], -1)
+            if spin>0:
+                vs0=self.backend.bk_reshape(vs0,[vs0.shape[0],vs0.shape[1],2,vs0.shape[2]//2])
+                s0=self.backend.bk_reshape(s0,[s0.shape[0],s0.shape[1],2,s0.shape[2]//2])
         #### COMPUTE S1, S2, S3 and S4
         nside_j3 = nside  # NSIDE start (nside_j3 = nside / 2^j3)
 
@@ -2848,13 +2855,13 @@ class funct(FOC.FoCUS):
             ####### S1 and S2
             ### Make the convolution I1 * Psi_j3
             conv1 = self.convol(
-                I1, axis=1, cell_ids=cell_ids_j3, nside=nside_j3
+                I1, cell_ids=cell_ids_j3, nside=nside_j3,
+                spin=spin
             )  # [Nbatch, Norient3 , Npix_j3]
 
             if cmat is not None:
-
                 tmp2 = self.backend.bk_repeat(conv1, self.NORIENT, axis=-2)
-
+                
                 conv1 = self.backend.bk_reduce_sum(
                     self.backend.bk_reshape(
                         cmat[j3] * tmp2,
@@ -2871,7 +2878,6 @@ class funct(FOC.FoCUS):
             M1 = self.backend.bk_L1(M1_square)  # [Nbatch, Npix_j3, Norient3]
             # Store M1_j3 in a dictionary
             M1_dic[j3] = M1
-
             if not cross:  # Auto
                 M1_square = self.backend.bk_real(M1_square)
 
@@ -2882,11 +2888,11 @@ class funct(FOC.FoCUS):
                 else:
                     if calc_var:
                         s2, vs2 = self.masked_mean(
-                            M1_square, vmask, axis=2, rank=j3, calc_var=True
+                            M1_square, vmask, rank=j3, calc_var=True
                         )
                     else:
-                        s2 = self.masked_mean(M1_square, vmask, axis=2, rank=j3)
-
+                        s2 = self.masked_mean(M1_square, vmask, rank=j3)
+                        
                 if cond_init_P1_dic:
                     # We fill P1_dic with S2 for normalisation of S3 and S4
                     P1_dic[j3] = self.backend.bk_real(s2)  # [Nbatch, Nmask, Norient3]
@@ -2929,11 +2935,11 @@ class funct(FOC.FoCUS):
                 else:
                     if calc_var:
                         s1, vs1 = self.masked_mean(
-                            M1, vmask, axis=2, rank=j3, calc_var=True
+                            M1, vmask, rank=j3, calc_var=True
                         )  # [Nbatch, Nmask, Norient3]
                     else:
                         s1 = self.masked_mean(
-                            M1, vmask, axis=2, rank=j3
+                            M1, vmask, rank=j3
                         )  # [Nbatch, Nmask, Norient3]
 
                 if return_data:
@@ -2967,7 +2973,8 @@ class funct(FOC.FoCUS):
             else:  # Cross
                 ### Make the convolution I2 * Psi_j3
                 conv2 = self.convol(
-                    I2, axis=2, cell_ids=cell_ids_j3, nside=nside_j3
+                    I2,  cell_ids=cell_ids_j3, nside=nside_j3,
+                    spin=spin
                 )  # [Nbatch, Npix_j3, Norient3]
                 if cmat is not None:
                     tmp2 = self.backend.bk_repeat(conv2, self.NORIENT, axis=-2)
@@ -3001,17 +3008,17 @@ class funct(FOC.FoCUS):
                     else:
                         if calc_var:
                             p1, vp1 = self.masked_mean(
-                                M1_square, vmask, axis=2, rank=j3, calc_var=True
+                                M1_square, vmask,  rank=j3, calc_var=True
                             )  # [Nbatch, Nmask, Norient3]
                             p2, vp2 = self.masked_mean(
-                                M2_square, vmask, axis=2, rank=j3, calc_var=True
+                                M2_square, vmask,  rank=j3, calc_var=True
                             )  # [Nbatch, Nmask, Norient3]
                         else:
                             p1 = self.masked_mean(
-                                M1_square, vmask, axis=2, rank=j3
+                                M1_square, vmask, rank=j3
                             )  # [Nbatch, Nmask, Norient3]
                             p2 = self.masked_mean(
-                                M2_square, vmask, axis=2, rank=j3
+                                M2_square, vmask, rank=j3
                             )  # [Nbatch, Nmask, Norient3]
                     # We fill P1_dic with S2 for normalisation of S3 and S4
                     P1_dic[j3] = self.backend.bk_real(p1)  # [Nbatch, Nmask, Norient3]
@@ -3027,10 +3034,10 @@ class funct(FOC.FoCUS):
                 else:
                     if calc_var:
                         s2, vs2 = self.masked_mean(
-                            s2, vmask, axis=2, rank=j3, calc_var=True
+                            s2, vmask, rank=j3, calc_var=True
                         )
                     else:
-                        s2 = self.masked_mean(s2, vmask, axis=2, rank=j3)
+                        s2 = self.masked_mean(s2, vmask, rank=j3)
 
                 if return_data:
                     if out_nside is not None and out_nside < nside_j3:
@@ -3072,11 +3079,11 @@ class funct(FOC.FoCUS):
                 else:
                     if calc_var:
                         s1, vs1 = self.masked_mean(
-                            MX, vmask, axis=2, rank=j3, calc_var=True
+                            MX, vmask, rank=j3, calc_var=True
                         )  # [Nbatch, Nmask, Norient3]
                     else:
                         s1 = self.masked_mean(
-                            MX, vmask, axis=2, rank=j3
+                            MX, vmask, rank=j3
                         )  # [Nbatch, Nmask, Norient3]
                 if return_data:
                     if out_nside is not None and out_nside < nside_j3:
@@ -3482,39 +3489,39 @@ class funct(FOC.FoCUS):
             ### Image I1,
             # downscale the I1 [Nbatch, Npix_j3]
             if j3 != Jmax - 1:
-                I1 = self.smooth(I1, axis=1, cell_ids=cell_ids_j3, nside=nside_j3)
+                I1 = self.smooth(I1, cell_ids=cell_ids_j3, nside=nside_j3)
                 I1, new_cell_ids_j3 = self.ud_grade_2(
-                    I1, axis=1, cell_ids=cell_ids_j3, nside=nside_j3
+                    I1, cell_ids=cell_ids_j3, nside=nside_j3
                 )
 
                 ### Image I2
                 if cross:
-                    I2 = self.smooth(I2, axis=1, cell_ids=cell_ids_j3, nside=nside_j3)
+                    I2 = self.smooth(I2,  cell_ids=cell_ids_j3, nside=nside_j3)
                     I2, new_cell_ids_j3 = self.ud_grade_2(
-                        I2, axis=1, cell_ids=cell_ids_j3, nside=nside_j3
+                        I2, cell_ids=cell_ids_j3, nside=nside_j3
                     )
 
                 ### Modules
                 for j2 in range(0, j3 + 1):  # j2 =< j3
                     ### Dictionary M1_dic[j2]
                     M1_smooth = self.smooth(
-                        M1_dic[j2], axis=2, cell_ids=cell_ids_j3, nside=nside_j3
+                        M1_dic[j2], cell_ids=cell_ids_j3, nside=nside_j3
                     )  # [Nbatch, Npix_j3, Norient3]
                     M1_dic[j2], new_cell_ids_j2 = self.ud_grade_2(
-                        M1_smooth, axis=2, cell_ids=cell_ids_j3, nside=nside_j3
+                        M1_smooth,  cell_ids=cell_ids_j3, nside=nside_j3
                     )  # [Nbatch, Npix_j3, Norient3]
 
                     ### Dictionary M2_dic[j2]
                     if cross:
                         M2_smooth = self.smooth(
-                            M2_dic[j2], axis=2, cell_ids=cell_ids_j3, nside=nside_j3
+                            M2_dic[j2],  cell_ids=cell_ids_j3, nside=nside_j3
                         )  # [Nbatch, Npix_j3, Norient3]
                         M2_dic[j2], new_cell_ids_j2 = self.ud_grade_2(
-                            M2_smooth, axis=2, cell_ids=cell_ids_j3, nside=nside_j3
+                            M2_smooth,  cell_ids=cell_ids_j3, nside=nside_j3
                         )  # [Nbatch, Npix_j3, Norient3]
                 ### Mask
                 vmask, new_cell_ids_j3 = self.ud_grade_2(
-                    vmask, axis=1, cell_ids=cell_ids_j3, nside=nside_j3
+                    vmask, cell_ids=cell_ids_j3, nside=nside_j3
                 )
 
                 if self.mask_thres is not None:
@@ -3529,33 +3536,21 @@ 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)
+            S1 = self.backend.bk_concat(S1, -2)
+            S2 = self.backend.bk_concat(S2, -2)
+            S3 = self.backend.bk_concat(S3, -3)
+            S4 = self.backend.bk_concat(S4, -4)
             if cross:
-                S3P = self.backend.bk_concat(S3P, 2)
+                S3P = self.backend.bk_concat(S3P, -3)
             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)
+                VS1 = self.backend.bk_concat(VS1, -2)
+                VS2 = self.backend.bk_concat(VS2, -2)
+                VS3 = self.backend.bk_concat(VS3, -3)
+                VS4 = self.backend.bk_concat(VS4, -4)
                 if cross:
-                    VS3P = self.backend.bk_concat(VS3P, 2)
+                    VS3P = self.backend.bk_concat(VS3P, -3)
         if calc_var:
             if not cross:
                 return scat_cov(
@@ -3637,7 +3632,7 @@ class funct(FOC.FoCUS):
         ### Compute |I1 * Psi_j2| * Psi_j3 = M1_j2 * Psi_j3
         # Warning: M1_dic[j2] is already at j3 resolution [Nbatch, Norient3, Npix_j3]
         MconvPsi = self.convol(
-            M_dic[j2], axis=2, cell_ids=cell_ids, nside=nside_j2
+            M_dic[j2], cell_ids=cell_ids, nside=nside_j2
         )  # [Nbatch,   Norient3, Norient2, Npix_j3]
 
         if cmat2 is not None:
@@ -3674,12 +3669,12 @@ class funct(FOC.FoCUS):
         else:
             if calc_var:
                 s3, vs3 = self.masked_mean(
-                    s3, vmask, axis=3, rank=j2, calc_var=True
+                    s3, vmask, rank=j2, calc_var=True
                 )  # [Nbatch, Nmask, Norient3, Norient2]
                 return s3, vs3
             else:
                 s3 = self.masked_mean(
-                    s3, vmask, axis=3, rank=j2
+                    s3, vmask, rank=j2
                 )  # [Nbatch, Nmask, Norient3, Norient2]
             return s3
 
@@ -3717,12 +3712,12 @@ class funct(FOC.FoCUS):
         else:
             if calc_var:
                 s4, vs4 = self.masked_mean(
-                    s4, vmask, axis=4, rank=j2, calc_var=True
+                    s4, vmask, rank=j2, calc_var=True
                 )  # [Nbatch, Nmask, Norient3, Norient2, Norient1]
                 return s4, vs4
             else:
                 s4 = self.masked_mean(
-                    s4, vmask, axis=4, rank=j2
+                    s4, vmask, rank=j2
                 )  # [Nbatch, Nmask, Norient3, Norient2, Norient1]
                 return s4
 
@@ -3922,20 +3917,20 @@ class funct(FOC.FoCUS):
     #
     # ---------------------------------------------------------------------------
     def scattering_cov(
-        self,
-        data,
-        data2=None,
-        Jmax=None,
-        if_large_batch=False,
-        S4_criteria=None,
-        use_ref=False,
-        normalization="S2",
-        edge=False,
-        in_mask=None,
-        pseudo_coef=1,
-        get_variance=False,
-        ref_sigma=None,
-        iso_ang=False,
+            self,
+            data,
+            data2=None,
+            Jmax=None,
+            if_large_batch=False,
+            S4_criteria=None,
+            use_ref=False,
+            normalization="S2",
+            edge=False,
+            in_mask=None,
+            pseudo_coef=1,
+            get_variance=False,
+            ref_sigma=None,
+            iso_ang=False,
     ):
         """
         Calculates the scattering correlations for a batch of images, including:
@@ -4048,7 +4043,10 @@ class funct(FOC.FoCUS):
                 if data2 is not None:
                     N_image2 = data2.shape[0]
 
-            nside = int(np.sqrt(npix // 12))
+            if spin==0:
+                nside = int(np.sqrt(npix // 12))
+            else:
+                nside = int(np.sqrt(npix // 24))
 
             J = int(np.log(nside) / np.log(2))  # Number of j scales
 
@@ -5781,7 +5779,9 @@ class funct(FOC.FoCUS):
 
     def from_gaussian(self, x):
 
-        x = self.backend.bk_clip_by_value(x, self.val_min, self.val_max)
+        x = self.backend.bk_clip_by_value(x,
+                                          self.val_min+1E-4*abs(self.val_min),
+                                          self.val_max-1E-4*abs(self.val_max))
         return self.f_gaussian(self.backend.to_numpy(x))
 
     def square(self, x):
@@ -6079,7 +6079,38 @@ class funct(FOC.FoCUS):
         return scat_cov(
             s0, s2, s3, s4, s1=s1, s3p=s3p, backend=self.backend, use_1D=self.use_1D
         )
-
+    def calc_matrix_orientation(self,noise_map,image2=None):
+        # Décalage circulaire par matrice de permutation
+        def circ_shift_matrix(N,k):
+            return np.roll(np.eye(N), shift=-k, axis=1)
+        Norient = self.NORIENT
+        im=self.convol(noise_map)
+        if image2 is None:
+            mm=np.mean(abs(self.backend.to_numpy(im)),0)
+        else:
+            im2=self.convol(self.backend.bk_cast(image2))
+            mm=np.mean(self.backend.to_numpy(
+                self.backend.bk_L1(im*self.backend.bk_conjugate(im2))).real,0)
+        
+        Norient=mm.shape[0]
+        xx=np.cos(np.arange(Norient)/Norient*2*np.pi)
+        yy=np.sin(np.arange(Norient)/Norient*2*np.pi)
+
+        a=np.sum(mm*xx[:,None],0)
+        b=np.sum(mm*yy[:,None],0)
+        
+        o=np.fmod(Norient*np.arctan2(-b,a)/(2*np.pi)+Norient,Norient)
+        xx=np.arange(Norient)
+        alpha = o[None,:]-xx[:,None]
+        beta = np.fmod(1+o[None,:]-xx[:,None],Norient)
+        alpha=(1-alpha)*(alpha<1)*(alpha>0)+beta*(beta<1)*(beta>0)
+        
+        m=np.zeros([Norient,Norient,mm.shape[1]])
+        for k in range(Norient):
+            m[k,:,:]=np.roll(alpha,k,0)
+        #m=np.mean(m,0)
+        return self.backend.bk_cast(m)
+        
     def synthesis(
         self,
         image_target,