Fixes by format action

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pyGroupedTransforms"
-version = "1.0.0"
+version = "1.1.0"
 authors = [
   { name="Felix Wirth", email="fwi012001@gmail.com" },
 ]

src/pyGroupedTransforms/GroupedTransform.py

@@ -1,7 +1,12 @@
 import numpy as np
+import pykeops
+import torch
+from pykeops.torch import LazyTensor
 
 from pyGroupedTransforms import *
 
+from .NFFTtools import index_set_without_zeros
+
 # All code that is linked to NFMTtools or to system = "mixed" is not tested yet....
 
 systems = {
@@ -183,7 +188,7 @@ def __init__(
         system,
         X,
         settings=[],
-        fastmult=True,
+        algorithm="nfft",
         parallel=True,
         basis_vect=[],
         N=[],
@@ -227,9 +232,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 +245,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 +270,91 @@ 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))
+            ]
+
+            D = X.shape[1]
+
+            freq_list = []
+
+            for s in self.settings:
+
+                if len(s.bandwidths) == 0:
+                    full = np.zeros((1, D), dtype=np.int32)
+
+                else:
+                    local = index_set_without_zeros(
+                        np.array(s.bandwidths, dtype=np.int32)
+                    ).T
+
+                    full = np.zeros((local.shape[0], D), dtype=np.int32)
+
+                    for i, dim in enumerate(s.u):
+                        full[:, dim] = local[:, i]
+
+                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")
+
+            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
+
+            self.transforms = [
+                DeferredLinearOperator(
+                    dtype=np.complex128,
+                    shape=(X.shape[0], len(freq)),
+                    mfunc=trafo,
+                    rmfunc=adjoint,
+                )
+            ]
+
         else:
             self.transforms = []
             s1 = self.settings[0]
@@ -321,9 +411,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 +434,22 @@ def adjoint_worker(i):
                         adjoint_worker(i)
 
                 return fhat
-            else:
+
+            elif self.algorithm == "keops":
+                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(
                     "The GroupedTransform and the GroupedCoefficients have different settings"
                 )
 
-            if self.fastmult:
+            if self.algorithm == "nfft":
                 results = []
 
                 def worker(i):
@@ -376,7 +470,11 @@ def worker(i):
                         worker(i)
 
                 return sum(results)
-            else:
+
+            elif self.algorithm == "keops":
+                return self.transforms[0] @ other.data
+
+            elif self.algorithm == "direct":
                 return self.matrix @ other.data
         else:
             raise ValueError("Wrong input data type")
@@ -412,7 +510,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/NFCTtools.py

@@ -77,8 +77,7 @@ def nfct_index_set(
     if d == 1:
         return np.array([i for i in range(0, bandwidths[0])], "int")
 
-    bandwidths = bandwidths[::-1]
-    tmp = tuple([list(range(0, bw)) for bw in bandwidths])
+    tmp = tuple([list(range(0, bw)) for bw in bandwidths[::-1]])
     tmp = itertools.product(*(tmp[::-1]))
 
     freq = np.empty((d, np.prod(bandwidths)), dtype=int)
@@ -127,7 +126,7 @@ def get_transform(
     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) = X.shape
+    M, d = X.shape
 
     if len(bandwidths) == 0:
         return DeferredLinearOperator(

src/pyGroupedTransforms/NFFTtools.py

@@ -1,4 +1,7 @@
 import numpy as np
+import pykeops
+import torch
+from pykeops.torch import LazyTensor
 
 from pyGroupedTransforms import *
 
@@ -88,7 +91,7 @@ def nfft_index_set(
             list(range(int(-bandwidths[0] / 2), int(bandwidths[0] / 2))), dtype="int"
         )
 
-    tmp = [list(range(int(-bw / 2), int(bw / 2))) for bw in bandwidths]
+    tmp = [list(range(int(-bw / 2), int(bw / 2))) for bw in bandwidths[::-1]]
     tmp = itertools.product(*(tmp[::-1]))
 
     freq = np.empty((d, np.prod(bandwidths)), dtype=int)
@@ -134,11 +137,10 @@ def get_transform(
     # Output:
      * `F::LinearMap{ComplexF64}` ... Linear map of the Fourier-transform implemented by the NFFT
     """
-
     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)
+    M, d = np.shape(X)
 
     if len(bandwidths) == 0:
         return DeferredLinearOperator(