bug: use casadi MX.interpn_linear function instead of plugin #3783

pybamm-team · May 9, 2024 · 6e4ddbb · 6e4ddbb
1 parent c0ff64f
commit 6e4ddbb
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/pybamm/expression_tree/interpolant.py b/pybamm/expression_tree/interpolant.py
@@ -126,9 +126,10 @@ def __init__(
                     fill_value_1 = "extrapolate"
                 interpolating_function = interpolate.interp1d(
                     x1,
-                    y.T,
+                    y,
                     bounds_error=False,
                     fill_value=fill_value_1,
+                    axis=0,
                 )
             elif interpolator == "cubic":
                 interpolating_function = interpolate.CubicSpline(

diff --git a/pybamm/expression_tree/operations/convert_to_casadi.py b/pybamm/expression_tree/operations/convert_to_casadi.py
@@ -157,15 +157,31 @@ def _convert(self, symbol, t, y, y_dot, inputs):
                     )
 
                 if len(converted_children) == 1:
-                    return casadi.interpolant(
-                        "LUT", solver, symbol.x, symbol.y.flatten()
-                    )(*converted_children)
+                    if solver == "linear":
+                        test = casadi.MX.interpn_linear(
+                            symbol.x, symbol.y.flatten(), converted_children
+                        )
+                        if test.shape[0] == 1 and test.shape[1] > 1:
+                            # for some reason, pybamm.Interpolant always returns a column vector, so match that
+                            test = test.T
+                        return test
+                    else:
+                        return casadi.interpolant(
+                            "LUT", solver, symbol.x, symbol.y.flatten()
+                        )(*converted_children)
                 elif len(converted_children) in [2, 3]:
-                    LUT = casadi.interpolant(
-                        "LUT", solver, symbol.x, symbol.y.ravel(order="F")
-                    )
-                    res = LUT(casadi.hcat(converted_children).T).T
-                    return res
+                    if solver == "linear":
+                        return casadi.MX.interpn_linear(
+                            symbol.x,
+                            symbol.y.ravel(order="F"),
+                            casadi.hcat(converted_children).T,
+                        ).T
+                    else:
+                        LUT = casadi.interpolant(
+                            "LUT", solver, symbol.x, symbol.y.ravel(order="F")
+                        )
+                        res = LUT(casadi.hcat(converted_children).T).T
+                        return res
                 else:  # pragma: no cover
                     raise ValueError(
                         f"Invalid converted_children count: {len(converted_children)}"