foscat 2025.10.2__py3-none-any.whl → 2026.1.1__py3-none-any.whl

foscat/SphericalStencil.py

@@ -2,13 +2,8 @@
 # Author: J.-M. Delouis
 import numpy as np
 import healpy as hp
-import foscat.scat_cov as sc
 import torch
 
-import numpy as np
-import torch
-import healpy as hp
-
 
 class SphericalStencil:
     """
@@ -61,8 +56,7 @@ class SphericalStencil:
             device=None,
             dtype=None,
             n_gauges=1,
-            gauge_type='cosmo',
-            scat_op=None,
+            gauge_type='phi',
     ):
         assert kernel_sz >= 1 and int(kernel_sz) == kernel_sz
         assert kernel_sz % 2 == 1, "kernel_sz must be odd"
@@ -75,10 +69,6 @@ class SphericalStencil:
         self.gauge_type=gauge_type
         
         self.nest = bool(nest)
-        if scat_op is None:
-            self.f=sc.funct(KERNELSZ=self.KERNELSZ)
-        else:
-            self.f=scat_op
             
         # Torch defaults
         if device is None:
@@ -354,10 +344,27 @@ class SphericalStencil:
         # --- build the local (P,3) stencil once on device
         P = self.P
         vec_np = np.zeros((P, 3), dtype=float)
-        grid = (np.arange(self.KERNELSZ) - self.KERNELSZ // 2) / self.nside
-        vec_np[:, 0] = np.tile(grid, self.KERNELSZ)
-        vec_np[:, 1] = np.repeat(grid, self.KERNELSZ)
-        vec_np[:, 2] = 1.0 - np.sqrt(vec_np[:, 0]**2 + vec_np[:, 1]**2)
+        grid = (np.arange(self.KERNELSZ) - self.KERNELSZ // 2)
+         
+        # NEW: angular offsets
+        xx,yy=np.meshgrid(grid,grid)
+        s=1.0 # could be modified
+        alpha_pix = hp.nside2resol(self.nside, arcmin=False)  # ~ taille angulaire typique
+        dtheta = (np.sqrt(xx**2+yy**2) * alpha_pix * s).ravel()
+        dphi   = (np.arctan2(yy,xx)).ravel()
+        # local spherical displacement
+        # convert to unit vectors
+        x = np.sin(dtheta) * np.cos(dphi)
+        y = np.sin(dtheta) * np.sin(dphi)
+        z = np.cos(dtheta)  
+        #print(self.nside*x.reshape(self.KERNELSZ,self.KERNELSZ))
+        #print(self.nside*y.reshape(self.KERNELSZ,self.KERNELSZ))
+        #print(self.nside*z.reshape(self.KERNELSZ,self.KERNELSZ))
+        vec_np = np.stack([x, y, z], axis=-1)
+        
+        #vec_np[:, 0] = np.tile(grid, self.KERNELSZ)
+        #vec_np[:, 1] = np.repeat(grid, self.KERNELSZ)
+        #vec_np[:, 2] = 1.0 - np.sqrt(vec_np[:, 0]**2 + vec_np[:, 1]**2)
         vec_t = torch.as_tensor(vec_np, device=self.device, dtype=self.dtype)     # (P,3)
 
         # --- rotation matrices for all targets & gauges: (K,G,3,3)
@@ -371,7 +378,7 @@ class SphericalStencil:
             th, ph, alpha, G=self.G, gauge_cosmo=(self.gauge_type=='cosmo'),
             device=self.device, dtype=self.dtype
         )  # shape (K,G,3,3)
-
+        
         # --- rotate stencil for each (target, gauge): (K,G,P,3)
         #     einsum over local stencil (P,3) with rotation (K,G,3,3)
         rotated = torch.einsum('kgij,pj->kgpi', R_t, vec_t)  # (K,G,P,3)
@@ -568,119 +575,6 @@ class SphericalStencil:
         self.dtype  = dtype
 
     
-    '''
-    def bind_support_torch_multi(self, ids_sorted_np, *, device=None, dtype=None):
-        """
-        Multi-gauge sparse binding (Step B).
-        Uses self.idx_t_multi / self.w_t_multi prepared by prepare_torch(..., G>1)
-        and builds, for each gauge g, (pos_safe, w_norm, present).
-
-        Parameters
-        ----------
-        ids_sorted_np : np.ndarray (K,)
-            Sorted pixel ids for available samples (matches the last axis of your data).
-        device, dtype : torch device/dtype for the produced mapping tensors.
-
-        Side effects
-        ------------
-        Sets:
-          - self.ids_sorted_np  : (K,)
-          - self.pos_safe_t_multi : (G, 4, K*P)  LongTensor
-          - self.w_norm_t_multi   : (G, 4, K*P)  Tensor
-          - self.present_t_multi  : (G, 4, K*P)  BoolTensor
-          - (and mirrors device/dtype in self.device/self.dtype)
-        """
-        assert hasattr(self, 'idx_t_multi') and self.idx_t_multi is not None, \
-            "Call prepare_torch(..., G>0) before bind_support_torch_multi(...)"
-        assert hasattr(self, 'w_t_multi') and self.w_t_multi is not None
-
-        if device is None: device = self.device
-        if dtype  is None: dtype  = self.dtype
-
-        self.ids_sorted_np = np.asarray(ids_sorted_np, dtype=np.int64).reshape(-1)
-        ids_sorted = torch.as_tensor(self.ids_sorted_np, device=device, dtype=torch.long)
-
-        G, _, M = self.idx_t_multi.shape
-        K = self.Kb
-        P = self.P
-        assert M == K*P, "idx_t_multi second axis must have K*P columns"
-
-        pos_list, present_list, wnorm_list = [], [], []
-
-        for g in range(G):
-            idx = self.idx_t_multi[g].to(device=device, dtype=torch.long)   # (4, M)
-            w   = self.w_t_multi[g].to(device=device, dtype=dtype)          # (4, M)
-
-            pos = torch.searchsorted(ids_sorted, idx.reshape(-1)).view(4, M)
-            in_range = pos < ids_sorted.numel()
-            cmp_vals = torch.full_like(idx, -1)
-            cmp_vals[in_range] = ids_sorted[pos[in_range]]
-            present = (cmp_vals == idx)
-
-            # normalize weights per column after masking
-            w = w * present
-            colsum = w.sum(dim=0, keepdim=True).clamp_min(1e-12)
-            w_norm = w / colsum
-
-            pos_safe = torch.where(present, pos, torch.zeros_like(pos))
-
-            pos_list.append(pos_safe)
-            present_list.append(present)
-            wnorm_list.append(w_norm)
-
-        self.pos_safe_t_multi = torch.stack(pos_list, dim=0)     # (G, 4, M)
-        self.present_t_multi  = torch.stack(present_list, dim=0) # (G, 4, M)
-        self.w_norm_t_multi   = torch.stack(wnorm_list, dim=0)   # (G, 4, M)
-
-        # mirror runtime placement
-        self.device = device
-        self.dtype  = dtype
-    
-    # ------------------------------------------------------------------
-    # Step B: bind support Torch
-    # ------------------------------------------------------------------
-    def bind_support_torch(self, ids_sorted_np, *, device=None, dtype=None):
-        """
-        Map HEALPix neighbor indices (from Step A) to actual data samples
-        sorted by pixel id. Produces pos_safe and normalized weights.
-
-        Parameters
-        ----------
-        ids_sorted_np : np.ndarray (K,)
-            Sorted pixel ids for available data.
-        device, dtype : Torch device/dtype for results.
-        """
-        if device is None:
-            device = self.device
-        if dtype is None:
-            dtype = self.dtype
-
-        self.ids_sorted_np = np.asarray(ids_sorted_np, dtype=np.int64)
-        ids_sorted = torch.as_tensor(self.ids_sorted_np, device=device, dtype=torch.long)
-
-        idx = self.idx_t.to(device=device, dtype=torch.long)
-        w   = self.w_t.to(device=device, dtype=dtype)
-
-        M = self.Kb * self.P
-        idx = idx.view(4, M)
-        w   = w.view(4, M)
-
-        pos = torch.searchsorted(ids_sorted, idx.reshape(-1)).view(4, M)
-        in_range = pos < ids_sorted.shape[0]
-        cmp_vals = torch.full_like(idx, -1)
-        cmp_vals[in_range] = ids_sorted[pos[in_range]]
-        present = (cmp_vals == idx)
-
-        w = w * present
-        colsum = w.sum(dim=0, keepdim=True).clamp_min(1e-12)
-        w_norm = w / colsum
-
-        self.pos_safe_t = torch.where(present, pos, torch.zeros_like(pos))
-        self.w_norm_t   = w_norm
-        self.present_t  = present
-        self.device = device
-        self.dtype  = dtype
-    '''
     # ------------------------------------------------------------------
     # Step C: apply convolution (already Torch in your code)
     # ------------------------------------------------------------------
@@ -1215,7 +1109,7 @@ class SphericalStencil:
         vals = torch.cat(vals_all, dim=0)
 
         
-        indices = torch.stack([cols, rows], dim=0)             # (2, nnz) invert rows/cols for foscat needs
+        indices = torch.stack([cols, rows], dim=0)             
 
         if return_sparse_tensor:
             M = torch.sparse_coo_tensor(indices, vals, size=shape, device=device, dtype=k_dtype).coalesce()
@@ -1224,123 +1118,10 @@ class SphericalStencil:
             return vals, indices, shape
 
 
-    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=max_poll)
-            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):
-        return torch.tensor(x,device=self.device,dtype=self.dtype)
-    
+        return torch.tensor(x,device='cuda')
+
     def to_numpy(self,x):
         if isinstance(x,np.ndarray):
             return x
-        return x.cpu().numpy()
-        
+        return x.cpu().numpy()

foscat/alm_loc.py

@@ -0,0 +1,270 @@
+
+import numpy as np
+import healpy as hp
+
+from foscat.alm import alm as _alm
+import torch
+
+class alm_loc(_alm):
+    """
+    Local/partial-sky variant of foscat.alm.alm.
+
+    Key design choice (to match alm.py exactly when full-sky is provided):
+    - Reuse *all* Legendre/normalization machinery from the parent class (alm),
+      i.e. shift_ph(), compute_legendre_m(), ratio_mm, A/B recurrences, etc.
+      This is critical for matching alm.map2alm() numerically.
+
+    Differences vs alm.map2alm():
+    - Input map is [..., n] with explicit (nside, cell_ids)
+    - Only rings touched by cell_ids are processed.
+    - For rings with full coverage, we run the exact same FFT+tiling logic as alm.comp_tf()
+      (but only for those rings) -> bitwise comparable up to backend FFT differences.
+    - For rings with partial coverage, we compute a *partial DFT* for m=0..mmax,
+      using the same phase convention as alm.comp_tf():
+          FFT kernel uses exp(-i 2pi (m mod Nring) j / Nring)
+          then apply the per-ring shift exp(-i m phi0) via self.matrix_shift_ph
+    """
+
+    def __init__(self, backend=None, lmax=24, limit_range=1e10):
+        super().__init__(backend=backend, lmax=lmax, nside=None, limit_range=limit_range)
+
+    # --------- helpers: ring layout identical to alm.ring_th/ring_ph ----------
+    @staticmethod
+    def _ring_starts_sizes(nside: int):
+        starts = []
+        sizes = []
+        n = 0
+        for k in range(nside - 1):
+            N = 4 * (k + 1)
+            starts.append(n); sizes.append(N)
+            n += N
+        for _ in range(2 * nside + 1):
+            N = 4 * nside
+            starts.append(n); sizes.append(N)
+            n += N
+        for k in range(nside - 1):
+            N = 4 * (nside - 1 - k)
+            starts.append(n); sizes.append(N)
+            n += N
+        return np.asarray(starts, np.int64), np.asarray(sizes, np.int32)
+
+    def _to_ring_ids(self, nside: int, cell_ids: np.ndarray, nest: bool) -> np.ndarray:
+        if nest:
+            return hp.nest2ring(nside, cell_ids)
+        return cell_ids
+
+    def _group_by_ring(self, nside: int, ring_ids: np.ndarray):
+        """
+        Returns:
+          ring_idx: ring number (0..4*nside-2) per pixel
+          pos:      position along ring (0..Nring-1) per pixel
+          order:    sort order grouping by ring then pos
+          starts,sizes: ring layout
+        """
+        starts, sizes = self._ring_starts_sizes(nside)
+
+        # ring index = last start <= ring_id
+        ring_idx = np.searchsorted(starts, ring_ids, side="right") - 1
+        ring_idx = ring_idx.astype(np.int32)
+
+        pos = (ring_ids - starts[ring_idx]).astype(np.int32)
+
+        order = np.lexsort((pos, ring_idx))
+        return ring_idx, pos, order, starts, sizes
+
+    # ------------------ local Fourier transform per ring ---------------------
+    def comp_tf_loc(self, im, nside: int, cell_ids, nest: bool = False, realfft: bool = True, mmax=None):
+        """
+        Returns:
+          rings_used: 1D np.ndarray of ring indices present
+          ft: backend tensor of shape [..., nrings_used, mmax+1] (complex)
+               where last axis is m, ring axis matches rings_used order.
+        """
+        nside = int(nside)
+        cell_ids = np.asarray(cell_ids, dtype=np.int64)
+        if mmax is None:
+            mmax = min(self.lmax, 3 * nside - 1)
+        mmax = int(mmax)
+
+        # Ensure parent caches for this nside exist (matrix_shift_ph, A/B, ratio_mm, etc.)
+        self.shift_ph(nside)
+
+        ring_ids = self._to_ring_ids(nside, cell_ids, nest)
+        ring_idx, pos, order, starts, sizes = self._group_by_ring(nside, ring_ids)
+
+        ring_idx = ring_idx[order]
+        pos = pos[order]
+
+        i_im = self.backend.bk_cast(im)
+        i_im = self.backend.bk_gather(i_im, order, axis=-1)  # reorder last axis
+
+        rings_used, start_ptr, counts = np.unique(ring_idx, return_index=True, return_counts=True)
+
+        # Build output per ring as list then concat
+        out_per_ring = []
+        for r, s0, cnt in zip(rings_used.tolist(), start_ptr.tolist(), counts.tolist()):
+            Nring = int(sizes[r])
+            p = pos[s0:s0+cnt]
+
+            v = self.backend.bk_gather(i_im, np.arange(s0, s0+cnt, dtype=np.int64), axis=-1)
+
+            if cnt == Nring:
+                # Full ring: exact same FFT+tiling logic as alm.comp_tf for 1 ring
+                # Need data ordered by pos (already grouped, but ensure pos is 0..N-1)
+                if not np.all(p == np.arange(Nring, dtype=p.dtype)):
+                    # reorder within ring
+                    sub_order = np.argsort(p)
+                    v = self.backend.bk_gather(v, sub_order, axis=-1)
+
+                if realfft:
+                    tmp = self.rfft2fft(v)
+                else:
+                    tmp = self.backend.bk_fft(v)
+
+                l_n = tmp.shape[-1]
+                if l_n < mmax + 1:
+                    repeat_n = (mmax // l_n) + 1
+                    tmp = self.backend.bk_tile(tmp, repeat_n, axis=-1)
+
+                tmp = tmp[..., :mmax+1]
+
+                # Apply per-ring shift exp(-i m phi0) exactly like alm.comp_tf
+                shift = self.matrix_shift_ph[nside][r, :mmax+1]  # [m]
+                tmp = tmp * shift
+                out_per_ring.append(self.backend.bk_expand_dims(tmp, axis=-2))  # [...,1,m]
+            else:
+                # Partial ring: partial DFT for required m, using same aliasing as FFT branch
+                m_vec = np.arange(mmax+1, dtype=np.int64)
+                m_mod = (m_vec % Nring).astype(np.int64)
+
+                # angles: 2pi * pos * m_mod / Nring
+                ang = (2.0 * np.pi / Nring) * p.astype(np.float64)[:, None] * m_mod[None, :].astype(np.float64)
+                ker = np.exp(-1j * ang).astype(np.complex128)  # [cnt, m]
+
+                ker_bk = self.backend.bk_cast(ker)
+
+                # v is [..., cnt]; we want [..., m] = sum_cnt v*ker
+                tmp = self.backend.bk_reduce_sum(
+                    self.backend.bk_expand_dims(v, axis=-1) * ker_bk,
+                    axis=-2
+                )  # [..., m]
+
+                shift = self.matrix_shift_ph[nside][r, :mmax+1]  # [m] true m shift
+                tmp = tmp * shift
+                out_per_ring.append(self.backend.bk_expand_dims(tmp, axis=-2))  # [...,1,m]
+
+        ft = self.backend.bk_concat(out_per_ring, axis=-2)  # [..., nrings, m]
+        return np.asarray(rings_used, dtype=np.int32), ft
+
+    # ---------------------------- map -> alm --------------------------------
+    def map2alm_loc(self, im, nside: int, cell_ids, nest: bool = False, lmax=None):
+        nside = int(nside)
+        if lmax is None:
+            lmax = min(self.lmax, 3 * nside - 1)
+        lmax = int(lmax)
+
+        # Ensure a batch dimension like alm.map2alm expects
+        _added_batch = False
+        if hasattr(im, 'ndim') and im.ndim == 1:
+            im = im[None, :]
+            _added_batch = True
+        elif (not hasattr(im, 'ndim')) and len(im.shape) == 1:
+            im = im[None, :]
+            _added_batch = True
+
+        rings_used, ft = self.comp_tf_loc(im, nside=nside, cell_ids=cell_ids, nest=nest, realfft=True, mmax=lmax)
+
+        # cos(theta) on used rings
+        co_th = np.cos(self.ring_th(nside)[rings_used])
+
+        # ft is [..., R, m]
+        alm_out = None
+        
+        
+
+        for m in range(lmax + 1):
+            # IMPORTANT: reuse alm.compute_legendre_m and its normalization exactly
+            plm = self.compute_legendre_m(co_th, m, lmax, nside) / (12 * nside**2)  # [L,R]
+            plm_bk = self.backend.bk_cast(plm)
+
+            ft_m = ft[..., :, m]  # [..., R]
+            tmp = self.backend.bk_reduce_sum(
+                self.backend.bk_expand_dims(ft_m, axis=-2) * plm_bk,
+                axis=-1
+            )  # [..., L]
+            l_vals = np.arange(m, lmax + 1, dtype=np.float64)
+            scale = np.sqrt(2.0 * l_vals + 1.0)
+
+            # convertir scale en backend tensor (torch) sur le bon device
+            scale_t = self.backend.bk_cast(scale)  # ou un helper équivalent
+            # reshape pour broadcast si nécessaire: [1, L] ou [L]
+            shape = (1,) * (tmp.ndim - 1) + (scale_t.shape[0],)
+            scale_t = scale_t.reshape(shape)
+            
+            tmp = tmp * scale_t
+            if m == 0:
+                alm_out = tmp
+            else:
+                alm_out = self.backend.bk_concat([alm_out, tmp], axis=-1)
+        if _added_batch:
+            alm_out = alm_out[0]
+        return alm_out
+
+    # ---------------------------- alm -> Cl ---------------------------------
+    def anafast_loc(self, im, nside: int, cell_ids, nest: bool = False, lmax=None):
+        
+        if lmax is None:
+            lmax = min(self.lmax, 3 * nside - 1)
+        lmax = int(lmax)
+
+        alm = self.map2alm_loc(im, nside=nside, cell_ids=cell_ids, nest=nest, lmax=lmax)
+
+        # cl has same batch dims as alm, plus ell dim
+        batch_shape = alm.shape[:-1]
+        cl = torch.zeros(batch_shape + (lmax + 1,), dtype=torch.float64, device=alm.device)
+
+        idx = 0
+        for m in range(lmax + 1):
+            L = lmax - m + 1
+            a = alm[..., idx:idx+L]   # shape: batch + (L,)
+            idx += L
+
+            p = self.backend.bk_real(a * self.backend.bk_conjugate(a))  # batch + (L,)
+
+            if m == 0:
+                cl[..., m:] += p
+            else:
+                cl[..., m:] += 2.0 * p
+
+        # divide by (2l+1), broadcast over batch dims
+        denom = (2 * torch.arange(lmax + 1, dtype=cl.dtype, device=alm.device) + 1)  # (lmax+1,)
+        denom = denom.reshape((1,) * len(batch_shape) + (lmax + 1,))                  # batch-broadcast
+        cl = cl / denom
+        return cl
+    '''
+    def anafast_loc(self, im, nside: int, cell_ids, nest: bool = False, lmax=None):
+        if lmax is None:
+            lmax = min(self.lmax, 3 * nside - 1)
+        lmax = int(lmax)
+
+        alm = self.map2alm_loc(im, nside=nside, cell_ids=cell_ids, nest=nest, lmax=lmax)
+        
+        # Unpack and compute Cl with correct real-field folding:
+        cl = torch.zeros((lmax + 1,), dtype=alm.dtype, device=alm.device)
+        
+        idx = 0
+        for m in range(lmax + 1):
+            L = lmax - m + 1
+            a = alm[..., idx:idx+L]
+            idx += L
+            p = self.backend.bk_real(a * self.backend.bk_conjugate(a))
+            # sum over any batch dims
+            p = self.backend.bk_reduce_sum(p, axis=tuple(range(p.ndim-1))) if p.ndim > 1 else p
+            if m == 0:
+                cl[m:] += p
+            else:
+                cl[m:] += 2.0 * p
+        denom = (2*torch.arange(lmax+1,dtype=p.dtype, device=alm.device)+1)
+        cl = cl / denom
+        return cl
+    '''

foscat/scat.py

@@ -1659,7 +1659,7 @@ class funct(FOC.FoCUS):
         s2j2 = None
         l2_image = None
         for j1 in range(jmax):
-            if j1 < jmax - self.OSTEP:  # stop to add scales
+            if j1 < jmax:  # stop to add scales
                 # Convol image along the axis defined by 'axis' using the wavelet defined at
                 # the foscat initialisation
                 # c_image_real is [....,Npix_j1,....,Norient]

foscat/scat1D.py

@@ -1282,7 +1282,7 @@ class funct(FOC.FoCUS):
         l2_image = None
 
         for j1 in range(jmax):
-            if j1 < jmax - self.OSTEP:  # stop to add scales
+            if j1 < jmax:  # stop to add scales
                 # Convol image along the axis defined by 'axis' using the wavelet defined at
                 # the foscat initialisation
                 # c_image_real is [....,Npix_j1,....,Norient]