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

foscat/BkTorch.py

@@ -149,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)
@@ -246,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:
@@ -391,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)
@@ -429,6 +425,7 @@ class BkTorch(BackendBase.BackendBase):
                 xr = self.backend.concat(
                     [self.bk_real(data[k]) for k in range(ndata)], axis=axis
                 )
+                    
                 xi = self.backend.concat(
                     [self.bk_imag(data[k]) for k in range(ndata)], axis=axis
                 )
@@ -440,7 +437,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]
@@ -448,7 +447,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