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

foscat/BkBase.py

@@ -1,14 +1,15 @@
 import numpy as np
 
+
 class BackendBase:
-    
+
     def __init__(self, name, mpi_rank=0, all_type="float64", gpupos=0, silent=False):
-        
-        self.BACKEND=name
-        self.mpi_rank=mpi_rank
-        self.all_type=all_type
-        self.gpupos=gpupos
-        self.silent=silent
+
+        self.BACKEND = name
+        self.mpi_rank = mpi_rank
+        self.all_type = all_type
+        self.gpupos = gpupos
+        self.silent = silent
         # ---------------------------------------------−---------
         # table use to compute the iso orientation rotation
         self._iso_orient = {}
@@ -101,6 +102,7 @@ class BackendBase:
         return self.bk_reshape(
             self.backend.matmul(self.bk_reshape(x, oshape), lmat), oshape2
         )
+
     def calc_iso_orient(self, norient):
         tmp = np.zeros([norient * norient, norient])
         for i in range(norient):
@@ -188,7 +190,9 @@ class BackendBase:
                     tmp[:, :, k, l_orient] = np.cos(x * k) * np.cos((x.T) * l_orient)
 
             self._fft_2_orient[(norient, nharm, imaginary)] = self.bk_cast(
-                self.bk_constant(tmp.reshape(norient * norient, (1 + nharm) * (1 + nharm)))
+                self.bk_constant(
+                    tmp.reshape(norient * norient, (1 + nharm) * (1 + nharm))
+                )
             )
             self._fft_2_orient_C[(norient, nharm, imaginary)] = self.bk_complex(
                 self._fft_2_orient[(norient, nharm, imaginary)],
@@ -343,7 +347,7 @@ class BackendBase:
                 self._fft_3_orient[(norient, nharm, imaginary)],
                 0 * self._fft_3_orient[(norient, nharm, imaginary)],
             )
-            
+
     # ---------------------------------------------−---------
     # --             BACKEND DEFINITION                    --
     # ---------------------------------------------−---------
@@ -355,7 +359,7 @@ class BackendBase:
 
     def bk_sparse_dense_matmul(self, smat, mat):
         raise NotImplementedError("This is an abstract class.")
-    
+
     def conv2d(self, x, w, strides=[1, 1, 1, 1], padding="SAME"):
         raise NotImplementedError("This is an abstract class.")
 
@@ -379,9 +383,6 @@ class BackendBase:
 
     def bk_flattenR(self, x):
         raise NotImplementedError("This is an abstract class.")
-               
-    def bk_flatten(self, x):
-        raise NotImplementedError("This is an abstract class.")
 
     def bk_flatten(self, x):
         raise NotImplementedError("This is an abstract class.")
@@ -442,7 +443,7 @@ class BackendBase:
 
     def bk_tensor(self, data):
         raise NotImplementedError("This is an abstract class.")
-        
+
     def bk_shape_tensor(self, shape):
         raise NotImplementedError("This is an abstract class.")
 
@@ -460,7 +461,7 @@ class BackendBase:
 
     def bk_tanh(self, data):
         raise NotImplementedError("This is an abstract class.")
-        
+
     def bk_max(self, data):
         raise NotImplementedError("This is an abstract class.")
 
@@ -499,11 +500,11 @@ class BackendBase:
 
     def bk_fft(self, data):
         raise NotImplementedError("This is an abstract class.")
-            
-    def bk_fftn(self, data,dim=None):
+
+    def bk_fftn(self, data, dim=None):
         raise NotImplementedError("This is an abstract class.")
 
-    def bk_ifftn(self, data,dim=None,norm=None):
+    def bk_ifftn(self, data, dim=None, norm=None):
         raise NotImplementedError("This is an abstract class.")
 
     def bk_rfft(self, data):
@@ -524,23 +525,32 @@ class BackendBase:
     def bk_relu(self, x):
         raise NotImplementedError("This is an abstract class.")
 
-    def bk_clip_by_value(self, x,xmin,xmax):
+    def bk_clip_by_value(self, x, xmin, xmax):
         raise NotImplementedError("This is an abstract class.")
 
     def bk_cast(self, x):
         raise NotImplementedError("This is an abstract class.")
-            
-    def bk_variable(self,x):
+
+    def bk_variable(self, x):
+        raise NotImplementedError("This is an abstract class.")
+
+    def bk_assign(self, x, y):
+        raise NotImplementedError("This is an abstract class.")
+
+    def bk_constant(self, x):
+        raise NotImplementedError("This is an abstract class.")
+
+    def bk_cos(self, x):
         raise NotImplementedError("This is an abstract class.")
-        
-    def bk_assign(self,x,y):
+
+    def bk_sin(self, x):
         raise NotImplementedError("This is an abstract class.")
-            
-    def bk_constant(self,x):
+
+    def bk_arctan2(self, c, s):
         raise NotImplementedError("This is an abstract class.")
-            
-    def bk_empty(self,x):
+
+    def bk_empty(self, list):
         raise NotImplementedError("This is an abstract class.")
-        
-    def to_numpy(self,x):
+
+    def to_numpy(self, x):
         raise NotImplementedError("This is an abstract class.")

foscat/BkNumpy.py

@@ -1,11 +1,13 @@
-import foscat.BkBase as BackendBase
 import numpy as np
 
+import foscat.BkBase as BackendBase
+
+
 class BkNumpy(BackendBase.BackendBase):
-    
+
     def __init__(self, *args, **kwargs):
         # Impose que use_2D=True pour la classe scat
-        super().__init__(name='tensorflow', *args, **kwargs)
+        super().__init__(name="tensorflow", *args, **kwargs)
 
         # ===========================================================================
         # INIT
@@ -31,13 +33,14 @@ class BkNumpy(BackendBase.BackendBase):
                 self.all_bk_type = self.backend.float64
                 self.all_cbk_type = self.backend.complex128
             else:
-                print("ERROR INIT FOCUS ", all_type, " should be float32 or float64")
+                print(
+                    "ERROR INIT FOCUS ", self.all_type, " should be float32 or float64"
+                )
                 return None
-                
+
         # ===========================================================================
         # INIT
 
-        gpus = []
         gpuname = "CPU:0"
         self.gpulist = {}
         self.gpulist[0] = gpuname
@@ -49,27 +52,25 @@ class BkNumpy(BackendBase.BackendBase):
     def bk_SparseTensor(self, indice, w, dense_shape=[]):
         return self.scipy.sparse.coo_matrix(
             (w, (indice[:, 0], indice[:, 1])), shape=dense_shape
-            )
+        )
 
     def bk_stack(self, list, axis=0):
         return self.backend.stack(list, axis=axis)
 
     def bk_sparse_dense_matmul(self, smat, mat):
         return smat.dot(mat)
-    
+
     def conv2d(self, x, w, strides=[1, 1, 1, 1], padding="SAME"):
-        res = np.zeros(
-                [x.shape[0], x.shape[1], x.shape[2], w.shape[3]], dtype=x.dtype
-            )
+        res = np.zeros([x.shape[0], x.shape[1], x.shape[2], w.shape[3]], dtype=x.dtype)
         for k in range(w.shape[2]):
             for l_orient in range(w.shape[3]):
                 for j in range(res.shape[0]):
                     tmp = self.scipy.signal.convolve2d(
-                            x[j, :, :, k],
-                            w[:, :, k, l_orient],
-                            mode="same",
-                            boundary="symm",
-                        )
+                        x[j, :, :, k],
+                        w[:, :, k, l_orient],
+                        mode="same",
+                        boundary="symm",
+                    )
                     res[j, :, :, l_orient] += tmp
                     del tmp
         return res
@@ -79,8 +80,8 @@ class BkNumpy(BackendBase.BackendBase):
         for k in range(w.shape[2]):
             for j in range(res.shape[0]):
                 tmp = self.scipy.signal.convolve1d(
-                        x[j, :, k], w[:, k], mode="same", boundary="symm"
-                 )
+                    x[j, :, k], w[:, k], mode="same", boundary="symm"
+                )
                 res[j, :, :] += tmp
                 del tmp
         return res
@@ -114,10 +115,6 @@ class BkNumpy(BackendBase.BackendBase):
             return np.concatenate([x.real.flatten(), x.imag.flatten()], 0)
         else:
             return x.flatten()
-                
-    
-    def bk_flatten(self, x):
-        return x.flatten()
 
     def bk_flatten(self, x):
         return x.flatten()
@@ -144,10 +141,6 @@ class BkNumpy(BackendBase.BackendBase):
     # ---------------------------------------------−---------
     # return a tensor size
 
-    def bk_size(self, data):
-        return data.size
-
-
     def bk_reduce_mean(self, data, axis=None):
 
         if axis is None:
@@ -223,7 +216,7 @@ class BkNumpy(BackendBase.BackendBase):
 
     def bk_tensor(self, data):
         return data
-        
+
     def bk_shape_tensor(self, shape):
         return np.zeros(shape)
 
@@ -284,19 +277,27 @@ class BkNumpy(BackendBase.BackendBase):
     def bk_zeros(self, shape, dtype=None):
         return np.zeros(shape, dtype=dtype)
 
-    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 np.reverse(data, axis=axis)
 
     def bk_fft(self, data):
         return self.backend.fft.fft(data)
-            
-    def bk_fftn(self, data,dim=None):
+
+    def bk_fftn(self, data, dim=None):
         return self.backend.fft.fftn(data)
 
-    def bk_ifftn(self, data,dim=None,norm=None):
+    def bk_ifftn(self, data, dim=None, norm=None):
         return self.backend.fft.ifftn(data)
 
     def bk_rfft(self, data):
@@ -318,8 +319,8 @@ class BkNumpy(BackendBase.BackendBase):
     def bk_relu(self, x):
         return (x > 0) * x
 
-    def bk_clip_by_value(self, x,xmin,xmax):
-        return self.backend.clip(x,xmin,xmax)
+    def bk_clip_by_value(self, x, xmin, xmax):
+        return self.backend.clip(x, xmin, xmax)
 
     def bk_cast(self, x):
         if isinstance(x, np.float64):
@@ -355,20 +356,28 @@ class BkNumpy(BackendBase.BackendBase):
             out_type = self.all_bk_type
 
         return x.astype(out_type)
-            
-    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):
         return x

foscat/BkTensorflow.py

@@ -1,15 +1,16 @@
 import sys
 
-import foscat.BkBase as BackendBase
 import numpy as np
 import tensorflow as tf
-        
+
+import foscat.BkBase as BackendBase
+
+
 class BkTensorflow(BackendBase.BackendBase):
-    
+
     def __init__(self, *args, **kwargs):
         # Impose que use_2D=True pour la classe scat
-        super().__init__(name='tensorflow', *args, **kwargs)
-
+        super().__init__(name="tensorflow", *args, **kwargs)
 
         # ===========================================================================
         # INIT
@@ -30,12 +31,14 @@ class BkTensorflow(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"
+            )
+
         if self.mpi_rank == 0:
             if not self.silent:
                 print(
@@ -59,12 +62,10 @@ class BkTensorflow(BackendBase.BackendBase):
                 # Currently, memory growth needs to be the same across GPUs
                 for gpu in gpus:
                     self.backend.config.experimental.set_memory_growth(gpu, True)
-                logical_gpus = (
-                     self.backend.config.experimental.list_logical_devices("GPU")
-                )
-                print(
-                    len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs"
+                logical_gpus = self.backend.config.experimental.list_logical_devices(
+                    "GPU"
                 )
+                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
@@ -92,10 +93,10 @@ class BkTensorflow(BackendBase.BackendBase):
         return self.backend.sparse.sparse_dense_matmul(smat, mat)
 
     # for tensorflow wrapping only
-    def periodic_pad(self,x, pad_height, pad_width):
+    def periodic_pad(self, x, pad_height, pad_width):
         """
         Applies periodic ('wrap') padding to a 4D TensorFlow tensor (N, H, W, C).
-        
+
         Args:
         x (tf.Tensor): Input tensor with shape (batch_size, height, width, channels).
             pad_height (tuple): Tuple (top, bottom) defining the vertical padding size.
@@ -104,23 +105,27 @@ class BkTensorflow(BackendBase.BackendBase):
         Returns:
             tf.Tensor: Tensor with periodic padding applied.
         """
-        #Vertical padding: take slices from bottom and top to wrap around
-        top_pad = x[:, -pad_height:, :, :]    # Top padding from the bottom rows
+        # Vertical padding: take slices from bottom and top to wrap around
+        top_pad = x[:, -pad_height:, :, :]  # Top padding from the bottom rows
         bottom_pad = x[:, :pad_height, :, :]  # Bottom padding from the top rows
-        x_padded = self.backend.concat([top_pad, x, bottom_pad], axis=1)  # Concatenate vertically
+        x_padded = self.backend.concat(
+            [top_pad, x, bottom_pad], axis=1
+        )  # Concatenate vertically
+
+        # Horizontal padding: take slices from right and left to wrap around
+        left_pad = x_padded[:, :, -pad_width:, :]  # Left padding from right columns
+        right_pad = x_padded[:, :, :pad_width, :]  # Right padding from left columns
 
-        #Horizontal padding: take slices from right and left to wrap around
-        left_pad = x_padded[:, :, -pad_width:, :]   # Left padding from right columns
-        right_pad = x_padded[:, :, :pad_width, :]   # Right padding from left columns
-        
-        x_padded = self.backend.concat([left_pad, x_padded, right_pad], axis=2)  # Concatenate horizontally
+        x_padded = self.backend.concat(
+            [left_pad, x_padded, right_pad], axis=2
+        )  # Concatenate horizontally
 
         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) 
+        x_padded = self.periodic_pad(x, kx // 2, ky // 2)
         return self.backend.nn.conv2d(x_padded, w, strides=strides, padding="VALID")
 
     def conv1d(self, x, w, strides=[1, 1, 1], padding="SAME"):
@@ -133,11 +138,9 @@ class BkTensorflow(BackendBase.BackendBase):
     def bk_threshold(self, x, threshold, greater=True):
 
         return self.backend.cast(x > threshold, x.dtype) * x
-        
 
     def bk_maximum(self, x1, x2):
         return self.backend.maximum(x1, x2)
-        
 
     def bk_device(self, device_name):
         return self.backend.device(device_name)
@@ -153,22 +156,18 @@ class BkTensorflow(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)))]
-                )
+                self.bk_imag(x), [np.prod(np.array(list(x.shape)))]
+            )
             return self.bk_concat([rr, ii], axis=0)
         else:
             return self.backend.reshape(x, [np.prod(np.array(list(x.shape)))])
-                
-    
+
     def bk_flatten(self, x):
         return self.backend.flatten(x)
 
-    def bk_size(self, x):
-        return self.backend.size(x)
-
     def bk_resize_image(self, x, shape):
         return self.bk_cast(self.backend.image.resize(x, shape, method="bilinear"))
 
@@ -282,7 +281,7 @@ class BkTensorflow(BackendBase.BackendBase):
 
     def bk_tensor(self, data):
         return self.backend.constant(data)
-        
+
     def bk_shape_tensor(self, shape):
         return self.backend.tensor(shape=shape)
 
@@ -320,10 +319,10 @@ class BkTensorflow(BackendBase.BackendBase):
         return self.backend.repeat(data, nn, axis=axis)
 
     def bk_tile(self, data, nn, axis=0):
-        order=[1 for k in data.shape]
-        order[axis]=nn
+        order = [1 for k in data.shape]
+        order[axis] = nn
         return self.backend.tile(data, self.backend.constant(order, tf.int32))
-        
+
     def bk_roll(self, data, nn, axis=0):
         return self.backend.roll(data, nn, axis=axis)
 
@@ -338,12 +337,8 @@ class BkTensorflow(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)
@@ -351,10 +346,10 @@ class BkTensorflow(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:
@@ -363,38 +358,42 @@ class BkTensorflow(BackendBase.BackendBase):
     def bk_zeros(self, shape, dtype=None):
         return self.backend.zeros(shape, dtype=dtype)
 
-    def bk_gather(self, data, idx):
-        return self.backend.gather(data, idx)
+    def bk_gather(self, data, idx, axis=0):
+        return self.backend.gather(data, idx, axis=axis)
 
     def bk_reverse(self, data, axis=0):
         return self.backend.reverse(data, axis=[axis])
 
     def bk_fft(self, data):
         return self.backend.signal.fft(data)
-    
 
-    def bk_fftn(self, data,dim=None):
-        #Equivalent of torch.fft.fftn(x, dim=dims) in TensorFlow
-        if len(dim)==2:
-            return self.backend.signal.fft2d(self.bk_complex(data, 0*data))
+    def bk_fftn(self, data, dim=None):
+        # Equivalent of torch.fft.fftn(x, dim=dims) in TensorFlow
+        if len(dim) == 2:
+            return self.backend.signal.fft2d(self.bk_complex(data, 0 * data))
         else:
-            return self.backend.signal.fft1d(self.bk_complex(data, 0*data))
+            return self.backend.signal.fft1d(self.bk_complex(data, 0 * data))
 
-    def bk_ifftn(self, data,dim=None,norm=None):
+    def bk_ifftn(self, data, dim=None, norm=None):
         if norm is not None:
-            if len(dim)==2:
-                normalization=self.backend.sqrt(self.backend.cast(data.shape[dim[0]]*data.shape[dim[1]], self.all_cbk_type))
-                return self.backend.signal.ifft2d(data)*normalization
-                    
+            if len(dim) == 2:
+                normalization = self.backend.sqrt(
+                    self.backend.cast(
+                        data.shape[dim[0]] * data.shape[dim[1]], self.all_cbk_type
+                    )
+                )
+                return self.backend.signal.ifft2d(data) * normalization
+
             else:
-                normalization=self.backend.sqrt(self.backend.cast(data.shape[dim[0]], self.all_cbk_type))
-                return self.backend.signal.ifft1d(data)*normalization
+                normalization = self.backend.sqrt(
+                    self.backend.cast(data.shape[dim[0]], self.all_cbk_type)
+                )
+                return self.backend.signal.ifft1d(data) * normalization
         else:
-            if len(dim)==2:
+            if len(dim) == 2:
                 return self.backend.signal.ifft2d(data)
             else:
                 return self.backend.signal.ifft1d(data)
-        
 
     def bk_rfft(self, data):
         return self.backend.signal.rfft(data)
@@ -420,10 +419,10 @@ class BkTensorflow(BackendBase.BackendBase):
         else:
             return self.backend.nn.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)
-        return self.backend.clip_by_value(x,xmin,xmax)
+            x = np.clip(x, xmin, xmax)
+        return self.backend.clip_by_value(x, xmin, xmax)
 
     def bk_cast(self, x):
         if isinstance(x, np.float64):
@@ -459,21 +458,30 @@ class BkTensorflow(BackendBase.BackendBase):
             out_type = self.all_bk_type
 
         return self.backend.cast(x, out_type)
-            
-    def bk_variable(self,x):
+
+    def bk_variable(self, x):
         return self.backend.Variable(x)
-        
-    def bk_assign(self,x,y):
-        x.assign(y)
-            
-    def bk_constant(self,x):
+
+    def bk_assign(self, x, y):
+        return x.assign(y)
+
+    def bk_constant(self, x):
         return self.backend.constant(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.constant(list)
-        
-    def to_numpy(self,x):
+
+    def to_numpy(self, x):
         if isinstance(x, np.ndarray):
             return x
-        
+
         return x.numpy()