foscat 2025.11.1__tar.gz → 2026.1.1__tar.gz

{foscat-2025.11.1/src/foscat.egg-info → foscat-2026.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: foscat
-Version: 2025.11.1
+Version: 2026.1.1
 Summary: Generate synthetic Healpix or 2D data using Cross Scattering Transform
 Author-email: Jean-Marc DELOUIS <jean.marc.delouis@ifremer.fr>
 Maintainer-email: Theo Foulquier <theo.foulquier@ifremer.fr>
@@ -91,74 +91,6 @@ git clone https://github.com/jmdelouis/FOSCAT_DEMO.git
 
 ```
 
-# Spherical data example
-
-## compute a synthetic image
-
-```
-python demo.py -n=32 -k -c -s=100
-```
-
-The _demo.py_ script serves as a demonstration of the capabilities of the foscat library. It utilizes the Cross Wavelet Scattering Transform to generate a Healpix map that possesses the same characteristics as a specified input map.
-
-- `-n=32` computes map with nside=32.
-- `-k` uses 5x5 kernel.
-- `-c` uses Scattering Covariance.
-- `-l` uses LBFGS minimizer.
-- `-s=100` computes 100 steps.
-
-```
-python demo.py -n=8 [-c|--cov][-s|--steps=3000][-S=1234|--seed=1234][-k|--k5x5][-d|--data][-o|--out][-r|--orient] [-p|--path][-a|--adam]
-
-```
-
-- The "-n" option specifies the nside of the input map. The maximum nside value is 256 with the default map.
-- The "--cov" option (optional) uses scat_cov instead of scat.
-- The "--steps" option (optional) specifies the number of iterations. If not specified, the default value is 1000.
-- The "--seed" option (optional) specifies the seed of the random generator.
-- The "--path" option (optional) allows you to define the path where the output files will be written. The default path is "data".
-- The "--k5x5" option (optional) uses a 5x5 kernel instead of a 3x3.
-- The "--data" option (optional) specifies the input data file to be used. If not specified, the default file "LSS_map_nside128.npy" will be used.
-- The "--out" option (optional) specifies the output file name. If not specified, the output file will be saved in "demo".
-- The "--orient" option (optional) specifies the number of orientations. If not specified, the default value is 4.
-- The "--adam" option (optional) makes the synthesis using the ADAM optimizer instead of the L_BFGS.
-
-## plot the result
-
-The following script generates a series of plots that showcase different aspects of the synthesis process using the _demo.py_ script.
-
-> python test2D.py
-
-```
-python plotdemo.py -n=32 -c
-```
-
-# 2D field demo
-
-> python test2Dplot.py
-
-# compute a synthetic turbulent field
-
-The python scripts _demo2D.py_ included in this package demonstrate how to use the foscat library to generate a 2D synthetic fields that have patterns with the same statistical properties as a specified 2D image. In this particular case, the input field is a sea surface temperature extracted from a north atlantic ocean simulation.
-
-> python testHealpix.py
-
-```
-python demo2d.py -n=32 -k -c
-```
-
-> python testHplot.py
-
-The following script generates a series of plots that showcase different aspects of the synthesis process using the _demo2D.py_ script.
-
-```
-python plotdemo2d.py -n=32 -c
-```
-
-For more information, see the [documentation](https://foscat-documentation.readthedocs.io/en/latest/index.html).
-
-> mpirun -np 3 testHealpix_mpi.py
-
 ## Authors and acknowledgment
 
 Authors: J.-M. Delouis, P. Campeti, T. Foulquier, J. Mangin, L. Mousset, T. Odaka, F. Paul, E. Allys

{foscat-2025.11.1 → foscat-2026.1.1}/README.md

@@ -62,74 +62,6 @@ git clone https://github.com/jmdelouis/FOSCAT_DEMO.git
 
 ```
 
-# Spherical data example
-
-## compute a synthetic image
-
-```
-python demo.py -n=32 -k -c -s=100
-```
-
-The _demo.py_ script serves as a demonstration of the capabilities of the foscat library. It utilizes the Cross Wavelet Scattering Transform to generate a Healpix map that possesses the same characteristics as a specified input map.
-
-- `-n=32` computes map with nside=32.
-- `-k` uses 5x5 kernel.
-- `-c` uses Scattering Covariance.
-- `-l` uses LBFGS minimizer.
-- `-s=100` computes 100 steps.
-
-```
-python demo.py -n=8 [-c|--cov][-s|--steps=3000][-S=1234|--seed=1234][-k|--k5x5][-d|--data][-o|--out][-r|--orient] [-p|--path][-a|--adam]
-
-```
-
-- The "-n" option specifies the nside of the input map. The maximum nside value is 256 with the default map.
-- The "--cov" option (optional) uses scat_cov instead of scat.
-- The "--steps" option (optional) specifies the number of iterations. If not specified, the default value is 1000.
-- The "--seed" option (optional) specifies the seed of the random generator.
-- The "--path" option (optional) allows you to define the path where the output files will be written. The default path is "data".
-- The "--k5x5" option (optional) uses a 5x5 kernel instead of a 3x3.
-- The "--data" option (optional) specifies the input data file to be used. If not specified, the default file "LSS_map_nside128.npy" will be used.
-- The "--out" option (optional) specifies the output file name. If not specified, the output file will be saved in "demo".
-- The "--orient" option (optional) specifies the number of orientations. If not specified, the default value is 4.
-- The "--adam" option (optional) makes the synthesis using the ADAM optimizer instead of the L_BFGS.
-
-## plot the result
-
-The following script generates a series of plots that showcase different aspects of the synthesis process using the _demo.py_ script.
-
-> python test2D.py
-
-```
-python plotdemo.py -n=32 -c
-```
-
-# 2D field demo
-
-> python test2Dplot.py
-
-# compute a synthetic turbulent field
-
-The python scripts _demo2D.py_ included in this package demonstrate how to use the foscat library to generate a 2D synthetic fields that have patterns with the same statistical properties as a specified 2D image. In this particular case, the input field is a sea surface temperature extracted from a north atlantic ocean simulation.
-
-> python testHealpix.py
-
-```
-python demo2d.py -n=32 -k -c
-```
-
-> python testHplot.py
-
-The following script generates a series of plots that showcase different aspects of the synthesis process using the _demo2D.py_ script.
-
-```
-python plotdemo2d.py -n=32 -c
-```
-
-For more information, see the [documentation](https://foscat-documentation.readthedocs.io/en/latest/index.html).
-
-> mpirun -np 3 testHealpix_mpi.py
-
 ## Authors and acknowledgment
 
 Authors: J.-M. Delouis, P. Campeti, T. Foulquier, J. Mangin, L. Mousset, T. Odaka, F. Paul, E. Allys

{foscat-2025.11.1 → foscat-2026.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "foscat"
-version = "2025.11.1"
+version = "2026.01.1"
 description = "Generate synthetic Healpix or 2D data using Cross Scattering Transform"
 readme = "README.md"
 license = { text = "BSD-3-Clause" }

{foscat-2025.11.1 → foscat-2026.1.1}/src/foscat/FoCUS.py

@@ -5,8 +5,11 @@ import healpy as hp
 import numpy as np
 import foscat.HealSpline as HS
 from scipy.interpolate import griddata
+from foscat.SphereDownGeo import SphereDownGeo
+from foscat.SphereUpGeo import SphereUpGeo
+import torch
 
-TMPFILE_VERSION = "V10_0"
+TMPFILE_VERSION = "V12_0"
 
 
 class FoCUS:
@@ -36,7 +39,7 @@ class FoCUS:
             mpi_rank=0
     ):
 
-        self.__version__ = "2025.11.1"
+        self.__version__ = "2026.01.1"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -57,7 +60,8 @@ class FoCUS:
         self.kernelR_conv = {}
         self.kernelI_conv = {}
         self.padding_conv = {}
-
+        self.down = {}
+        self.up = {}
         if not self.silent:
             print("================================================")
             print("          START FOSCAT CONFIGURATION")
@@ -648,6 +652,7 @@ class FoCUS:
         return rim
 
     # --------------------------------------------------------
+    
     def ud_grade_2(self, im, axis=0, cell_ids=None, nside=None,max_poll=False):
 
         if self.use_2D:
@@ -721,6 +726,22 @@ class FoCUS:
 
         else:
             shape = list(im.shape)
+            if nside is None:
+                l_nside=int(np.sqrt(shape[-1]//12))
+            else:
+                l_nside=nside
+            
+            nbatch=1
+            for k in range(len(shape)-1):
+                nbatch*=shape[k]
+            if l_nside not in self.down:
+                print('initialise down', l_nside)
+                self.down[l_nside] = SphereDownGeo(nside_in=l_nside, dtype=self.all_bk_type,mode="smooth", in_cell_ids=cell_ids)
+            
+            res,out_cell=self.down[l_nside](self.backend.bk_reshape(im,[nbatch,1,shape[-1]]))
+            
+            return self.backend.bk_reshape(res,shape[:-1]+[out_cell.shape[0]]),out_cell
+            '''
             if self.use_median:
                 if cell_ids is not None:
                     sim, new_cell_ids = self.backend.binned_mean(im, cell_ids,reduce='median')
@@ -747,6 +768,7 @@ class FoCUS:
                 return self.backend.bk_reduce_mean(
                     self.backend.bk_reshape(im, shape[0:-1]+[shape[-1]//4,4]), axis=-1
                 ),None
+            '''
 
     # --------------------------------------------------------
     def up_grade(self, im, nout,
@@ -836,6 +858,7 @@ class FoCUS:
             else:
                 lout = nside
                 
+            '''
             if (lout,nout) not in self.pix_interp_val or force_init_index:
                 if not self.silent:
                     print("compute lout nout", lout, nout)
@@ -926,12 +949,32 @@ class FoCUS:
                     
                     del w
                     del p
-
-            if lout == nout:
-                imout = im
-            else:
-                # work only on the last column
-
+            '''
+            shape=list(im.shape)
+            nbatch=1
+            for k in range(len(shape)-1):
+                nbatch*=shape[k]
+            
+            im=self.backend.bk_reshape(im,[nbatch,1,shape[-1]])
+            
+            while lout<nout:
+                if lout not in self.up:
+                    if o_cell_ids is None:
+                        l_o_cell_ids=torch.tensor(np.arange(12*(lout**2),dtype='int'),device=im.device)
+                    else:
+                        l_o_cell_ids=o_cell_ids
+                    self.up[lout] = SphereUpGeo(nside_out=lout, 
+                                                dtype=self.all_bk_type,
+                                                 cell_ids_out=l_o_cell_ids,
+                                                 up_norm="col_l1")
+                im, fine_ids = self.up[lout](self.backend.bk_cast(im))
+                lout*=2
+                if lout<nout and o_cell_ids is not None:
+                    o_cell_ids=torch.repeat(fine_ids,4)*4+ \
+                                torch.tile(torch.tensor([0,1,2,3],device=fine_ids.device,dtype=fine_ids.dtype),fine_ids.shape[0])
+            
+            return self.backend.bk_reshape(im,shape[:-1]+[im.shape[-1]])
+            '''
                 ndata = 1
                 for k in range(len(ishape)-1):
                     ndata = ndata * ishape[k]
@@ -960,6 +1003,7 @@ class FoCUS:
                     return self.backend.bk_reshape(
                         imout, ishape[0:-1]+[imout.shape[-1]]
                     )
+            '''
         return imout
 
     # --------------------------------------------------------
@@ -1354,7 +1398,9 @@ class FoCUS:
                     else:
                         l_cell_ids=cell_ids
                         
-                    nvalid=self.KERNELSZ**2
+                    nvalid=4*self.KERNELSZ**2
+                    if nvalid>12*nside**2:
+                        nvalid=12*nside**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')           
@@ -1500,7 +1546,7 @@ class FoCUS:
                     
                 else:
                     if l_kernel == 5:
-                        pw = 0.5
+                        pw = 0.75
                         pw2 = 0.5
                         threshold = 2e-5
 

foscat-2026.1.1/src/foscat/SphereDownGeo.py

@@ -0,0 +1,380 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+import healpy as hp
+
+
+class SphereDownGeo(nn.Module):
+    """
+    Geometric HEALPix downsampling operator (NESTED indexing).
+
+    This module reduces resolution by a factor 2:
+        nside_out = nside_in // 2
+
+    Input conventions
+    -----------------
+    - If in_cell_ids is None:
+        x is expected to be full-sphere: [B, C, N_in]
+        output is [B, C, K_out] with K_out = len(cell_ids_out) (or N_out if None).
+    - If in_cell_ids is provided (fine pixels at nside_in, NESTED):
+        x can be either:
+          * compact: [B, C, K_in] where K_in = len(in_cell_ids), aligned with in_cell_ids order
+          * full-sphere: [B, C, N_in] (also supported)
+        output is [B, C, K_out] where cell_ids_out is derived as unique(in_cell_ids // 4),
+        unless you explicitly pass cell_ids_out (then it will be intersected with the derived set).
+
+    Modes
+    -----
+    - mode="smooth": linear downsampling y = M @ x  (M sparse)
+    - mode="maxpool": non-linear max over available children (fast)
+    """
+
+    def __init__(
+        self,
+        nside_in: int,
+        mode: str = "smooth",
+        radius_deg: float | None = None,
+        sigma_deg: float | None = None,
+        weight_norm: str = "l1",
+        cell_ids_out: np.ndarray | list[int] | None = None,
+        in_cell_ids: np.ndarray | list[int] | torch.Tensor | None = None,
+        use_csr=True,
+        device=None,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device
+        self.dtype = dtype
+
+        self.nside_in = int(nside_in)
+        assert (self.nside_in & (self.nside_in - 1)) == 0, "nside_in must be a power of 2."
+        self.nside_out = self.nside_in // 2
+        assert self.nside_out >= 1, "nside_out must be >= 1."
+
+        self.N_in = 12 * self.nside_in * self.nside_in
+        self.N_out = 12 * self.nside_out * self.nside_out
+
+        self.mode = str(mode).lower()
+        assert self.mode in ("smooth", "maxpool"), "mode must be 'smooth' or 'maxpool'."
+
+        self.weight_norm = str(weight_norm).lower()
+        assert self.weight_norm in ("l1", "l2"), "weight_norm must be 'l1' or 'l2'."
+
+        # ---- Handle reduced-domain inputs (fine pixels) ----
+        self.in_cell_ids = self._validate_in_cell_ids(in_cell_ids)
+        self.has_in_subset = self.in_cell_ids is not None
+        if self.has_in_subset:
+            # derive parents
+            derived_out = np.unique(self.in_cell_ids // 4).astype(np.int64)
+            if cell_ids_out is None:
+                self.cell_ids_out = derived_out
+            else:
+                req_out = self._validate_cell_ids_out(cell_ids_out)
+                # keep only those compatible with derived_out (otherwise they'd be all-zero)
+                self.cell_ids_out = np.intersect1d(req_out, derived_out, assume_unique=False)
+                if self.cell_ids_out.size == 0:
+                    raise ValueError(
+                        "After intersecting cell_ids_out with unique(in_cell_ids//4), "
+                        "no coarse pixel remains. Check your inputs."
+                    )
+        else:
+            self.cell_ids_out = self._validate_cell_ids_out(cell_ids_out)
+
+        self.K_out = int(self.cell_ids_out.size)
+
+        # Column basis for smooth matrix:
+        # - full sphere: columns are 0..N_in-1
+        # - subset: columns are 0..K_in-1 aligned to self.in_cell_ids
+        self.K_in = int(self.in_cell_ids.size) if self.has_in_subset else self.N_in
+
+        if self.mode == "smooth":
+            if radius_deg is None:
+                # default: include roughly the 4 children footprint
+                # (healpy pixel size ~ sqrt(4pi/N), coarse pixel is 4x area)
+                radius_deg = 2.0 * hp.nside2resol(self.nside_out, arcmin=True) / 60.0
+            if sigma_deg is None:
+                sigma_deg = max(radius_deg / 2.0, 1e-6)
+
+            self.radius_deg = float(radius_deg)
+            self.sigma_deg = float(sigma_deg)
+            self.radius_rad = self.radius_deg * np.pi / 180.0
+            self.sigma_rad = self.sigma_deg * np.pi / 180.0
+                                        
+            M = self._build_down_matrix()  # shape (K_out, K_in or N_in)
+              
+            self.M = M.coalesce()
+            
+            if use_csr:
+                self.M = self.M.to_sparse_csr().to(self.device)
+
+            self.M_size = M.size()
+
+        else:
+            # Precompute children indices for maxpool
+            # For subset mode, store mapping from each parent to indices in compact vector,
+            # with -1 for missing children.
+            children = np.stack(
+                [4 * self.cell_ids_out + i for i in range(4)],
+                axis=1,
+            ).astype(np.int64)  # [K_out, 4] in fine pixel ids (full indexing)
+
+            if self.has_in_subset:
+                # map each child pixel id to position in in_cell_ids (compact index)
+                pos = self._positions_in_sorted(self.in_cell_ids, children.reshape(-1))
+                children_compact = pos.reshape(self.K_out, 4).astype(np.int64)  # -1 if missing
+                self.register_buffer(
+                    "children_compact",
+                    torch.tensor(children_compact, dtype=torch.long, device=self.device),
+                )
+            else:
+                self.register_buffer(
+                    "children_full",
+                    torch.tensor(children, dtype=torch.long, device=self.device),
+                )
+
+        # expose ids as torch buffers for convenience
+        self.register_buffer(
+            "cell_ids_out_t",
+            torch.tensor(self.cell_ids_out.astype(np.int64), dtype=torch.long, device=self.device),
+        )
+        if self.has_in_subset:
+            self.register_buffer(
+                "in_cell_ids_t",
+                torch.tensor(self.in_cell_ids.astype(np.int64), dtype=torch.long, device=self.device),
+            )
+
+    # ---------------- validation helpers ----------------
+    def _validate_cell_ids_out(self, cell_ids_out):
+        """Return a 1D np.int64 array of coarse cell ids (nside_out)."""
+        if cell_ids_out is None:
+            return np.arange(self.N_out, dtype=np.int64)
+
+        arr = np.asarray(cell_ids_out, dtype=np.int64).reshape(-1)
+        if arr.size == 0:
+            raise ValueError("cell_ids_out is empty: provide at least one coarse pixel id.")
+        arr = np.unique(arr)
+        if arr.min() < 0 or arr.max() >= self.N_out:
+            raise ValueError(f"cell_ids_out must be in [0, {self.N_out-1}] for nside_out={self.nside_out}.")
+        return arr
+
+    def _validate_in_cell_ids(self, in_cell_ids):
+        """Return a 1D np.int64 array of fine cell ids (nside_in) or None."""
+        if in_cell_ids is None:
+            return None
+        if torch.is_tensor(in_cell_ids):
+            arr = in_cell_ids.detach().cpu().numpy()
+        else:
+            arr = np.asarray(in_cell_ids)
+        arr = np.asarray(arr, dtype=np.int64).reshape(-1)
+        if arr.size == 0:
+            raise ValueError("in_cell_ids is empty: provide at least one fine pixel id or None.")
+        arr = np.unique(arr)
+        if arr.min() < 0 or arr.max() >= self.N_in:
+            raise ValueError(f"in_cell_ids must be in [0, {self.N_in-1}] for nside_in={self.nside_in}.")
+        return arr
+
+    @staticmethod
+    def _positions_in_sorted(sorted_ids: np.ndarray, query_ids: np.ndarray) -> np.ndarray:
+        """
+        For each query_id, return its index in sorted_ids if present, else -1.
+        sorted_ids must be sorted ascending unique.
+        """
+        q = np.asarray(query_ids, dtype=np.int64)
+        idx = np.searchsorted(sorted_ids, q)
+        ok = (idx >= 0) & (idx < sorted_ids.size) & (sorted_ids[idx] == q)
+        out = np.full(q.shape, -1, dtype=np.int64)
+        out[ok] = idx[ok]
+        return out
+
+    # ---------------- weights and matrix build ----------------
+    def _normalize_weights(self, w: np.ndarray) -> np.ndarray:
+        w = np.asarray(w, dtype=np.float64)
+        if w.size == 0:
+            return w
+        w = np.maximum(w, 0.0)
+
+        if self.weight_norm == "l1":
+            s = w.sum()
+            if s <= 0.0:
+                return np.ones_like(w) / max(w.size, 1)
+            return w / s
+
+        # l2
+        s2 = (w * w).sum()
+        if s2 <= 0.0:
+            return np.ones_like(w) / max(np.sqrt(w.size), 1.0)
+        return w / np.sqrt(s2)
+
+    def _build_down_matrix(self) -> torch.Tensor:
+        """Construct sparse matrix M (K_out, K_in or N_in) for the selected coarse pixels."""
+        nside_in = self.nside_in
+        nside_out = self.nside_out
+
+        radius_rad = self.radius_rad
+        sigma_rad = self.sigma_rad
+
+        rows: list[int] = []
+        cols: list[int] = []
+        vals: list[float] = []
+
+        # For subset columns, we use self.in_cell_ids as the basis
+        subset_cols = self.has_in_subset
+        in_ids = self.in_cell_ids  # np.ndarray or None
+
+        for r, p_out in enumerate(self.cell_ids_out.tolist()):
+            theta0, phi0 = hp.pix2ang(nside_out, int(p_out), nest=True)
+            vec0 = hp.ang2vec(theta0, phi0)
+
+            neigh = hp.query_disc(nside_in, vec0, radius_rad, inclusive=True, nest=True)
+            neigh = np.asarray(neigh, dtype=np.int64)
+
+            if subset_cols:
+                # keep only valid fine pixels
+                # neigh is not sorted; intersect1d expects sorted
+                neigh_sorted = np.sort(neigh)
+                keep = np.intersect1d(neigh_sorted, in_ids, assume_unique=False)
+                neigh = keep
+
+            # Fallback: if radius query returns nothing in subset mode, at least try the 4 children
+            if neigh.size == 0:
+                children = (4 * int(p_out) + np.arange(4, dtype=np.int64))
+                if subset_cols:
+                    pos = self._positions_in_sorted(in_ids, children)
+                    ok = pos >= 0
+                    if np.any(ok):
+                        neigh = children[ok]
+                    else:
+                        # nothing to connect -> row stays zero
+                        continue
+                else:
+                    neigh = children
+
+            theta, phi = hp.pix2ang(nside_in, neigh, nest=True)
+            vec = hp.ang2vec(theta, phi)
+
+            # angular distance via dot product
+            dots = np.clip(np.dot(vec, vec0), -1.0, 1.0)
+            ang = np.arccos(dots)
+            w = np.exp(- 2*(ang / sigma_rad) ** 2)
+
+            w = self._normalize_weights(w)
+
+            if subset_cols:
+                pos = self._positions_in_sorted(in_ids, neigh)
+                # all should be present due to filtering, but guard anyway
+                ok = pos >= 0
+                neigh_pos = pos[ok]
+                w = w[ok]
+                if neigh_pos.size == 0:
+                    continue
+                for c, v in zip(neigh_pos.tolist(), w.tolist()):
+                    rows.append(r)
+                    cols.append(int(c))
+                    vals.append(float(v))
+            else:
+                for c, v in zip(neigh.tolist(), w.tolist()):
+                    rows.append(r)
+                    cols.append(int(c))
+                    vals.append(float(v))
+
+        if len(rows) == 0:
+            # build an all-zero sparse tensor
+            indices = torch.zeros((2, 0), dtype=torch.long, device=self.device)
+            vals_t = torch.zeros((0,), dtype=self.dtype, device=self.device)
+            return torch.sparse_coo_tensor(
+                indices, vals_t, size=(self.K_out, self.K_in), device=self.device, dtype=self.dtype
+            ).coalesce()
+
+        rows_t = torch.tensor(rows, dtype=torch.long, device=self.device)
+        cols_t = torch.tensor(cols, dtype=torch.long, device=self.device)
+        vals_t = torch.tensor(vals, dtype=self.dtype, device=self.device)
+
+        indices = torch.stack([rows_t, cols_t], dim=0)
+        M = torch.sparse_coo_tensor(
+            indices,
+            vals_t,
+            size=(self.K_out, self.K_in),
+            device=self.device,
+            dtype=self.dtype,
+        ).coalesce()
+        return M
+
+    # ---------------- forward ----------------
+    def forward(self, x: torch.Tensor):
+        """
+        Parameters
+        ----------
+        x : torch.Tensor
+            If has_in_subset:
+                - [B,C,K_in] (compact, aligned with in_cell_ids) OR [B,C,N_in] (full sphere)
+            Else:
+                - [B,C,N_in] (full sphere)
+
+        Returns
+        -------
+        y : torch.Tensor
+            [B,C,K_out]
+        cell_ids_out : torch.Tensor
+            [K_out] coarse pixel ids (nside_out), aligned with y last dimension.
+        """
+        if x.dim() != 3:
+            raise ValueError("x must be [B, C, N]")
+
+        B, C, N = x.shape
+        if self.has_in_subset:
+            if N not in (self.K_in, self.N_in):
+                raise ValueError(
+                    f"x last dim must be K_in={self.K_in} (compact) or N_in={self.N_in} (full), got {N}"
+                )
+        else:
+            if N != self.N_in:
+                raise ValueError(f"x last dim must be N_in={self.N_in}, got {N}")
+
+        if self.mode == "smooth":
+
+            # If x is full-sphere but M is subset-based, gather compact inputs
+            if self.has_in_subset and N == self.N_in:
+                x_use = x.index_select(dim=2, index=self.in_cell_ids_t.to(x.device))
+            else:
+                x_use = x
+
+            # sparse mm expects 2D: (K_out, K_in) @ (K_in, B*C)
+            x2 = x_use.reshape(B * C, -1).transpose(0, 1).contiguous()
+            y2 = torch.sparse.mm(self.M, x2)
+            y = y2.transpose(0, 1).reshape(B, C, self.K_out).contiguous()
+            return y, self.cell_ids_out_t.to(x.device)
+
+        # maxpool
+        if self.has_in_subset and N == self.N_in:
+            x_use = x.index_select(dim=2, index=self.in_cell_ids_t.to(x.device))
+        else:
+            x_use = x
+
+        if self.has_in_subset:
+            # children_compact: [K_out, 4] indices in 0..K_in-1 or -1
+            ch = self.children_compact.to(x.device)  # [K_out,4]
+            # gather with masking
+            # We build y by iterating 4 children with max
+            y = None
+            for j in range(4):
+                idx = ch[:, j]  # [K_out]
+                mask = idx >= 0
+                # start with very negative so missing children don't win
+                tmp = torch.full((B, C, self.K_out), -torch.inf, device=x.device, dtype=x.dtype)
+                if mask.any():
+                    tmp[:, :, mask] = x_use.index_select(dim=2, index=idx[mask]).reshape(B, C, -1)
+                y = tmp if y is None else torch.maximum(y, tmp)
+            # If a parent had no valid children at all, it is -inf -> set to 0
+            y = torch.where(torch.isfinite(y), y, torch.zeros_like(y))
+            return y, self.cell_ids_out_t.to(x.device)
+
+        else:
+            ch = self.children_full.to(x.device)  # [K_out,4] full indices
+            # gather children and max
+            xch = x_use.index_select(dim=2, index=ch.reshape(-1)).reshape(B, C, self.K_out, 4)
+            y = xch.max(dim=3).values
+            return y, self.cell_ids_out_t.to(x.device)