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

foscat/healpix_unet_torch.py

@@ -0,0 +1,717 @@
+"""
+HEALPix U-Net (nested) with Foscat + PyTorch niceties
+----------------------------------------------------
+GPU by default (when available), with graceful CPU fallback if Foscat ops are CPU-only.
+
+- 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
+
+Shape convention: (B, C, Npix)
+
+Requirements: foscat (scat_cov.funct + HOrientedConvol.Convol_torch must be differentiable on torch tensors)
+"""
+from __future__ import annotations
+from typing import List, Optional, Literal, Tuple
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import foscat.scat_cov as sc
+import foscat.HOrientedConvol as ho
+import matplotlib.pyplot as plt
+
+
+class HealpixUNet(nn.Module):
+    """U-Net-like architecture on the HEALPix sphere using Foscat oriented convolutions.
+
+    Parameters
+    ----------
+    in_nside : int
+        Input HEALPix nside (nested scheme).
+    n_chan_in : int
+        Number of input channels.
+    chanlist : list[int]
+        Channels per encoder level (depth = len(chanlist)). Example: [16, 32, 64].
+    cell_ids : np.ndarray
+        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.
+    task : {'regression','segmentation'}, default 'regression'
+        Chooses the head and default activation.
+    out_channels : int, default 1
+        Number of output channels (e.g. num_classes for segmentation).
+    final_activation : {'none','sigmoid','softmax'} | None
+        If None, uses sensible default per task: 'none' for regression, 'softmax' for segmentation (multi-class),
+        'sigmoid' for segmentation when out_channels==1.
+    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.
+    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.
+
+    Notes
+    -----
+    - 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``.
+    - Foscat ops device is auto-probed to avoid CPU/CUDA mismatches.
+    """
+
+    def __init__(
+            self,
+            *,
+            in_nside: int,
+            n_chan_in: int,
+            chanlist: List[int],
+            cell_ids: np.ndarray,
+            KERNELSZ: int = 3,
+            task: Literal['regression', 'segmentation'] = 'regression',
+            out_channels: int = 1,
+            final_activation: Optional[Literal['none', 'sigmoid', 'softmax']] = None,
+            device: Optional[torch.device | str] = None,
+            prefer_foscat_gpu: bool = True,
+            dtype: Literal['float32','float64'] = 'float32'
+    ) -> None:
+        super().__init__()
+
+        self.dtype=dtype
+        if dtype=='float32':
+            self.np_dtype=np.float32
+        else:
+            self.np_dtype=np.float64
+            
+        if cell_ids is None:
+            raise ValueError("cell_ids must be provided for the finest resolution.")
+        if len(chanlist) == 0:
+            raise ValueError("chanlist must be non-empty (depth >= 1).")
+
+        self.in_nside = int(in_nside)
+        self.n_chan_in = int(n_chan_in)
+        self.chanlist = list(map(int, chanlist))
+        self.KERNELSZ = int(KERNELSZ)
+        self.task = task
+        self.out_channels = int(out_channels)
+        self.prefer_foscat_gpu = bool(prefer_foscat_gpu)
+
+        # Choose default final activation if not given
+        if final_activation is None:
+            if task == 'regression':
+                self.final_activation = 'none'
+            else:  # segmentation
+                self.final_activation = 'sigmoid' if out_channels == 1 else 'softmax'
+        else:
+            self.final_activation = final_activation
+
+        # Resolve preferred device
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = torch.device(device)
+
+        # Foscat functional wrapper (backend + grade ops)
+        self.f = sc.funct(KERNELSZ=self.KERNELSZ)
+
+        # ---------- Build multi-resolution bookkeeping ----------
+        depth = len(self.chanlist)
+        self.l_cell_ids: List[np.ndarray] = [None] * (depth + 1)  # per encoder level + bottom
+        self.l_cell_ids[0] = np.asarray(cell_ids)
+
+        enc_nsides: List[int] = [self.in_nside]
+        current_nside = self.in_nside
+
+        # ---------- Oriented convolutions per level (encoder & decoder) ----------
+        self.hconv_enc: List[ho.HOrientedConvol] = []
+        self.hconv_dec: List[ho.HOrientedConvol] = []
+        
+        # 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()
+            
+            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
+            )
+            self.l_cell_ids[l + 1] = self.f.backend.to_numpy(next_ids)
+            current_nside //= 2
+            enc_nsides.append(current_nside)
+
+        # encoder conv weights and BN
+        self.enc_w1 = nn.ParameterList()
+        self.enc_bn1 = nn.ModuleList()
+        self.enc_w2 = nn.ParameterList()
+        self.enc_bn2 = nn.ModuleList()
+
+        self.enc_nsides = enc_nsides  # [in, in/2, ..., in/2**depth]
+        
+        inC = self.n_chan_in
+        for l, outC in enumerate(self.chanlist):
+
+            # 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))
+            self.enc_w1.append(nn.Parameter(w1))
+            self.enc_bn1.append(nn.BatchNorm1d(outC))
+
+            # 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))
+            self.enc_w2.append(nn.Parameter(w2))
+            self.enc_bn2.append(nn.BatchNorm1d(outC))
+
+            inC = outC  # next level input channels
+
+        # decoder conv weights and BN (mirrored levels)
+        self.dec_w1 = nn.ParameterList()
+        self.dec_bn1 = nn.ModuleList()
+        self.dec_w2 = nn.ParameterList()
+        self.dec_bn2 = nn.ModuleList()
+
+        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.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
+
+            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))
+            self.dec_w1.append(nn.Parameter(w1))
+            self.dec_bn1.append(nn.BatchNorm1d(outC_dec))
+
+            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))
+            self.dec_w2.append(nn.Parameter(w2))
+            self.dec_bn2.append(nn.BatchNorm1d(outC_dec))
+
+        # 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
+
+        # ---- Decide runtime device (probe Foscat on CUDA, else CPU) ----
+        self.runtime_device = self._probe_and_set_runtime_device(self.device)
+
+    # -------------------------- 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."""
+        for name, val in list(vars(hc).items()):
+            try:
+                if torch.is_tensor(val):
+                    setattr(hc, name, val.to(device))
+                elif isinstance(val, (list, tuple)) and val and torch.is_tensor(val[0]):
+                    setattr(hc, name, type(val)([v.to(device) for v in val]))
+            except Exception:
+                # silently ignore non-tensor or protected attributes
+                pass
+
+    @torch.no_grad()
+    def _probe_and_set_runtime_device(self, preferred: torch.device) -> torch.device:
+        """Try to run a tiny Foscat conv on preferred device; fallback to CPU if it fails."""
+        if preferred.type == 'cuda' and self.prefer_foscat_gpu:
+            try:
+                # move module params/buffers first
+                super().to(preferred)
+                # move Foscat operator internals
+                for hc in self.hconv_enc + self.hconv_dec + [self.head_hconv]:
+                    self._move_hconv_tensors(hc, preferred)
+                # dry run on level 0
+                npix0 = int(len(self.l_cell_ids[0]))
+                x_try = torch.zeros(1, self.n_chan_in, npix0, device=preferred)
+                y_try = self.hconv_enc[0].Convol_torch(x_try, self.enc_w1[0])
+                # success -> stay on CUDA
+                self._foscat_device = preferred
+                return preferred
+            except Exception as e:
+                # fallback to CPU; keep error for info
+                self._gpu_probe_error = repr(e)
+                pass
+        # CPU fallback
+        cpu = torch.device('cpu')
+        super().to(cpu)
+        for hc in self.hconv_enc + self.hconv_dec + [self.head_hconv]:
+            self._move_hconv_tensors(hc, cpu)
+        self._foscat_device = cpu
+        return cpu
+
+    def set_device(self, device: torch.device | str) -> torch.device:
+        """Request a (re)device; will probe Foscat and return the actual runtime device used."""
+        device = torch.device(device)
+        self.device = device
+        self.runtime_device = self._probe_and_set_runtime_device(device)
+        return self.runtime_device
+
+    # -------------------------- forward --------------------------
+    def forward(self, x: torch.Tensor,cell_ids: Optional[np.ndarray ] = None) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor, shape (B, C_in, Npix)
+            Input tensor on `in_nside` grid.
+        cell_ids : np.ndarray (B, Npix) optional, use another cell_ids than the initial one.
+                   if None use the initial cell_ids.
+        """
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("Input must be a torch.Tensor")
+        if x.dim() != 3:
+            raise ValueError("Input must be (B, C, Npix)")
+
+        # Ensure input lives on the runtime (probed) device
+        x = x.to(self.runtime_device)
+
+        B, C, N = x.shape
+        if C != self.n_chan_in:
+            raise ValueError(f"Expected {self.n_chan_in} input channels, got {C}")
+
+        # Encoder
+        skips: List[torch.Tensor] = []
+        l_data = x
+        current_nside = self.in_nside
+        l_cell_ids=cell_ids
+        
+        if cell_ids is not None:
+            t_cell_ids={}
+            t_cell_ids[0]=l_cell_ids
+        else:
+            t_cell_ids=self.l_cell_ids
+            
+        for l, outC in enumerate(self.chanlist):
+            # conv1 + BN + ReLU
+            l_data = self.hconv_enc[l].Convol_torch(l_data,
+                                                    self.enc_w1[l],
+                                                    cell_ids=l_cell_ids)
+            l_data = self.enc_bn1[l](l_data)
+            l_data = F.relu(l_data, inplace=True)
+
+            # conv2 + BN + ReLU
+            l_data = self.hconv_enc[l].Convol_torch(l_data,
+                                                    self.enc_w2[l],
+                                                    cell_ids=l_cell_ids)
+            l_data = self.enc_bn2[l](l_data)
+            l_data = F.relu(l_data, inplace=True)
+
+            # save skip at this resolution
+            skips.append(l_data)
+
+            # 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
+                )
+                if cell_ids is not None:
+                    t_cell_ids[l+1]=l_cell_ids
+                else:
+                    l_cell_ids=None
+                    
+                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
+
+            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(
+                    l_data,
+                    cell_ids=t_cell_ids[level + 1],  # source (coarser) ids
+                    o_cell_ids=t_cell_ids[level],     # target (finer) ids
+                    nside=src_nside,
+                )
+                
+                if isinstance(l_data, torch.Tensor) and l_data.device != self.runtime_device:
+                    l_data = l_data.to(self.runtime_device)
+
+            # concat with skip features at this resolution
+            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
+            hc = self.hconv_dec[d]
+            l_data = hc.Convol_torch(l_data,
+                                     self.dec_w1[d],
+                                     cell_ids=l_cell_ids)
+            
+            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 = 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)
+        if self.final_activation == 'sigmoid':
+            out = torch.sigmoid(out)
+        elif self.final_activation == 'softmax':
+            out = torch.softmax(out, dim=1)
+        return out
+
+    # -------------------------- 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]))
+                
+        return torch.cat(outs, dim=0)
+
+    # -----------------------------
+    # Kernel extraction & plotting
+    # -----------------------------
+    def _arch_shapes(self):
+        """Return expected (in_c, out_c) per conv for encoder/decoder.
+
+        Returns
+        -------
+        enc_shapes : list[tuple[tuple[int,int], tuple[int,int]]]
+            For each level `l`, ((in1, out1), (in2, out2)) for the two encoder convs.
+        dec_shapes : list[tuple[tuple[int,int], tuple[int,int]]]
+            For each level `l`, ((in1, out1), (in2, out2)) for the two decoder convs.
+        """
+        nlayer = len(self.chanlist)
+        enc_shapes = []
+        l_chan = self.n_chan_in
+        for l in range(nlayer):
+            enc_shapes.append(((l_chan, self.chanlist[l]), (self.chanlist[l], self.chanlist[l])))
+            l_chan = self.chanlist[l] + 1
+
+        dec_shapes = []
+        l_chan = self.chanlist[-1] + 1
+        for l in range(nlayer):
+            in1 = l_chan + 1
+            out2 = 1 + (self.chanlist[nlayer - 1 - l] if (nlayer - 1 - l) > 0 else 0)
+            dec_shapes.append(((in1, in1), (in1, out2)))
+            l_chan = out2
+        return enc_shapes, dec_shapes
+
+    def extract_kernels(self, stage: str = "encoder", layer: int = 0, conv: int = 0):
+        """Extract raw convolution kernels for a given stage/level/conv.
+
+        Parameters
+        ----------
+        stage : {"encoder", "decoder"}
+            Which part of the network to inspect.
+        layer : int
+            Pyramid level (0 = finest encoder level / bottommost decoder level).
+        conv : int
+            0 for the first conv at that level, 1 for the second conv.
+
+        Returns
+        -------
+        np.ndarray
+            Array of shape (in_c, out_c, K, K) containing the spatial kernels.
+        """
+        assert stage in {"encoder", "decoder"}
+        assert conv in {0, 1}
+        K = self.KERNELSZ
+        enc_shapes, dec_shapes = self._arch_shapes()
+
+        if stage == "encoder":
+            if conv==0:
+                w = self.enc_w1[layer]
+            else:
+                w = self.enc_w2[layer]
+        else:
+            if conv==0:
+                w = self.dec_w1[layer]
+            else:
+                w = self.dec_w2[layer]
+
+        w_np = self.f.backend.to_numpy(w.detach())
+        return w_np.reshape(w.shape[0],w.shape[1],K,K)
+
+    def plot_kernels(
+        self,
+        stage: str = "encoder",
+        layer: int = 0,
+        conv: int = 0,
+        fixed: str = "in",
+        index: int = 0,
+        max_tiles: int = 16,
+    ):
+        """Quick visualization of kernels on a grid using matplotlib.
+
+        Parameters
+        ----------
+        stage : {"encoder", "decoder"}
+            Which tower to visualize.
+        layer : int
+            Level to visualize.
+        conv : int
+            0 or 1: first or second conv in the level.
+        fixed : {"in", "out"}
+            If "in", show kernels for a fixed input channel across many outputs.
+            If "out", show kernels for a fixed output channel across many inputs.
+        index : int
+            Channel index to fix (according to `fixed`).
+        max_tiles : int
+            Maximum number of tiles to display.
+        """
+        import math
+        import matplotlib.pyplot as plt
+
+        W = self.extract_kernels(stage=stage, layer=layer, conv=conv)
+        ic, oc, K,_ = W.shape
+
+        if fixed == "in":
+            idx = min(index, ic - 1)
+            tiles = [W[idx, j] for j in range(oc)]
+            title = f"{stage} L{layer} C{conv} | in={idx}"
+        else:
+            idx = min(index, oc - 1)
+            tiles = [W[i, idx] for i in range(ic)]
+            title = f"{stage} L{layer} C{conv} | out={idx}"
+
+        tiles = tiles[:max_tiles]
+        n = len(tiles)
+        cols = int(math.ceil(math.sqrt(n)))
+        rows = int(math.ceil(n / cols))
+
+        plt.figure(figsize=(2.5 * cols, 2.5 * rows))
+        for i, ker in enumerate(tiles, 1):
+            ax = plt.subplot(rows, cols, i)
+            ax.imshow(ker)
+            ax.set_xticks([])
+            ax.set_yticks([])
+        plt.suptitle(title)
+        plt.tight_layout()
+        plt.show()
+
+# -----------------------------
+# Unit tests (smoke tests)
+# -----------------------------
+# Run with:  python UNET.py  (or)  python UNET.py -q  for quieter output
+# These tests assume Foscat and its dependencies are installed.
+
+
+def _dummy_cell_ids(nside: int) -> np.ndarray:
+    """Return a simple identity mapping for HEALPix nested pixel IDs.
+
+    Notes
+    -----
+    Replace with your pipeline's real `cell_ids` if you have a precomputed
+    mapping consistent with Foscat/HEALPix nested ordering.
+    """
+    return np.arange(12 * nside * nside, dtype=np.int64)
+
+
+if __name__ == "__main__":
+    import unittest
+
+    class TestUNET(unittest.TestCase):
+        """Lightweight smoke tests for shape and parameter plumbing."""
+
+        def setUp(self):
+            self.nside = 4  # small grid for fast tests (npix = 192)
+            self.chanlist = [4, 8]  # two-level encoder/decoder
+            self.batch = 2
+            self.channels = 1
+            self.npix = 12 * self.nside * self.nside
+            self.cell_ids = _dummy_cell_ids(self.nside)
+            self.net = UNET(
+                in_nside=self.nside,
+                n_chan_in=self.channels,
+                chanlist=self.chanlist,
+                cell_ids=self.cell_ids,
+            )
+
+        def test_forward_shape(self):
+            # random input
+            x = np.random.randn(self.batch, self.channels, self.npix).astype(self.np_dtype)
+            x = self.net.f.backend.bk_cast(x)
+            y = self.net.eval(x)
+            # expected output: same npix, 1 channel at the very top
+            self.assertEqual(y.shape[0], self.batch)
+            self.assertEqual(y.shape[1], 1)
+            self.assertEqual(y.shape[2], self.npix)
+            # sanity: no NaNs
+            y_np = self.net.f.backend.to_numpy(y)
+            self.assertFalse(np.isnan(y_np).any())
+
+        def test_param_roundtrip_and_determinism(self):
+            x = np.random.randn(self.batch, self.channels, self.npix).astype(self.np_dtype)
+            x = self.net.f.backend.bk_cast(x)
+
+            # forward twice -> identical outputs with fixed params
+            y1 = self.net.eval(x)
+            y2 = self.net.eval(x)
+            y1_np = self.net.f.backend.to_numpy(y1)
+            y2_np = self.net.f.backend.to_numpy(y2)
+            np.testing.assert_allclose(y1_np, y2_np, rtol=0, atol=0)
+
+            # perturb parameters -> output should (very likely) change
+            p = self.net.get_param()
+            p_np = self.net.f.backend.to_numpy(p).copy()
+            if p_np.size > 0:
+                p_np[0] += 1.0
+                self.net.set_param(p_np)
+                y3 = self.net.eval(x)
+                y3_np = self.net.f.backend.to_numpy(y3)
+                with self.assertRaises(AssertionError):
+                    np.testing.assert_allclose(y1_np, y3_np, rtol=0, atol=0)
+
+    unittest.main()
+
+def fit(
+        model: HealpixUNet,
+        x_train: torch.Tensor | np.ndarray,
+        y_train: torch.Tensor | np.ndarray,
+        *,
+        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,
+        verbose: bool = True,
+        optimizer:  Literal['ADAM', 'LBFGS'] = 'LBFGS',
+) -> 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
+
+    # Ensure model is on its runtime device (already probed in __init__)
+    model.to(model.runtime_device)
+
+    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()
+    else:
+        if model.out_channels == 1:
+            criterion = nn.BCEWithLogitsLoss() if model.final_activation == 'none' else nn.BCELoss()
+        else:
+            if y_t.dim() == 3:
+                y_t = y_t.argmax(dim=1)
+            criterion = nn.CrossEntropyLoss()
+
+    ds = TensorDataset(x_t, y_t)
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+
+    if optimizer=='ADAM':
+        optim = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+        l_n_epoch=n_epoch
+        n_inter=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
+
+    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():
+                        optim.zero_grad(set_to_none=True)
+                        preds = model(xb)
+                        loss = criterion(preds, yb)
+                        loss.backward()
+                        return loss
+
+                    _ = 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)
+
+                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
+
+            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}")
+
+    return {"loss": history}
+
+
+__all__ = ["HealpixUNet", "fit"]