foscat 3.8.0__py3-none-any.whl → 3.9.0__py3-none-any.whl

foscat/BkTorch.py

@@ -1,15 +1,20 @@
 import sys
-import foscat.BkBase as BackendBase
+
 import numpy as np
 import torch
 
+import foscat.BkBase as BackendBase
+
+
 class BkTorch(BackendBase.BackendBase):
-    
+
     def __init__(self, *args, **kwargs):
         # Impose que use_2D=True pour la classe scat
-        super().__init__(name='torch', *args, **kwargs)
+        super().__init__(name="torch", *args, **kwargs)
         self.backend = torch
-        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.device = (
+            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        )
 
         self.float64 = self.backend.float64
         self.float32 = self.backend.float32
@@ -22,11 +27,13 @@ class BkTorch(BackendBase.BackendBase):
             "float32": (self.backend.float32, self.backend.complex64),
             "float64": (self.backend.float64, self.backend.complex128),
         }
-        
+
         if self.all_type in dtype_map:
             self.all_bk_type, self.all_cbk_type = dtype_map[self.all_type]
         else:
-            raise ValueError(f"ERROR INIT foscat: {all_type} should be float32 or float64")
+            raise ValueError(
+                f"ERROR INIT foscat: {self.all_type} should be float32 or float64"
+            )
 
         # ===========================================================================
         # INIT
@@ -50,14 +57,20 @@ class BkTorch(BackendBase.BackendBase):
             except RuntimeError as e:
                 # Memory growth must be set before GPUs have been initialized
                 print(e)
-                
-        self.torch_device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+        self.torch_device = (
+            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        )
 
     # ---------------------------------------------−---------
     # --             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)
@@ -65,21 +78,20 @@ 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
+
         lx = x.permute(0, 3, 1, 2)
         wx = w.permute(3, 2, 0, 1)  # de (5, 5, 1, 4) à (4, 1, 5, 5)
 
         # Calculer le padding symétrique
         kx, ky = w.shape[0], w.shape[1]
-            
+
         # Appliquer le padding
-        x_padded = F.pad(lx, (ky // 2, ky // 2, kx // 2, kx // 2), mode='circular')
+        x_padded = F.pad(lx, (ky // 2, ky // 2, kx // 2, kx // 2), mode="circular")
 
         # Appliquer la convolution
-        return F.conv2d(x_padded, wx, stride=1, padding=0).permute(0,2,3,1)
+        return F.conv2d(x_padded, wx, stride=1, padding=0).permute(0, 2, 3, 1)
 
     def conv1d(self, x, w, strides=[1, 1, 1], padding="SAME"):
         # to be written!!!
@@ -108,7 +120,7 @@ class BkTorch(BackendBase.BackendBase):
     def bk_flattenR(self, x):
         if self.bk_is_complex(x):
             rr = self.backend.reshape(
-                 self.bk_real(x), [np.prod(np.array(list(x.shape)))]
+                self.bk_real(x), [np.prod(np.array(list(x.shape)))]
             )
             ii = self.backend.reshape(
                 self.bk_imag(x), [np.prod(np.array(list(x.shape)))]
@@ -122,18 +134,18 @@ class BkTorch(BackendBase.BackendBase):
 
     def bk_resize_image(self, x, shape):
         tmp = self.backend.nn.functional.interpolate(
-                x.permute(0,3,1,2), size=shape, mode="bilinear", align_corners=False
-            )
-        return self.bk_cast(tmp.permute(0,2,3,1))
+            x.permute(0, 3, 1, 2), size=shape, mode="bilinear", align_corners=False
+        )
+        return self.bk_cast(tmp.permute(0, 2, 3, 1))
 
     def bk_L1(self, x):
         if x.dtype == self.all_cbk_type:
             xr = self.bk_real(x)
-            xi = self.bk_imag(x)
+            # xi = self.bk_imag(x)
 
             r = self.backend.sign(xr) * self.backend.sqrt(self.backend.sign(xr) * xr)
             # return r
-            i = self.backend.sign(xi) * self.backend.sqrt(self.backend.sign(xi) * xi)
+            # i = self.backend.sign(xi) * self.backend.sqrt(self.backend.sign(xi) * xi)
 
             return r
         else:
@@ -163,7 +175,6 @@ class BkTorch(BackendBase.BackendBase):
     def bk_size(self, data):
         return data.numel()
 
-
     def constant(self, data):
         return data
 
@@ -194,7 +205,7 @@ class BkTorch(BackendBase.BackendBase):
             r = self.backend.std(data)
         else:
             r = self.backend.std(data, axis)
-            
+
         if self.bk_is_complex(data):
             return self.bk_complex(r, 0 * r)
         else:
@@ -245,7 +256,7 @@ class BkTorch(BackendBase.BackendBase):
 
     def bk_tensor(self, data):
         return self.backend.constant(data).to(self.torch_device)
-        
+
     def bk_shape_tensor(self, shape):
         return self.backend.tensor(shape=shape).to(self.torch_device)
 
@@ -282,7 +293,7 @@ class BkTorch(BackendBase.BackendBase):
         return data.view(shape)
 
     def bk_repeat(self, data, nn, axis=0):
-        return self.backend.repeat(data, nn, axis=axis)
+        return self.backend.repeat_interleave(data, repeats=nn, dim=axis)
 
     def bk_tile(self, data, nn, axis=0):
 
@@ -304,12 +315,8 @@ class BkTorch(BackendBase.BackendBase):
         if axis is None:
             if data[0].dtype == self.all_cbk_type:
                 ndata = len(data)
-                xr = self.backend.concat(
-                        [self.bk_real(data[k]) for k in range(ndata)]
-                    )
-                xi = self.backend.concat(
-                        [self.bk_imag(data[k]) for k in range(ndata)]
-                    )
+                xr = self.backend.concat([self.bk_real(data[k]) for k in range(ndata)])
+                xi = self.backend.concat([self.bk_imag(data[k]) for k in range(ndata)])
                 return self.bk_complex(xr, xi)
             else:
                 return self.backend.concat(data)
@@ -317,11 +324,11 @@ class BkTorch(BackendBase.BackendBase):
             if data[0].dtype == self.all_cbk_type:
                 ndata = len(data)
                 xr = self.backend.concat(
-                        [self.bk_real(data[k]) for k in range(ndata)], axis=axis
-                    )
+                    [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
-                    )
+                    [self.bk_imag(data[k]) for k in range(ndata)], axis=axis
+                )
                 return self.bk_complex(xr, xi)
             else:
                 return self.backend.concat(data, axis=axis)
@@ -329,20 +336,28 @@ class BkTorch(BackendBase.BackendBase):
     def bk_zeros(self, shape, dtype=None):
         return self.backend.zeros(shape, dtype=dtype).to(self.torch_device)
 
-    def bk_gather(self, data, idx):
-        return data[idx]
+    def bk_gather(self, data, idx, axis=0):
+        if axis == 0:
+            return data[idx]
+        elif axis == 1:
+            return data[:, idx]
+        elif axis == 2:
+            return data[:, :, idx]
+        elif axis == 3:
+            return data[:, :, :, idx]
+        return data[:, :, :, :, idx]
 
     def bk_reverse(self, data, axis=0):
         return self.backend.flip(data, dims=[axis])
 
     def bk_fft(self, data):
         return self.backend.fft.fft(data)
-            
-    def bk_fftn(self, data,dim=None):
-        return self.backend.fft.fftn(data,dim=dim)
 
-    def bk_ifftn(self, data,dim=None,norm=None):
-        return self.backend.fft.ifftn(data,dim=dim,norm=norm)
+    def bk_fftn(self, data, dim=None):
+        return self.backend.fft.fftn(data, dim=dim)
+
+    def bk_ifftn(self, data, dim=None, norm=None):
+        return self.backend.fft.ifftn(data, dim=dim, norm=norm)
 
     def bk_rfft(self, data):
         return self.backend.fft.rfft(data)
@@ -366,12 +381,24 @@ class BkTorch(BackendBase.BackendBase):
     def bk_relu(self, x):
         return self.backend.relu(x)
 
-    def bk_clip_by_value(self, x,xmin,xmax):
+    def bk_clip_by_value(self, x, xmin, xmax):
         if isinstance(x, np.ndarray):
-            x = np.clip(x,xmin,xmax)
-        x = self.backend.tensor(x, dtype=self.backend.float32) if not isinstance(x, self.backend.Tensor) else x
-        xmin = self.backend.tensor(xmin, dtype=self.backend.float32) if not isinstance(xmin, self.backend.Tensor) else xmin
-        xmax = self.backend.tensor(xmax, dtype=self.backend.float32) if not isinstance(xmax, self.backend.Tensor) else xmax
+            x = np.clip(x, xmin, xmax)
+        x = (
+            self.backend.tensor(x, dtype=self.backend.float32)
+            if not isinstance(x, self.backend.Tensor)
+            else x
+        )
+        xmin = (
+            self.backend.tensor(xmin, dtype=self.backend.float32)
+            if not isinstance(xmin, self.backend.Tensor)
+            else xmin
+        )
+        xmax = (
+            self.backend.tensor(xmax, dtype=self.backend.float32)
+            if not isinstance(xmax, self.backend.Tensor)
+            else xmax
+        )
         return self.backend.clamp(x, min=xmin, max=xmax)
 
     def bk_cast(self, x):
@@ -416,22 +443,31 @@ class BkTorch(BackendBase.BackendBase):
             out_type = self.all_bk_type
 
         return x.type(out_type).to(self.torch_device)
-            
-    def bk_variable(self,x):
+
+    def bk_variable(self, x):
         return self.bk_cast(x)
-        
-    def bk_assign(self,x,y):
-        x=y
-            
-    def bk_constant(self,x):
-            
+
+    def bk_assign(self, x, y):
+        return y
+
+    def bk_constant(self, x):
+
         return self.bk_cast(x)
-        
-    def bk_empty(self,list):
+
+    def bk_cos(self, x):
+        return self.backend.cos(x)
+
+    def bk_sin(self, x):
+        return self.backend.sin(x)
+
+    def bk_arctan2(self, c, s):
+        return self.backend.arctan2(c, s)
+
+    def bk_empty(self, list):
         return self.backend.empty(list)
-        
-    def to_numpy(self,x):
+
+    def to_numpy(self, x):
         if isinstance(x, np.ndarray):
             return x
-            
+
         return x.cpu().numpy()

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.0"
+        self.__version__ = "3.9.0"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -45,7 +44,7 @@ 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
@@ -105,31 +104,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 +174,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 +211,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 +239,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 +253,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 +263,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,7 +737,7 @@ 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
 
@@ -1794,14 +1808,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 +1882,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,:]
 
@@ -2707,9 +2721,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: