diff --git a/src/pyGroupedTransforms/GroupedTransform.py b/src/pyGroupedTransforms/GroupedTransform.py index 57bf99a..da40e01 100644 --- a/src/pyGroupedTransforms/GroupedTransform.py +++ b/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,8 @@ def __init__( if len(s.bandwidths) == 0: full = np.zeros((1, D), dtype=np.int32) - - - else: + + else: local = index_set_without_zeros( np.array(s.bandwidths, dtype=np.int32) ).T @@ -302,15 +299,11 @@ 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" - ) - - def trafo(fhat): + 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) @@ -327,15 +320,19 @@ def trafo(fhat): 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() - + 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() + kernel = (-2 * torch.pi * phase).cos() - 1j * ( + -2 * torch.pi * phase + ).sin() f_i = LazyTensor(f_torch[None, :, None].contiguous()) @@ -343,16 +340,21 @@ def adjoint(f): 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 - )] - + self.transforms = [ + DeferredLinearOperator( + dtype=np.complex128, + shape=(X.shape[0], len(freq)), + mfunc=trafo, + rmfunc=adjoint, + ) + ] + else: self.transforms = [] s1 = self.settings[0] @@ -432,15 +434,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( @@ -468,10 +470,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: diff --git a/src/pyGroupedTransforms/NFFTtools.py b/src/pyGroupedTransforms/NFFTtools.py index d0b02c6..432ac47 100644 --- a/src/pyGroupedTransforms/NFFTtools.py +++ b/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}