foscat 2025.11.1__py3-none-any.whl → 2026.2.1__py3-none-any.whl

foscat/SphereUpGeo.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from foscat.SphereDownGeo import SphereDownGeo
+
+
+class SphereUpGeo(nn.Module):
+    """Geometric HEALPix upsampling operator using the transpose of SphereDownGeo.
+
+    `cell_ids_out` (coarse pixels at nside_out, NESTED) is mandatory.
+    Forward expects x of shape [B, C, K_out] aligned with that order.
+    Output is a full fine-grid map [B, C, N_in] at nside_in = 2*nside_out.
+
+    Normalization (diagonal corrections):
+      - up_norm='adjoint': x_up = M^T x
+      - up_norm='col_l1':  x_up = (M^T x) / col_sum, col_sum[i] = sum_k M[k,i]
+      - up_norm='diag_l2': x_up = (M^T x) / col_l2,  col_l2[i]  = sum_k M[k,i]^2
+    """
+
+    def __init__(
+        self,
+        nside_out: int,
+        cell_ids_out,
+        radius_deg: float | None = None,
+        sigma_deg: float | None = None,
+        weight_norm: str = "l1",
+        up_norm: str = "col_l1",
+        eps: float = 1e-12,
+        device=None,
+        dtype=torch.float32,
+    ):
+        super().__init__()
+
+        if cell_ids_out is None:
+            raise ValueError("cell_ids_out is mandatory (1D list/np/tensor of coarse HEALPix ids at nside_out).")
+
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device
+        self.dtype = dtype
+
+        self.nside_out = int(nside_out)
+        assert (self.nside_out & (self.nside_out - 1)) == 0, "nside_out must be a power of 2."
+        self.nside_in = self.nside_out * 2
+
+        self.N_out = 12 * self.nside_out * self.nside_out
+        self.N_in = 12 * self.nside_in * self.nside_in
+
+        up_norm = str(up_norm).lower().strip()
+        if up_norm not in ("adjoint", "col_l1", "diag_l2"):
+            raise ValueError("up_norm must be 'adjoint', 'col_l1', or 'diag_l2'.")
+        self.up_norm = up_norm
+        self.eps = float(eps)
+
+        # Coarse ids in user-provided order (must be unique for alignment)
+        if isinstance(cell_ids_out, torch.Tensor):
+            cell_ids_out_np = cell_ids_out.detach().cpu().numpy().astype(np.int64)
+        else:
+            cell_ids_out_np = np.asarray(cell_ids_out, dtype=np.int64)
+
+        if cell_ids_out_np.ndim != 1:
+            raise ValueError("cell_ids_out must be 1D")
+        if cell_ids_out_np.size == 0:
+            raise ValueError("cell_ids_out must be non-empty")
+        if cell_ids_out_np.min() < 0 or cell_ids_out_np.max() >= self.N_out:
+            raise ValueError("cell_ids_out contains out-of-bounds ids for this nside_out")
+        if np.unique(cell_ids_out_np).size != cell_ids_out_np.size:
+            raise ValueError("cell_ids_out must not contain duplicates (order matters for alignment).")
+
+        self.cell_ids_out_np = cell_ids_out_np
+        self.K_out = int(cell_ids_out_np.size)
+        self.register_buffer("cell_ids_out_t", torch.as_tensor(cell_ids_out_np, dtype=torch.long, device=self.device))
+
+        # Build the FULL down operator at fine resolution (nside_in -> nside_out)
+        tmp_down = SphereDownGeo(
+            nside_in=self.nside_in,
+            mode="smooth",
+            radius_deg=radius_deg,
+            sigma_deg=sigma_deg,
+            weight_norm=weight_norm,
+            device=self.device,
+            dtype=self.dtype,
+            use_csr=False,
+        )
+
+        M_down_full = torch.sparse_coo_tensor(
+            tmp_down.M.indices(),
+            tmp_down.M.values(),
+            size=(tmp_down.N_out, tmp_down.N_in),
+            device=self.device,
+            dtype=self.dtype,
+        ).coalesce()
+
+        # Extract ONLY the requested coarse rows, in the provided order.
+        # We do this on CPU with numpy for simplicity and speed at init.
+        idx = M_down_full.indices().cpu().numpy()
+        vals = M_down_full.values().cpu().numpy()
+        rows = idx[0]
+        cols = idx[1]
+
+        # Map original row id -> new row position [0..K_out-1]
+        row_map = {int(r): i for i, r in enumerate(cell_ids_out_np.tolist())}
+        mask = np.fromiter((r in row_map for r in rows), dtype=bool, count=rows.size)
+
+        rows_sel = rows[mask]
+        cols_sel = cols[mask]
+        vals_sel = vals[mask]
+
+        new_rows = np.fromiter((row_map[int(r)] for r in rows_sel), dtype=np.int64, count=rows_sel.size)
+
+        M_down_sub = torch.sparse_coo_tensor(
+            torch.as_tensor(np.stack([new_rows, cols_sel], axis=0), dtype=torch.long),
+            torch.as_tensor(vals_sel, dtype=self.dtype),
+            size=(self.K_out, self.N_in),
+            device=self.device,
+            dtype=self.dtype,
+        ).coalesce()
+
+        # Store M^T (sparse) so forward is just sparse.mm
+        M_up = self._transpose_sparse(M_down_sub)  # [N_in, K_out]
+        self.register_buffer("M_indices", M_up.indices())
+        self.register_buffer("M_values", M_up.values())
+        self.M_size = M_up.size()
+
+        # Diagonal normalizers (length N_in), based on the selected coarse rows only
+        idx_sub = M_down_sub.indices()
+        vals_sub = M_down_sub.values()
+        fine_cols = idx_sub[1]
+
+        col_sum = torch.zeros(self.N_in, device=self.device, dtype=self.dtype)
+        col_l2 = torch.zeros(self.N_in, device=self.device, dtype=self.dtype)
+        col_sum.scatter_add_(0, fine_cols, vals_sub)
+        col_l2.scatter_add_(0, fine_cols, vals_sub * vals_sub)
+
+        self.register_buffer("col_sum", col_sum)
+        self.register_buffer("col_l2", col_l2)
+
+        # Fine ids (full sphere)
+        self.register_buffer("cell_ids_in_t", torch.arange(self.N_in, dtype=torch.long, device=self.device))
+
+        self.M_T =  torch.sparse_coo_tensor(
+            self.M_indices.to(device=self.device),
+            self.M_values.to(device=self.device, dtype=self.dtype),
+            size=self.M_size,
+            device=self.device,
+            dtype=self.dtype,
+        ).coalesce().to_sparse_csr().to(self.device)
+        
+    @staticmethod
+    def _transpose_sparse(M: torch.Tensor) -> torch.Tensor:
+        M = M.coalesce()
+        idx = M.indices()
+        vals = M.values()
+        R, C = M.size()
+        idx_T = torch.stack([idx[1], idx[0]], dim=0)
+        return torch.sparse_coo_tensor(idx_T, vals, size=(C, R), device=M.device, dtype=M.dtype).coalesce()
+
+    def forward(self, x: torch.Tensor):
+        """x: [B, C, K_out] -> x_up: [B, C, N_in]."""
+        B, C, K_out = x.shape
+        assert K_out == self.K_out, f"Expected K_out={self.K_out}, got {K_out}"
+
+        x_bc = x.reshape(B * C, K_out)
+        x_up_bc_T = torch.sparse.mm(self.M_T, x_bc.T)    # [N_in, B*C]
+        x_up = x_up_bc_T.T.reshape(B, C, self.N_in) # [B, C, N_in]
+
+        if self.up_norm == "col_l1":
+            denom = self.col_sum.to(device=x.device, dtype=x.dtype).clamp_min(self.eps)
+            x_up = x_up / denom.view(1, 1, -1)
+        elif self.up_norm == "diag_l2":
+            denom = self.col_l2.to(device=x.device, dtype=x.dtype).clamp_min(self.eps)
+            x_up = x_up / denom.view(1, 1, -1)
+
+        return x_up, self.cell_ids_in_t.to(device=x.device)

foscat/SphericalStencil.py

@@ -2,13 +2,8 @@
 # Author: J.-M. Delouis
 import numpy as np
 import healpy as hp
-import foscat.scat_cov as sc
 import torch
 
-import numpy as np
-import torch
-import healpy as hp
-
 
 class SphericalStencil:
     """
@@ -61,8 +56,7 @@ class SphericalStencil:
             device=None,
             dtype=None,
             n_gauges=1,
-            gauge_type='cosmo',
-            scat_op=None,
+            gauge_type='phi',
     ):
         assert kernel_sz >= 1 and int(kernel_sz) == kernel_sz
         assert kernel_sz % 2 == 1, "kernel_sz must be odd"
@@ -75,10 +69,6 @@ class SphericalStencil:
         self.gauge_type=gauge_type
         
         self.nest = bool(nest)
-        if scat_op is None:
-            self.f=sc.funct(KERNELSZ=self.KERNELSZ)
-        else:
-            self.f=scat_op
             
         # Torch defaults
         if device is None:
@@ -354,10 +344,27 @@ class SphericalStencil:
         # --- build the local (P,3) stencil once on device
         P = self.P
         vec_np = np.zeros((P, 3), dtype=float)
-        grid = (np.arange(self.KERNELSZ) - self.KERNELSZ // 2) / self.nside
-        vec_np[:, 0] = np.tile(grid, self.KERNELSZ)
-        vec_np[:, 1] = np.repeat(grid, self.KERNELSZ)
-        vec_np[:, 2] = 1.0 - np.sqrt(vec_np[:, 0]**2 + vec_np[:, 1]**2)
+        grid = (np.arange(self.KERNELSZ) - self.KERNELSZ // 2)
+         
+        # NEW: angular offsets
+        xx,yy=np.meshgrid(grid,grid)
+        s=1.0 # could be modified
+        alpha_pix = hp.nside2resol(self.nside, arcmin=False)  # ~ taille angulaire typique
+        dtheta = (np.sqrt(xx**2+yy**2) * alpha_pix * s).ravel()
+        dphi   = (np.arctan2(yy,xx)).ravel()
+        # local spherical displacement
+        # convert to unit vectors
+        x = np.sin(dtheta) * np.cos(dphi)
+        y = np.sin(dtheta) * np.sin(dphi)
+        z = np.cos(dtheta)  
+        #print(self.nside*x.reshape(self.KERNELSZ,self.KERNELSZ))
+        #print(self.nside*y.reshape(self.KERNELSZ,self.KERNELSZ))
+        #print(self.nside*z.reshape(self.KERNELSZ,self.KERNELSZ))
+        vec_np = np.stack([x, y, z], axis=-1)
+        
+        #vec_np[:, 0] = np.tile(grid, self.KERNELSZ)
+        #vec_np[:, 1] = np.repeat(grid, self.KERNELSZ)
+        #vec_np[:, 2] = 1.0 - np.sqrt(vec_np[:, 0]**2 + vec_np[:, 1]**2)
         vec_t = torch.as_tensor(vec_np, device=self.device, dtype=self.dtype)     # (P,3)
 
         # --- rotation matrices for all targets & gauges: (K,G,3,3)
@@ -371,7 +378,7 @@ class SphericalStencil:
             th, ph, alpha, G=self.G, gauge_cosmo=(self.gauge_type=='cosmo'),
             device=self.device, dtype=self.dtype
         )  # shape (K,G,3,3)
-
+        
         # --- rotate stencil for each (target, gauge): (K,G,P,3)
         #     einsum over local stencil (P,3) with rotation (K,G,3,3)
         rotated = torch.einsum('kgij,pj->kgpi', R_t, vec_t)  # (K,G,P,3)
@@ -568,119 +575,6 @@ class SphericalStencil:
         self.dtype  = dtype
 
     
-    '''
-    def bind_support_torch_multi(self, ids_sorted_np, *, device=None, dtype=None):
-        """
-        Multi-gauge sparse binding (Step B).
-        Uses self.idx_t_multi / self.w_t_multi prepared by prepare_torch(..., G>1)
-        and builds, for each gauge g, (pos_safe, w_norm, present).
-
-        Parameters
-        ----------
-        ids_sorted_np : np.ndarray (K,)
-            Sorted pixel ids for available samples (matches the last axis of your data).
-        device, dtype : torch device/dtype for the produced mapping tensors.
-
-        Side effects
-        ------------
-        Sets:
-          - self.ids_sorted_np  : (K,)
-          - self.pos_safe_t_multi : (G, 4, K*P)  LongTensor
-          - self.w_norm_t_multi   : (G, 4, K*P)  Tensor
-          - self.present_t_multi  : (G, 4, K*P)  BoolTensor
-          - (and mirrors device/dtype in self.device/self.dtype)
-        """
-        assert hasattr(self, 'idx_t_multi') and self.idx_t_multi is not None, \
-            "Call prepare_torch(..., G>0) before bind_support_torch_multi(...)"
-        assert hasattr(self, 'w_t_multi') and self.w_t_multi is not None
-
-        if device is None: device = self.device
-        if dtype  is None: dtype  = self.dtype
-
-        self.ids_sorted_np = np.asarray(ids_sorted_np, dtype=np.int64).reshape(-1)
-        ids_sorted = torch.as_tensor(self.ids_sorted_np, device=device, dtype=torch.long)
-
-        G, _, M = self.idx_t_multi.shape
-        K = self.Kb
-        P = self.P
-        assert M == K*P, "idx_t_multi second axis must have K*P columns"
-
-        pos_list, present_list, wnorm_list = [], [], []
-
-        for g in range(G):
-            idx = self.idx_t_multi[g].to(device=device, dtype=torch.long)   # (4, M)
-            w   = self.w_t_multi[g].to(device=device, dtype=dtype)          # (4, M)
-
-            pos = torch.searchsorted(ids_sorted, idx.reshape(-1)).view(4, M)
-            in_range = pos < ids_sorted.numel()
-            cmp_vals = torch.full_like(idx, -1)
-            cmp_vals[in_range] = ids_sorted[pos[in_range]]
-            present = (cmp_vals == idx)
-
-            # normalize weights per column after masking
-            w = w * present
-            colsum = w.sum(dim=0, keepdim=True).clamp_min(1e-12)
-            w_norm = w / colsum
-
-            pos_safe = torch.where(present, pos, torch.zeros_like(pos))
-
-            pos_list.append(pos_safe)
-            present_list.append(present)
-            wnorm_list.append(w_norm)
-
-        self.pos_safe_t_multi = torch.stack(pos_list, dim=0)     # (G, 4, M)
-        self.present_t_multi  = torch.stack(present_list, dim=0) # (G, 4, M)
-        self.w_norm_t_multi   = torch.stack(wnorm_list, dim=0)   # (G, 4, M)
-
-        # mirror runtime placement
-        self.device = device
-        self.dtype  = dtype
-    
-    # ------------------------------------------------------------------
-    # Step B: bind support Torch
-    # ------------------------------------------------------------------
-    def bind_support_torch(self, ids_sorted_np, *, device=None, dtype=None):
-        """
-        Map HEALPix neighbor indices (from Step A) to actual data samples
-        sorted by pixel id. Produces pos_safe and normalized weights.
-
-        Parameters
-        ----------
-        ids_sorted_np : np.ndarray (K,)
-            Sorted pixel ids for available data.
-        device, dtype : Torch device/dtype for results.
-        """
-        if device is None:
-            device = self.device
-        if dtype is None:
-            dtype = self.dtype
-
-        self.ids_sorted_np = np.asarray(ids_sorted_np, dtype=np.int64)
-        ids_sorted = torch.as_tensor(self.ids_sorted_np, device=device, dtype=torch.long)
-
-        idx = self.idx_t.to(device=device, dtype=torch.long)
-        w   = self.w_t.to(device=device, dtype=dtype)
-
-        M = self.Kb * self.P
-        idx = idx.view(4, M)
-        w   = w.view(4, M)
-
-        pos = torch.searchsorted(ids_sorted, idx.reshape(-1)).view(4, M)
-        in_range = pos < ids_sorted.shape[0]
-        cmp_vals = torch.full_like(idx, -1)
-        cmp_vals[in_range] = ids_sorted[pos[in_range]]
-        present = (cmp_vals == idx)
-
-        w = w * present
-        colsum = w.sum(dim=0, keepdim=True).clamp_min(1e-12)
-        w_norm = w / colsum
-
-        self.pos_safe_t = torch.where(present, pos, torch.zeros_like(pos))
-        self.w_norm_t   = w_norm
-        self.present_t  = present
-        self.device = device
-        self.dtype  = dtype
-    '''
     # ------------------------------------------------------------------
     # Step C: apply convolution (already Torch in your code)
     # ------------------------------------------------------------------
@@ -1215,7 +1109,7 @@ class SphericalStencil:
         vals = torch.cat(vals_all, dim=0)
 
         
-        indices = torch.stack([cols, rows], dim=0)             # (2, nnz) invert rows/cols for foscat needs
+        indices = torch.stack([cols, rows], dim=0)             
 
         if return_sparse_tensor:
             M = torch.sparse_coo_tensor(indices, vals, size=shape, device=device, dtype=k_dtype).coalesce()
@@ -1224,123 +1118,10 @@ class SphericalStencil:
             return vals, indices, shape
 
 
-    def _to_numpy_1d(self, ids):
-        """Return a 1D numpy array of int64 for a single set of cell ids."""
-        import numpy as np, torch
-        if isinstance(ids, np.ndarray):
-            return ids.reshape(-1).astype(np.int64, copy=False)
-        if torch.is_tensor(ids):
-            return ids.detach().cpu().to(torch.long).view(-1).numpy()
-        # python list/tuple of ints
-        return np.asarray(ids, dtype=np.int64).reshape(-1)
-    
-    def _is_varlength_batch(self, ids):
-        """
-        True if ids is a list/tuple of per-sample id arrays (var-length batch).
-        False if ids is a single array/tensor of ids (shared for whole batch).
-        """
-        import numpy as np, torch
-        if isinstance(ids, (list, tuple)):
-            return True
-        if isinstance(ids, np.ndarray) and ids.ndim == 2:
-            # This would be a dense (B, Npix) matrix -> NOT var-length list
-            return False
-        if torch.is_tensor(ids) and ids.dim() == 2:
-            return False
-        return False
-    
-    def Down(self, im, cell_ids=None, nside=None,max_poll=False):
-        """
-        If `cell_ids` is a single set of ids -> return a single (Tensor, Tensor).
-        If `cell_ids` is a list (var-length)           -> return (list[Tensor], list[Tensor]).
-        """
-        if self.f is None:
-            if self.dtype==torch.float64:
-                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float64')
-            else:
-                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float32')
-            
-        if cell_ids is None:
-            dim,cdim = self.f.ud_grade_2(im,cell_ids=self.cell_ids,nside=self.nside,max_poll=max_poll)
-            return dim,cdim
-        
-        if nside is None:
-            nside = self.nside
-
-        # var-length mode: list/tuple of ids, one per sample
-        if self._is_varlength_batch(cell_ids):
-            outs, outs_ids = [], []
-            B = len(cell_ids)
-            for b in range(B):
-                cid_b = self._to_numpy_1d(cell_ids[b])
-                # extraire le bon échantillon d'`im`
-                if torch.is_tensor(im):
-                    xb = im[b:b+1]  # (1, C, N_b)
-                    yb, ids_b = self.f.ud_grade_2(xb, cell_ids=cid_b, nside=nside,max_poll=max_poll)
-                    outs.append(yb.squeeze(0))  # (C, N_b')
-                else:
-                    # si im est déjà une liste de (C, N_b)
-                    xb = im[b]
-                    yb, ids_b = self.f.ud_grade_2(xb[None, ...], cell_ids=cid_b, nside=nside,max_poll=max_poll)
-                    outs.append(yb.squeeze(0))
-                outs_ids.append(torch.as_tensor(ids_b, device=outs[-1].device, dtype=torch.long))
-            return outs, outs_ids
-
-        # grille commune (un seul vecteur d'ids)
-        cid = self._to_numpy_1d(cell_ids)
-        return self.f.ud_grade_2(im, cell_ids=cid, nside=nside,max_poll=False)
-
-    def Up(self, im, cell_ids=None, nside=None, o_cell_ids=None):
-        """
-        If `cell_ids` / `o_cell_ids` are single arrays  -> return Tensor.
-        If they are lists (var-length per sample)       -> return list[Tensor].
-        """
-        if self.f is None:
-            if self.dtype==torch.float64:
-                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float64')
-            else:
-                self.f=sc.funct(KERNELSZ=self.KERNELSZ,all_type='float32')
-                
-        if cell_ids is None:
-            dim = self.f.up_grade(im,self.nside*2,cell_ids=self.cell_ids,nside=self.nside)
-            return dim
-        
-        if nside is None:
-            nside = self.nside
-
-        # var-length: listes parallèles
-        if self._is_varlength_batch(cell_ids):
-            assert isinstance(o_cell_ids, (list, tuple)) and len(o_cell_ids) == len(cell_ids), \
-                "In var-length mode, `o_cell_ids` must be a list with same length as `cell_ids`."
-            outs = []
-            B = len(cell_ids)
-            for b in range(B):
-                cid_b  = self._to_numpy_1d(cell_ids[b])      # coarse ids
-                ocid_b = self._to_numpy_1d(o_cell_ids[b])    # fine   ids
-                if torch.is_tensor(im):
-                    xb = im[b:b+1]  # (1, C, N_b_coarse)
-                    yb = self.f.up_grade(xb, nside*2, cell_ids=cid_b, nside=nside,
-                                         o_cell_ids=ocid_b, force_init_index=True)
-                    outs.append(yb.squeeze(0))  # (C, N_b_fine)
-                else:
-                    xb = im[b]  # (C, N_b_coarse)
-                    yb = self.f.up_grade(xb[None, ...], nside*2, cell_ids=cid_b, nside=nside,
-                                         o_cell_ids=ocid_b, force_init_index=True)
-                    outs.append(yb.squeeze(0))
-            return outs
-
-        # grille commune
-        cid  = self._to_numpy_1d(cell_ids)
-        ocid = self._to_numpy_1d(o_cell_ids) if o_cell_ids is not None else None
-        return self.f.up_grade(im, nside*2, cell_ids=cid, nside=nside,
-                               o_cell_ids=ocid, force_init_index=True)
-
-
     def to_tensor(self,x):
-        return torch.tensor(x,device=self.device,dtype=self.dtype)
-    
+        return torch.tensor(x,device='cuda')
+
     def to_numpy(self,x):
         if isinstance(x,np.ndarray):
             return x
-        return x.cpu().numpy()
-        
+        return x.cpu().numpy()