Compute and use the gradients
from .core import MetricObject
from interpolation.splines import eval_linear, UCGrid
from interpolation.splines import extrap_options as xto
+
import numpy as np
+from copy import copy
class LinearFast(MetricObject):
@@ -10,6 +12,7 @@ class LinearFast(MetricObject):
def __init__(self, f_val, grids, extrap_options=None):
+ self.dim = len(grids)
self.f_val = f_val
self.grid_list = grids
self.Grid = UCGrid(*grids)
@@ -34,20 +37,50 @@ class LinearFastDecay(LinearFast):
self, f_val, grids, limit_func, decay_weights=None, extrap_gradient=False
):
+ super().__init__(f_val, grids)
+
self.limit_func = limit_func
self.upper_limits = np.array([x[-1] for x in grids])
if decay_weights is None:
- self.decay_weights = np.ones(len(grids))
+ self.decay_weights = np.ones(self.dim)
else:
self.decay_weights = decay_weights
self.extrap_gradient = extrap_gradient
- super().__init__(f_val, grids)
+ def _grad(self, points):
+
+ # Save in a matrix: rows are points, columns are dimensions
+ col_args = copy(points)
+
+ # Preallocate gradient
+ grads = np.empty((col_args.shape[0], self.dim))
+
+ for i, grid in enumerate(self.grid_list):
+
+ col = copy(col_args[:, i])
+
+ lower_inds = np.searchsorted(grid, col) - 1
+ lower_inds = np.minimum(np.maximum(lower_inds, 0), len(grid) - 2)
+
+ x_low = grid[lower_inds]
+ x_upp = grid[lower_inds + 1]
+
+ # Lower values
+ col_args[:, i] = x_low
+ vals_low = eval_linear(self.Grid, self.f_val, col_args,)
+ # Upper values
+ col_args[:, i] = x_upp
+ vals_upp = eval_linear(self.Grid, self.f_val, col_args,)
+
+ # Restore column
+ col_args[:, i] = col
+
+ # Get gradient
+ grads[:, i] = (vals_upp - vals_low) / (x_upp - x_low)
- def grad(self, *args):
- pass
+ return grads
def decay(self, x, closest_x):
@@ -84,9 +117,11 @@ class LinearFastDecay(LinearFast):
# Add gradient term if requested
if self.extrap_gradient:
- grad = self.grad(upper_ex_nearest)
- f[upper_ex_inds] += decay * np.dot(
- upper_ex_points - upper_ex_f_nearest, grad
+ grad = self._grad(upper_ex_nearest)
+ f[upper_ex_inds] += decay * np.sum(
+ (upper_ex_points - upper_ex_nearest) * grad,
+ axis = 1,
+ keepdims=False,
)
return np.reshape(f, argshape)
