foscat 2025.5.0__py3-none-any.whl → 2025.6.1__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):
         """
@@ -134,11 +149,7 @@ class BkTorch(BackendBase.BackendBase):
     # --             BACKEND DEFINITION                    --
     # ---------------------------------------------−---------
     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.T, w, dense_shape).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)
@@ -146,20 +157,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
+
+        # 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]
+
+        # Reshape filters to match conv2d format [O_c, 1, wx, wy]
+        w = w[:, None, :, :]  # [O_c, 1, wx, wy]
 
-        lx = x.permute(0, 3, 1, 2)
-        wx = w.permute(3, 2, 0, 1)  # de (5, 5, 1, 4) à (4, 1, 5, 5)
+        pad_x = wx // 2
+        pad_y = wy // 2
 
-        # Calculer le padding symétrique
-        kx, ky = w.shape[0], w.shape[1]
+        # Reflective padding to reduce edge artifacts
+        x_padded = F.pad(x, (pad_y, pad_y, pad_x, pad_x), mode="reflect")
 
-        # Appliquer le padding
-        x_padded = F.pad(lx, (ky // 2, ky // 2, kx // 2, kx // 2), mode="circular")
+        # Apply convolution
+        y = F.conv2d(x_padded, w)  # [B, O_c, Nx, Ny]
 
-        # Appliquer la convolution
-        return F.conv2d(x_padded, wx, stride=1, padding=0).permute(0, 2, 3, 1)
+        # 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!!!
@@ -211,13 +242,13 @@ class BkTorch(BackendBase.BackendBase):
             xr = self.bk_real(x)
             # xi = self.bk_imag(x)
 
-            r = self.backend.sign(xr) * self.backend.sqrt(self.backend.sign(xr) * xr)
+            r = self.backend.sign(xr) * self.backend.sqrt(self.backend.sign(xr) * xr + 1E-16)
             # return r
             # i = self.backend.sign(xi) * self.backend.sqrt(self.backend.sign(xi) * xi)
 
             return r
         else:
-            return self.backend.sign(x) * self.backend.sqrt(self.backend.sign(x) * x)
+            return self.backend.sign(x) * self.backend.sqrt(self.backend.sign(x) * x + 1E-16)
 
     def bk_square_comp(self, x):
         if x.dtype == self.all_cbk_type:
@@ -356,9 +387,9 @@ class BkTorch(BackendBase.BackendBase):
         return self.backend.argmax(data)
 
     def bk_reshape(self, data, shape):
-        if isinstance(data, np.ndarray):
-            return data.reshape(shape)
-        return data.view(shape)
+        #if isinstance(data, np.ndarray):
+        #    return data.reshape(shape)
+        return data.reshape(shape)
 
     def bk_repeat(self, data, nn, axis=0):
         return self.backend.repeat_interleave(data, repeats=nn, dim=axis)
@@ -376,7 +407,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):
 
@@ -405,7 +436,9 @@ class BkTorch(BackendBase.BackendBase):
         return self.backend.zeros(shape, dtype=dtype).to(self.torch_device)
 
     def bk_gather(self, data, idx, axis=0):
-        if axis == 0:
+        if axis == -1:
+            return data[...,idx]
+        elif axis == 0:
             return data[idx]
         elif axis == 1:
             return data[:, idx]
@@ -413,7 +446,7 @@ class BkTorch(BackendBase.BackendBase):
             return data[:, :, idx]
         elif axis == 3:
             return data[:, :, :, idx]
-        return data[:, :, :, :, idx]
+        return data[idx,...]
 
     def bk_reverse(self, data, axis=0):
         return self.backend.flip(data, dims=[axis])

foscat/CNN.py

@@ -9,13 +9,12 @@ class CNN:
 
     def __init__(
         self,
-        scat_operator=None,
         nparam=1,
-        nscale=1,
+        KERNELSZ=3,
+        NORIENT=4,
         chanlist=[],
         in_nside=1,
         n_chan_in=1,
-        nbatch=1,
         SEED=1234,
         filename=None,
     ):
@@ -31,31 +30,30 @@ class CNN:
             self.in_nside = outlist[4]
             self.nbatch = outlist[1]
             self.n_chan_in = outlist[8]
+            self.NORIENT = outlist[9]
             self.x = self.scat_operator.backend.bk_cast(outlist[6])
             self.out_nside = self.in_nside // (2**self.nscale)
         else:
-            self.nscale = nscale
-            self.nbatch = nbatch
+            self.nscale = len(chanlist)-1
             self.npar = nparam
             self.n_chan_in = n_chan_in
             self.scat_operator = scat_operator
-            if len(chanlist) != nscale + 1:
-                print(
-                    "len of chanlist (here %d) should of nscale+1 (here %d)"
-                    % (len(chanlist), nscale + 1)
-                )
-                return None
+            if self.scat_operator is None:
+                self.scat_operator = sc.funct(
+                    KERNELSZ=KERNELSZ,
+                    NORIENT=NORIENT)
 
             self.chanlist = chanlist
-            self.KERNELSZ = scat_operator.KERNELSZ
-            self.all_type = scat_operator.all_type
+            self.KERNELSZ = self.scat_operator.KERNELSZ
+            self.NORIENT = self.scat_operator.NORIENT
+            self.all_type = self.scat_operator.all_type
             self.in_nside = in_nside
             self.out_nside = self.in_nside // (2**self.nscale)
-
+            self.backend = self.scat_operator.backend
             np.random.seed(SEED)
-            self.x = scat_operator.backend.bk_cast(
-                np.random.randn(self.get_number_of_weights())
-                / (self.KERNELSZ * self.KERNELSZ)
+            self.x = self.scat_operator.backend.bk_cast(
+                np.random.rand(self.get_number_of_weights())
+                / (self.KERNELSZ * (self.KERNELSZ//2+1)*self.NORIENT)
             )
 
     def save(self, filename):
@@ -70,6 +68,7 @@ class CNN:
             self.get_weights().numpy(),
             self.all_type,
             self.n_chan_in,
+            self.NORIENT,
         ]
 
         myout = open("%s.pkl" % (filename), "wb")
@@ -82,8 +81,8 @@ class CNN:
             totnchan = totnchan + self.chanlist[i] * self.chanlist[i + 1]
         return (
             self.npar * 12 * self.out_nside**2 * self.chanlist[self.nscale]
-            + totnchan * self.KERNELSZ * self.KERNELSZ
-            + self.KERNELSZ * self.KERNELSZ * self.n_chan_in * self.chanlist[0]
+            + totnchan * self.KERNELSZ * (self.KERNELSZ//2+1)
+            + self.KERNELSZ * (self.KERNELSZ//2+1) * self.n_chan_in * self.chanlist[0]
         )
 
     def set_weights(self, x):
@@ -95,30 +94,32 @@ class CNN:
     def eval(self, im, indices=None, weights=None):
 
         x = self.x
-        ww = self.scat_operator.backend.bk_reshape(
-            x[0 : self.KERNELSZ * self.KERNELSZ * self.n_chan_in * self.chanlist[0]],
-            [self.KERNELSZ * self.KERNELSZ, self.n_chan_in, self.chanlist[0]],
+        ww = self.backend.bk_reshape(
+            x[0 : self.KERNELSZ * (self.KERNELSZ//2+1) * self.n_chan_in * self.chanlist[0]],
+            [self.n_chan_in, self.KERNELSZ * (self.KERNELSZ//2+1),  self.chanlist[0]],
         )
-        nn = self.KERNELSZ * self.KERNELSZ * self.n_chan_in * self.chanlist[0]
+        nn = self.KERNELSZ * (self.KERNELSZ//2+1) * self.n_chan_in * self.chanlist[0]
 
         im = self.scat_operator.healpix_layer(im, ww)
-        im = self.scat_operator.backend.bk_relu(im)
+        im = self.backend.bk_relu(im)
+        
+        im = self.backend.bk_reduce_mean(self.backend.bk_reshape(im,[im.shape[0],im.shape[1],im.shape[2]//4,4]),3)
 
         for k in range(self.nscale):
             ww = self.scat_operator.backend.bk_reshape(
                 x[
                     nn : nn
                     + self.KERNELSZ
-                    * self.KERNELSZ
+                    * (self.KERNELSZ//2+1)
                     * self.chanlist[k]
                     * self.chanlist[k + 1]
                 ],
-                [self.KERNELSZ * self.KERNELSZ, self.chanlist[k], self.chanlist[k + 1]],
+                [self.chanlist[k], self.KERNELSZ * (self.KERNELSZ//2+1),  self.chanlist[k + 1]],
             )
             nn = (
                 nn
                 + self.KERNELSZ
-                * self.KERNELSZ
+                * (self.KERNELSZ//2)
                 * self.chanlist[k]
                 * self.chanlist[k + 1]
             )
@@ -129,7 +130,7 @@ class CNN:
                     im, ww, indices=indices[k], weights=weights[k]
                 )
             im = self.scat_operator.backend.bk_relu(im)
-            im = self.scat_operator.ud_grade_2(im, axis=0)
+            im = self.backend.bk_reduce_mean(self.backend.bk_reshape(im,[im.shape[0],im.shape[1],im.shape[2]//4,4]),3)
 
         ww = self.scat_operator.backend.bk_reshape(
             x[
@@ -141,11 +142,11 @@ class CNN:
 
         im = self.scat_operator.backend.bk_matmul(
             self.scat_operator.backend.bk_reshape(
-                im, [1, 12 * self.out_nside**2 * self.chanlist[self.nscale]]
+                im, [im.shape[0], im.shape[1] * im.shape[2]]
             ),
             ww,
         )
-        im = self.scat_operator.backend.bk_reshape(im, [self.npar])
+        #im = self.scat_operator.backend.bk_reshape(im, [self.npar])
         im = self.scat_operator.backend.bk_relu(im)
 
         return im