Auto code format

Author: texlucas

src/pyGroupedTransforms/GroupedTransform.py

@@ -1,12 +1,11 @@
 import numpy as np
+import pykeops
 import torch
+from pykeops.torch import LazyTensor
 
 from pyGroupedTransforms import *
-from .NFFTtools import index_set_without_zeros
-
-import pykeops
-from pykeops.torch import LazyTensor
 
+from .NFFTtools import index_set_without_zeros
 
 # All code that is linked to NFMTtools or to system = "mixed" is not tested yet....
 
@@ -272,13 +271,12 @@ def __init__(
                         bandwidths=s.bandwidths, X=np.copy(X[:, u], order="C")
                     )
         elif algorithm == "keops":
-            self.matrix = np.empty((0,0), dtype=object)
-            
+            self.matrix = np.empty((0, 0), dtype=object)
+
             self.transforms = [
-                DeferredLinearOperator()
-                for _ in range(len(self.settings))
+                DeferredLinearOperator() for _ in range(len(self.settings))
             ]
-            
+
             D = X.shape[1]
 
             freq_list = []
@@ -287,9 +285,9 @@ def __init__(
 
                 if len(s.bandwidths) == 0:
                     full = np.zeros((1, D), dtype=np.int32)
-                    
+
                 elif self.system == "cos":
-                    
+
                     local = np.atleast_2d(
                         s.mode.index_set_without_zeros(
                             np.array(s.bandwidths, dtype=np.int32)
@@ -300,7 +298,7 @@ def __init__(
 
                     for i, dim in enumerate(s.u):
                         full[:, dim] = local[:, i]
-                        
+
                 else:
                     local = index_set_without_zeros(
                         np.array(s.bandwidths, dtype=np.int32)
@@ -314,13 +312,9 @@ def __init__(
                 freq_list.append(full)
 
             freq = np.vstack(freq_list)
-            
+
             X_torch = torch.tensor(X, dtype=torch.float64, device="cuda")
-            I_torch = torch.tensor(
-                freq,
-                dtype=torch.float64,
-                device="cuda"
-            )
+            I_torch = torch.tensor(freq, dtype=torch.float64, device="cuda")
 
             def trafo(fhat):
                 if self.system == "cos" or self.system == "cheb":
@@ -329,8 +323,8 @@ def trafo(fhat):
 
                     kernel = 1.0
                     for i in range(D):
-                        Xi = LazyTensor(X_torch[:, None, i:i+1].contiguous())
-                        Ki = LazyTensor(I_torch[None, :, i:i+1].contiguous())
+                        Xi = LazyTensor(X_torch[:, None, i : i + 1].contiguous())
+                        Ki = LazyTensor(I_torch[None, :, i : i + 1].contiguous())
                         kernel = kernel * (2 * torch.pi * Xi * Ki).cos()
 
                     fhat_torch = (
@@ -354,7 +348,9 @@ def trafo(fhat):
                     two_pi_phase = -2 * torch.pi * phase_fwd
                     kernel = two_pi_phase.cos() + 1j * two_pi_phase.sin()
 
-                    fhat_torch = torch.tensor(fhat, dtype=torch.complex128, device="cuda")
+                    fhat_torch = torch.tensor(
+                        fhat, dtype=torch.complex128, device="cuda"
+                    )
                     fhat_j = LazyTensor(fhat_torch[None, :, None].contiguous())
 
                     try:
@@ -370,16 +366,16 @@ def adjoint(f):
                 if self.system == "cos" or self.system == "cheb":
                     mult = np.sqrt(2.0) ** np.count_nonzero(freq, axis=1)
                     mult_torch = torch.tensor(mult, dtype=torch.float64, device="cuda")
-                    
+
                     f_torch = (
                         torch.as_tensor(f, dtype=torch.float64, device="cuda")
                         * mult_torch
                     ).contiguous()
 
                     kernel = 1.0
                     for i in range(D):
-                        Ki = LazyTensor(I_torch[:, None, i:i+1].contiguous())
-                        Xi = LazyTensor(X_torch[None, :, i:i+1].contiguous())
+                        Ki = LazyTensor(I_torch[:, None, i : i + 1].contiguous())
+                        Xi = LazyTensor(X_torch[None, :, i : i + 1].contiguous())
                         kernel = kernel * (2 * torch.pi * Xi * Ki).cos()
 
                     f_i = LazyTensor(f_torch[None, :, None].contiguous())
@@ -393,7 +389,9 @@ def adjoint(f):
                         print("Error in KeOps adjoint:", e)
                         return None
                 else:
-                    f_torch = torch.tensor(f, dtype=torch.complex128, device="cuda").contiguous()
+                    f_torch = torch.tensor(
+                        f, dtype=torch.complex128, device="cuda"
+                    ).contiguous()
 
                     X_i = LazyTensor(X_torch[None, :, :].contiguous())
                     K_j = LazyTensor(I_torch[:, None, :].contiguous())
@@ -415,10 +413,15 @@ def adjoint(f):
                         return None
 
             keops_dtype = np.float64 if self.system == "cos" else np.complex128
-            self.transforms = [DeferredLinearOperator(
-                dtype=keops_dtype, shape=(X.shape[0], len(freq)), mfunc=trafo, rmfunc=adjoint
-            )]
-    
+            self.transforms = [
+                DeferredLinearOperator(
+                    dtype=keops_dtype,
+                    shape=(X.shape[0], len(freq)),
+                    mfunc=trafo,
+                    rmfunc=adjoint,
+                )
+            ]
+
         else:
             self.transforms = []
             s1 = self.settings[0]
@@ -498,15 +501,15 @@ def adjoint_worker(i):
                         adjoint_worker(i)
 
                 return fhat
-            
+
             elif self.algorithm == "keops":
-                return GroupedCoefficients(self.settings, self.transforms[0].H @ other)  
-             
+                return GroupedCoefficients(self.settings, self.transforms[0].H @ other)
+
             elif self.algorithm == "direct":
                 return GroupedCoefficients(
                     self.settings, (self.matrix.conj()).T @ other
                 )
-                
+
         elif isinstance(other, GC):  # `f = F*fhat`    (fhat = other)
             if self.settings != other.settings:
                 raise ValueError(
@@ -534,10 +537,10 @@ def worker(i):
                         worker(i)
 
                 return sum(results)
-            
+
             elif self.algorithm == "keops":
                 return self.transforms[0] @ other.data
-            
+
             elif self.algorithm == "direct":
                 return self.matrix @ other.data
         else:

src/pyGroupedTransforms/NFFTtools.py

@@ -1,11 +1,10 @@
 import numpy as np
-
-from pyGroupedTransforms import *
-
-import torch
 import pykeops
+import torch
 from pykeops.torch import LazyTensor
 
+from pyGroupedTransforms import *
+
 
 def datalength(
     bandwidths: np.ndarray,
@@ -173,6 +172,7 @@ def adjoint(f):  # function adjoint(f::Vector{ComplexF64})::Vector{ComplexF64}
         dtype=np.complex128, shape=(M, N), mfunc=trafo, rmfunc=adjoint
     )
 
+
 def get_matrix(
     bandwidths, X
 ):  # get_matrix(bandwidths::Vector{Int}, X::Array{Float64})::Array{ComplexF64}