Spectral differentiation fix

Author: Maslyaev

epde/preprocessing/deriv_calculators.py

@@ -176,18 +176,20 @@ def spectral_derivative_1d(self, func: np.ndarray, grid: np.ndarray, n=None, ste
         frequencies = np.fft.rfftfreq(spacing_vector.size, d=(spacing_vector[0][1] - spacing_vector[0][0]))
         frequencies_filtered = self.butterworth_filter(frequencies, n, steepness)
         return np.real(np.fft.irfft(1j * 2 * np.pi * frequencies_filtered * func_projection_filtered))
-
-    def spectral_derivative_nd(self, func: np.ndarray, grid: list, n=None, steepness=1,
-                               deriv_hist: list = []):
-        '''Многомерная спектральная производная,принимает на вход количество частот по каждой размерности и крутизну для фильтра Баттерворта, если они не указаны-фильтрация не производится'''
 
+    def spectral_derivative_nd(self, func : np.ndarray, grid : list, n = None, steepness = 1, 
+                               deriv_hist : list = []):
+        '''Многомерная спектральная производная,принимает на вход количество частот по каждой размерности и крутизну для фильтра Баттерворта, если они не указаны-фильтрация не производится'''
+
         if isinstance(n, int):
             n = np.full(shape=len(grid), fill_value=n)
-        if n is None:
+        if isinstance(n, type(None)):
             n = np.min(func.shape)
         all_dim_derivative = []
-        func_projection = np.fft.fftn(func, axes=[0,1])
-            
+        func_projection = np.fft.fftn(func)
+        print(func_projection.shape) #marker
+        inverter = lambda x: 1 if x == 0 else (x if x != 1 else 0)
+        
         for counter, i in enumerate(grid):
             spacing_vector = np.reshape(grid[counter], (1, grid[counter].size))
             frequencies = np.fft.fftfreq(spacing_vector.size, d=(spacing_vector[0][1] - spacing_vector[0][0]))
@@ -197,9 +199,9 @@ def spectral_derivative_nd(self, func: np.ndarray, grid: list, n=None, steepness
             derivative = np.real(np.fft.ifftn(derivative*1j * 2 * np.pi))
             all_dim_derivative.append((deriv_descr, derivative))
         return all_dim_derivative
-
-    def differentiate(self, field: np.ndarray, grid: list, max_order: Union[int, list],
-                      mixed: bool = False, n=None, steepness=1, deriv_hist: list = []) -> list:
+    
+    def differentiate(self, field : np.ndarray, grid : list, max_order : Union[int, list], 
+                      mixed : bool = False, n = None, steepness = 1, deriv_hist : list = []) -> list:
         if isinstance(max_order, int):
             max_order = [max_order,] * field.ndim
         else:
@@ -210,70 +212,74 @@ def differentiate(self, field: np.ndarray, grid: list, max_order: Union[int, lis
             def num_of_derivs_with_ord(ords):
                 temp = ords + field.ndim - 1
                 numerator = np.math.factorial(temp)
-                denominator = np.math.factorial(
-                    ords) * np.math.factorial(field.ndim - 1)
+                denominator = np.math.factorial(ords) * np.math.factorial(field.ndim - 1)
                 return int(numerator / denominator)
-            expeced_num_of_derivs = sum([num_of_derivs_with_ord(
-                cur_ord + 1) for cur_ord in range(max_order[0])])
+            expeced_num_of_derivs = sum([num_of_derivs_with_ord(cur_ord + 1) for cur_ord in range(max_order[0])])
         else:
             expeced_num_of_derivs = sum(max_order)
-        derivatives = {}
-
+        derivatives = {}     
+        
         if mixed:
             for axis in range(field.ndim):
-                higher_ord_derivs = self.spectral_derivative_nd(field, grid, n=n, steepness=steepness,
-                                                                deriv_hist=deriv_hist)
+                higher_ord_derivs = self.spectral_derivative_nd(field, grid, n = n, steepness = steepness, 
+                                                                deriv_hist = deriv_hist)
                 part_derivs = []
                 for history, field in higher_ord_derivs:
-                    part_derivs.extend(self.differentiate(
-                        field, grid, max_order - 1, deriv_hist=history))
+                    part_derivs.extend(self.differentiate(field, grid, max_order - 1, deriv_hist=history))
                 # self.differentiate(field, grid, max_order)
                 for key, deriv in part_derivs:
                     if key in derivatives.keys():
-                        assert np.all(np.isclose(
-                            deriv, derivatives[key])), 'Shuffle in differentiation orders shall not affect the values.'
+                        assert np.all(np.isclose(deriv, derivatives[key])), 'Shuffle in differentiation orders shall not affect the values.'
                     derivatives[key] = deriv
         else:
             for axis in range(field.ndim):
-                axis_derivs = self.spectral_derivative_high_ord(field, grid, axis=axis, max_order=max_order[axis],
-                                                                n=n, steepness=steepness)
+                axis_derivs = self.spectral_derivative_high_ord(field, grid, axis = axis, max_order = max_order[axis],
+                                                                n = n, steepness = steepness)
                 for key, deriv in axis_derivs:
                     derivatives[key] = deriv
-
-        # print(f'derivatives orders are {[deriv[0] for deriv in derivatives]}')
+                
+        print(f'derivatives orders are {[deriv[0] for deriv in derivatives]}')
         if len(derivatives) != expeced_num_of_derivs:
-            raise Exception(
-                f'Expected number of derivatives {expeced_num_of_derivs} does not match obtained {len(derivatives)}')
+            raise Exception(f'Expected number of derivatives {expeced_num_of_derivs} does not match obtained {len(derivatives)}')
         return derivatives
-
-    def __call__(self, data: np.ndarray, grid: list, max_order: Union[int, list],
-                 mixed: bool = False, n=None, steepness=1) -> np.ndarray:
-        def make_unsparse_sparse(*grids):  # TODO^ find more elegant solution
+    
+    def __call__(self, data : np.ndarray, grid : list, max_order : Union[int, list], 
+                 mixed : bool = False, n = None, steepness = 1) -> np.ndarray:
+        def make_unsparse_sparse(*grids): # TODO^ find more e;egant solution
             unique_vals = [np.unique(grid) for grid in grids]
-            return np.meshgrid(*unique_vals, sparse=True, indexing='ij')
-
+            return np.meshgrid(*unique_vals, sparse = True, indexing = 'ij')
+        
         if len(grid) > 1 and grid[0].shape == grid[1].shape:
             grid = make_unsparse_sparse(*grid)
-        if isinstance(n, int) or n is None:
-            n = np.full(shape=len(grid), fill_value=n)
-
+        if isinstance(n, int) or isinstance(n, type(None)):
+            n = np.full(shape=len(grid), fill_value=n)        
+        
         derivatives = self.differentiate(data, grid, max_order, mixed, n, steepness).values()
         derivatives = np.vstack([der.reshape(-1) for der in derivatives]).T
-
+        
         return derivatives
-
-    def spectral_derivative_high_ord(self, func: np.ndarray, grid: list, axis: int = 0,
-                                     max_order: int = 1, n: np.ndarray = None, steepness=1) -> list:
+    
+    def spectral_derivative_high_ord(self, func : np.ndarray, grid : list, axis : int = 0, 
+                                     max_order : int = 1, n = None, steepness = 1) -> list:
         derivs = []
         cur_deriv = func
+        func_projection = np.fft.fftn(func)
+        spacing_vector = np.reshape(grid[axis], (1, grid[axis].size))
+        frequencies = np.fft.fftfreq(spacing_vector.size, d=(spacing_vector[0][1] - spacing_vector[0][0]))
+        print(func_projection.shape) #marker
         # inverter = lambda x: 1 if x == 0 else (x if x != 1 else 0)
-
-        for deriv_idx in range(max_order):
-            deriv_descr = tuple([axis,] * (deriv_idx + 1))
-            cur_deriv = np.apply_along_axis(self.spectral_derivative_1d, axis, cur_deriv, grid[axis],
-                                            n[axis], steepness)
+        if isinstance(n, int):
+            n = np.full(shape=len(grid), fill_value=n)
+        if isinstance(n, type(None)):
+            n = np.min(func.shape)
+        frequencies_filtered = self.butterworth_filter(frequencies, n[axis], steepness)
+        for deriv_idx in range(1,max_order+1):
+            deriv_descr = tuple([axis,] * (deriv_idx)) # inverter(axis) V
+            derivative = np.apply_along_axis(np.multiply, axis, func_projection,
+                                             np.power(frequencies_filtered*1j * 2 * np.pi, deriv_idx,dtype = complex))
+            cur_deriv= np.real(np.fft.ifftn(derivative))
             derivs.append((deriv_descr, cur_deriv))
-
+            
         return derivs
 
 class TotalVariation(AbstractDeriv):