From 24effcf1691a6960eb5634002cceb853085f6c65 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 24 Jul 2024 14:31:15 -0400 Subject: [PATCH 1/2] minimal patch to support cupy install and disabled cufinufft --- src/aspire/nufft/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/nufft/__init__.py b/src/aspire/nufft/__init__.py index 07d92c736c..d38a0df96e 100644 --- a/src/aspire/nufft/__init__.py +++ b/src/aspire/nufft/__init__.py @@ -198,7 +198,7 @@ def anufft(sig_f, fourier_pts, sz, real=False, epsilon=1e-8): adjoint = adjoint.real if real else adjoint - if cp and not _keep_on_gpu: + if cp and isinstance(adjoint, cp.ndarray) and not _keep_on_gpu: adjoint = adjoint.get() return adjoint @@ -259,7 +259,7 @@ def nufft(sig_f, fourier_pts, real=False, epsilon=1e-8): transform = transform.real if real else transform - if cp and not _keep_on_gpu: + if cp and isinstance(transform, cp.ndarray) and not _keep_on_gpu: transform = transform.get() return transform From 456aa22b037abc4f089acc55d2dc04f5148c6673 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 24 Jul 2024 15:53:32 -0400 Subject: [PATCH 2/2] skip enormous FFB2D test on GPU --- tests/test_FFBbasis2D.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_FFBbasis2D.py b/tests/test_FFBbasis2D.py index 8acf7201d1..c3ee42dd75 100644 --- a/tests/test_FFBbasis2D.py +++ b/tests/test_FFBbasis2D.py @@ -6,6 +6,7 @@ from scipy.special import jv from aspire.basis import Coef, FFBBasis2D +from aspire.nufft import all_backends from aspire.source import Simulation from aspire.utils.misc import grid_2d from aspire.volume import Volume @@ -126,6 +127,9 @@ def testShift(self, basis): params = [pytest.param(512, np.float32, marks=pytest.mark.expensive)] +@pytest.mark.skipif( + all_backends()[0] == "cufinufft", reason="Not enough memory to run via GPU" +) @pytest.mark.parametrize( "L, dtype", params, @@ -136,6 +140,7 @@ def testHighResFFBBasis2D(L, dtype): sim = Simulation( n=1, L=L, + C=1, dtype=dtype, amplitudes=1, offsets=0, @@ -149,4 +154,6 @@ def testHighResFFBBasis2D(L, dtype): # Mask to compare inside disk of radius 1. mask = grid_2d(L, normalized=True)["r"] < 1 - assert np.allclose(im_ffb.asnumpy()[0][mask], im.asnumpy()[0][mask], atol=1e-4) + np.testing.assert_allclose( + im_ffb.asnumpy()[0][mask], im.asnumpy()[0][mask], rtol=1e-05, atol=1e-4 + )