Don't use the Jacobian for prediction (#149)

All gridders create a Jacobian matrix to make predictions (using a dot product with the estimated
parameters). This is very inefficient with regards to memory. Generating largeish grids requires a
huge amount of RAM.

Replace these with a for loop summation. Accelerate them with numba to achieve the same 
performance. Refactor the Green's functions calculations into separate functions that can be
JITed. The pure Python version is slower than the dot product but at least it won't blow up 
memory.

Author: leouieda

verde/spline.py

@@ -18,6 +18,10 @@
     from .utils import dummy_jit as jit
 
 
+# Default arguments for numba.jit
+JIT_ARGS = dict(nopython=True, target="cpu", fastmath=True, parallel=True)
+
+
 class Spline(BaseGridder):
     r"""
     Biharmonic spline interpolation using Green's functions.
@@ -71,14 +75,14 @@ class Spline(BaseGridder):
         then will be set to the data coordinates the first time
         :meth:`~verde.Spline.fit` is called.
     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
+        Computation engine for the Jacobian matrix and prediction. 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 usually faster, which
         makes fitting and predicting faster.
 
     Attributes
     ----------
-    forces_ : array
+    force_ : array
         The estimated forces that fit the observed data.
     region_ : tuple
         The boundaries (``[W, E, S, N]``) of the data used to fit the
@@ -152,9 +156,19 @@ def predict(self, coordinates):
 
         """
         check_is_fitted(self, ["force_"])
-        jac = self.jacobian(coordinates[:2], self.force_coords)
         shape = np.broadcast(*coordinates[:2]).shape
-        return jac.dot(self.force_).reshape(shape)
+        force_east, force_north = n_1d_arrays(self.force_coords, n=2)
+        east, north = n_1d_arrays(coordinates, n=2)
+        data = np.empty(east.size, dtype=east.dtype)
+        if parse_engine(self.engine) == "numba":
+            data = predict_numba(
+                east, north, force_east, force_north, self.mindist, self.force_, data
+            )
+        else:
+            data = predict_numpy(
+                east, north, force_east, force_north, self.mindist, self.force_, data
+            )
+        return data.reshape(shape)
 
     def jacobian(self, coordinates, force_coords, dtype="float64"):
         """
@@ -183,12 +197,14 @@ def jacobian(self, coordinates, force_coords, dtype="float64"):
         """
         force_east, force_north = n_1d_arrays(force_coords, n=2)
         east, north = n_1d_arrays(coordinates, n=2)
+        jac = np.empty((east.size, force_east.size), dtype=dtype)
         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)
+            jac = jacobian_numba(
+                east, north, force_east, force_north, self.mindist, jac
+            )
         else:
             jac = jacobian_numpy(
-                east, north, force_east, force_north, self.mindist, dtype
+                east, north, force_east, force_north, self.mindist, jac
             )
         return jac
 
@@ -214,34 +230,59 @@ def warn_weighted_exact_solution(spline, weights):
         )
 
 
-@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.
-    """
+def greens_func(east, north, mindist):
+    "Calculate the Green's function for the Bi-Harmonic Spline"
+    distance = np.sqrt(east ** 2 + north ** 2)
+    # 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)
+
+
+def predict_numpy(east, north, force_east, force_north, mindist, forces, result):
+    "Calculate the predicted data using numpy."
+    result[:] = 0
+    for j in range(forces.size):
+        green = greens_func(east - force_east[j], north - force_north[j], mindist)
+        result += green * forces[j]
+    return result
+
+
+def jacobian_numpy(east, north, force_east, force_north, mindist, jac):
+    "Calculate the Jacobian using numpy broadcasting."
+    # Reshaping the data to a column vector will automatically build a distance matrix
+    # between each data point and force.
+    jac[:] = greens_func(
+        east.reshape((east.size, 1)) - force_east,
+        north.reshape((north.size, 1)) - force_north,
+        mindist,
+    )
+    return jac
+
+
+@jit(**JIT_ARGS)
+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
+        result[i] = 0
+        for j in range(forces.size):
+            green = GREENS_FUNC_JIT(
+                east[i] - force_east[j], north[i] - force_north[j], mindist
+            )
+            result[i] += green * forces[j]
+    return result
+
+
+@jit(**JIT_ARGS)
+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 j in range(force_east.size):
-            distance = np.sqrt(
-                (east[i] - force_east[j]) ** 2 + (north[i] - force_north[j]) ** 2
+            jac[i, j] = GREENS_FUNC_JIT(
+                east[i] - force_east[j], north[i] - force_north[j], mindist
             )
-            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)
+# Jit compile the Green's functions for use in the numba functions
+GREENS_FUNC_JIT = jit(**JIT_ARGS)(greens_func)

verde/tests/test_spline.py

@@ -116,3 +116,14 @@ def test_spline_jacobian_implementations():
     jac_numpy = Spline(engine="numpy").jacobian(coords, coords)
     jac_numba = Spline(engine="numba").jacobian(coords, coords)
     npt.assert_allclose(jac_numpy, jac_numba)
+
+
+@requires_numba
+def test_spline_predict_implementations():
+    "Compare the numba and numpy implementations."
+    size = 500
+    data = CheckerBoard().scatter(size=size, random_state=1)
+    coords = (data.easting, data.northing)
+    pred_numpy = Spline(engine="numpy").fit(coords, data.scalars).predict(coords)
+    pred_numba = Spline(engine="numba").fit(coords, data.scalars).predict(coords)
+    npt.assert_allclose(pred_numpy, pred_numba)

verde/tests/test_vector.py

@@ -18,7 +18,7 @@
 def data2d():
     "Make 2 component vector data"
     synth = CheckerBoard()
-    coordinates = grid_coordinates(synth.region, shape=(30, 25))
+    coordinates = grid_coordinates(synth.region, shape=(15, 20))
     data = tuple(synth.predict(coordinates).ravel() for i in range(2))
     return tuple(i.ravel() for i in coordinates), data
 
@@ -38,7 +38,7 @@ def test_vector2d(data2d):
 def test_vector2d_weights(data2d):
     "Use unit weights and a regular grid solution"
     coords, data = data2d
-    outlier = 100
+    outlier = 10
     outlier_value = 100000
     data_outlier = tuple(i.copy() for i in data)
     data_outlier[0][outlier] += outlier_value
@@ -70,6 +70,15 @@ def test_vector2d_jacobian_implementations(data2d):
     npt.assert_allclose(jac_numpy, jac_numba)
 
 
+@requires_numba
+def test_vector2d_predict_implementations(data2d):
+    "Compare the numba and numpy implementations."
+    coords, data = data2d
+    pred_numpy = VectorSpline2D(engine="numpy").fit(coords, data).predict(coords)
+    pred_numba = VectorSpline2D(engine="numba").fit(coords, data).predict(coords)
+    npt.assert_allclose(pred_numpy, pred_numba, atol=1e-4)
+
+
 ########################################################################################
 # Test the Vector meta-gridder
 

verde/trend.py

@@ -6,6 +6,7 @@
 
 from .base import BaseGridder, check_fit_input, least_squares
 from .coordinates import get_region
+from .utils import n_1d_arrays
 
 
 class Trend(BaseGridder):
@@ -91,8 +92,9 @@ def fit(self, coordinates, data, weights=None):
 
         """
         coordinates, data, weights = check_fit_input(coordinates, data, weights)
-        self.region_ = get_region(coordinates[:2])
-        jac = self.jacobian(coordinates[:2], dtype=data.dtype)
+        easting, northing = n_1d_arrays(coordinates, 2)
+        self.region_ = get_region((easting, northing))
+        jac = self.jacobian((easting, northing), dtype=data.dtype)
         self.coef_ = least_squares(jac, data, weights, damping=None)
         return self
 
@@ -117,9 +119,13 @@ def predict(self, coordinates):
 
         """
         check_is_fitted(self, ["coef_"])
-        jac = self.jacobian(coordinates[:2])
+        easting, northing = n_1d_arrays(coordinates, 2)
         shape = np.broadcast(*coordinates[:2]).shape
-        return jac.dot(self.coef_).reshape(shape)
+        data = np.zeros(easting.size, dtype=easting.dtype)
+        combinations = polynomial_power_combinations(self.degree)
+        for coef, (i, j) in zip(self.coef_, combinations):
+            data += (easting ** i) * (northing ** j) * coef
+        return data.reshape(shape)
 
     def jacobian(self, coordinates, dtype="float64"):
         """
@@ -162,15 +168,15 @@ def jacobian(self, coordinates, dtype="float64"):
          [ 1  4 -1 16 -4  1]]
 
         """
-        easting, northing = coordinates[:2]
+        easting, northing = n_1d_arrays(coordinates, 2)
         if easting.shape != northing.shape:
             raise ValueError("Coordinate arrays must have the same shape.")
         combinations = polynomial_power_combinations(self.degree)
         ndata = easting.size
         nparams = len(combinations)
         out = np.empty((ndata, nparams), dtype=dtype)
         for col, (i, j) in enumerate(combinations):
-            out[:, col] = (easting.ravel() ** i) * (northing.ravel() ** j)
+            out[:, col] = (easting ** i) * (northing ** j)
         return out