foscat 2025.7.2__py3-none-any.whl → 2025.8.3__py3-none-any.whl

foscat/FoCUS.py

@@ -3,9 +3,10 @@ import sys
 
 import healpy as hp
 import numpy as np
+import foscat.HealSpline as HS
 from scipy.interpolate import griddata
 
-TMPFILE_VERSION = "V6_0"
+TMPFILE_VERSION = "V7_0"
 
 
 class FoCUS:
@@ -35,7 +36,7 @@ class FoCUS:
             mpi_rank=0
     ):
 
-        self.__version__ = "2025.07.2"
+        self.__version__ = "2025.08.3"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -637,11 +638,13 @@ class FoCUS:
         return image.numpy()[self.ring2nest[lout]]
 
     # --------------------------------------------------------
-    def ud_grade(self, im, j, axis=0):
+    def ud_grade(self, im, j, axis=0, cell_ids=None, nside=None):
         rim = im
         for k in range(j):
             # rim = self.smooth(rim, axis=axis)
-            rim = self.ud_grade_2(rim, axis=axis)
+            rim = self.ud_grade_2(rim, axis=axis,
+                                  cell_ids=cell_ids,
+                                  nside=nside)
         return rim
 
     # --------------------------------------------------------
@@ -654,55 +657,39 @@ class FoCUS:
                     print("Use of 2D scat with data that has less than 2D")
                 return None, None
 
-            npix = im.shape[axis]
-            npiy = im.shape[axis + 1]
-            odata = 1
-            if len(ishape) > axis + 2:
-                for k in range(axis + 2, len(ishape)):
-                    odata = odata * ishape[k]
+            npix = im.shape[-2]
+            npiy = im.shape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(im.shape)-2):
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
-                self.backend.bk_cast(im), [ndata, npix, npiy, odata]
+                self.backend.bk_cast(im), [ndata, npix, npiy, 1]
             )
             tim = self.backend.bk_reshape(
                 tim[:, 0 : 2 * (npix // 2), 0 : 2 * (npiy // 2), :],
-                [ndata, npix // 2, 2, npiy // 2, 2, odata],
+                [ndata, npix // 2, 2, npiy // 2, 2, 1],
             )
 
             res = self.backend.bk_reduce_sum(self.backend.bk_reduce_sum(tim, 4), 2) / 4
 
-            if axis == 0:
-                if len(ishape) == 2:
-                    return self.backend.bk_reshape(res, [npix // 2, npiy // 2]), None
-                else:
-                    return (
-                        self.backend.bk_reshape(
-                            res, [npix // 2, npiy // 2] + ishape[axis + 2 :]
-                        ),
-                        None,
-                    )
+            if len(ishape) == 2:
+                return (
+                    self.backend.bk_reshape(
+                        res, [npix // 2, npiy // 2]
+                    ),
+                    None,
+                )
             else:
-                if len(ishape) == axis + 2:
-                    return (
-                        self.backend.bk_reshape(
-                            res, ishape[0:axis] + [npix // 2, npiy // 2]
-                        ),
-                        None,
-                    )
-                else:
-                    return (
-                        self.backend.bk_reshape(
-                            res,
-                            ishape[0:axis]
-                            + [npix // 2, npiy // 2]
-                            + ishape[axis + 2 :],
-                        ),
-                        None,
-                    )
+                return (
+                    self.backend.bk_reshape(
+                        res,
+                        ishape[0:-2]
+                        + [npix // 2, npiy // 2],
+                    ),
+                    None,
+                )
 
             return self.backend.bk_reshape(res, [npix // 2, npiy // 2]), None
         elif self.use_1D:
@@ -733,11 +720,16 @@ class FoCUS:
                 ),None
 
     # --------------------------------------------------------
-    def up_grade(self, im, nout, axis=-1, nouty=None):
+    def up_grade(self, im, nout,
+                 axis=-1,
+                 nouty=None,
+                 cell_ids=None,
+                 o_cell_ids=None,
+                 nside=None):
 
         ishape = list(im.shape)
         if self.use_2D:
-            if len(ishape) < axis + 2:
+            if len(ishape) < 2:
                 if not self.silent:
                     print("Use of 2D scat with data that has less than 2D")
                 return None
@@ -745,39 +737,28 @@ class FoCUS:
             if nouty is None:
                 nouty = nout
 
-            if ishape[axis] == nout and ishape[axis + 1] == nouty:
+            if ishape[-2] == nout and ishape[-1] == nouty:
                 return im
 
-            npix = im.shape[axis]
-            npiy = im.shape[axis + 1]
-            odata = 1
+            npix = im.shape[-2]
+            npiy = im.shape[-1]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(im.shape)-2):
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
-                self.backend.bk_cast(im), [ndata, npix, npiy, odata]
+                self.backend.bk_cast(im), [ndata, npix, npiy,1]
             )
 
             res = self.backend.bk_resize_image(tim, [nout, nouty])
 
-            if axis == 0:
-                if len(ishape) == 2:
-                    return self.backend.bk_reshape(res, [nout, nouty])
-                else:
-                    return self.backend.bk_reshape(
-                        res, [nout, nouty] + ishape[axis + 2 :]
-                    )
+            if len(ishape) == 2:
+                return self.backend.bk_reshape(res, [nout, nouty])
             else:
-                if len(ishape) == axis + 2:
-                    return self.backend.bk_reshape(res, ishape[0:axis] + [nout, nouty])
-                else:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [nout, nouty] 
-                    )
-
-            return self.backend.bk_reshape(res, [nout, nouty])
+                return self.backend.bk_reshape(
+                        res, ishape[0:-2] + [nout, nouty] 
+                )
 
         elif self.use_1D:
             if len(ishape) < axis + 1:
@@ -790,13 +771,11 @@ class FoCUS:
 
             npix = im.shape[axis]
             odata = 1
-            if len(ishape) > axis + 1:
-                for k in range(axis + 1, len(ishape)):
-                    odata = odata * ishape[k]
-
+            
             ndata = 1
-            for k in range(axis):
-                ndata = ndata * ishape[k]
+            if len(ishape)>1:
+                for k in range(len(ishape)-1):
+                    ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
                 self.backend.bk_cast(im), [ndata, npix, odata]
@@ -819,54 +798,107 @@ class FoCUS:
                     self.backend.bk_concat([res1, res2], -2),
                     [ndata, tim.shape[1] * 2, odata],
                 )
+            return self.backend.bk_reshape(tim, ishape[0:-1] + [nout])
 
-            if axis == 0:
-                if len(ishape) == 1:
-                    return self.backend.bk_reshape(tim, [nout])
-                else:
-                    return self.backend.bk_reshape(tim, [nout] + ishape[axis + 1 :])
+        else:
+            if nside is None:
+                lout = int(np.sqrt(im.shape[-1] // 12))
             else:
-                if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(tim, ishape[0:axis] + [nout])
+                lout = nside
+                
+            if (lout,nout) not in self.pix_interp_val:
+                if not self.silent:
+                    print("compute lout nout", lout, nout)
+                if cell_ids is None:
+                    o_cell_ids=np.arange(12 * nout**2, dtype="int")
+                    i_npix=12*lout**2
+                
+                    #level=int(np.log2(lout)) # nside=128
+
+                    #sp = HS.heal_spline(level,gamma=2.0)
+
+                    th, ph = hp.pix2ang(
+                        nout, o_cell_ids, nest=True
+                    )
+                    
+                    all_idx,www=hp.get_interp_weights(lout,th,ph,nest=True)
+                                          
+                    #www,all_idx,hidx=sp.ang2weigths(th,ph,nest=True)
+
+                    w=www.T
+                    p=all_idx.T
+                    
+                    w=w.flatten()
+                    p=p.flatten()
+                    
+                    indice = np.zeros([o_cell_ids.shape[0] * 4, 2], dtype="int")
+                    indice[:, 1] = np.repeat(np.arange(o_cell_ids.shape[0]), 4)
+                    indice[:, 0] = p
+                    
+                    self.pix_interp_val[(lout,nout)] = 1
+                    self.weight_interp_val[(lout,nout)] = self.backend.bk_SparseTensor(
+                        self.backend.bk_constant(indice),
+                        self.backend.bk_constant(self.backend.bk_cast(w)),
+                        dense_shape=[i_npix,o_cell_ids.shape[0]],
+                    )
+
                 else:
-                    return self.backend.bk_reshape(
-                        tim, ishape[0:axis] + [nout] + ishape[axis + 1 :]
+                    ratio=(nout//lout)**2
+                    if o_cell_ids is None:
+                        o_cell_ids=np.tile(cell_ids,ratio)*ratio+np.repeat(np.arange(ratio),cell_ids.shape[0])
+                    i_npix=cell_ids.shape[0]
+
+                    #level=int(np.log2(lout)) # nside=128
+
+                    #sp = HS.heal_spline(level,gamma=2.0)
+
+                    th, ph = hp.pix2ang(
+                        nout, o_cell_ids, nest=True
                     )
 
-            return self.backend.bk_reshape(tim, [nout])
+                    all_idx,www=hp.get_interp_weights(lout,th,ph,nest=True)
+                    #www,all_idx,hidx=sp.ang2weigths(th,ph,nest=True)
 
-        else:
+                    hidx,inv_idx = np.unique(all_idx,
+                                    return_inverse=True)
+                    all_idx = inv_idx
+                    
+                    sorter = np.argsort(hidx)
+                    
+                    index=sorter[np.searchsorted(hidx,
+                                                 cell_ids,
+                                                 sorter=sorter)]
+                    
+                    mask        = -np.ones([hidx.shape[0]])
+                    
+                    mask[index] = np.arange(index.shape[0],dtype='int')
 
-            lout = int(np.sqrt(im.shape[-1] // 12))
+                    all_idx=mask[all_idx]
 
-            if (lout,nout) not in self.pix_interp_val:
-                if not self.silent:
-                    print("compute lout nout", lout, nout)
-                th, ph = hp.pix2ang(
-                    nout, np.arange(12 * nout**2, dtype="int"), nest=True
-                )
-                p, w = hp.get_interp_weights(lout, th, ph, nest=True)
-                del th
-                del ph
-
-                indice = np.zeros([12 * nout * nout * 4, 2], dtype="int")
-                p = p.T
-                w = w.T
-                t = np.argsort(
-                    p, 1
-                ).flatten()  # to make oder indices for sparsematrix computation
-                t = t + np.repeat(np.arange(12 * nout * nout) * 4, 4)
-                p = p.flatten()[t]
-                w = w.flatten()[t]
-                indice[:, 1] = np.repeat(np.arange(12 * nout**2), 4)
-                indice[:, 0] = p
-
-                self.pix_interp_val[(lout,nout)] = 1
-                self.weight_interp_val[(lout,nout)] = self.backend.bk_SparseTensor(
-                    self.backend.bk_constant(indice),
-                    self.backend.bk_constant(self.backend.bk_cast(w.flatten())),
-                    dense_shape=[12 * lout**2,12 * nout**2],
-                )
+                    www[all_idx==-1]=0.0
+                    www/=np.sum(www,0)[None,:]
+                    
+                    all_idx[all_idx==-1]=0
+               
+                    w=www.T
+                    p=all_idx.T
+                    
+                    w=w.flatten()
+                    p=p.flatten()
+                    
+                    indice = np.zeros([o_cell_ids.shape[0] * 4, 2], dtype="int")
+                    indice[:, 1] = np.repeat(np.arange(o_cell_ids.shape[0]), 4)
+                    indice[:, 0] = p
+                    
+                    self.pix_interp_val[(lout,nout)] = 1
+                    self.weight_interp_val[(lout,nout)] = self.backend.bk_SparseTensor(
+                        self.backend.bk_constant(indice),
+                        self.backend.bk_constant(self.backend.bk_cast(w)),
+                        dense_shape=[i_npix,o_cell_ids.shape[0]],
+                    )
+                    
+                    del w
+                    del p
 
             if lout == nout:
                 imout = im
@@ -877,7 +909,7 @@ class FoCUS:
                 for k in range(len(ishape)-1):
                     ndata = ndata * ishape[k]
                 tim = self.backend.bk_reshape(
-                    self.backend.bk_cast(im), [ndata, 12 * lout**2]
+                    self.backend.bk_cast(im), [ndata, ishape[-1]]
                 )
                 if tim.dtype == self.all_cbk_type:
                     rr = self.backend.bk_sparse_dense_matmul(
@@ -894,12 +926,12 @@ class FoCUS:
                         tim,
                         self.weight_interp_val[(lout,nout)],
                     )
-                    
+
                 if len(ishape) == 1:
-                    return self.backend.bk_reshape(imout, [12 * nout**2])
+                    return self.backend.bk_reshape(imout, [imout.shape[-1]])
                 else:
                     return self.backend.bk_reshape(
-                        imout, ishape[0:axis]+[12 * nout**2]
+                        imout, ishape[0:-1]+[imout.shape[-1]]
                     )
         return imout
 
@@ -1191,7 +1223,7 @@ class FoCUS:
                                       % (self.TEMPLATE_PATH, l_kernel**2,TMPFILE_VERSION, nside)
                 )
             else:
-                if cell_ids is not None and nside>512:
+                if cell_ids is not None and spin==0:
                     tmp = self.read_index(
                         "%s/XXXX_%s_W%d_%d_%d_PIDX.fst"  # can not work
                         % (
@@ -1204,6 +1236,7 @@ class FoCUS:
                     )
 
                 else:
+                    '''
                     print('LOAD ',"%s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN%d.fst"
                         % (
                             self.TEMPLATE_PATH,
@@ -1212,6 +1245,7 @@ class FoCUS:
                             self.NORIENT,
                             nside,spin  # if cell_ids computes the index
                         ))
+                    '''
                     tmp = self.read_index(
                         "%s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN%d.fst"
                         % (
@@ -1224,20 +1258,21 @@ class FoCUS:
                     )
                         
         except:
-            if cell_ids is not None and nside<=512:
+            if cell_ids is not None and spin!=0:
                 self.init_index(nside, kernel=kernel, spin=spin)
                 
             if not self.use_2D:
-                print('NOT FOUND THEN COMPUTE %s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN%d.fst'
-                    % (
-                        self.TEMPLATE_PATH,
-                        TMPFILE_VERSION,
-                        l_kernel**2,
-                        self.NORIENT,
-                        nside,spin  # if cell_ids computes the index
-                    )
-                      )
                 if spin!=0:
+                    # keep the print here as spin!=0 can be long
+                    print('NOT FOUND THEN COMPUTE %s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN%d.fst'
+                        % (
+                            self.TEMPLATE_PATH,
+                            TMPFILE_VERSION,
+                            l_kernel**2,
+                            self.NORIENT,
+                            nside,spin  # if cell_ids computes the index
+                        )
+                      )
                     try:
                         tmp = self.read_index(
                             "%s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN0.fst"
@@ -1250,6 +1285,7 @@ class FoCUS:
                             )
                         )
                     except:
+                        '''
                         print('NOT FOUND THEN COMPUTE %s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN0.fst'
                               % (
                                   self.TEMPLATE_PATH,
@@ -1259,7 +1295,7 @@ class FoCUS:
                                   nside
                               )
                               )
-                        
+                        '''
                         self.init_index(nside, kernel=kernel, spin=0)
                         
                         tmp = self.read_index(
@@ -1282,6 +1318,7 @@ class FoCUS:
                                         )
                                     )
 
+                    '''
                     nn=self.NORIENT*12*nside**2
                     idxEB=np.concatenate([tmp,tmp,tmp,tmp],0)
                     idxEB[tmp.shape[0]:2*tmp.shape[0],0]+=12*nside**2
@@ -1289,7 +1326,156 @@ class FoCUS:
                     idxEB[2*tmp.shape[0]:,1]+=nn
 
                     tmpEB=np.zeros([tmpw.shape[0]*4],dtype='complex')
+                    '''
+                    import foscat.HOrientedConvol as hs
 
+                    hconvol=hs.HOrientedConvol(nside,3*self.KERNELSZ,cell_ids=cell_ids)
+                    
+                    if cell_ids is None:
+                        l_cell_ids=np.arange(12*nside**2)
+                    else:
+                        l_cell_ids=cell_ids
+                        
+                    nvalid=self.KERNELSZ**2
+                    idxEB=hconvol.idx_nn[:,0:nvalid]
+                    tmpEB=np.zeros([self.NORIENT,4,l_cell_ids.shape[0],nvalid],dtype='complex')       
+                    tmpS=np.zeros([4,l_cell_ids.shape[0],nvalid],dtype='float')           
+                    
+                    idx={}
+                    nn=0
+                    nn2=1
+                    if nside<64:
+                        pp=10
+                    else:
+                        pp=1
+                    while nn2>0:
+                        idx2={}
+                        nn2=0
+                        im=np.zeros([12*nside**2])
+                        for n in range(l_cell_ids.shape[0]):
+                            if im[hconvol.idx_nn[n,0]]==0 and n not in idx:
+                                im[hconvol.idx_nn[n,:]]=1.0
+                                idx[hconvol.idx_nn[n,0]]=1.0
+                                idx2[hconvol.idx_nn[n,0]]=1.0
+                                nn+=1
+                                nn2+=1
+                        im=np.zeros([12*nside**2])
+                        for k in idx2:
+                            im[k]=1.0
+                        r=self.convol(im)
+                        for k in range(self.NORIENT):
+                            ralm=hp.map2alm(hp.reorder(r[k].cpu().numpy().real,n2r=True))[None,:]
+                            ialm=hp.map2alm(hp.reorder(r[k].cpu().numpy().imag,n2r=True))[None,:]
+                            
+                            alm=np.concatenate([ralm,0*ralm,0*ralm],0)
+                            rqe,rue,rie=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                            alm=np.concatenate([ialm,0*ialm,0*ialm],0)
+                            iqe,iue,iie=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                            
+                            alm=np.concatenate([0*ralm,ralm,0*ralm],0)
+                            rqb,rub,rib=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                            alm=np.concatenate([0*ialm,ialm,0*ialm],0)
+                            iqb,iub,iib=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                            
+                            rqe=hp.reorder(rqe,r2n=True)
+                            rue=hp.reorder(rue,r2n=True)
+                            rqb=hp.reorder(rqb,r2n=True)
+                            rub=hp.reorder(rub,r2n=True)
+                            
+                            iqe=hp.reorder(iqe,r2n=True)
+                            iue=hp.reorder(iue,r2n=True)
+                            iqb=hp.reorder(iqb,r2n=True)
+                            iub=hp.reorder(iub,r2n=True)
+                            
+                            for l in idx2:
+                                tmpEB[k,0,l]=rqe[idxEB[l,:]]+1J*iqe[idxEB[l,:]]
+                                tmpEB[k,1,l]=rue[idxEB[l,:]]+1J*iue[idxEB[l,:]]
+                                tmpEB[k,2,l]=rqb[idxEB[l,:]]+1J*iqb[idxEB[l,:]]
+                                tmpEB[k,3,l]=rub[idxEB[l,:]]+1J*iub[idxEB[l,:]]
+                                
+                        r=self.smooth(im)
+                        
+                        ralm=hp.map2alm(hp.reorder(r.cpu().numpy(),n2r=True))[None,:]
+                            
+                        alm=np.concatenate([ralm,0*ralm,0*ralm],0)
+                        rqe,rue,rie=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                            
+                        alm=np.concatenate([0*ralm,ralm,0*ralm],0)
+                        rqb,rub,rib=hp.alm2map_spin(alm,nside,spin,3*nside-1)
+                        
+                        rqe=hp.reorder(rqe,r2n=True)
+                        rue=hp.reorder(rue,r2n=True)
+                        rqb=hp.reorder(rqb,r2n=True)
+                        rub=hp.reorder(rub,r2n=True)
+                            
+                        for l in idx2:
+                            tmpS[0,l,:]=rqe[idxEB[l,:]]
+                            tmpS[1,l,:]=rue[idxEB[l,:]]
+                            tmpS[2,l,:]=rqb[idxEB[l,:]]
+                            tmpS[3,l,:]=rub[idxEB[l,:]]
+                        if 100*nn/(l_cell_ids.shape[0])>pp:
+                            if nside<64:
+                                pp+=10
+                            else:
+                                pp+=1
+                            print('%.2f%% Done'%(100*nn/(l_cell_ids.shape[0])))
+                        
+                    wav=tmpEB.flatten()
+                    wwav=tmpS.flatten()
+                    ndata=l_cell_ids.shape[0]*nvalid
+                    indice_1_1=np.tile(idxEB.flatten(),4*self.NORIENT)
+                    for k in range(self.NORIENT):
+                        indice_1_1[(4*k+1)*ndata:(4*k+2)*ndata]+=l_cell_ids.shape[0]
+                        indice_1_1[(4*k+3)*ndata:(4*k+4)*ndata]+=l_cell_ids.shape[0]
+                    
+                    indice_1_0=np.tile(np.tile(np.repeat(np.arange(l_cell_ids.shape[0]),nvalid),4),self.NORIENT)
+                    for k in range(self.NORIENT):
+                        indice_1_0[(4*k+2)*ndata:(4*k+4)*ndata]+=self.NORIENT*l_cell_ids.shape[0]
+                        indice_1_0[(4*k)*ndata:(4*k+4)*ndata]+=k*l_cell_ids.shape[0]
+                    '''
+                    import matplotlib.pyplot as plt
+                    plt.figure()
+                    plt.subplot(2,2,1)
+                    plt.plot(indice_1_0)
+                    plt.subplot(2,2,2)
+                    plt.plot(indice_1_1)
+                    plt.subplot(2,2,3)
+                    plt.plot(wav.real)
+                    plt.subplot(2,2,4)
+                    plt.plot(abs(wav))
+                    
+                    iarg=np.argsort(indice_1_0)
+                    indice_1_1=indice_1_1[iarg]
+                    indice_1_0=indice_1_0[iarg]
+                    wav=wav[iarg]
+                    '''
+                    
+                    indice=np.concatenate([indice_1_1[:,None],indice_1_0[:,None]],1)
+                    
+                    indice_2_1=np.tile(idxEB.flatten(),4)
+                    indice_2_1[ndata:2*ndata]+=l_cell_ids.shape[0]
+                    indice_2_1[3*ndata:4*ndata]+=l_cell_ids.shape[0]
+                    indice_2_0=np.tile(np.repeat(np.arange(l_cell_ids.shape[0]),nvalid),4)
+                    indice_2_0[2*ndata:]+=l_cell_ids.shape[0]
+                    '''
+                    plt.figure()
+                    plt.subplot(2,2,1)
+                    plt.plot(indice_2_0)
+                    plt.subplot(2,2,2)
+                    plt.plot(indice_2_1)
+                    plt.subplot(2,2,3)
+                    plt.plot(wav.real)
+                    plt.subplot(2,2,4)
+                    plt.plot(wwav)
+                    
+                    iarg=np.argsort(indice_2_0)
+                    indice_2_1=indice_2_1[iarg]
+                    indice_2_0=indice_2_0[iarg]
+                    wwav=wwav[iarg]
+                    '''
+                    indice2=np.concatenate([indice_2_1[:,None],indice_2_0[:,None]],1)
+                    
+                    '''
                     for k in range(self.NORIENT*12*nside**2):
                         if k%(nside**2)==0:
                             print('Init index 1/2 spin=%d Please wait %d done against %d nside=%d kernel=%d'%(spin,k//(nside**2),
@@ -1316,7 +1502,7 @@ class FoCUS:
                         tmpEB[idx+2*tmp.shape[0]]=hp.reorder(i,r2n=True)[tmp[idx,0]]+1J*hp.reorder(i2,r2n=True)[tmp[idx,0]]
                         tmpEB[idx+3*tmp.shape[0]]=hp.reorder(q,r2n=True)[tmp[idx,0]]+1J*hp.reorder(q2,r2n=True)[tmp[idx,0]]
 
-                    
+                    '''
                     self.save_index("%s/FOSCAT_%s_W%d_%d_%d_PIDX-SPIN%d.fst"% (self.TEMPLATE_PATH,
                                                                        self.TMPFILE_VERSION,
                                                                        self.KERNELSZ**2,
@@ -1324,7 +1510,7 @@ class FoCUS:
                                                                        nside,
                                                                        spin
                                                                        ),
-                                    idxEB
+                                    indice
                                     )
                     self.save_index("%s/FOSCAT_%s_W%d_%d_%d_WAVE-SPIN%d.fst"% (self.TEMPLATE_PATH,
                                                                        self.TMPFILE_VERSION,
@@ -1333,9 +1519,9 @@ class FoCUS:
                                                                        nside,
                                                                        spin,
                                                                        ),
-                                    tmpEB
+                                    wav
                                     )
-                    
+                    '''
                     tmp = self.read_index(
                             "%s/FOSCAT_%s_W%d_%d_%d_PIDX2-SPIN0.fst"
                             % (
@@ -1377,7 +1563,7 @@ class FoCUS:
                         tmpEB[idx+2*tmp.shape[0]]=hp.reorder(i,r2n=True)[tmp[idx,0]]
                         tmpEB[idx+3*tmp.shape[0]]=hp.reorder(q,r2n=True)[tmp[idx,0]]
 
-                    
+                    '''
                     self.save_index("%s/FOSCAT_%s_W%d_%d_%d_PIDX2-SPIN%d.fst"% (self.TEMPLATE_PATH,
                                                                        self.TMPFILE_VERSION,
                                                                        self.KERNELSZ**2,
@@ -1385,7 +1571,7 @@ class FoCUS:
                                                                        nside,
                                                                        spin
                                                                        ),
-                                    idxEB
+                                    indice2
                                     )
                     self.save_index("%s/FOSCAT_%s_W%d_%d_%d_SMOO-SPIN%d.fst"% (self.TEMPLATE_PATH,
                                                                        self.TMPFILE_VERSION,
@@ -1394,10 +1580,11 @@ class FoCUS:
                                                                        nside,
                                                                        spin,
                                                                        ),
-                                    tmpEB
+                                    wwav
                                     )
+                    
                 else:
-
+                    '''
                     if l_kernel == 5:
                         pw = 0.5
                         pw2 = 0.5
@@ -1412,8 +1599,20 @@ class FoCUS:
                         pw = 0.5
                         pw2 = 0.25
                         threshold = 4e-5
+                    '''
+                    import foscat.HOrientedConvol as hs
 
-                    if cell_ids is not None and nside>512:
+                    hconvol=hs.HOrientedConvol(nside,l_kernel,cell_ids=cell_ids)
+                    
+                    orientations=np.pi*np.arange(self.NORIENT)/self.NORIENT
+                    
+                    wav,indice,wwav,indice2=hconvol.make_wavelet_matrix(orientations,
+                                                                        polar=True,
+                                                                        return_index=True,
+                                                                        return_smooth=True)
+                    
+                    '''
+                    if cell_ids is not None and nside>256:
                         if not isinstance(cell_ids, np.ndarray):
                             cell_ids = self.backend.to_numpy(cell_ids)
                         th, ph = hp.pix2ang(nside, cell_ids, nest=True)
@@ -1537,8 +1736,7 @@ class FoCUS:
                     wav = wav[:iv]
                     indice2 = indice2[:iv2, :]
                     wwav = wwav[:iv2]
-                    if not self.silent:
-                        print("Kernel Size ", iv / (self.NORIENT * 12 * nside * nside))
+                    '''
 
                     if cell_ids is None:
                         if not self.silent:
@@ -1613,7 +1811,7 @@ class FoCUS:
                             )
                         return None
 
-        if cell_ids is None or nside<=512:
+        if cell_ids is None or spin!=0:
             self.barrier()
             if self.use_2D:
                 tmp = self.read_index(
@@ -1719,6 +1917,20 @@ class FoCUS:
                 tmp2[lidx,0]=0
                 tmp2[:,1]+=i_id*lcell_ids.shape[0]
                 tmp2[:,0]+=i_id2*lcell_ids.shape[0]
+
+                #add normalisation
+                ww=np.bincount(tmp2[:,1],weights=ws)
+                ws/=ww[tmp2[:,1]]
+
+                wh=np.bincount(tmp[:,1])
+                ww=np.bincount(tmp[:,1],weights=wr)
+                wr-=(ww/wh)[tmp[:,1]]
+                ww=np.bincount(tmp[:,1],weights=wi)
+                wi-=(ww/wh)[tmp[:,1]]
+                
+                ww=np.bincount(tmp[:,1],weights=np.sqrt(wr*wr+wi*wi))
+                wr/=ww[tmp[:,1]]
+                wi/=ww[tmp[:,1]]
                 
         else:
             tmp = indice
@@ -2116,8 +2328,8 @@ class FoCUS:
             ichannel = 1
             for i in range(1, len(shape) - 1):
                 ichannel *= shape[i]
-
-            l_x = self.backend.bk_reshape(x, [shape[0], 1, ichannel, shape[-1]])
+            
+            l_x = self.backend.bk_reshape(x, [shape[0], 1, ichannel,shape[-1]])
 
             if self.padding == "VALID":
                 oshape = [k for k in shape]
@@ -2168,13 +2380,8 @@ class FoCUS:
             res = v1 / vh
 
             oshape = [x.shape[0]] + [mask.shape[0]]
-            if axis > 0:
-                oshape = oshape + list(x.shape[1:axis])
-
-            if len(x.shape[axis:-2]) > 0:
-                oshape = oshape + list(x.shape[axis:-2])
-            else:
-                oshape = oshape + [1]
+            if len(x.shape)>1:
+                oshape = oshape + list(x.shape[1:-2])
 
             if calc_var:
                 if self.backend.bk_is_complex(vtmp):
@@ -2204,9 +2411,9 @@ class FoCUS:
         elif self.use_1D:
             mtmp = l_mask
             vtmp = l_x
-            v1 = self.backend.bk_reduce_sum(l_mask[1,:,...,:] * vtmp, axis=-1)
-            v2 = self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=-1)
-            vh = self.backend.bk_reduce_sum(mtmp, axis=-1)
+            v1 = self.backend.bk_reduce_sum(l_mask * vtmp, axis=-1)
+            v2 = self.backend.bk_reduce_sum(l_mask * vtmp * vtmp, axis=-1)
+            vh = self.backend.bk_reduce_sum(l_mask , axis=-1)
 
             res = v1 / vh
 
@@ -2576,15 +2783,16 @@ class FoCUS:
                     print("Use of 2D scat with data that has less than 2D")
                 return None
 
-            npix = ishape[axis]
-            npiy = ishape[axis + 1]
+            npix = ishape[-2]
+            npiy = ishape[-1]
+            
             odata = 1
-            if len(ishape) > axis + 2:
-                for k in range(axis + 2, len(ishape)):
+            if len(ishape) > 1:
+                for k in range(len(ishape)-2):
                     odata = odata * ishape[k]
 
             ndata = 1
-            for k in range(axis):
+            for k in range(len(ishape)-2):
                 ndata = ndata * ishape[k]
 
             tim = self.backend.bk_reshape(
@@ -2605,7 +2813,7 @@ class FoCUS:
             ishape = list(in_image.shape)
 
             npix = ishape[-1]
-
+            
             ndata = 1
             for k in range(len(ishape) - 1):
                 ndata = ndata * ishape[k]
@@ -2618,7 +2826,7 @@ class FoCUS:
                 res = self.backend.bk_complex(rr, ii)
             else:
                 res = self.backend.conv1d(tim, self.ww_SmoothT[1])
-
+                
             return self.backend.bk_reshape(res, ishape)
 
         else:
@@ -2646,16 +2854,30 @@ class FoCUS:
                 odata = odata * ishape[k]
 
             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)
+            if spin==0:
+                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)
             else:
-                res = self.backend.bk_sparse_dense_matmul(tim, l_w_smooth)
+                tim=self.backend.bk_reshape(tim,[odata//2,2*tim.shape[-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: