diff --git a/src/pyGroupedTransforms/GroupedTransform.py b/src/pyGroupedTransforms/GroupedTransform.py index 76ed90d..da40e01 100644 --- a/src/pyGroupedTransforms/GroupedTransform.py +++ b/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 = { @@ -266,24 +271,90 @@ 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)) - ] + self.transforms = [ + DeferredLinearOperator() for _ in range(len(self.settings)) + ] - for idx, s in enumerate(self.settings): + D = X.shape[1] - if len(s.bandwidths) == 0: - u = (0,) - else: - u = s.u + 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, + ) + ] - 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] @@ -363,19 +434,15 @@ def adjoint_worker(i): adjoint_worker(i) return fhat + 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 GroupedCoefficients(self.settings, self.transforms[0].H @ other) - 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( @@ -403,15 +470,10 @@ def worker(i): worker(i) return sum(results) + 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 self.transforms[0] @ other.data - return sum(results) elif self.algorithm == "direct": return self.matrix @ other.data else: diff --git a/src/pyGroupedTransforms/NFFTtools.py b/src/pyGroupedTransforms/NFFTtools.py index 10872ae..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, @@ -172,80 +171,6 @@ 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(