foscat 3.8.2__py3-none-any.whl → 2025.3.0__py3-none-any.whl

foscat/FoCUS.py

@@ -1,5 +1,4 @@
 import os
-import os
 import sys
 
 import healpy as hp
@@ -11,32 +10,32 @@ TMPFILE_VERSION = "V4_0"
 
 class FoCUS:
     def __init__(
-            self,
-            NORIENT=4,
-            LAMBDA=1.2,
-            KERNELSZ=3,
-            slope=1.0,
-            all_type="float32",
-            nstep_max=16,
-            padding="SAME",
-            gpupos=0,
-            mask_thres=None,
-            mask_norm=False,
-            isMPI=False,
-            TEMPLATE_PATH="data",
-            BACKEND="tensorflow",
-            use_2D=False,
-            use_1D=False,
-            return_data=False,
-            JmaxDelta=0,
-            DODIV=False,
-            InitWave=None,
-            silent=True,
-            mpi_size=1,
-            mpi_rank=0,
+        self,
+        NORIENT=4,
+        LAMBDA=1.2,
+        KERNELSZ=3,
+        slope=1.0,
+        all_type="float32",
+        nstep_max=16,
+        padding="SAME",
+        gpupos=0,
+        mask_thres=None,
+        mask_norm=False,
+        isMPI=False,
+        TEMPLATE_PATH="data",
+        BACKEND="tensorflow",
+        use_2D=False,
+        use_1D=False,
+        return_data=False,
+        JmaxDelta=0,
+        DODIV=False,
+        InitWave=None,
+        silent=True,
+        mpi_size=1,
+        mpi_rank=0,
     ):
 
-        self.__version__ = "3.8.2"
+        self.__version__ = "2025.03.0"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -45,12 +44,18 @@ class FoCUS:
         self.mask_thres = mask_thres
         self.mask_norm = mask_norm
         self.InitWave = InitWave
-        self.mask_mask=None
+        self.mask_mask = None
         self.mpi_size = mpi_size
         self.mpi_rank = mpi_rank
         self.return_data = return_data
         self.silent = silent
 
+        self.kernel_smooth = {}
+        self.padding_smooth = {}
+        self.kernelR_conv = {}
+        self.kernelI_conv = {}
+        self.padding_conv = {}
+
         if not self.silent:
             print("================================================")
             print("          START FOSCAT CONFIGURATION")
@@ -69,10 +74,7 @@ class FoCUS:
                 if not self.silent:
                     print("The directory %s is created")
             except:
-                if not self.silent:
-                    print(
-                        "Impossible to create the directory %s" % (self.TEMPLATE_PATH)
-                    )
+                print("Impossible to create the directory %s" % (self.TEMPLATE_PATH))
                 return None
 
         self.number_of_loss = 0
@@ -82,10 +84,9 @@ class FoCUS:
         self.padding = padding
 
         if JmaxDelta != 0:
-            if not self.silent:
-                print(
-                    "OPTION JmaxDelta is not avialable anymore after version 3.6.2. Please use Jmax option in eval function"
-                )
+            print(
+                "OPTION JmaxDelta is not avialable anymore after version 3.6.2. Please use Jmax option in eval function"
+            )
             return None
 
         self.OSTEP = JmaxDelta
@@ -105,31 +106,34 @@ class FoCUS:
 
         self.all_type = all_type
         self.BACKEND = BACKEND
-        
-        if BACKEND=='torch':
+
+        if BACKEND == "torch":
             from foscat.BkTorch import BkTorch
+
             self.backend = BkTorch(
                 all_type=all_type,
                 mpi_rank=mpi_rank,
                 gpupos=gpupos,
                 silent=self.silent,
-                )
-        elif BACKEND=='tensorflow':
+            )
+        elif BACKEND == "tensorflow":
             from foscat.BkTensorflow import BkTensorflow
+
             self.backend = BkTensorflow(
                 all_type=all_type,
                 mpi_rank=mpi_rank,
                 gpupos=gpupos,
                 silent=self.silent,
-                )
+            )
         else:
             from foscat.BkNumpy import BkNumpy
+
             self.backend = BkNumpy(
                 all_type=all_type,
                 mpi_rank=mpi_rank,
                 gpupos=gpupos,
                 silent=self.silent,
-                )
+            )
 
         self.all_bk_type = self.backend.all_bk_type
         self.all_cbk_type = self.backend.all_cbk_type
@@ -172,9 +176,9 @@ class FoCUS:
         self.Y_CNN = {}
         self.Z_CNN = {}
 
-        self.filters_set={}
-        self.edge_masks={}
-        
+        self.filters_set = {}
+        self.edge_masks = {}
+
         wwc = np.zeros([KERNELSZ**2, l_NORIENT]).astype(all_type)
         wws = np.zeros([KERNELSZ**2, l_NORIENT]).astype(all_type)
 
@@ -209,15 +213,27 @@ class FoCUS:
             w_smooth = w_smooth.flatten()
         else:
             for i in range(NORIENT):
-                a = (NORIENT-1-i) / float(NORIENT) * np.pi # get the same angle number than scattering lib
-                if KERNELSZ<5:
-                    xx = (3 / float(KERNELSZ)) * LAMBDA * (x * np.cos(a) + y * np.sin(a))
-                    yy = (3 / float(KERNELSZ)) * LAMBDA * (x * np.sin(a) - y * np.cos(a))
+                a = (
+                    (NORIENT - 1 - i) / float(NORIENT) * np.pi
+                )  # get the same angle number than scattering lib
+                if KERNELSZ < 5:
+                    xx = (
+                        (3 / float(KERNELSZ)) * LAMBDA * (x * np.cos(a) + y * np.sin(a))
+                    )
+                    yy = (
+                        (3 / float(KERNELSZ)) * LAMBDA * (x * np.sin(a) - y * np.cos(a))
+                    )
                 else:
-                    xx = (3 /5) * LAMBDA * (x * np.cos(a) + y * np.sin(a))
-                    yy = (3 /5) * LAMBDA * (x * np.sin(a) - y * np.cos(a))
+                    xx = (3 / 5) * LAMBDA * (x * np.cos(a) + y * np.sin(a))
+                    yy = (3 / 5) * LAMBDA * (x * np.sin(a) - y * np.cos(a))
                 if KERNELSZ == 5:
-                    w_smooth=np.exp(-2*((3.0/float(KERNELSZ)*xx)**2+(3.0/float(KERNELSZ)*yy)**2))
+                    w_smooth = np.exp(
+                        -2
+                        * (
+                            (3.0 / float(KERNELSZ) * xx) ** 2
+                            + (3.0 / float(KERNELSZ) * yy) ** 2
+                        )
+                    )
                 else:
                     w_smooth = np.exp(-0.5 * (xx**2 + yy**2))
                 tmp1 = np.cos(yy * np.pi) * w_smooth
@@ -225,7 +241,7 @@ class FoCUS:
 
                 wwc[:, i] = tmp1.flatten() - tmp1.mean()
                 wws[:, i] = tmp2.flatten() - tmp2.mean()
-                #sigma = np.sqrt((wwc[:, i] ** 2).mean())
+                # sigma = np.sqrt((wwc[:, i] ** 2).mean())
                 sigma = np.mean(w_smooth)
                 wwc[:, i] /= sigma
                 wws[:, i] /= sigma
@@ -239,7 +255,7 @@ class FoCUS:
 
                     wwc[:, NORIENT] = tmp1.flatten() - tmp1.mean()
                     wws[:, NORIENT] = tmp2.flatten() - tmp2.mean()
-                    #sigma = np.sqrt((wwc[:, NORIENT] ** 2).mean())
+                    # sigma = np.sqrt((wwc[:, NORIENT] ** 2).mean())
                     sigma = np.mean(w_smooth)
 
                     wwc[:, NORIENT] /= sigma
@@ -249,13 +265,13 @@ class FoCUS:
 
                     wwc[:, NORIENT + 1] = tmp1.flatten() - tmp1.mean()
                     wws[:, NORIENT + 1] = tmp2.flatten() - tmp2.mean()
-                    #sigma = np.sqrt((wwc[:, NORIENT + 1] ** 2).mean())
+                    # sigma = np.sqrt((wwc[:, NORIENT + 1] ** 2).mean())
                     sigma = np.mean(w_smooth)
                     wwc[:, NORIENT + 1] /= sigma
                     wws[:, NORIENT + 1] /= sigma
 
                 w_smooth = w_smooth.flatten()
-            
+
         if self.use_1D:
             KERNELSZ = 5
 
@@ -723,19 +739,19 @@ class FoCUS:
     def ud_grade(self, im, j, axis=0):
         rim = im
         for k in range(j):
-            #rim = self.smooth(rim, axis=axis)
+            # rim = self.smooth(rim, axis=axis)
             rim = self.ud_grade_2(rim, axis=axis)
         return rim
 
     # --------------------------------------------------------
-    def ud_grade_2(self, im, axis=0):
+    def ud_grade_2(self, im, axis=0, cell_ids=None, nside=None):
 
         if self.use_2D:
             ishape = list(im.shape)
             if len(ishape) < axis + 2:
                 if not self.silent:
                     print("Use of 2D scat with data that has less than 2D")
-                return None
+                return None, None
 
             npix = im.shape[axis]
             npiy = im.shape[axis + 1]
@@ -760,29 +776,40 @@ class FoCUS:
 
             if axis == 0:
                 if len(ishape) == 2:
-                    return self.backend.bk_reshape(res, [npix // 2, npiy // 2])
+                    return self.backend.bk_reshape(res, [npix // 2, npiy // 2]), None
                 else:
-                    return self.backend.bk_reshape(
-                        res, [npix // 2, npiy // 2] + ishape[axis + 2 :]
+                    return (
+                        self.backend.bk_reshape(
+                            res, [npix // 2, npiy // 2] + ishape[axis + 2 :]
+                        ),
+                        None,
                     )
             else:
                 if len(ishape) == axis + 2:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [npix // 2, npiy // 2]
+                    return (
+                        self.backend.bk_reshape(
+                            res, ishape[0:axis] + [npix // 2, npiy // 2]
+                        ),
+                        None,
                     )
                 else:
-                    return self.backend.bk_reshape(
-                        res,
-                        ishape[0:axis] + [npix // 2, npiy // 2] + ishape[axis + 2 :],
+                    return (
+                        self.backend.bk_reshape(
+                            res,
+                            ishape[0:axis]
+                            + [npix // 2, npiy // 2]
+                            + ishape[axis + 2 :],
+                        ),
+                        None,
                     )
 
-            return self.backend.bk_reshape(res, [npix // 2, npiy // 2])
+            return self.backend.bk_reshape(res, [npix // 2, npiy // 2]), None
         elif self.use_1D:
             ishape = list(im.shape)
             if len(ishape) < axis + 1:
                 if not self.silent:
                     print("Use of 1D scat with data that has less than 1D")
-                return None
+                return None, None
 
             npix = im.shape[axis]
             odata = 1
@@ -805,23 +832,33 @@ class FoCUS:
 
             if axis == 0:
                 if len(ishape) == 1:
-                    return self.backend.bk_reshape(res, [npix // 2])
+                    return self.backend.bk_reshape(res, [npix // 2]), None
                 else:
-                    return self.backend.bk_reshape(
-                        res, [npix // 2] + ishape[axis + 1 :]
+                    return (
+                        self.backend.bk_reshape(res, [npix // 2] + ishape[axis + 1 :]),
+                        None,
                     )
             else:
                 if len(ishape) == axis + 1:
-                    return self.backend.bk_reshape(res, ishape[0:axis] + [npix // 2])
+                    return (
+                        self.backend.bk_reshape(res, ishape[0:axis] + [npix // 2]),
+                        None,
+                    )
                 else:
-                    return self.backend.bk_reshape(
-                        res, ishape[0:axis] + [npix // 2] + ishape[axis + 1 :]
+                    return (
+                        self.backend.bk_reshape(
+                            res, ishape[0:axis] + [npix // 2] + ishape[axis + 1 :]
+                        ),
+                        None,
                     )
 
-            return self.backend.bk_reshape(res, [npix // 2])
+            return self.backend.bk_reshape(res, [npix // 2]), None
 
         else:
             shape = list(im.shape)
+            if cell_ids is not None:
+                sim, new_cell_ids = self.backend.binned_mean(im, cell_ids)
+                return sim, new_cell_ids
 
             lout = int(np.sqrt(shape[axis] // 12))
             if im.__class__ == np.zeros([0]).__class__:
@@ -840,8 +877,11 @@ class FoCUS:
                 if len(shape) > axis:
                     oshape = oshape + shape[axis + 1 :]
 
-            return self.backend.bk_reduce_mean(
-                self.backend.bk_reshape(im, oshape), axis=axis + 1
+            return (
+                self.backend.bk_reduce_mean(
+                    self.backend.bk_reshape(im, oshape), axis=axis + 1
+                ),
+                None,
             )
 
     # --------------------------------------------------------
@@ -1794,14 +1834,14 @@ class FoCUS:
             if self.padding == "VALID":
                 l_mask = l_mask[
                     :,
-                    self.KERNELSZ // 2 : -self.KERNELSZ // 2+1,
-                    self.KERNELSZ // 2 : -self.KERNELSZ // 2+1,
+                    self.KERNELSZ // 2 : -self.KERNELSZ // 2 + 1,
+                    self.KERNELSZ // 2 : -self.KERNELSZ // 2 + 1,
                 ]
                 if shape[axis] != l_mask.shape[1]:
                     l_mask = l_mask[
                         :,
-                        self.KERNELSZ // 2 : -self.KERNELSZ // 2+1,
-                        self.KERNELSZ // 2 : -self.KERNELSZ // 2+1,
+                        self.KERNELSZ // 2 : -self.KERNELSZ // 2 + 1,
+                        self.KERNELSZ // 2 : -self.KERNELSZ // 2 + 1,
                     ]
 
             ichannel = 1
@@ -1868,10 +1908,10 @@ class FoCUS:
             l_mask = self.backend.bk_complex(l_mask, self.backend.bk_cast(0.0 * l_mask))
 
         if self.use_2D:
-            #if self.padding == "VALID":
+            # if self.padding == "VALID":
             mtmp = l_mask
             vtmp = l_x
-            #else:
+            # else:
             #    mtmp = l_mask[:,self.KERNELSZ // 2 : -self.KERNELSZ // 2,self.KERNELSZ // 2 : -self.KERNELSZ // 2,:]
             #    vtmp = l_x[:,self.KERNELSZ // 2 : -self.KERNELSZ // 2,self.KERNELSZ // 2 : -self.KERNELSZ // 2,:]
 
@@ -2125,7 +2165,7 @@ class FoCUS:
                 return self.backend.bk_reduce_sum(r)
 
     # ---------------------------------------------−---------
-    def convol(self, in_image, axis=0):
+    def convol(self, in_image, axis=0, cell_ids=None, nside=None):
 
         image = self.backend.bk_cast(in_image)
 
@@ -2290,6 +2330,61 @@ class FoCUS:
             return self.backend.bk_reshape(res, in_image.shape + [self.NORIENT])
 
         else:
+            ishape = list(image.shape)
+
+            if cell_ids is not None:
+                if cell_ids.shape[0] not in self.padding_conv:
+                    import healpix_convolution as hc
+                    from xdggs.healpix import HealpixInfo
+
+                    res = self.backend.bk_zeros(
+                        ishape + [self.NORIENT], dtype=self.backend.all_cbk_type
+                    )
+
+                    grid_info = HealpixInfo(
+                        level=int(np.log(nside) / np.log(2)), indexing_scheme="nested"
+                    )
+
+                    for k in range(self.NORIENT):
+                        kernelR, kernelI = hc.kernels.wavelet_kernel(
+                            cell_ids, grid_info=grid_info, orientation=k, is_torch=True
+                        )
+                        self.kernelR_conv[(cell_ids.shape[0], k)] = kernelR.to(
+                            self.backend.all_bk_type
+                        ).to(image.device)
+                        self.kernelI_conv[(cell_ids.shape[0], k)] = kernelI.to(
+                            self.backend.all_bk_type
+                        ).to(image.device)
+                        self.padding_conv[(cell_ids.shape[0], k)] = hc.pad(
+                            cell_ids,
+                            grid_info=grid_info,
+                            ring=5 // 2,  # wavelet kernel_size=5 is hard coded
+                            mode="mean",
+                            constant_value=0,
+                        )
+
+                for k in range(self.NORIENT):
+
+                    kernelR = self.kernelR_conv[(cell_ids.shape[0], k)]
+                    kernelI = self.kernelI_conv[(cell_ids.shape[0], k)]
+                    padding = self.padding_conv[(cell_ids.shape[0], k)]
+                    if len(ishape) == 2:
+                        for l in range(ishape[0]):
+                            padded_data = padding.apply(image[l], is_torch=True)
+                            res[l, :, k] = kernelR.matmul(
+                                padded_data
+                            ) + 1j * kernelI.matmul(padded_data)
+                    else:
+                        for l in range(ishape[0]):
+                            for k2 in range(ishape[2]):
+                                padded_data = padding.apply(
+                                    image[l, :, k2], is_torch=True
+                                )
+                                res[l, :, k2, k] = kernelR.matmul(
+                                    padded_data
+                                ) + 1j * kernelI.matmul(padded_data)
+                return res
+
             nside = int(np.sqrt(image.shape[axis] // 12))
 
             if self.Idx_Neighbours[nside] is None:
@@ -2306,7 +2401,6 @@ class FoCUS:
             l_ww_real = self.ww_Real[nside]
             l_ww_imag = self.ww_Imag[nside]
 
-            ishape = list(image.shape)
             odata = 1
             for k in range(axis + 1, len(ishape)):
                 odata = odata * ishape[k]
@@ -2460,7 +2554,7 @@ class FoCUS:
         return res
 
     # ---------------------------------------------−---------
-    def smooth(self, in_image, axis=0):
+    def smooth(self, in_image, axis=0, cell_ids=None, nside=None):
 
         image = self.backend.bk_cast(in_image)
 
@@ -2589,6 +2683,50 @@ class FoCUS:
             return self.backend.bk_reshape(res, in_image.shape)
 
         else:
+
+            ishape = list(image.shape)
+
+            if cell_ids is not None:
+                if cell_ids.shape[0] not in self.padding_smooth:
+                    import healpix_convolution as hc
+                    from xdggs.healpix import HealpixInfo
+
+                    grid_info = HealpixInfo(
+                        level=int(np.log(nside) / np.log(2)), indexing_scheme="nested"
+                    )
+
+                    kernel = hc.kernels.wavelet_smooth_kernel(
+                        cell_ids, grid_info=grid_info, is_torch=True
+                    )
+
+                    self.kernel_smooth[cell_ids.shape[0]] = kernel.to(
+                        self.backend.all_bk_type
+                    ).to(image.device)
+
+                    self.padding_smooth[cell_ids.shape[0]] = hc.pad(
+                        cell_ids,
+                        grid_info=grid_info,
+                        ring=5 // 2,  # wavelet kernel_size=5 is hard coded
+                        mode="mean",
+                        constant_value=0,
+                    )
+
+                kernel = self.kernel_smooth[cell_ids.shape[0]]
+                padding = self.padding_smooth[cell_ids.shape[0]]
+
+                res = self.backend.bk_zeros(ishape, dtype=self.backend.all_cbk_type)
+
+                if len(ishape) == 2:
+                    for l in range(ishape[0]):
+                        padded_data = padding.apply(image[l], is_torch=True)
+                        res[l] = kernel.matmul(padded_data)
+                else:
+                    for l in range(ishape[0]):
+                        for k2 in range(ishape[2]):
+                            padded_data = padding.apply(image[l, :, k2], is_torch=True)
+                            res[l, :, k2] = kernel.matmul(padded_data)
+                return res
+
             nside = int(np.sqrt(image.shape[axis] // 12))
 
             if self.Idx_Neighbours[nside] is None:
@@ -2604,7 +2742,6 @@ class FoCUS:
                 self.w_smooth[nside] = ws
 
             l_w_smooth = self.w_smooth[nside]
-            ishape = list(image.shape)
 
             odata = 1
             for k in range(axis + 1, len(ishape)):
@@ -2707,9 +2844,11 @@ class FoCUS:
     # ---------------------------------------------−---------
     def get_ww(self, nside=1):
         if self.use_2D:
-            
-            return (self.ww_RealT[1].reshape(self.KERNELSZ*self.KERNELSZ,self.NORIENT),
-                    self.ww_ImagT[1].reshape(self.KERNELSZ*self.KERNELSZ,self.NORIENT))
+
+            return (
+                self.ww_RealT[1].reshape(self.KERNELSZ * self.KERNELSZ, self.NORIENT),
+                self.ww_ImagT[1].reshape(self.KERNELSZ * self.KERNELSZ, self.NORIENT),
+            )
         else:
             return (self.ww_Real[nside], self.ww_Imag[nside])
 

foscat/Synthesis.py

@@ -240,9 +240,9 @@ class Synthesis:
         grd_mask = self.grd_mask
 
         if grd_mask is not None:
-            g_tot = grd_mask * self.to_numpy(g_tot)
+            g_tot = self.operation.backend.to_numpy(g_tot * grd_mask)
         else:
-            g_tot = self.to_numpy(g_tot)
+            g_tot = self.operation.backend.to_numpy(g_tot)
 
         g_tot[np.isnan(g_tot)] = 0.0
 
@@ -426,7 +426,7 @@ class Synthesis:
                 factr=factr,
                 maxiter=maxitt,
             )
-            print('Final Loss ',loss)
+            print("Final Loss ", loss)
             # update bias input data
             if iteration < NUM_STEP_BIAS - 1:
                 # if self.mpi_rank==0: