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

foscat/BkTorch.py

@@ -63,70 +63,329 @@ 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_old(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).
+        import torch, numpy as np
 
-        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).
+        # ---- Tensorize & device/dtype plumbing ----
+        if isinstance(data, np.ndarray):
+            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=data.device)
 
-        Returns:
-            torch.Tensor: Tensor of shape [..., n_bins], with averaged values per group of 4 cells.
+        data = data.to(device=getattr(self, 'torch_device', data.device))
+        cell_ids = cell_ids.to(device=data.device, dtype=torch.long)
+
+        if data.ndim < 1:
+            raise ValueError("`data` must have at least 1 dimension (last is N).")
+        N = data.shape[-1]
+
+        # 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)
+
+            return mean_pad, groups_pad, mask
+
+        else:
+            raise ValueError("`cell_ids` must be of shape [N] or [B, N].")
+
+    def binned_mean(
+                self,
+                data,
+                cell_ids,
+                *,
+                reduce: str = "mean",          # <-- NEW: "mean" (par défaut) ou "max"
+                padded: bool = False,
+                fill_value: float = float("nan"),
+    ):
+        """
+            Reduce values over parent HEALPix pixels (nested) when downgrading nside→nside/2.
+
+        Parameters
+            ----------
+            data : torch.Tensor | np.ndarray
+            Shape [..., N] or [B, ..., N].
+            cell_ids : torch.LongTensor | np.ndarray
+            Shape [N] or [B, N] (nested indexing at the child resolution).
+            reduce : {"mean","max"}, default "mean"
+            Aggregation to apply within each parent group of 4 children.
+            padded : bool, default False
+            Only used when `cell_ids` is [B, N]. If False, returns ragged Python lists.
+            If True, returns padded tensors + mask.
+            fill_value : float, default NaN
+            Padding value when `padded=True`.
+
+        Returns
+            -------
+            # idem à ta doc existante, mais la valeur est une moyenne (reduce="mean")
+            # ou un maximum (reduce="max").
         """
+        
+        # ---- Tensorize & device/dtype plumbing ----
         if isinstance(data, np.ndarray):
             data = torch.from_numpy(data).to(
-                dtype=torch.float32, device=self.torch_device
+                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
+                dtype=torch.long, device=data.device
             )
+            data = data.to(device=getattr(self, "torch_device", data.device))
+            cell_ids = cell_ids.to(device=data.device, dtype=torch.long)
 
-        # Compute supercell ids by grouping 4 nested cells together
-        groups = cell_ids // 4
-
-        # Get unique group ids and inverse mapping
-        unique_groups, inverse_indices = torch.unique(groups, return_inverse=True)
-        n_bins = unique_groups.shape[0]
-
-        # 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]
-
-        # 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]
-
-        # Flatten everything for scatter
-        idx_offset_flat = idx_offset.flatten()
-        data_flat_flat = data_flat.flatten()
-
-        # 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)
 
-        # 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)
-
-        # Compute mean
-        mean = out / counts  # Shape: [B * n_bins]
-        mean = mean.view(B, n_bins)
+        # Utilitaires pour 'max' (fallback si scatter_reduce_ indisponible)
+        def _segment_max(vals_flat, idx_flat, out_size):
+            """Retourne out[out_idx] = max(vals[ idx==out_idx ]), vectorisé si possible."""
+            # PyTorch >= 1.12 / 2.0: scatter_reduce_ disponible
+            if hasattr(torch.Tensor, "scatter_reduce_"):
+                out = torch.full((out_size,), -float("inf"),
+                                 dtype=vals_flat.dtype, device=vals_flat.device)
+                out.scatter_reduce_(0, idx_flat, vals_flat, reduce="amax", include_self=True)
+                return out
+            # Fallback simple (boucle sur indices uniques) – OK pour du downsample
+            out = torch.full((out_size,), -float("inf"),
+                             dtype=vals_flat.dtype, device=vals_flat.device)
+            uniq = torch.unique(idx_flat)
+            for u in uniq.tolist():
+                m = (idx_flat == u)
+                # éviter max() sur tensor vide
+                if m.any():
+                    out[u] = torch.max(vals_flat[m])
+            return out
+
+        # ---- Flatten leading dims for scatter convenience ----
+        orig = data.shape[:-1]
+        if cell_ids.ndim == 1:
+            # Shared mapping for all rows
+            groups = (cell_ids // 4).to(torch.long)                      # [N]
+            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 = torch.arange(R, device=data.device).unsqueeze(1) * n_bins
+            idx = inv.unsqueeze(0).expand(R, -1) + row_offsets            # [R, N]
+
+            vals_flat = data_flat.reshape(-1)
+            idx_flat = idx.reshape(-1)
+            out_size = R * n_bins
+
+            if reduce == "mean":
+                out_sum = torch.zeros(out_size, 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)
+                reduced = out_sum / out_cnt
+            elif reduce == "max":
+                reduced = _segment_max(vals_flat, idx_flat, out_size)
+            else:
+                raise ValueError("reduce must be 'mean' or 'max'.")
+
+            output = reduced.view(*orig, n_bins)
+            return output, parents
+
+        elif cell_ids.ndim == 2:
+            # Per-batch mapping
+            B = cell_ids.shape[0]
+            R = int(np.prod(orig)) if len(orig) > 0 else 1
+            data_flat = data.reshape(R, N)                                # [R, N]
+            B_data = data.shape[0] if len(orig) > 0 else 1
+            if B_data % B != 0:
+                raise ValueError(
+                    f"Leading dim of data ({B_data}) must be a multiple of cell_ids batch ({B})."
+                )
+            # T = repeats per batch row (product of extra leading dims)
+            T = (R // B_data) if B_data > 0 else 1
+
+            means_list, groups_list = [], []
+            max_bins = 0
+
+            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, stop = b * T, (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_size = T * n_bins_b
+
+                if reduce == "mean":
+                    out_sum = torch.zeros(out_size, 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)
+                    reduced_bt = (out_sum / out_cnt).view(T, n_bins_b)
+                elif reduce == "max":
+                    reduced_bt = _segment_max(vals_flat, idx_flat, out_size).view(T, n_bins_b)
+                else:
+                    raise ValueError("reduce must be 'mean' or 'max'.")
+
+                means_list.append(reduced_bt)
+                groups_list.append(parents_b)
+
+            if not padded:
+                return means_list, groups_list
+
+            # Padded output (B, T, max_bins) [+ mask]
+            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 dims), max_bins] si besoin
+            if len(orig) > 1:
+                extra = orig[1:]
+                mean_pad = mean_pad.view(B, *extra, max_bins)
+            else:
+                mean_pad = mean_pad.view(B, max_bins)
 
-        # Restore original leading dimensions
-        return mean.view(*original_shape, n_bins), unique_groups
+            return mean_pad, groups_pad, mask
 
+        else:
+            raise ValueError("`cell_ids` must be of shape [N] or [B, N].")
+        
     def average_by_cell_group(data, cell_ids):
         """
         data: tensor of shape [..., N, ...] (ex: [B, N, C])
@@ -154,7 +413,7 @@ class BkTorch(BackendBase.BackendBase):
         return S.numel()
     
     def bk_SparseTensor(self, indice, w, dense_shape=[]):
-        return self.backend.sparse_coo_tensor(indice.T, w, dense_shape).to_sparse_csr().to(self.torch_device)
+        return self.backend.sparse_coo_tensor(indice, w, dense_shape).coalesce().to_sparse_csr().to(self.torch_device)
 
     def bk_stack(self, list, axis=0):
         return self.backend.stack(list, axis=axis).to(self.torch_device)