foscat 2025.9.1__py3-none-any.whl → 2025.9.4__py3-none-any.whl

foscat/healpix_unet_torch.py

@@ -6,11 +6,11 @@ GPU by default (when available), with graceful CPU fallback if Foscat ops are CP
 - ReLU + BatchNorm after each convolution (encoder & decoder)
 - Segmentation/Regression heads with optional final activation
 - PyTorch-ified: inherits from nn.Module, standard state_dict
-- Device management: tries CUDA first; if Foscat HOrientedConvol cannot run on CUDA, falls back to CPU
+- Device management: tries CUDA first; if Foscat SphericalStencil cannot run on CUDA, falls back to CPU
 
 Shape convention: (B, C, Npix)
 
-Requirements: foscat (scat_cov.funct + HOrientedConvol.Convol_torch must be differentiable on torch tensors)
+Requirements: foscat (scat_cov.funct + SphericalStencil.Convol_torch must be differentiable on torch tensors)
 """
 from __future__ import annotations
 from typing import List, Optional, Literal, Tuple
@@ -18,13 +18,15 @@ import numpy as np
 
 import torch
 import torch.nn as nn
+import healpy as hp
+
 import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, TensorDataset
 
 import foscat.scat_cov as sc
-import foscat.HOrientedConvol as ho
+import foscat.SphericalStencil as ho
 import matplotlib.pyplot as plt
 
-
 class HealpixUNet(nn.Module):
     """U-Net-like architecture on the HEALPix sphere using Foscat oriented convolutions.
 
@@ -40,6 +42,13 @@ class HealpixUNet(nn.Module):
         Cell indices for the finest resolution (nside = in_nside) in nested scheme.
     KERNELSZ : int, default 3
         Spatial kernel size K (K x K) for oriented convolution.
+    gauge_type : str
+        Type of gauge :
+        'cosmo' use the same definition than
+           https://www.aanda.org/articles/aa/abs/2022/12/aa44566-22/aa44566-22.html
+        'phi' is define at the pole, could be better for earth observation not using intensivly the pole
+    G : int, default 1
+        Number of gauges for the orientation definition.
     task : {'regression','segmentation'}, default 'regression'
         Chooses the head and default activation.
     out_channels : int, default 1
@@ -50,6 +59,8 @@ class HealpixUNet(nn.Module):
     device : str | torch.device | None, default: 'cuda' if available else 'cpu'
         Preferred device. The module will probe whether Foscat ops can run on CUDA; if not,
         it will fallback to CPU and keep all parameters/buffers on CPU for consistency.
+    down_type:
+        {"mean","max"}, default "max". Equivalent of max poll during down  
     prefer_foscat_gpu : bool, default True
         When device is CUDA, try to move Foscat operators (internal tensors) to CUDA and do a dry-run.
         If the dry-run fails, everything falls back to CPU.
@@ -58,7 +69,7 @@ class HealpixUNet(nn.Module):
     -----
     - Two oriented convolutions per level. After each conv: BatchNorm1d + ReLU.
     - Downsampling uses foscat ``ud_grade_2``; upsampling uses ``up_grade``.
-    - Convolution kernels are explicit parameters (shape [C_in, C_out, K*K]) and applied via ``HOrientedConvol.Convol_torch``.
+    - Convolution kernels are explicit parameters (shape [C_in, C_out, K*K]) and applied via ``SphericalStencil.Convol_torch``.
     - Foscat ops device is auto-probed to avoid CPU/CUDA mismatches.
     """
 
@@ -75,16 +86,28 @@ class HealpixUNet(nn.Module):
             final_activation: Optional[Literal['none', 'sigmoid', 'softmax']] = None,
             device: Optional[torch.device | str] = None,
             prefer_foscat_gpu: bool = True,
-            dtype: Literal['float32','float64'] = 'float32'
+            gauge_type: Optional[Literal['cosmo','phi']] = 'cosmo',
+            G: int =1,
+            down_type: Optional[Literal['mean','max']] = 'max',  
+            dtype: Literal['float32','float64'] = 'float32',
+            head_reduce: Literal['mean','learned']='mean'
     ) -> None:
         super().__init__()
-
+        
         self.dtype=dtype
         if dtype=='float32':
             self.np_dtype=np.float32
+            self.torch_dtype=torch.float32
         else:
             self.np_dtype=np.float64
-            
+            self.torch_dtype=torch.float32
+
+        self.gauge_type=gauge_type
+        self.G = int(G)
+
+        if self.G < 1:
+            raise ValueError("G must be >= 1")
+    
         if cell_ids is None:
             raise ValueError("cell_ids must be provided for the finest resolution.")
         if len(chanlist) == 0:
@@ -93,11 +116,16 @@ class HealpixUNet(nn.Module):
         self.in_nside = int(in_nside)
         self.n_chan_in = int(n_chan_in)
         self.chanlist = list(map(int, chanlist))
+        self.chanlist = [self.chanlist[k]*self.G for k in range(len(self.chanlist))]
         self.KERNELSZ = int(KERNELSZ)
         self.task = task
-        self.out_channels = int(out_channels)
+        self.out_channels = int(out_channels)*self.G
         self.prefer_foscat_gpu = bool(prefer_foscat_gpu)
-
+        if down_type == 'max':
+            self.max_poll = True
+        else:
+            self.max_poll = False
+        
         # Choose default final activation if not given
         if final_activation is None:
             if task == 'regression':
@@ -124,25 +152,26 @@ class HealpixUNet(nn.Module):
         current_nside = self.in_nside
 
         # ---------- Oriented convolutions per level (encoder & decoder) ----------
-        self.hconv_enc: List[ho.HOrientedConvol] = []
-        self.hconv_dec: List[ho.HOrientedConvol] = []
+        self.hconv_enc: List[ho.SphericalStencil] = []
+        self.hconv_dec: List[ho.SphericalStencil] = []
         
         # dummy data to propagate shapes/ids through ud_grade_2
         l_data = self.f.backend.bk_cast(np.zeros((1, 1, cell_ids.shape[0]), dtype=self.np_dtype))
 
         for l in range(depth):
             # operator at encoder level l
-            hc = ho.HOrientedConvol(current_nside,
-                                    self.KERNELSZ,
-                                    cell_ids=self.l_cell_ids[l],
-                                    dtype=self.dtype)
-            #hc.make_idx_weights()
+            hc = ho.SphericalStencil(current_nside,
+                                     self.KERNELSZ,
+                                     n_gauges = self.G,
+                                     gauge_type=self.gauge_type,
+                                     cell_ids=self.l_cell_ids[l],
+                                     dtype=self.torch_dtype)
             
             self.hconv_enc.append(hc)
             
             # downsample once to get next level ids and new data shape
             l_data, next_ids = hc.Down(
-                l_data, cell_ids=self.l_cell_ids[l], nside=current_nside
+                l_data, cell_ids=self.l_cell_ids[l], nside=current_nside,max_poll=self.max_poll
             )
             self.l_cell_ids[l + 1] = self.f.backend.to_numpy(next_ids)
             current_nside //= 2
@@ -158,20 +187,23 @@ class HealpixUNet(nn.Module):
         
         inC = self.n_chan_in
         for l, outC in enumerate(self.chanlist):
+            if outC % self.G != 0:
+                raise ValueError(f"chanlist[{l}] = {outC} must be divisible by G={self.G}")
+            outC_g = outC // self.G
 
-            # conv1: inC -> outC
-            w1 = torch.empty(inC, outC, self.KERNELSZ * self.KERNELSZ)
-            nn.init.kaiming_uniform_(w1.view(inC * outC, -1), a=np.sqrt(5))
+            # conv1: inC -> outC (via multi-gauge => noyau (Ci, Co_g, P))
+            w1 = torch.empty(inC, outC_g, self.KERNELSZ * self.KERNELSZ)
+            nn.init.kaiming_uniform_(w1.view(inC * outC_g, -1), a=np.sqrt(5))
             self.enc_w1.append(nn.Parameter(w1))
-            self.enc_bn1.append(nn.BatchNorm1d(outC))
+            self.enc_bn1.append(self._norm_1d(outC, kind="group"))
 
-            # conv2: outC -> outC
-            w2 = torch.empty(outC, outC, self.KERNELSZ * self.KERNELSZ)
-            nn.init.kaiming_uniform_(w2.view(outC * outC, -1), a=np.sqrt(5))
+            # conv2: outC -> outC  (entrée = total outC ; noyau (outC, outC_g, P))
+            w2 = torch.empty(outC, outC_g, self.KERNELSZ * self.KERNELSZ)
+            nn.init.kaiming_uniform_(w2.view(outC * outC_g, -1), a=np.sqrt(5))
             self.enc_w2.append(nn.Parameter(w2))
-            self.enc_bn2.append(nn.BatchNorm1d(outC))
+            self.enc_bn2.append(self._norm_1d(outC, kind="group"))
 
-            inC = outC  # next level input channels
+            inC = outC  # next layer sees total channels
 
         # decoder conv weights and BN (mirrored levels)
         self.dec_w1 = nn.ParameterList()
@@ -181,45 +213,88 @@ class HealpixUNet(nn.Module):
 
         for d in range(depth):
             level = depth - 1 - d  # encoder level we are going back to
-            hc = ho.HOrientedConvol(self.enc_nsides[level],
-                                    self.KERNELSZ,
-                                    cell_ids=self.l_cell_ids[level],
-                                    dtype=self.dtype)
+            hc = ho.SphericalStencil(self.enc_nsides[level],
+                                     self.KERNELSZ,
+                                     n_gauges = self.G,
+                                     gauge_type=self.gauge_type,
+                                     cell_ids=self.l_cell_ids[level],
+                                     dtype=self.torch_dtype)
             #hc.make_idx_weights()
             self.hconv_dec.append(hc)
 
             upC = self.chanlist[level + 1] if level + 1 < depth else self.chanlist[level]
             skipC = self.chanlist[level]
-            inC_dec = upC + skipC
-            outC_dec = skipC
+            inC_dec  = upC + skipC           # total en entrée
+            outC_dec = skipC                 # total en sortie (ce que tu avais déjà)
+
+            if outC_dec % self.G != 0:
+                raise ValueError(f"decoder outC at level {level} = {outC_dec} must be divisible by G={self.G}")
+            outC_dec_g = outC_dec // self.G
 
-            w1 = torch.empty(inC_dec, outC_dec, self.KERNELSZ * self.KERNELSZ)
-            nn.init.kaiming_uniform_(w1.view(inC_dec * outC_dec, -1), a=np.sqrt(5))
+            w1 = torch.empty(inC_dec, outC_dec_g, self.KERNELSZ * self.KERNELSZ)
+            nn.init.kaiming_uniform_(w1.view(inC_dec * outC_dec_g, -1), a=np.sqrt(5))
             self.dec_w1.append(nn.Parameter(w1))
-            self.dec_bn1.append(nn.BatchNorm1d(outC_dec))
+            self.dec_bn1.append(self._norm_1d(outC_dec, kind="group"))
 
-            w2 = torch.empty(outC_dec, outC_dec, self.KERNELSZ * self.KERNELSZ)
-            nn.init.kaiming_uniform_(w2.view(outC_dec * outC_dec, -1), a=np.sqrt(5))
+            w2 = torch.empty(outC_dec, outC_dec_g, self.KERNELSZ * self.KERNELSZ)
+            nn.init.kaiming_uniform_(w2.view(outC_dec * outC_dec_g, -1), a=np.sqrt(5))
             self.dec_w2.append(nn.Parameter(w2))
-            self.dec_bn2.append(nn.BatchNorm1d(outC_dec))
+            self.dec_bn2.append(self._norm_1d(outC_dec, kind="group"))
 
         # Output head (on finest grid, channels = chanlist[0])
-        self.head_hconv = ho.HOrientedConvol(self.in_nside,
-                                             self.KERNELSZ,
-                                             cell_ids=self.l_cell_ids[0],
-                                             dtype=self.dtype)
-        
-        head_inC = self.chanlist[0]
-        self.head_w = nn.Parameter(torch.empty(head_inC, self.out_channels, self.KERNELSZ * self.KERNELSZ))
-        nn.init.kaiming_uniform_(self.head_w.view(head_inC * self.out_channels, -1), a=np.sqrt(5))
-        self.head_bn = nn.BatchNorm1d(self.out_channels) if self.task == 'segmentation' else None
+        self.head_hconv = ho.SphericalStencil(self.in_nside,
+                                              self.KERNELSZ,
+                                              n_gauges=self.G,   #Mandatory for the output
+                                              gauge_type=self.gauge_type,
+                                              cell_ids=self.l_cell_ids[0],
+                                              dtype=self.torch_dtype)
 
+        head_inC = self.chanlist[0]
+        if self.out_channels % self.G != 0:
+            raise ValueError(f"out_channels={self.out_channels} must be divisible by G={self.G}")
+        outC_head_g = self.out_channels // self.G
+
+        self.head_w = nn.Parameter(
+            torch.empty(head_inC, outC_head_g, self.KERNELSZ * self.KERNELSZ)
+        )
+        nn.init.kaiming_uniform_(self.head_w.view(head_inC * outC_head_g, -1), a=np.sqrt(5))
+        self.head_bn = self._norm_1d(self.out_channels, kind="group") if self.task == 'segmentation' else None
+        
+        # Choose how to reduce across gauges at head:
+        # 'sum' (default), 'mean', or 'learned' (via 1x1 conv).
+        self.head_reduce = getattr(self, 'head_reduce', 'mean')  # you can turn this into a ctor arg if you like
+        if self.head_reduce == 'learned':
+            # Mixer takes G*outC_head_g -> out_channels (K-wise 1x1)
+            self.head_mixer = nn.Conv1d(self.G * outC_head_g, self.out_channels, kernel_size=1, bias=True)
+        else:
+            self.head_mixer = None
+    
         # ---- Decide runtime device (probe Foscat on CUDA, else CPU) ----
         self.runtime_device = self._probe_and_set_runtime_device(self.device)
 
+    # -------------------------- define local batchnorm/group -------------------
+    def _norm_1d(self, C: int, kind: str = "group", **kwargs) -> nn.Module:
+        """
+        Return a normalization layer for (B, C, N) tensors.
+        kind: "group" | "instance" | "batch"
+        kwargs: extra args (e.g., num_groups for GroupNorm)
+        """
+        if kind == "group":
+            num_groups = kwargs.get("num_groups", min(8, max(1, C // 8)) or 1)
+            # s’assurer que num_groups divise C
+            while C % num_groups != 0 and num_groups > 1:
+                num_groups //= 2
+            return nn.GroupNorm(num_groups=num_groups, num_channels=C)
+        elif kind == "instance":
+            return nn.InstanceNorm1d(C, affine=True, track_running_stats=False)
+        elif kind == "batch":
+            return nn.BatchNorm1d(C)
+        else:
+            raise ValueError(f"Unknown norm kind: {kind}")
+            
     # -------------------------- device plumbing --------------------------
-    def _move_hconv_tensors(self, hc: ho.HOrientedConvol, device: torch.device) -> None:
-        """Best-effort: move any torch.Tensor attribute of HOrientedConvol to device."""
+    def _move_hconv_tensors(self, hc: ho.SphericalStencil, device: torch.device) -> None:
+        """Best-effort: move any torch.Tensor attribute of SphericalStencil to device."""
         for name, val in list(vars(hc).items()):
             try:
                 if torch.is_tensor(val):
@@ -265,6 +340,78 @@ class HealpixUNet(nn.Module):
         self.device = device
         self.runtime_device = self._probe_and_set_runtime_device(device)
         return self.runtime_device
+    
+    # --- inside HealpixUNet class, add a single-sample forward helper ---
+    def _forward_one(self, x1: torch.Tensor, cell_ids1=None) -> torch.Tensor:
+        """
+        Single-sample forward. x1: (1, C_in, Npix_1). Returns (1, out_channels, Npix_1).
+        `cell_ids1` can be None or a 1D array (Npix_1,) for this sample.
+        """
+        if x1.dim() != 3 or x1.shape[0] != 1:
+            raise ValueError(f"_forward_one expects (1, C, Npix), got {tuple(x1.shape)}")
+        # Reuse existing forward by calling it with B=1 (your code already supports per-sample ids)
+        if cell_ids1 is None:
+            return super().forward(x1)
+        else:
+            # normalize ids to numpy 1D
+            if isinstance(cell_ids1, torch.Tensor):
+                cell_ids1 = cell_ids1.detach().cpu().numpy()
+            elif isinstance(cell_ids1, list):
+                cell_ids1 = np.asarray(cell_ids1)
+            if cell_ids1.ndim == 1:
+                ci = cell_ids1[None, :]  # (1, Npix_1) so the current code path is happy
+            else:
+                ci = cell_ids1
+            return super().forward(x1, cell_ids=ci)
+        
+    def _as_tensor_batch(self, x):
+        """
+        Ensure a (B, C, N) tensor.
+        - If x is a list of tensors, concatenate if all N are equal.
+        - If len==1, keep a batch dim (1, C, N).
+        - If x is already a tensor, return as-is.
+        """
+        if isinstance(x, list):
+            if len(x) == 1:
+                t = x[0]
+                # If t is (C, N) -> make it (1, C, N)
+                return t.unsqueeze(0) if t.dim() == 2 else t
+            # all same length -> concat along batch
+            Ns = [t.shape[-1] for t in x]
+            if all(n == Ns[0] for n in Ns):
+                return torch.cat([t if t.dim() == 3 else t.unsqueeze(0) for t in x], dim=0)
+            # variable-length with B>1 not supported in a single tensor
+            raise ValueError("Variable-length batch detected; use batch_size=1 or loop per-sample.")
+        return x
+    
+    # --- replace your current `forward` signature/body with a dispatcher ---
+    def forward_any(self, x, cell_ids: Optional[np.ndarray] = None):
+        """
+        If `x` is a Tensor (B,C,N): standard batched path (requires same N for all).
+        If `x` is a list of Tensors: variable-length per-sample path, returns a list of outputs.
+        """
+        # Variable-length list path
+        if isinstance(x, (list, tuple)):
+            outs = []
+            if cell_ids is None or isinstance(cell_ids, (list, tuple)):
+                cids = cell_ids if isinstance(cell_ids, (list, tuple)) else [None] * len(x)
+            else:
+                raise ValueError("When x is a list, cell_ids must be a list of same length or None.")
+
+            for xb, cb in zip(x, cids):
+                if not torch.is_tensor(xb):
+                    xb = torch.as_tensor(xb, dtype=torch.float32, device=self.runtime_device)
+                if xb.dim() == 2:
+                    xb = xb.unsqueeze(0)  # (1,C,Nb)
+                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._forward_one(xb.to(self.runtime_device), cell_ids1=cb)  # (1,Co,Nb)
+                outs.append(yb.squeeze(0))  # -> (Co, Nb)
+            return outs  # List[Tensor] (each length Nb)
+
+        # Fixed-length tensor path (your current implementation)
+        return super().forward(x, cell_ids=cell_ids)
 
     # -------------------------- forward --------------------------
     def forward(self, x: torch.Tensor,cell_ids: Optional[np.ndarray ] = None) -> torch.Tensor:
@@ -306,6 +453,7 @@ class HealpixUNet(nn.Module):
             l_data = self.hconv_enc[l].Convol_torch(l_data,
                                                     self.enc_w1[l],
                                                     cell_ids=l_cell_ids)
+            l_data = self._as_tensor_batch(l_data)
             l_data = self.enc_bn1[l](l_data)
             l_data = F.relu(l_data, inplace=True)
 
@@ -313,6 +461,7 @@ class HealpixUNet(nn.Module):
             l_data = self.hconv_enc[l].Convol_torch(l_data,
                                                     self.enc_w2[l],
                                                     cell_ids=l_cell_ids)
+            l_data = self._as_tensor_batch(l_data)
             l_data = self.enc_bn2[l](l_data)
             l_data = F.relu(l_data, inplace=True)
 
@@ -321,12 +470,10 @@ class HealpixUNet(nn.Module):
 
             # downsample (except bottom level) -> ensure output is on runtime_device
             if l < len(self.chanlist) - 1:
-                #l_data, l_cell_ids = self.f.ud_grade_2(
-                #    l_data, cell_ids=t_cell_ids[l], nside=current_nside
-                #)
                 l_data, l_cell_ids = self.hconv_enc[l].Down(
-                    l_data, cell_ids=t_cell_ids[l], nside=current_nside
+                    l_data, cell_ids=t_cell_ids[l], nside=current_nside,max_poll=self.max_poll
                 )
+                l_data = self._as_tensor_batch(l_data)
                 if cell_ids is not None:
                     t_cell_ids[l+1]=l_cell_ids
                 else:
@@ -335,34 +482,24 @@ class HealpixUNet(nn.Module):
                 if isinstance(l_data, torch.Tensor) and l_data.device != self.runtime_device:
                     l_data = l_data.to(self.runtime_device)
                 current_nside //= 2
-
+                
         # Decoder
         for d in range(len(self.chanlist)):
-            level = len(self.chanlist) - 1 - d  # corresponding encoder level
+            level = len(self.chanlist) - 1 - d  # encoder level we are going back to
 
             if level < len(self.chanlist) - 1:
-                # upsample to next finer grid
-                # upsample to next finer grid (from level+1 -> level)
-                src_nside = self.enc_nsides[level + 1]    # current (coarser)
-                tgt_nside = self.enc_nsides[level]        # next finer (== src*2)
-                # Foscat up_grade signature expects current (coarse) ids in `cell_ids`
-                # and target (fine) ids in `o_cell_ids` (matching original UNET code).
-                '''
-                l_data = self.f.up_grade(
-                    l_data,
-                    tgt_nside,
-                    cell_ids=t_cell_ids[level + 1],  # source (coarser) ids
-                    o_cell_ids=t_cell_ids[level],     # target (finer) ids
-                    nside=src_nside,
-                )
-                '''
-                l_data = self.hconv_enc[l].Up(
+                # upsample: from encoder level (level+1) [coarser] -> level [finer]
+                src_nside = self.enc_nsides[level + 1]    # coarse
+
+                # Use the **decoder** operator at this step (consistent with your hconv_dec stack)
+                l_data = self.hconv_dec[d].Up(
                     l_data,
-                    cell_ids=t_cell_ids[level + 1],  # source (coarser) ids
-                    o_cell_ids=t_cell_ids[level],     # target (finer) ids
+                    cell_ids=t_cell_ids[level + 1],   # source/coarse IDs
+                    o_cell_ids=t_cell_ids[level],     # target/fine IDs
                     nside=src_nside,
                 )
-                
+                l_data = self._as_tensor_batch(l_data)
+
                 if isinstance(l_data, torch.Tensor) and l_data.device != self.runtime_device:
                     l_data = l_data.to(self.runtime_device)
 
@@ -370,49 +507,85 @@ class HealpixUNet(nn.Module):
             concat = self.f.backend.bk_concat([skips[level], l_data], 1)
             l_data = concat.to(self.runtime_device) if torch.is_tensor(concat) else concat
 
-            if cell_ids is not None:
-                l_cell_ids = t_cell_ids[level]
-                
-            # apply decoder convs on this grid
+            # choose the right cell_ids for convolutions at this resolution
+            l_cell_ids = t_cell_ids[level] if (cell_ids is not None) else None
+
+            # apply decoder convs on this grid using the matching decoder operator
             hc = self.hconv_dec[d]
-            l_data = hc.Convol_torch(l_data,
-                                     self.dec_w1[d],
-                                     cell_ids=l_cell_ids)
-            
+            l_data = hc.Convol_torch(l_data, self.dec_w1[d], cell_ids=l_cell_ids)
+            l_data = self._as_tensor_batch(l_data)
             l_data = self.dec_bn1[d](l_data)
             l_data = F.relu(l_data, inplace=True)
 
-            l_data = hc.Convol_torch(l_data,
-                                     self.dec_w2[d],
-                                     cell_ids=l_cell_ids)
-            
+            l_data = hc.Convol_torch(l_data, self.dec_w2[d], cell_ids=l_cell_ids)
+            l_data = self._as_tensor_batch(l_data)
             l_data = self.dec_bn2[d](l_data)
             l_data = F.relu(l_data, inplace=True)
 
         # Head on finest grid
-        out = self.head_hconv.Convol_torch(l_data, self.head_w)
-        if self.head_bn is not None:
-            out = self.head_bn(out)
+        # y_head_raw: (B, G*outC_head_g, K)
+        y_head_raw = self.head_hconv.Convol_torch(l_data, self.head_w, cell_ids=l_cell_ids)
+
+        B, Ctot, K = y_head_raw.shape
+        outC_head_g = int(self.out_channels)//self.G
+        assert Ctot == self.G * outC_head_g, \
+            f"Head expects G*outC_head_g channels, got {Ctot} != {self.G}*{outC_head_g}"
+
+        if self.head_mixer is not None and self.head_reduce == 'learned':
+            # 1x1 learned mixing across G*outC_head_g -> out_channels
+            y = self.head_mixer(y_head_raw)  # (B, out_channels, K)
+        else:
+            # reshape to (B, G, outC_head_g, K) then reduce across G
+            y_g = y_head_raw.view(B, self.G, outC_head_g, K)
+
+            y = y_g.mean(dim=1)   # (B, outC_head_g, K)
+
+        y = self._as_tensor_batch(y)
+        
+        # Optional BN + activation as before
+        if self.task == 'segmentation' and self.head_bn is not None:
+            y = self.head_bn(y)
+            
         if self.final_activation == 'sigmoid':
-            out = torch.sigmoid(out)
+            y = torch.sigmoid(y)
+            
         elif self.final_activation == 'softmax':
-            out = torch.softmax(out, dim=1)
-        return out
+            y = torch.softmax(y, dim=1)
+            
+        return y
 
     # -------------------------- utilities --------------------------
     @torch.no_grad()
     def predict(self, x: torch.Tensor, batch_size: int = 8,cell_ids: Optional[np.ndarray ] = None) -> torch.Tensor:
         self.eval()
         outs = []
-        for i in range(0, x.shape[0], batch_size):
-            if cell_ids is not None:
-                outs.append(self.forward(x[i : i + batch_size],
-                                         cell_ids=cell_ids[i : i + batch_size]))
-            else:
-                outs.append(self.forward(x[i : i + batch_size]))
+        if isinstance(x,np.ndarray):
+            x=self.to_Tensor(x)
+            
+        if not isinstance(x, torch.Tensor):
+            for i in range(len(x)):
+                if cell_ids is not None:
+                    outs.append(self.forward(x[i][None,:],cell_ids=cell_ids[i][:]))
+                else:
+                    outs.append(self.forward(x[i][None,:]))
+        else:
+            for i in range(0, x.shape[0], batch_size):
+                if cell_ids is not None:
+                    outs.append(self.forward(x[i : i + batch_size],
+                                             cell_ids=cell_ids[i : i + batch_size]))
+                else:
+                    outs.append(self.forward(x[i : i + batch_size]))
                 
         return torch.cat(outs, dim=0)
 
+    def to_tensor(self,x):
+        return self.hconv_enc[0].f.backend.bk_cast(x)
+    
+    def to_numpy(self,x):
+        if isinstance(x,np.ndarray):
+            return x
+        return x.cpu().numpy()
+    
     # -----------------------------
     # Kernel extraction & plotting
     # -----------------------------
@@ -609,109 +782,421 @@ if __name__ == "__main__":
                     np.testing.assert_allclose(y1_np, y3_np, rtol=0, atol=0)
 
     unittest.main()
+    
+from torch.utils.data import Dataset
+# 1) Dataset that omits cell_ids when None
+from torch.utils.data import Dataset, DataLoader, TensorDataset
+
+class HealpixDataset(Dataset):
+    """
+    Returns (x, y, cell_ids) per-sample if cell_ids is given, else (x, y).
+    Shapes:
+      x: (C, Npix)
+      y: (C_out or 1, Npix)
+      cell_ids: (Npix,) per-sample (or broadcasted from (Npix,))
+    """
+    def __init__(self, x, y, cell_ids=None, dtype=torch.float32):
+        self.x = torch.as_tensor(x, dtype=dtype)
+        self.y = torch.as_tensor(y, dtype=dtype)
+        assert self.x.shape[0] == self.y.shape[0], "x and y must share batch size"
+        self._has_cids = cell_ids is not None
+        if self._has_cids:
+            cid = torch.as_tensor(cell_ids, dtype=torch.long)
+            if cid.dim() == 1:
+                cid = cid.unsqueeze(0).expand(self.x.shape[0], -1)
+            assert cid.shape[0] == self.x.shape[0], "cell_ids must match batch size"
+            self.cids = cid
+        else:
+            self.cids = None
+
+    def __len__(self):
+        return self.x.shape[0]
+
+    def __getitem__(self, i):
+        if self._has_cids:
+            return self.x[i], self.y[i], self.cids[i]
+        else:
+            return self.x[i], self.y[i]
+
+# ---------------------------
+# Datasets / Collate helpers
+# ---------------------------
+
+class HealpixDataset(Dataset):
+    """
+    Fixed-grid dataset (common Npix for all samples).
+    Returns (x, y) if cell_ids is None, else (x, y, cell_ids).
+
+    x: (B, C, Npix)
+    y: (B, C_out or 1, Npix) or class indices depending on task
+    cell_ids: (Npix,) or (B, Npix)
+    """
+    def __init__(self,
+                 x: torch.Tensor,
+                 y: torch.Tensor,
+                 cell_ids: Optional[Union[np.ndarray, torch.Tensor]] = None,
+                 dtype: torch.dtype = torch.float32):
+        x = torch.as_tensor(x, dtype=dtype)
+        y = torch.as_tensor(y, dtype=dtype if y.ndim == 3 else torch.long)
+        assert x.shape[0] == y.shape[0], "x and y must share batch size"
+        self.x, self.y = x, y
+        self._has_cids = cell_ids is not None
+        if self._has_cids:
+            c = torch.as_tensor(cell_ids, dtype=torch.long)
+            if c.ndim == 1:  # broadcast single (Npix,) to (B, Npix)
+                c = c.unsqueeze(0).expand(x.shape[0], -1)
+            assert c.shape == (x.shape[0], x.shape[2]), "cell_ids must be (B,Npix) or (Npix,)"
+            self.cids = c
+        else:
+            self.cids = None
+
+    def __len__(self) -> int: return self.x.shape[0]
+
+    def __getitem__(self, i: int):
+        if self._has_cids:
+            return self.x[i], self.y[i], self.cids[i]
+        return self.x[i], self.y[i]
+
+# ---------------------------
+# Datasets / Collate helpers
+# ---------------------------
+
+class HealpixDataset(Dataset):
+    """
+    Fixed-grid dataset (common Npix for all samples).
+    Returns (x, y) if cell_ids is None, else (x, y, cell_ids).
+
+    x: (B, C, Npix)
+    y: (B, C_out or 1, Npix) or class indices depending on task
+    cell_ids: (Npix,) or (B, Npix)
+    """
+    def __init__(self,
+                 x: torch.Tensor,
+                 y: torch.Tensor,
+                 cell_ids: Optional[Union[np.ndarray, torch.Tensor]] = None,
+                 dtype: torch.dtype = torch.float32):
+        x = torch.as_tensor(x, dtype=dtype)
+        y = torch.as_tensor(y, dtype=dtype if y.ndim == 3 else torch.long)
+        assert x.shape[0] == y.shape[0], "x and y must share batch size"
+        self.x, self.y = x, y
+        self._has_cids = cell_ids is not None
+        if self._has_cids:
+            c = torch.as_tensor(cell_ids, dtype=torch.long)
+            if c.ndim == 1:  # broadcast single (Npix,) to (B, Npix)
+                c = c.unsqueeze(0).expand(x.shape[0], -1)
+            assert c.shape == (x.shape[0], x.shape[2]), "cell_ids must be (B,Npix) or (Npix,)"
+            self.cids = c
+        else:
+            self.cids = None
+
+    def __len__(self) -> int: return self.x.shape[0]
+
+    def __getitem__(self, i: int):
+        if self._has_cids:
+            return self.x[i], self.y[i], self.cids[i]
+        return self.x[i], self.y[i]
+
+
+class VarLenHealpixDataset(Dataset):
+    """
+    Variable-length per-sample dataset.
+
+    x_list[b]: (C, Npix_b) or (1, C, Npix_b)
+    y_list[b]: (C_out or 1, Npix_b) or (1, C_out, Npix_b)  (regression/segmentation targets)
+               For multi-class segmentation with CrossEntropyLoss, you may pass
+               class indices of shape (Npix_b,) or (1, Npix_b) (we’ll squeeze later).
+    cids_list[b]: (Npix_b,) or None
+    """
+    def __init__(self,
+                 x_list: List[Union[np.ndarray, torch.Tensor]],
+                 y_list: List[Union[np.ndarray, torch.Tensor]],
+                 cids_list: Optional[List[Union[np.ndarray, torch.Tensor]]] = None,
+                 dtype: torch.dtype = torch.float32):
+        assert len(x_list) == len(y_list), "x_list and y_list must have the same length"
+        self.x = [torch.as_tensor(x, dtype=dtype) for x in x_list]
+        # y can be float (regression) or long (class indices); we’ll coerce later per task
+        self.y = [torch.as_tensor(y) for y in y_list]
+        if cids_list is not None:
+            assert len(cids_list) == len(x_list), "cids_list must match x_list length"
+            self.c = [torch.as_tensor(c, dtype=torch.long) for c in cids_list]
+        else:
+            self.c = None
+
+    def __len__(self) -> int: return len(self.x)
+
+    def __getitem__(self, i: int):
+        ci = None if self.c is None else self.c[i]
+        return self.x[i], self.y[i], ci
+
+from torch.utils.data import Dataset, DataLoader
+
+class VarLenHealpixDataset(Dataset):
+    """
+    x_list: list of (C, Npix_b) tensors or arrays
+    y_list: list of (C_out or 1, Npix_b) tensors or arrays
+    cids_list: optional list of (Npix_b,) arrays
+    """
+    def __init__(self, x_list, y_list, cids_list=None, dtype=torch.float32):
+        assert len(x_list) == len(y_list)
+        self.x = [torch.as_tensor(x, dtype=dtype) for x in x_list]
+        self.y = [torch.as_tensor(y, dtype=dtype) for y in y_list]
+        self.c = None
+        if cids_list is not None:
+            assert len(cids_list) == len(x_list)
+            self.c = [np.asarray(c) for c in cids_list]
+
+    def __len__(self): return len(self.x)
+
+    def __getitem__(self, i):
+        if self.c is None:
+            return self.x[i], self.y[i], None
+        return self.x[i], self.y[i], self.c[i]
+
+def varlen_collate(batch):
+    # Just return lists; do not stack.
+    xs, ys, cs = zip(*batch)  # tuples of length B
+    # keep None if all Nones, else list
+    c_out = None if all(c is None for c in cs) else list(cs)
+    return list(xs), list(ys), c_out
+
+def varlen_collate(batch):
+    """
+    Collate for variable-length samples: keep lists, do NOT stack.
+    Returns lists: xs, ys, cs (cs can be None).
+    """
+    xs, ys, cs = zip(*batch)
+    c_out = None if all(c is None for c in cs) else list(cs)
+    return list(xs), list(ys), c_out
+
+
+# ---------------------------
+# Training function
+# ---------------------------
 
 def fit(
-        model: HealpixUNet,
-        x_train: torch.Tensor | np.ndarray,
-        y_train: torch.Tensor | np.ndarray,
+        model,
+        x_train: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        y_train: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
         *,
+        cell_ids_train: Optional[Union[np.ndarray, torch.Tensor, List[Union[np.ndarray, torch.Tensor]]]] = None,
         n_epoch: int = 10,
         view_epoch: int = 10,
         batch_size: int = 16,
         lr: float = 1e-3,
         weight_decay: float = 0.0,
-        clip_grad_norm: float | None = None,
+        clip_grad_norm: Optional[float] = None,
         verbose: bool = True,
-        optimizer:  Literal['ADAM', 'LBFGS'] = 'LBFGS',
+        optimizer: Literal['ADAM', 'LBFGS'] = 'ADAM',
 ) -> dict:
-    """Train helper for the torch-ified HEALPix U-Net.
-
-    Optimizes all registered parameters (kernels + BN affine) with Adam on MSE for regression,
-    or CrossEntropy/BCE for segmentation.
-
-    Device policy
-    -------------
-    Uses the model's probed runtime device (CUDA if Foscat conv works there; otherwise CPU).
     """
-    import numpy as _np
-    from torch.utils.data import TensorDataset, DataLoader
+    Train helper that supports:
+      - Fixed-grid tensors (B,C,N) with optional (B,N) or (N,) cell_ids.
+      - Variable-length lists: x=[(C,N_b)], y=[...], cell_ids=[(N_b,)], returning per-sample grids.
 
-    # Ensure model is on its runtime device (already probed in __init__)
-    model.to(model.runtime_device)
+    ADAM: standard minibatch update.
+    LBFGS: uses a closure that sums losses over the current (variable-length) mini-batch.
 
-    def _to_t(x):
-        if isinstance(x, torch.Tensor):
-            return x.float().to(model.runtime_device)
-        return torch.from_numpy(_np.asarray(x)).float().to(model.runtime_device)
-
-    x_t = _to_t(x_train)
-    y_t = _to_t(y_train)
-
-    # choose loss
-    if model.task == 'regression':
-        criterion = nn.MSELoss()
+    Notes
+    -----
+    - For segmentation with multiple classes, pass integer class targets for y:
+      fixed-grid: (B, N) int64; variable-length: each y[b] of shape (N_b,) or (1,N_b).
+    - For regression, pass float targets with the same (C_out, N) channeling.
+    """
+    device = model.runtime_device if hasattr(model, "runtime_device") else (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+    model.to(device)
+
+    # Detect variable-length mode
+    varlen_mode = isinstance(x_train, (list, tuple))
+
+    # ----- Build DataLoader
+    if not varlen_mode:
+        # Fixed-grid path
+        x_t = torch.as_tensor(x_train, dtype=torch.float32, device=device)
+        y_is_class = (model.task != 'regression' and getattr(model, "out_channels", 1) > 1)
+        y_dtype = torch.long if y_is_class and (not torch.is_tensor(y_train) or y_train.ndim != 3) else torch.float32
+        y_t = torch.as_tensor(y_train, dtype=y_dtype, device=device)
+
+        if cell_ids_train is None:
+            ds = TensorDataset(x_t, y_t)
+            loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+            with_cell_ids = False
+        else:
+            ds = HealpixDataset(x_t, y_t, cell_ids=cell_ids_train, dtype=torch.float32)
+            loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+            with_cell_ids = True
+    else:
+        # Variable-length path
+        ds = VarLenHealpixDataset(x_train, y_train, cids_list=cell_ids_train, dtype=torch.float32)
+        loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False, collate_fn=varlen_collate)
+        with_cell_ids = cell_ids_train is not None
+
+    # ----- Loss
+    if getattr(model, "task", "regression") == 'regression':
+        criterion = nn.MSELoss(reduction='mean')
+        seg_multiclass = False
     else:
-        if model.out_channels == 1:
-            criterion = nn.BCEWithLogitsLoss() if model.final_activation == 'none' else nn.BCELoss()
+        # segmentation
+        if getattr(model, "out_channels", 1) == 1:
+            # binary
+            # assume model head returns logits if final_activation == 'none'
+            criterion = nn.BCEWithLogitsLoss() if getattr(model, "final_activation", "none") == 'none' else nn.BCELoss()
+            seg_multiclass = False
         else:
-            if y_t.dim() == 3:
-                y_t = y_t.argmax(dim=1)
             criterion = nn.CrossEntropyLoss()
+            seg_multiclass = True
 
-    ds = TensorDataset(x_t, y_t)
-    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
-
-    if optimizer=='ADAM':
+    # ----- Optimizer
+    if optimizer.upper() == 'ADAM':
         optim = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
-        l_n_epoch=n_epoch
-        n_inter=1
+        outer = n_epoch
+        inner = 1
     else:
-        optim = torch.optim.LBFGS(model.parameters(), lr=lr, max_iter=20, history_size=n_epoch*5, line_search_fn="strong_wolfe")
-        l_n_epoch=n_epoch//20
-        n_inter=20
+        optim = torch.optim.LBFGS(model.parameters(), lr=lr, max_iter=20,
+                                  history_size=max(10, n_epoch * 5), line_search_fn="strong_wolfe")
+        # emulate "epochs" with multiple inner LBFGS steps
+        outer = max(1, n_epoch // 20)
+        inner = 20
 
+    # ----- Training loop
     history: List[float] = []
     model.train()
 
-    for epoch in range(l_n_epoch):
-        for k in range(n_inter):
-            epoch_loss = 0.0
-            n_samples = 0
-            for xb, yb in loader:
-                # LBFGS a besoin d'un closure qui recalcule loss et gradients
-                if isinstance(optim, torch.optim.LBFGS):
-                    def closure():
+    for epoch in range(outer):
+        for _ in range(inner):
+            epoch_loss, n_samples = 0.0, 0
+
+            for batch in loader:
+                if not varlen_mode:
+                    # -------- fixed-grid
+                    if with_cell_ids:
+                        xb, yb, cb = batch
+                        cb_np = cb.detach().cpu().numpy()
+                    else:
+                        xb, yb = batch
+                        cb_np = None
+
+                    xb = xb.to(device, dtype=torch.float32, non_blocking=True)
+                    # y type: float for regression or binary; long for CrossEntropy
+                    yb = yb.to(device, non_blocking=True)
+
+                    if isinstance(optim, torch.optim.LBFGS):
+                        def closure():
+                            optim.zero_grad(set_to_none=True)
+                            preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
+                            loss = criterion(preds, yb)
+                            loss.backward()
+                            return loss
+                        _ = optim.step(closure)
+                        with torch.no_grad():
+                            preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
+                            loss = criterion(preds, yb)
+                    else:
                         optim.zero_grad(set_to_none=True)
-                        preds = model(xb)
+                        preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
                         loss = criterion(preds, yb)
                         loss.backward()
-                        return loss
+                        if clip_grad_norm is not None:
+                            nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
+                        optim.step()
 
-                    _ = optim.step(closure)  # LBFGS appelle plusieurs fois le closure
-                    # on recalcule la loss finale pour l’agg (sans gradient)
-                    with torch.no_grad():
-                        preds = model(xb)
-                        loss = criterion(preds, yb)
+                    bs = xb.shape[0]
+                    epoch_loss += float(loss.item()) * bs
+                    n_samples += bs
 
                 else:
-                    optim.zero_grad(set_to_none=True)
-                    preds = model(xb)
-                    loss = criterion(preds, yb)
-                    loss.backward()
-                    if clip_grad_norm is not None:
-                        nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
-                    optim.step()
-
-                bs = xb.shape[0]
-                epoch_loss += loss.item() * bs
-                n_samples += bs
+                    # -------- variable-length (lists)
+                    xs, ys, cs = batch  # lists
+
+                    def _prep_xyc(i):
+                        # x_i : (C, N_i)  -> (1, C, N_i)
+                        xb = torch.as_tensor(xs[i], device=device, dtype=torch.float32)
+                        if xb.dim() == 2:
+                            xb = xb.unsqueeze(0)
+                        elif xb.dim() != 3 or xb.shape[0] != 1:
+                            raise ValueError("Each x[i] must be (C,N) or (1,C,N)")
+
+                        # y_i :
+                        yb = torch.as_tensor(ys[i], device=device)
+                        if seg_multiclass:
+                            # class indices: (N_i,) ou (1, N_i)
+                            if yb.dim() == 2 and yb.shape[0] == 1:
+                                yb = yb.squeeze(0)           # -> (N_i,)
+                            elif yb.dim() != 1:
+                                raise ValueError("For multiclass CE, y[i] must be (N,) or (1,N)")
+                            # le critère CE recevra (1,C_out,N_i) et (N_i,)
+                        else:
+                            # régression / binaire: cible de forme (1, C_out, N_i)
+                            if yb.dim() == 2:
+                                yb = yb.unsqueeze(0)
+                            elif yb.dim() != 3 or yb.shape[0] != 1:
+                                raise ValueError("For regression/binary, y[i] must be (C_out,N) or (1,C_out,N)")
+
+                        # cell_ids : (N_i,) -> (1, N_i) en numpy (le forward les attend en np.ndarray)
+                        if cs is None or cs[i] is None:
+                            cb_np = None
+                        else:
+                            c = cs[i].detach().cpu().numpy() if torch.is_tensor(cs[i]) else np.asarray(cs[i])
+                            if c.ndim == 1:
+                                c = c[None, :]               # -> (1, N_i)
+                            cb_np = c
+                        return xb, yb, cb_np
+
+                    if isinstance(optim, torch.optim.LBFGS):
+                        def closure():
+                            optim.zero_grad(set_to_none=True)
+                            total = 0.0
+                            for i in range(len(xs)):
+                                xb, yb, cb_np = _prep_xyc(i)
+                                preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
+                                # adapter la cible à la sortie
+                                if seg_multiclass:
+                                    loss_i = criterion(preds, yb)       # preds: (1,C_out,N_i), yb: (N_i,)
+                                else:
+                                    loss_i = criterion(preds, yb)       # preds: (1,C_out,N_i), yb: (1,C_out,N_i)
+                                loss_i.backward()
+                                total += float(loss_i.item())
+                            # retourner un scalaire Tensor pour LBFGS
+                            return torch.tensor(total / max(1, len(xs)), device=device, dtype=torch.float32)
+
+                        _ = optim.step(closure)
+                        # logging (sans grad)
+                        with torch.no_grad():
+                            total = 0.0
+                            for i in range(len(xs)):
+                                xb, yb, cb_np = _prep_xyc(i)
+                                preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
+                                if seg_multiclass:
+                                    loss_i = criterion(preds, yb)
+                                else:
+                                    loss_i = criterion(preds, yb)
+                                total += float(loss_i.item())
+                            loss_val = total / max(1, len(xs))
+                    else:
+                        optim.zero_grad(set_to_none=True)
+                        total = 0.0
+                        for i in range(len(xs)):
+                            xb, yb, cb_np = _prep_xyc(i)
+                            preds = model(xb, cell_ids=cb_np) if cb_np is not None else model(xb)
+                            if seg_multiclass:
+                                loss_i = criterion(preds, yb)
+                            else:
+                                loss_i = criterion(preds, yb)
+                            loss_i.backward()
+                            total += float(loss_i.item())
+                        if clip_grad_norm is not None:
+                            nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
+                        optim.step()
+                        loss_val = total / max(1, len(xs))
+
+                    epoch_loss += loss_val * max(1, len(xs))
+                    n_samples  += max(1, len(xs))
+
+
 
             epoch_loss /= max(1, n_samples)
             history.append(epoch_loss)
-            if verbose and ((epoch*n_inter+k+1)%view_epoch==0 or epoch*n_inter+k==0):
-                print(f"[epoch {epoch*n_inter+k+1}/{l_n_epoch*n_inter}] loss={epoch_loss:.6f}")
+            # print every view_epoch logical step
+            if verbose and ((len(history) % view_epoch == 0) or (len(history) == 1)):
+                print(f"[epoch {len(history)}] loss={epoch_loss:.6f}")
 
     return {"loss": history}
-
-
-__all__ = ["HealpixUNet", "fit"]