Explicitly set numba config for each jit function (#221)

Using `parallel=True` on the Green's function was raising a warning
because there is nothing to parallelize. Set arguments individually to
avoid this. Add missing `numba.prange` to Spline Jacobian calculation to
actually run it in parallel. A combination of Dask + numba makes some 
tests hang on CI randomly even if passing locally.Run the tests that use 
dask with the numpy engine to avoid this.

Author: leouieda

.azure-pipelines.yml

@@ -94,16 +94,12 @@ jobs:
         python.version: '3.7'
         PYTHON: '3.7'
         CONDA_INSTALL_EXTRA: "codecov pykdtree numba"
-        NUMBA_NUM_THREADS: 1
-        OMP_NUM_THREADS: 1
-        MKL_NUM_THREADS: 1
+        DASK_SCHEDULER: "synchronous"
       Python36-Optional:
         python.version: '3.6'
         PYTHON: '3.6'
         CONDA_INSTALL_EXTRA: "codecov pykdtree numba"
-        NUMBA_NUM_THREADS: 1
-        OMP_NUM_THREADS: 1
-        MKL_NUM_THREADS: 1
+        DASK_SCHEDULER: "synchronous"
 
   steps:
 
@@ -203,16 +199,12 @@ jobs:
         python.version: '3.7'
         PYTHON: '3.7'
         CONDA_INSTALL_EXTRA: "codecov pykdtree numba"
-        NUMBA_NUM_THREADS: 1
-        OMP_NUM_THREADS: 1
-        MKL_NUM_THREADS: 1
+        DASK_SCHEDULER: "synchronous"
       Python36-Optional:
         python.version: '3.6'
         PYTHON: '3.6'
         CONDA_INSTALL_EXTRA: "codecov pykdtree numba"
-        NUMBA_NUM_THREADS: 1
-        OMP_NUM_THREADS: 1
-        MKL_NUM_THREADS: 1
+        DASK_SCHEDULER: "synchronous"
 
   steps:
 

.travis.yml

@@ -56,17 +56,13 @@ matrix:
           os: linux
           env:
               - PYTHON=3.7
-              - NUMBA_NUM_THREADS=1
-              - OMP_NUM_THREADS=1
-              - MKL_NUM_THREADS=1
+              - DASK_SCHEDULER=synchronous
               - CONDA_INSTALL_EXTRA="codecov pykdtree numba"
         - name: "Linux - Python 3.6 (optional deps)"
           os: linux
           env:
               - PYTHON=3.6
-              - NUMBA_NUM_THREADS=1
-              - OMP_NUM_THREADS=1
-              - MKL_NUM_THREADS=1
+              - DASK_SCHEDULER=synchronous
               - CONDA_INSTALL_EXTRA="codecov pykdtree numba"
 
 # Setup the build environment

verde/spline.py

@@ -23,9 +23,6 @@
 # Otherwise, DeprecationWarning won't be shown, kind of defeating the purpose.
 warnings.simplefilter("default")
 
-# Default arguments for numba.jit
-JIT_ARGS = dict(nopython=True, target="cpu", fastmath=True, parallel=True)
-
 
 class SplineCV(BaseGridder):
     r"""
@@ -538,7 +535,7 @@ def jacobian_numpy(east, north, force_east, force_north, mindist, jac):
     return jac
 
 
-@jit(**JIT_ARGS)
+@jit(nopython=True, fastmath=True, parallel=True)
 def predict_numba(east, north, force_east, force_north, mindist, forces, result):
     "Calculate the predicted data using numba to speed things up."
     for i in numba.prange(east.size):  # pylint: disable=not-an-iterable
@@ -551,10 +548,10 @@ def predict_numba(east, north, force_east, force_north, mindist, forces, result)
     return result
 
 
-@jit(**JIT_ARGS)
+@jit(nopython=True, fastmath=True, parallel=True)
 def jacobian_numba(east, north, force_east, force_north, mindist, jac):
     "Calculate the Jacobian matrix using numba to speed things up."
-    for i in range(east.size):
+    for i in numba.prange(east.size):  # pylint: disable=not-an-iterable
         for j in range(force_east.size):
             jac[i, j] = GREENS_FUNC_JIT(
                 east[i] - force_east[j], north[i] - force_north[j], mindist
@@ -563,4 +560,4 @@ def jacobian_numba(east, north, force_east, force_north, mindist, jac):
 
 
 # Jit compile the Green's functions for use in the numba functions
-GREENS_FUNC_JIT = jit(**JIT_ARGS)(greens_func)
+GREENS_FUNC_JIT = jit(nopython=True, fastmath=True)(greens_func)

verde/tests/test_spline.py

@@ -15,11 +15,15 @@
 
 
 @pytest.mark.parametrize(
-    "delayed,client",
-    [(False, None), (True, None), (False, Client(processes=False))],
+    "delayed,client,engine",
+    [
+        (False, None, "auto"),
+        (True, None, "numpy"),
+        (False, Client(processes=False), "numpy"),
+    ],
     ids=["serial", "delayed", "distributed"],
 )
-def test_spline_cv(delayed, client):
+def test_spline_cv(delayed, client, engine):
     "See if the cross-validated spline solution matches the synthetic data"
     region = (100, 500, -800, -700)
     synth = CheckerBoard(region=region)
@@ -33,6 +37,7 @@ def test_spline_cv(delayed, client):
         cv=ShuffleSplit(n_splits=1, random_state=0),
         delayed=delayed,
         client=client,
+        engine=engine,
     ).fit(coords, data.scalars)
     if client is not None:
         client.close()

verde/vector.py

@@ -22,9 +22,6 @@
 # Otherwise, DeprecationWarning won't be shown, kind of defeating the purpose.
 warnings.simplefilter("default")
 
-# Default arguments for numba.jit
-JIT_ARGS = dict(nopython=True, target="cpu", fastmath=True, parallel=True)
-
 
 class Vector(BaseGridder):
     """
@@ -444,7 +441,7 @@ def jacobian_2d_numpy(east, north, force_east, force_north, mindist, poisson, ja
     return jac
 
 
-@jit(**JIT_ARGS)
+@jit(nopython=True, fastmath=True, parallel=True)
 def predict_2d_numba(
     east, north, force_east, force_north, mindist, poisson, forces, vec_east, vec_north
 ):
@@ -462,7 +459,7 @@ def predict_2d_numba(
     return vec_east, vec_north
 
 
-@jit(**JIT_ARGS)
+@jit(nopython=True, fastmath=True, parallel=True)
 def jacobian_2d_numba(east, north, force_east, force_north, mindist, poisson, jac):
     "Calculate the Jacobian matrix using numba to speed things up."
     nforces = force_east.size
@@ -480,4 +477,4 @@ def jacobian_2d_numba(east, north, force_east, force_north, mindist, poisson, ja
 
 
 # JIT compile the Greens functions for use in numba functions
-GREENS_FUNC_2D_JIT = jit(**JIT_ARGS)(greens_func_2d)
+GREENS_FUNC_2D_JIT = jit(nopython=True, fastmath=True)(greens_func_2d)