[WIP] Vector-valued interpolation in econforgeinterp

HARK/econforgeinterp.py

@@ -24,10 +24,12 @@ class LinearFast(MetricObject):
 
     def __init__(self, f_val, grids, extrap_mode="linear"):
         """
-        f_val: numpy.array
-            An array containing the values of the function at the grid points.
-            It's i-th dimension must be of the same lenght as the i-th grid.
+        f_val: numpy.array, or [numpy.array]
+            An array or list of arrays containing the values of the function(s) at the grid points.
+            If it is an array, it's i-th dimension must be of the same lenght as the i-th grid.
             f_val[i,j,k] must be f(grids[0][i], grids[1][j], grids[2][k]).
+            If it is a list of arrays, each array must have the same shape, and
+            f_val[n][i,j,k] must be f[n](grids[0][i], grids[1][j], grids[2][k]).
         grids: [numpy.array]
             One-dimensional list of numpy arrays. It's i-th entry must be the grid
             to be used for the i-th independent variable.
@@ -36,7 +38,18 @@ def __init__(self, f_val, grids, extrap_mode="linear"):
             constant extrapolation. The default is multilinear.
         """
         self.dim = len(grids)
-        self.f_val = f_val
+
+        if isinstance(f_val, list):
+            # First check that all arrays have the same shape
+            if not all([x.shape == f_val[0].shape for x in f_val]):
+                raise ValueError("All arrays must have the same shape.")
+            # Stack arrays in the list across a new dimension
+            self.f_val = np.stack(f_val, axis=-1)
+            self.output_dim = len(f_val)
+        else:
+            self.f_val = f_val
+            self.output_dim = 1
+
         self.grid_list = grids
         self.Grid = CGrid(*grids)
 
@@ -69,7 +82,12 @@ def __call__(self, *args):
         )
 
         # Reshape the output to the shape of inputs
-        return np.reshape(f, array_args[0].shape)
+        if self.output_dim == 1:
+            return np.reshape(f, array_args[0].shape)
+        else:
+            return tuple(
+                np.reshape(f[:, j], array_args[0].shape) for j in range(self.output_dim)
+            )
 
     def _derivs(self, deriv_tuple, *args):
         """
@@ -93,9 +111,12 @@ def _derivs(self, deriv_tuple, *args):
 
         Returns
         -------
-        [numpy.array]
+        [numpy.array] or [[numpy.array]]
             List of the derivatives that were requested in the same order
             as deriv_tuple. Each element has the shape of items in args.
+            If the interpolator represents a function of n variables, the
+            output is a list of length n, where the i-th element has the
+            requested derivatives of the i-th output.
         """
 
         # Format arguments
@@ -113,9 +134,18 @@ def _derivs(self, deriv_tuple, *args):
         )
 
         # Reshape
-        derivs = [
-            derivs[:, j].reshape(args[0].shape) for j, tup in enumerate(deriv_tuple)
-        ]
+        if self.output_dim == 1:
+            derivs = [
+                derivs[:, j].reshape(args[0].shape) for j, tup in enumerate(deriv_tuple)
+            ]
+        else:
+            derivs = [
+                [
+                    derivs[:, i, j].reshape(args[0].shape)
+                    for j, tup in enumerate(deriv_tuple)
+                ]
+                for i in range(self.output_dim)
+            ]
 
         return derivs
 
@@ -132,12 +162,14 @@ def gradient(self, *args):
 
         Returns
         -------
-        [numpy.array]
+        [numpy.array] or [[numpy.array]]
             List of the derivatives of the function with respect to each
             input, evaluated at the given points. E.g. if the interpolator
             represents 3D function f, f.gradient(x,y,z) will return
             [df/dx(x,y,z), df/dy(x,y,z), df/dz(x,y,z)]. Each element has the
             shape of items in args.
+            If the fundtion has multiple outputs, the output will be a list
+            that will have the gradient for the ith output as the ith element.
         """
         # Form a tuple that indicates which derivatives to get
         # in the way eval_linear expects
@@ -179,7 +211,12 @@ def _eval_and_grad(self, *args):
 
         results = self._derivs(eval_tup + deriv_tup, *args)
 
-        return (results[0], results[1:])
+        if self.output_dim == 1:
+            return (results[0], results[1:])
+        else:
+            levels = [x[0] for x in results]
+            grads = [x[1:] for x in results]
+            return (levels, grads)
 
 
 class DecayInterp(MetricObject):
@@ -225,6 +262,7 @@ def __init__(
         self.limit_grad = limit_grad
 
         self.grid_list = self.interp.grid_list
+        self.output_dim = self.interp.output_dim
 
         self.upper_limits = np.array([x[-1] for x in self.grid_list])
         self.dim = len(self.grid_list)
@@ -263,13 +301,19 @@ def __call__(self, *args):
         upper_ex_points = col_args[upper_ex_inds,]
         upper_ex_nearest = np.minimum(upper_ex_points, self.upper_limits[None, :])
 
-        # Find function evaluations with regular extrapolation
+        # Find function evaluations with regular interpolator
         f = self.interp(*[col_args[:, i] for i in range(self.dim)])
 
         # Find extrapolated values with chosen method
-        f[upper_ex_inds] = self.extrap_fun(upper_ex_points, upper_ex_nearest)
+        if self.output_dim == 1:
+            f[upper_ex_inds] = self.extrap_fun(upper_ex_points, upper_ex_nearest)
+            return np.reshape(f, argshape)
+        else:
+            extrap_vals = self.extrap_fun(upper_ex_points, upper_ex_nearest)
+            for i in range(self.output_dim):
+                f[i][upper_ex_inds] = extrap_vals[i]
 
-        return np.reshape(f, argshape)
+            return tuple(np.reshape(f[i], argshape) for i in range(self.output_dim))
 
     def extrap_decay_prop(self, x, closest_x):
         """
@@ -293,7 +337,13 @@ def extrap_decay_prop(self, x, closest_x):
         dist = np.dot(np.abs(x - closest_x), decay_weights)
         weight = np.exp(-1 * dist)
 
-        return weight * f_val_x + (1 - weight) * g_val_x
+        if self.output_dim == 1:
+            return weight * f_val_x + (1 - weight) * g_val_x
+        else:
+            return tuple(
+                weight * f_val_x[i] + (1 - weight) * g_val_x[i]
+                for i in range(self.output_dim)
+            )
 
     def extrap_decay_hark(self, x, closest_x):
         """
@@ -308,6 +358,9 @@ def extrap_decay_hark(self, x, closest_x):
             the closest point that falls inside the grid.
         """
 
+        if self.output_dim > 1:
+            raise NotImplementedError("extrap_decay_hark only works for 1D outputs")
+
         # Evaluate limiting function at x
         g_val_x = self.limit_fun(*[x[:, i] for i in range(self.dim)])
 
@@ -356,4 +409,10 @@ def extrap_paste(self, x, closest_x):
         # Evaluate limit function at x
         g_val_x = self.limit_fun(*[x[:, i] for i in range(self.dim)])
 
-        return f_val_closest + (g_val_x - g_val_closest)
+        if self.output_dim == 1:
+            return g_val_x + (f_val_closest - g_val_closest)
+        else:
+            return tuple(
+                g_val_x[i] + (f_val_closest[i] - g_val_closest[i])
+                for i in range(self.output_dim)
+            )

HARK/tests/test_econforgeinterp.py

@@ -375,3 +375,130 @@ def test_compare_smooth_with_LinearInterp(self):
         efor_vals = efor_lim_interp(x_eval)
 
         self.assertTrue(np.allclose(base_vals, efor_vals))
+
+
+class Test2Dto2DInterp(unittest.TestCase):
+    def f1(self, x, y):
+        return 2 * x + y
+
+    def f2(self, x, y):
+        return 3 * x + 2 * y
+
+    def setUp(self):
+        self.x = np.linspace(0, 10, 11)
+        self.y = np.linspace(0, 10, 11)
+        x_t, y_t = np.meshgrid(self.x, self.y, indexing="ij")
+
+        # Create interpolator
+        self.interp = LinearFast(
+            [self.f1(x_t, y_t), self.f2(x_t, y_t)],
+            [self.x, self.y],
+            extrap_mode="linear",
+        )
+
+        # Create points for evaluation
+        self.x_eval = np.random.rand(100) * 10
+        self.y_eval = np.random.rand(100) * 10
+
+    def test_outputs(self):
+        # Evaluete functions
+        f1_eval = self.f1(self.x_eval, self.y_eval)
+        f2_eval = self.f2(self.x_eval, self.y_eval)
+
+        # Evaluate interpolator
+        f1_inter, f2_inter = self.interp(self.x_eval, self.y_eval)
+
+        # Compare outputs
+        self.assertTrue(np.allclose(f1_eval, f1_inter))
+        self.assertTrue(np.allclose(f2_eval, f2_inter))
+
+    def test_derivatives(self):
+        # Evaluate interpolator
+        gradient = self.interp.gradient(self.x_eval, self.y_eval)
+
+        f1_derivs = gradient[0]
+        f2_derivs = gradient[1]
+
+        # Compare outputs
+        self.assertTrue(np.allclose(f1_derivs[0], 2 * np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f1_derivs[1], np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f2_derivs[0], 3 * np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f2_derivs[1], 2 * np.ones_like(self.x_eval)))
+
+    def test_eval_and_derivs(self):
+        levels, gradients = self.interp._eval_and_grad(self.x_eval, self.y_eval)
+        f1 = levels[0]
+        f2 = levels[1]
+        f1_derivs = gradients[0]
+        f2_derivs = gradients[1]
+
+        # Compare outputs
+        self.assertTrue(np.allclose(f1, self.f1(self.x_eval, self.y_eval)))
+        self.assertTrue(np.allclose(f2, self.f2(self.x_eval, self.y_eval)))
+        self.assertTrue(np.allclose(f1_derivs[0], 2 * np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f1_derivs[1], np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f2_derivs[0], 3 * np.ones_like(self.x_eval)))
+        self.assertTrue(np.allclose(f2_derivs[1], 2 * np.ones_like(self.x_eval)))
+
+
+class TestMultiouputDecay(unittest.TestCase):
+    def limit_fun(self, x):
+        return (np.sqrt(x), np.zeros_like(x))
+
+    def setUp(self):
+        # Create grid
+        self.x = np.linspace(0, 10, 11)
+        # Values for the two interpolated functions
+        y = self.x * 3 - 2
+        z = -5 * self.x + 3
+        # Create interpolator with prop method
+        self.prop_interp = DecayInterp(
+            LinearFast([y, z], [self.x], extrap_mode="linear"),
+            limit_fun=self.limit_fun,
+            extrap_method="decay_prop",
+        )
+        # Create interpolator with paste method
+        self.paste_interp = DecayInterp(
+            LinearFast([y, z], [self.x], extrap_mode="linear"),
+            limit_fun=self.limit_fun,
+            extrap_method="paste",
+        )
+
+    def test_interpolation(self):
+        # Create points for evaluation
+        x_eval = np.array([0.5, 3.5, 5.5, 7.5, 9.5])
+        # Evaluate interpolator
+        prop_vals = self.prop_interp(x_eval)
+        paste_vals = self.paste_interp(x_eval)
+        # Compare outputs
+        self.assertTrue(np.allclose(prop_vals[0], 3 * x_eval - 2))
+        self.assertTrue(np.allclose(prop_vals[1], -5 * x_eval + 3))
+        self.assertTrue(np.allclose(paste_vals[0], 3 * x_eval - 2))
+        self.assertTrue(np.allclose(paste_vals[1], -5 * x_eval + 3))
+
+    def test_extrapolation(self):
+        # Two points outside the grid, one far away and one close
+        x = np.array([10.001, 200.0])
+
+        # Evaluate interpolators
+        prop_vals = self.prop_interp(x)
+        paste_vals = self.paste_interp(x)
+
+        # Compare outputs
+
+        # Near points
+        self.assertAlmostEqual(prop_vals[0][0], 3 * x[0] - 2, places=2)
+        self.assertAlmostEqual(paste_vals[0][0], 3 * x[0] - 2, places=2)
+        self.assertAlmostEqual(prop_vals[1][0], -5 * x[0] + 3, places=2)
+        self.assertAlmostEqual(paste_vals[1][0], -5 * x[0] + 3, places=2)
+        # Far points
+        self.assertAlmostEqual(prop_vals[0][1], self.limit_fun(x[1])[0], places=2)
+        self.assertAlmostEqual(prop_vals[1][1], self.limit_fun(x[1])[1], places=2)
+
+        # The limit changes for paste
+        step_0 = self.limit_fun(x[1])[0] - self.limit_fun(10.0)[0]
+        step_1 = self.limit_fun(x[1])[1] - self.limit_fun(10.0)[1]
+        f_0 = 3 * 10.0 - 2
+        f_1 = -5 * 10.0 + 3
+        self.assertAlmostEqual(paste_vals[0][1], f_0 + step_0)
+        self.assertAlmostEqual(paste_vals[1][1], f_1 + step_1)