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

foscat/BkTorch.py

@@ -63,70 +63,262 @@ class BkTorch(BackendBase.BackendBase):
             torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
         )
 
-    import torch
-
-    def binned_mean(self, data, cell_ids):
+    # ---------------------------------
+    # HEALPix binning utilities (nested)
+    # ---------------------------------
+    # Robust binned_mean that supports arbitrary subsets (N not divisible by 4)
+    # and batched cell_ids of shape [B, N]. It returns compact per-parent means
+    # even when some parents are missing (sparse coverage).
+
+    def binned_mean(self, data, cell_ids, *, padded: bool = False, fill_value: float = float("nan")):
+        """Average values over parent HEALPix pixels (nested) when downgrading nside→nside/2.
+
+        Works with full-sky or sparse subsets (no need for N to be divisible by 4).
+
+        Parameters
+        ----------
+        data : torch.Tensor or np.ndarray
+            Shape ``[..., N]`` or ``[B, ..., N]``.
+        cell_ids : torch.LongTensor or np.ndarray
+            Shape ``[N]`` or ``[B, N]`` (nested indexing at the *child* resolution).
+        padded : bool, optional (default: False)
+            Only used when ``cell_ids`` is ``[B, N]``. If ``False``, returns Python
+            lists (ragged) of per-batch results. If ``True``, returns padded tensors
+            plus a boolean mask of valid bins.
+        fill_value : float, optional
+            Value used for padding when ``padded=True``.
+
+        Returns
+        -------
+        If ``cell_ids`` is ``[N]``:
+            mean  : torch.Tensor, shape ``[..., n_bins]``
+            groups: torch.LongTensor, shape ``[n_bins]``  (sorted unique parents)
+
+        If ``cell_ids`` is ``[B, N]`` and ``padded=False``:
+            means_list  : List[torch.Tensor] of length B, each shape ``[T, n_bins_b]``
+                          where ``T = prod(data.shape[1:-1])`` (or 1 if none).
+            groups_list : List[torch.LongTensor] of length B, each shape ``[n_bins_b]``
+
+        If ``cell_ids`` is ``[B, N]`` and ``padded=True``:
+            mean_padded : torch.Tensor, shape ``[B, T, max_bins]`` (or ``[B, max_bins]`` if T==1)
+            groups_pad  : torch.LongTensor, shape ``[B, max_bins]`` (parents, padded with -1)
+            mask        : torch.BoolTensor, shape ``[B, max_bins]`` (True where valid)
         """
-        Compute the mean over groups of 4 nested HEALPix cells (nside → nside/2).
-
-        Args:
-            data (torch.Tensor): Tensor of shape [..., N], where N is the number of HEALPix cells.
-            cell_ids (torch.LongTensor): Tensor of shape [N], with cell indices (nested ordering).
+        import torch, numpy as np
 
-        Returns:
-            torch.Tensor: Tensor of shape [..., n_bins], with averaged values per group of 4 cells.
-        """
+        # ---- Tensorize & device/dtype plumbing ----
         if isinstance(data, np.ndarray):
-            data = torch.from_numpy(data).to(
-                dtype=torch.float32, device=self.torch_device
-            )
+            data = torch.from_numpy(data).to(dtype=torch.float32, device=getattr(self, 'torch_device', 'cpu'))
         if isinstance(cell_ids, np.ndarray):
-            cell_ids = torch.from_numpy(cell_ids).to(
-                dtype=torch.long, device=self.torch_device
-            )
-
-        # Compute supercell ids by grouping 4 nested cells together
-        groups = cell_ids // 4
+            cell_ids = torch.from_numpy(cell_ids).to(dtype=torch.long, device=data.device)
 
-        # Get unique group ids and inverse mapping
-        unique_groups, inverse_indices = torch.unique(groups, return_inverse=True)
-        n_bins = unique_groups.shape[0]
+        data = data.to(device=getattr(self, 'torch_device', data.device))
+        cell_ids = cell_ids.to(device=data.device, dtype=torch.long)
 
-        # Flatten all leading dimensions into a single batch dimension
-        original_shape = data.shape[:-1]
+        if data.ndim < 1:
+            raise ValueError("`data` must have at least 1 dimension (last is N).")
         N = data.shape[-1]
-        data_flat = data.reshape(-1, N)  # Shape: [B, N]
 
-        # Prepare to compute sums using scatter_add
-        B = data_flat.shape[0]
-
-        # Repeat inverse indices for each batch element
-        idx = inverse_indices.repeat(B, 1)  # Shape: [B, N]
+        # Flatten leading dims (rows) for scatter convenience
+        orig = data.shape[:-1]
+        T = int(np.prod(orig[1:])) if len(orig) > 1 else 1  # repeats per batch row
+        if cell_ids.ndim == 1:
+            # Shared mapping for all rows
+            groups = (cell_ids // 4).to(torch.long)  # [N]
+            # Unique parent ids + inverse indices
+            parents, inv = torch.unique(groups, sorted=True, return_inverse=True)
+            n_bins = parents.numel()
+
+            R = int(np.prod(orig)) if len(orig) > 0 else 1
+            data_flat = data.reshape(R, N)  # [R, N]
+
+            # Row offsets -> independent bins per row
+            row_offsets = torch.arange(R, device=data.device).unsqueeze(1) * n_bins  # [R,1]
+            idx = inv.unsqueeze(0).expand(R, -1) + row_offsets                         # [R,N]
+
+            vals_flat = data_flat.reshape(-1)
+            idx_flat  = idx.reshape(-1)
+
+            out_sum = torch.zeros(R * n_bins, dtype=data.dtype, device=data.device)
+            out_cnt = torch.zeros_like(out_sum)
+            out_sum.scatter_add_(0, idx_flat, vals_flat)
+            out_cnt.scatter_add_(0, idx_flat, torch.ones_like(vals_flat))
+            out_cnt.clamp_(min=1)
+
+            mean = (out_sum / out_cnt).view(*orig, n_bins)
+            return mean, parents
+
+        elif cell_ids.ndim == 2:
+            B = cell_ids.shape[0]
+            if data.shape[0] % B != 0:
+                raise ValueError(f"Leading dim of data ({data.shape[0]}) must be a multiple of cell_ids batch ({B}).")
+            R = int(np.prod(orig)) if len(orig) > 0 else 1
+            data_flat = data.reshape(R, N)  # [R, N]
+            B_data = data.shape[0]
+            T = R // B_data                 # repeats per batch row (product of extra leading dims)
+
+            means_list, groups_list = [], []
+            max_bins = 0
+            # First pass: compute per-batch parents/inv and scatter means
+            for b in range(B):
+                groups_b = (cell_ids[b] // 4).to(torch.long)  # [N]
+                parents_b, inv_b = torch.unique(groups_b, sorted=True, return_inverse=True)
+                n_bins_b = parents_b.numel()
+                max_bins = max(max_bins, n_bins_b)
+
+                # rows for this batch in data_flat
+                start = b * T
+                stop  = (b + 1) * T
+                rows  = slice(start, stop)                    # T rows
+
+                row_offsets = (torch.arange(T, device=data.device).unsqueeze(1) * n_bins_b)
+                idx = inv_b.unsqueeze(0).expand(T, -1) + row_offsets  # [T, N]
+
+                vals_flat = data_flat[rows].reshape(-1)
+                idx_flat  = idx.reshape(-1)
+
+                out_sum = torch.zeros(T * n_bins_b, dtype=data.dtype, device=data.device)
+                out_cnt = torch.zeros_like(out_sum)
+                out_sum.scatter_add_(0, idx_flat, vals_flat)
+                out_cnt.scatter_add_(0, idx_flat, torch.ones_like(vals_flat))
+                out_cnt.clamp_(min=1)
+                mean_bt = (out_sum / out_cnt).view(T, n_bins_b)  # [T, n_bins_b]
+
+                means_list.append(mean_bt)
+                groups_list.append(parents_b)
+
+            if not padded:
+                return means_list, groups_list
+
+            # Padded output
+            # mean_padded: [B, T, max_bins]; groups_pad: [B, max_bins]; mask: [B, max_bins]
+            mean_pad = torch.full((B, T, max_bins), fill_value, dtype=data.dtype, device=data.device)
+            groups_pad = torch.full((B, max_bins), -1, dtype=torch.long, device=data.device)
+            mask = torch.zeros((B, max_bins), dtype=torch.bool, device=data.device)
+            for b, (m_b, g_b) in enumerate(zip(means_list, groups_list)):
+                nb = g_b.numel()
+                mean_pad[b, :, :nb] = m_b
+                groups_pad[b, :nb] = g_b
+                mask[b, :nb] = True
+
+            # Reshape back to [B, (*extra leading dims), max_bins] if needed
+            if len(orig) > 1:
+                extra = orig[1:]  # e.g., (D1, D2, ...)
+                mean_pad = mean_pad.view(B, *extra, max_bins)
+            else:
+                mean_pad = mean_pad.view(B, max_bins)
 
-        # Offset indices to simulate a per-batch scatter into [B * n_bins]
-        batch_offsets = torch.arange(B, device=data.device).unsqueeze(1) * n_bins
-        idx_offset = idx + batch_offsets  # Shape: [B, N]
+            return mean_pad, groups_pad, mask
 
-        # Flatten everything for scatter
-        idx_offset_flat = idx_offset.flatten()
-        data_flat_flat = data_flat.flatten()
+        else:
+            raise ValueError("`cell_ids` must be of shape [N] or [B, N].")
+    '''
+    def binned_mean(self, data, cell_ids):
+        """
+        Moyenne par groupes de 4 pixels HEALPix nested (nside -> nside/2),
+        fonctionne avec un sous-ensemble arbitraire de pixels.
+
+        Args
+        ----
+        data: torch.Tensor | np.ndarray, shape [..., N]  ou  [B, ..., N]
+        cell_ids: torch.LongTensor | np.ndarray, shape [N] ou [B, N] (nested)
+
+        Returns
+        -------
+        mean: torch.Tensor, shape [..., G] ou [B, ..., G]
+        groups_out: torch.LongTensor, shape [G]  (ids HEALPix parents à nside/2)
+        """
 
-        # Accumulate sums per bin
-        out = torch.zeros(B * n_bins, dtype=data.dtype, device=data.device)
-        out = out.scatter_add(0, idx_offset_flat, data_flat_flat)
+        # --- to tensors on device ---
+        if isinstance(data, np.ndarray):
+            data = torch.from_numpy(data).to(dtype=torch.float32, device=self.torch_device)
+        if isinstance(cell_ids, np.ndarray):
+            cell_ids = torch.from_numpy(cell_ids).to(dtype=torch.long, device=self.torch_device)
+        data = data.to(self.torch_device)
+        cell_ids = cell_ids.to(self.torch_device, dtype=torch.long)
 
-        # Count number of elements per bin (to compute mean)
-        ones = torch.ones_like(data_flat_flat)
-        counts = torch.zeros(B * n_bins, dtype=data.dtype, device=data.device)
-        counts = counts.scatter_add(0, idx_offset_flat, ones)
+        # --- shapes ---
+        if data.ndim < 1:
+            raise ValueError("`data` must have at least 1 dim; last is N.")
+        N = data.shape[-1]
+        if N % 4 != 0:
+            raise ValueError(f"N={N} must be divisible by 4 for nested groups of 4.")
+
+        # --- parent groups @ nside/2, accept [N] or [B,N] ---
+        if cell_ids.ndim == 1:
+            if cell_ids.shape[0] != N:
+                raise ValueError(f"cell_ids shape {tuple(cell_ids.shape)} incompatible with N={N}.")
+            groups_parent = (cell_ids // 4).long()                  # [N]
+            # densification -> [0..G-1]
+            unique_groups, inverse = torch.unique(groups_parent, return_inverse=True)  # [G], [N]
+            # mapping identique pour toutes les lignes/rows de data
+            B_ids = 1
+        elif cell_ids.ndim == 2:
+            B_ids, N_ids = cell_ids.shape
+            if N_ids != N:
+                raise ValueError(f"cell_ids last dim {N_ids} must equal N={N}.")
+            # vérif compatibilité batch (data doit commencer par B ou par un multiple de B)
+            leading = data.shape[:-1]
+            if len(leading) == 0:
+                raise ValueError("`data` must have a leading dim to match [B, N] cell_ids.")
+            B_data = leading[0]
+            if B_data % B_ids != 0:
+                raise ValueError(f"Leading batch of data ({B_data}) must be a multiple of cell_ids batch ({B_ids}).")
+
+            # Construire un mapping DENSE par batch, mais on impose que la topologie
+            # des parents soit la même pour tous les batches -> on se base sur le batch 0
+            groups_parent0 = (cell_ids[0] // 4).long()              # [N]
+            unique_groups, inverse0 = torch.unique(groups_parent0, return_inverse=True)  # [G], [N]
+
+            # Vérification (optionnelle mais sûre) : chaque batch a les mêmes parents (ordre potentiellement différent OK)
+            # -> ici on exige même l'égalité stricte pour éviter les surprises ;
+            #    sinon on pourrait densifier par-batch et retourner une liste de groups_out.
+            for b in range(1, B_ids):
+                if not torch.equal(groups_parent0, (cell_ids[b] // 4).long()):
+                    raise ValueError("All batches in cell_ids must share the same parent groups (order & content).")
+
+            # Construire l'inverse pour tous les batches en répliquant celui du batch 0
+            inverse = inverse0.unsqueeze(0).expand(B_ids, -1)       # [B_ids, N]
+        else:
+            raise ValueError("`cell_ids` must be [N] or [B, N].")
 
-        # Compute mean
-        mean = out / counts  # Shape: [B * n_bins]
-        mean = mean.view(B, n_bins)
+        G = unique_groups.numel()  # nb de groupes parents (nside/2)
 
-        # Restore original leading dimensions
-        return mean.view(*original_shape, n_bins), unique_groups
+        # --- aplatir data en lignes ---
+        original_shape = data.shape[:-1]                            # e.g. [B, D1, ...]
+        R = int(np.prod(original_shape)) if original_shape else 1
+        data_flat = data.reshape(R, N)                              # [R, N]
 
+        # --- construire indices de bins par ligne ---
+        if cell_ids.ndim == 1:
+            idx = inverse.expand(R, -1)                             # [R, N], dans [0..G-1]
+        else:
+            # cell_ids est [B_ids, N]; on doit « étirer » chaque ligne de mapping
+            # pour couvrir les R lignes de data_flat en respectant la 1ère dim (B_data)
+            B_data = original_shape[0]
+            T = R // B_data                                         # répétitions par batch-row
+            idx = inverse.repeat_interleave(T, dim=0)               # [B_ids*T, N] == [R, N]
+
+        # --- scatter add (somme et compte) par ligne ---
+        device = data.device
+        row_offsets = torch.arange(R, device=device).unsqueeze(1) * G
+        idx_offset = idx.to(torch.long) + row_offsets               # [R,N]
+        idx_offset_flat = idx_offset.reshape(-1)
+        vals_flat = data_flat.reshape(-1)
+
+        out_sum = torch.zeros(R * G, dtype=data.dtype, device=device)
+        out_sum.scatter_add_(0, idx_offset_flat, vals_flat)
+
+        ones = torch.ones_like(vals_flat, dtype=data.dtype, device=device)
+        out_cnt = torch.zeros(R * G, dtype=data.dtype, device=device)
+        out_cnt.scatter_add_(0, idx_offset_flat, ones)
+        out_cnt = torch.clamp(out_cnt, min=1)
+
+        mean = (out_sum / out_cnt).view(R, G).view(*original_shape, G)
+
+        # On retourne les VRAIS ids HEALPix parents à nside/2
+        return mean, unique_groups
+    '''
     def average_by_cell_group(data, cell_ids):
         """
         data: tensor of shape [..., N, ...] (ex: [B, N, C])

foscat/FoCUS.py

@@ -36,7 +36,7 @@ class FoCUS:
             mpi_rank=0
     ):
 
-        self.__version__ = "2025.08.4"
+        self.__version__ = "2025.09.1"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None