From 943b0a74b00ef2e5cfce8db309351ad7d209f56f Mon Sep 17 00:00:00 2001 From: Sepand KASHANI Date: Tue, 2 Apr 2024 18:38:24 +0200 Subject: [PATCH] [BugFix] Robustize cuFINUFFT Python guru interface. As per the docs: * finufft.Plan(n_modes_or_dim) accepts (int, tuple[int]) * cufinufft.Plan(n_modes) accepts (int, tuple[int]) In practice `n_modes_or_dim` is transformed internally to a NumPy array, hence providing any iterable to finufft.Plan() works. However `n_modes` in cufinufft.Plan() is more restrictive: it works with tuple[int], fails planning if list[int], and wrongfully passes planning if ndarray[int], but then seg-faults at execute(). This commit increases the robustness of cufinufft.Plan() to match the flexibility of the CPU interface. --- python/cufinufft/cufinufft/_plan.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/cufinufft/cufinufft/_plan.py b/python/cufinufft/cufinufft/_plan.py index 70973ba00..263c12cd7 100644 --- a/python/cufinufft/cufinufft/_plan.py +++ b/python/cufinufft/cufinufft/_plan.py @@ -5,6 +5,8 @@ """ import atexit +import collections.abc +import numbers import sys import warnings @@ -108,8 +110,12 @@ def __init__(self, nufft_type, n_modes, n_trans=1, eps=1e-6, isign=None, else: raise TypeError("Expected complex64 or complex128.") - if isinstance(n_modes, int): + if isinstance(n_modes, numbers.Integral): n_modes = (n_modes,) + elif isinstance(n_modes, collections.abc.Iterable): + n_modes = tuple(n_modes) + else: + raise ValueError(f"Invalid n_modes '{n_modes}'") self.dim = len(n_modes) self.type = nufft_type