Repository: fatiando/verde
Pull request #110 (2018-08-09 00:40:31)

Use numba to optimize Sline.jacobian
Calculating the Jacobian using numba with parallel `for` loops results 
in 10x speedup compared to numpy. This translates to a 40-45% 
speedup in the whole process.
Keep numba as optional and only use it when installed or requested.
Refactor some of the underlying spline code to minimize repetition.
Move least-squares fitting to a separate function to be shared by all
classes.
Move methods shared between `Spline` and `VectorSpline2D` to 
functions and remove inheritance. 

TODO:
- [x] Delete notebook
- [x] Document `engine` argument
- [x] Add numba to install docs


Patch:
diff --git a/.travis.yml b/.travis.yml
index 9ad3b1e..c785c6f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -32,6 +32,7 @@ env:
         - LINKCHECK=false
         - DEPLOY_DOCS=false
         - DEPLOY_PYPI=false
+        - EXTRAS=false
 
 # Specify the build configurations. Be sure to only deploy from a single build.
 matrix:
@@ -48,6 +49,11 @@ matrix:
           env:
               - PYTHON=3.5
               - BUILD_DOCS=true
+        - os: linux
+          env:
+              - PYTHON=3.6
+              - COVERAGE=true
+              - EXTRAS=true
         - os: osx
           env:
               - PYTHON=3.6
@@ -74,6 +80,10 @@ before_install:
     - if [ "$CHECK_FORMAT" == "true" ]; then
           conda install --yes --channel conda-forge black python=$PYTHON;
       fi
+    # Install extra dependencies
+    - if [ "$EXTRAS" == "true" ]; then
+          conda install --yes --channel conda-forge pykdtree numba python=$PYTHON;
+      fi
     # Show installed pkg information for postmortem diagnostic
     - conda list
 
diff --git a/doc/install.rst b/doc/install.rst
index 0a770e4..3266a19 100644
--- a/doc/install.rst
+++ b/doc/install.rst
@@ -25,7 +25,15 @@ Dependencies
 * `xarray <http://xarray.pydata.org/>`__
 * `scikit-learn <http://scikit-learn.org/>`__
 * `pooch <http://www.fatiando.org/pooch/>`__
-* `pykdtree <https://github.com/storpipfugl/pykdtree>`__ (optional)
+
+The following are optional dependencies that can make some parts of the code faster if
+they are installed:
+
+* `pykdtree <https://github.com/storpipfugl/pykdtree>`__: replaces
+  :class:`scipy.spatial.cKDTree` in :class:`verde.BlockReduce` and
+  :func:`verde.distance_mask` for better performance.
+* `numba <https://numba.pydata.org/>`__: replaces numpy calculations of Jacobian
+  matrices in splines with multi-threaded versions.
 
 Most of the examples in the :ref:`gallery` and :ref:`tutorials` also use:
 
@@ -34,11 +42,6 @@ Most of the examples in the :ref:`gallery` and :ref:`tutorials` also use:
 * `pyproj <https://jswhit.github.io/pyproj/>`__ for cartographic projections
 * `dask <https://dask.pydata.org/>`__ for parallelism
 
-.. note::
-
-    PyKDTree is an optional dependency and will be used instead of
-    :class:`scipy.spatial.cKDTree` if available for better performance.
-
 
 Installing with conda
 ---------------------
diff --git a/environment.yml b/environment.yml
index 7e9db1e..6367dd0 100644
--- a/environment.yml
+++ b/environment.yml
@@ -12,6 +12,9 @@ dependencies:
     - scikit-learn
     - pip:
         - pooch
+    # Optional
+    - pykdtree
+    - numba
     # Development requirements
     - dask
     - pyproj
diff --git a/examples/trend.py b/examples/trend.py
index 9d7f110..7840d60 100644
--- a/examples/trend.py
+++ b/examples/trend.py
@@ -24,7 +24,7 @@ print("\nTrend estimator:", trend)
 
 # Add the estimated trend and the residual data to the DataFrame
 data["trend"] = trend.predict(coordinates)
-data["residual"] = trend.residual_
+data["residual"] = data.total_field_anomaly_nt - data.trend
 print("\nUpdated DataFrame:")
 print(data.head())
 
diff --git a/verde/base.py b/verde/base.py
index fc75309..92c60e7 100644
--- a/verde/base.py
+++ b/verde/base.py
@@ -1,11 +1,15 @@
 """
 Base classes for all gridders.
 """
+from warnings import warn
+
 import xarray as xr
 import pandas as pd
 import numpy as np
 from sklearn.base import BaseEstimator
 from sklearn.metrics import r2_score
+from sklearn.preprocessing import StandardScaler
+from sklearn.linear_model import LinearRegression, Ridge
 
 from .coordinates import grid_coordinates, profile_coordinates, scatter_points
 
@@ -620,3 +624,26 @@ def check_fit_input(coordinates, data, weights, unpack=True):
         if len(data) == 1:
             data = data[0]
     return coordinates, data, weights
+
+
+def least_squares(jacobian, data, weights, damping=None):
+    """
+    Estimate forces that fit the data using least-squares. Scales the
+    Jacobian matrix to have unit standard deviation. This helps balance the
+    regularization and the difference between forces.
+    """
+    if jacobian.shape[0] < jacobian.shape[1]:
+        warn(
+            "Under-determined problem detected (ndata, nparams)={}.".format(
+                jacobian.shape
+            )
+        )
+    scaler = StandardScaler(copy=False, with_mean=False, with_std=True)
+    jacobian = scaler.fit_transform(jacobian)
+    if damping is None:
+        regr = LinearRegression(fit_intercept=False, normalize=False)
+    else:
+        regr = Ridge(alpha=damping, fit_intercept=False, normalize=False)
+    regr.fit(jacobian, data.ravel(), sample_weight=weights)
+    params = regr.coef_ / scaler.scale_
+    return params
diff --git a/verde/spline.py b/verde/spline.py
index 5bd6ee9..1a4a008 100644
--- a/verde/spline.py
+++ b/verde/spline.py
@@ -5,11 +5,17 @@ from warnings import warn
 
 import numpy as np
 from sklearn.utils.validation import check_is_fitted
-from sklearn.preprocessing import StandardScaler
-from sklearn.linear_model import LinearRegression, Ridge
 
-from .base import BaseGridder, check_fit_input
+from .base import BaseGridder, check_fit_input, least_squares
 from .coordinates import grid_coordinates, get_region
+from .utils import n_1d_arrays, parse_engine
+
+try:
+    import numba
+    from numba import jit
+except ImportError:
+    numba = None
+    from .utils import dummy_jit as jit
 
 
 class Spline(BaseGridder):
@@ -70,6 +76,11 @@ class Spline(BaseGridder):
         If not None, then the boundaries (``[W, E, S, N]``) used to generate a regular
         grid of forces. If None is given, then will use the boundaries of data given to
         the first call to :meth:`~verde.Spline.fit`.
+    engine : str
+        Computation engine for the Jacobian matrix. Can be ``'auto'``, ``'numba'``, or
+        ``'numpy'``. If ``'auto'``, will use numba if it is installed or numpy
+        otherwise. The numba version is multi-threaded and considerably faster, which
+        makes fitting and predicting faster.
 
     Attributes
     ----------
@@ -85,13 +96,20 @@ class Spline(BaseGridder):
     """
 
     def __init__(
-        self, mindist=1e-5, damping=None, shape=None, spacing=None, region=None
+        self,
+        mindist=1e-5,
+        damping=None,
+        shape=None,
+        spacing=None,
+        region=None,
+        engine="auto",
     ):
         self.damping = damping
         self.shape = shape
         self.spacing = spacing
         self.mindist = mindist
         self.region = region
+        self.engine = engine
 
     def fit(self, coordinates, data, weights=None):
         """
@@ -121,78 +139,15 @@ class Spline(BaseGridder):
             Returns this estimator instance for chaining operations.
 
         """
-        # Remove pre-existing force coordinates when fitting for a second time.
-        if hasattr(self, "force_coords_"):
-            del self.force_coords_
         coordinates, data, weights = check_fit_input(coordinates, data, weights)
-        self._check_weighted_exact_solution(weights)
+        warn_weighted_exact_solution(self, weights)
         # Capture the data region to use as a default when gridding.
         self.region_ = get_region(coordinates[:2])
-        jacobian = self.jacobian(coordinates[:2])
-        self.force_ = self._estimate_forces(jacobian, data, weights)
+        self.force_coords_ = get_force_coordinates(self, coordinates)
+        jacobian = self.jacobian(coordinates[:2], self.force_coords_)
+        self.force_ = least_squares(jacobian, data, weights, self.damping)
         return self
 
-    def _estimate_forces(self, jacobian, data, weights):
-        """
-        Estimate forces that fit the data using least-squares. Scales the
-        Jacobian matrix to have unit standard deviation. This helps balance the
-        regularization and the difference between forces.
-        """
-        if jacobian.shape[0] < jacobian.shape[1]:
-            warn(
-                " ".join(
-                    [
-                        "Under-determined problem detected",
-                        "(ndata, nparams)={}.".format(jacobian.shape),
-                        "Configuration of forces:",
-                        "spacing={} shape={}".format(self.spacing, self.shape),
-                    ]
-                )
-            )
-        scaler = StandardScaler(copy=False, with_mean=False, with_std=True)
-        jacobian = scaler.fit_transform(jacobian)
-        if self.damping is None:
-            regr = LinearRegression(fit_intercept=False, normalize=False)
-        else:
-            regr = Ridge(alpha=self.damping, fit_intercept=False, normalize=False)
-        regr.fit(jacobian, data.ravel(), sample_weight=weights)
-        # Undo the scaling so that we can use forces on the unscaled Jacobian
-        # later on.
-        forces = regr.coef_ / scaler.scale_
-        return forces
-
-    def _get_force_coordinates(self, coordinates):
-        """
-        Generate arrays for the force coordinates depending on the
-        configuration of the spline.
-        """
-        # Set the force positions. If no shape or spacing are given, then they
-        # will coincide with the data points. This will only happen once so the
-        # model won't change during later fits.
-        if self.shape is None and self.spacing is None:
-            coords = tuple(i.copy() for i in coordinates[:2])
-        else:
-            if self.region is None:
-                self.region = get_region(coordinates)
-            coords = grid_coordinates(
-                self.region, shape=self.shape, spacing=self.spacing
-            )
-        return coords
-
-    def _check_weighted_exact_solution(self, weights):
-        """
-        If a weighted exact solution is requested, warn the user that it won't
-        work.
-        """
-        # Check if we're using weights with the exact solution and warn the
-        # user that it won't have any effect.
-        exact = all(i is None for i in [self.damping, self.spacing, self.shape])
-        if weights is not None and exact:
-            warn(
-                "Weights are ignored for the exact solution. "
-                "Use damping or specify a shape/spacing for the forces."
-            )
-
     def predict(self, coordinates):
         """
         Evaluate the estimated spline on the given set of points.
@@ -214,11 +169,11 @@ class Spline(BaseGridder):
 
         """
         check_is_fitted(self, ["force_", "force_coords_"])
-        jac = self.jacobian(coordinates[:2])
+        jac = self.jacobian(coordinates[:2], self.force_coords_)
         shape = np.broadcast(*coordinates[:2]).shape
         return jac.dot(self.force_).reshape(shape)
 
-    def jacobian(self, coordinates, dtype="float64"):
+    def jacobian(self, coordinates, force_coordinates, dtype="float64"):
         """
         Make the Jacobian matrix for the 2D biharmonic spline.
 
@@ -231,6 +186,9 @@ class Spline(BaseGridder):
             Arrays with the coordinates of each data point. Should be in the
             following order: (easting, northing, vertical, ...). Only easting and
             northing will be used, all subsequent coordinates will be ignored.
+        force_coordinates : tuple of arrays
+            Arrays with the coordinates for the forces. Should be in the same order as
+            the coordinate arrays.
         dtype : str or numpy dtype
             The type of the Jacobian array.
 
@@ -240,20 +198,113 @@ class Spline(BaseGridder):
             The (n_data, n_forces) Jacobian matrix.
 
         """
-        if not hasattr(self, "force_coords_"):
-            self.force_coords_ = self._get_force_coordinates(coordinates)
-        force_easting, force_northing = (
-            np.atleast_1d(i).ravel() for i in self.force_coords_[:2]
+        force_east, force_north = n_1d_arrays(force_coordinates, n=2)
+        east, north = n_1d_arrays(coordinates, n=2)
+        if parse_engine(self.engine) == "numba":
+            jac = np.empty((east.size, force_east.size), dtype=dtype)
+            jacobian_numba(east, north, force_east, force_north, self.mindist, jac)
+        else:
+            jac = jacobian_numpy(
+                east, north, force_east, force_north, self.mindist, dtype
+            )
+        return jac
+
+
+def warn_weighted_exact_solution(spline, weights):
+    """
+    Warn the user that a weights doesn't work for the exact solution.
+
+    Parameters
+    ----------
+    spline : estimator
+        The spline instance that we'll check. Needs to have ``shape``, ``spacing``, and
+        ``damping`` attributes.
+    weights : array or None
+        The weights given to fit.
+
+    """
+    # Check if we're using weights with the exact solution and warn the
+    # user that it won't have any effect.
+    exact = all(i is None for i in [spline.damping, spline.spacing, spline.shape])
+    if weights is not None and exact:
+        warn(
+            "Weights are ignored for the exact solution. "
+            "Use damping or specify a shape/spacing for the forces."
         )
-        easting, northing = (np.atleast_1d(i).ravel() for i in coordinates[:2])
-        # Reshaping the data to a column vector will automatically build a
-        # distance matrix between each data point and force.
-        distance = np.hypot(
-            easting.reshape((easting.size, 1)) - force_easting.ravel(),
-            northing.reshape((northing.size, 1)) - force_northing.ravel(),
-            dtype=dtype,
+
+
+def get_force_coordinates(spline, coordinates):
+    """
+    Make arrays for force coordinates depending on the spline configuration.
+
+    If no ``shape`` and ``spacing`` are given, then will use the data coordinates as
+    force coordinates. Otherwise, will generate a grid based on the given spacing/shape
+    and the data region.
+
+    Parameters
+    ----------
+    spline : estimator
+        The spline instance. Needs to have ``shape`` and ``spacing`` attributes. If
+        ``region`` is present, will use it to define the grid region. If not, will use
+        the data region.
+    coordinates : tuple of arrays
+        Arrays with the coordinates of each data point. Should be in the
+        following order: (easting, northing, vertical, ...). Only easting and
+        northing will be used, all subsequent coordinates will be ignored.
+
+    Returns
+    -------
+    force_coordinates : tuple of arrays
+        The coordinate arrays for the forces.
+
+    """
+    # Set the force positions. If no shape or spacing are given, then they
+    # will coincide with the data points. This will only happen once so the
+    # model won't change during later fits.
+    if spline.shape is None and spline.spacing is None:
+        coords = tuple(i.copy() for i in n_1d_arrays(coordinates, n=2))
+    else:
+        if spline.region is None:
+            region = get_region(coordinates)
+        else:
+            region = spline.region
+        coords = tuple(
+            i.ravel()
+            for i in grid_coordinates(
+                region, shape=spline.shape, spacing=spline.spacing
+            )
         )
-        # The mindist factor helps avoid singular matrices when the force and
-        # computation point are too close
-        distance += self.mindist
-        return (distance ** 2) * (np.log(distance) - 1)
+    return coords
+
+
+@jit(nopython=True, target="cpu", 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 numba.prange(east.size):  # pylint: disable=not-an-iterable
+        for j in range(force_east.size):
+            distance = np.sqrt(
+                (east[i] - force_east[j]) ** 2 + (north[i] - force_north[j]) ** 2
+            )
+            distance += mindist
+            jac[i, j] = (distance ** 2) * (np.log(distance) - 1)
+    return jac
+
+
+def jacobian_numpy(east, north, force_east, force_north, mindist, dtype):
+    """
+    Calculate the Jacobian using numpy broadcasting. Is slightly slower than the numba
+    version.
+    """
+    # Reshaping the data to a column vector will automatically build a
+    # distance matrix between each data point and force.
+    distance = np.hypot(
+        east.reshape((east.size, 1)) - force_east,
+        north.reshape((north.size, 1)) - force_north,
+        dtype=dtype,
+    )
+    # The mindist factor helps avoid singular matrices when the force and
+    # computation point are too close
+    distance += mindist
+    return (distance ** 2) * (np.log(distance) - 1)
diff --git a/verde/trend.py b/verde/trend.py
index 4989eae..bee89d2 100644
--- a/verde/trend.py
+++ b/verde/trend.py
@@ -3,10 +3,8 @@
 """
 import numpy as np
 from sklearn.utils.validation import check_is_fitted
-from sklearn.preprocessing import StandardScaler
-from sklearn.linear_model import LinearRegression
 
-from .base import BaseGridder, check_fit_input
+from .base import BaseGridder, check_fit_input, least_squares
 from .coordinates import get_region
 
 
@@ -34,9 +32,6 @@ class Trend(BaseGridder):
         The boundaries (``[W, E, S, N]``) of the data used to fit the
         interpolator. Used as the default region for the
         :meth:`~verde.Trend.grid` and :meth:`~verde.Trend.scatter` methods.
-    residual_ : array
-        The difference between the input data and the predicted polynomial
-        trend.
 
     Examples
     --------
@@ -50,8 +45,6 @@ class Trend(BaseGridder):
     >>> import numpy as np
     >>> np.allclose(trend.predict(coordinates), data)
     True
-    >>> np.allclose(trend.residual_, 0, atol=1e-5)
-    True
     >>> # Use weights to account for outliers
     >>> data_out = data.copy()
     >>> data_out[2, 2] += 500
@@ -60,7 +53,8 @@ class Trend(BaseGridder):
     >>> trend_out = Trend(degree=1).fit(coordinates, data_out, weights)
     >>> print(', '.join(['{:.1f}'.format(i) for i in trend_out.coef_]))
     10.0, 2.0, -0.4
-    >>> print('{:.2f}'.format(trend_out.residual_[2, 2]))
+    >>> residual = data_out - trend_out.predict(coordinates)
+    >>> print('{:.2f}'.format(residual[2, 2]))
     500.00
 
     """
@@ -99,12 +93,7 @@ class Trend(BaseGridder):
         coordinates, data, weights = check_fit_input(coordinates, data, weights)
         self.region_ = get_region(coordinates[:2])
         jac = self.jacobian(coordinates[:2], dtype=data.dtype)
-        scaler = StandardScaler(copy=False, with_mean=False, with_std=True)
-        jac = scaler.fit_transform(jac)
-        regr = LinearRegression(fit_intercept=False, normalize=False)
-        regr.fit(jac, data.ravel(), sample_weight=weights)
-        self.residual_ = data - jac.dot(regr.coef_).reshape(data.shape)
-        self.coef_ = regr.coef_ / scaler.scale_
+        self.coef_ = least_squares(jac, data, weights, damping=None)
         return self
 
     def predict(self, coordinates):
diff --git a/verde/utils.py b/verde/utils.py
index e53f97e..16b547e 100644
--- a/verde/utils.py
+++ b/verde/utils.py
@@ -1,11 +1,101 @@
 """
 General utilities.
 """
+import functools
+
 import numpy as np
 
 from .base import check_data
 
 
+def parse_engine(engine):
+    """
+    Choose the best engine available and check if it's valid.
+
+    Parameters
+    ----------
+    engine : str
+        The name of the engine. If ``"auto"`` will favor numba if it's available.
+
+    Returns
+    -------
+    engine : str
+        The name of the engine that should be used.
+
+    """
+    engines = {"auto", "numba", "numpy"}
+    if engine not in engines:
+        raise ValueError("Invalid engine '{}'. Must be in {}.".format(engine, engines))
+    if engine == "auto":
+        try:
+            import numba  # pylint: disable=unused-variable
+
+            return "numba"
+        except ImportError:
+            return "numpy"
+    return engine
+
+
+def dummy_jit(**kwargs):  # pylint: disable=unused-argument
+    """
+    Replace numba.jit if not installed with a function that raises RunTimeError.
+
+    Use as a decorator.
+
+    Parameters
+    ----------
+    function
+        A function that you would decorate with :func:`numba.jit`.
+
+    Returns
+    -------
+    function
+        A function that raises :class:`RunTimeError` warning that numba isn't installed.
+
+    """
+
+    def dummy_decorator(function):
+        "The actual decorator"
+
+        @functools.wraps(function)
+        def dummy_function(*args, **kwargs):  # pylint: disable=unused-argument
+            "Just raise an exception."
+            raise RuntimeError("Could not find numba.")
+
+        return dummy_function
+
+    return dummy_decorator
+
+
+def n_1d_arrays(arrays, n):
+    """
+    Get the first n elements from a tuple/list, make sure they are arrays, and ravel.
+
+    Parameters
+    ----------
+    arrays : tuple of arrays
+        The arrays. Can be lists or anything that can be converted to a numpy array
+        (including numpy arrays).
+    n : int
+        How many arrays to return.
+
+    Returns
+    -------
+    1darrays : tuple of arrays
+        The converted 1D numpy arrays.
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> arrays = [np.arange(4).reshape(2, 2)]*3
+    >>> n_1d_arrays(arrays, n=2)
+    (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
+
+    """
+    return tuple(np.atleast_1d(i).ravel() for i in arrays[:n])
+
+
 def variance_to_weights(variance, tol=1e-15, dtype="float64"):
     """
     Converts data variances to weights for gridding.
diff --git a/verde/vector.py b/verde/vector.py
index 5049d9b..6c989ca 100644
--- a/verde/vector.py
+++ b/verde/vector.py
@@ -5,12 +5,13 @@ Vector gridding using elasticity Green's functions from Sandwell and Wessel
 import numpy as np
 from sklearn.utils.validation import check_is_fitted
 
-from .base import check_fit_input
-from .spline import Spline
+from .base import check_fit_input, least_squares, BaseGridder
+from .spline import warn_weighted_exact_solution, get_force_coordinates
+from .utils import n_1d_arrays
 from .coordinates import get_region
 
 
-class VectorSpline2D(Spline):
+class VectorSpline2D(BaseGridder):
     r"""
     Elastically coupled interpolation of 2-component vector data.
 
@@ -94,13 +95,11 @@ class VectorSpline2D(Spline):
         region=None,
     ):
         self.poisson = poisson
-        super().__init__(
-            mindist=mindist,
-            damping=damping,
-            shape=shape,
-            spacing=spacing,
-            region=region,
-        )
+        self.damping = damping
+        self.shape = shape
+        self.spacing = spacing
+        self.mindist = mindist
+        self.region = region
 
     def fit(self, coordinates, data, weights=None):
         """
@@ -133,9 +132,6 @@ class VectorSpline2D(Spline):
             Returns this estimator instance for chaining operations.
 
         """
-        # Remove pre-existing force coordinates when fitting for a second time.
-        if hasattr(self, "force_coords_"):
-            del self.force_coords_
         coordinates, data, weights = check_fit_input(
             coordinates, data, weights, unpack=False
         )
@@ -149,10 +145,11 @@ class VectorSpline2D(Spline):
             weights = np.concatenate([i.ravel() for i in weights])
         else:
             weights = None
-        self._check_weighted_exact_solution(weights)
+        warn_weighted_exact_solution(self, weights)
         data = np.concatenate([i.ravel() for i in data])
-        jacobian = self.jacobian(coordinates[:2])
-        self.force_ = self._estimate_forces(jacobian, data, weights)
+        self.force_coords_ = get_force_coordinates(self, coordinates)
+        jacobian = self.jacobian(coordinates[:2], self.force_coords_)
+        self.force_ = least_squares(jacobian, data, weights, self.damping)
         return self
 
     def predict(self, coordinates):
@@ -177,13 +174,13 @@ class VectorSpline2D(Spline):
 
         """
         check_is_fitted(self, ["force_", "force_coords_"])
-        jac = self.jacobian(coordinates[:2])
+        jac = self.jacobian(coordinates[:2], self.force_coords_)
         cast = np.broadcast(*coordinates[:2])
         npoints = cast.size
         components = jac.dot(self.force_).reshape((2, npoints))
         return tuple(comp.reshape(cast.shape) for comp in components)
 
-    def jacobian(self, coordinates, dtype="float64"):
+    def jacobian(self, coordinates, force_coordinates, dtype="float64"):
         """
         Make the Jacobian matrix for the 2D coupled elastic deformation.
 
@@ -202,6 +199,9 @@ class VectorSpline2D(Spline):
             Arrays with the coordinates of each data point. Should be in the
             following order: (easting, northing, vertical, ...). Only easting and
             northing will be used, all subsequent coordinates will be ignored.
+        force_coordinates : tuple of arrays
+            Arrays with the coordinates for the forces. Should be in the same order as
+            the coordinate arrays.
         dtype : str or numpy dtype
             The type of the Jacobian array.
 
@@ -211,10 +211,8 @@ class VectorSpline2D(Spline):
             The (n_data*2, n_forces*2) Jacobian matrix.
 
         """
-        if not hasattr(self, "force_coords_"):
-            self.force_coords_ = self._get_force_coordinates(coordinates)
-        force_coordinates = [np.atleast_1d(i).ravel() for i in self.force_coords_[:2]]
-        coordinates = [np.atleast_1d(i).ravel() for i in coordinates[:2]]
+        force_coordinates = n_1d_arrays(force_coordinates, n=2)
+        coordinates = n_1d_arrays(coordinates, n=2)
         npoints = coordinates[0].size
         nforces = force_coordinates[0].size
         # Reshaping the data coordinates to a column vector will automatically


Test patch:
diff --git a/verde/tests/test_spline.py b/verde/tests/test_spline.py
index 8e97e73..fd09508 100644
--- a/verde/tests/test_spline.py
+++ b/verde/tests/test_spline.py
@@ -8,6 +8,8 @@ import numpy.testing as npt
 
 from ..spline import Spline
 from ..datasets.synthetic import CheckerBoard
+from ..utils import n_1d_arrays
+from .utils import requires_numba
 
 
 def test_spline():
@@ -37,13 +39,13 @@ def test_spline():
 def test_spline_twice():
     "Check that the forces are updated when fitting twice"
     grid = CheckerBoard(region=(1000, 5000, -8000, -6000)).grid(shape=(10, 10))
-    coords = np.meshgrid(grid.easting, grid.northing)
+    coords = n_1d_arrays(np.meshgrid(grid.easting, grid.northing), n=2)
     spline = Spline()
-    spline.fit(coords, grid.scalars.values)
+    spline.fit(coords, grid.scalars.values.ravel())
     npt.assert_allclose(spline.force_coords_, coords)
     grid2 = CheckerBoard(region=(-15, -5, 10, 20)).grid(shape=(10, 10))
-    coords2 = np.meshgrid(grid2.easting, grid2.northing)
-    spline.fit(coords2, grid2.scalars.values)
+    coords2 = n_1d_arrays(np.meshgrid(grid2.easting, grid2.northing), n=2)
+    spline.fit(coords2, grid2.scalars.values.ravel())
     npt.assert_allclose(spline.force_coords_, coords2)
     assert not np.allclose(spline.force_coords_, coords)
 
@@ -122,3 +124,13 @@ def test_spline_warns_underdetermined():
         assert len(warn) == 1
         assert issubclass(warn[-1].category, UserWarning)
         assert str(warn[-1].message).startswith("Under-determined problem")
+
+
+@requires_numba
+def test_spline_jacobian_implementations():
+    "Compare the numba and numpy implementations."
+    data = CheckerBoard().scatter(size=1500, random_state=1)
+    coords = (data.easting, data.northing)
+    jac_numpy = Spline(engine="numpy").jacobian(coords, coords)
+    jac_numba = Spline(engine="numba").jacobian(coords, coords)
+    npt.assert_allclose(jac_numpy, jac_numba)
diff --git a/verde/tests/test_utils.py b/verde/tests/test_utils.py
new file mode 100644
index 0000000..df35bb5
--- /dev/null
+++ b/verde/tests/test_utils.py
@@ -0,0 +1,37 @@
+"""
+Test the utility functions.
+"""
+import sys
+from unittest import mock
+
+import pytest
+
+from ..utils import parse_engine, dummy_jit
+
+
+def test_parse_engine():
+    "Check that it works for common input"
+    assert parse_engine("numba") == "numba"
+    assert parse_engine("numpy") == "numpy"
+    with mock.patch.dict(sys.modules, {"numba": None}):
+        assert parse_engine("auto") == "numpy"
+    with mock.patch.dict(sys.modules, {"numba": mock.MagicMock()}):
+        assert parse_engine("auto") == "numba"
+
+
+def test_parse_engine_fails():
+    "Check that the exception is raised for invalid engines"
+    with pytest.raises(ValueError):
+        parse_engine("some invalid engine")
+
+
+def test_dummy_jit():
+    "Make sure the dummy function raises an exception"
+
+    @dummy_jit(target="cpt")
+    def function():
+        "Some random function"
+        return 0
+
+    with pytest.raises(RuntimeError):
+        function()
diff --git a/verde/tests/test_vector.py b/verde/tests/test_vector.py
index b956388..39e6f01 100644
--- a/verde/tests/test_vector.py
+++ b/verde/tests/test_vector.py
@@ -9,6 +9,7 @@ import numpy.testing as npt
 from ..datasets.synthetic import CheckerBoard
 from ..coordinates import grid_coordinates
 from ..vector import VectorSpline2D
+from ..utils import n_1d_arrays
 
 
 @pytest.fixture
@@ -29,18 +30,18 @@ def test_vector2d(data2d):
     # There should be 1 force per data point
     assert data[0].size == spline.force_coords_[0].size
     assert data[0].size * 2 == spline.force_.size
-    npt.assert_allclose(spline.force_coords_, coords)
+    npt.assert_allclose(spline.force_coords_, n_1d_arrays(coords, n=2))
 
 
 def test_vector2d_twice(data2d):
     "Check that the force coordinates are updated if fitting a second time"
     coords, data = data2d
     spline = VectorSpline2D().fit(coords, data)
-    npt.assert_allclose(spline.force_coords_, coords)
+    npt.assert_allclose(spline.force_coords_, n_1d_arrays(coords, n=2))
     coords2 = tuple(i * 1000 for i in coords)
     spline.fit(coords2, data)
-    npt.assert_allclose(spline.force_coords_, coords2)
-    assert not np.allclose(spline.force_coords_, coords)
+    npt.assert_allclose(spline.force_coords_, n_1d_arrays(coords2, n=2))
+    assert not np.allclose(spline.force_coords_, n_1d_arrays(coords, n=2))
 
 
 def test_vector2d_weights(data2d):
diff --git a/verde/tests/utils.py b/verde/tests/utils.py
new file mode 100644
index 0000000..4a1618a
--- /dev/null
+++ b/verde/tests/utils.py
@@ -0,0 +1,16 @@
+"""
+Testing utilities.
+"""
+import pytest
+
+
+def requires_numba(function):
+    """
+    Skip the decorated test if numba is not installed.
+    """
+    try:
+        import numba
+    except ImportError:
+        numba = None
+    mark = pytest.mark.skipif(numba is None, reason="requires numba")
+    return mark(function)