feat: add keops algorithm

Author: texlucas

src/pyGroupedTransforms/GroupedTransform.py

@@ -183,7 +183,7 @@ def __init__(
         system,
         X,
         settings=[],
-        fastmult=True,
+        algorithm="nfft",
         parallel=True,
         basis_vect=[],
         N=[],
@@ -227,9 +227,9 @@ def __init__(
                     )
 
         if system in {"chui1", "chui2", "chui3", "chui4"}:
-            fastmult = True
+            algorithm = "nfft"
 
-        self.fastmult = fastmult
+        self.algorithm = algorithm
         self.basis_vect = basis_vect
         self.system = system
         self.X = X
@@ -240,7 +240,7 @@ def __init__(
         else:
             self.settings = settings
 
-        if fastmult:
+        if algorithm == "nfft":
             self.matrix = np.empty((0, 0), dtype=object)
             self.transforms = [
                 DeferredLinearOperator() for i in range(len(self.settings))
@@ -265,6 +265,25 @@ def __init__(
                     self.transforms[idx] = s.mode.get_transform(
                         bandwidths=s.bandwidths, X=np.copy(X[:, u], order="C")
                     )
+        elif algorithm == "keops":
+                self.matrix = np.empty((0,0), dtype=object)
+
+                self.transforms = [
+                    DeferredLinearOperator()
+                    for _ in range(len(self.settings))
+                ]
+
+                for idx, s in enumerate(self.settings):
+
+                    if len(s.bandwidths) == 0:
+                        u = (0,)
+                    else:
+                        u = s.u
+
+                    self.transforms[idx] = s.mode.get_transform_keops(
+                        bandwidths=s.bandwidths,
+                        X=np.copy(X[:,u], order="C")
+                    )
         else:
             self.transforms = []
             s1 = self.settings[0]
@@ -321,9 +340,8 @@ def __mul__(self, other):
         If other (= f) is an numpy.ndarray, this function
         overloads the * notation in order to achieve the adjoint transform `f = F*f`.
         """
-
         if isinstance(other, np.ndarray):  # `f = F*f`    (f = other)
-            if self.fastmult:
+            if self.algorithm == "nfft":
                 fhat = GroupedCoefficients(self.settings)
 
                 def adjoint_worker(i):
@@ -345,17 +363,26 @@ def adjoint_worker(i):
                         adjoint_worker(i)
 
                 return fhat
-            else:
+            elif self.algorithm == "keops":
+                fhat = GroupedCoefficients(self.settings)       
+                
+                for i in range(len(self.transforms)):
+                    adjoint_result = self.transforms[i].H @ other
+                    fhat[self.settings[i].u] = adjoint_result
+
+                return fhat
+            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(
                     "The GroupedTransform and the GroupedCoefficients have different settings"
                 )
 
-            if self.fastmult:
+            if self.algorithm == "nfft":
                 results = []
 
                 def worker(i):
@@ -376,7 +403,16 @@ def worker(i):
                         worker(i)
 
                 return sum(results)
-            else:
+            elif self.algorithm == "keops":
+                results = []
+        
+                for i in range(len(self.transforms)):
+                    u = self.settings[i].u
+                    result = self.transforms[i] @ other[u]
+                    results.append(result)
+
+                return sum(results)
+            elif self.algorithm == "direct":
                 return self.matrix @ other.data
         else:
             raise ValueError("Wrong input data type")
@@ -412,7 +448,7 @@ def __getitem__(self, u):
 
         if idx is None:
             raise ValueError("This term is not contained")
-        elif self.fastmult:
+        elif self.algorithm == "nfft":
             return self.transforms[idx]
         else:
             return self.get_matrix()

src/pyGroupedTransforms/NFFTtools.py

@@ -2,6 +2,10 @@
 
 from pyGroupedTransforms import *
 
+import torch
+import pykeops
+from pykeops.torch import LazyTensor
+
 
 def datalength(
     bandwidths: np.ndarray,
@@ -168,6 +172,80 @@ def adjoint(f):  # function adjoint(f::Vector{ComplexF64})::Vector{ComplexF64}
     return DeferredLinearOperator(
         dtype=np.complex128, shape=(M, N), mfunc=trafo, rmfunc=adjoint
     )
+    
+def get_transform_keops(
+    bandwidths: np.array, X: np.array
+):
+    if bandwidths.ndim > 1 or bandwidths.dtype != "int32":
+        return "Please use an zero or one-dimensional numpy.array with dtype 'int32' as input"
+
+    M, d = np.shape(X)
+
+    if len(bandwidths) == 0:
+        return DeferredLinearOperator(
+            dtype=np.complex128,
+            shape=(M, 1),
+            mfunc=lambda fhat: np.full(M, fhat[0]),
+            rmfunc=lambda f: np.array([np.sum(f)]),
+        )
+
+    X_torch = torch.tensor(X, dtype=torch.float64, device="cuda")
+    freq = index_set_without_zeros(
+        np.array(bandwidths, dtype=np.int32)
+    ).T
+    
+    I_torch = torch.tensor(
+        freq,
+        dtype=torch.float64,
+        device="cuda"
+    )
+
+ 
+    def trafo(fhat): 
+        X_i = LazyTensor(X_torch[:, None, :].contiguous())
+        K_j = LazyTensor(I_torch[None, :, :].contiguous())
+        phase_fwd = (X_i * K_j).sum(-1)
+        two_pi_phase = -2 * torch.pi * phase_fwd
+        kernel_fwd = two_pi_phase.cos() + 1j * two_pi_phase.sin()
+        fhat_torch = torch.tensor(fhat, dtype=torch.complex128, device="cuda")
+        fhat_j = LazyTensor(fhat_torch[None, :, None].contiguous())
+
+        try:
+            result = (kernel_fwd * fhat_j).sum(dim=1)
+            torch.cuda.synchronize()
+            result = result.squeeze().cpu().numpy()
+            return result
+        except Exception as e:
+            print("Error in KeOps trafo:", e)
+            return None
+        
+    def adjoint(f):
+        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())
+
+        phase = (K_j * X_i).sum(-1)
+        kernel = (-2*torch.pi*phase).cos() - 1j*(-2*torch.pi*phase).sin()
+
+        f_i = LazyTensor(f_torch[None, :, None].contiguous())
+
+        try:
+            result = (kernel * f_i).sum(dim=1)
+            torch.cuda.synchronize()
+            result = result.squeeze().cpu().numpy()
+           
+            return result
+        except Exception as e:
+            print("Error in KeOps adjoint:", e)
+            return None
+            
+
+    Ncoeffs = np.prod(bandwidths - 1)
+    
+    return DeferredLinearOperator(
+        dtype=np.complex128, shape=(M, Ncoeffs), mfunc=trafo, rmfunc=adjoint
+    )
 
 
 def get_matrix(