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

foscat/planar_vit.py

@@ -0,0 +1,206 @@
+# healpix_unet_torch.py
+# (Planar Vision Transformer baseline for lat–lon grids)
+from __future__ import annotations
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------
+# Building blocks
+# ---------------------------
+
+class _MLP(nn.Module):
+    """ViT MLP: Linear -> GELU -> Dropout -> Linear -> Dropout."""
+    def __init__(self, dim: int, mlp_ratio: float = 4.0, drop: float = 0.1):
+        super().__init__()
+        hidden = int(dim * mlp_ratio)
+        self.fc1 = nn.Linear(dim, hidden)
+        self.fc2 = nn.Linear(hidden, dim)
+        self.act = nn.GELU()
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.drop(self.act(self.fc1(x)))
+        x = self.drop(self.fc2(x))
+        return x
+
+
+class _ViTBlock(nn.Module):
+    """
+    Transformer block (Pre-LN):
+      x = x + Drop(MHA(LN(x)))
+      x = x + Drop(MLP(LN(x)))
+    """
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, drop: float = 0.1):
+        super().__init__()
+        assert dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn  = nn.MultiheadAttention(dim, num_heads, dropout=drop, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp   = _MLP(dim, mlp_ratio, drop)
+        self.drop_path = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Multi-head self-attention
+        x = x + self.drop_path(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
+        # Feed-forward
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+# ---------------------------
+# Planar ViT (lat–lon images)
+# ---------------------------
+
+class PlanarViT(nn.Module):
+    """
+    Vision Transformer for 2D lat–lon grids (planar baseline).
+
+    Input : (B, C=T_in, H, W)
+    Output: (B, out_ch, H, W)   # dense per-pixel prediction
+
+    Pipeline
+    --------
+    1) Patch embedding via Conv2d(kernel_size=patch, stride=patch) -> embed_dim
+    2) Optional CLS token (disabled by default for dense output)
+    3) Learned positional embeddings (or none)
+    4) Stack of Transformer blocks
+    5) Linear head per token, then nearest upsample back to (H, W)
+
+    Notes
+    -----
+    - Keep H, W divisible by `patch`.
+    - For residual-of-persistence training (recommended for monthly SST):
+        pred = x[:, -1:, ...] + model(x)
+      and train the loss on `pred` vs target.
+    """
+    def __init__(
+        self,
+        in_ch: int,          # e.g., T_in months
+        H: int,
+        W: int,
+        *,
+        embed_dim: int = 384,
+        depth: int = 8,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        patch: int = 4,
+        out_ch: int = 1,
+        dropout: float = 0.1,
+        cls_token: bool = False,    # keep False for dense prediction
+        pos_embed: str = "learned", # or "none"
+    ):
+        super().__init__()
+        assert H % patch == 0 and W % patch == 0, "H and W must be divisible by patch"
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+
+        self.H, self.W = H, W
+        self.patch = patch
+        self.embed_dim = embed_dim
+        self.cls_token_enabled = bool(cls_token)
+        self.use_pos_embed = (pos_embed == "learned")
+
+        # 1) Patch embedding (Conv2d with stride=patch) → tokens
+        self.patch_embed = nn.Conv2d(in_ch, embed_dim, kernel_size=patch, stride=patch)
+
+        # 2) Token bookkeeping & positional embeddings
+        Hp, Wp = H // patch, W // patch
+        self.num_tokens = Hp * Wp + (1 if self.cls_token_enabled else 0)
+
+        if self.cls_token_enabled:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        else:
+            self.cls_token = None
+
+        if self.use_pos_embed:
+            self.pos_embed = nn.Parameter(torch.zeros(1, self.num_tokens, embed_dim))
+            nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        else:
+            self.pos_embed = None
+
+        # 3) Transformer encoder
+        self.blocks = nn.ModuleList([
+            _ViTBlock(embed_dim, num_heads, mlp_ratio=mlp_ratio, drop=dropout)
+            for _ in range(depth)
+        ])
+
+        # 4) Patch-wise head (token -> out_ch)
+        self.head = nn.Linear(embed_dim, out_ch)
+
+        # Store for unpatching
+        self.Hp, self.Wp = Hp, Wp
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (B, C, H, W)  with H,W fixed to construction-time H,W
+        returns: (B, out_ch, H, W)
+        """
+        B, C, H, W = x.shape
+        if (H != self.H) or (W != self.W):
+            raise ValueError(f"Input H,W must be ({self.H},{self.W}), got ({H},{W}).")
+
+        # Patch embedding → (B, E, Hp, Wp) → (B, Np, E)
+        z = self.patch_embed(x)                 # (B, E, Hp, Wp)
+        z = z.flatten(2).transpose(1, 2)        # (B, Np, E)
+
+        # Optional CLS
+        if self.cls_token_enabled:
+            cls = self.cls_token.expand(B, -1, -1)  # (B,1,E)
+            z = torch.cat([cls, z], dim=1)          # (B,1+Np,E)
+
+        # Positional embedding
+        if self.pos_embed is not None:
+            z = z + self.pos_embed[:, :z.shape[1], :]
+
+        # Transformer
+        for blk in self.blocks:
+            z = blk(z)                             # (B, N, E)
+
+        # Drop CLS for dense output
+        if self.cls_token_enabled:
+            tokens = z[:, 1:, :]                   # (B, Np, E)
+        else:
+            tokens = z
+
+        # Token head → (B, Np, out_ch) → (B, out_ch, Hp, Wp) → upsample to (H, W)
+        y_tok = self.head(tokens).transpose(1, 2)  # (B, out_ch, Np)
+        y = y_tok.reshape(B, -1, self.Hp, self.Wp) # (B, out_ch, Hp, Wp)
+        y = F.interpolate(y, scale_factor=self.patch, mode="nearest")
+        return y
+
+
+# ---------------------------
+# Utilities
+# ---------------------------
+
+def count_parameters(model: nn.Module) -> tuple[int, int]:
+    """Return (total_params, trainable_params)."""
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total, trainable
+
+
+# ---------------------------
+# Smoke test
+# ---------------------------
+
+if __name__ == "__main__":
+    # Example: T_in=6, grid 128x256, predict 1 channel
+    B, C, H, W = 2, 6, 128, 256
+    x = torch.randn(B, C, H, W)
+
+    model = PlanarViT(
+        in_ch=C, H=H, W=W,
+        embed_dim=384, depth=8, num_heads=12,
+        mlp_ratio=4.0, patch=4, out_ch=1, dropout=0.1,
+        cls_token=False, pos_embed="learned"
+    )
+    y = model(x)
+    tot, trn = count_parameters(model)
+    print("Output:", tuple(y.shape))
+    print("Params:", f"total={tot:,}", f"trainable={trn:,}")
+    

foscat/unet_2_d_from_healpix_params.py

@@ -0,0 +1,421 @@
+from __future__ import annotations
+from typing import List, Optional, Literal, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from contextlib import nullcontext
+
+class PlanarUNet(nn.Module):
+    """
+    U-Net 2D (images HxW) mirroring the parameterization of the HealpixUNet.
+
+    Key compat points with HealpixUNet:
+      - Same constructor fields: in_nside, n_chan_in, chanlist, KERNELSZ, task,
+        out_channels, final_activation, device, down_type, dtype, head_reduce.
+      - Two convs per level (encoder & decoder), GroupNorm + ReLU after each conv.
+      - Downsampling by factor 2 at each level; upsampling mirrors back.
+      - Head produces `out_channels` with optional BN and final activation.
+
+    Differences vs sphere version:
+      - Operates on regular 2D images of size (3*in_nside, 4*in_nside).
+      - Standard Conv2d instead of custom spherical stencil.
+      - No gauges (G=1 implicit) and no cell_ids.
+
+    Shapes
+    ------
+    Input  : (B, C_in,  3*in_nside, 4*in_nside)
+    Output : (B, C_out, 3*in_nside, 4*in_nside)
+
+    Constraints
+    -----------
+    `in_nside` must be divisible by 2**depth, where depth == len(chanlist).
+    """
+
+    def __init__(
+        self,
+        *,
+        in_nside: int,
+        n_chan_in: int,
+        chanlist: List[int],
+        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,
+        down_type: Optional[Literal['mean','max']] = 'max',
+        dtype: Literal['float32','float64'] = 'float32',
+        head_reduce: Literal['mean','learned'] = 'mean',  # kept for API symmetry
+    ) -> None:
+        super().__init__()
+
+        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.down_type = down_type
+        self.dtype = torch.float32 if dtype == 'float32' else torch.float64
+        self.head_reduce = head_reduce
+        
+        # default final activation consistent with HealpixUNet
+        if final_activation is None:
+            if task == 'regression':
+                self.final_activation = 'none'
+            else:
+                self.final_activation = 'sigmoid' if out_channels == 1 else 'softmax'
+        else:
+            self.final_activation = final_activation
+
+        # Resolve device
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = torch.device(device)
+
+        depth = len(self.chanlist)
+        # geometry
+        H0, W0 = 3 * self.in_nside, 4 * self.in_nside
+        # ensure divisibility by 2**depth
+        if (self.in_nside % (2 ** depth)) != 0:
+            raise ValueError(
+                f"in_nside={self.in_nside} must be divisible by 2**depth where depth={depth}"
+            )
+
+        padding = self.KERNELSZ // 2
+
+        # --- Encoder ---
+        enc_layers = []
+        inC = self.n_chan_in
+        self.skips_channels: List[int] = []
+        for outC in self.chanlist:
+            block = nn.Sequential(
+                nn.Conv2d(inC, outC, kernel_size=self.KERNELSZ, padding=padding, bias=False),
+                _norm_2d(outC, kind="group"),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(outC, outC, kernel_size=self.KERNELSZ, padding=padding, bias=False),
+                _norm_2d(outC, kind="group"),
+                nn.ReLU(inplace=True),
+            )
+            enc_layers.append(block)
+            inC = outC
+            self.skips_channels.append(outC)
+        self.encoder = nn.ModuleList(enc_layers)
+
+        # Pools
+        if self.down_type == 'max':
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        else:
+            self.pool = nn.AvgPool2d(kernel_size=2, stride=2)
+
+        # --- Decoder ---
+        dec_layers = []
+        upconvs = []
+        for l in reversed(range(depth)):
+            skipC = self.skips_channels[l]
+            upC = self.skips_channels[l + 1] if (l + 1) < depth else self.skips_channels[l]
+            inC_dec = upC + skipC
+            outC_dec = skipC
+
+            upconvs.append(
+                nn.ConvTranspose2d(upC, upC, kernel_size=2, stride=2)
+            )
+            dec_layers.append(
+                nn.Sequential(
+                    nn.Conv2d(inC_dec, outC_dec, kernel_size=self.KERNELSZ, padding=padding, bias=False),
+                    _norm_2d(outC_dec, kind="group"),
+                    nn.ReLU(inplace=True),
+                    nn.Conv2d(outC_dec, outC_dec, kernel_size=self.KERNELSZ, padding=padding, bias=False),
+                    _norm_2d(outC_dec, kind="group"),
+                    nn.ReLU(inplace=True),
+                )
+            )
+        self.upconvs = nn.ModuleList(upconvs)
+        self.decoder = nn.ModuleList(dec_layers)
+
+        # --- Head ---
+        head_inC = self.chanlist[0]
+        self.head_conv = nn.Conv2d(head_inC, self.out_channels, kernel_size=self.KERNELSZ, padding=padding)
+        self.head_bn = _norm_2d(self.out_channels, kind="group") if self.task == 'segmentation' else None
+
+        # optional learned mixer kept for API compatibility (no gauges here)
+        self.head_mixer = nn.Identity()
+
+        self.to(self.device, dtype=self.dtype)
+
+    def to_tensor(self,x):
+        return torch.tensor(x,device=self.device)
+    
+    def to_numpy(self,x):
+        if isinstance(x,np.ndarray):
+            return x
+        return x.cpu().numpy()
+    
+    # -------------------------- forward --------------------------
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, C_in, H, W) with H=3*in_nside, W=4*in_nside"""
+        if x.dim() != 4:
+            raise ValueError("Input must be (B, C, H, W)")
+        if x.shape[1] != self.n_chan_in:
+            raise ValueError(f"Expected {self.n_chan_in} input channels, got {x.shape[1]}")
+
+        x = x.to(self.device, dtype=self.dtype)
+
+        skips = []
+        z = x
+        for l, block in enumerate(self.encoder):
+            z = block(z)
+            skips.append(z)
+            if l < len(self.encoder) - 1:
+                z = self.pool(z)
+
+        # Decoder
+        for d, l in enumerate(reversed(range(len(self.chanlist)))):
+            if l < len(self.chanlist) - 1:
+                z = self.upconvs[d](z)
+                # pad if odd due to pooling/upsampling asymmetry (shouldn't happen given divisibility)
+                sh = skips[l].shape
+                if z.shape[-2:] != sh[-2:]:
+                    z = _pad_to_match(z, sh[-2], sh[-1])
+            z = torch.cat([skips[l], z], dim=1)
+            z = self.decoder[d](z)
+
+        y = self.head_conv(z)
+        if self.task == 'segmentation' and self.head_bn is not None:
+            y = self.head_bn(y)
+
+        if self.final_activation == 'sigmoid':
+            y = torch.sigmoid(y)
+        elif self.final_activation == 'softmax':
+            y = torch.softmax(y, dim=1)
+        return y
+
+    @torch.no_grad()
+    def predict(self, x: torch.Tensor, batch_size: int = 8) -> torch.Tensor:
+        self.eval()
+        outs = []
+        for i in range(0, x.shape[0], batch_size):
+            xb = x[i:i+batch_size]
+            outs.append(self.forward(xb))
+        return torch.cat(outs, dim=0)
+
+    @torch.no_grad()
+    def predict(
+            self,
+            x: torch.Tensor,
+            batch_size: int = 8,
+            *,
+            amp: bool = False,
+            out_device: Optional[str] = 'cpu',
+            out_dtype: Literal['float32','float16'] = 'float32',
+            show_pbar: bool = False,
+    ) -> torch.Tensor:
+        """Memory-safe prediction.
+        - Streams mini-batches avec torch.inference_mode() + AMP optionnel.
+        - Déplace chaque batch de sorties sur `out_device` (CPU par défaut) pour libérer la VRAM.
+        - Vérifie et explicite les erreurs de shape.
+        """
+        self.eval()
+
+        # --- checks & normalisation d'entrée ---
+        x = x if torch.is_tensor(x) else torch.as_tensor(x)
+        if x.ndim != 4:
+            raise ValueError(f"predict expects (N,C,H,W), got {tuple(getattr(x,'shape',()))}")
+        if x.shape[1] != self.n_chan_in:
+            raise ValueError(f"predict expected {self.n_chan_in} channels, got {x.shape[1]}")
+        n = int(x.shape[0])
+        if n == 0:
+            H, W = int(x.shape[-2]), int(x.shape[-1])
+            return torch.empty((0, self.out_channels, H, W), device=out_device or self.device)
+
+        # --- préparation ---
+        dtype_map = {'float32': torch.float32, 'float16': torch.float16}
+        out_dtype_t = dtype_map[out_dtype]
+        use_cuda = (self.device.type == 'cuda')
+        if use_cuda:
+            torch.backends.cudnn.benchmark = True
+
+        from math import ceil
+        nb = ceil(n / batch_size)
+        rng = range(0, n, batch_size)
+        if show_pbar:
+            try:
+                from tqdm import tqdm  # type: ignore
+                rng = tqdm(rng, total=nb, desc='predict')
+            except Exception:
+                pass
+
+        # --- inférence batch par batch ---
+        out_list: List[torch.Tensor] = []
+        with torch.inference_mode():
+            ctx = (torch.cuda.amp.autocast() if (amp and use_cuda) else nullcontext())
+            for i in rng:
+                xb = x[i:i+batch_size].to(self.device, dtype=self.dtype, non_blocking=True)
+                with ctx:
+                    yb = self.forward(xb)
+                # Déplacer la sortie vers l'appareil voulu (CPU par défaut)
+                yb = yb.to(out_device, dtype=out_dtype_t) if out_device is not None else yb.to(dtype=out_dtype_t)
+                out_list.append(yb)
+                del xb, yb
+                if use_cuda:
+                    torch.cuda.empty_cache()
+
+        if not out_list:
+            raise RuntimeError(f"predict produced no outputs; check input shape {tuple(x.shape)} and batch_size={batch_size}")
+        return torch.cat(out_list, dim=0)
+
+# -----------------------------
+# Helpers
+# -----------------------------
+
+def _norm_2d(C: int, kind: str = "group", **kwargs) -> nn.Module:
+    if kind == "group":
+        num_groups = kwargs.get("num_groups", min(8, max(1, C // 8)) or 1)
+        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.InstanceNorm2d(C, affine=True, track_running_stats=False)
+    elif kind == "batch":
+        return nn.BatchNorm2d(C)
+    else:
+        raise ValueError(f"Unknown norm kind: {kind}")
+
+
+def _pad_to_match(x: torch.Tensor, H: int, W: int) -> torch.Tensor:
+    """Pad x (B,C,h,w) with zeros on right/bottom to reach (H,W)."""
+    _, _, h, w = x.shape
+    ph = max(0, H - h)
+    pw = max(0, W - w)
+    if ph == 0 and pw == 0:
+        return x
+    return F.pad(x, (0, pw, 0, ph), mode='constant', value=0)
+
+
+# -----------------------------
+# Training utilities (mirror of Healpix fit)
+# -----------------------------
+from typing import Union
+import numpy as np
+from torch.utils.data import DataLoader, TensorDataset
+
+def fit(
+        model: nn.Module,
+        x_train: Union[torch.Tensor, np.ndarray],
+        y_train: Union[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: Optional[float] = None,
+        verbose: bool = True,
+        optimizer: Literal['ADAM', 'LBFGS'] = 'ADAM',
+    ) -> dict:
+    """Training loop *miroir* de `healpix_unet_torch.fit`, adapté aux images 2D.
+
+    - Entrées fixes: tensors/ndarrays de même taille (B, C, H, W) avec H=3*nside, W=4*nside
+    - Perte: MSE (regression) / BCE(BCEWithLogits si final_activation='none') / CrossEntropy (multiclasses)
+    - Optimiseur: ADAM ou LBFGS avec closure
+    - Logs: renvoie {"loss": history}
+    """
+    device = next(model.parameters()).device
+    model.to(device)
+
+    # ---- DataLoader
+    x_t = torch.as_tensor(x_train, dtype=torch.float32, device=device)
+    y_is_class = (getattr(model, 'task', 'regression') != '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 == x_t.ndim - 1) else torch.float32
+    y_t = torch.as_tensor(y_train, dtype=y_dtype, device=device)
+
+    ds = TensorDataset(x_t, y_t)
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+
+    # ---- Loss
+    if getattr(model, 'task', 'regression') == 'regression':
+        criterion = nn.MSELoss(reduction='mean')
+        seg_multiclass = False
+    else:
+        if getattr(model, 'out_channels', 1) == 1:
+            criterion = nn.BCEWithLogitsLoss() if getattr(model, 'final_activation', 'none') == 'none' else nn.BCELoss()
+            seg_multiclass = False
+        else:
+            criterion = nn.CrossEntropyLoss()
+            seg_multiclass = True
+
+    # ---- Optim
+    if optimizer.upper() == 'ADAM':
+        optim = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+        outer, inner = n_epoch, 1
+    elif optimizer.upper() == 'LBFGS':
+        optim = torch.optim.LBFGS(model.parameters(), lr=lr, max_iter=20, history_size=50, line_search_fn='strong_wolfe')
+        outer, inner = max(1, n_epoch // 20), 20
+    else:
+        raise ValueError("optimizer must be 'ADAM' or 'LBFGS'")
+
+    # ---- Train
+    history: List[float] = []
+    model.train()
+
+    for epoch in range(outer):
+        for _ in range(inner):
+            epoch_loss, n_samples = 0.0, 0
+            for xb, yb in loader:
+                xb = xb.to(device, dtype=torch.float32, non_blocking=True)
+                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)
+                        if seg_multiclass:
+                            loss = criterion(preds, yb)
+                        else:
+                            loss = criterion(preds, yb)
+                        loss.backward()
+                        return loss
+                    loss_val = float(optim.step(closure).item())
+                else:
+                    optim.zero_grad(set_to_none=True)
+                    preds = model(xb)
+                    if seg_multiclass:
+                        loss = criterion(preds, yb)
+                    else:
+                        loss = criterion(preds, yb)
+                    loss.backward()
+                    if clip_grad_norm is not None:
+                        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
+                    optim.step()
+                    loss_val = float(loss.item())
+
+                epoch_loss += loss_val * xb.shape[0]
+                n_samples  += xb.shape[0]
+
+            epoch_loss /= max(1, n_samples)
+            history.append(epoch_loss)
+            if verbose and ((len(history) % view_epoch == 0) or (len(history) == 1)):
+                print(f"[epoch {len(history)}] loss={epoch_loss:.6f}")
+
+    return {"loss": history}
+
+
+# -----------------------------
+# Minimal smoke test
+# -----------------------------
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    nside = 32
+    chanlist = [16, 32, 64]
+    net = PlanarUNet(
+        in_nside=nside,
+        n_chan_in=3,
+        chanlist=chanlist,
+        KERNELSZ=3,
+        task='regression',
+        out_channels=1,
+    )
+    x = torch.randn(2, 3, 3*nside, 4*nside)
+    y = net(x)
+    print('input:', x.shape, 'output:', y.shape)

{foscat-2025.9.4.dist-info → foscat-2025.10.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: foscat
-Version: 2025.9.4
+Version: 2025.10.2
 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>

{foscat-2025.9.4.dist-info → foscat-2025.10.2.dist-info}/RECORD

@@ -4,12 +4,12 @@ foscat/BkTensorflow.py,sha256=iIdLx6VTOfOEocfZBOGyizQn5geDLTfdWWAwDeQr9YA,20056
 foscat/BkTorch.py,sha256=W3n3XkFw9oecmyfSWt2JbP5L5eKWn8AAu0RZx1yLQb0,31975
 foscat/CNN.py,sha256=4vky7jqTshL1aYLWsc-hQwf7gDjTVjL7I6HZiAsa6x4,5158
 foscat/CircSpline.py,sha256=CXi49FxF8ZoeZ17Ua8c1AZXe2B5ICEC9aCXb97atB3s,4028
-foscat/FoCUS.py,sha256=DFPzRMRHE7eUfXfqoqHiM0f4mAiOs_6GmUieipXplo4,105645
+foscat/FoCUS.py,sha256=9GHjHUBM8sfmYwZw7u-xGeOOjFCpbbvM-5m3yjPSt5g,105752
 foscat/GCNN.py,sha256=q7yWHCMJpP7-m3WvR3OQnp5taeYWaMxIY2hQ6SIb9gs,4487
 foscat/HOrientedConvol.py,sha256=xMaS-zzoUyXisBCPsHBVpn54tuA9Qv3na-tT86Cwn7U,38744
 foscat/HealBili.py,sha256=YRPk9PO5G8NdwKeb33xiJs3_pMPAgIv5phCs8GT6LN0,12943
 foscat/HealSpline.py,sha256=YRotJ1NQuXYFyFiM8fp6qkATIwRJ8lqIVo4vGXpHO-w,7472
-foscat/Plot.py,sha256=a_NKE91F82qBeLNL3gy0qE06GCL-JbHgeIQ2iwIRPjY,47326
+foscat/Plot.py,sha256=bpohWGsblTBxMrqE_X-iRvuvT-YyHDcgfWB4iYk5l10,49218
 foscat/Softmax.py,sha256=UDZGrTroYtmGEyokGUVpwNO_cgbICi9QVuRr8Yx52_k,2917
 foscat/SphericalStencil.py,sha256=DFipxQWupYPBa_62CkuMe9K7HkLFiHteF7Q6lAs3TLQ,56714
 foscat/Spline1D.py,sha256=rKzzenduaZZ-yBDJd35it6Gyrj1spqb7hoIaUgISPzY,2983
@@ -20,9 +20,13 @@ foscat/alm.py,sha256=XkK4rFVRoO-oJpr74iBffKt7hdS_iJkR016IlYm10gQ,33832
 foscat/backend.py,sha256=l3aMwDyXP6jURMIvratFMGWCTcQpaR68KnUuuGDezqE,45418
 foscat/backend_tens.py,sha256=9Dp136m9frkclkwifJQLLbIpl3ETI3_txdPUZcKfuMw,1618
 foscat/heal_NN.py,sha256=krEHM9NMZ74T9HUf-qK5td0tFttBA5SbaRgzThM2GYs,16943
-foscat/healpix_unet_torch.py,sha256=LXu5pDptDQN-mn3Fv7-I0Q_g_i4gPoEd5b9BmUTTSdU,51279
+foscat/healpix_unet_torch.py,sha256=CDdrWyPJqF_gUT6rwB-TnghsbtkJ89WFw-K14m37DDQ,52221
+foscat/healpix_vit_skip.py,sha256=26qpYoX7W1vCJujqtYUiRPUmwrDf_UJSN5kbL7DVV8I,20359
+foscat/healpix_vit_torch-old.py,sha256=_PJecWRIWJc2FTQB_rthEeqLKYJ_UIdTN8ib_JuQ_xw,28985
+foscat/healpix_vit_torch.py,sha256=XOqEOazob6WRptSD5dh5acWM_1_uErKXaVWPm08X-O0,22198
 foscat/loss_backend_tens.py,sha256=dCOVN6faDtIpN3VO78HTmYP2i5fnFAf-Ddy5qVBlGrM,1783
 foscat/loss_backend_torch.py,sha256=k3z18Dj3SaLKK6ZIKcm7GO4U_YKYVP6LtHG1aIbxkYk,1627
+foscat/planar_vit.py,sha256=lQqwyz_P8G-Dav2vLqgkssDfeSe15YmjFzP5W-otjs0,6888
 foscat/scat.py,sha256=qGYiBIysPt65MdmF07WWA4piVlTfA9-lFDTaicnqC2w,72822
 foscat/scat1D.py,sha256=W5Uu6wdQ4ZsFKXpof0f1OBl-1wjJmW7ruvddRWxe7uM,53726
 foscat/scat2D.py,sha256=boKj0ASqMMSy7uQLK6hPniG87m3hZGJBYBiq5v8F9IQ,532
@@ -31,8 +35,9 @@ foscat/scat_cov1D.py,sha256=XOxsZZ5TYq8f34i2tUgIfzyaqaTDlICB3HzD2l_puro,531
 foscat/scat_cov2D.py,sha256=pAm0fKw8wyXram0TFbtw8tGcc8QPKuPXpQk0kh10r4U,7078
 foscat/scat_cov_map.py,sha256=9MzpwT2g9S3dmnjHEMK7PPLQ27oGQg2VFVsP_TDUU5E,2869
 foscat/scat_cov_map2D.py,sha256=zaIIYshXCqAeZ04I158GhD-Op4aoMlLnLEy7rxckVYY,2842
-foscat-2025.9.4.dist-info/licenses/LICENSE,sha256=i0ukIr8ZUpkSY2sZaE9XZK-6vuSU5iG6IgX_3pjatP8,1505
-foscat-2025.9.4.dist-info/METADATA,sha256=MznoPNHG--ttgyql7lKM2YZrrYz3AiHLag3yn2C04iM,7215
-foscat-2025.9.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-foscat-2025.9.4.dist-info/top_level.txt,sha256=AGySXBBAlJgb8Tj8af6m_F-aiNg2zNTcybCUPVOKjAg,7
-foscat-2025.9.4.dist-info/RECORD,,
+foscat/unet_2_d_from_healpix_params.py,sha256=r8hN-s091f3yHYlvAAiBbLOvtsz9vPrdwrWPM0ULR2Q,15949
+foscat-2025.10.2.dist-info/licenses/LICENSE,sha256=i0ukIr8ZUpkSY2sZaE9XZK-6vuSU5iG6IgX_3pjatP8,1505
+foscat-2025.10.2.dist-info/METADATA,sha256=WTeul0o0mER67oWvadXX0JV9hn1R3yAesPHOgPn8dys,7216
+foscat-2025.10.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+foscat-2025.10.2.dist-info/top_level.txt,sha256=AGySXBBAlJgb8Tj8af6m_F-aiNg2zNTcybCUPVOKjAg,7
+foscat-2025.10.2.dist-info/RECORD,,