foscat 2025.5.0__py3-none-any.whl → 2025.6.1__py3-none-any.whl

foscat/FoCUS.py

@@ -5,7 +5,7 @@ import healpy as hp
 import numpy as np
 from scipy.interpolate import griddata
 
-TMPFILE_VERSION = "V4_0"
+TMPFILE_VERSION = "V5_0"
 
 
 class FoCUS:
@@ -35,7 +35,7 @@ class FoCUS:
         mpi_rank=0,
     ):
 
-        self.__version__ = "2025.05.0"
+        self.__version__ = "2025.06.1"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -176,11 +176,14 @@ class FoCUS:
         self.Y_CNN = {}
         self.Z_CNN = {}
 
+        self.Idx_CNN = {}
+        self.Idx_WCNN = {}
+        
         self.filters_set = {}
         self.edge_masks = {}
 
-        wwc = np.zeros([KERNELSZ**2, l_NORIENT]).astype(all_type)
-        wws = np.zeros([KERNELSZ**2, l_NORIENT]).astype(all_type)
+        wwc = np.zeros([l_NORIENT, KERNELSZ**2]).astype(all_type)
+        wws = np.zeros([l_NORIENT, KERNELSZ**2]).astype(all_type)
 
         x = np.repeat(np.arange(KERNELSZ) - KERNELSZ // 2, KERNELSZ).reshape(
             KERNELSZ, KERNELSZ
@@ -203,12 +206,12 @@ class FoCUS:
                     -0.5 * (xx**2 + yy**2)
                 )
 
-            wwc[:, 0] = tmp.flatten() - tmp.mean()
+            wwc[0] = tmp.flatten() - tmp.mean()
             tmp = 0 * w_smooth
-            wws[:, 0] = tmp.flatten()
+            wws[0] = tmp.flatten()
             sigma = np.sqrt((wwc[:, 0] ** 2).mean())
-            wwc[:, 0] /= sigma
-            wws[:, 0] /= sigma
+            wwc[0] /= sigma
+            wws[0] /= sigma
 
             w_smooth = w_smooth.flatten()
         else:
@@ -239,12 +242,12 @@ class FoCUS:
                 tmp1 = np.cos(yy * np.pi) * w_smooth
                 tmp2 = np.sin(yy * np.pi) * w_smooth
 
-                wwc[:, i] = tmp1.flatten() - tmp1.mean()
-                wws[:, i] = tmp2.flatten() - tmp2.mean()
+                wwc[i] = tmp1.flatten() - tmp1.mean()
+                wws[i] = tmp2.flatten() - tmp2.mean()
                 # sigma = np.sqrt((wwc[:, i] ** 2).mean())
                 sigma = np.mean(w_smooth)
-                wwc[:, i] /= sigma
-                wws[:, i] /= sigma
+                wwc[i] /= sigma
+                wws[i] /= sigma
 
                 if DODIV and i == 0:
                     r = xx**2 + yy**2
@@ -253,22 +256,22 @@ class FoCUS:
                     tmp1 = r * np.cos(2 * theta) * w_smooth
                     tmp2 = r * np.sin(2 * theta) * w_smooth
 
-                    wwc[:, NORIENT] = tmp1.flatten() - tmp1.mean()
-                    wws[:, NORIENT] = tmp2.flatten() - tmp2.mean()
+                    wwc[NORIENT] = tmp1.flatten() - tmp1.mean()
+                    wws[NORIENT] = tmp2.flatten() - tmp2.mean()
                     # sigma = np.sqrt((wwc[:, NORIENT] ** 2).mean())
                     sigma = np.mean(w_smooth)
 
-                    wwc[:, NORIENT] /= sigma
-                    wws[:, NORIENT] /= sigma
+                    wwc[NORIENT] /= sigma
+                    wws[NORIENT] /= sigma
                     tmp1 = r * np.cos(2 * theta + np.pi)
                     tmp2 = r * np.sin(2 * theta + np.pi)
 
-                    wwc[:, NORIENT + 1] = tmp1.flatten() - tmp1.mean()
-                    wws[:, NORIENT + 1] = tmp2.flatten() - tmp2.mean()
+                    wwc[NORIENT + 1] = tmp1.flatten() - tmp1.mean()
+                    wws[NORIENT + 1] = tmp2.flatten() - tmp2.mean()
                     # sigma = np.sqrt((wwc[:, NORIENT + 1] ** 2).mean())
                     sigma = np.mean(w_smooth)
-                    wwc[:, NORIENT + 1] /= sigma
-                    wws[:, NORIENT + 1] /= sigma
+                    wwc[NORIENT + 1] /= sigma
+                    wws[NORIENT + 1] /= sigma
 
                 w_smooth = w_smooth.flatten()
 
@@ -316,14 +319,14 @@ class FoCUS:
                 r = np.sum(np.sqrt(c * c + s * s))
                 c = c / r
                 s = s / r
-                self.ww_RealT[1] = self.backend.bk_constant(
-                    np.array(c).reshape(xx.shape[0], 1, 1)
+                self.ww_RealT[1] = self.backend.bk_cast(
+                    self.backend.bk_constant(np.array(c).reshape(xx.shape[0]))
                 )
-                self.ww_ImagT[1] = self.backend.bk_constant(
-                    np.array(s).reshape(xx.shape[0], 1, 1)
+                self.ww_ImagT[1] = self.backend.bk_cast(
+                    self.backend.bk_constant(np.array(s).reshape(xx.shape[0]))
                 )
-                self.ww_SmoothT[1] = self.backend.bk_constant(
-                    np.array(w).reshape(xx.shape[0], 1, 1)
+                self.ww_SmoothT[1] = self.backend.bk_cast(
+                    self.backend.bk_constant(np.array(w).reshape(xx.shape[0]))
                 )
 
         else:
@@ -333,22 +336,16 @@ class FoCUS:
             self.ww_SmoothT = {}
 
             self.ww_SmoothT[1] = self.backend.bk_constant(
-                self.w_smooth.reshape(KERNELSZ, KERNELSZ, 1, 1)
-            )
-            www = np.zeros([KERNELSZ, KERNELSZ, NORIENT, NORIENT], dtype=self.all_type)
-            for k in range(NORIENT):
-                www[:, :, k, k] = self.w_smooth.reshape(KERNELSZ, KERNELSZ)
-            self.ww_SmoothT[NORIENT] = self.backend.bk_constant(
-                www.reshape(KERNELSZ, KERNELSZ, NORIENT, NORIENT)
+                self.w_smooth.reshape(1, KERNELSZ, KERNELSZ)
             )
             self.ww_RealT[1] = self.backend.bk_constant(
                 self.backend.bk_reshape(
-                    wwc.astype(self.all_type), [KERNELSZ, KERNELSZ, 1, NORIENT]
+                    wwc.astype(self.all_type), [NORIENT, KERNELSZ, KERNELSZ]
                 )
             )
             self.ww_ImagT[1] = self.backend.bk_constant(
                 self.backend.bk_reshape(
-                    wws.astype(self.all_type), [KERNELSZ, KERNELSZ, 1, NORIENT]
+                    wws.astype(self.all_type), [NORIENT, KERNELSZ, KERNELSZ]
                 )
             )
 
@@ -505,211 +502,27 @@ class FoCUS:
             )
         return indices, weights, xc, yc, zc
 
-    # ---------------------------------------------−---------
-    def calc_orientation(self, im):  # im is [Ndata,12*Nside**2]
-        nside = int(np.sqrt(im.shape[1] // 12))
-        l_kernel = self.KERNELSZ * self.KERNELSZ
-        norient = 32
-        w = np.zeros([l_kernel, 1, 2 * norient])
-        ca = np.cos(np.arange(norient) / norient * np.pi)
-        sa = np.sin(np.arange(norient) / norient * np.pi)
-        stat = np.zeros([12 * nside**2, norient])
-
-        if self.ww_CNN[nside] is None:
-            self.init_CNN_index(nside, transpose=False)
-
-        y = self.Y_CNN[nside]
-        z = self.Z_CNN[nside]
-
-        for k in range(norient):
-            w[:, 0, k] = np.exp(-0.5 * nside**2 * ((y) ** 2 + (z) ** 2)) * np.cos(
-                nside * (y * ca[k] + z * sa[k]) * np.pi / 2
-            )
-            w[:, 0, k + norient] = np.exp(
-                -0.5 * nside**2 * ((y) ** 2 + (z) ** 2)
-            ) * np.sin(nside * (y * ca[k] + z * sa[k]) * np.pi / 2)
-            w[:, 0, k] = w[:, 0, k] - np.mean(w[:, 0, k])
-            w[:, 0, k + norient] = w[:, 0, k] - np.mean(w[:, 0, k + norient])
-
-        for k in range(im.shape[0]):
-            tmp = im[k].reshape(12 * nside**2, 1)
-            im2 = self.healpix_layer(tmp, w)
-            stat = stat + im2[:, 0:norient] ** 2 + im2[:, norient:] ** 2
-
-        rotation = (np.argmax(stat, 1)).astype("float") / 32.0 * 180.0
-
-        indices, weights, x, y, z = self.calc_indices_convol(
-            nside, 9, rotation=rotation
-        )
-
-        return indices, weights
-
-    def init_CNN_index(self, nside, transpose=False):
-        l_kernel = int(self.KERNELSZ * self.KERNELSZ)
-        try:
-            indices = np.load(
-                "%s/FOSCAT_%s_I%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside)
-            )
-            weights = np.load(
-                "%s/FOSCAT_%s_W%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside)
-            )
-            xc = np.load(
-                "%s/FOSCAT_%s_X%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside)
-            )
-            yc = np.load(
-                "%s/FOSCAT_%s_Y%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside)
-            )
-            zc = np.load(
-                "%s/FOSCAT_%s_Z%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside)
-            )
-        except:
-            indices, weights, xc, yc, zc = self.calc_indices_convol(nside, l_kernel)
-            np.save(
-                "%s/FOSCAT_%s_I%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside),
-                indices,
-            )
-            np.save(
-                "%s/FOSCAT_%s_W%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside),
-                weights,
-            )
-            np.save(
-                "%s/FOSCAT_%s_X%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside),
-                xc,
-            )
-            np.save(
-                "%s/FOSCAT_%s_Y%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside),
-                yc,
-            )
-            np.save(
-                "%s/FOSCAT_%s_Z%d_%d_%d_CNNV3.npy"
-                % (self.TEMPLATE_PATH, TMPFILE_VERSION, l_kernel, self.NORIENT, nside),
-                zc,
-            )
-            if not self.silent:
-                print(
-                    "Write %s/FOSCAT_%s_W%d_%d_%d_CNNV2.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        l_kernel,
-                        self.NORIENT,
-                        nside,
-                    )
-                )
-
-        self.X_CNN[nside] = xc
-        self.Y_CNN[nside] = yc
-        self.Z_CNN[nside] = zc
-        self.ww_CNN[nside] = self.backend.bk_SparseTensor(
-            indices, weights, [12 * nside * nside * l_kernel, 12 * nside * nside]
-        )
-
-    # ---------------------------------------------−---------
-    def healpix_layer_coord(self, im, axis=0):
-        nside = int(np.sqrt(im.shape[axis] // 12))
-        if self.ww_CNN[nside] is None:
-            self.init_CNN_index(nside)
-        return self.X_CNN[nside], self.Y_CNN[nside], self.Z_CNN[nside]
-
-    # ---------------------------------------------−---------
-    def healpix_layer_transpose(self, im, ww, indices=None, weights=None, axis=0):
-        nside = int(np.sqrt(im.shape[axis] // 12))
-
-        if im.shape[1 + axis] != ww.shape[1]:
-            if not self.silent:
-                print("Weights channels should be equal to the input image channels")
-            return -1
-        if axis == 1:
-            results = []
-
-            for k in range(im.shape[0]):
-
-                tmp = self.healpix_layer(
-                    im[k], ww, indices=indices, weights=weights, axis=0
-                )
-                tmp = self.backend.bk_reshape(
-                    self.up_grade(tmp, 2 * nside), [12 * 4 * nside**2, ww.shape[2]]
-                )
-
-                results.append(tmp)
-
-            return self.backend.bk_stack(results, axis=0)
-        else:
-            tmp = self.healpix_layer(
-                im, ww, indices=indices, weights=weights, axis=axis
-            )
-
-            return self.up_grade(tmp, 2 * nside)
-
     # ---------------------------------------------−---------
     # ---------------------------------------------−---------
-    def healpix_layer(self, im, ww, indices=None, weights=None, axis=0):
-        nside = int(np.sqrt(im.shape[axis] // 12))
-        l_kernel = self.KERNELSZ * self.KERNELSZ
-
-        if im.shape[1 + axis] != ww.shape[1]:
-            if not self.silent:
-                print("Weights channels should be equal to the input image channels")
-            return -1
-
+    def healpix_layer(self, im, ww, indices=None, weights=None):
+        #ww [N_i,NORIENT,KERNELSZ*KERNELSZ//2,N_o,NORIENT]
+        #im [N_batch,N_i,  NORIENT,N]
+        nside=int(np.sqrt(im.shape[-1]//12))
         if indices is None:
-            if self.ww_CNN[nside] is None:
-                self.init_CNN_index(nside, transpose=False)
-            mat = self.ww_CNN[nside]
-        else:
-            if weights is None:
-                print(
-                    "healpix_layer : If indices is not none weights should be specify"
-                )
-                return 0
-
-            mat = self.backend.bk_SparseTensor(
-                indices, weights, [12 * nside * nside * l_kernel, 12 * nside * nside]
-            )
-
-        if axis == 1:
-            results = []
-
-            for k in range(im.shape[0]):
-
-                tmp = self.backend.bk_sparse_dense_matmul(mat, im[k])
-
-                density = self.backend.bk_reshape(
-                    tmp, [12 * nside * nside, l_kernel * im.shape[1 + axis]]
-                )
-
-                density = self.backend.bk_matmul(
-                    density,
-                    self.backend.bk_reshape(
-                        ww, [l_kernel * im.shape[1 + axis], ww.shape[2]]
-                    ),
-                )
-
-                results.append(
-                    self.backend.bk_reshape(density, [12 * nside**2, ww.shape[2]])
-                )
-
-            return self.backend.bk_stack(results, axis=0)
-        else:
-            tmp = self.backend.bk_sparse_dense_matmul(mat, im)
-
-            density = self.backend.bk_reshape(
-                tmp, [12 * nside * nside, l_kernel * im.shape[1]]
-            )
-
-            return self.backend.bk_matmul(
-                density,
-                self.backend.bk_reshape(ww, [l_kernel * im.shape[1], ww.shape[2]]),
-            )
+            if (nside,self.NORIENT,self.KERNELSZ) not in self.ww_CNN:
+                self.init_index_cnn(nside,self.NORIENT)
+            indices = self.Idx_CNN[(nside,self.NORIENT,self.KERNELSZ)]
+            mat = self.Idx_WCNN[(nside,self.NORIENT,self.KERNELSZ)]
+            
+        wim = self.backend.bk_gather(im,indices.flatten(),axis=3) #[N_batch,N_i,NORIENT,K*(K+1),N_o,NORIENT,N,N_w]
+        
+        wim = self.backend.bk_reshape(wim,[im.shape[0],im.shape[1],im.shape[2]]+list(indices.shape))*mat[None,...]
+        #win is [N_batch,N_i,  NORIENT,K*(K+1),1,  NORIENT,N,N_w]
+        #ww is  [1,      N_i,  NORIENT,K*(K+1),N_o,NORIENT]
+        wim = self.backend.bk_reduce_sum(wim[:,:,:,:,None]*ww[None,:,:,:,:,:,None,None],[1,2,3])
+
+        wim = self.backend.bk_reduce_sum(wim,-1)
+        return self.backend.bk_reshape(wim,[im.shape[0],ww.shape[3],ww.shape[4],im.shape[-1]])
 
     # ---------------------------------------------−---------
 
@@ -806,53 +619,20 @@ class FoCUS:
             return self.backend.bk_reshape(res, [npix // 2, npiy // 2]), None
         elif self.use_1D:
             ishape = list(im.shape)
-            if len(ishape) < axis + 1:
-                if not self.silent:
-                    print("Use of 1D scat with data that has less than 1D")
-                return None, None
 
-            npix = im.shape[axis]
-            odata = 1
-            if len(ishape) > axis + 1:
-                for k in range(axis + 1, len(ishape)):
-                    odata = odata * ishape[k]
+            npix = ishape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(ishape) - 1):
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
-                self.backend.bk_cast(im), [ndata, npix, odata]
-            )
-            tim = self.backend.bk_reshape(
-                tim[:, 0 : 2 * (npix // 2), :], [ndata, npix // 2, 2, odata]
+                self.backend.bk_cast(im), [ndata, npix // 2, 2]
             )
 
-            res = self.backend.bk_reduce_mean(tim, 2)
-
-            if axis == 0:
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [npix // 2]), None
-                else:
-                    return (
-                        self.backend.bk_reshape(res, [npix // 2] + ishape[axis + 1 :]),
-                        None,
-                    )
-            else:
-                if len(ishape) == axis + 1:
-                    return (
-                        self.backend.bk_reshape(res, ishape[0:axis] + [npix // 2]),
-                        None,
-                    )
-                else:
-                    return (
-                        self.backend.bk_reshape(
-                            res, ishape[0:axis] + [npix // 2] + ishape[axis + 1 :]
-                        ),
-                        None,
-                    )
+            res = self.backend.bk_reduce_mean(tim, -1)
 
-            return self.backend.bk_reshape(res, [npix // 2]), None
+            return self.backend.bk_reshape(res, ishape[0:-1] + [npix // 2]), None
 
         else:
             shape = list(im.shape)
@@ -1384,13 +1164,18 @@ class FoCUS:
         return res
 
     # ---------------------------------------------−---------
-    def init_index(self, nside, kernel=-1):
+    def init_index(self, nside, kernel=-1, cell_ids=None):
 
         if kernel == -1:
             l_kernel = self.KERNELSZ
         else:
             l_kernel = kernel
 
+        if cell_ids is not None:
+            ncell = cell_ids.shape[0]
+        else:
+            ncell = 12 * nside * nside
+
         try:
             if self.use_2D:
                 tmp = np.load(
@@ -1398,16 +1183,29 @@ class FoCUS:
                     % (self.TEMPLATE_PATH, l_kernel**2, TMPFILE_VERSION, nside)
                 )
             else:
-                tmp = np.load(
-                    "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        l_kernel**2,
-                        self.NORIENT,
-                        nside,
+                if cell_ids is not None:
+                    tmp = np.load(
+                        "%s/XXXX_%s_W%d_%d_%d_PIDX.npy"  # can not work
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            l_kernel**2,
+                            self.NORIENT,
+                            nside,  # if cell_ids computes the index
+                        )
+                    )
+
+                else:
+                    tmp = np.load(
+                        "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            l_kernel**2,
+                            self.NORIENT,
+                            nside,  # if cell_ids computes the index
+                        )
                     )
-                )
         except:
             if not self.use_2D:
 
@@ -1426,36 +1224,64 @@ class FoCUS:
                     pw2 = 0.25
                     threshold = 4e-5
 
-                th, ph = hp.pix2ang(nside, np.arange(12 * nside**2), nest=True)
-                x, y, z = hp.pix2vec(nside, np.arange(12 * nside**2), nest=True)
+                if cell_ids is not None:
+                    if not isinstance(cell_ids, np.ndarray):
+                        cell_ids = self.backend.to_numpy(cell_ids)
+                    th, ph = hp.pix2ang(nside, cell_ids, nest=True)
+                    x, y, z = hp.pix2vec(nside, cell_ids, nest=True)
 
-                t, p = hp.pix2ang(nside, np.arange(12 * nside * nside), nest=True)
-                phi = [p[k] / np.pi * 180 for k in range(12 * nside * nside)]
-                thi = [t[k] / np.pi * 180 for k in range(12 * nside * nside)]
+                    t, p = hp.pix2ang(nside, cell_ids, nest=True)
+                    phi = [p[k] / np.pi * 180 for k in range(ncell)]
+                    thi = [t[k] / np.pi * 180 for k in range(ncell)]
 
-                indice2 = np.zeros([12 * nside * nside * 64, 2], dtype="int")
-                indice = np.zeros(
-                    [12 * nside * nside * 64 * self.NORIENT, 2], dtype="int"
-                )
-                wav = np.zeros(
-                    [12 * nside * nside * 64 * self.NORIENT], dtype="complex"
-                )
-                wwav = np.zeros([12 * nside * nside * 64 * self.NORIENT], dtype="float")
+                    indice2 = np.zeros([ncell * 64, 2], dtype="int")
+                    indice = np.zeros([ncell * 64 * self.NORIENT, 2], dtype="int")
+                    wav = np.zeros([ncell * 64 * self.NORIENT], dtype="complex")
+                    wwav = np.zeros([ncell * 64 * self.NORIENT], dtype="float")
 
+                else:
+
+                    th, ph = hp.pix2ang(nside, np.arange(12 * nside**2), nest=True)
+                    x, y, z = hp.pix2vec(nside, np.arange(12 * nside**2), nest=True)
+
+                    t, p = hp.pix2ang(nside, np.arange(12 * nside * nside), nest=True)
+                    phi = [p[k] / np.pi * 180 for k in range(12 * nside * nside)]
+                    thi = [t[k] / np.pi * 180 for k in range(12 * nside * nside)]
+
+                    indice2 = np.zeros([12 * nside * nside * 64, 2], dtype="int")
+                    indice = np.zeros(
+                        [12 * nside * nside * 64 * self.NORIENT, 2], dtype="int"
+                    )
+                    wav = np.zeros(
+                        [12 * nside * nside * 64 * self.NORIENT], dtype="complex"
+                    )
+                    wwav = np.zeros(
+                        [12 * nside * nside * 64 * self.NORIENT], dtype="float"
+                    )
                 iv = 0
                 iv2 = 0
-                for iii in range(12 * nside * nside):
 
-                    if iii % (nside * nside) == nside * nside - 1:
-                        if not self.silent:
-                            print(
-                                "Pre-compute nside=%6d %.2f%%"
-                                % (nside, 100 * iii / (12 * nside * nside))
-                            )
+                for iii in range(ncell):
+                    if cell_ids is None:
+                        if iii % (nside * nside) == nside * nside - 1:
+                            if not self.silent:
+                                print(
+                                    "Pre-compute nside=%6d %.2f%%"
+                                    % (nside, 100 * iii / (12 * nside * nside))
+                                )
 
-                    hidx = hp.query_disc(
-                        nside, [x[iii], y[iii], z[iii]], 2 * np.pi / nside, nest=True
-                    )
+                    if cell_ids is not None:
+                        hidx = np.where(
+                            (x - x[iii]) ** 2 + (y - y[iii]) ** 2 + (z - z[iii]) ** 2
+                            < (2 * np.pi / nside) ** 2
+                        )[0]
+                    else:
+                        hidx = hp.query_disc(
+                            nside,
+                            [x[iii], y[iii], z[iii]],
+                            2 * np.pi / nside,
+                            nest=True,
+                        )
 
                     R = hp.Rotator(rot=[phi[iii], -thi[iii]], eulertype="ZYZ")
 
@@ -1474,8 +1300,8 @@ class FoCUS:
                     )
                     idx = np.where((ww**2) > threshold)[0]
                     nval2 = len(idx)
-                    indice2[iv2 : iv2 + nval2, 0] = iii
-                    indice2[iv2 : iv2 + nval2, 1] = hidx[idx]
+                    indice2[iv2 : iv2 + nval2, 1] = iii
+                    indice2[iv2 : iv2 + nval2, 0] = hidx[idx]
                     wwav[iv2 : iv2 + nval2] = ww[idx] / np.sum(ww[idx])
                     iv2 += nval2
 
@@ -1497,15 +1323,18 @@ class FoCUS:
                         idx = np.where(vnorm > threshold)[0]
 
                         nval = len(idx)
-                        indice[iv : iv + nval, 0] = iii * 4 + l_rotation
-                        indice[iv : iv + nval, 1] = hidx[idx]
+                        indice[iv : iv + nval, 1] = iii + l_rotation * ncell
+                        indice[iv : iv + nval, 0] = hidx[idx]
                         # print([hidx[k] for k in idx])
                         # print(hp.query_disc(nside, [x[iii],y[iii],z[iii]], np.pi/nside,nest=True))
                         normr = np.mean(wresr[idx])
                         normi = np.mean(wresi[idx])
 
                         val = wresr[idx] - normr + 1j * (wresi[idx] - normi)
-                        val = val / abs(val).sum()
+                        r = abs(val).sum()
+
+                        if r > 0:
+                            val = val / r
 
                         wav[iv : iv + nval] = val
                         iv += nval
@@ -1609,56 +1438,57 @@ class FoCUS:
                 wav=w.flatten()
                 wwav=wwav.flatten()
                 """
-                if not self.silent:
-                    print(
-                        "Write FOSCAT_%s_W%d_%d_%d_PIDX.npy"
-                        % (TMPFILE_VERSION, self.KERNELSZ**2, self.NORIENT, nside)
+                if cell_ids is None:
+                    if not self.silent:
+                        print(
+                            "Write FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                            % (TMPFILE_VERSION, self.KERNELSZ**2, self.NORIENT, nside)
+                        )
+                    np.save(
+                        "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            self.NORIENT,
+                            nside,
+                        ),
+                        indice,
+                    )
+                    np.save(
+                        "%s/FOSCAT_%s_W%d_%d_%d_WAVE.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            self.NORIENT,
+                            nside,
+                        ),
+                        wav,
+                    )
+                    np.save(
+                        "%s/FOSCAT_%s_W%d_%d_%d_PIDX2.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            self.NORIENT,
+                            nside,
+                        ),
+                        indice2,
+                    )
+                    np.save(
+                        "%s/FOSCAT_%s_W%d_%d_%d_SMOO.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            self.NORIENT,
+                            nside,
+                        ),
+                        wwav,
                     )
-                np.save(
-                    "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        self.KERNELSZ**2,
-                        self.NORIENT,
-                        nside,
-                    ),
-                    indice,
-                )
-                np.save(
-                    "%s/FOSCAT_%s_W%d_%d_%d_WAVE.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        self.KERNELSZ**2,
-                        self.NORIENT,
-                        nside,
-                    ),
-                    wav,
-                )
-                np.save(
-                    "%s/FOSCAT_%s_W%d_%d_%d_PIDX2.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        self.KERNELSZ**2,
-                        self.NORIENT,
-                        nside,
-                    ),
-                    indice2,
-                )
-                np.save(
-                    "%s/FOSCAT_%s_W%d_%d_%d_SMOO.npy"
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        self.KERNELSZ**2,
-                        self.NORIENT,
-                        nside,
-                    ),
-                    wwav,
-                )
-            else:
+            if self.use_2D:
                 if l_kernel**2 == 9:
                     if self.rank == 0:
                         self.comp_idx_w9(nside)
@@ -1674,23 +1504,24 @@ class FoCUS:
                             )
                         return None
 
-        self.barrier()
-        if self.use_2D:
-            tmp = np.load(
-                "%s/W%d_%s_%d_IDX.npy"
-                % (self.TEMPLATE_PATH, l_kernel**2, TMPFILE_VERSION, nside)
-            )
-        else:
-            tmp = np.load(
-                "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
-                % (
-                    self.TEMPLATE_PATH,
-                    TMPFILE_VERSION,
-                    self.KERNELSZ**2,
-                    self.NORIENT,
-                    nside,
+        if cell_ids is None:
+            self.barrier()
+            if self.use_2D:
+                tmp = np.load(
+                    "%s/W%d_%s_%d_IDX.npy"
+                    % (self.TEMPLATE_PATH, l_kernel**2, TMPFILE_VERSION, nside)
+                )
+            else:
+                tmp = np.load(
+                    "%s/FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                    % (
+                        self.TEMPLATE_PATH,
+                        TMPFILE_VERSION,
+                        self.KERNELSZ**2,
+                        self.NORIENT,
+                        nside,
+                    )
                 )
-            )
             tmp2 = np.load(
                 "%s/FOSCAT_%s_W%d_%d_%d_PIDX2.npy"
                 % (
@@ -1731,22 +1562,28 @@ class FoCUS:
                     nside,
                 )
             )
-
-            wr = self.backend.bk_SparseTensor(
-                self.backend.bk_constant(tmp),
-                self.backend.bk_constant(self.backend.bk_cast(wr)),
-                dense_shape=[12 * nside**2 * self.NORIENT, 12 * nside**2],
-            )
-            wi = self.backend.bk_SparseTensor(
-                self.backend.bk_constant(tmp),
-                self.backend.bk_constant(self.backend.bk_cast(wi)),
-                dense_shape=[12 * nside**2 * self.NORIENT, 12 * nside**2],
-            )
-            ws = self.backend.bk_SparseTensor(
-                self.backend.bk_constant(tmp2),
-                self.backend.bk_constant(self.backend.bk_cast(ws)),
-                dense_shape=[12 * nside**2, 12 * nside**2],
-            )
+        else:
+            tmp = indice
+            tmp2 = indice2
+            wr = wav.real
+            wi = wav.imag
+            ws = self.slope * wwav
+
+        wr = self.backend.bk_SparseTensor(
+            self.backend.bk_constant(tmp),
+            self.backend.bk_constant(self.backend.bk_cast(wr)),
+            dense_shape=[ncell, self.NORIENT * ncell],
+        )
+        wi = self.backend.bk_SparseTensor(
+            self.backend.bk_constant(tmp),
+            self.backend.bk_constant(self.backend.bk_cast(wi)),
+            dense_shape=[ncell, self.NORIENT * ncell],
+        )
+        ws = self.backend.bk_SparseTensor(
+            self.backend.bk_constant(tmp2),
+            self.backend.bk_constant(self.backend.bk_cast(ws)),
+            dense_shape=[ncell, ncell],
+        )
 
         if kernel == -1:
             self.Idx_Neighbours[nside] = tmp
@@ -1757,42 +1594,268 @@ class FoCUS:
 
         return wr, wi, ws, tmp
 
+
     # ---------------------------------------------−---------
-    # convert swap axes tensor x [....,a,....,b,....] to [....,b,....,a,....]
-    def swapaxes(self, x, axis1, axis2):
-        shape = list(x.shape)
-        if axis1 < 0:
-            laxis1 = len(shape) + axis1
-        else:
-            laxis1 = axis1
-        if axis2 < 0:
-            laxis2 = len(shape) + axis2
+    def init_index_cnn(self, nside, NORIENT=4,kernel=-1, cell_ids=None):
+
+        if kernel == -1:
+            l_kernel = self.KERNELSZ
         else:
-            laxis2 = axis2
+            l_kernel = kernel
 
-        naxes = len(shape)
-        thelist = [i for i in range(naxes)]
-        thelist[laxis1] = laxis2
-        thelist[laxis2] = laxis1
-        return self.backend.bk_transpose(x, thelist)
+        if cell_ids is not None:
+            ncell = cell_ids.shape[0]
+        else:
+            ncell = 12 * nside * nside
 
-    # ---------------------------------------------−---------
-    # Mean using mask x [....,Npix,....], mask[Nmask,Npix]  to [....,Nmask,....]
-    # if use_2D
-    # Mean using mask x [....,12,Nside+2*off,Nside+2*off,....], mask[Nmask,12,Nside+2*off,Nside+2*off]  to [....,Nmask,....]
-    def masked_mean(self, x, mask, axis=0, rank=0, calc_var=False):
+        try:
+            
+            if cell_ids is not None:
+                tmp = np.load(
+                      "%s/XXXX_%s_W%d_%d_%d_PIDX.npy"  # can not work
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            l_kernel**2,
+                            NORIENT,
+                            nside,  # if cell_ids computes the index
+                        )
+                    )
 
-        # ==========================================================================
-        # in input data=[Nbatch,...,X[,Y],NORIENT[,NORIENT]]
-        # in input mask=[Nmask,X[,Y]]
-        # if self.use_2D :  X[,Y]] = [X,Y]
-        # if second level:  NORIENT[,NORIENT]= NORIENT,NORIENT
-        # ==========================================================================
+            else:
+                    tmp = np.load(
+                        "%s/CNN_FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            l_kernel**2,
+                            NORIENT,
+                            nside,  # if cell_ids computes the index
+                        )
+                    )
+        except:
 
-        shape = list(x.shape)
+            pw = 8.0 
+            pw2 = 1.0
+            threshold = 1e-3
+            
+            if l_kernel == 5:
+                    pw = 8.0
+                    pw2 = 0.5
+                    threshold = 2e-4
 
-        if not self.use_2D:
-            nside = int(np.sqrt(x.shape[axis] // 12))
+            elif l_kernel == 3:
+                    pw = 8.0 
+                    pw2 = 1.0
+                    threshold = 1e-3
+
+            elif l_kernel == 7:
+                    pw = 8.0
+                    pw2 = 0.25
+                    threshold = 4e-5
+            
+            n_weights = self.KERNELSZ*(self.KERNELSZ//2+1)
+            
+            if cell_ids is not None:
+                    if not isinstance(cell_ids, np.ndarray):
+                        cell_ids = self.backend.to_numpy(cell_ids)
+                    th, ph = hp.pix2ang(nside, cell_ids, nest=True)
+                    x, y, z = hp.pix2vec(nside, cell_ids, nest=True)
+
+                    t, p = hp.pix2ang(nside, cell_ids, nest=True)
+                    phi = [p[k] / np.pi * 180 for k in range(ncell)]
+                    thi = [t[k] / np.pi * 180 for k in range(ncell)]
+
+                    indice = np.zeros([n_weights, NORIENT, ncell,4], dtype="int")
+                    
+                    wav = np.zeros([n_weights, NORIENT, ncell,4], dtype="float")
+
+            else:
+
+                    th, ph = hp.pix2ang(nside, np.arange(12 * nside**2), nest=True)
+                    x, y, z = hp.pix2vec(nside, np.arange(12 * nside**2), nest=True)
+
+                    t, p = hp.pix2ang(nside, np.arange(12 * nside * nside), nest=True)
+                    phi = [p[k] / np.pi * 180 for k in range(12 * nside * nside)]
+                    thi = [t[k] / np.pi * 180 for k in range(12 * nside * nside)]
+
+                    indice = np.zeros(
+                        [n_weights, NORIENT, 12 * nside * nside,4], dtype="int"
+                    )
+                    wav = np.zeros(
+                        [n_weights, NORIENT, 12 * nside * nside,4], dtype="float"
+                    )
+            iv = 0
+            iv2 = 0
+
+            for iii in range(ncell):
+                if cell_ids is None:
+                    if iii % (nside * nside) == nside * nside - 1:
+                        if not self.silent:
+                                print(
+                                    "Pre-compute nside=%6d %.2f%%"
+                                    % (nside, 100 * iii / (12 * nside * nside))
+                                )
+
+                if cell_ids is not None:
+                        hidx = np.where(
+                            (x - x[iii]) ** 2 + (y - y[iii]) ** 2 + (z - z[iii]) ** 2
+                            < (2 * np.pi / nside) ** 2
+                        )[0]
+                else:
+                        hidx = hp.query_disc(
+                            nside,
+                            [x[iii], y[iii], z[iii]],
+                            2 * np.pi / nside,
+                            nest=True,
+                        )
+
+                R = hp.Rotator(rot=[phi[iii], -thi[iii]], eulertype="ZYZ")
+
+                t2, p2 = R(th[hidx], ph[hidx])
+
+                vec2 = hp.ang2vec(t2, p2)
+
+                x2 = vec2[:, 0]
+                y2 = vec2[:, 1]
+                z2 = vec2[:, 2]
+
+                for l_rotation in range(NORIENT):
+
+                        angle = (
+                            l_rotation / 4.0 * np.pi
+                            - phi[iii] / 180.0 * np.pi * (z[hidx] > 0)
+                            - (180.0 - phi[iii]) / 180.0 * np.pi * (z[hidx] < 0)
+                        )
+
+
+                        axes  = y2 * np.cos(angle)  - x2 * np.sin(angle)
+                        axes2 = -y2 * np.sin(angle) - x2 * np.cos(angle)
+                        
+                        for k_weights in range(self.KERNELSZ//2+1):
+                            for l_weights in range(self.KERNELSZ):
+                                
+                                val=np.exp(-(pw*(axes2*(nside)-(k_weights-self.KERNELSZ//2))**2+pw*(axes*(nside)-(l_weights-self.KERNELSZ//2))**2))+ \
+                                    np.exp(-(pw*(axes2*(nside)+(k_weights-self.KERNELSZ//2))**2+pw*(axes*(nside)-(l_weights-self.KERNELSZ//2))**2))
+
+                                idx = np.argsort(-val)
+                                idx = idx[0:4]
+                                
+                                nval = len(idx)
+                                val=val[idx]
+                                
+                                r = abs(val).sum()
+
+                                if r > 0:
+                                    val = val / r
+                                    
+                                indice[k_weights*self.KERNELSZ+l_weights,l_rotation,iii,:] = hidx[idx]
+                                wav[k_weights*self.KERNELSZ+l_weights,l_rotation,iii,:] = val
+                
+            if not self.silent:
+                    print("Kernel Size ", iv / (NORIENT * 12 * nside * nside))
+                
+            if cell_ids is None:
+                    if not self.silent:
+                        print(
+                            "Write FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                            % (TMPFILE_VERSION, self.KERNELSZ**2, NORIENT, nside)
+                        )
+                    np.save(
+                        "%s/CNN_FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            NORIENT,
+                            nside,
+                        ),
+                        indice,
+                    )
+                    np.save(
+                        "%s/CNN_FOSCAT_%s_W%d_%d_%d_WAVE.npy"
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            self.KERNELSZ**2,
+                            NORIENT,
+                            nside,
+                        ),
+                        wav,
+                    )
+
+        if cell_ids is None:
+            self.barrier()
+            if self.use_2D:
+                tmp = np.load(
+                    "%s/W%d_%s_%d_IDX.npy"
+                    % (self.TEMPLATE_PATH, l_kernel**2, TMPFILE_VERSION, nside)
+                )
+            else:
+                tmp = np.load(
+                    "%s/CNN_FOSCAT_%s_W%d_%d_%d_PIDX.npy"
+                    % (
+                        self.TEMPLATE_PATH,
+                        TMPFILE_VERSION,
+                        self.KERNELSZ**2,
+                        NORIENT,
+                        nside,
+                    )
+                )
+            wav = np.load(
+                "%s/CNN_FOSCAT_%s_W%d_%d_%d_WAVE.npy"
+                % (
+                    self.TEMPLATE_PATH,
+                    TMPFILE_VERSION,
+                    self.KERNELSZ**2,
+                    NORIENT,
+                    nside,
+                )
+            )
+        else:
+            tmp = indice
+        
+        self.Idx_CNN[(nside,NORIENT,self.KERNELSZ)] = tmp
+        self.Idx_WCNN[(nside,NORIENT,self.KERNELSZ)] = self.backend.bk_cast(wav)
+
+        return wav, tmp
+        
+    # ---------------------------------------------−---------
+    # convert swap axes tensor x [....,a,....,b,....] to [....,b,....,a,....]
+    def swapaxes(self, x, axis1, axis2):
+        shape = list(x.shape)
+        if axis1 < 0:
+            laxis1 = len(shape) + axis1
+        else:
+            laxis1 = axis1
+        if axis2 < 0:
+            laxis2 = len(shape) + axis2
+        else:
+            laxis2 = axis2
+
+        naxes = len(shape)
+        thelist = [i for i in range(naxes)]
+        thelist[laxis1] = laxis2
+        thelist[laxis2] = laxis1
+        return self.backend.bk_transpose(x, thelist)
+
+    # ---------------------------------------------−---------
+    # Mean using mask x [....,Npix,....], mask[Nmask,Npix]  to [....,Nmask,....]
+    # if use_2D
+    # Mean using mask x [....,12,Nside+2*off,Nside+2*off,....], mask[Nmask,12,Nside+2*off,Nside+2*off]  to [....,Nmask,....]
+    def masked_mean(self, x, mask, axis=0, rank=0, calc_var=False):
+
+        # ==========================================================================
+        # in input data=[Nbatch,...,NORIENT[,NORIENT],X[,Y]]
+        # in input mask=[Nmask,X[,Y]]
+        # if self.use_2D :  X[,Y]] = [X,Y]
+        # if second level:  NORIENT[,NORIENT]= NORIENT,NORIENT
+        # ==========================================================================
+
+        shape = list(x.shape)
+
+        if not self.use_2D and not self.use_1D:
+            nside = int(np.sqrt(x.shape[axis] // 12))
 
         l_mask = mask
         if self.mask_norm:
@@ -1802,6 +1865,7 @@ class FoCUS:
                 ),
                 1,
             )
+
             if not self.use_2D:
                 l_mask = (
                     12
@@ -1845,13 +1909,11 @@ class FoCUS:
                     ]
 
             ichannel = 1
-            for i in range(axis):
+            for i in range(1, len(shape) - 2):
                 ichannel *= shape[i]
-            ochannel = 1
-            for i in range(axis + 2, len(shape)):
-                ochannel *= shape[i]
+
             l_x = self.backend.bk_reshape(
-                x, [ichannel, 1, shape[axis], shape[axis + 1], ochannel]
+                x, [shape[0], 1, ichannel, shape[-2], shape[-1]]
             )
 
             if self.padding == "VALID":
@@ -1876,12 +1938,10 @@ class FoCUS:
                     l_mask = l_mask[:, self.KERNELSZ // 2 : -self.KERNELSZ // 2 + 1]
 
             ichannel = 1
-            for i in range(axis):
+            for i in range(1, len(shape) - 1):
                 ichannel *= shape[i]
-            ochannel = 1
-            for i in range(axis + 1, len(shape)):
-                ochannel *= shape[i]
-            l_x = self.backend.bk_reshape(x, [ichannel, 1, shape[axis], ochannel])
+
+            l_x = self.backend.bk_reshape(x, [shape[0], 1, ichannel, shape[-1]])
 
             if self.padding == "VALID":
                 oshape = [k for k in shape]
@@ -1891,18 +1951,14 @@ class FoCUS:
                 )
         else:
             ichannel = 1
-            for i in range(axis):
+            for i in range(len(shape) - 1):
                 ichannel *= shape[i]
-            ochannel = 1
-            for i in range(axis + 1, len(shape)):
-                ochannel *= shape[i]
-            l_x = self.backend.bk_reshape(x, [ichannel, 1, shape[axis], ochannel])
 
-        # data=[Nbatch,...,X[,Y],NORIENT[,NORIENT]] => data=[Nbatch,1,...,X[,Y],NORIENT[,NORIENT]]
-        # mask=[Nmask,X[,Y]] => mask=[1,Nmask,X[,Y]]
-        l_mask = self.backend.bk_expand_dims(l_mask, 0)
-        # mask=[1,Nmask,X[,Y]] => mask=[1,Nmask,X[,Y],1]
-        l_mask = self.backend.bk_expand_dims(l_mask, -1)
+            l_x = self.backend.bk_reshape(x, [ichannel, 1, shape[-1]])
+
+        # data=[Nbatch,...,NORIENT[,NORIENT],X[,Y]] => data=[Nbatch,1,...,NORIENT[,NORIENT],X[,Y]]
+        # mask=[Nmask,X[,Y]] => mask=[1,Nmask,....,X[,Y]]
+        l_mask = self.backend.bk_expand_dims(self.backend.bk_expand_dims(l_mask, 0), 0)
 
         if l_x.dtype == self.all_cbk_type:
             l_mask = self.backend.bk_complex(l_mask, self.backend.bk_cast(0.0 * l_mask))
@@ -1916,21 +1972,23 @@ class FoCUS:
             #    vtmp = l_x[:,self.KERNELSZ // 2 : -self.KERNELSZ // 2,self.KERNELSZ // 2 : -self.KERNELSZ // 2,:]
 
             v1 = self.backend.bk_reduce_sum(
-                self.backend.bk_reduce_sum(mtmp * vtmp, axis=2), 2
+                self.backend.bk_reduce_sum(mtmp * vtmp, axis=-1), -1
             )
             v2 = self.backend.bk_reduce_sum(
-                self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=2), 2
+                self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=-1), -1
+            )
+            vh = self.backend.bk_reduce_sum(
+                self.backend.bk_reduce_sum(mtmp, axis=-1), -1
             )
-            vh = self.backend.bk_reduce_sum(self.backend.bk_reduce_sum(mtmp, axis=2), 2)
 
             res = v1 / vh
 
-            oshape = []
+            oshape = [x.shape[0]] + [mask.shape[0]]
             if axis > 0:
-                oshape = oshape + list(x.shape[0:axis])
-            oshape = oshape + [mask.shape[0]]
-            if axis + 1 < len(x.shape):
-                oshape = oshape + list(x.shape[axis + 2 :])
+                oshape = oshape + list(x.shape[1:axis])
+
+            if len(x.shape[axis:-2]) > 0:
+                oshape = oshape + list(x.shape[axis:-2])
 
             if calc_var:
                 if self.backend.bk_is_complex(vtmp):
@@ -1960,19 +2018,15 @@ class FoCUS:
         elif self.use_1D:
             mtmp = l_mask
             vtmp = l_x
-
-            v1 = self.backend.bk_reduce_sum(mtmp * vtmp, axis=2)
-            v2 = self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=2)
-            vh = self.backend.bk_reduce_sum(mtmp, axis=2)
+            v1 = self.backend.bk_reduce_sum(mtmp * vtmp, axis=-1)
+            v2 = self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=-1)
+            vh = self.backend.bk_reduce_sum(mtmp, axis=-1)
 
             res = v1 / vh
 
-            oshape = []
-            if axis > 0:
-                oshape = oshape + list(x.shape[0:axis])
-            oshape = oshape + [mask.shape[0]]
-            if axis + 1 < len(x.shape):
-                oshape = oshape + list(x.shape[axis + 1 :])
+            oshape = [x.shape[0]] + [mask.shape[0]]
+            if len(x.shape) > 1:
+                oshape = oshape + list(x.shape[1:-1])
 
             if calc_var:
                 if self.backend.bk_is_complex(vtmp):
@@ -1991,7 +2045,6 @@ class FoCUS:
                     )
                 else:
                     res2 = self.backend.bk_sqrt((v2 / vh - res * res) / (vh))
-
                 res = self.backend.bk_reshape(res, oshape)
                 res2 = self.backend.bk_reshape(res2, oshape)
                 return res, res2
@@ -2000,18 +2053,20 @@ class FoCUS:
                 return res
 
         else:
-            v1 = self.backend.bk_reduce_sum(l_mask * l_x, axis=2)
-            v2 = self.backend.bk_reduce_sum(l_mask * l_x * l_x, axis=2)
-            vh = self.backend.bk_reduce_sum(l_mask, axis=2)
+            v1 = self.backend.bk_reduce_sum(l_mask * l_x, axis=-1)
+            v2 = self.backend.bk_reduce_sum(l_mask * l_x * l_x, axis=-1)
+            vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
 
             res = v1 / vh
 
             oshape = []
             if axis > 0:
-                oshape = oshape + list(x.shape[0:axis])
+                oshape = [x.shape[0]]
+            else:
+                oshape = [1]
             oshape = oshape + [mask.shape[0]]
-            if axis + 1 < len(x.shape):
-                oshape = oshape + list(x.shape[axis + 1 :])
+            if axis > 1:
+                oshape = oshape + list(x.shape[1:-1])
 
             if calc_var:
                 if self.backend.bk_is_complex(l_x):
@@ -2176,169 +2231,67 @@ class FoCUS:
                     print("Use of 2D scat with data that has less than 2D")
                 return None
 
-            npix = ishape[axis]
-            npiy = ishape[axis + 1]
-            odata = 1
-            if len(ishape) > axis + 2:
-                for k in range(axis + 2, len(ishape)):
-                    odata = odata * ishape[k]
+            npix = ishape[-2]
+            npiy = ishape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(ishape) - 2):
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
-                self.backend.bk_cast(in_image), [ndata, npix, npiy, odata]
+                self.backend.bk_cast(in_image), [ndata, npix, npiy]
             )
 
             if self.backend.bk_is_complex(tim):
-                rr1 = self.backend.conv2d(
-                    self.backend.bk_real(tim),
-                    self.ww_RealT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
-                ii1 = self.backend.conv2d(
-                    self.backend.bk_real(tim),
-                    self.ww_ImagT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
-                rr2 = self.backend.conv2d(
-                    self.backend.bk_imag(tim),
-                    self.ww_RealT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
-                ii2 = self.backend.conv2d(
-                    self.backend.bk_imag(tim),
-                    self.ww_ImagT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
+                rr1 = self.backend.conv2d(self.backend.bk_real(tim), self.ww_RealT[1])
+                ii1 = self.backend.conv2d(self.backend.bk_real(tim), self.ww_ImagT[1])
+                rr2 = self.backend.conv2d(self.backend.bk_imag(tim), self.ww_RealT[1])
+                ii2 = self.backend.conv2d(self.backend.bk_imag(tim), self.ww_ImagT[1])
                 res = self.backend.bk_complex(rr1 - ii2, ii1 + rr2)
             else:
-                rr = self.backend.conv2d(
-                    tim,
-                    self.ww_RealT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
-                ii = self.backend.conv2d(
-                    tim,
-                    self.ww_ImagT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
+                rr = self.backend.conv2d(tim, self.ww_RealT[1])
+                ii = self.backend.conv2d(tim, self.ww_ImagT[1])
                 res = self.backend.bk_complex(rr, ii)
 
-            if axis == 0:
-                if len(ishape) == 2:
-                    return self.backend.bk_reshape(
-                        res, [res.shape[1], res.shape[2], self.NORIENT]
-                    )
-                else:
-                    return self.backend.bk_reshape(
-                        res,
-                        [res.shape[1], res.shape[2], self.NORIENT] + ishape[axis + 2 :],
-                    )
-            else:
-                if len(ishape) == axis + 2:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [res.shape[1], res.shape[2], self.NORIENT]
-                    )
-                else:
-                    return self.backend.bk_reshape(
-                        res,
-                        ishape[0:axis]
-                        + [res.shape[1], res.shape[2], self.NORIENT]
-                        + ishape[axis + 2 :],
-                    )
+            return self.backend.bk_reshape(
+                res, ishape[0:-2] + [self.NORIENT, npix, npiy]
+            )
 
-            return self.backend.bk_reshape(res, in_image.shape + [self.NORIENT])
         elif self.use_1D:
             ishape = list(in_image.shape)
-            if len(ishape) < axis + 1:
-                if not self.silent:
-                    print("Use of 1D scat with data that has less than 1D")
-                return None
 
-            npix = ishape[axis]
-            odata = 1
-            if len(ishape) > axis + 1:
-                for k in range(axis + 1, len(ishape)):
-                    odata = odata * ishape[k]
+            npix = ishape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(ishape) - 1):
                 ndata = ndata * ishape[k]
 
-            tim = self.backend.bk_reshape(
-                self.backend.bk_cast(in_image), [ndata, npix, odata]
-            )
+            tim = self.backend.bk_reshape(self.backend.bk_cast(in_image), [ndata, npix])
 
             if self.backend.bk_is_complex(tim):
-                rr1 = self.backend.conv1d(
-                    self.backend.bk_real(tim),
-                    self.ww_RealT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
-                ii1 = self.backend.conv1d(
-                    self.backend.bk_real(tim),
-                    self.ww_ImagT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
-                rr2 = self.backend.conv1d(
-                    self.backend.bk_imag(tim),
-                    self.ww_RealT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
-                ii2 = self.backend.conv1d(
-                    self.backend.bk_imag(tim),
-                    self.ww_ImagT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
+                rr1 = self.backend.conv1d(self.backend.bk_real(tim), self.ww_RealT[1])
+                ii1 = self.backend.conv1d(self.backend.bk_real(tim), self.ww_ImagT[1])
+                rr2 = self.backend.conv1d(self.backend.bk_imag(tim), self.ww_RealT[1])
+                ii2 = self.backend.conv1d(self.backend.bk_imag(tim), self.ww_ImagT[1])
                 res = self.backend.bk_complex(rr1 - ii2, ii1 + rr2)
             else:
-                rr = self.backend.conv1d(
-                    tim, self.ww_RealT[odata], strides=[1, 1, 1], padding=self.padding
-                )
-                ii = self.backend.conv1d(
-                    tim, self.ww_ImagT[odata], strides=[1, 1, 1], padding=self.padding
-                )
+                rr = self.backend.conv1d(tim, self.ww_RealT[1])
+                ii = self.backend.conv1d(tim, self.ww_ImagT[1])
                 res = self.backend.bk_complex(rr, ii)
 
-            if axis == 0:
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [res.shape[1]])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [res.shape[1]] + ishape[axis + 2 :]
-                    )
-            else:
-                if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(res, ishape[0:axis] + [res.shape[1]])
-                else:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [res.shape[1]] + ishape[axis + 1 :]
-                    )
-
-            return self.backend.bk_reshape(res, in_image.shape + [self.NORIENT])
+            return self.backend.bk_reshape(res, ishape)
 
         else:
             ishape = list(image.shape)
-
+            """
             if cell_ids is not None:
                 if cell_ids.shape[0] not in self.padding_conv:
+                    print(image.shape,cell_ids.shape)
                     import healpix_convolution as hc
                     from xdggs.healpix import HealpixInfo
 
                     res = self.backend.bk_zeros(
-                        ishape + [self.NORIENT], dtype=self.backend.all_cbk_type
+                        ishape[0:-1] + [self.NORIENT]+ishape[-1], dtype=self.backend.all_cbk_type
                     )
 
                     grid_info = HealpixInfo(
@@ -2384,14 +2337,15 @@ class FoCUS:
                                     padded_data
                                 ) + 1j * kernelI.matmul(padded_data)
                 return res
-
-            nside = int(np.sqrt(image.shape[axis] // 12))
+            """
+            if nside is None:
+                nside = int(np.sqrt(image.shape[-1] // 12))
 
             if self.Idx_Neighbours[nside] is None:
                 if self.InitWave is None:
-                    wr, wi, ws, widx = self.init_index(nside)
+                    wr, wi, ws, widx = self.init_index(nside, cell_ids=cell_ids)
                 else:
-                    wr, wi, ws, widx = self.InitWave(self, nside)
+                    wr, wi, ws, widx = self.InitWave(nside, cell_ids=cell_ids)
 
                 self.Idx_Neighbours[nside] = 1  # self.backend.bk_constant(tmp)
                 self.ww_Real[nside] = wr
@@ -2401,156 +2355,63 @@ class FoCUS:
             l_ww_real = self.ww_Real[nside]
             l_ww_imag = self.ww_Imag[nside]
 
-            odata = 1
-            for k in range(axis + 1, len(ishape)):
-                odata = odata * ishape[k]
+            # always convolve the last dimension
 
-            if axis > 0:
-                ndata = 1
-                for k in range(axis):
+            ndata = 1
+            if len(ishape) > 1:
+                for k in range(len(ishape) - 1):
                     ndata = ndata * ishape[k]
-                tim = self.backend.bk_reshape(
-                    self.backend.bk_cast(image), [ndata, 12 * nside**2, odata]
-                )
-                if tim.dtype == self.all_cbk_type:
-                    rr1 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_real, self.backend.bk_real(tim[0])
-                        ),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii1 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_imag, self.backend.bk_real(tim[0])
-                        ),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    rr2 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_real, self.backend.bk_imag(tim[0])
-                        ),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii2 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_imag, self.backend.bk_imag(tim[0])
-                        ),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    res = self.backend.bk_complex(rr1 - ii2, ii1 + rr2)
-                else:
-                    rr = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(l_ww_real, tim[0]),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(l_ww_imag, tim[0]),
-                        [1, 12 * nside**2, self.NORIENT, odata],
-                    )
-                    res = self.backend.bk_complex(rr, ii)
-
-                for k in range(1, ndata):
-                    if tim.dtype == self.all_cbk_type:
-                        rr1 = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_ww_real, self.backend.bk_real(tim[k])
-                            ),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        ii1 = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_ww_imag, self.backend.bk_real(tim[k])
-                            ),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        rr2 = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_ww_real, self.backend.bk_imag(tim[k])
-                            ),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        ii2 = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_ww_imag, self.backend.bk_imag(tim[k])
-                            ),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        res = self.backend.bk_concat(
-                            [res, self.backend.bk_complex(rr1 - ii2, ii1 + rr2)], 0
-                        )
-                    else:
-                        rr = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(l_ww_real, tim[k]),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        ii = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(l_ww_imag, tim[k]),
-                            [1, 12 * nside**2, self.NORIENT, odata],
-                        )
-                        res = self.backend.bk_concat(
-                            [res, self.backend.bk_complex(rr, ii)], 0
-                        )
-
-                if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [12 * nside**2, self.NORIENT]
-                    )
-                else:
-                    return self.backend.bk_reshape(
-                        res,
-                        ishape[0:axis]
-                        + [12 * nside**2]
-                        + ishape[axis + 1 :]
-                        + [self.NORIENT],
-                    )
+            tim = self.backend.bk_reshape(
+                self.backend.bk_cast(image), [ndata, ishape[-1]]
+            )
 
-            if axis == 0:
-                tim = self.backend.bk_reshape(
-                    self.backend.bk_cast(image), [12 * nside**2, odata]
+            if tim.dtype == self.all_cbk_type:
+                rr1 = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(
+                        self.backend.bk_real(tim),
+                        l_ww_real,
+                    ),
+                    [ndata, self.NORIENT, ishape[-1]],
                 )
-                if tim.dtype == self.all_cbk_type:
-                    rr1 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_real, self.backend.bk_real(tim)
-                        ),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii1 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_imag, self.backend.bk_real(tim)
-                        ),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    rr2 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_real, self.backend.bk_imag(tim)
-                        ),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii2 = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_ww_imag, self.backend.bk_imag(tim)
-                        ),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    res = self.backend.bk_complex(rr1 - ii2, ii1 + rr2)
-                else:
-                    rr = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(l_ww_real, tim),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    ii = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(l_ww_imag, tim),
-                        [12 * nside**2, self.NORIENT, odata],
-                    )
-                    res = self.backend.bk_complex(rr, ii)
+                ii1 = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(
+                        self.backend.bk_real(tim),
+                        l_ww_imag,
+                    ),
+                    [ndata, self.NORIENT, ishape[-1]],
+                )
+                rr2 = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(
+                        self.backend.bk_imag(tim),
+                        l_ww_real,
+                    ),
+                    [ndata, self.NORIENT, ishape[-1]],
+                )
+                ii2 = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(
+                        self.backend.bk_imag(tim),
+                        l_ww_imag,
+                    ),
+                    [ndata, self.NORIENT, ishape[-1]],
+                )
+                res = self.backend.bk_complex(rr1 - ii2, ii1 + rr2)
+            else:
+                rr = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(tim, l_ww_real),
+                    [ndata, self.NORIENT, ishape[-1]],
+                )
+                ii = self.backend.bk_reshape(
+                    self.backend.bk_sparse_dense_matmul(tim, l_ww_imag),
+                    [ndata, self.NORIENT, ishape[-1]],
+                )
+                res = self.backend.bk_complex(rr, ii)
+            if len(ishape) > 1:
+                return self.backend.bk_reshape(
+                    res, ishape[0:-1] + [self.NORIENT, ishape[-1]]
+                )
+            else:
+                return self.backend.bk_reshape(res, [self.NORIENT, ishape[-1]])
 
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [12 * nside**2, self.NORIENT])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [12 * nside**2] + ishape[axis + 1 :] + [self.NORIENT]
-                    )
         return res
 
     # ---------------------------------------------−---------
@@ -2578,114 +2439,43 @@ class FoCUS:
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
-                self.backend.bk_cast(in_image), [ndata, npix, npiy, odata]
+                self.backend.bk_cast(in_image), [ndata, npix, npiy]
             )
 
             if self.backend.bk_is_complex(tim):
-                rr = self.backend.conv2d(
-                    self.backend.bk_real(tim),
-                    self.ww_SmoothT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
-                ii = self.backend.conv2d(
-                    self.backend.bk_imag(tim),
-                    self.ww_SmoothT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
+                rr = self.backend.conv2d(self.backend.bk_real(tim), self.ww_SmoothT[1])
+                ii = self.backend.conv2d(self.backend.bk_imag(tim), self.ww_SmoothT[1])
                 res = self.backend.bk_complex(rr, ii)
             else:
-                res = self.backend.conv2d(
-                    tim,
-                    self.ww_SmoothT[odata],
-                    strides=[1, 1, 1, 1],
-                    padding=self.padding,
-                )
+                res = self.backend.conv2d(tim, self.ww_SmoothT[1])
 
-            if axis == 0:
-                if len(ishape) == 2:
-                    return self.backend.bk_reshape(res, [res.shape[1], res.shape[2]])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [res.shape[1], res.shape[2]] + ishape[axis + 2 :]
-                    )
-            else:
-                if len(ishape) == axis + 2:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [res.shape[1], res.shape[2]]
-                    )
-                else:
-                    return self.backend.bk_reshape(
-                        res,
-                        ishape[0:axis]
-                        + [res.shape[1], res.shape[2]]
-                        + ishape[axis + 2 :],
-                    )
+            return self.backend.bk_reshape(res, ishape)
 
-            return self.backend.bk_reshape(res, in_image.shape)
         elif self.use_1D:
 
             ishape = list(in_image.shape)
-            if len(ishape) < axis + 1:
-                if not self.silent:
-                    print("Use of 1D scat with data that has less than 1D")
-                return None
 
-            npix = ishape[axis]
-            odata = 1
-            if len(ishape) > axis + 1:
-                for k in range(axis + 1, len(ishape)):
-                    odata = odata * ishape[k]
+            npix = ishape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(ishape) - 1):
                 ndata = ndata * ishape[k]
 
-            tim = self.backend.bk_reshape(
-                self.backend.bk_cast(in_image), [ndata, npix, odata]
-            )
+            tim = self.backend.bk_reshape(self.backend.bk_cast(in_image), [ndata, npix])
 
             if self.backend.bk_is_complex(tim):
-                rr = self.backend.conv1d(
-                    self.backend.bk_real(tim),
-                    self.ww_SmoothT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
-                ii = self.backend.conv1d(
-                    self.backend.bk_imag(tim),
-                    self.ww_SmoothT[odata],
-                    strides=[1, 1, 1],
-                    padding=self.padding,
-                )
+                rr = self.backend.conv1d(self.backend.bk_real(tim), self.ww_SmoothT[1])
+                ii = self.backend.conv1d(self.backend.bk_imag(tim), self.ww_SmoothT[1])
                 res = self.backend.bk_complex(rr, ii)
             else:
-                res = self.backend.conv1d(
-                    tim, self.ww_SmoothT[odata], strides=[1, 1, 1], padding=self.padding
-                )
-
-            if axis == 0:
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [res.shape[1]])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [res.shape[1]] + ishape[axis + 1 :]
-                    )
-            else:
-                if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(res, ishape[0:axis] + [res.shape[1]])
-                else:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [res.shape[1]] + ishape[axis + 1 :]
-                    )
+                res = self.backend.conv1d(tim, self.ww_SmoothT[1])
 
-            return self.backend.bk_reshape(res, in_image.shape)
+            return self.backend.bk_reshape(res, ishape)
 
         else:
 
             ishape = list(image.shape)
-
+            """
             if cell_ids is not None:
                 if cell_ids.shape[0] not in self.padding_smooth:
                     import healpix_convolution as hc
@@ -2726,15 +2516,16 @@ class FoCUS:
                             padded_data = padding.apply(image[l, :, k2], is_torch=True)
                             res[l, :, k2] = kernel.matmul(padded_data)
                 return res
-
-            nside = int(np.sqrt(image.shape[axis] // 12))
+            """
+            if nside is None:
+                nside = int(np.sqrt(image.shape[-1] // 12))
 
             if self.Idx_Neighbours[nside] is None:
 
                 if self.InitWave is None:
-                    wr, wi, ws, widx = self.init_index(nside)
+                    wr, wi, ws, widx = self.init_index(nside, cell_ids=cell_ids)
                 else:
-                    wr, wi, ws, widx = self.InitWave(self, nside)
+                    wr, wi, ws, widx = self.InitWave(self, nside, cell_ids=cell_ids)
 
                 self.Idx_Neighbours[nside] = 1
                 self.ww_Real[nside] = wr
@@ -2744,92 +2535,24 @@ class FoCUS:
             l_w_smooth = self.w_smooth[nside]
 
             odata = 1
-            for k in range(axis + 1, len(ishape)):
+            for k in range(0, len(ishape) - 1):
                 odata = odata * ishape[k]
 
-            if axis == 0:
-                tim = self.backend.bk_reshape(image, [12 * nside**2, odata])
-                if tim.dtype == self.all_cbk_type:
-                    rr = self.backend.bk_sparse_dense_matmul(
-                        l_w_smooth, self.backend.bk_real(tim)
-                    )
-                    ri = self.backend.bk_sparse_dense_matmul(
-                        l_w_smooth, self.backend.bk_imag(tim)
-                    )
-                    res = self.backend.bk_complex(rr, ri)
-                else:
-                    res = self.backend.bk_sparse_dense_matmul(l_w_smooth, tim)
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [12 * nside**2])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [12 * nside**2] + ishape[axis + 1 :]
-                    )
-
-            if axis > 0:
-                ndata = ishape[0]
-                for k in range(1, axis):
-                    ndata = ndata * ishape[k]
-                tim = self.backend.bk_reshape(image, [ndata, 12 * nside**2, odata])
-                if tim.dtype == self.all_cbk_type:
-                    rr = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_w_smooth, self.backend.bk_real(tim[0])
-                        ),
-                        [1, 12 * nside**2, odata],
-                    )
-                    ri = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(
-                            l_w_smooth, self.backend.bk_imag(tim[0])
-                        ),
-                        [1, 12 * nside**2, odata],
-                    )
-                    res = self.backend.bk_complex(rr, ri)
-                else:
-                    res = self.backend.bk_reshape(
-                        self.backend.bk_sparse_dense_matmul(l_w_smooth, tim[0]),
-                        [1, 12 * nside**2, odata],
-                    )
-
-                for k in range(1, ndata):
-                    if tim.dtype == self.all_cbk_type:
-                        rr = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_w_smooth, self.backend.bk_real(tim[k])
-                            ),
-                            [1, 12 * nside**2, odata],
-                        )
-                        ri = self.backend.bk_reshape(
-                            self.backend.bk_sparse_dense_matmul(
-                                l_w_smooth, self.backend.bk_imag(tim[k])
-                            ),
-                            [1, 12 * nside**2, odata],
-                        )
-                        res = self.backend.bk_concat(
-                            [res, self.backend.bk_complex(rr, ri)], 0
-                        )
-                    else:
-                        res = self.backend.bk_concat(
-                            [
-                                res,
-                                self.backend.bk_reshape(
-                                    self.backend.bk_sparse_dense_matmul(
-                                        l_w_smooth, tim[k]
-                                    ),
-                                    [1, 12 * nside**2, odata],
-                                ),
-                            ],
-                            0,
-                        )
-
-                if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [12 * nside**2]
-                    )
-                else:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [12 * nside**2] + ishape[axis + 1 :]
-                    )
+            tim = self.backend.bk_reshape(image, [odata, ishape[-1]])
+            if tim.dtype == self.all_cbk_type:
+                rr = self.backend.bk_sparse_dense_matmul(
+                    self.backend.bk_real(tim), l_w_smooth
+                )
+                ri = self.backend.bk_sparse_dense_matmul(
+                    self.backend.bk_imag(tim), l_w_smooth
+                )
+                res = self.backend.bk_complex(rr, ri)
+            else:
+                res = self.backend.bk_sparse_dense_matmul(tim, l_w_smooth)
+            if len(ishape) == 1:
+                return self.backend.bk_reshape(res, [ishape[-1]])
+            else:
+                return self.backend.bk_reshape(res, ishape[0:-1] + [ishape[-1]])
 
         return res