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

foscat/BkTorch.py

@@ -62,158 +62,118 @@ class BkTorch(BackendBase.BackendBase):
         self.torch_device = (
             torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
         )
-
-    # ---------------------------------
-    # 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).
+        
+    def downsample_mean_2x2(self,tim: torch.Tensor) -> torch.Tensor:
+        """
+        Average-pool tensor tim over non-overlapping 2x2 spatial blocks.
 
         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``.
+        tim : torch.Tensor
+            Tensor of shape [a, N1, N2, b].
 
         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)
+        torch.Tensor
+            Downsampled tensor of shape [a, N1//2, N2//2, b],
+            each element being the mean of a 2x2 block.
         """
-        import torch, numpy as np
-
-        # ---- 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)
-
-        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
+        a, N1, N2, b = tim.shape
+        # Ensure even sizes
+        N1_2 = N1 // 2
+        N2_2 = N2 // 2
+
+        # reshape to group 2x2 patches
+        tim_reshaped = tim[:, :2*N1_2, :2*N2_2, :].reshape(a, N1_2, 2, N2_2, 2, b)
+        # mean over the two small dims (2x2)
+        out = tim_reshaped.mean(dim=(2, 4))
+        return out
+
+    def downsample_median_2x2(self,tim: torch.Tensor) -> torch.Tensor:
+        """
+        2x2 block median downsampling on spatial axes (N1, N2).
+
+        Input:
+            tim: [a, N1, N2, b]  (real or complex)
+        Output:
+            out: [a, N1//2, N2//2, b]
+              each value is the median over the corresponding 2x2 block.
+            - For complex inputs: median is taken by sorting the 4 values by |.|,
+              returning the complex sample at the lower median rank.
+        """
+        a, N1, N2, b = tim.shape
+        N1_2 = N1 // 2
+        N2_2 = N2 // 2
+        # On ignore la dernière ligne/colonne si N1/N2 sont impairs
+        x = tim[:, :2*N1_2, :2*N2_2, :]  # [a, 2*N1_2, 2*N2_2, b]
+
+        # Regrouper les blocs 2x2 -> construire une dernière dimension de taille 4
+        # Réarrange: [a, N1_2, 2, N2_2, 2, b] -> [a, N1_2, N2_2, b, 4]
+        x = x.reshape(a, N1_2, 2, N2_2, 2, b).permute(0, 1, 3, 5, 2, 4).reshape(a, N1_2, N2_2, b, 4)
+
+        if not torch.is_complex(x):
+            # Réel : médiane le long de la dernière dim (taille 4)
+            med, _ = torch.median(x, dim=-1)      # [a, N1_2, N2_2, b]
+            return med
+        else:
+            # Complexe : trier par module puis prendre l'élément de rang 1 (médiane inférieure)
+            mags = x.abs()                        # [a, N1_2, N2_2, b, 4]
+            sorted_mag, idx = torch.sort(mags, dim=-1)   # idx: indices triés par |.| croissant
+            # Récupérer l'indice de médiane inférieure (pour 4 éléments -> position 1)
+            med_rank = 1
+            gather_idx = idx[..., med_rank:med_rank+1]   # [a, N1_2, N2_2, b, 1]
+            # Sélectionner la valeur complexe correspondante
+            med = torch.gather(x, dim=-1, index=gather_idx).squeeze(-1)   # [a, N1_2, N2_2, b]
+            return med
+
+    def downsample_mean_1d(self,tim: torch.Tensor) -> torch.Tensor:
+        """
+        Downsample tensor tim [a, N1] by averaging non-overlapping 2-element blocks.
+        Output shape: [a, N1//2]
+        """
+        a, N1 = tim.shape
+        N1_2 = N1 // 2
 
-            # 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)
+        # Ignore the last element if N1 is odd
+        x = tim[:, :2 * N1_2]
 
-            return mean_pad, groups_pad, mask
+        # Reshape to group pairs of 2 and take mean
+        x = x.reshape(a, N1_2, 2)
+        out = x.mean(dim=-1)  # [a, N1//2]
+        return out
 
+    def downsample_median_1d(self,tim: torch.Tensor) -> torch.Tensor:
+        """
+        Downsample tensor tim [a, N1] by taking the median of non-overlapping pairs (2 values).
+        Output shape: [a, N1//2]
+          - For real inputs: median of the two values.
+          - For complex inputs: pick the complex value with the smallest |.| among the two.
+        """
+        a, N1 = tim.shape
+        N1_2 = N1 // 2
+        x = tim[:, :2 * N1_2].reshape(a, N1_2, 2)  # group 2 by 2
+
+        if not torch.is_complex(x):
+            # Sort values in ascending order, then take mean of the two (true median for 2 samples)
+            x_sorted, _ = torch.sort(x, dim=-1)
+            med = x_sorted.mean(dim=-1)  # [a, N1//2]
+            return med
         else:
-            raise ValueError("`cell_ids` must be of shape [N] or [B, N].")
+            # Complex: sort by magnitude
+            mags = x.abs()
+            sorted_mags, idx = torch.sort(mags, dim=-1)
+            # Take the one with smallest magnitude (lower median)
+            med = torch.gather(x, dim=-1, index=idx[..., 0:1]).squeeze(-1)
+            return med
+    # ---------------------------------
+    # 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(
+
+    def binned_mean_old(
                 self,
                 data,
                 cell_ids,
@@ -385,7 +345,268 @@ class BkTorch(BackendBase.BackendBase):
 
         else:
             raise ValueError("`cell_ids` must be of shape [N] or [B, N].")
-        
+
+
+    def binned_mean(  # (garde ton nom si besoin de compat)
+        self,
+        data,
+        cell_ids,
+        *,
+        reduce: str = "mean",          # "mean" | "max" | "median"
+        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","median"}, default "mean"
+            Aggregation within each parent group of 4 children.
+        padded : bool, default False
+            Only 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
+        -------
+        As in your original function, with aggregation set by `reduce`.
+        """
+
+        # ---- 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=getattr(self, "torch_device", data.device)
+            )
+        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]
+        orig = data.shape[:-1]
+
+        # ---- Utilities ----
+        def _segment_max(vals_flat, idx_flat, out_size):
+            """Compute out[g] = max(vals[idx==g]); vectorized when possible."""
+            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
+            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)
+                if m.any():
+                    out[u] = torch.max(vals_flat[m])
+            return out
+
+        def _median_from_four_real(v4: torch.Tensor) -> torch.Tensor:
+            """
+            v4: [..., 4] real with NaN for missing.
+            Returns true median: average of the 2 middle finite values.
+            """
+            # Sort with NaN-last: replace NaN by +inf for sorting, then restore
+            is_nan = torch.isnan(v4)
+            v4_sortkey = torch.where(is_nan, torch.full_like(v4, float('inf')), v4)
+            v_sorted, _ = torch.sort(v4_sortkey, dim=-1)  # NaN (inf) at the end
+
+            # Count finite per group
+            k = torch.sum(~is_nan, dim=-1)  # [...]
+
+            # For k==0 -> NaN; k==1 -> the single value; k>=2 -> average middle two
+            # Indices for middle-two among the first k finite values: m-1 and m (with m = k//2)
+            m = torch.clamp(k // 2, min=1)  # ensure >=1 for gather
+            idx_lo = (m - 1).unsqueeze(-1)
+            idx_hi = torch.clamp(m, max=3).unsqueeze(-1)  # upper middle (cap at 3)
+
+            # Gather from sorted finite section
+            gather_lo = torch.gather(v_sorted, -1, idx_lo).squeeze(-1)
+            gather_hi = torch.gather(v_sorted, -1, idx_hi).squeeze(-1)
+            med = 0.5 * (gather_lo + gather_hi)
+
+            # Handle k==1: both idx point to same single finite value -> OK
+            # Handle k==0: set NaN
+            med = torch.where(k > 0, med, torch.full_like(med, float('nan')))
+            return med
+
+        def _median_from_four_complex(v4: torch.Tensor) -> torch.Tensor:
+            """
+            v4: [..., 4] complex with NaN for missing.
+            Returns lower median by magnitude (rank 1 among finite elements).
+            """
+            mags = v4.abs()
+            # NaN mags -> set to +inf so they go last
+            mags_key = torch.where(torch.isnan(mags), torch.full_like(mags, float('inf')), mags)
+            mags_sorted, idx = torch.sort(mags_key, dim=-1)
+            # Count finite elements
+            k = torch.sum(torch.isfinite(mags), dim=-1)  # [...]
+            # lower median rank = max(0, k//2 - 1)
+            rank = torch.clamp(k // 2 - 1, min=0).unsqueeze(-1)
+            pick = torch.gather(idx, -1, rank)
+            med = torch.gather(v4, -1, pick).squeeze(-1)
+            # If k==0 -> NaN+NaNj
+            med = torch.where(
+                k > 0,
+                med,
+                torch.full_like(med, complex(float('nan'), float('nan')))
+            )
+            return med
+
+        # ---- Branch: cell_ids shape [N] (shared mapping) ----
+        if cell_ids.ndim == 1:
+            groups = (cell_ids // 4).to(torch.long)         # [N] parent ids (global)
+            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]
+
+            if reduce in ("mean", "max"):
+                # Vectorized scatter path (same as before)
+                row_offsets = torch.arange(R, device=data.device).unsqueeze(1) * n_bins
+                idx = inv.unsqueeze(0).expand(R, -1) + row_offsets
+                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
+                else:  # "max"
+                    reduced = _segment_max(vals_flat, idx_flat, out_size)
+
+                output = reduced.view(*orig, n_bins)
+                return output, parents
+
+            elif reduce == "median":
+                # Build a 4-slot array per parent using child offset = cell_ids % 4
+                off = (cell_ids % 4).to(torch.long)         # [N] in {0,1,2,3}
+                out4 = torch.full((R, n_bins, 4),
+                                  torch.nan,
+                                  dtype=data.dtype,
+                                  device=data.device)
+                # flat indexing to scatter
+                base = torch.arange(R, device=data.device).unsqueeze(1) * (n_bins * 4)
+                flat_index = base + (inv.unsqueeze(0) * 4) + off.unsqueeze(0)  # [R, N]
+                out4 = out4.reshape(-1)
+                out4.scatter_(0, flat_index.reshape(-1), data_flat.reshape(-1))
+                out4 = out4.view(R, n_bins, 4)  # [R, n_bins, 4]
+
+                if torch.is_complex(data):
+                    med = _median_from_four_complex(out4)        # [R, n_bins]
+                else:
+                    med = _median_from_four_real(out4)           # [R, n_bins]
+
+                output = med.view(*orig, n_bins)
+                return output, parents
+            else:
+                raise ValueError("reduce must be 'mean', 'max', or 'median'.")
+
+        # ---- Branch: cell_ids shape [B, N] (per-batch mapping) ----
+        elif cell_ids.ndim == 2:
+            B = cell_ids.shape[0]
+            R = int(np.prod(orig)) if len(orig) > 0 else 1
+            data_flat = data.reshape(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 = (R // B_data) if B_data > 0 else 1  # repeats per batch row
+
+            outs_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 of data_flat that correspond to this batch row
+                start, stop = b * T, (b + 1) * T
+                rows = slice(start, stop)  # T rows -> [T, N]
+                vals = data_flat[rows]
+
+                if reduce in ("mean", "max"):
+                    row_offsets = torch.arange(T, device=data.device).unsqueeze(1) * n_bins_b
+                    idx = inv_b.unsqueeze(0).expand(T, -1) + row_offsets
+                    vals_flat = vals.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)
+                    else:
+                        reduced_bt = _segment_max(vals_flat, idx_flat, out_size).view(T, n_bins_b)
+                    outs_list.append(reduced_bt)
+                    groups_list.append(parents_b)
+                elif reduce == "median":
+                    off_b = (cell_ids[b] % 4).to(torch.long)  # [N] in {0,1,2,3}
+                    out4 = torch.full((T, n_bins_b, 4),
+                                      torch.nan,
+                                      dtype=data.dtype,
+                                      device=data.device)
+                    base = torch.arange(T, device=data.device).unsqueeze(1) * (n_bins_b * 4)
+                    flat_index = base + (inv_b.unsqueeze(0) * 4) + off_b.unsqueeze(0)  # [T, N]
+                    out4 = out4.reshape(-1)
+                    out4.scatter_(0, flat_index.reshape(-1), vals.reshape(-1))
+                    out4 = out4.view(T, n_bins_b, 4)  # [T, n_bins_b, 4]
+
+                    if torch.is_complex(data):
+                        reduced_bt = _median_from_four_complex(out4)  # [T, n_bins_b]
+                    else:
+                        reduced_bt = _median_from_four_real(out4)     # [T, n_bins_b]
+
+                    outs_list.append(reduced_bt)
+                    groups_list.append(parents_b)
+                else:
+                    raise ValueError("reduce must be 'mean', 'max', or 'median'.")
+
+            if not padded:
+                return outs_list, groups_list
+
+            # Padded output (B, T, max_bins) [+ mask]
+            out_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, (o_b, g_b) in enumerate(zip(outs_list, groups_list)):
+                nb = g_b.numel()
+                out_pad[b, :, :nb] = o_b
+                groups_pad[b, :nb] = g_b
+                mask[b, :nb] = True
+
+            if len(orig) > 1:
+                extra = orig[1:]
+                out_pad = out_pad.view(B, *extra, max_bins)
+            else:
+                out_pad = out_pad.view(B, max_bins)
+
+            return out_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])
@@ -404,6 +625,271 @@ class BkTorch(BackendBase.BackendBase):
 
         return torch.bincount(group_indices, weights=data) / counts, unique_groups
 
+    def bk_masked_median(self, x: torch.Tensor, mask: torch.Tensor,
+                                      max_iter: int = 100, tol: float = 1e-6, eps: float = 1e-12):
+        """
+        Masked geometric median over the last axis using Weiszfeld iteration (1D case).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Shape [a, b, c, N]. Can be real or complex.
+        mask : torch.Tensor
+            Binary mask of shape [a, b, 1, N]; broadcast across 'c'.
+        max_iter : int
+            Max number of Weiszfeld iterations.
+        tol : float
+            Convergence tolerance on the max absolute update per voxel.
+        eps : float
+            Small value to avoid division-by-zero in the weights.
+
+        Returns
+        -------
+        med  : torch.Tensor, shape [a, b, c]
+            Geometric median of x along the last axis where mask == 1.
+            - For complex x: distances use the complex magnitude |x - y|.
+              The returned median is complex.
+        med2 : torch.Tensor, shape [a, b, c]
+            Geometric median of squared values along the last axis where mask == 1.
+            - If x is real : median of x**2 (real).
+            - If x is complex : median of |x|**2 (real).
+        """
+
+        # --- helpers ---
+        def _nan_like(y: torch.Tensor) -> torch.Tensor:
+            """Return a NaN tensor with the same shape/dtype/device as y."""
+            if torch.is_complex(y):
+                return torch.full_like(y, complex(float('nan'), float('nan')))
+            else:
+                return torch.full_like(y, float('nan'))
+
+        def safe_nanmax(t: torch.Tensor) -> torch.Tensor:
+            """
+            Backward-compatible replacement for torch.nanmax.
+            Assumes t is real-valued (we only call it on absolute updates).
+            """
+            if torch.isnan(t).all():
+                return torch.tensor(float('nan'), device=t.device, dtype=t.dtype)
+            if torch.isnan(t).any():
+                return torch.max(t[~torch.isnan(t)])
+            return torch.max(t)
+
+        # --- prep shapes & mask ---
+        # Broadcast mask to x's shape [a,b,c,N]
+        mask_bool = mask.to(torch.bool).expand_as(x)               # [a,b,c,N]
+        m_float = mask_bool.to(dtype=x.real.dtype)                  # weights need real dtype
+
+        # Count valid samples per voxel
+        valid_counts = mask_bool.sum(dim=-1)                        # [a,b,c]
+        zero_valid = (valid_counts == 0)
+
+        # Denominator for masked mean initialization (avoid div-by-zero with clamp_min)
+        denom = valid_counts.clamp_min(1).to(dtype=x.real.dtype)    # real
+
+        # --- initialize y with masked mean (good starting point) ---
+        if torch.is_complex(x):
+            # (m_float*x) promotes to complex; denom to complex for division
+            y = (m_float * x).sum(dim=-1) / denom.to(dtype=x.dtype)  # [a,b,c], complex
+        else:
+            y = (m_float * x).sum(dim=-1) / denom                    # [a,b,c], real
+
+        # Put NaNs where there are no valid samples
+        y = torch.where(zero_valid, _nan_like(y), y)
+
+        # --- Weiszfeld iterations for x -> med ---
+        # y_{k+1} = sum_i (x_i / ||x_i - y_k||) / sum_i (1/||x_i - y_k||), masked
+        for _ in range(max_iter):
+            if torch.all(zero_valid):
+                break
+
+            diff = x - y.unsqueeze(-1)                              # [a,b,c,N]
+            dist = diff.abs()                                       # real, [a,b,c,N]
+            w = m_float * (1.0 / torch.clamp(dist, min=eps))        # real weights
+            w_sum = w.sum(dim=-1)                                   # [a,b,c], real
+            y_new = (w * x).sum(dim=-1) / w_sum.clamp_min(eps)      # [a,b,c], real/complex
+
+            # Keep NaNs on zero-valid voxels
+            y_new = torch.where(zero_valid, _nan_like(y_new), y_new)
+
+            # Convergence
+            upd = (y_new - y).abs()                                 # real
+            if safe_nanmax(upd).item() <= tol:
+                y = y_new
+                break
+            y = y_new
+
+        med = y  # [a,b,c]
+
+        # --- Weiszfeld iterations for squared values -> med2 ---
+        # For complex: use |x|^2; for real: x^2
+        s = (x.abs() ** 2) if torch.is_complex(x) else (x ** 2)     # [a,b,c,N], real
+        # Init with masked mean of s
+        z = (m_float * s).sum(dim=-1) / denom                        # [a,b,c], real
+        z = torch.where(zero_valid, _nan_like(z), z)
+
+        # Weiszfeld on real scalars s
+        for _ in range(max_iter):
+            if torch.all(zero_valid):
+                break
+            diff_s = s - z.unsqueeze(-1)                             # [a,b,c,N]
+            dist_s = diff_s.abs().clamp_min(eps)                     # real
+            w_s = m_float * (1.0 / dist_s)                           # real
+            w_s_sum = w_s.sum(dim=-1)                                # [a,b,c]
+            z_new = (w_s * s).sum(dim=-1) / w_s_sum.clamp_min(eps)   # [a,b,c]
+            z_new = torch.where(zero_valid, _nan_like(z_new), z_new)
+
+            upd_s = (z_new - z).abs()
+            if safe_nanmax(upd_s).item() <= tol:
+                z = z_new
+                break
+            z = z_new
+
+        med2 = z  # [a,b,c], real
+
+        return med, med2
+
+    def bk_masked_median_2d_weiszfeld(self,
+                                      x: torch.Tensor,
+                                      mask: torch.Tensor,
+                                      max_iter: int = 100,
+                                      tol: float = 1e-6,
+                                      eps: float = 1e-12):
+        """
+        Masked geometric median over 2D spatial axes using Weiszfeld iteration.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input of shape [a, b, c, N1, N2]. Can be real or complex.
+        mask : torch.Tensor
+            Binary mask of shape [a, b, 1, N1, N2]; broadcasted over 'c'.
+        max_iter : int
+            Maximum number of Weiszfeld iterations.
+        tol : float
+            Stopping tolerance on the max absolute update per voxel.
+        eps : float
+            Small positive value to avoid division by zero in weights.
+
+        Returns
+        -------
+        med : torch.Tensor, shape [a, b, c]
+            Geometric median of x over (N1, N2) where mask == 1.
+            - If x is complex, distances are magnitudes |x - y| in the complex plane,
+              and the returned value is the complex sample estimate (not its magnitude).
+        med2 : torch.Tensor, shape [a, b, c]
+            Geometric median of squared values over (N1, N2) where mask == 1.
+            - If x is real : median of x**2 (via Weiszfeld in 1D).
+            - If x is complex : median of |x|**2 (real, non-negative).
+
+        Notes
+        -----
+        - Voxels with zero valid samples return NaN (NaN+NaNj for complex med).
+        - Weiszfeld update: y_{k+1} = sum_i w_i x_i / sum_i w_i with w_i = 1 / ||x_i - y_k||.
+          Here ||.|| is |.| for real numbers and the complex magnitude for complex numbers.
+        """
+        # Broadcast mask to x's shape [a,b,c,N1,N2] and flatten spatial dims to N
+        mask_bool = mask.to(torch.bool).expand_as(x)
+        a, b, c, N1, N2 = x.shape
+        N = N1 * N2
+        x_flat = x.reshape(a, b, c, N)
+        m_flat = mask_bool.reshape(a, b, c, N)
+
+        # Count valid samples per voxel
+        valid_counts = m_flat.sum(dim=-1)  # [a,b,c]
+
+        # Helper to create NaN of the right dtype
+        def _nan_like(y):
+            if torch.is_complex(y):
+                return torch.full_like(y, complex(float('nan'), float('nan')))
+            else:
+                return torch.full_like(y, float('nan'))
+
+        # --- Geometric median of x (real or complex) ---
+        # Initialize y0: masked mean (robust enough as a starting point)
+        # y0 = sum(mask * x) / sum(mask)
+        denom = valid_counts.clamp_min(1).to(x.dtype)
+        if torch.is_complex(x):
+            denom_c = denom.to(x.dtype)
+            y = ( (m_flat * x_flat).sum(dim=-1) / denom_c )  # [a,b,c]
+        else:
+            y = ( (m_flat * x_flat).sum(dim=-1) / denom )    # [a,b,c]
+
+        # Where there are zero valid samples, set to NaN now (and keep NaN through)
+        zero_valid = (valid_counts == 0)
+        if torch.is_complex(x):
+            y = torch.where(zero_valid, torch.full_like(y, complex(float('nan'), float('nan'))), y)
+        else:
+            y = torch.where(zero_valid, torch.full_like(y, float('nan')), y)
+
+        # helper: nanmax replacement
+        def safe_nanmax(t):
+            if torch.isnan(t).all():
+                return torch.tensor(float('nan'), device=t.device, dtype=torch.float32)
+            return torch.max(t[~torch.isnan(t)]) if torch.isnan(t).any() else torch.max(t)
+    
+        # Iterate Weiszfeld
+        for _ in range(max_iter):
+            # Skip voxels with no valid samples
+            if torch.all(zero_valid):
+                break
+
+            # diff: [a,b,c,N], distances are |diff|
+            diff = x_flat - y.unsqueeze(-1)                      # broadcast y over N
+            dist = diff.abs()                                    # real tensor, [a,b,c,N]
+
+            # weights w = mask / max(dist, eps)
+            w = m_flat * (1.0 / torch.clamp(dist, min=eps))      # [a,b,c,N]
+            w_sum = w.sum(dim=-1)                                # [a,b,c], real
+
+            # Next iterate y_new = sum(w * x) / sum(w)
+            # For complex x, w is real so (w*x) is complex — OK.
+            y_new = (w * x_flat).sum(dim=-1) / w_sum.clamp_min(eps)
+
+            # Keep NaN on zero-valid voxels
+            if torch.is_complex(x):
+                y_new = torch.where(zero_valid, torch.full_like(y_new, complex(float('nan'), float('nan'))), y_new)
+            else:
+                y_new = torch.where(zero_valid, torch.full_like(y_new, float('nan')), y_new)
+
+            # Convergence check (max absolute update over all voxels)
+            upd = (y_new - y).abs()
+            if safe_nanmax(upd).item() <= tol:
+                y = y_new
+                break
+            y = y_new
+
+        med = y  # [a,b,c]
+
+        # --- Geometric median of squared values (med2) ---
+        if torch.is_complex(x):
+            s_flat = (x_flat.abs() ** 2)                         # [a,b,c,N], real
+        else:
+            s_flat = (x_flat ** 2)                               # [a,b,c,N], real
+
+        # Initialize z0 = masked mean of s
+        z = (m_flat * s_flat).sum(dim=-1) / denom               # [a,b,c], real
+        z = torch.where(zero_valid, torch.full_like(z, float('nan')), z)
+
+        # Weiszfeld on scalars (1D) for s: distances are |s_i - z|
+        for _ in range(max_iter):
+            if torch.all(zero_valid):
+                break
+            diff_s = s_flat - z.unsqueeze(-1)                    # [a,b,c,N]
+            dist_s = diff_s.abs().clamp_min(eps)                 # avoid div-by-zero
+            w_s = m_flat * (1.0 / dist_s)
+            w_s_sum = w_s.sum(dim=-1)
+            z_new = (w_s * s_flat).sum(dim=-1) / w_s_sum.clamp_min(eps)
+            z_new = torch.where(zero_valid, torch.full_like(z_new, float('nan')), z_new)
+            upd_s = (z_new - z).abs()
+            if safe_nanmax(upd_s).item() <= tol:
+                z = z_new
+                break
+            z = z_new
+
+        med2 = z  # [a,b,c], real
+
+        return med, med2
+    
     # ---------------------------------------------−---------
     # --             BACKEND DEFINITION                    --
     # ---------------------------------------------−---------
@@ -578,6 +1064,14 @@ class BkTorch(BackendBase.BackendBase):
             return self.backend.mean(data)
         else:
             return self.backend.mean(data, axis)
+        
+    def bk_reduce_median(self, data, axis=None):
+
+        if axis is None:
+            res,_ = self.backend.median(data)
+        else:
+            res,_ = self.backend.median(data, axis)
+        return res
 
     def bk_reduce_min(self, data, axis=None):