From 8ca012cd9d96ebb4488eb291a3401d3584189171 Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 14:45:12 +0200 Subject: [PATCH 01/40] Added --- .../operators/interfaces/nudft_numpy.py | 40 +++-- tests/test_autodiff.py | 145 ++++++++++++++++++ tests/test_ndft.py | 59 +++---- 3 files changed, 205 insertions(+), 39 deletions(-) create mode 100644 tests/test_autodiff.py diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py index 31a9fa18..9d26d02f 100644 --- a/src/mrinufft/operators/interfaces/nudft_numpy.py +++ b/src/mrinufft/operators/interfaces/nudft_numpy.py @@ -8,6 +8,7 @@ from ..base import FourierOperatorCPU +<<<<<<< Updated upstream def get_fourier_matrix(ktraj, shape): """Get the NDFT Fourier Matrix.""" n = np.prod(shape) @@ -17,47 +18,64 @@ def get_fourier_matrix(ktraj, shape): grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape))) traj_grid = ktraj @ grid_r matrix = np.exp(-2j * np.pi * traj_grid) +======= +def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False): + """Get the NDFT Fourier Matrix.""" + n = np.prod(shape) + ndim = len(shape) + matrix = np.zeros((len(ktraj), n), dtype=dtype) + r = [np.linspace(-s/2, s/2-1, s) for s in shape] + grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape))) + traj_grid = ktraj @ grid_r + matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype) + if normalize: + matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) +>>>>>>> Stashed changes return matrix -def implicit_type2_ndft(ktraj, image, shape): +def implicit_type2_ndft(ktraj, image, shape, normalize=False): """Compute the NDFT using the implicit type 2 (image -> kspace) algorithm.""" - r = [np.arange(s) for s in shape] + r = [np.linspace(-s/2, s/2-1, s) for s in shape] grid_r = np.reshape( np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(image.shape)) ) res = np.zeros(len(ktraj), dtype=image.dtype) for j in range(np.prod(image.shape)): res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j]) + if normalize: + matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) return res -def implicit_type1_ndft(ktraj, coeffs, shape): +def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False): """Compute the NDFT using the implicit type 1 (kspace -> image) algorithm.""" - r = [np.arange(s) for s in shape] + r = [np.linspace(-s/2, s/2-1, s) for s in shape] grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(shape))) res = np.zeros(np.prod(shape), dtype=coeffs.dtype) for i in range(len(ktraj)): res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r) + if normalize: + matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) return res -def get_implicit_matrix(ktraj, shape): +def get_implicit_matrix(ktraj, shape, normalize=False): """Get the NDFT Fourier Matrix as implicit operator. This is more memory efficient than the explicit matrix. """ return sp.sparse.linalg.LinearOperator( (len(ktraj), np.prod(shape)), - matvec=lambda x: implicit_type2_ndft(ktraj, x, shape), - rmatvec=lambda x: implicit_type1_ndft(ktraj, x, shape), + matvec=lambda x: implicit_type2_ndft(ktraj, x, shape, normalize), + rmatvec=lambda x: implicit_type1_ndft(ktraj, x, shape, normalize), ) class RawNDFT: """Implementation of the NUDFT using numpy.""" - def __init__(self, samples, shape, explicit_matrix=True): + def __init__(self, samples, shape, explicit_matrix=True, normalize=False): self.samples = samples self.shape = shape self.n_samples = len(samples) @@ -65,13 +83,13 @@ def __init__(self, samples, shape, explicit_matrix=True): if explicit_matrix: try: self._fourier_matrix = sp.sparse.linalg.aslinearoperator( - get_fourier_matrix(self.samples, self.shape) + get_fourier_matrix(self.samples, self.shape, normalize=normalize) ) except MemoryError: warnings.warn("Not enough memory, using an implicit definition anyway") - self._fourier_matrix = get_implicit_matrix(self.samples, self.shape) + self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize) else: - self._fourier_matrix = get_implicit_matrix(self.samples, self.shape) + self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize) def op(self, coeffs, image): """Compute the forward NUDFT.""" diff --git a/tests/test_autodiff.py b/tests/test_autodiff.py new file mode 100644 index 00000000..01de2f91 --- /dev/null +++ b/tests/test_autodiff.py @@ -0,0 +1,145 @@ +"""Test the autodiff functionnality.""" + +import numpy as np +from mrinufft.operators.interfaces.nudft_numpy import get_fourier_matrix +import pytest +from pytest_cases import parametrize_with_cases, parametrize, fixture +from case_trajectories import CasesTrajectories +from mrinufft.operators import get_operator + + +from helpers import ( + kspace_from_op, + image_from_op, + to_interface, +) + + +TORCH_AVAILABLE = True +try: + import torch + import torch.testing as tt +except ImportError: + TORCH_AVAILABLE = False + + +@fixture(scope="module") +@parametrize(backend=["cufinufft", "finufft"]) +@parametrize(autograd=["data"]) +@parametrize_with_cases( + "kspace_loc, shape", + cases=[ + CasesTrajectories.case_grid2D, + CasesTrajectories.case_nyquist_radial2D, + ], # 2D cases only for reduced memory footprint. +) +def operator(kspace_loc, shape, backend, autograd): + """Create NUFFT operator with autodiff capabilities.""" + kspace_loc = kspace_loc.astype(np.float32) + + nufft = get_operator(backend_name=backend, autograd=autograd)( + samples=kspace_loc, + shape=shape, + smaps=None, + ) + + return nufft + + +@fixture(scope="module") +def ndft_matrix(operator): + """Get the NDFT matrix from the operator.""" + return get_fourier_matrix(operator.samples, operator.shape, normalize=True) + + +@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) +@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") +def test_adjoint_and_grad(operator, ndft_matrix, interface): + """Test the adjoint and gradient of the operator.""" + if operator.backend == "finufft" and "gpu" in interface: + pytest.skip("GPU not supported for finufft backend") + ndft_matrix_torch = to_interface(ndft_matrix, interface=interface) + ksp_data = to_interface(kspace_from_op(operator), interface=interface) + img_data = to_interface(image_from_op(operator), interface=interface) + ksp_data.requires_grad = True + + with torch.autograd.set_detect_anomaly(True): + adj_data = operator.adj_op(ksp_data).reshape(img_data.shape) + adj_data_ndft = (ndft_matrix_torch.conj().T @ ksp_data.flatten()).reshape( + adj_data.shape + ) + loss_nufft = torch.mean(torch.abs(adj_data) ** 2) + loss_ndft = torch.mean(torch.abs(adj_data_ndft) ** 2) + + # Check if nufft and ndft are close in the backprop + grad_ndft_kdata = torch.autograd.grad(loss_ndft, ksp_data, retain_graph=True)[0] + grad_nufft_kdata = torch.autograd.grad(loss_nufft, ksp_data, retain_graph=True)[0] + tt.assert_close(grad_ndft_kdata, grad_nufft_kdata, rtol=1, atol=1) + + +@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) +@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") +def test_adjoint_and_gradauto(operator, ndft_matrix, interface): + """Test the adjoint and gradient of the operator using autograd gradcheck.""" + if operator.backend == "finufft" and "gpu" in interface: + pytest.skip("GPU not supported for finufft backend") + + ksp_data = to_interface(kspace_from_op(operator), interface=interface) + ksp_data = torch.ones(ksp_data.shape, requires_grad=True, dtype=ksp_data.dtype) + print(ksp_data.shape) + # todo: tighten the tolerance + assert torch.autograd.gradcheck( + operator.adjoint, + ksp_data, + eps=1e-10, + rtol=1, + atol=1, + nondet_tol=1, + raise_exception=True, + ) + + +@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) +@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") +def test_forward_and_grad(operator, ndft_matrix, interface): + """Test the adjoint and gradient of the operator.""" + if operator.backend == "finufft" and "gpu" in interface: + pytest.skip("GPU not supported for finufft backend") + + ndft_matrix_torch = to_interface(ndft_matrix, interface=interface) + ksp_data_ref = to_interface(kspace_from_op(operator), interface=interface) + img_data = to_interface(image_from_op(operator), interface=interface) + img_data.requires_grad = True + + with torch.autograd.set_detect_anomaly(True): + ksp_data = operator.op(img_data).reshape(ksp_data_ref.shape) + ksp_data_ndft = (ndft_matrix_torch @ img_data.flatten()).reshape(ksp_data.shape) + loss_nufft = torch.mean(torch.abs(ksp_data - ksp_data_ref) ** 2) + loss_ndft = torch.mean(torch.abs(ksp_data_ndft - ksp_data_ref) ** 2) + + # Check if nufft and ndft are close in the backprop + grad_ndft_kdata = torch.autograd.grad(loss_ndft, img_data, retain_graph=True)[0] + grad_nufft_kdata = torch.autograd.grad(loss_nufft, img_data, retain_graph=True)[0] + assert torch.allclose(grad_ndft_kdata, grad_nufft_kdata, atol=6e-3) + + +@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) +@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") +def test_forward_and_gradauto(operator, ndft_matrix, interface): + """Test the forward and gradient of the operator using autograd gradcheck.""" + if operator.backend == "finufft" and "gpu" in interface: + pytest.skip("GPU not supported for finufft backend") + + img_data = to_interface(image_from_op(operator), interface=interface) + img_data = torch.ones(img_data.shape, requires_grad=True, dtype=img_data.dtype) + print(img_data.shape) + # todo: tighten the tolerance + assert torch.autograd.gradcheck( + operator.adjoint, + img_data, + eps=1e-10, + rtol=1, + atol=1, + nondet_tol=1, + raise_exception=True, + ) diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 7ae595a8..5bcb3e8c 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -13,32 +13,9 @@ from case_trajectories import CasesTrajectories, case_grid1D from helpers import assert_almost_allclose +from mrinufft import get_operator -@parametrize_with_cases( - "kspace_grid, shape", - cases=[ - case_grid1D, - CasesTrajectories.case_grid2D, - ], # 3D is ignored (too much possibility for the reordering) -) -def test_ndft_grid_matrix(kspace_grid, shape): - """Test that the ndft matrix is a good matrix for doing fft.""" - ndft_matrix = get_fourier_matrix(kspace_grid, shape) - # Create a random image - fft_matrix = [None] * len(shape) - for i in range(len(shape)): - fft_matrix[i] = sp.fft.fft(np.eye(shape[i])) - fft_mat = fft_matrix[0] - if len(shape) == 2: - fft_mat = fft_matrix[0].flatten()[:, None] @ fft_matrix[1].flatten()[None, :] - fft_mat = ( - fft_mat.reshape(shape * 2) - .transpose(2, 0, 1, 3) - .reshape((np.prod(shape),) * 2) - ) - assert np.allclose(ndft_matrix, fft_mat) - @parametrize_with_cases( "kspace, shape", @@ -56,7 +33,7 @@ def test_ndft_implicit2(kspace, shape): linop_coef = implicit_type2_ndft(kspace, random_image.flatten(), shape) matrix_coef = matrix @ random_image.flatten() - assert np.allclose(linop_coef, matrix_coef) + assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5) @parametrize_with_cases( @@ -76,7 +53,32 @@ def test_ndft_implicit1(kspace, shape): linop_coef = implicit_type1_ndft(kspace, random_kspace.flatten(), shape) matrix_coef = matrix.conj().T @ random_kspace.flatten() - assert np.allclose(linop_coef, matrix_coef) + assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5) + +@parametrize_with_cases( + "kspace, shape", + cases=[ + CasesTrajectories.case_random2D, + CasesTrajectories.case_grid2D, + CasesTrajectories.case_grid3D, + ], +) +def test_ndft_nufft(kspace, shape): + "Test that NDFT matches NUFFT" + ndft_op = RawNDFT(kspace, shape, normalize=True) + random_kspace = 1j * np.random.randn(len(kspace)) + random_kspace += np.random.randn(len(kspace)) + random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape) + operator = get_operator("pynfft")(kspace, shape) # FIXME: @PAC, we need to get ref here + nufft_k = operator.op(random_image) + nufft_i = operator.adj_op(random_kspace) + + ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype) + ndft_i = np.empty(shape, dtype=random_kspace.dtype) + ndft_op.op(ndft_k, random_image) + ndft_op.adj_op(random_kspace, ndft_i) + assert_almost_allclose(nufft_k, ndft_k, atol=1e-4, rtol=1e-4, mismatch=5) + assert_almost_allclose(nufft_i, ndft_i, atol=1e-4, rtol=1e-4, mismatch=5) @parametrize_with_cases( @@ -98,6 +100,7 @@ def test_ndft_fft(kspace_grid, shape): kspace = kspace.reshape(img.shape) if len(shape) >= 2: kspace = kspace.swapaxes(0, 1) - kspace_fft = sp.fft.fftn(img) + kspace_fft = sp.fft.fftn(sp.fft.fftshift(img)) + + assert_almost_allclose(kspace, kspace_fft, atol=1e-4, rtol=1e-4, mismatch=5) - assert_almost_allclose(kspace, kspace_fft, atol=1e-5, rtol=1e-5, mismatch=5) From 093edfbe3def12e22c4aec3657f70377c238525b Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 14:46:50 +0200 Subject: [PATCH 02/40] Fix --- src/mrinufft/operators/interfaces/nudft_numpy.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py index 9d26d02f..68690cc1 100644 --- a/src/mrinufft/operators/interfaces/nudft_numpy.py +++ b/src/mrinufft/operators/interfaces/nudft_numpy.py @@ -8,17 +8,6 @@ from ..base import FourierOperatorCPU -<<<<<<< Updated upstream -def get_fourier_matrix(ktraj, shape): - """Get the NDFT Fourier Matrix.""" - n = np.prod(shape) - ndim = len(shape) - matrix = np.zeros((len(ktraj), n), dtype=complex) - r = [np.arange(shape[i]) for i in range(ndim)] - grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape))) - traj_grid = ktraj @ grid_r - matrix = np.exp(-2j * np.pi * traj_grid) -======= def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False): """Get the NDFT Fourier Matrix.""" n = np.prod(shape) @@ -30,7 +19,6 @@ def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False): matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype) if normalize: matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) ->>>>>>> Stashed changes return matrix From 74c1ecd6d0a07dcc141156aa506be9d1967d07e5 Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 14:48:12 +0200 Subject: [PATCH 03/40] Remove bymistake add --- tests/test_autodiff.py | 145 ----------------------------------------- 1 file changed, 145 deletions(-) delete mode 100644 tests/test_autodiff.py diff --git a/tests/test_autodiff.py b/tests/test_autodiff.py deleted file mode 100644 index 01de2f91..00000000 --- a/tests/test_autodiff.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Test the autodiff functionnality.""" - -import numpy as np -from mrinufft.operators.interfaces.nudft_numpy import get_fourier_matrix -import pytest -from pytest_cases import parametrize_with_cases, parametrize, fixture -from case_trajectories import CasesTrajectories -from mrinufft.operators import get_operator - - -from helpers import ( - kspace_from_op, - image_from_op, - to_interface, -) - - -TORCH_AVAILABLE = True -try: - import torch - import torch.testing as tt -except ImportError: - TORCH_AVAILABLE = False - - -@fixture(scope="module") -@parametrize(backend=["cufinufft", "finufft"]) -@parametrize(autograd=["data"]) -@parametrize_with_cases( - "kspace_loc, shape", - cases=[ - CasesTrajectories.case_grid2D, - CasesTrajectories.case_nyquist_radial2D, - ], # 2D cases only for reduced memory footprint. -) -def operator(kspace_loc, shape, backend, autograd): - """Create NUFFT operator with autodiff capabilities.""" - kspace_loc = kspace_loc.astype(np.float32) - - nufft = get_operator(backend_name=backend, autograd=autograd)( - samples=kspace_loc, - shape=shape, - smaps=None, - ) - - return nufft - - -@fixture(scope="module") -def ndft_matrix(operator): - """Get the NDFT matrix from the operator.""" - return get_fourier_matrix(operator.samples, operator.shape, normalize=True) - - -@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) -@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") -def test_adjoint_and_grad(operator, ndft_matrix, interface): - """Test the adjoint and gradient of the operator.""" - if operator.backend == "finufft" and "gpu" in interface: - pytest.skip("GPU not supported for finufft backend") - ndft_matrix_torch = to_interface(ndft_matrix, interface=interface) - ksp_data = to_interface(kspace_from_op(operator), interface=interface) - img_data = to_interface(image_from_op(operator), interface=interface) - ksp_data.requires_grad = True - - with torch.autograd.set_detect_anomaly(True): - adj_data = operator.adj_op(ksp_data).reshape(img_data.shape) - adj_data_ndft = (ndft_matrix_torch.conj().T @ ksp_data.flatten()).reshape( - adj_data.shape - ) - loss_nufft = torch.mean(torch.abs(adj_data) ** 2) - loss_ndft = torch.mean(torch.abs(adj_data_ndft) ** 2) - - # Check if nufft and ndft are close in the backprop - grad_ndft_kdata = torch.autograd.grad(loss_ndft, ksp_data, retain_graph=True)[0] - grad_nufft_kdata = torch.autograd.grad(loss_nufft, ksp_data, retain_graph=True)[0] - tt.assert_close(grad_ndft_kdata, grad_nufft_kdata, rtol=1, atol=1) - - -@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) -@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") -def test_adjoint_and_gradauto(operator, ndft_matrix, interface): - """Test the adjoint and gradient of the operator using autograd gradcheck.""" - if operator.backend == "finufft" and "gpu" in interface: - pytest.skip("GPU not supported for finufft backend") - - ksp_data = to_interface(kspace_from_op(operator), interface=interface) - ksp_data = torch.ones(ksp_data.shape, requires_grad=True, dtype=ksp_data.dtype) - print(ksp_data.shape) - # todo: tighten the tolerance - assert torch.autograd.gradcheck( - operator.adjoint, - ksp_data, - eps=1e-10, - rtol=1, - atol=1, - nondet_tol=1, - raise_exception=True, - ) - - -@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) -@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") -def test_forward_and_grad(operator, ndft_matrix, interface): - """Test the adjoint and gradient of the operator.""" - if operator.backend == "finufft" and "gpu" in interface: - pytest.skip("GPU not supported for finufft backend") - - ndft_matrix_torch = to_interface(ndft_matrix, interface=interface) - ksp_data_ref = to_interface(kspace_from_op(operator), interface=interface) - img_data = to_interface(image_from_op(operator), interface=interface) - img_data.requires_grad = True - - with torch.autograd.set_detect_anomaly(True): - ksp_data = operator.op(img_data).reshape(ksp_data_ref.shape) - ksp_data_ndft = (ndft_matrix_torch @ img_data.flatten()).reshape(ksp_data.shape) - loss_nufft = torch.mean(torch.abs(ksp_data - ksp_data_ref) ** 2) - loss_ndft = torch.mean(torch.abs(ksp_data_ndft - ksp_data_ref) ** 2) - - # Check if nufft and ndft are close in the backprop - grad_ndft_kdata = torch.autograd.grad(loss_ndft, img_data, retain_graph=True)[0] - grad_nufft_kdata = torch.autograd.grad(loss_nufft, img_data, retain_graph=True)[0] - assert torch.allclose(grad_ndft_kdata, grad_nufft_kdata, atol=6e-3) - - -@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"]) -@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed") -def test_forward_and_gradauto(operator, ndft_matrix, interface): - """Test the forward and gradient of the operator using autograd gradcheck.""" - if operator.backend == "finufft" and "gpu" in interface: - pytest.skip("GPU not supported for finufft backend") - - img_data = to_interface(image_from_op(operator), interface=interface) - img_data = torch.ones(img_data.shape, requires_grad=True, dtype=img_data.dtype) - print(img_data.shape) - # todo: tighten the tolerance - assert torch.autograd.gradcheck( - operator.adjoint, - img_data, - eps=1e-10, - rtol=1, - atol=1, - nondet_tol=1, - raise_exception=True, - ) From 0250aa8d3753bad0191cdc5f42cd1c112f589f44 Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 15:38:27 +0200 Subject: [PATCH 04/40] Fix --- tests/test_ndft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 5bcb3e8c..3d972ff5 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -64,7 +64,7 @@ def test_ndft_implicit1(kspace, shape): ], ) def test_ndft_nufft(kspace, shape): - "Test that NDFT matches NUFFT" + """Test that NDFT matches NUFFT""" ndft_op = RawNDFT(kspace, shape, normalize=True) random_kspace = 1j * np.random.randn(len(kspace)) random_kspace += np.random.randn(len(kspace)) From 060a8bdd125d2140b82dfaa6a492ec78682a80bf Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 15:39:16 +0200 Subject: [PATCH 05/40] Fixed lint --- .../operators/interfaces/nudft_numpy.py | 20 +++++++++++-------- tests/test_ndft.py | 9 +++++---- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py index 68690cc1..3e8e81aa 100644 --- a/src/mrinufft/operators/interfaces/nudft_numpy.py +++ b/src/mrinufft/operators/interfaces/nudft_numpy.py @@ -13,18 +13,18 @@ def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False): n = np.prod(shape) ndim = len(shape) matrix = np.zeros((len(ktraj), n), dtype=dtype) - r = [np.linspace(-s/2, s/2-1, s) for s in shape] + r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape] grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape))) traj_grid = ktraj @ grid_r matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype) if normalize: - matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) + matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) return matrix def implicit_type2_ndft(ktraj, image, shape, normalize=False): """Compute the NDFT using the implicit type 2 (image -> kspace) algorithm.""" - r = [np.linspace(-s/2, s/2-1, s) for s in shape] + r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape] grid_r = np.reshape( np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(image.shape)) ) @@ -32,19 +32,19 @@ def implicit_type2_ndft(ktraj, image, shape, normalize=False): for j in range(np.prod(image.shape)): res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j]) if normalize: - matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) + matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) return res def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False): """Compute the NDFT using the implicit type 1 (kspace -> image) algorithm.""" - r = [np.linspace(-s/2, s/2-1, s) for s in shape] + r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape] grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(shape))) res = np.zeros(np.prod(shape), dtype=coeffs.dtype) for i in range(len(ktraj)): res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r) if normalize: - matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))) + matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) return res @@ -75,9 +75,13 @@ def __init__(self, samples, shape, explicit_matrix=True, normalize=False): ) except MemoryError: warnings.warn("Not enough memory, using an implicit definition anyway") - self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize) + self._fourier_matrix = get_implicit_matrix( + self.samples, self.shape, normalize + ) else: - self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize) + self._fourier_matrix = get_implicit_matrix( + self.samples, self.shape, normalize + ) def op(self, coeffs, image): """Compute the forward NUDFT.""" diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 3d972ff5..7f90d14e 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -16,7 +16,6 @@ from mrinufft import get_operator - @parametrize_with_cases( "kspace, shape", cases=[ @@ -55,6 +54,7 @@ def test_ndft_implicit1(kspace, shape): assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5) + @parametrize_with_cases( "kspace, shape", cases=[ @@ -69,10 +69,12 @@ def test_ndft_nufft(kspace, shape): random_kspace = 1j * np.random.randn(len(kspace)) random_kspace += np.random.randn(len(kspace)) random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape) - operator = get_operator("pynfft")(kspace, shape) # FIXME: @PAC, we need to get ref here + operator = get_operator("pynfft")( + kspace, shape + ) # FIXME: @PAC, we need to get ref here nufft_k = operator.op(random_image) nufft_i = operator.adj_op(random_kspace) - + ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype) ndft_i = np.empty(shape, dtype=random_kspace.dtype) ndft_op.op(ndft_k, random_image) @@ -103,4 +105,3 @@ def test_ndft_fft(kspace_grid, shape): kspace_fft = sp.fft.fftn(sp.fft.fftshift(img)) assert_almost_allclose(kspace, kspace_fft, atol=1e-4, rtol=1e-4, mismatch=5) - From aecb844c74ae53ae67deb852204bc9e647ac28fd Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 15:40:50 +0200 Subject: [PATCH 06/40] Lint --- src/mrinufft/operators/interfaces/nudft_numpy.py | 4 ++-- tests/test_ndft.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py index 3e8e81aa..bcc6c033 100644 --- a/src/mrinufft/operators/interfaces/nudft_numpy.py +++ b/src/mrinufft/operators/interfaces/nudft_numpy.py @@ -32,7 +32,7 @@ def implicit_type2_ndft(ktraj, image, shape, normalize=False): for j in range(np.prod(image.shape)): res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j]) if normalize: - matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) + res /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) return res @@ -44,7 +44,7 @@ def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False): for i in range(len(ktraj)): res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r) if normalize: - matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) + res /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)) return res diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 7f90d14e..fa66b8b2 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -64,7 +64,7 @@ def test_ndft_implicit1(kspace, shape): ], ) def test_ndft_nufft(kspace, shape): - """Test that NDFT matches NUFFT""" + """Test that NDFT matches NUFFT.""" ndft_op = RawNDFT(kspace, shape, normalize=True) random_kspace = 1j * np.random.randn(len(kspace)) random_kspace += np.random.randn(len(kspace)) From 3130bc1c5f443294a2f71dcae30178bb8357d392 Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 17:15:00 +0200 Subject: [PATCH 07/40] Added refbackend --- tests/test_ndft.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_ndft.py b/tests/test_ndft.py index fa66b8b2..57aedfa6 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -63,18 +63,18 @@ def test_ndft_implicit1(kspace, shape): CasesTrajectories.case_grid3D, ], ) -def test_ndft_nufft(kspace, shape): +def test_ndft_nufft(kspace, shape, request): """Test that NDFT matches NUFFT.""" ndft_op = RawNDFT(kspace, shape, normalize=True) random_kspace = 1j * np.random.randn(len(kspace)) random_kspace += np.random.randn(len(kspace)) random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape) - operator = get_operator("pynfft")( + operator = get_operator(request.config.getoption("ref"))( kspace, shape - ) # FIXME: @PAC, we need to get ref here + ) nufft_k = operator.op(random_image) nufft_i = operator.adj_op(random_kspace) - + ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype) ndft_i = np.empty(shape, dtype=random_kspace.dtype) ndft_op.op(ndft_k, random_image) From bc014b8973e3b355f17365a9eb933cc57b92fb4b Mon Sep 17 00:00:00 2001 From: chaithyagr Date: Fri, 26 Apr 2024 17:17:48 +0200 Subject: [PATCH 08/40] Fix NDFT --- tests/test_ndft.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 57aedfa6..7a157d34 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -69,12 +69,10 @@ def test_ndft_nufft(kspace, shape, request): random_kspace = 1j * np.random.randn(len(kspace)) random_kspace += np.random.randn(len(kspace)) random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape) - operator = get_operator(request.config.getoption("ref"))( - kspace, shape - ) + operator = get_operator(request.config.getoption("ref"))(kspace, shape) nufft_k = operator.op(random_image) nufft_i = operator.adj_op(random_kspace) - + ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype) ndft_i = np.empty(shape, dtype=random_kspace.dtype) ndft_op.op(ndft_k, random_image) From 0cc73c41cf743ea19ffa053f0cd43b54f43f192e Mon Sep 17 00:00:00 2001 From: Pierre-antoine Comby Date: Mon, 29 Apr 2024 10:48:25 +0200 Subject: [PATCH 09/40] feat: use finufft as ref backend. --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 69598fdb..4e89f0ed 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,7 @@ def pytest_addoption(parser): ) parser.addoption( "--ref", - default="pynfft", + default="finufft", help="Reference backend on which the tests are performed.", ) From 21e090f21803e9e57cc721010661d30449ce1a0b Mon Sep 17 00:00:00 2001 From: Pierre-antoine Comby Date: Mon, 29 Apr 2024 10:49:37 +0200 Subject: [PATCH 10/40] feat(tests): move ndft vs nufft tests to own file. --- tests/operators/test_operator_ref.py | 74 ++++++++++++++++++++++++++++ tests/test_ndft.py | 26 ---------- 2 files changed, 74 insertions(+), 26 deletions(-) create mode 100644 tests/operators/test_operator_ref.py diff --git a/tests/operators/test_operator_ref.py b/tests/operators/test_operator_ref.py new file mode 100644 index 00000000..b51e1633 --- /dev/null +++ b/tests/operators/test_operator_ref.py @@ -0,0 +1,74 @@ +"""Tests for the reference backend.""" + +from pytest_cases import parametrize_with_cases, fixture +from case_trajectories import CasesTrajectories + +from mrinufft import get_operator +from mrinufft.operators.interfaces.nudft_numpy import MRInumpy +from helpers import assert_almost_allclose, kspace_from_op, image_from_op + + +@fixture(scope="session", autouse=True) +def ref_backend(request): + """Get the reference backend from the CLI.""" + return request.config.getoption("ref") + + +@fixture(scope="module") +@parametrize_with_cases( + "kspace, shape", + cases=[ + CasesTrajectories.case_random2D, + CasesTrajectories.case_grid2D, + CasesTrajectories.case_grid3D, + ], +) +def ref_operator(request, ref_backend, kspace, shape): + """Generate a NFFT operator, matching the property of the first operator.""" + return get_operator(ref_backend)(kspace, shape) + + +@fixture(scope="module") +def ndft_operator(ref_operator): + """Get a NDFT operator matching the reference operator.""" + return MRInumpy(ref_operator.samples, ref_operator.shape) + + +@fixture(scope="module") +def image_data(ref_operator): + """Generate a random image. Remains constant for the module.""" + return image_from_op(ref_operator) + + +@fixture(scope="module") +def kspace_data(ref_operator): + """Generate a random kspace. Remains constant for the module.""" + return kspace_from_op(ref_operator) + + +def test_ref_nufft_forward(ref_operator, ndft_operator, image_data): + """Test that the reference nufft matches the NDFT.""" + nufft_ksp = ref_operator.op(image_data) + ndft_ksp = ndft_operator.op(image_data) + + assert_almost_allclose( + nufft_ksp, + ndft_ksp, + atol=1e-4, + rtol=1e-4, + mismatch=5, + ) + + +def test_ref_nufft_adjoint(ref_operator, ndft_operator, kspace_data): + """Test that the reference nufft matches the NDFT adjoint.""" + nufft_img = ref_operator.adj_op(kspace_data) + ndft_img = ndft_operator.adj_op(kspace_data) + + assert_almost_allclose( + nufft_img, + ndft_img, + atol=1e-4, + rtol=1e-4, + mismatch=5, + ) diff --git a/tests/test_ndft.py b/tests/test_ndft.py index 7a157d34..cd21622e 100644 --- a/tests/test_ndft.py +++ b/tests/test_ndft.py @@ -55,32 +55,6 @@ def test_ndft_implicit1(kspace, shape): assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5) -@parametrize_with_cases( - "kspace, shape", - cases=[ - CasesTrajectories.case_random2D, - CasesTrajectories.case_grid2D, - CasesTrajectories.case_grid3D, - ], -) -def test_ndft_nufft(kspace, shape, request): - """Test that NDFT matches NUFFT.""" - ndft_op = RawNDFT(kspace, shape, normalize=True) - random_kspace = 1j * np.random.randn(len(kspace)) - random_kspace += np.random.randn(len(kspace)) - random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape) - operator = get_operator(request.config.getoption("ref"))(kspace, shape) - nufft_k = operator.op(random_image) - nufft_i = operator.adj_op(random_kspace) - - ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype) - ndft_i = np.empty(shape, dtype=random_kspace.dtype) - ndft_op.op(ndft_k, random_image) - ndft_op.adj_op(random_kspace, ndft_i) - assert_almost_allclose(nufft_k, ndft_k, atol=1e-4, rtol=1e-4, mismatch=5) - assert_almost_allclose(nufft_i, ndft_i, atol=1e-4, rtol=1e-4, mismatch=5) - - @parametrize_with_cases( "kspace_grid, shape", cases=[ From 3da762fe4eaa4699006ff3bcc6f66bc305694916 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Sep 2024 11:59:55 +0200 Subject: [PATCH 11/40] Add support for pipe --- .github/workflows/test-ci.yml | 6 +-- examples/GPU/example_density.py | 20 ++++++- .../GPU/example_learn_samples_multicoil.py | 2 +- pyproject.toml | 4 +- .../operators/interfaces/cufinufft.py | 54 +++++++++++++++++-- tests/operators/test_density_for_op.py | 29 ++++++---- 6 files changed, 91 insertions(+), 24 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index ed3ea469..e6843a95 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -137,8 +137,6 @@ jobs: pip install torchkbnufft elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then pip install tensorflow-mri==0.21.0 tensorflow-probability==0.17.0 tensorflow-io==0.27.0 matplotlib==3.7 - elif [[ ${{ matrix.backend }} == "cufinufft" ]]; then - pip install "cufinufft<2.3" else pip install ${{ matrix.backend }} fi @@ -215,7 +213,7 @@ jobs: export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch - python -m pip install gpuNUFFT "cufinufft<2.3" sigpy scikit-image + python -m pip install gpuNUFFT cufinufft sigpy scikit-image - name: Run examples shell: bash @@ -326,7 +324,7 @@ jobs: export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch - python -m pip install gpuNUFFT "cufinufft<2.3" + python -m pip install gpuNUFFT cufinufft - name: Build API documentation run: | diff --git a/examples/GPU/example_density.py b/examples/GPU/example_density.py index 03c0235a..2b09b22d 100644 --- a/examples/GPU/example_density.py +++ b/examples/GPU/example_density.py @@ -142,7 +142,7 @@ # If this method is widely used in the literature, there exists no convergence guarantees for it. # # .. note:: -# The Pipe method is currently only implemented for gpuNUFFT. +# The Pipe method is currently only implemented for gpuNUFFT and cufinufft backend. # %% flat_traj = traj.reshape(-1, 2) @@ -158,3 +158,21 @@ axs[2].imshow(abs(adjoint_manual)) axs[2].set_title("Pipe density compensation") print(nufft.density) + +# %% +# We can also do density compensation using cufinufft backend + +# %% +flat_traj = traj.reshape(-1, 2) +nufft = get_operator("cufinufft")( + traj, shape=mri_2D.shape, density={"name": "pipe", "osf": 2} +) +adjoint_manual = nufft.adj_op(kspace) +fig, axs = plt.subplots(1, 3, figsize=(15, 5)) +axs[0].imshow(abs(mri_2D)) +axs[0].set_title("Ground Truth") +axs[1].imshow(abs(adjoint)) +axs[1].set_title("no density compensation") +axs[2].imshow(abs(adjoint_manual)) +axs[2].set_title("Pipe density compensation") +print(nufft.density) \ No newline at end of file diff --git a/examples/GPU/example_learn_samples_multicoil.py b/examples/GPU/example_learn_samples_multicoil.py index da8198ec..c61d4282 100644 --- a/examples/GPU/example_learn_samples_multicoil.py +++ b/examples/GPU/example_learn_samples_multicoil.py @@ -75,7 +75,7 @@ def __init__(self, inital_trajectory, n_coils, img_size=(256, 256)): squeeze_dims=False, ) # A simple density compensated adjoint SENSE operator with sensitivity maps `smaps`. - self.sense_op = get_operator("gpunufft", wrt_data=True, wrt_traj=True)( + self.sense_op = get_operator("cufinufft", wrt_data=True, wrt_traj=True)( sample_points, shape=img_size, density=True, diff --git a/pyproject.toml b/pyproject.toml index bfe9bcd7..49d85072 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,8 @@ dynamic = ["version"] gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"] torchkbnufft = ["torchkbnufft", "cupy-cuda12x"] -cufinufft = ["cufinufft<2.3", "cupy-cuda12x"] -finufft = ["finufft"] +cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"] +finufft = ["finufft>=2.3"] pynfft = ["pynfft2>=1.4.3", "numpy>=2.0.0"] pynufft = ["pynufft"] io = ["pymapvbvd"] diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 40aa20af..558647c4 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -40,11 +40,6 @@ DTYPE_R2C = {"float32": "complex64", "float64": "complex128"} -def _error_check(ier, msg): - if ier != 0: - raise RuntimeError(msg) - - class RawCufinufftPlan: """Light wrapper around the guru interface of finufft.""" @@ -836,3 +831,52 @@ def toggle_grad_traj(self): if self.uses_sense: self.smaps = self.smaps.conj() self.raw_op.toggle_grad_traj() + + + @classmethod + def pipe( + cls, + kspace_loc, + volume_shape, + num_iterations=10, + osf=2, + normalize=True, + **kwargs, + ): + """Compute the density compensation weights for a given set of kspace locations. + + Parameters + ---------- + kspace_loc: np.ndarray + the kspace locations + volume_shape: np.ndarray + the volume shape + num_iterations: int default 10 + the number of iterations for density estimation + osf: float or int + The oversampling factor the volume shape + normalize: bool + Whether to normalize the density compensation. + We normalize such that the energy of PSF = 1 + """ + if CUFINUFFT_AVAILABLE is False: + raise ValueError( + "gpuNUFFT is not available, cannot " "estimate the density compensation" + ) + volume_shape = np.array([int(osf * i) for i in volume_shape]) + grid_op = MRICufiNUFFT( + samples=kspace_loc, + shape=volume_shape, + upsampfac=1, + gpu_spreadinterponly=1, + gpu_kerevalmeth=0, + **kwargs, + ) + density_comp = cp.ones(kspace_loc.shape[0], dtype=grid_op.cpx_dtype) + for _ in range(num_iterations): + density_comp /= cp.abs( + grid_op.op( + grid_op.adj_op(density_comp.astype(grid_op.cpx_dtype)) + ).squeeze() + ) + return density_comp.squeeze() \ No newline at end of file diff --git a/tests/operators/test_density_for_op.py b/tests/operators/test_density_for_op.py index 3e24cecf..3161c8b6 100644 --- a/tests/operators/test_density_for_op.py +++ b/tests/operators/test_density_for_op.py @@ -25,21 +25,28 @@ def radial_distance(traj, shape): CasesTrajectories.case_nyquist_radial3D, ], ) -@parametrize(backend=["gpunufft", "tensorflow"]) +@parametrize(backend=["gpunufft", "tensorflow", "cufinufft"]) def test_pipe(backend, traj, shape, osf): """Test the pipe method.""" distance = radial_distance(traj, shape) if osf != 2 and backend == "tensorflow": pytest.skip("OSF < 2 not supported for tensorflow.") - result = pipe(traj, shape, backend, osf=osf, num_iterations=10) + result = pipe(traj, shape, backend=backend, osf=osf, num_iterations=10) + if backend == "cufinufft": + result = result.get() result = result / np.mean(result) distance = distance / np.mean(distance) - if backend == "tensorflow": - # If tensorflow, we dont perfectly estimate, but we still want to ensure - # we can get density - assert_correlate(result, distance, slope=1, slope_err=None, r_value_err=0.5) - elif osf != 2: - # If OSF < 2, we dont perfectly estimate - assert_correlate(result, distance, slope=1, slope_err=None, r_value_err=0.2) - else: - assert_correlate(result, distance, slope=1, slope_err=0.1, r_value_err=0.1) + r_err = 0.2 + slope_err = None + if osf == 2: + r_err = 0.1 + slope_err = 0.1 + if backend == "cufinufft": + r_err *= 2 + slope_err = slope_err * 2 if slope_err is not None else None + if len(shape)==3: + r_err *= 2 + slope_err = slope_err * 2 if slope_err is not None else None + elif backend == "tensorflow": + r_err = 0.5 + assert_correlate(result, distance, slope=1, slope_err=slope_err, r_value_err=r_err) From d644cf670584400561355857d455c2ede990bc2b Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Sep 2024 13:19:32 +0200 Subject: [PATCH 12/40] \!docs_build try to run cufinufft tests --- .github/workflows/test-ci.yml | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index e6843a95..6dbacabb 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -137,6 +137,9 @@ jobs: pip install torchkbnufft elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then pip install tensorflow-mri==0.21.0 tensorflow-probability==0.17.0 tensorflow-io==0.27.0 matplotlib==3.7 + elif [[ ${{ matrix.backend }} == "cufinufft" ]]; then + git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + pip install finufft/python/cufinufft else pip install ${{ matrix.backend }} fi @@ -213,7 +216,9 @@ jobs: export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch - python -m pip install gpuNUFFT cufinufft sigpy scikit-image + python -m pip install gpuNUFFT sigpy scikit-image + git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + pip install finufft/python/cufinufft - name: Run examples shell: bash @@ -324,8 +329,10 @@ jobs: export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch - python -m pip install gpuNUFFT cufinufft - + python -m pip install gpuNUFFT + git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + pip install finufft/python/cufinufft + - name: Build API documentation run: | python -m sphinx docs docs_build From 0dca8f65bff521c0a79d33199840d6dd651966c0 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Sep 2024 13:31:31 +0200 Subject: [PATCH 13/40] \!docs_build fix style --- src/mrinufft/operators/interfaces/cufinufft.py | 5 ++--- tests/operators/test_density_for_op.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 558647c4..ad03ec13 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -832,7 +832,6 @@ def toggle_grad_traj(self): self.smaps = self.smaps.conj() self.raw_op.toggle_grad_traj() - @classmethod def pipe( cls, @@ -868,7 +867,7 @@ def pipe( samples=kspace_loc, shape=volume_shape, upsampfac=1, - gpu_spreadinterponly=1, + gpu_spreadinterponly=1, gpu_kerevalmeth=0, **kwargs, ) @@ -879,4 +878,4 @@ def pipe( grid_op.adj_op(density_comp.astype(grid_op.cpx_dtype)) ).squeeze() ) - return density_comp.squeeze() \ No newline at end of file + return density_comp.squeeze() diff --git a/tests/operators/test_density_for_op.py b/tests/operators/test_density_for_op.py index 3161c8b6..30bd3708 100644 --- a/tests/operators/test_density_for_op.py +++ b/tests/operators/test_density_for_op.py @@ -44,7 +44,7 @@ def test_pipe(backend, traj, shape, osf): if backend == "cufinufft": r_err *= 2 slope_err = slope_err * 2 if slope_err is not None else None - if len(shape)==3: + if len(shape) == 3: r_err *= 2 slope_err = slope_err * 2 if slope_err is not None else None elif backend == "tensorflow": From 643e1e9d09c091d24dc5d5a01d41623c3e90c617 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Mon, 23 Sep 2024 11:47:23 +0200 Subject: [PATCH 14/40] Added next235 for stability --- .../operators/interfaces/cufinufft.py | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index ad03ec13..18ea1c1e 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -40,6 +40,32 @@ DTYPE_R2C = {"float32": "complex64", "float64": "complex128"} +def _next235beven(n, b): + """Find the next even integer not less than n. + + This function finds the next even integer not less than n, with prime factors no + larger than 5, and is a multiple of b (where b is a number that only + has prime factors 2, 3, and 5). + It is used in particular with `pipe` density compensation estimation. + """ + if n <= 2: + return 2 + if n % 2 == 1: + n += 1 # make it even + nplus = n - 2 # to cancel out the +=2 at start of loop + numdiv = 2 # a dummy that is >1 + while numdiv > 1 or nplus % b != 0: + nplus += 2 # stays even + numdiv = nplus + while numdiv % 2 == 0: + numdiv //= 2 # remove all factors of 2, 3, 5... + while numdiv % 3 == 0: + numdiv //= 3 + while numdiv % 5 == 0: + numdiv //= 5 + return nplus + + class RawCufinufftPlan: """Light wrapper around the guru interface of finufft.""" @@ -862,7 +888,7 @@ def pipe( raise ValueError( "gpuNUFFT is not available, cannot " "estimate the density compensation" ) - volume_shape = np.array([int(osf * i) for i in volume_shape]) + volume_shape = np.array([_next235beven(int(osf * i), 1) for i in volume_shape]) grid_op = MRICufiNUFFT( samples=kspace_loc, shape=volume_shape, From af6bbfa06001a31f83dd2ee1a3d19f713d00ea26 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Mon, 23 Sep 2024 12:16:39 +0200 Subject: [PATCH 15/40] Fix lint --- examples/GPU/example_density.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/GPU/example_density.py b/examples/GPU/example_density.py index 2b09b22d..24e51743 100644 --- a/examples/GPU/example_density.py +++ b/examples/GPU/example_density.py @@ -175,4 +175,4 @@ axs[1].set_title("no density compensation") axs[2].imshow(abs(adjoint_manual)) axs[2].set_title("Pipe density compensation") -print(nufft.density) \ No newline at end of file +print(nufft.density) From 02c834f197be3e7f0984b891c6f10bf2fc5dbb05 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Mon, 23 Sep 2024 16:28:24 +0200 Subject: [PATCH 16/40] Fix CUPY --- src/mrinufft/operators/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py index 1fe79e4a..d84a493e 100644 --- a/src/mrinufft/operators/base.py +++ b/src/mrinufft/operators/base.py @@ -441,7 +441,9 @@ def compute_density(self, method=None): if `backend` is `tensorflow`, `gpunufft` or `torchkbnufft-cpu` or `torchkbnufft-gpu`. """ - if isinstance(method, np.ndarray): + if isinstance(method, np.ndarray) or ( + CUPY_AVAILABLE and isinstance(method, cp.ndarray) + ): self.density = method return None if not method: From 8cfd4275baa44c5d582276a3971273fd7ec889d3 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 24 Oct 2024 17:35:12 +0200 Subject: [PATCH 17/40] WIP --- examples/GPU/example_learn_samples.py | 53 ++++++++++--------- .../operators/interfaces/cufinufft.py | 2 +- 2 files changed, 28 insertions(+), 27 deletions(-) diff --git a/examples/GPU/example_learn_samples.py b/examples/GPU/example_learn_samples.py index 631a7a4b..19e83477 100644 --- a/examples/GPU/example_learn_samples.py +++ b/examples/GPU/example_learn_samples.py @@ -47,10 +47,10 @@ def __init__(self, inital_trajectory): data=torch.Tensor(inital_trajectory), requires_grad=True, ) - self.operator = get_operator("gpunufft", wrt_data=True, wrt_traj=True)( + self.operator = get_operator("cufinufft", wrt_data=True, wrt_traj=True)( self.trajectory.detach().cpu().numpy(), shape=(256, 256), - density=True, + density=False, squeeze_dims=False, ) @@ -60,11 +60,12 @@ def forward(self, x): self.operator.samples = self.trajectory.clone() # A simple acquisition model simulated with a forward NUFFT operator - kspace = self.operator.op(x) + #kspace = self.operator.op(x) # A simple density compensated adjoint operator - adjoint = self.operator.adj_op(kspace) - return adjoint / torch.linalg.norm(adjoint) + #adjoint = self.operator.adj_op(kspace) + #return adjoint / torch.linalg.norm(adjoint) + return # %% @@ -113,8 +114,8 @@ def plot_state(axs, mri_2D, traj, recon, loss=None, save_name=None): mri_2D = mri_2D / torch.linalg.norm(mri_2D) model.eval() recon = model(mri_2D) -fig, axs = plt.subplots(1, 3, figsize=(15, 5)) -plot_state(axs, mri_2D, init_traj, recon) +#fig, axs = plt.subplots(1, 3, figsize=(15, 5)) +#plot_state(axs, mri_2D, init_traj, recon) # %% # Start training loop @@ -122,34 +123,34 @@ def plot_state(axs, mri_2D, traj, recon, loss=None, save_name=None): losses = [] image_files = [] model.train() -with tqdm(range(100), unit="steps") as tqdms: +with tqdm(range(1000), unit="steps") as tqdms: for i in tqdms: out = model(mri_2D) - loss = torch.norm(out - mri_2D[None]) - numpy_loss = loss.detach().cpu().numpy() - tqdms.set_postfix({"loss": numpy_loss}) - losses.append(numpy_loss) - optimizer.zero_grad() - loss.backward() - optimizer.step() + #loss = torch.norm(out - mri_2D[None]) + #numpy_loss = loss.detach().cpu().numpy() + #tqdms.set_postfix({"loss": numpy_loss}) + #losses.append(numpy_loss) + #optimizer.zero_grad() + #loss.backward() + #optimizer.step() with torch.no_grad(): # Clamp the value of trajectory between [-0.5, 0.5] for param in model.parameters(): param.clamp_(-0.5, 0.5) - schedulder.step() + #schedulder.step() # Generate images for gif hashed = joblib.hash((i, "learn_traj", time.time())) filename = "/tmp/" + f"{hashed}.png" - fig, axs = plt.subplots(2, 2, figsize=(10, 10)) - plot_state( - axs, - mri_2D, - model.trajectory.detach().cpu().numpy(), - out, - losses, - save_name=filename, - ) - image_files.append(filename) + #fig, axs = plt.subplots(2, 2, figsize=(10, 10)) + #plot_state( + # axs, + # mri_2D, + # model.trajectory.detach().cpu().numpy(), + # out, + # losses, + # save_name=filename, + #) + #image_files.append(filename) # Make a GIF of all images. diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 18ea1c1e..0eb0cf62 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -298,7 +298,7 @@ def samples(self, samples): if typ == "grad" and not self._grad_wrt_traj: continue self.raw_op._set_pts(typ, samples) - self.compute_density(self._density_method) + #self.compute_density(self._density_method) @with_numpy_cupy @nvtx_mark() From 3c3f1c811ea477fa6cc7b2fddafaa1bd0d9cb73e Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 24 Oct 2024 18:34:25 +0200 Subject: [PATCH 18/40] Updates --- src/mrinufft/operators/base.py | 4 ++- .../operators/interfaces/cufinufft.py | 27 +++++++++++-------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py index 34f676d2..ce0e4121 100644 --- a/src/mrinufft/operators/base.py +++ b/src/mrinufft/operators/base.py @@ -13,7 +13,7 @@ import numpy as np -from mrinufft._array_compat import with_numpy, with_numpy_cupy, AUTOGRAD_AVAILABLE +from mrinufft._array_compat import with_numpy, with_numpy_cupy, AUTOGRAD_AVAILABLE, CUPY_AVAILABLE from mrinufft._utils import auto_cast, power_method from mrinufft.density import get_density from mrinufft.extras import get_smaps @@ -21,6 +21,8 @@ if AUTOGRAD_AVAILABLE: from mrinufft.operators.autodiff import MRINufftAutoGrad +if CUPY_AVAILABLE: + import cupy as cp # Mapping between numpy float and complex types. diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 8e94b5e9..5fc81949 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -86,10 +86,12 @@ def __init__( # and type 2 with 2. self.plans = [None, None, None] self.grad_plan = None - + self._kx = cp.array(samples[:, 0], copy=False) + self._ky = cp.array(samples[:, 1], copy=False) + self._kz = cp.array(samples[:, 2], copy=False) if self.ndim == 3 else None for i in [1, 2]: self._make_plan(i, **kwargs) - self._set_pts(i, samples) + self._set_pts(i) @property def dtype(self): @@ -108,14 +110,16 @@ def _make_plan(self, typ, **kwargs): dtype=DTYPE_R2C[str(self._dtype)], **kwargs, ) - - def _set_pts(self, typ, samples): + + def _set_kxyz(self, samples): + self._kx.set(samples[:, 0]) + self._ky.set(samples[:, 1]) + if self.ndim == 3: + self._kz.set(samples[:, 2]) + + def _set_pts(self, typ): plan = self.grad_plan if typ == "grad" else self.plans[typ] - plan.setpts( - cp.array(samples[:, 0], copy=False), - cp.array(samples[:, 1], copy=False), - cp.array(samples[:, 2], copy=False) if self.ndim == 3 else None, - ) + plan.setpts(self._kx, self._ky, self._kz) def _destroy_plan(self, typ): if self.plans[typ] is not None: @@ -295,10 +299,11 @@ def samples(self, samples): self._samples = np.asfortranarray( proper_trajectory(samples, normalize="pi").astype(np.float32, copy=False) ) + self.raw_op._set_kxyz(self._samples) for typ in [1, 2, "grad"]: if typ == "grad" and not self._grad_wrt_traj: continue - self.raw_op._set_pts(typ, self._samples) + self.raw_op._set_pts(typ) self.compute_density(self._density_method) @FourierOperatorBase.density.setter @@ -831,7 +836,7 @@ def _make_plan_grad(self, **kwargs): isign=1, **kwargs, ) - self.raw_op._set_pts(typ="grad", samples=self.samples) + self.raw_op._set_pts(typ="grad") def get_lipschitz_cst(self, max_iter=10, **kwargs): """Return the Lipschitz constant of the operator. From bb28eb9d4cf0c580923824cd197790296f5f9d1c Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 25 Oct 2024 11:42:22 +0200 Subject: [PATCH 19/40] fix back learn examples --- examples/GPU/example_learn_samples.py | 53 +++++++++++++-------------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/examples/GPU/example_learn_samples.py b/examples/GPU/example_learn_samples.py index 19e83477..631a7a4b 100644 --- a/examples/GPU/example_learn_samples.py +++ b/examples/GPU/example_learn_samples.py @@ -47,10 +47,10 @@ def __init__(self, inital_trajectory): data=torch.Tensor(inital_trajectory), requires_grad=True, ) - self.operator = get_operator("cufinufft", wrt_data=True, wrt_traj=True)( + self.operator = get_operator("gpunufft", wrt_data=True, wrt_traj=True)( self.trajectory.detach().cpu().numpy(), shape=(256, 256), - density=False, + density=True, squeeze_dims=False, ) @@ -60,12 +60,11 @@ def forward(self, x): self.operator.samples = self.trajectory.clone() # A simple acquisition model simulated with a forward NUFFT operator - #kspace = self.operator.op(x) + kspace = self.operator.op(x) # A simple density compensated adjoint operator - #adjoint = self.operator.adj_op(kspace) - #return adjoint / torch.linalg.norm(adjoint) - return + adjoint = self.operator.adj_op(kspace) + return adjoint / torch.linalg.norm(adjoint) # %% @@ -114,8 +113,8 @@ def plot_state(axs, mri_2D, traj, recon, loss=None, save_name=None): mri_2D = mri_2D / torch.linalg.norm(mri_2D) model.eval() recon = model(mri_2D) -#fig, axs = plt.subplots(1, 3, figsize=(15, 5)) -#plot_state(axs, mri_2D, init_traj, recon) +fig, axs = plt.subplots(1, 3, figsize=(15, 5)) +plot_state(axs, mri_2D, init_traj, recon) # %% # Start training loop @@ -123,34 +122,34 @@ def plot_state(axs, mri_2D, traj, recon, loss=None, save_name=None): losses = [] image_files = [] model.train() -with tqdm(range(1000), unit="steps") as tqdms: +with tqdm(range(100), unit="steps") as tqdms: for i in tqdms: out = model(mri_2D) - #loss = torch.norm(out - mri_2D[None]) - #numpy_loss = loss.detach().cpu().numpy() - #tqdms.set_postfix({"loss": numpy_loss}) - #losses.append(numpy_loss) - #optimizer.zero_grad() - #loss.backward() - #optimizer.step() + loss = torch.norm(out - mri_2D[None]) + numpy_loss = loss.detach().cpu().numpy() + tqdms.set_postfix({"loss": numpy_loss}) + losses.append(numpy_loss) + optimizer.zero_grad() + loss.backward() + optimizer.step() with torch.no_grad(): # Clamp the value of trajectory between [-0.5, 0.5] for param in model.parameters(): param.clamp_(-0.5, 0.5) - #schedulder.step() + schedulder.step() # Generate images for gif hashed = joblib.hash((i, "learn_traj", time.time())) filename = "/tmp/" + f"{hashed}.png" - #fig, axs = plt.subplots(2, 2, figsize=(10, 10)) - #plot_state( - # axs, - # mri_2D, - # model.trajectory.detach().cpu().numpy(), - # out, - # losses, - # save_name=filename, - #) - #image_files.append(filename) + fig, axs = plt.subplots(2, 2, figsize=(10, 10)) + plot_state( + axs, + mri_2D, + model.trajectory.detach().cpu().numpy(), + out, + losses, + save_name=filename, + ) + image_files.append(filename) # Make a GIF of all images. From cdf75af63fdc87894047c883f168ae36b6f5ac41 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 25 Oct 2024 12:12:32 +0200 Subject: [PATCH 20/40] move tto flatiron --- .github/workflows/test-ci.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index af51571a..904a21fa 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -138,7 +138,7 @@ jobs: elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then pip install tensorflow-mri==0.21.0 tensorflow-probability==0.17.0 tensorflow-io==0.27.0 matplotlib==3.7 elif [[ ${{ matrix.backend }} == "cufinufft" ]]; then - git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + git clone https://github.com/flatironinstitute/finufft pip install finufft/python/cufinufft else pip install ${{ matrix.backend }} @@ -217,7 +217,7 @@ jobs: export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch python -m pip install gpuNUFFT sigpy scikit-image - git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + git clone https://github.com/flatironinstitute/finufft pip install finufft/python/cufinufft - name: Run examples @@ -330,7 +330,7 @@ jobs: export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch python -m pip install gpuNUFFT - git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly + git clone https://github.com/flatironinstitute/finufft pip install finufft/python/cufinufft - name: Build API documentation From 986fb9636a81f975e19f5a60b84c64681373931c Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 25 Oct 2024 12:14:21 +0200 Subject: [PATCH 21/40] fix black --- src/mrinufft/operators/base.py | 7 ++++++- src/mrinufft/operators/interfaces/cufinufft.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py index ce0e4121..5a129016 100644 --- a/src/mrinufft/operators/base.py +++ b/src/mrinufft/operators/base.py @@ -13,7 +13,12 @@ import numpy as np -from mrinufft._array_compat import with_numpy, with_numpy_cupy, AUTOGRAD_AVAILABLE, CUPY_AVAILABLE +from mrinufft._array_compat import ( + with_numpy, + with_numpy_cupy, + AUTOGRAD_AVAILABLE, + CUPY_AVAILABLE, +) from mrinufft._utils import auto_cast, power_method from mrinufft.density import get_density from mrinufft.extras import get_smaps diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 5fc81949..a0f00005 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -110,7 +110,7 @@ def _make_plan(self, typ, **kwargs): dtype=DTYPE_R2C[str(self._dtype)], **kwargs, ) - + def _set_kxyz(self, samples): self._kx.set(samples[:, 0]) self._ky.set(samples[:, 1]) From d4edc58c68bbd570a63b5ca6e435b487562db451 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 8 Nov 2024 09:47:25 +0100 Subject: [PATCH 22/40] Move to test on GPU --- examples/GPU/example_learn_samples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/GPU/example_learn_samples.py b/examples/GPU/example_learn_samples.py index 631a7a4b..a827e527 100644 --- a/examples/GPU/example_learn_samples.py +++ b/examples/GPU/example_learn_samples.py @@ -47,7 +47,7 @@ def __init__(self, inital_trajectory): data=torch.Tensor(inital_trajectory), requires_grad=True, ) - self.operator = get_operator("gpunufft", wrt_data=True, wrt_traj=True)( + self.operator = get_operator("cufinufft", wrt_data=True, wrt_traj=True)( self.trajectory.detach().cpu().numpy(), shape=(256, 256), density=True, From 1d014845bb95e1c512c43b01d13bac446d385f1d Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Wed, 13 Nov 2024 11:13:40 +0100 Subject: [PATCH 23/40] Update pyproject toml and use it in test-ci, to prevent duplication of dependencies and actually test them! --- .github/workflows/test-ci.yml | 49 +++++------------------------------ pyproject.toml | 12 ++++++++- 2 files changed, 18 insertions(+), 43 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 904a21fa..50b8ca13 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -46,28 +46,10 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install -e .[test] - - - name: Install pynfft - if: ${{ matrix.backend == 'pynfft' || env.ref_backend == 'pynfft' }} - shell: bash - run: | - python -m pip install "pynfft2>=1.4.3" - - - name: Install pynufft - if: ${{ matrix.backend == 'pynufft-cpu' || env.ref_backend == 'pynufft-cpu' }} - run: python -m pip install pynufft - - - name: Install finufft - if: ${{ matrix.backend == 'finufft' || env.ref_backend == 'finufft'}} - shell: bash - run: python -m pip install finufft - - - name: Install Sigpy - if: ${{ matrix.backend == 'sigpy' || env.ref_backend == 'sigpy'}} - shell: bash - run: python -m pip install sigpy - - - name: Install BART + python -m pip install -e .[${{ env.ref_backend }}] + python -m pip install -e .[${{ matrix.backend }}] + + - name: Install BART if needed if: ${{ matrix.backend == 'bart' || env.ref_backend == 'bart'}} shell: bash run: | @@ -79,11 +61,6 @@ jobs: make echo $PWD >> $GITHUB_PATH - - name: Install torchkbnufft-cpu - if: ${{ matrix.backend == 'torchkbnufft-cpu' || env.ref_backend == 'torchkbnufft-cpu'}} - run: python -m pip install torchkbnufft - - - name: Run Tests shell: bash run: | @@ -117,7 +94,6 @@ jobs: source $RUNNER_WORKSPACE/venv/bin/activate pip install --upgrade pip wheel pip install -e mri-nufft[test] - pip install cupy-cuda12x finufft "numpy<2.0" - name: Install torch with CUDA 12.1 shell: bash @@ -133,16 +109,7 @@ jobs: export CUDA_BIN_PATH=/usr/local/cuda-12.1/ export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} - if [[ ${{ matrix.backend }} == "torchkbnufft-gpu" ]]; then - pip install torchkbnufft - elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then - pip install tensorflow-mri==0.21.0 tensorflow-probability==0.17.0 tensorflow-io==0.27.0 matplotlib==3.7 - elif [[ ${{ matrix.backend }} == "cufinufft" ]]; then - git clone https://github.com/flatironinstitute/finufft - pip install finufft/python/cufinufft - else - pip install ${{ matrix.backend }} - fi + pip install -e .[${{ matrix.backend }}] - name: Run Tests shell: bash @@ -217,8 +184,7 @@ jobs: export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch python -m pip install gpuNUFFT sigpy scikit-image - git clone https://github.com/flatironinstitute/finufft - pip install finufft/python/cufinufft + pip install "cufinufft>=2.3.1" - name: Run examples shell: bash @@ -330,8 +296,7 @@ jobs: export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} pip install cupy-cuda12x torch python -m pip install gpuNUFFT - git clone https://github.com/flatironinstitute/finufft - pip install finufft/python/cufinufft + pip install "cufinufft>=2.3.1" - name: Build API documentation run: | diff --git a/pyproject.toml b/pyproject.toml index 49d85072..c4edd4fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,11 +12,21 @@ dynamic = ["version"] [project.optional-dependencies] gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"] + torchkbnufft = ["torchkbnufft", "cupy-cuda12x"] -cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"] +torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"] +torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"] + +cufinufft = ["cufinufft>=2.3.1", "cupy-cuda12x"] +tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"] finufft = ["finufft>=2.3"] +sigpy = ["sigpy"] pynfft = ["pynfft2>=1.4.3", "numpy>=2.0.0"] + pynufft = ["pynufft"] +pynufft-cpu = ["pynufft"] +pynufft-gpu = ["pynufft"] + io = ["pymapvbvd"] smaps = ["scikit-image"] autodiff = ["torch"] From 9714ca97b5d202b0868c32dc0f1b06dca954c4cd Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Wed, 13 Nov 2024 11:22:52 +0100 Subject: [PATCH 24/40] Make CI build shorter --- .github/workflows/test-ci.yml | 46 ++++++++++++----------------------- 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 50b8ca13..63bbbdb8 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -45,9 +45,7 @@ jobs: shell: bash run: | python -m pip install --upgrade pip - python -m pip install -e .[test] - python -m pip install -e .[${{ env.ref_backend }}] - python -m pip install -e .[${{ matrix.backend }}] + python -m pip install -e .[test,${{ env.ref_backend }},${{ matrix.backend }}] - name: Install BART if needed if: ${{ matrix.backend == 'bart' || env.ref_backend == 'bart'}} @@ -95,12 +93,6 @@ jobs: pip install --upgrade pip wheel pip install -e mri-nufft[test] - - name: Install torch with CUDA 12.1 - shell: bash - if: ${{ matrix.backend != 'tensorflow'}} - run: | - source $RUNNER_WORKSPACE/venv/bin/activate - pip install torch - name: Install backend shell: bash @@ -109,7 +101,7 @@ jobs: export CUDA_BIN_PATH=/usr/local/cuda-12.1/ export PATH=/usr/local/cuda-12.1/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} - pip install -e .[${{ matrix.backend }}] + pip install -e .[${{ matrix.backend }},autodiff] - name: Run Tests shell: bash @@ -170,21 +162,18 @@ jobs: path: ~/.cache/brainweb key: ${{ runner.os }}-Brainweb + - name: Point to CUDA 12.1 #TODO: This can be combined from other jobs + run: | + export CUDA_BIN_PATH=/usr/local/cuda-12.1/ + export PATH=/usr/local/cuda-12.1/bin/:${PATH} + export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} + - name: Install Python deps shell: bash run: | python -m pip install --upgrade pip - python -m pip install -e .[test,dev] - python -m pip install finufft pooch brainweb-dl torch + python -m pip install -e .[test,dev,finufft,cufinufft,gpuNUFFT,sigpy,smaps,autodiff,doc] - - name: Install GPU related interfaces - run: | - export CUDA_BIN_PATH=/usr/local/cuda-12.1/ - export PATH=/usr/local/cuda-12.1/bin/:${PATH} - export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} - pip install cupy-cuda12x torch - python -m pip install gpuNUFFT sigpy scikit-image - pip install "cufinufft>=2.3.1" - name: Run examples shell: bash @@ -282,21 +271,18 @@ jobs: with: python-version: "3.10" + - name: Point to CUDA 12.1 + run: | + export CUDA_BIN_PATH=/usr/local/cuda-12.1/ + export PATH=/usr/local/cuda-12.1/bin/:${PATH} + export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} + - name: Install dependencies shell: bash -l {0} run: | python -m pip install --upgrade pip - python -m pip install .[doc] - python -m pip install finufft + python -m pip install .[doc,finufft,autodiff,gpunufft,cufinufft] - - name: Install GPU related interfaces - run: | - export CUDA_BIN_PATH=/usr/local/cuda-12.1/ - export PATH=/usr/local/cuda-12.1/bin/:${PATH} - export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} - pip install cupy-cuda12x torch - python -m pip install gpuNUFFT - pip install "cufinufft>=2.3.1" - name: Build API documentation run: | From 78c60f92ae5ac99be3313f4aa33b74931b7f0396 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Wed, 13 Nov 2024 11:31:47 +0100 Subject: [PATCH 25/40] Test run to run --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c4edd4fe..d227367a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ torchkbnufft = ["torchkbnufft", "cupy-cuda12x"] torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"] torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"] -cufinufft = ["cufinufft>=2.3.1", "cupy-cuda12x"] +cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"] tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"] finufft = ["finufft>=2.3"] sigpy = ["sigpy"] From d8448169f291e878d35019f6ffd072f072c9287d Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Wed, 13 Nov 2024 12:36:23 +0100 Subject: [PATCH 26/40] \!docs_build Added --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d227367a..69c87581 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"] tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"] finufft = ["finufft>=2.3"] sigpy = ["sigpy"] -pynfft = ["pynfft2>=1.4.3", "numpy>=2.0.0"] +pynfft = ["pynfft2>=1.4.3", "numpy<2.0.0"] pynufft = ["pynufft"] pynufft-cpu = ["pynufft"] From 7a6f4a031eeb10b29a6428e25febfa2a587d61f6 Mon Sep 17 00:00:00 2001 From: Asma TANABENE Date: Thu, 5 Dec 2024 15:01:18 +0100 Subject: [PATCH 27/40] adding density normalization --- src/mrinufft/operators/interfaces/cufinufft.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index a0f00005..e6f58518 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -906,6 +906,7 @@ def pipe( raise ValueError( "gpuNUFFT is not available, cannot " "estimate the density compensation" ) + original_shape = volume_shape volume_shape = np.array([_next235beven(int(osf * i), 1) for i in volume_shape]) grid_op = MRICufiNUFFT( samples=kspace_loc, @@ -922,4 +923,9 @@ def pipe( grid_op.adj_op(density_comp.astype(grid_op.cpx_dtype)) ).squeeze() ) + if normalize: + test_op = MRICufiNUFFT(samples=kspace_loc, shape=original_shape, **kwargs) + test_im = cp.ones(original_shape, dtype=test_op.cpx_dtype) + test_im_recon = test_op.adj_op(density_comp * test_op.op(test_im)) + density_comp /= cp.mean(cp.abs(test_im_recon)) return density_comp.squeeze() From a94d5307141e18d15e0e6607098b7c99040d53bd Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Wed, 11 Dec 2024 16:46:15 +0100 Subject: [PATCH 28/40] Add support for 2.3.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1df5f983..66c146be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ torchkbnufft = ["torchkbnufft", "cupy-cuda12x"] torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"] torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"] -cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"] +cufinufft = ["cufinufft>=2.3.1", "cupy-cuda12x"] tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"] finufft = ["finufft>=2.3"] sigpy = ["sigpy"] From 67d6c754b826b09a40ed58f64b9bd734f72f3761 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Dec 2024 14:31:11 +0100 Subject: [PATCH 29/40] Multiple fixes --- examples/GPU/example_learn_samples.py | 2 +- pyproject.toml | 9 +++++---- src/mrinufft/density/utils.py | 9 +++++++++ src/mrinufft/operators/interfaces/cufinufft.py | 12 ++++++++---- src/mrinufft/operators/interfaces/gpunufft.py | 14 +++++++++----- 5 files changed, 32 insertions(+), 14 deletions(-) diff --git a/examples/GPU/example_learn_samples.py b/examples/GPU/example_learn_samples.py index c81ff77a..81f12ed4 100644 --- a/examples/GPU/example_learn_samples.py +++ b/examples/GPU/example_learn_samples.py @@ -24,7 +24,7 @@ # .. colab-link:: # :needs_gpu: 1 # -# !pip install mri-nufft[gpunufft] scikit-image +# !pip install mri-nufft[cufinufft] scikit-image import time import joblib diff --git a/pyproject.toml b/pyproject.toml index c0c25f29..fccd0cc0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,8 +14,8 @@ dynamic = ["version"] gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"] torchkbnufft = ["torchkbnufft", "cupy-cuda12x"] -torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"] -torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"] +torchkbnufft-cpu = ["mri-nufft[torchkbnufft]"] +torchkbnufft-gpu = ["mri-nufft[torchkbnufft]"] cufinufft = ["cufinufft>=2.3.1", "cupy-cuda12x"] tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"] @@ -24,8 +24,9 @@ sigpy = ["sigpy"] pynfft = ["pynfft2>=1.4.3; python_version < '3.12'", "numpy>=2.0.0; python_version < '3.12'"] pynufft = ["pynufft"] -pynufft-cpu = ["pynufft"] -pynufft-gpu = ["pynufft"] +pynufft-cpu = ["mri-nufft[pynufft]"] +pynufft-gpu = ["mri-nufft[pynufft]"] + extra = ["pymapvbvd", "scikit-image", "scikit-learn", "pywavelets"] autodiff = ["torch"] diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index 0b42ac91..ab65f4b3 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -5,6 +5,7 @@ import numpy as np from mrinufft._utils import MethodRegister, proper_trajectory +from mrinufft import get_operator register_density = MethodRegister("density_compensation") @@ -51,3 +52,11 @@ def normalize_weights(weights): """ inv_weights = np.sum(weights) / weights return inv_weights / (np.sum(inv_weights)) + + +def normalize_density(kspace_loc, shape, density, backend, **kwargs): + """Normalize the density to ensure that the reconstruction is stable.""" + test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) + test_im = np.ones(shape, dtype=test_op.cpx_dtype) + test_im_recon = test_op.adj_op(density * test_op.op(test_im)) + density /= np.mean(np.abs(test_im_recon)) \ No newline at end of file diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index e6f58518..8aaa7877 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -18,6 +18,7 @@ pin_memory, sizeof_fmt, ) +from mrinufft.density.utils import normalize_density CUFINUFFT_AVAILABLE = CUPY_AVAILABLE try: @@ -924,8 +925,11 @@ def pipe( ).squeeze() ) if normalize: - test_op = MRICufiNUFFT(samples=kspace_loc, shape=original_shape, **kwargs) - test_im = cp.ones(original_shape, dtype=test_op.cpx_dtype) - test_im_recon = test_op.adj_op(density_comp * test_op.op(test_im)) - density_comp /= cp.mean(cp.abs(test_im_recon)) + density_comp = normalize_density( + kspace_loc=kspace_loc, + shape=original_shape, + density=density_comp, + backend=cls.backend, + **kwargs + ) return density_comp.squeeze() diff --git a/src/mrinufft/operators/interfaces/gpunufft.py b/src/mrinufft/operators/interfaces/gpunufft.py index b5d65af6..118c94d2 100644 --- a/src/mrinufft/operators/interfaces/gpunufft.py +++ b/src/mrinufft/operators/interfaces/gpunufft.py @@ -6,6 +6,7 @@ from ..base import FourierOperatorBase, with_numpy_cupy from mrinufft._utils import proper_trajectory, get_array_module, auto_cast from mrinufft.operators.interfaces.utils import is_cuda_array, is_host_array +from mrinufft.density.utils import normalize_density GPUNUFFT_AVAILABLE = True try: @@ -596,6 +597,7 @@ def pipe( raise ValueError( "gpuNUFFT is not available, cannot " "estimate the density compensation" ) + original_shape = volume_shape volume_shape = (np.array(volume_shape) * osf).astype(int) grid_op = MRIGpuNUFFT( samples=kspace_loc, @@ -607,11 +609,13 @@ def pipe( max_iter=num_iterations ) if normalize: - spike = np.zeros(volume_shape) - mid_loc = tuple(v // 2 for v in volume_shape) - spike[mid_loc] = 1 - psf = grid_op.adj_op(grid_op.op(spike)) - density_comp /= np.linalg.norm(psf) + density_comp = normalize_density( + kspace_loc=kspace_loc, + shape=original_shape, + density=density_comp, + backend=cls.backend, + **kwargs + ) return density_comp.squeeze() def get_lipschitz_cst(self, max_iter=10, tolerance=1e-5, **kwargs): From 906f363f4a9a685bdd6a6ce01d71614b584bb4ec Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Dec 2024 14:32:39 +0100 Subject: [PATCH 30/40] Prevent circular import --- src/mrinufft/density/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index ab65f4b3..fb52d822 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -5,7 +5,6 @@ import numpy as np from mrinufft._utils import MethodRegister, proper_trajectory -from mrinufft import get_operator register_density = MethodRegister("density_compensation") @@ -56,6 +55,7 @@ def normalize_weights(weights): def normalize_density(kspace_loc, shape, density, backend, **kwargs): """Normalize the density to ensure that the reconstruction is stable.""" + from mrinufft import get_operator test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) test_im = np.ones(shape, dtype=test_op.cpx_dtype) test_im_recon = test_op.adj_op(density * test_op.op(test_im)) From 9917045a1706897b7bb5ebce4f25e72d31831ab4 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Dec 2024 14:45:14 +0100 Subject: [PATCH 31/40] fix Lint --- src/mrinufft/density/utils.py | 3 ++- src/mrinufft/operators/interfaces/cufinufft.py | 2 +- src/mrinufft/operators/interfaces/gpunufft.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index fb52d822..1a774dfc 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -56,7 +56,8 @@ def normalize_weights(weights): def normalize_density(kspace_loc, shape, density, backend, **kwargs): """Normalize the density to ensure that the reconstruction is stable.""" from mrinufft import get_operator + test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) test_im = np.ones(shape, dtype=test_op.cpx_dtype) test_im_recon = test_op.adj_op(density * test_op.op(test_im)) - density /= np.mean(np.abs(test_im_recon)) \ No newline at end of file + density /= np.mean(np.abs(test_im_recon)) diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 8aaa7877..77cd96fb 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -930,6 +930,6 @@ def pipe( shape=original_shape, density=density_comp, backend=cls.backend, - **kwargs + **kwargs, ) return density_comp.squeeze() diff --git a/src/mrinufft/operators/interfaces/gpunufft.py b/src/mrinufft/operators/interfaces/gpunufft.py index 118c94d2..02ef22f7 100644 --- a/src/mrinufft/operators/interfaces/gpunufft.py +++ b/src/mrinufft/operators/interfaces/gpunufft.py @@ -614,7 +614,7 @@ def pipe( shape=original_shape, density=density_comp, backend=cls.backend, - **kwargs + **kwargs, ) return density_comp.squeeze() From 119ee92cf5bca65779987674686462984d33ca16 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Dec 2024 14:52:51 +0100 Subject: [PATCH 32/40] Fix --- src/mrinufft/density/utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index 1a774dfc..b11b696f 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -56,8 +56,12 @@ def normalize_weights(weights): def normalize_density(kspace_loc, shape, density, backend, **kwargs): """Normalize the density to ensure that the reconstruction is stable.""" from mrinufft import get_operator - + xp = np + if backend == "cufinufft": + import cupy as cp + xp = cp test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) - test_im = np.ones(shape, dtype=test_op.cpx_dtype) + test_im = xp.ones(shape, dtype=test_op.cpx_dtype) test_im_recon = test_op.adj_op(density * test_op.op(test_im)) - density /= np.mean(np.abs(test_im_recon)) + density /= xp.mean(xp.abs(test_im_recon)) + return density From ba53be942638f4fd62a87732af922cb724200409 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Fri, 20 Dec 2024 16:20:59 +0100 Subject: [PATCH 33/40] Some fixes and updates --- src/mrinufft/density/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index b11b696f..e1992fed 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -56,9 +56,11 @@ def normalize_weights(weights): def normalize_density(kspace_loc, shape, density, backend, **kwargs): """Normalize the density to ensure that the reconstruction is stable.""" from mrinufft import get_operator + xp = np if backend == "cufinufft": import cupy as cp + xp = cp test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) test_im = xp.ones(shape, dtype=test_op.cpx_dtype) From 3ef45a6319c077b86733d48fe4e45664d49f009f Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 2 Jan 2025 09:12:21 +0100 Subject: [PATCH 34/40] Update test-ci.yml --- .github/workflows/test-ci.yml | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 22f833a6..a6ec31a3 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -98,9 +98,9 @@ jobs: shell: bash run: | source $RUNNER_WORKSPACE/venv/bin/activate - export CUDA_BIN_PATH=/usr/local/cuda-12.1/ - export PATH=/usr/local/cuda-12.1/bin/:${PATH} - export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} + export CUDA_BIN_PATH=/usr/local/cuda-12.4/ + export PATH=/usr/local/cuda-12.4/bin/:${PATH} + export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH} pip install -e .[${{ matrix.backend }},autodiff] - name: Run Tests @@ -169,11 +169,11 @@ jobs: python -m pip install -e .[extra,test,dev] python -m pip install finufft pooch brainweb-dl torch fastmri - - name: Point to CUDA 12.1 #TODO: This can be combined from other jobs + - name: Point to CUDA 12.4 #TODO: This can be combined from other jobs run: | - export CUDA_BIN_PATH=/usr/local/cuda-12.1/ - export PATH=/usr/local/cuda-12.1/bin/:${PATH} - export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} + export CUDA_BIN_PATH=/usr/local/cuda-12.4/ + export PATH=/usr/local/cuda-12.4/bin/:${PATH} + export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH} - name: Install Python deps shell: bash @@ -277,11 +277,11 @@ jobs: with: python-version: "3.10" - - name: Point to CUDA 12.1 + - name: Point to CUDA 12.4 run: | - export CUDA_BIN_PATH=/usr/local/cuda-12.1/ - export PATH=/usr/local/cuda-12.1/bin/:${PATH} - export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH} + export CUDA_BIN_PATH=/usr/local/cuda-12.4/ + export PATH=/usr/local/cuda-12.4/bin/:${PATH} + export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH} - name: Install dependencies shell: bash -l {0} From 7989caa2495391f47b8db410d61a44aabe389744 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 2 Jan 2025 09:47:10 +0100 Subject: [PATCH 35/40] Update test-ci.yml --- .github/workflows/test-ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index a6ec31a3..6e6cfa77 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -122,11 +122,11 @@ jobs: if: always() shell: bash run: | - cd $RUNNER_WORKSPACE + #cd $RUNNER_WORKSPACE ls -al - rm -rf finufft - rm -rf gpuNUFFT - rm -rf venv + #rm -rf finufft + #rm -rf gpuNUFFT + #rm -rf venv get-commit-message: runs-on: ubuntu-latest From f29a0201311e3e60892f8c88bc6b8b1fa14c7460 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 2 Jan 2025 10:36:45 +0100 Subject: [PATCH 36/40] Updates --- .github/workflows/test-ci.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 22f833a6..066b8307 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -46,6 +46,10 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install -e .[test,${{ env.ref_backend }},${{ matrix.backend }}] + if [[ ${{ matrix.backend }} == "finufft" ]]; then + # Install torch for testing autodiff + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + fi - name: Install BART if needed if: ${{ matrix.backend == 'bart' || env.ref_backend == 'bart'}} @@ -91,7 +95,7 @@ jobs: python -m venv venv source $RUNNER_WORKSPACE/venv/bin/activate pip install --upgrade pip wheel - pip install -e mri-nufft[test] + pip install -e mri-nufft[test,${{ env.ref_backend }}] - name: Install backend From 477349eba3e1ed49fdb8f094c4ef566690a73a60 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 2 Jan 2025 10:37:21 +0100 Subject: [PATCH 37/40] Make more fixes --- .github/workflows/test-ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 1058c78b..88fdbc79 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -126,11 +126,11 @@ jobs: if: always() shell: bash run: | - #cd $RUNNER_WORKSPACE + cd $RUNNER_WORKSPACE ls -al - #rm -rf finufft - #rm -rf gpuNUFFT - #rm -rf venv + rm -rf finufft + rm -rf gpuNUFFT + rm -rf venv get-commit-message: runs-on: ubuntu-latest From 5f21f6ebab0653699dc46d24ed52dca961af2736 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Thu, 2 Jan 2025 10:49:29 +0100 Subject: [PATCH 38/40] Careful repro test --- .github/workflows/test-ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 88fdbc79..330f5d97 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -127,10 +127,10 @@ jobs: shell: bash run: | cd $RUNNER_WORKSPACE - ls -al - rm -rf finufft - rm -rf gpuNUFFT - rm -rf venv + # ls -al + # rm -rf finufft + # rm -rf gpuNUFFT + # rm -rf venv get-commit-message: runs-on: ubuntu-latest From 134f70a60e16a557caa086ce16e160fdbcd78d09 Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Mon, 6 Jan 2025 13:31:23 +0100 Subject: [PATCH 39/40] Fixes --- .github/workflows/test-ci.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 330f5d97..2f9e0d43 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -105,7 +105,7 @@ jobs: export CUDA_BIN_PATH=/usr/local/cuda-12.4/ export PATH=/usr/local/cuda-12.4/bin/:${PATH} export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH} - pip install -e .[${{ matrix.backend }},autodiff] + pip install --no-binary cufinufft -e .[${{ matrix.backend }},autodiff] - name: Run Tests shell: bash @@ -183,7 +183,7 @@ jobs: shell: bash run: | python -m pip install --upgrade pip - python -m pip install -e .[test,dev,extra,finufft,cufinufft,gpuNUFFT,sigpy,autodiff,doc] + python -m pip install -e --no-binary cufinufft .[test,dev,extra,finufft,cufinufft,gpuNUFFT,sigpy,autodiff,doc] - name: Run examples shell: bash @@ -291,7 +291,7 @@ jobs: shell: bash -l {0} run: | python -m pip install --upgrade pip - python -m pip install .[doc,finufft,autodiff,gpunufft,cufinufft] + python -m pip install --no-binary cufinufft .[doc,finufft,autodiff,gpunufft,cufinufft] - name: Build API documentation From 35ca5b8850c2442f83f6f09e4be3755699ee4b1a Mon Sep 17 00:00:00 2001 From: Chaithya G R Date: Mon, 6 Jan 2025 13:58:08 +0100 Subject: [PATCH 40/40] Fixes --- .github/workflows/test-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 2f9e0d43..9984fb3a 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -183,7 +183,7 @@ jobs: shell: bash run: | python -m pip install --upgrade pip - python -m pip install -e --no-binary cufinufft .[test,dev,extra,finufft,cufinufft,gpuNUFFT,sigpy,autodiff,doc] + python -m pip install --no-binary cufinufft -e .[test,dev,extra,finufft,cufinufft,gpuNUFFT,sigpy,autodiff,doc] - name: Run examples shell: bash