Skip to content

Commit

Permalink
[BugFix] Robustize cuFINUFFT Python guru interface.
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
SepandKashani committed Apr 2, 2024
1 parent 9744d14 commit 943b0a7
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion python/cufinufft/cufinufft/_plan.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
"""

import atexit
import collections.abc
import numbers
import sys
import warnings

Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 943b0a7

Please sign in to comment.