API: Wrap around negative workers from cpu_count()+1

scipy · Aug 8, 2019 · a620d98 · a620d98
1 parent efc6991
commit a620d98
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 58 deletions.
diff --git a/scipy/fft/_basic.py b/scipy/fft/_basic.py
@@ -49,8 +49,9 @@ def fft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See the notes below for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or -1 for
-        unlimited. Must be -1 or >0.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -193,8 +194,9 @@ def ifft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -277,8 +279,9 @@ def rfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -372,8 +375,9 @@ def irfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -456,8 +460,9 @@ def hfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See `fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -529,8 +534,9 @@ def ihfft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See `fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -594,8 +600,9 @@ def fftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -699,8 +706,9 @@ def ifftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -787,8 +795,9 @@ def fft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -886,8 +895,9 @@ def ifft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -970,8 +980,9 @@ def rfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1048,8 +1059,9 @@ def rfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1110,8 +1122,9 @@ def irfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1191,8 +1204,9 @@ def irfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False,
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1248,8 +1262,9 @@ def hfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1337,8 +1352,9 @@ def hfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See `fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1391,8 +1407,9 @@ def ihfftn(x, s=None, axes=None, norm=None, overwrite_x=False, workers=None):
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -1467,8 +1484,9 @@ def ihfft2(x, s=None, axes=(-2, -1), norm=None, overwrite_x=False,
         If True, the contents of `x` can be destroyed; the default is False.
         See :func:`fft` for more details.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------

diff --git a/scipy/fft/_pocketfft/helper.py b/scipy/fft/_pocketfft/helper.py
@@ -4,10 +4,12 @@
 from .pypocketfft import good_size
 import operator
 import sys
+import os
 
 
 # TODO: Add configuration support
 _default_workers = 1
+_cpu_count = os.cpu_count()
 
 
 def _iterable_of_int(x, name=None):
@@ -156,10 +158,13 @@ def _workers(workers):
     if workers is None:
         return _default_workers
 
-    if workers == 0 or workers < -1:
-        raise ValueError("workers must be > 0, or -1")
+    if workers < 0:
+        workers += 1 + _cpu_count
 
-    return workers if workers != -1 else 0
+    if workers <= 0:
+        raise ValueError("workers value out of range")
+
+    return workers
 
 
 def next_fast_len(target, kind='C2C'):

diff --git a/scipy/fft/_realtransforms.py b/scipy/fft/_realtransforms.py
@@ -34,8 +34,9 @@ def dctn(x, type=2, s=None, axes=None, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -91,8 +92,9 @@ def idctn(x, type=2, s=None, axes=None, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -148,8 +150,9 @@ def dstn(x, type=2, s=None, axes=None, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -205,8 +208,9 @@ def idstn(x, type=2, s=None, axes=None, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -256,8 +260,9 @@ def dct(x, type=2, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -406,8 +411,9 @@ def idct(x, type=2, n=None, axis=-1, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -468,8 +474,9 @@ def dst(x, type=2, n=None, axis=-1, norm=None, overwrite_x=False, workers=None):
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------
@@ -585,8 +592,9 @@ def idst(x, type=2, n=None, axis=-1, norm=None, overwrite_x=False,
     overwrite_x : bool, optional
         If True, the contents of `x` can be destroyed; the default is False.
     workers : int, optional
-        Maximum number of workers to use for parallel computation, or 0 for
-        unlimited. Must be -1 or >0. See `fft` for more details.
+        Maximum number of workers to use for parallel computation. If negative,
+        the value wraps around from ``os.cpu_count()``. See `fft` for more
+        details.
 
     Returns
     -------

diff --git a/scipy/fft/tests/test_multithreading.py b/scipy/fft/tests/test_multithreading.py
@@ -41,9 +41,13 @@ def test_mixed_threads_processes(x):
     fft.fft(x, workers=2)
 
 def test_invalid_workers(x):
+    import os
+    cpus = os.cpu_count()
 
-    with pytest.raises(ValueError, match='workers must be'):
+    fft.ifft([1], workers=-cpus)
+
+    with pytest.raises(ValueError, match='workers value out of range'):
         fft.fft(x, workers=0)
 
-    with pytest.raises(ValueError, match='workers must be'):
-        fft.ifft(x, workers=-2)
+    with pytest.raises(ValueError, match='workers value out of range'):
+        fft.ifft(x, workers=-cpus-1)