Repository: pybamm-team/PyBaMM (python)

Issue:
Compatibility with NumPy v2
NumPy version 2.0.0 has been released today – we'll need to fix the failing unit tests. The tests are coming from `autograd` and the authors mention that they don't intend to maintain the dependency anymore so we should aim to remove it altogether, either now or later

Pull request description:
#4183 remove autograd
# Description

Removes autograd dependency and thus enables update to numpy 2.0.
The "differentiate numpy function directly" functionality, which relied on autograd, is deprecated, but I am not aware of any cases where it was being used, since we switched to implementing our own functions and specifying their derivatives.

Fixes #4183 

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [ ] Bug fix (non-breaking change which fixes an issue)

# Key checklist:

- [ ] No style issues: `$ pre-commit run` (or `$ nox -s pre-commit`) (see [CONTRIBUTING.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CONTRIBUTING.md#installing-and-using-pre-commit) for how to set this up to run automatically when committing locally, in just two lines of code)
- [ ] All tests pass: `$ python run-tests.py --all` (or `$ nox -s tests`)
- [ ] The documentation builds: `$ python run-tests.py --doctest` (or `$ nox -s doctests`)

You can run integration tests, unit tests, and doctests together at once, using `$ python run-tests.py --quick` (or `$ nox -s quick`).

## Further checks:

- [ ] Code is commented, particularly in hard-to-understand areas
- [ ] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e4638dff..fe78866a3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,6 +50,7 @@
 
 ## Breaking changes
 
+- Functions that are created using `pybamm.Function(function_object, children)` can no longer be differentiated symbolically (e.g. to compute the Jacobian). This should affect no users, since function derivatives for all "standard" functions are explicitly implemented ([#4196](https://github.com/pybamm-team/PyBaMM/pull/4196))
 - Removed data files under `pybamm/input` and released them in a separate repository upstream at [pybamm-data](https://github.com/pybamm-team/pybamm-data/releases/tag/v1.0.0). Note that data files under `pybamm/input/parameters` have not been removed. ([#4098](https://github.com/pybamm-team/PyBaMM/pull/4098))
 - Removed `check_model` argument from `Simulation.solve`. To change the `check_model` option, use `Simulation(..., discretisation_kwargs={"check_model": False})`. ([#4020](https://github.com/pybamm-team/PyBaMM/pull/4020))
 - Removed multiple Docker images. Here on, a single Docker image tagged `pybamm/pybamm:latest` will be provided with both solvers (`IDAKLU` and `JAX`) pre-installed. ([#3992](https://github.com/pybamm-team/PyBaMM/pull/3992))
diff --git a/asv.conf.json b/asv.conf.json
index 5117cf75d..bf9935ce9 100644
--- a/asv.conf.json
+++ b/asv.conf.json
@@ -82,7 +82,6 @@
       "wget": [],
       "cmake": [],
       "anytree": [],
-      "autograd": [],
       "scikit-fem": [],
       "imageio": [],
       "pybtex": [],
diff --git a/pybamm/expression_tree/functions.py b/pybamm/expression_tree/functions.py
index 5f686539b..4e087e972 100644
--- a/pybamm/expression_tree/functions.py
+++ b/pybamm/expression_tree/functions.py
@@ -11,7 +11,6 @@ from collections.abc import Sequence
 from typing_extensions import TypeVar
 
 import pybamm
-from pybamm.util import import_optional_dependency
 
 
 class Function(pybamm.Symbol):
@@ -26,9 +25,6 @@ class Function(pybamm.Symbol):
         func(child0.evaluate(t, y, u), child1.evaluate(t, y, u), etc).
     children : :class:`pybamm.Symbol`
         The children nodes to apply the function to
-    derivative : str, optional
-        Which derivative to use when differentiating ("autograd" or "derivative").
-        Default is "autograd".
     differentiated_function : method, optional
         The function which was differentiated to obtain this one. Default is None.
     """
@@ -38,7 +34,6 @@ class Function(pybamm.Symbol):
         function: Callable,
         *children: pybamm.Symbol,
         name: str | None = None,
-        derivative: str | None = "autograd",
         differentiated_function: Callable | None = None,
     ):
         # Turn numbers into scalars
@@ -57,7 +52,6 @@ class Function(pybamm.Symbol):
         domains = self.get_children_domains(children)
 
         self.function = function
-        self.derivative = derivative
         self.differentiated_function = differentiated_function
 
         super().__init__(name, children=children, domains=domains)
@@ -99,30 +93,10 @@ class Function(pybamm.Symbol):
         Derivative with respect to child number 'idx'.
         See :meth:`pybamm.Symbol._diff()`.
         """
-        autograd = import_optional_dependency("autograd")
-        # Store differentiated function, needed in case we want to convert to CasADi
-        if self.derivative == "autograd":
-            return Function(
-                autograd.elementwise_grad(self.function, idx),
-                *children,
-                differentiated_function=self.function,
-            )
-        elif self.derivative == "derivative":
-            if len(children) > 1:
-                raise ValueError(
-                    """
-                    differentiation using '.derivative()' not implemented for functions
-                    with more than one child
-                    """
-                )
-            else:
-                # keep using "derivative" as derivative
-                return pybamm.Function(
-                    self.function.derivative(),  # type: ignore[attr-defined]
-                    *children,
-                    derivative="derivative",
-                    differentiated_function=self.function,
-                )
+        raise NotImplementedError(
+            "Derivative of base Function class is not implemented. "
+            "Please implement in child class."
+        )
 
     def _function_jac(self, children_jacs):
         """Calculate the Jacobian of a function."""
@@ -190,7 +164,6 @@ class Function(pybamm.Symbol):
                 self.function,
                 *children,
                 name=self.name,
-                derivative=self.derivative,
                 differentiated_function=self.differentiated_function,
             )
         else:
@@ -217,7 +190,6 @@ class Function(pybamm.Symbol):
                 self.function,
                 *children,
                 name=self.name,
-                derivative=self.derivative,
                 differentiated_function=self.differentiated_function,
             )
         )
diff --git a/pybamm/expression_tree/interpolant.py b/pybamm/expression_tree/interpolant.py
index 84c982342..48ef49f87 100644
--- a/pybamm/expression_tree/interpolant.py
+++ b/pybamm/expression_tree/interpolant.py
@@ -50,6 +50,7 @@ class Interpolant(pybamm.Function):
         interpolator: str | None = "linear",
         extrapolate: bool = True,
         entries_string: str | None = None,
+        _num_derivatives: int = 0,
     ):
         # Check interpolator is valid
         if interpolator not in ["linear", "cubic", "pchip"]:
@@ -189,9 +190,13 @@ class Interpolant(pybamm.Function):
         self.x = x
         self.y = y
         self.entries_string = entries_string
-        super().__init__(
-            interpolating_function, *children, name=name, derivative="derivative"
-        )
+
+        # Differentiate the interpolating function if necessary
+        self._num_derivatives = _num_derivatives
+        for _ in range(_num_derivatives):
+            interpolating_function = interpolating_function.derivative()
+
+        super().__init__(interpolating_function, *children, name=name)
 
         # Store information as attributes
         self.interpolator = interpolator
@@ -213,6 +218,7 @@ class Interpolant(pybamm.Function):
             name=snippet["name"],
             interpolator=snippet["interpolator"],
             extrapolate=snippet["extrapolate"],
+            _num_derivatives=snippet["_num_derivatives"],
         )
 
     @property
@@ -241,6 +247,7 @@ class Interpolant(pybamm.Function):
                 self.entries_string,
                 *tuple([child.id for child in self.children]),
                 *tuple(self.domain),
+                self._num_derivatives,
             )
         )
 
@@ -256,6 +263,7 @@ class Interpolant(pybamm.Function):
             interpolator=self.interpolator,
             extrapolate=self.extrapolate,
             entries_string=self.entries_string,
+            _num_derivatives=self._num_derivatives,
         )
 
     def _function_evaluate(self, evaluated_children):
@@ -311,6 +319,27 @@ class Interpolant(pybamm.Function):
         else:  # pragma: no cover
             raise ValueError(f"Invalid dimension: {self.dimension}")
 
+    def _function_diff(self, children: Sequence[pybamm.Symbol], idx: float):
+        """
+        Derivative with respect to child number 'idx'.
+        See :meth:`pybamm.Symbol._diff()`.
+        """
+        if len(children) > 1:
+            raise NotImplementedError(
+                "differentiation not implemented for functions with more than one child"
+            )
+        else:
+            # keep using "derivative" as derivative
+            return Interpolant(
+                self.x,
+                self.y,
+                children,
+                name=self.name,
+                interpolator=self.interpolator,
+                extrapolate=self.extrapolate,
+                _num_derivatives=self._num_derivatives + 1,
+            )
+
     def to_json(self):
         """
         Method to serialise an Interpolant object into JSON.
@@ -323,6 +352,7 @@ class Interpolant(pybamm.Function):
             "y": self.y.tolist(),
             "interpolator": self.interpolator,
             "extrapolate": self.extrapolate,
+            "_num_derivatives": self._num_derivatives,
         }
 
         return json_dict
diff --git a/pyproject.toml b/pyproject.toml
index 9bb2ae9c3..7de16acca 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -121,7 +121,6 @@ jax = [
 ]
 # Contains all optional dependencies, except for jax and dev dependencies
 all = [
-    "autograd>=1.6.2",
     "scikit-fem>=8.1.0",
     "pybamm[examples,plot,cite,bpx,tqdm]",
 ]


Test patch:
diff --git a/tests/unit/test_expression_tree/test_functions.py b/tests/unit/test_expression_tree/test_functions.py
index 6c6881c37..16626f751 100644
--- a/tests/unit/test_expression_tree/test_functions.py
+++ b/tests/unit/test_expression_tree/test_functions.py
@@ -13,7 +13,6 @@ import sympy
 from tests import (
     function_test,
     multi_var_function_test,
-    multi_var_function_cube_test,
 )
 
 
@@ -52,57 +51,12 @@ class TestFunction(TestCase):
 
     def test_diff(self):
         a = pybamm.StateVector(slice(0, 1))
-        b = pybamm.StateVector(slice(1, 2))
-        y = np.array([5])
         func = pybamm.Function(function_test, a)
-        self.assertEqual(func.diff(a).evaluate(y=y), 2)
-        self.assertEqual(func.diff(func).evaluate(), 1)
-        func = pybamm.sin(a)
-        self.assertEqual(func.evaluate(y=y), np.sin(a.evaluate(y=y)))
-        self.assertEqual(func.diff(a).evaluate(y=y), np.cos(a.evaluate(y=y)))
-        func = pybamm.exp(a)
-        self.assertEqual(func.evaluate(y=y), np.exp(a.evaluate(y=y)))
-        self.assertEqual(func.diff(a).evaluate(y=y), np.exp(a.evaluate(y=y)))
-
-        # multiple variables
-        func = pybamm.Function(multi_var_function_test, 4 * a, 3 * a)
-        self.assertEqual(func.diff(a).evaluate(y=y), 7)
-        func = pybamm.Function(multi_var_function_test, 4 * a, 3 * b)
-        self.assertEqual(func.diff(a).evaluate(y=np.array([5, 6])), 4)
-        self.assertEqual(func.diff(b).evaluate(y=np.array([5, 6])), 3)
-        func = pybamm.Function(multi_var_function_cube_test, 4 * a, 3 * b)
-        self.assertEqual(func.diff(a).evaluate(y=np.array([5, 6])), 4)
-        self.assertEqual(
-            func.diff(b).evaluate(y=np.array([5, 6])), 3 * 3 * (3 * 6) ** 2
-        )
-
-        # exceptions
-        func = pybamm.Function(
-            multi_var_function_cube_test, 4 * a, 3 * b, derivative="derivative"
-        )
-        with self.assertRaises(ValueError):
+        with self.assertRaisesRegex(
+            NotImplementedError, "Derivative of base Function class is not implemented"
+        ):
             func.diff(a)
 
-    def test_function_of_multiple_variables(self):
-        a = pybamm.Variable("a")
-        b = pybamm.Parameter("b")
-        func = pybamm.Function(multi_var_function_test, a, b)
-        self.assertEqual(func.name, "function (multi_var_function_test)")
-        self.assertEqual(str(func), "multi_var_function_test(a, b)")
-        self.assertEqual(func.children[0].name, a.name)
-        self.assertEqual(func.children[1].name, b.name)
-
-        # test eval and diff
-        a = pybamm.StateVector(slice(0, 1))
-        b = pybamm.StateVector(slice(1, 2))
-        y = np.array([5, 2])
-        func = pybamm.Function(multi_var_function_test, a, b)
-
-        self.assertEqual(func.evaluate(y=y), 7)
-        self.assertEqual(func.diff(a).evaluate(y=y), 1)
-        self.assertEqual(func.diff(b).evaluate(y=y), 1)
-        self.assertEqual(func.diff(func).evaluate(), 1)
-
     def test_exceptions(self):
         a = pybamm.Variable("a", domain="something")
         b = pybamm.Variable("b", domain="something else")
diff --git a/tests/unit/test_expression_tree/test_interpolant.py b/tests/unit/test_expression_tree/test_interpolant.py
index 00a0f3bb4..f5ded9cf8 100644
--- a/tests/unit/test_expression_tree/test_interpolant.py
+++ b/tests/unit/test_expression_tree/test_interpolant.py
@@ -326,6 +326,20 @@ class TestInterpolant(TestCase):
                 decimal=3,
             )
 
+        # test 2D interpolation diff fails
+        x = (np.arange(-5.01, 5.01, 0.05), np.arange(-5.01, 5.01, 0.01))
+        xx, yy = np.meshgrid(x[0], x[1], indexing="ij")
+        z = np.sin(xx**2 + yy**2)
+        var1 = pybamm.StateVector(slice(0, 1))
+        var2 = pybamm.StateVector(slice(1, 2))
+        # linear
+        interp = pybamm.Interpolant(x, z, (var1, var2), interpolator="linear")
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            "differentiation not implemented for functions with more than one child",
+        ):
+            interp.diff(var1)
+
     def test_processing(self):
         x = np.linspace(0, 1, 200)
         y = pybamm.StateVector(slice(0, 2))
@@ -369,6 +383,7 @@ class TestInterpolant(TestCase):
             ],
             "interpolator": "linear",
             "extrapolate": True,
+            "_num_derivatives": 0,
         }
 
         # check correct writing to json
diff --git a/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py b/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py
index 2b9aa0847..20c2e40db 100644
--- a/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py
+++ b/tests/unit/test_expression_tree/test_operations/test_convert_to_casadi.py
@@ -314,22 +314,6 @@ class TestCasadiConverter(TestCase):
         y_eval = np.linspace(0, 1, expr.size)
         self.assert_casadi_equal(f(y_eval), casadi.SX(expr.evaluate(y=y_eval)))
 
-    def test_convert_differentiated_function(self):
-        a = pybamm.InputParameter("a")
-        b = pybamm.InputParameter("b")
-
-        def myfunction(x, y):
-            return x + y**3
-
-        f = pybamm.Function(myfunction, a, b).diff(a)
-        self.assert_casadi_equal(
-            f.to_casadi(inputs={"a": 1, "b": 2}), casadi.DM(1), evalf=True
-        )
-        f = pybamm.Function(myfunction, a, b).diff(b)
-        self.assert_casadi_equal(
-            f.to_casadi(inputs={"a": 1, "b": 2}), casadi.DM(12), evalf=True
-        )
-
     def test_convert_input_parameter(self):
         casadi_t = casadi.MX.sym("t")
         casadi_y = casadi.MX.sym("y", 10)
diff --git a/tests/unit/test_expression_tree/test_operations/test_jac.py b/tests/unit/test_expression_tree/test_operations/test_jac.py
index 64280fd93..bec38d724 100644
--- a/tests/unit/test_expression_tree/test_operations/test_jac.py
+++ b/tests/unit/test_expression_tree/test_operations/test_jac.py
@@ -8,7 +8,6 @@ import numpy as np
 import unittest
 from scipy.sparse import eye
 from tests import get_mesh_for_testing
-from tests import multi_var_function_test
 
 
 class TestJacobian(TestCase):
@@ -213,12 +212,6 @@ class TestJacobian(TestCase):
         dfunc_dy = func.jac(y).evaluate(y=y0)
         np.testing.assert_array_equal(0, dfunc_dy)
 
-        # several children
-        func = pybamm.Function(multi_var_function_test, 2 * y, 3 * y)
-        jacobian = np.diag(5 * np.ones(4))
-        dfunc_dy = func.jac(y).evaluate(y=y0)
-        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
-
     def test_index(self):
         vec = pybamm.StateVector(slice(0, 5))
         ind = pybamm.Index(vec, 3)
diff --git a/tests/unit/test_expression_tree/test_operations/test_jac_2D.py b/tests/unit/test_expression_tree/test_operations/test_jac_2D.py
index 6df2c9f77..2be71d85e 100644
--- a/tests/unit/test_expression_tree/test_operations/test_jac_2D.py
+++ b/tests/unit/test_expression_tree/test_operations/test_jac_2D.py
@@ -9,7 +9,6 @@ import unittest
 from scipy.sparse import eye
 from tests import (
     get_1p1d_discretisation_for_testing,
-    multi_var_function_test,
 )
 
 
@@ -200,12 +199,6 @@ class TestJacobian(TestCase):
         dfunc_dy = func.jac(y).evaluate(y=y0)
         np.testing.assert_array_equal(0, dfunc_dy)
 
-        # several children
-        func = pybamm.Function(multi_var_function_test, 2 * y, 3 * y)
-        jacobian = np.diag(5 * np.ones(8))
-        dfunc_dy = func.jac(y).evaluate(y=y0)
-        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
-
     def test_jac_of_domain_concatenation(self):
         # create mesh
         disc = get_1p1d_discretisation_for_testing()