foscat 2025.5.0__py3-none-any.whl → 2025.5.2__py3-none-any.whl

foscat/BkTensorflow.py

@@ -68,10 +68,16 @@ class BkTensorflow(BackendBase.BackendBase):
                 print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
                 sys.stdout.flush()
                 self.ngpu = len(logical_gpus)
-                gpuname = logical_gpus[self.gpupos % self.ngpu].name
-                self.gpulist = {}
-                for i in range(self.ngpu):
-                    self.gpulist[i] = logical_gpus[i].name
+                if self.ngpu > 0:
+                    gpuname = logical_gpus[self.gpupos % self.ngpu].name
+                    self.gpulist = {}
+                    for i in range(self.ngpu):
+                        self.gpulist[i] = logical_gpus[i].name
+                else:
+                    gpuname = "CPU:0"
+                    self.gpulist = {}
+                    self.gpulist[0] = gpuname
+                    self.ngpu = 1
 
             except RuntimeError as e:
                 # Memory growth must be set before GPUs have been initialized
@@ -122,18 +128,136 @@ class BkTensorflow(BackendBase.BackendBase):
 
         return x_padded
 
-    def conv2d(self, x, w, strides=[1, 1, 1, 1], padding="SAME"):
-        kx = w.shape[0]
-        ky = w.shape[1]
-        x_padded = self.periodic_pad(x, kx // 2, ky // 2)
-        return self.backend.nn.conv2d(x_padded, w, strides=strides, padding="VALID")
+    def binned_mean(self, data, cell_ids):
+        """
+        data: Tensor of shape [..., N] (float32 or float64)
+        cell_ids: Tensor of shape [N], int indices in [0, n_bins)
+        Returns: mean per bin, shape [..., n_bins]
+        """
+        ishape = list(data.shape)
+        A = 1
+        for k in range(len(ishape) - 1):
+            A *= ishape[k]
+        N = tf.shape(data)[-1]
+
+        # Step 1: group indices
+        groups = tf.math.floordiv(cell_ids, 4)  # [N]
+        unique_groups, I = tf.unique(groups)  # I: [N]
+        n_bins = tf.shape(unique_groups)[0]
+
+        # Step 2: build I_tiled with batch + channel offsets
+        I_tiled = tf.tile(I[None, :], [A, 1])  # shape [, N]
+
+        # Offset index to flatten across [A, n_bins]
+        batch_channel_offsets = tf.range(A)[:, None] * n_bins
+        I_offset = I_tiled + batch_channel_offsets  # shape [A, N]]
+
+        # Step 3: flatten data to shape [A, N]
+        data_reshaped = tf.reshape(data, [A, N])  # shape [A, N]
+
+        # Flatten all for scatter_nd
+        indices = tf.reshape(I_offset, [-1])  # [A*N]
+        values = tf.reshape(data_reshaped, [-1])  # [A*N]
+
+        """
+        # Prepare for scatter: indices → [A*N, 1]
+        scatter_indices = tf.expand_dims(indices, axis=1)
+        scatter_indices = tf.cast(scatter_indices, tf.int64)
+        """
+        total_bins = A * n_bins
+
+        # Step 4: sum per bin
+        sum_per_bin = tf.math.unsorted_segment_sum(values, indices, total_bins)
+        sum_per_bin = tf.reshape(sum_per_bin, ishape[0:-1] + [n_bins])  # [A, n_bins]
+
+        # Step 5: count per bin (same indices)
+        counts = tf.math.unsorted_segment_sum(1.0 + 0 * values, indices, total_bins)
+        # counts = tf.math.bincount(indices, minlength=total_bins, maxlength=total_bins)
+        counts = tf.reshape(counts, ishape[0:-1] + [n_bins])
+        # counts = tf.maximum(counts, 1)  # Avoid division by zero
+        # counts = tf.cast(counts, dtype=data.dtype)
+
+        # Step 6: mean
+        mean_per_bin = sum_per_bin / counts  # [B, A, n_bins]
+
+        return mean_per_bin, unique_groups
+
+    def conv2d(self, x, w):
+        """
+        Perform 2D convolution using TensorFlow.
+
+        Args:
+            x: Tensor of shape [..., Nx, Ny] – input
+            w: Tensor of shape [O_c, wx, wy] – conv weights
+
+        Returns:
+            Tensor of shape [..., O_c, Nx, Ny]
+        """
+        # Extract shape
+        *leading_dims, Nx, Ny = x.shape
+        O_c, wx, wy = w.shape
+
+        # Flatten leading dims into a batch dimension
+        B = tf.reduce_prod(leading_dims) if leading_dims else 1
+        x = tf.reshape(x, [B, Nx, Ny, 1])  # TensorFlow format: [B, H, W, C_in=1]
+
+        # Reshape weights to [wx, wy, in_channels=1, out_channels]
+        w = tf.reshape(w, [O_c, wx, wy])
+        w = tf.transpose(w, perm=[1, 2, 0])  # [wx, wy, O_c]
+        w = tf.reshape(w, [wx, wy, 1, O_c])  # [wx, wy, C_in=1, C_out]
+
+        # Apply 'reflect' padding manually
+        pad_x = wx // 2
+        pad_y = wy // 2
+        x_padded = tf.pad(
+            x, [[0, 0], [pad_x, pad_x], [pad_y, pad_y], [0, 0]], mode="REFLECT"
+        )
+
+        # Perform convolution
+        y = tf.nn.conv2d(
+            x_padded, w, strides=[1, 1, 1, 1], padding="VALID"
+        )  # [B, Nx, Ny, O_c]
+
+        # Transpose back to match original format: [..., O_c, Nx, Ny]
+        y = tf.transpose(y, [0, 3, 1, 2])  # [B, O_c, Nx, Ny]
+        y = tf.reshape(y, [*leading_dims, O_c, Nx, Ny])
+
+        return y
+
+    def conv1d(self, x, w):
+        """
+        Perform 1D convolution using TensorFlow.
+
+        Args:
+            x: Tensor of shape [..., N] – input
+            w: Tensor of shape [k] – conv weights
+
+        Returns:
+            Tensor of shape [...,N]
+        """
+        # Extract shapes
+        *leading_dims, N = x.shape
+        k = w.shape[0]
+
+        # Flatten leading dims into batch dimension
+        B = tf.reduce_prod(leading_dims) if leading_dims else 1
+        x = tf.reshape(x, [B, N, 1])  # TensorFlow 1D format: [B, L, C=1]
+
+        # Prepare weights: [k, in_channels=1, out_channels=O_c]
+        w = tf.reshape(w, [k, 1, 1])
+
+        # Apply 'reflect' padding
+        pad = k // 2
+        x_padded = tf.pad(x, [[0, 0], [pad, pad], [0, 0]], mode="REFLECT")
+
+        # Perform convolution
+        y = tf.nn.conv1d(x_padded, w, stride=1, padding="VALID")  # [B, N, O_c]
 
-    def conv1d(self, x, w, strides=[1, 1, 1], padding="SAME"):
-        kx = w.shape[0]
-        paddings = self.backend.constant([[0, 0], [kx // 2, kx // 2], [0, 0]])
-        tmp = self.backend.pad(x, paddings, "SYMMETRIC")
+        # Transpose to [B, O_c, N] and reshape back
+        y = tf.transpose(y, [0, 2, 1])  # [B, 1, N]
+        y = tf.reshape(y, [*leading_dims, N])  # [..., N]
 
-        return self.backend.nn.conv1d(tmp, w, stride=strides, padding="VALID")
+        return y
 
     def bk_threshold(self, x, threshold, greater=True):
 

foscat/BkTorch.py

@@ -2,6 +2,7 @@ import sys
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 
 import foscat.BkBase as BackendBase
 
@@ -62,55 +63,69 @@ 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):
         """
-        data: Tensor of shape [B, N, A]
-        I: Tensor of shape [N], integer indices in [0, n_bins)
-        Returns: mean per bin, shape [B, n_bins, A]
+        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).
+
+        Returns:
+            torch.Tensor: Tensor of shape [..., n_bins], with averaged values per group of 4 cells.
         """
-        groups = cell_ids // 4  # [N]
+        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
+            )
 
-        unique_groups, I = np.unique(groups, return_inverse=True)
+        # 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]
 
-        B = data.shape[0]
+        # Flatten all leading dimensions into a single batch dimension
+        original_shape = data.shape[:-1]
+        N = data.shape[-1]
+        data_flat = data.reshape(-1, N)  # Shape: [B, N]
 
-        counts = torch.bincount(torch.tensor(I).to(data.device))[None, :]
+        # Prepare to compute sums using scatter_add
+        B = data_flat.shape[0]
 
-        I = np.tile(I, B) + np.tile(n_bins * np.arange(B, dtype="int"), data.shape[1])
+        # Repeat inverse indices for each batch element
+        idx = inverse_indices.repeat(B, 1)  # Shape: [B, N]
 
-        if len(data.shape) == 3:
-            A = data.shape[2]
-            I = np.repeat(I, A) * A + np.repeat(
-                np.arange(A, dtype="int"), data.shape[1] * B
-            )
+        # 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]
 
-        I = torch.tensor(I).to(data.device)
+        # Flatten everything for scatter
+        idx_offset_flat = idx_offset.flatten()
+        data_flat_flat = data_flat.flatten()
 
-        # Comptage par bin
-        if len(data.shape) == 2:
-            sum_per_bin = torch.zeros(
-                [B * n_bins], dtype=data.dtype, device=data.device
-            )
-            sum_per_bin = sum_per_bin.scatter_add(
-                0, I, self.bk_reshape(data, B * data.shape[1])
-            ).reshape(B, n_bins)
+        # 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)
 
-            mean_per_bin = sum_per_bin / counts  # [B, n_bins, A]
-        else:
-            sum_per_bin = torch.zeros(
-                [B * n_bins * A], dtype=data.dtype, device=data.device
-            )
-            sum_per_bin = sum_per_bin.scatter_add(
-                0, I, self.bk_reshape(data, B * data.shape[1] * A)
-            ).reshape(
-                B, n_bins, A
-            )  # [B, n_bins]
+        # 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)
 
-            mean_per_bin = sum_per_bin / counts[:, :, None]  # [B, n_bins, A]
+        # Compute mean
+        mean = out / counts  # Shape: [B * n_bins]
+        mean = mean.view(B, n_bins)
 
-        return mean_per_bin, unique_groups
+        # Restore original leading dimensions
+        return mean.view(*original_shape, n_bins), unique_groups
 
     def average_by_cell_group(data, cell_ids):
         """
@@ -146,20 +161,40 @@ class BkTorch(BackendBase.BackendBase):
     def bk_sparse_dense_matmul(self, smat, mat):
         return smat.matmul(mat)
 
-    def conv2d(self, x, w, strides=[1, 1, 1, 1], padding="SAME"):
-        import torch.nn.functional as F
+    def conv2d(self, x, w):
+        """
+        Perform 2D convolution using PyTorch format.
+
+        Args:
+            x: Tensor of shape [..., Nx, Ny] – input
+            w: Tensor of shape [O_c, wx, wy] – conv weights
+
+        Returns:
+            Tensor of shape [..., O_c, Nx, Ny]
+        """
+        *leading_dims, Nx, Ny = x.shape  # extract leading dims
+        O_c, wx, wy = w.shape
 
-        lx = x.permute(0, 3, 1, 2)
-        wx = w.permute(3, 2, 0, 1)  # de (5, 5, 1, 4) à (4, 1, 5, 5)
+        # Flatten leading dims into batch dimension
+        B = int(torch.prod(torch.tensor(leading_dims))) if leading_dims else 1
+        x = x.reshape(B, 1, Nx, Ny)  # [B, 1, Nx, Ny]
 
-        # Calculer le padding symétrique
-        kx, ky = w.shape[0], w.shape[1]
+        # Reshape filters to match conv2d format [O_c, 1, wx, wy]
+        w = w[:, None, :, :]  # [O_c, 1, wx, wy]
 
-        # Appliquer le padding
-        x_padded = F.pad(lx, (ky // 2, ky // 2, kx // 2, kx // 2), mode="circular")
+        pad_x = wx // 2
+        pad_y = wy // 2
 
-        # Appliquer la convolution
-        return F.conv2d(x_padded, wx, stride=1, padding=0).permute(0, 2, 3, 1)
+        # Reflective padding to reduce edge artifacts
+        x_padded = F.pad(x, (pad_y, pad_y, pad_x, pad_x), mode="reflect")
+
+        # Apply convolution
+        y = F.conv2d(x_padded, w)  # [B, O_c, Nx, Ny]
+
+        # Restore original leading dimensions
+        y = y.reshape(*leading_dims, O_c, Nx, Ny)
+
+        return y
 
     def conv1d(self, x, w, strides=[1, 1, 1], padding="SAME"):
         # to be written!!!
@@ -376,7 +411,7 @@ class BkTorch(BackendBase.BackendBase):
         return self.backend.unsqueeze(data, axis)
 
     def bk_transpose(self, data, thelist):
-        return self.backend.transpose(data, thelist)
+        return self.backend.transpose(data, thelist[0], thelist[1])
 
     def bk_concat(self, data, axis=None):