foscat 2025.8.4__py3-none-any.whl → 2025.9.3__py3-none-any.whl

foscat/HOrientedConvol.py

@@ -3,16 +3,37 @@ import matplotlib.pyplot as plt
 import healpy as hp
 from scipy.sparse import csr_array
 import torch
+import foscat.scat_cov as sc
 from scipy.spatial import cKDTree
 
 class HOrientedConvol:
-    def __init__(self,nside,KERNELSZ,cell_ids=None,nest=True):
+    def __init__(self,
+                 nside,
+                 KERNELSZ,
+                 cell_ids=None,
+                 nest=True,
+                 device='cuda',
+                 dtype='float64',
+                 polar=False,
+                 gamma=1.0,
+                 allow_extrapolation=True,
+                 no_cell_ids=False,
+                 ):
+
+
+        if dtype=='float64':
+            self.dtype=torch.float64
+        else:
+            self.dtype=torch.float32
         
         if KERNELSZ % 2 == 0:
             raise ValueError(f"N must be odd so that coordinates are integers from -K..K; got N={KERNELSZ}.")
 
         self.local_test=False
-        
+
+        if no_cell_ids==True:
+            cell_ids=np.arange(10)
+            
         if cell_ids is None:
             self.cell_ids=np.arange(12*nside**2)
             
@@ -28,37 +49,84 @@ class HOrientedConvol:
                 self.cell_ids=cell_ids
                 
             self.local_test=True
-            
-            idx_nn = self.knn_healpix_ckdtree(self.cell_ids, 
-                KERNELSZ*KERNELSZ, 
-                nside,
-                nest=nest,
-            )
 
+            if self.cell_ids.ndim==1:
+                idx_nn = self.knn_healpix_ckdtree(self.cell_ids, 
+                                                  KERNELSZ*KERNELSZ, 
+                                                  nside,
+                                                  nest=nest,
+                                                  )
+            else:
+                idx_nn = []
+                for k in range(self.cell_ids.shape[0]):
+                    idx_nn.append(self.knn_healpix_ckdtree(self.cell_ids[k], 
+                                                           KERNELSZ*KERNELSZ, 
+                                                           nside,
+                                                           nest=nest,
+                                                           ))
+                idx_nn=np.stack(idx_nn,0)
+        
+        if self.cell_ids.ndim==1:
+            mat_pt=self.rotation_matrices_from_healpix(nside,self.cell_ids,nest=nest)
             
-        mat_pt=self.rotation_matrices_from_healpix(nside,self.cell_ids,nest=nest)
+            if self.local_test:
+                t,p = hp.pix2ang(nside,self.cell_ids[idx_nn],nest=True)
+            else:
+                t,p = hp.pix2ang(nside,idx_nn,nest=True)
+
+            self.t=t[:,0]
+            self.p=p[:,0]
+            vec_orig=hp.ang2vec(t,p)
 
-        if self.local_test:
-            t,p = hp.pix2ang(nside,self.cell_ids[idx_nn],nest=True)
+            self.vec_rot = np.einsum('mki,ijk->kmj', vec_orig,mat_pt)
+
+            '''
+            if self.local_test:
+            idx_nn=self.remap_by_first_column(idx_nn)
+            '''
+            
+            del mat_pt
+            del vec_orig
         else:
-            t,p = hp.pix2ang(nside,idx_nn,nest=True)
             
-        vec_orig=hp.ang2vec(t,p)
+            t,p,vec_rot = [],[],[]
+            
+            for k in range(self.cell_ids.shape[0]):
+                mat_pt=self.rotation_matrices_from_healpix(nside,self.cell_ids[k],nest=nest)
+                
+                lt,lp = hp.pix2ang(nside,self.cell_ids[k,idx_nn[k]],nest=True)
+                
+                vec_orig=hp.ang2vec(lt,lp)
+                
+                l_vec_rot=np.einsum('mki,ijk->kmj', vec_orig,mat_pt)
+                vec_rot.append(l_vec_rot)
+                
+                del vec_orig
+                del mat_pt
+                
+                t.append(lt[:,0])
+                p.append(lp[:,0])
 
-        self.vec_rot = np.einsum('mki,ijk->kmj', vec_orig,mat_pt)
+                    
+            self.t=np.stack(t,0)
+            self.p=np.stack(p,0)
+            self.vec_rot=np.stack(vec_rot,0)
 
-        '''
-        if self.local_test:
-            idx_nn=self.remap_by_first_column(idx_nn)
-        '''
+            del t
+            del p
+            del vec_rot
+                
+        self.polar=polar
+        self.gamma=gamma
+        self.device=device
+        self.allow_extrapolation=allow_extrapolation
+        self.w_idx=None
         
-        del mat_pt
-        del vec_orig
-        self.t=t[:,0]
-        self.p=p[:,0]
         self.idx_nn=idx_nn
         self.nside=nside
         self.KERNELSZ=KERNELSZ
+        self.nest=nest
+        self.f=None
 
     def remap_by_first_column(self,idx: np.ndarray) -> np.ndarray:
         """
@@ -290,25 +358,160 @@ class HOrientedConvol:
         
         return csr_array((w, (indice_1_0, indice_1_1)), shape=(12*self.nside**2, 12*self.nside**2*NORIENT))
 
-    
-    def make_idx_weights(self,polar=False,gamma=1.0,device='cuda',allow_extrapolation=True):
+    def make_idx_weights_from_cell_ids(self,
+                                       i_cell_ids,
+                                       polar=False,
+                                       gamma=1.0,
+                                       device='cuda',
+                                       allow_extrapolation=True):
+        """
+        Accept 1D (Npix,) or 2D (B, Npix) cell_ids and return
+        tensors batched sur la 1ère dim (B, ...).
+        """
+        # → cast numpy
+        if torch.is_tensor(i_cell_ids):
+            cid = i_cell_ids.detach().cpu().numpy()
+        else:
+            cid = np.asarray(i_cell_ids)
+
+        # --- 1D: pas de boucle, on calcule une fois, puis on ajoute l'axe batch
+        if cid.ndim == 1:
+            l_idx_nn, l_w_idx, l_w_w = self.make_idx_weights_from_one_cell_ids(
+                cid, polar=polar, gamma=gamma, device=device,
+                allow_extrapolation=allow_extrapolation
+            )
+            idx_nn = torch.as_tensor(l_idx_nn, device=device, dtype=torch.long)[None, ...]  # (1, Npix, P)
+            w_idx  = torch.as_tensor(l_w_idx,  device=device, dtype=torch.long)[None, ...]  # (1, Npix, S, P) ou (1, Npix, P)
+            w_w    = torch.as_tensor(l_w_w,    device=device, dtype=self.dtype)[None, ...]  # (1, Npix, S, P) ou (1, Npix, P)
+            return idx_nn, w_idx, w_w
+
+        # --- 2D: boucle sur b, empilement en (B, ...)
+        elif cid.ndim == 2:
+            outs = [ self.make_idx_weights_from_one_cell_ids(
+                        cid[k], polar=polar, gamma=gamma, device=device,
+                        allow_extrapolation=allow_extrapolation)
+                     for k in range(cid.shape[0]) ]
+            idx_nn = torch.as_tensor(np.stack([o[0] for o in outs], axis=0), device=device, dtype=torch.long)
+            w_idx  = torch.as_tensor(np.stack([o[1] for o in outs], axis=0), device=device, dtype=torch.long)
+            w_w    = torch.as_tensor(np.stack([o[2] for o in outs], axis=0), device=device, dtype=self.dtype)
+            return idx_nn, w_idx, w_w
+
+        else:
+            raise ValueError(f"Unsupported cell_ids ndim={cid.ndim}; expected 1 or 2.")
+    '''
+    def make_idx_weights_from_cell_ids(self,i_cell_ids,
+                                       polar=False,
+                                       gamma=1.0,
+                                       device='cuda',
+                                       allow_extrapolation=True):
+        if len(i_cell_ids.shape)<2:
+            cell_ids=i_cell_ids
+            n_cids=1
+        else:
+            cell_ids=i_cell_ids[0]
+            n_cids=i_cell_ids.shape[0]
+            
+        idx_nn,w_idx,w_w    = [],[],[]
+        
+        for k in range(n_cids):
+            cell_ids=i_cell_ids[k]
+            l_idx_nn,l_w_idx,l_w_w = self.make_idx_weights_from_one_cell_ids(cell_ids,
+                                                                       polar=polar,
+                                                                       gamma=gamma,
+                                                                       device=device,
+                                                                       allow_extrapolation=allow_extrapolation)
+            idx_nn.append(l_idx_nn)
+            w_idx.append(l_w_idx)
+            w_w.append(l_w_w)
+            
+        idx_nn = torch.Tensor(np.stack(idx_nn,0)).to(device=device, dtype=torch.long)
+        w_idx  = torch.Tensor(np.stack(w_idx,0)).to(device=device, dtype=torch.long)
+        w_w    = torch.Tensor(np.stack(w_w,0)).to(device=device, dtype=self.dtype)
+        
+        return idx_nn,w_idx,w_w
+    '''
+
+    def make_idx_weights_from_one_cell_ids(self,
+                                           cell_ids,
+                                           polar=False,
+                                           gamma=1.0,
+                                           device='cuda',
+                                           allow_extrapolation=True):
+
+        idx_nn = self.knn_healpix_ckdtree(cell_ids, 
+                                          self.KERNELSZ*self.KERNELSZ, 
+                                          self.nside,
+                                          nest=self.nest,
+                                          )
+        
+        mat_pt=self.rotation_matrices_from_healpix(self.nside,cell_ids,nest=self.nest)
+
+        t,p = hp.pix2ang(self.nside,cell_ids[idx_nn],nest=self.nest)
+            
+        vec_orig=hp.ang2vec(t,p)
+
+        vec_rot = np.einsum('mki,ijk->kmj', vec_orig,mat_pt)
         
-        rotate=2*((self.t<np.pi/2)-0.5)[:,None]
+        del vec_orig
+        del mat_pt
+        
+        rotate=2*((t<np.pi/2)-0.5)[:,None]
         if polar:
-            xx=np.cos(self.p)[:,None]*self.vec_rot[:,:,0]-rotate*np.sin(self.p)[:,None]*self.vec_rot[:,:,1]
-            yy=-np.sin(self.p)[:,None]*self.vec_rot[:,:,0]-rotate*np.cos(self.p)[:,None]*self.vec_rot[:,:,1]
+            xx=np.cos(p)[:,None]*vec_rot[:,:,0]-rotate*np.sin(p)[:,None]*vec_rot[:,:,1]
+            yy=-np.sin(p)[:,None]*vec_rot[:,:,0]-rotate*np.cos(p)[:,None]*vec_rot[:,:,1]
         else:
-            xx=self.vec_rot[:,:,0]
-            yy=self.vec_rot[:,:,1]
+            xx=vec_rot[:,:,0]
+            yy=vec_rot[:,:,1]
+
+        del vec_rot
+        del rotate
+        del t
+        del p
         
-        self.w_idx,self.w_w = self.bilinear_weights_NxN(xx*self.nside*gamma,
-                                                        yy*self.nside*gamma,
-                                                        allow_extrapolation=allow_extrapolation)
+        w_idx,w_w = self.bilinear_weights_NxN(xx*self.nside*gamma,
+                                              yy*self.nside*gamma,
+                                              allow_extrapolation=allow_extrapolation)
+        '''
+        # calib : [Npix, K]
+        calib = np.zeros((w_idx.shape[0], w_idx.shape[2]))
+        # Hypothèses : 
+        # w_idx.shape == (Npix, M, K)  et  w_w.shape == (Npix, M, K)
+        Npix, M, K = w_idx.shape
+        nb_cols = K
+        
+        # 1) Accumulation par "bincount" avec décalage de ligne
+        row_ids = np.arange(Npix, dtype=np.int64)[:, None, None] * nb_cols
+        flat_idx = (row_ids + w_idx).ravel()           # indices dans [0, Npix*9)
+        weights  = w_w.ravel().astype(np.float64)      # ou dtype de ton choix
+        
+        calib = np.bincount(flat_idx, weights, minlength=Npix*nb_cols)\
+                  .reshape(Npix, nb_cols)
+        
+        # 2) Réinjection dans norm_a selon w_idx
+        norm_a = calib[np.arange(Npix)[:, None, None], w_idx]
+
+        w_w /= norm_a
+        w_w  = np.clip(w_w,0.0,1.0)
 
+        w_w[np.isnan(w_w)]=0.0
+        '''
+        #del xx
+        #del yy
+        
+        return idx_nn,w_idx,w_w,xx,yy
+        
+    def make_idx_weights(self,polar=False,gamma=1.0,device='cuda',allow_extrapolation=True,return_index=False):
+        
+        idx_nn,w_idx,w_w = self.make_idx_weights_from_one_cell_ids(self.cell_ids,
+                                                                  polar=polar,
+                                                                  gamma=gamma,
+                                                                  device=device,
+                                                                  allow_extrapolation=allow_extrapolation)
+        
         # Ensure types/devices
-        self.idx_nn = torch.Tensor(self.idx_nn).to(device=device, dtype=torch.long)
-        self.w_idx  = torch.Tensor(self.w_idx).to(device=device, dtype=torch.long)
-        self.w_w    = torch.Tensor(self.w_w).to(device=device, dtype=torch.float64)
+        self.idx_nn = torch.Tensor(idx_nn).to(device=device, dtype=torch.long)
+        self.w_idx  = torch.Tensor(w_idx).to(device=device, dtype=torch.long)
+        self.w_w    = torch.Tensor(w_w).to(device=device, dtype=self.dtype)
     
     def _grid_index(self, xi, yi):
         """
@@ -410,108 +613,169 @@ class HOrientedConvol:
         idx = np.stack([i00, i10, i01, i11], axis=1).astype(np.int64)
     
         return idx, w
+    
+    # --- Add inside class HOrientedConvol, just above Convol_torch ---
+    def _convol_single(self, im1: torch.Tensor, ww: torch.Tensor, cell_ids=None, nside=None):
+        """
+        Single-sample path. im1: (1, C_i, Npix_1). Returns (1, C_o, Npix_1).
+        """
+        if not isinstance(im1, torch.Tensor):
+            im1 = torch.as_tensor(im1, device=self.device, dtype=self.dtype)
+        if not isinstance(ww, torch.Tensor):
+            ww  = torch.as_tensor(ww, device=self.device, dtype=self.dtype)
+        assert im1.ndim == 3 and im1.shape[0] == 1, f"expected (1, C_i, Npix), got {tuple(im1.shape)}"
 
-    def Convol_torch(self, im, ww):
+        # Reuse the existing Convol_torch core by faking B=1 shapes.
+        # We call the existing (batched) implementation with B=1.
+        return self.Convol_torch(im1, ww, cell_ids=cell_ids, nside=nside)  # returns (1, C_o, Npix_1)
+
+    # --- Replace the first lines of Convol_torch with a dispatcher ---
+    def Convol_torch(self, im, ww, cell_ids=None, nside=None):
         """
-        Batched KERNELSZxKERNELSZ neighborhood aggregation in pure PyTorch (generalization of the 3x3 case).
-    
-        Parameters
-        ----------
-        im : Tensor, shape (B, C_i, Npix)
-            Input features per pixel for a batch of B samples.
-        ww : Tensor
-            Base mixing weights, indexed along its 'M' dimension by self.w_idx.
-            Supported shapes:
-              (C_i, C_o, M)
-              (C_i, C_o, M, S)
-              (B, C_i, C_o, M)
-              (B, C_i, C_o, M, S)
-    
-        Class members (already tensors; will be aligned to im.device/dtype):
-        -------------------------------------------------------------------
-        self.idx_nn : LongTensor, shape (Npix, P)
-            For each center pixel, the P neighbor indices into the Npix axis of `im`.
-            (P = K*K for a KxK neighborhood.)
-        self.w_idx  : LongTensor, shape (Npix, P) or (Npix, S, P)
-            Indices along the 'M' dimension of ww, per (center[, sector], neighbor).
-        self.w_w    : Tensor,     shape (Npix, P) or (Npix, S, P)
-            Additional scalar weights per neighbor (same layout as w_idx).
-    
-        Returns
-        -------
-        out : Tensor, shape (B, C_o, Npix)
-            Aggregated output per center pixel for each batch sample.
+        Batched KERNELSZxKERNELSZ aggregation.
+
+        Accepts either:
+          - im: Tensor (B, C_i, Npix)  with one shared or per-batch (B,Npix) cell_ids
+          - im: list/tuple of Tensors, each (C_i, Npix_b), with cell_ids a list of arrays
         """
-        # ---- Basic checks ----
+        import torch
+
+        # (A) Variable-length per-sample path: im is a list/tuple OR cell_ids is a list/tuple
+        if isinstance(im, (list, tuple)) or isinstance(cell_ids, (list, tuple)):
+            # Normalize to lists
+            im_list = im if isinstance(im, (list, tuple)) else [im]
+            cid_list = cell_ids if isinstance(cell_ids, (list, tuple)) else [cell_ids] * len(im_list)
+            assert len(im_list) == len(cid_list), "im list and cell_ids list must have same length"
+
+            outs = []
+            for xb, cb in zip(im_list, cid_list):
+                # xb: (C_i, Npix_b) -> (1, C_i, Npix_b)
+                if not torch.is_tensor(xb):
+                    xb = torch.as_tensor(xb, device=self.device, dtype=self.dtype)
+                if xb.dim() == 2:
+                    xb = xb.unsqueeze(0)
+                elif xb.dim() != 3 or xb.shape[0] != 1:
+                    raise ValueError(f"Each sample must be (C,N) or (1,C,N); got {tuple(xb.shape)}")
+
+                yb = self._convol_single(xb, ww, cell_ids=cb, nside=nside)  # (1, C_o, Npix_b)
+                outs.append(yb.squeeze(0))  # -> (C_o, Npix_b)
+            return outs  # List[Tensor], each (C_o, Npix_b)
+
+        # (B) Standard fixed-length batched path (your current implementation)
+        # ... keep your existing Convol_torch body from here unchanged ...
+        # (paste your current function body starting from the type casting and assertions)
+
+        # ---- Basic checks / casting ----
+        if not isinstance(im, torch.Tensor):
+            im = torch.as_tensor(im, device=self.device, dtype=self.dtype)
+        if not isinstance(ww, torch.Tensor):
+            ww = torch.as_tensor(ww, device=self.device, dtype=self.dtype)
+
         assert im.ndim == 3, f"`im` must be (B, C_i, Npix), got {tuple(im.shape)}"
-        assert ww.shape[2]==self.KERNELSZ*self.KERNELSZ, f"`ww` must be (C_i, C_o, KERNELSZ*KERNELSZ), got {tuple(ww.shape)}"
-        
         B, C_i, Npix = im.shape
         device = im.device
         dtype  = im.dtype
-        
-        # Align class tensors to device/dtype
-        idx_nn = self.idx_nn.to(device=device, dtype=torch.long)  # (Npix, P)
-        w_idx  = self.w_idx.to(device=device, dtype=torch.long)   # (Npix, P) or (Npix, S, P)
-        w_w    = self.w_w.to(device=device, dtype=dtype)          # (Npix, P) or (Npix, S, P)
-    
-        # Neighbor count P inferred from idx_nn
-        assert idx_nn.ndim == 2 and idx_nn.size(0) == Npix, \
-            f"`idx_nn` must be (Npix, P) with Npix={Npix}, got {tuple(idx_nn.shape)}"
-        P = idx_nn.size(1)
-    
-        # ---- 1) Gather neighbor values from im along the Npix dimension -> (B, C_i, Npix, P)
-        # im: (B,C_i,Npix) -> (B,C_i,Npix,1); idx: (1,1,Npix,P) broadcast over (B,C_i)
+
+        # ---- Recompute (idx_nn, w_idx, w_w) depending on cell_ids shape ----
+        # target shapes:
+        #   idx_nn_eff : (B, Npix, P)
+        #   w_idx_eff  : (B, Npix, S, P)
+        #   w_w_eff    : (B, Npix, S, P)
+        if cell_ids is not None:
+            # ---- Recompute (idx_nn, w_idx, w_w) depending on cell_ids shape ----
+            # Normaliser: accepter Tensor, ndarray ou list/tuple de 1 élément (cas var-length, B=1)
+            if isinstance(cell_ids, (list, tuple)):
+                # liste d'ids (souvent longueur 1 en var-length)
+                if len(cell_ids) == 1:
+                    cid = np.asarray(cell_ids[0])[None, :]  # -> (1, Npix)
+                else:
+                    # si jamais >1, on essaie d'empiler (doit avoir même Npix par élément)
+                    cid = np.stack([np.asarray(c) for c in cell_ids], axis=0)
+            elif torch.is_tensor(cell_ids):
+                c = cell_ids.detach().cpu().numpy()
+                cid = c if c.ndim != 1 else c[None, :]      # uniformiser en 2D quand B=1
+            else:
+                c = np.asarray(cell_ids)
+                cid = c if c.ndim != 1 else c[None, :]
+
+            # cid est maintenant (B, Npix)
+            idx_nn_eff, w_idx_eff, w_w_eff = self.make_idx_weights_from_cell_ids(
+                cid, nside, device=device
+            )
+            # shapes: (B, Npix, P), (B, Npix, S, P|P), (B, Npix, S, P|P)
+            P = idx_nn_eff.shape[-1]
+            S = w_idx_eff.shape[-2] if w_idx_eff.ndim == 4 else 1
+
+            # s’assurer des dtypes/devices
+            idx_nn_eff = torch.as_tensor(idx_nn_eff, device=device, dtype=torch.long)
+            w_idx_eff  = torch.as_tensor(w_idx_eff,  device=device, dtype=torch.long)
+            w_w_eff    = torch.as_tensor(w_w_eff,    device=device, dtype=dtype)
+        else:
+            # Use precomputed (shared for batch)
+            if self.w_idx is None:
+                if self.cell_ids.ndim==1:
+                    l_cell=self.cell_ids[None,:]
+                else:
+                    l_cell=self.cell_ids
+                    
+                idx_nn,w_idx,w_w = self.make_idx_weights_from_cell_ids(
+                    l_cell,
+                    polar=self.polar,
+                    gamma=self.gamma,
+                    device=self.device,
+                    allow_extrapolation=self.allow_extrapolation)
+
+                self.idx_nn = idx_nn
+                self.w_idx  = w_idx
+                self.w_w    = w_w
+            else:
+                idx_nn = self.idx_nn          # (Npix,P)
+                w_idx  = self.w_idx           # (Npix,P) or (Npix,S,P)
+                w_w    = self.w_w             # (Npix,P) or (Npix,S,P)
+
+            #assert idx_nn.ndim == 3 and idx_nn.size(1) == Npix, \
+            #    f"`idx_nn` must be (B,Npix,P) with Npix={Npix}, got {tuple(idx_nn.shape)}"
+
+            P = idx_nn.size(-1)
+
+            if w_idx.ndim == 3:
+                S = 1
+                w_idx_eff = w_idx[:, :, None, :]  # (B,Npix,1,P)
+                w_w_eff   = w_w[:, :, None, :]    # (B,Npix,1,P)
+            elif w_idx.ndim == 4:
+                S = w_idx.size(2)
+                w_idx_eff = w_idx    # (B,Npix,S,P)
+                w_w_eff   =  w_w    # (B,Npix,S,P)
+            else:
+                raise ValueError(f"Unsupported `w_idx` shape {tuple(w_idx.shape)}; expected (Npix,P) or (Npix,S,P)")
+            idx_nn_eff = idx_nn               # (B,Npix,P)
+
+        # ---- 1) Gather neighbor values from im along Npix -> (B, C_i, Npix, P)
         rim = torch.take_along_dim(
-            im.unsqueeze(-1),
-            idx_nn.unsqueeze(0).unsqueeze(0),
+            im.unsqueeze(-1),                  # (B, C_i, Npix, 1)
+            idx_nn_eff[:, None, :, :],         # (B, 1, Npix, P)
             dim=2
-        )  # (B, C_i, Npix, P)
-    
-        # ---- 2) Normalize w_idx / w_w to include a sector dim S ----
-        # Target layout: (Npix, S, P)
-        if w_idx.ndim == 2:
-            # (Npix, P) -> add sector dim S=1
-            assert w_idx.size(0) == Npix and w_idx.size(1) == P
-            w_idx_eff = w_idx.unsqueeze(1)  # (Npix, 1, P)
-            w_w_eff   = w_w.unsqueeze(1)    # (Npix, 1, P)
-            S = 1
-        elif w_idx.ndim == 3:
-            # (Npix, S, P)
-            Npix_, S, P_ = w_idx.shape
-            assert Npix_ == Npix and P_ == P, \
-                f"`w_idx` must be (Npix,S,P) with Npix={Npix}, P={P}, got {tuple(w_idx.shape)}"
-            assert w_w.shape == w_idx.shape, "`w_w` must match `w_idx` shape"
-            w_idx_eff = w_idx
-            w_w_eff   = w_w
-        else:
-            raise ValueError(f"Unsupported `w_idx` shape {tuple(w_idx.shape)}; expected (Npix,P) or (Npix,S,P)")
-    
-        # ---- 3) Normalize ww to (B, C_i, C_o, M, S) for uniform gather ----
+        )
+
+        # ---- 2) Normalize ww to (B, C_i, C_o, M, S)
         if ww.ndim == 3:
-            # (C_i, C_o, M) -> (B, C_i, C_o, M, S)
             C_i_w, C_o, M = ww.shape
             assert C_i_w == C_i, f"ww C_i mismatch: {C_i_w} vs im {C_i}"
             ww_eff = ww.unsqueeze(0).unsqueeze(-1).expand(B, -1, -1, -1, S)
-    
         elif ww.ndim == 4:
-            # Could be (C_i, C_o, M, S) or (B, C_i, C_o, M)
             if ww.shape[0] == C_i and ww.shape[1] != C_i:
-                # (C_i, C_o, M, S) -> (B, C_i, C_o, M, S)
+                # (C_i, C_o, M, S)
                 C_i_w, C_o, M, S_w = ww.shape
                 assert C_i_w == C_i, f"ww C_i mismatch: {C_i_w} vs im {C_i}"
                 assert S_w == S, f"ww S mismatch: {S_w} vs w_idx S {S}"
                 ww_eff = ww.unsqueeze(0).expand(B, -1, -1, -1, -1)
             elif ww.shape[0] == B:
-                # (B, C_i, C_o, M) -> (B, C_i, C_o, M, S)
+                # (B, C_i, C_o, M)
                 _, C_i_w, C_o, M = ww.shape
                 assert C_i_w == C_i, f"ww C_i mismatch: {C_i_w} vs im {C_i}"
                 ww_eff = ww.unsqueeze(-1).expand(-1, -1, -1, -1, S)
             else:
-                raise ValueError(
-                    f"Ambiguous 4D ww shape {tuple(ww.shape)}; expected (C_i,C_o,M,S) or (B,C_i,C_o,M)"
-                )
-    
+                raise ValueError(f"Ambiguous 4D ww shape {tuple(ww.shape)}; expected (C_i,C_o,M,S) or (B,C_i,C_o,M)")
         elif ww.ndim == 5:
             # (B, C_i, C_o, M, S)
             assert ww.shape[0] == B and ww.shape[1] == C_i, "ww batch/C_i mismatch"
@@ -520,27 +784,150 @@ class HOrientedConvol:
             ww_eff = ww
         else:
             raise ValueError(f"Unsupported ww shape {tuple(ww.shape)}")
-    
-        # ---- 4) Gather along M using w_idx_eff -> (B, C_i, C_o, Npix, S, P)
-        idx_exp = w_idx_eff.unsqueeze(0).unsqueeze(0).unsqueeze(0)     # (1,1,1,Npix,S,P)
+        
+        # --- Sanitize shapes: ensure w_idx_eff / w_w_eff == (B, Npix, S, P)
+        
+        # ---- 3) Gather along M using w_idx_eff -> (B, C_i, C_o, Npix, S, P)
+        idx_exp = w_idx_eff[:, None, None, :, :, :]            # (B,1,1,Npix,S,P)
         rw = torch.take_along_dim(
-            ww_eff.unsqueeze(-1),  # (B, C_i, C_o, M, S, 1)
-            idx_exp,               # (1,1,1,Npix,S,P) -> broadcast
-            dim=3                  # gather along M
+            ww_eff.unsqueeze(-1),                               # (B,C_i,C_o,M,S,1)
+            idx_exp,                                           # (B,1,1,Npix,S,P)
+            dim=3                                              # gather along M
         )  # -> (B, C_i, C_o, Npix, S, P)
+        # ---- 4) Apply extra neighbor weights ----
+        rw = rw * w_w_eff[:, None, None, :, :, :]              # (B, C_i, C_o, Npix, S, P)
+        # ---- 5) Combine neighbor values and weights ----
+        rim_exp = rim[:, :, None, :, None, :]                  # (B, C_i, 1, Npix, 1, P)
+        out_ci  = (rim_exp * rw).sum(dim=-1)                   # sum over P -> (B, C_i, C_o, Npix, S)
+        out_ci  = out_ci.sum(dim=-1)                           # sum over S -> (B, C_i, C_o, Npix)
+        out     = out_ci.sum(dim=1)                            # sum over C_i -> (B, C_o, Npix)
+
+        return out
     
-        # ---- 5) Apply extra neighbor weights ----
-        rw = rw * w_w_eff.unsqueeze(0).unsqueeze(0).unsqueeze(0)       # (B, C_i, C_o, Npix, S, P)
-    
-        # ---- 6) Combine neighbor values and weights ----
-        # rim: (B, C_i, Npix, P) -> expand to (B, C_i, 1, Npix, 1, P)
-        rim_exp = rim[:, :, None, :, None, :]
-        # sum over neighbors (P), then over sectors (S), then over input channels (C_i)
-        out_ci  = (rim_exp * rw).sum(dim=-1)    # (B, C_i, C_o, Npix, S)
-        out_ci  = out_ci.sum(dim=-1)            # (B, C_i, C_o, Npix)
-        out     = out_ci.sum(dim=1)             # (B, C_o, Npix)
+    def _to_numpy_1d(self, ids):
+        """Return a 1D numpy array of int64 for a single set of cell ids."""
+        import numpy as np, torch
+        if isinstance(ids, np.ndarray):
+            return ids.reshape(-1).astype(np.int64, copy=False)
+        if torch.is_tensor(ids):
+            return ids.detach().cpu().to(torch.long).view(-1).numpy()
+        # python list/tuple of ints
+        return np.asarray(ids, dtype=np.int64).reshape(-1)
+
+    def _is_varlength_batch(self, ids):
+        """
+        True if ids is a list/tuple of per-sample id arrays (var-length batch).
+        False if ids is a single array/tensor of ids (shared for whole batch).
+        """
+        import numpy as np, torch
+        if isinstance(ids, (list, tuple)):
+            return True
+        if isinstance(ids, np.ndarray) and ids.ndim == 2:
+            # This would be a dense (B, Npix) matrix -> NOT var-length list
+            return False
+        if torch.is_tensor(ids) and ids.dim() == 2:
+            return False
+        return False
+
+    def Down(self, im, cell_ids=None, nside=None,max_poll=False):
+        """
+        If `cell_ids` is a single set of ids -> return a single (Tensor, Tensor).
+        If `cell_ids` is a list (var-length)           -> return (list[Tensor], list[Tensor]).
+        """
+        if self.f is None:
+            if self.dtype==torch.float64:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float64')
+            else:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float32')
+                
+        if cell_ids is None:
+            dim,cdim = self.f.ud_grade_2(im,cell_ids=self.cell_ids,nside=self.nside,max_poll=False)
+            return dim,cdim
+        
+        if nside is None:
+            nside = self.nside
+
+        # var-length mode: list/tuple of ids, one per sample
+        if self._is_varlength_batch(cell_ids):
+            outs, outs_ids = [], []
+            B = len(cell_ids)
+            for b in range(B):
+                cid_b = self._to_numpy_1d(cell_ids[b])
+                # extraire le bon échantillon d'`im`
+                if torch.is_tensor(im):
+                    xb = im[b:b+1]  # (1, C, N_b)
+                    yb, ids_b = self.f.ud_grade_2(xb, cell_ids=cid_b, nside=nside,max_poll=max_poll)
+                    outs.append(yb.squeeze(0))  # (C, N_b')
+                else:
+                    # si im est déjà une liste de (C, N_b)
+                    xb = im[b]
+                    yb, ids_b = self.f.ud_grade_2(xb[None, ...], cell_ids=cid_b, nside=nside,max_poll=max_poll)
+                    outs.append(yb.squeeze(0))
+                outs_ids.append(torch.as_tensor(ids_b, device=outs[-1].device, dtype=torch.long))
+            return outs, outs_ids
+
+        # grille commune (un seul vecteur d'ids)
+        cid = self._to_numpy_1d(cell_ids)
+        return self.f.ud_grade_2(im, cell_ids=cid, nside=nside,max_poll=False)
+
+    def Up(self, im, cell_ids=None, nside=None, o_cell_ids=None):
+        """
+        If `cell_ids` / `o_cell_ids` are single arrays  -> return Tensor.
+        If they are lists (var-length per sample)       -> return list[Tensor].
+        """
+        if self.f is None:
+            if self.dtype==torch.float64:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float64')
+            else:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float32')
+                
+        if cell_ids is None:
+            dim = self.f.up_grade(im,self.nside*2,cell_ids=self.cell_ids,nside=self.nside)
+            return dim
+        
+        if nside is None:
+            nside = self.nside
+
+        # var-length: listes parallèles
+        if self._is_varlength_batch(cell_ids):
+            assert isinstance(o_cell_ids, (list, tuple)) and len(o_cell_ids) == len(cell_ids), \
+                "In var-length mode, `o_cell_ids` must be a list with same length as `cell_ids`."
+            outs = []
+            B = len(cell_ids)
+            for b in range(B):
+                cid_b  = self._to_numpy_1d(cell_ids[b])      # coarse ids
+                ocid_b = self._to_numpy_1d(o_cell_ids[b])    # fine   ids
+                if torch.is_tensor(im):
+                    xb = im[b:b+1]  # (1, C, N_b_coarse)
+                    yb = self.f.up_grade(xb, nside*2, cell_ids=cid_b, nside=nside,
+                                         o_cell_ids=ocid_b, force_init_index=True)
+                    outs.append(yb.squeeze(0))  # (C, N_b_fine)
+                else:
+                    xb = im[b]  # (C, N_b_coarse)
+                    yb = self.f.up_grade(xb[None, ...], nside*2, cell_ids=cid_b, nside=nside,
+                                         o_cell_ids=ocid_b, force_init_index=True)
+                    outs.append(yb.squeeze(0))
+            return outs
+
+        # grille commune
+        cid  = self._to_numpy_1d(cell_ids)
+        ocid = self._to_numpy_1d(o_cell_ids) if o_cell_ids is not None else None
+        return self.f.up_grade(im, nside*2, cell_ids=cid, nside=nside,
+                               o_cell_ids=ocid, force_init_index=True)
+
+    def to_tensor(self,x):
+        if self.f is None:
+            if self.dtype==torch.float64:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float64')
+            else:
+                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float32')
+        return self.f.backend.bk_cast(x)
     
-        return out
+    def to_numpy(self,x):
+        if isinstance(x,np.ndarray):
+            return x
+        return x.cpu().numpy()
+