Repository: pybamm-team/PyBaMM

Issue:
Function of many variables
Allow `Function` node to have an arbitrary number of children, and evaluate the `func` using all of those children. e.g. 
```python3 
def D(concentration, temperature):
    return (concentration * temperature) ** 2

c = pybamm.Variable("c")
T = pybamm.Variable("T")
f = pybamm.Function(D, [c,T])
# f.evaluate(t,y) = D(c.evaluate(t,y), T.evaluate(t,y))
```
This will eventually require re-writing the diff and jac methods (using chain rule on `Function`) so should be done after #197 

Discussion:

Patch:
diff --git a/docs/source/expression_tree/function.rst b/docs/source/expression_tree/function.rst
new file mode 100644
index 000000000..a27651f19
--- /dev/null
+++ b/docs/source/expression_tree/function.rst
@@ -0,0 +1,5 @@
+Function
+========
+
+.. autoclass:: pybamm.Function
+  :members:
\ No newline at end of file
diff --git a/docs/source/expression_tree/index.rst b/docs/source/expression_tree/index.rst
index 20c909560..57acd2fba 100644
--- a/docs/source/expression_tree/index.rst
+++ b/docs/source/expression_tree/index.rst
@@ -15,4 +15,5 @@ Expression Tree
   concatenations
   broadcasts
   simplify
+  function
   evaluate
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index ba9f37fe1..ec9c2f948 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -75,7 +75,6 @@ from .expression_tree.unary_operators import (
     UnaryOperator,
     Negate,
     AbsoluteValue,
-    Function,
     Index,
     SpatialOperator,
     Gradient,
@@ -91,6 +90,7 @@ from .expression_tree.unary_operators import (
     average,
     boundary_value,
 )
+from .expression_tree.function import Function
 from .expression_tree.parameter import Parameter, FunctionParameter
 from .expression_tree.broadcasts import Broadcast
 from .expression_tree.scalar import Scalar
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 9d019b804..0efd80e00 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -386,6 +386,7 @@ class Discretisation(object):
                 return spatial_method.process_binary_operators(
                     symbol, left, right, disc_left, disc_right
                 )
+
         elif isinstance(symbol, pybamm.UnaryOperator):
             child = symbol.child
             disc_child = self.process_symbol(child)
@@ -420,6 +421,10 @@ class Discretisation(object):
             else:
                 return symbol._unary_new_copy(disc_child)
 
+        elif isinstance(symbol, pybamm.Function):
+            disc_children = [self.process_symbol(child) for child in symbol.children]
+            return symbol._function_new_copy(disc_children)
+
         elif isinstance(symbol, pybamm.Variable):
             return pybamm.StateVector(self._y_slices[symbol.id], domain=symbol.domain)
 
diff --git a/pybamm/expression_tree/evaluate.py b/pybamm/expression_tree/evaluate.py
index 24da5a0b9..e1d6d40de 100644
--- a/pybamm/expression_tree/evaluate.py
+++ b/pybamm/expression_tree/evaluate.py
@@ -55,65 +55,76 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
         constant_symbols[symbol.id] = symbol.evaluate()
         return
 
-    # process children recursivly
+    # process children recursively
     for child in symbol.children:
         find_symbols(child, constant_symbols, variable_symbols)
 
-    # calculate the variable names that will hold the result of calcaulating the
+    # calculate the variable names that will hold the result of calculating the
     # children variables
-    children_vars = [id_to_python_variable(child.id, child.is_constant())
-                     for child in symbol.children]
+    children_vars = [
+        id_to_python_variable(child.id, child.is_constant())
+        for child in symbol.children
+    ]
 
     if isinstance(symbol, pybamm.BinaryOperator):
         # Multiplication and Division need special handling for scipy sparse matrices
         # TODO: we can pass through a dummy y and t to get the type and then hardcode
         # the right line, avoiding these checks
         if isinstance(symbol, pybamm.Multiplication):
-            symbol_str = \
-                "{0}.multiply({1}) if scipy.sparse.issparse({0}) else " \
-                "{1}.multiply({0}) if scipy.sparse.issparse({1}) else " \
+            symbol_str = (
+                "{0}.multiply({1}) if scipy.sparse.issparse({0}) else "
+                "{1}.multiply({0}) if scipy.sparse.issparse({1}) else "
                 "{0} * {1}".format(children_vars[0], children_vars[1])
+            )
         elif isinstance(symbol, pybamm.Division):
-            symbol_str = \
-                "{0}.multiply(1/{1}) if scipy.sparse.issparse({0}) else " \
+            symbol_str = (
+                "{0}.multiply(1/{1}) if scipy.sparse.issparse({0}) else "
                 "{0} / {1}".format(children_vars[0], children_vars[1])
+            )
         elif isinstance(symbol, pybamm.Outer):
-            symbol_str = "np.outer({}, {}).reshape(-1, 1)".format(children_vars[0],
-                                                                  children_vars[1])
+            symbol_str = "np.outer({}, {}).reshape(-1, 1)".format(
+                children_vars[0], children_vars[1]
+            )
         else:
-            symbol_str = children_vars[0] + ' ' + symbol.name + ' ' + children_vars[1]
+            symbol_str = children_vars[0] + " " + symbol.name + " " + children_vars[1]
 
     elif isinstance(symbol, pybamm.UnaryOperator):
-        # For a Function we create two lines of code, one in constant_symbols that
-        # contains the function handle, the other in variable_symbols that calls that
-        # function on the child variable
-        if isinstance(symbol, pybamm.Function):
-            constant_symbols[symbol.id] = symbol.func
-            funct_var = id_to_python_variable(symbol.id, True)
-            symbol_str = "{}({})".format(funct_var, children_vars[0])
-
         # Index has a different syntax than other univariate operations
-        elif isinstance(symbol, pybamm.Index):
-            symbol_str = "{}[{}:{}]".format(children_vars[0],
-                                            symbol.slice.start,
-                                            symbol.slice.stop)
+        if isinstance(symbol, pybamm.Index):
+            symbol_str = "{}[{}:{}]".format(
+                children_vars[0], symbol.slice.start, symbol.slice.stop
+            )
         else:
             symbol_str = symbol.name + children_vars[0]
 
+    # For a Function we create two lines of code, one in constant_symbols that
+    # contains the function handle, the other in variable_symbols that calls that
+    # function on the children variables
+    elif isinstance(symbol, pybamm.Function):
+        constant_symbols[symbol.id] = symbol.function
+        funct_var = id_to_python_variable(symbol.id, True)
+        children_str = ""
+        for child_var in children_vars:
+            if children_str == "":
+                children_str = child_var
+            else:
+                children_str += ", " + child_var
+        symbol_str = "{}({})".format(funct_var, children_str)
+
     elif isinstance(symbol, pybamm.Concatenation):
 
         # don't bother to concatenate if there is only a single child
         if isinstance(symbol, pybamm.NumpyConcatenation):
             if len(children_vars) > 1:
-                symbol_str = 'np.concatenate(({}))'.format(",".join(children_vars))
+                symbol_str = "np.concatenate(({}))".format(",".join(children_vars))
             else:
-                symbol_str = '{}'.format(",".join(children_vars))
+                symbol_str = "{}".format(",".join(children_vars))
 
         elif isinstance(symbol, pybamm.SparseStack):
             if len(children_vars) > 1:
                 symbol_str = "scipy.sparse.vstack(({}))".format(",".join(children_vars))
             else:
-                symbol_str = '{}'.format(",".join(children_vars))
+                symbol_str = "{}".format(",".join(children_vars))
 
         # DomainConcatenation specifies a particular ordering for the concatenation,
         # which we must follow
@@ -123,14 +134,16 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
             for child_var, slices in zip(children_vars, symbol._children_slices):
                 for child_dom, child_slice in slices.items():
                     slice_starts.append(symbol._slices[child_dom].start)
-                    child_vectors.append("{}[{}:{}]".format(
-                        child_var, child_slice.start, child_slice.stop
-                    ))
+                    child_vectors.append(
+                        "{}[{}:{}]".format(
+                            child_var, child_slice.start, child_slice.stop
+                        )
+                    )
             child_vectors = [v for _, v in sorted(zip(slice_starts, child_vectors))]
             if len(children_vars) > 1:
                 symbol_str = "np.concatenate(({}))".format(",".join(child_vectors))
             else:
-                symbol_str = '{}'.format(",".join(children_vars))
+                symbol_str = "{}".format(",".join(children_vars))
         else:
             raise NotImplementedError
 
@@ -139,7 +152,7 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
         symbol_str = symbol.name
 
     elif isinstance(symbol, pybamm.Time):
-        symbol_str = 't'
+        symbol_str = "t"
 
     else:
         raise NotImplementedError(
@@ -180,25 +193,18 @@ def to_python(symbol, debug=False):
 
     if debug:
         variable_lines = [
-            "print('{}'); ".format(line_format.format(
-                id_to_python_variable(symbol_id, False),
-                symbol_line
-            )) +
-            line_format.format(
-                id_to_python_variable(symbol_id, False),
-                symbol_line
+            "print('{}'); ".format(
+                line_format.format(id_to_python_variable(symbol_id, False), symbol_line)
+            )
+            + line_format.format(id_to_python_variable(symbol_id, False), symbol_line)
+            + "; print(type({0}),{0}.shape)".format(
+                id_to_python_variable(symbol_id, False)
             )
-            +
-            "; print(type({0}),{0}.shape)".format(
-                id_to_python_variable(symbol_id, False))
             for symbol_id, symbol_line in variable_symbols.items()
         ]
     else:
         variable_lines = [
-            line_format.format(
-                id_to_python_variable(symbol_id, False),
-                symbol_line
-            )
+            line_format.format(id_to_python_variable(symbol_id, False), symbol_line)
             for symbol_id, symbol_line in variable_symbols.items()
         ]
 
@@ -225,8 +231,9 @@ class EvaluatorPython:
         # store all the constant symbols in the tree as internal variables of this
         # object
         for symbol_id, value in constants.items():
-            setattr(self, id_to_python_variable(
-                symbol_id, True).replace("self.", ""), value)
+            setattr(
+                self, id_to_python_variable(symbol_id, True).replace("self.", ""), value
+            )
 
         # calculate the final variable that will output the result of calling `evaluate`
         # on `symbol`
@@ -234,11 +241,13 @@ class EvaluatorPython:
 
         # compile the generated python code
         self._variable_compiled = compile(
-            self._variable_function, self._result_var, 'exec')
+            self._variable_function, self._result_var, "exec"
+        )
 
         # compile the line that will return the output of `evaluate`
         self._return_compiled = compile(
-            self._result_var, 'return' + self._result_var, 'eval')
+            self._result_var, "return" + self._result_var, "eval"
+        )
 
     def evaluate(self, t=None, y=None, known_evals=None):
         """
diff --git a/pybamm/expression_tree/function.py b/pybamm/expression_tree/function.py
new file mode 100644
index 000000000..8eda5496d
--- /dev/null
+++ b/pybamm/expression_tree/function.py
@@ -0,0 +1,185 @@
+#
+# Function classes and methods
+#
+import autograd
+import numpy as np
+import pybamm
+from inspect import signature
+from scipy.sparse import csr_matrix
+
+
+class Function(pybamm.Symbol):
+    """A node in the expression tree representing an arbitrary function
+
+    Parameters
+    ----------
+    function : method
+        A function can have 0 or many inputs. If no inputs are given, self.evaluate()
+        simply returns func(). Otherwise, self.evaluate(t, y) returns
+        func(child0.evaluate(t, y), child1.evaluate(t, y), etc).
+    children : :class:`pybamm.Symbol`
+        The children nodes to apply the function to
+
+    **Extends:** :class:`UnaryOperator`
+    """
+
+    def __init__(self, function, *children):
+
+        name = "function ({})".format(function.__name__)
+        children_list = list(children)
+        domain = self.get_children_domains(children_list)
+
+        self.function = function
+
+        # hack to work out whether function takes any params
+        # (signature doesn't work for numpy)
+        if isinstance(function, np.ufunc):
+            self.takes_no_params = False
+        else:
+            self.takes_no_params = len(signature(function).parameters) == 0
+
+        super().__init__(name, children=children_list, domain=domain)
+
+    def get_children_domains(self, children_list):
+        """Obtains the unique domain of the children. If the
+        children have different domains then raise an error"""
+
+        domains = [child.domain for child in children_list if child.domain != []]
+
+        # check that there is one common domain amongst children
+        distinct_domains = set(tuple(dom) for dom in domains)
+
+        if len(distinct_domains) > 1:
+            raise pybamm.DomainError(
+                "Functions can only be applied to variables on the same domain"
+            )
+        elif len(distinct_domains) == 0:
+            domain = []
+        else:
+            domain = domains[0]
+
+        return domain
+
+    def diff(self, variable):
+        """ See :meth:`pybamm.Symbol.diff()`. """
+        if variable.id == self.id:
+            return pybamm.Scalar(1)
+        else:
+            children = self.orphans
+            partial_derivatives = [None] * len(children)
+            for i, child in enumerate(self.children):
+                # if variable appears in the function,use autograd to differentiate
+                # function, and apply chain rule
+                if variable.id in [symbol.id for symbol in child.pre_order()]:
+                    partial_derivatives[i] = child.diff(variable) * Function(
+                        autograd.grad(self.function), *children
+                    )
+
+            # remove None entries
+            partial_derivatives = list(filter(None, partial_derivatives))
+
+            derivative = sum(partial_derivatives)
+            if derivative == 0:
+                derivative = pybamm.Scalar(0)
+
+            return derivative
+
+    def jac(self, variable):
+        """ See :meth:`pybamm.Symbol.jac()`. """
+
+        if all(child.evaluates_to_number() for child in self.children):
+            # if children all evaluate to numbers the return zeros
+            # of right size
+            variable_y_indices = np.arange(
+                variable.y_slice.start, variable.y_slice.stop
+            )
+            jac = csr_matrix((1, np.size(variable_y_indices)))
+            jacobian = pybamm.Matrix(jac)
+        else:
+
+            # if at least one child contains variable dependence, then
+            # calculate the required partial jacobians and add them
+            jacobian = None
+            children = self.orphans
+            for child in children:
+                if not child.evaluates_to_number():
+                    jac_fun = Function(
+                        autograd.elementwise_grad(self.function), *children
+                    ) * child.jac(variable)
+
+                    jac_fun.domain = self.domain
+                    if jacobian is None:
+                        jacobian = jac_fun
+                    else:
+                        jacobian += jac_fun
+
+        return jacobian
+
+    def evaluate(self, t=None, y=None, known_evals=None):
+        """ See :meth:`pybamm.Symbol.evaluate()`. """
+        if known_evals is not None:
+            if self.id not in known_evals:
+                evaluated_children = [None] * len(self.children)
+                for i, child in enumerate(self.children):
+                    evaluated_children[i], known_evals = child.evaluate(
+                        t, y, known_evals
+                    )
+                known_evals[self.id] = self._function_evaluate(evaluated_children)
+            return known_evals[self.id], known_evals
+        else:
+            evaluated_children = [child.evaluate(t, y) for child in self.children]
+            return self._function_evaluate(evaluated_children)
+
+    def evaluate_for_shape(self, t=None, y=None):
+        """
+        Default behaviour: has same shape as all child
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        evaluated_children = [child.evaluate_for_shape() for child in self.children]
+        return self._function_evaluate(evaluated_children)
+
+    def _function_evaluate(self, evaluated_children):
+        if self.takes_no_params is True:
+            return self.function()
+        else:
+            return self.function(*evaluated_children)
+
+    def new_copy(self):
+        """ See :meth:`pybamm.Symbol.new_copy()`. """
+        return self._function_new_copy(self.orphans)
+
+    def _function_new_copy(self, children):
+        """Returns a new copy of the function.
+
+        Inputs
+        ------
+        children : : list
+            A list of the children of the function
+
+        Returns
+        -------
+            : :pybamm.Function
+            A new copy of the function
+        """
+        return pybamm.Function(self.function, *children)
+
+    def _function_simplify(self, simplified_children):
+        """
+        Simplifies the function.
+
+        Inputs
+        ------
+        simplified_children: : list
+            A list of simplified children of the function
+
+        Returns
+        -------
+         :: pybamm.Scalar() if no children
+         :: pybamm.Function if there are children
+        """
+        if self.takes_no_params is True:
+            # If self.func() takes no parameters then we can always simplify it
+            return pybamm.Scalar(self.function())
+        else:
+            return pybamm.Function(self.function, *simplified_children)
+
diff --git a/pybamm/expression_tree/parameter.py b/pybamm/expression_tree/parameter.py
index a235abf17..7adaf3605 100644
--- a/pybamm/expression_tree/parameter.py
+++ b/pybamm/expression_tree/parameter.py
@@ -35,7 +35,7 @@ class Parameter(pybamm.Symbol):
         return np.nan
 
 
-class FunctionParameter(pybamm.UnaryOperator):
+class FunctionParameter(pybamm.Symbol):
     """A node in the expression tree representing a function parameter
 
     This node will be replaced by a :class:`pybamm.Function` node if a callable function
@@ -56,10 +56,12 @@ class FunctionParameter(pybamm.UnaryOperator):
 
     """
 
-    def __init__(self, name, child, diff_variable=None):
+    def __init__(self, name, *children, diff_variable=None):
         # assign diff variable
         self.diff_variable = diff_variable
-        super().__init__(name, child)
+        children_list = list(children)
+        domain = self.get_children_domains(children_list)
+        super().__init__(name, children=children, domain=domain)
 
     def set_id(self):
         """See :meth:`pybamm.Symbol.set_id` """
@@ -69,12 +71,54 @@ class FunctionParameter(pybamm.UnaryOperator):
             + tuple(self.domain)
         )
 
+    def get_children_domains(self, children_list):
+        """Obtains the unique domain of the children. If the
+        children have different domains then raise an error"""
+        domains = [child.domain for child in children_list if child.domain != []]
+
+        # check that there is one common domain amongst children
+        distinct_domains = set(tuple(dom) for dom in domains)
+
+        if len(distinct_domains) > 1:
+            raise pybamm.DomainError(
+                "Functions can only be applied to variables on the same domain"
+            )
+        elif len(distinct_domains) == 0:
+            domain = []
+        else:
+            domain = domains[0]
+
+        return domain
+
     def diff(self, variable):
         """ See :meth:`pybamm.Symbol.diff()`. """
         # return a new FunctionParameter, that knows it will need to be differentiated
         # when the parameters are set
-        return FunctionParameter(self.name, self.orphans[0], diff_variable=variable)
+        return FunctionParameter(self.name, *self.orphans, diff_variable=variable)
+
+    def new_copy(self):
+        """ See :meth:`pybamm.Symbol.new_copy()`. """
+        return self._function_parameter_new_copy(self.orphans)
+
+    def _function_parameter_new_copy(self, children):
+        """Returns a new copy of the function parameter.
+
+        Inputs
+        ------
+        children : : list
+            A list of the children of the function
+
+        Returns
+        -------
+            : :pybamm.FunctionParameter
+            A new copy of the function parameter
+        """
+        return FunctionParameter(self.name, *children, diff_variable=self.diff_variable)
+
+    def evaluate_for_shape(self, t=None, y=None):
+        """
+        Returns the sum of the evaluated children
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return sum(child.evaluate_for_shape() for child in self.children)
 
-    def _unary_new_copy(self, child):
-        """ See :meth:`UnaryOperator._unary_new_copy()`. """
-        return FunctionParameter(self.name, child, diff_variable=self.diff_variable)
diff --git a/pybamm/expression_tree/simplify.py b/pybamm/expression_tree/simplify.py
index 90cd366a2..4deb3460b 100644
--- a/pybamm/expression_tree/simplify.py
+++ b/pybamm/expression_tree/simplify.py
@@ -580,6 +580,14 @@ class Simplification(object):
             new_symbol.domain = symbol.domain
             return simplify_if_constant(new_symbol)
 
+        elif isinstance(symbol, pybamm.Function):
+            simplified_children = [None] * len(symbol.children)
+            for i, child in enumerate(symbol.children):
+                simplified_children[i] = self.simplify(child)
+            new_symbol = symbol._function_simplify(simplified_children)
+            new_symbol.domain = symbol.domain
+            return simplify_if_constant(new_symbol)
+
         elif isinstance(symbol, pybamm.Concatenation):
             new_children = [self.simplify(child) for child in symbol.children]
             new_symbol = symbol._concatenation_simplify(new_children)
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 6563d9462..ff3758b82 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -1,11 +1,8 @@
 #
 # Unary operator classes and methods
 #
-import autograd
 import numpy as np
 import pybamm
-from inspect import signature
-from scipy.sparse import csr_matrix
 
 
 class UnaryOperator(pybamm.Symbol):
@@ -130,74 +127,6 @@ class AbsoluteValue(UnaryOperator):
         return np.abs(child)
 
 
-class Function(UnaryOperator):
-    """A node in the expression tree representing an arbitrary function
-
-    Parameters
-    ----------
-    func : method
-        A function that takes either 0 or 1 parameters. If func takes no parameters,
-        self.evaluate() return func(). Otherwise, self.evaluate(t,y) returns
-        func(child.evaluate(t,y))
-    child : :class:`pybamm.Symbol`
-        The child node to apply the function to
-
-    **Extends:** :class:`UnaryOperator`
-    """
-
-    def __init__(self, func, child):
-        """ See :meth:`pybamm.UnaryOperator.__init__()`. """
-        self.func = func
-        super().__init__("function ({})".format(func.__name__), child)
-        # hack to work out whether function takes any params
-        # (signature doesn't work for numpy)
-        if isinstance(func, np.ufunc):
-            self.takes_no_params = False
-        else:
-            self.takes_no_params = len(signature(func).parameters) == 0
-
-    def _diff(self, variable):
-        """ See :meth:`pybamm.Symbol._diff()`. """
-        child = self.orphans[0]
-        return child.diff(variable) * Function(autograd.grad(self.func), child)
-
-    def jac(self, variable):
-        """ See :meth:`pybamm.Symbol.jac()`. """
-        child = self.orphans[0]
-        if child.evaluates_to_number():
-            # return zeros of right size
-            variable_y_indices = np.arange(
-                variable.y_slice.start, variable.y_slice.stop
-            )
-            jac = csr_matrix((1, np.size(variable_y_indices)))
-            return pybamm.Matrix(jac)
-        else:
-            jac_fun = Function(autograd.elementwise_grad(self.func), child) * child.jac(
-                variable
-            )
-            jac_fun.domain = self.domain
-            return jac_fun
-
-    def _unary_evaluate(self, child):
-        """ See :meth:`UnaryOperator._unary_evaluate()`. """
-        if self.takes_no_params:
-            return self.func()
-        else:
-            return self.func(child)
-
-    def _unary_new_copy(self, child):
-        """ See :meth:`UnaryOperator._unary_new_copy()`. """
-        return pybamm.Function(self.func, child)
-
-    def _unary_simplify(self, simplified_child):
-        """ See :meth:`UnaryOperator._unary_simplify()`. """
-        if self.takes_no_params:
-            # If self.func() takes no parameters then we can always simplify it
-            return pybamm.Scalar(self.func())
-        else:
-            return pybamm.Function(self.func, simplified_child)
-
-
 class Index(UnaryOperator):
     """A node in the expression tree, which stores the index that should be
     extracted from its child after the child has been evaluated.
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index 5291368b4..03a430e1c 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -214,14 +214,14 @@ class ParameterValues(dict):
             return pybamm.Scalar(value, name=symbol.name, domain=symbol.domain)
 
         elif isinstance(symbol, pybamm.FunctionParameter):
-            new_child = self.process_symbol(symbol.children[0])
+            new_children = [self.process_symbol(child) for child in symbol.children]
             function_name = self[symbol.name]
 
             if callable(function_name):
-                function = pybamm.Function(function_name, new_child)
+                function = pybamm.Function(function_name, *new_children)
             else:
                 function = pybamm.Function(
-                    pybamm.load_function(function_name), new_child
+                    pybamm.load_function(function_name), *new_children
                 )
 
             if symbol.diff_variable is None:
@@ -240,12 +240,11 @@ class ParameterValues(dict):
             new_symbol.domain = symbol.domain
             return new_symbol
 
+        # Unary operators
         elif isinstance(symbol, pybamm.UnaryOperator):
             new_child = self.process_symbol(symbol.child)
             if isinstance(symbol, pybamm.Broadcast):
                 new_symbol = pybamm.Broadcast(new_child, symbol.domain)
-            elif isinstance(symbol, pybamm.Function):
-                new_symbol = pybamm.Function(symbol.func, new_child)
             elif isinstance(symbol, pybamm.Integral):
                 new_symbol = symbol.__class__(new_child, symbol.integration_variable)
             elif isinstance(symbol, pybamm.BoundaryOperator):
@@ -257,6 +256,13 @@ class ParameterValues(dict):
             new_symbol.domain = symbol.domain
             return new_symbol
 
+        # Functions
+        elif isinstance(symbol, pybamm.Function):
+            new_children = [None] * len(symbol.children)
+            for i, child in enumerate(symbol.children):
+                new_children[i] = self.process_symbol(child)
+            return pybamm.Function(symbol.function, *new_children)
+
         # Concatenations
         elif isinstance(symbol, pybamm.Concatenation):
             new_children = []


Test patch:
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index feff31a4c..af915a779 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -160,6 +160,28 @@ class TestDiscretise(unittest.TestCase):
         self.assertIsInstance(un2_disc, pybamm.AbsoluteValue)
         self.assertIsInstance(un2_disc.children[0], pybamm.Scalar)
 
+        # function of one variable
+        def myfun(x):
+            return np.exp(x)
+        func = pybamm.Function(myfun, var)
+        func_disc = disc.process_symbol(func)
+        self.assertIsInstance(func_disc, pybamm.Function)
+        self.assertIsInstance(func_disc.children[0], pybamm.StateVector)
+
+        func = pybamm.Function(myfun, scal)
+        func_disc = disc.process_symbol(func)
+        self.assertIsInstance(func_disc, pybamm.Function)
+        self.assertIsInstance(func_disc.children[0], pybamm.Scalar)
+
+        # function of multiple variables
+        def myfun(x, y):
+            return np.exp(x) * y
+        func = pybamm.Function(myfun, var, scal)
+        func_disc = disc.process_symbol(func)
+        self.assertIsInstance(func_disc, pybamm.Function)
+        self.assertIsInstance(func_disc.children[0], pybamm.StateVector)
+        self.assertIsInstance(func_disc.children[1], pybamm.Scalar)
+
         # boundary value
         bv_left = pybamm.BoundaryValue(var_vec, "left")
         bv_left_disc = disc.process_symbol(bv_left)
diff --git a/tests/unit/test_expression_tree/test_function.py b/tests/unit/test_expression_tree/test_function.py
new file mode 100644
index 000000000..fd8b7351d
--- /dev/null
+++ b/tests/unit/test_expression_tree/test_function.py
@@ -0,0 +1,88 @@
+#
+# Tests for the Function classes
+#
+import pybamm
+
+import unittest
+import numpy as np
+import autograd.numpy as auto_np
+
+
+def test_function(arg):
+    return arg + arg
+
+
+def test_const_function():
+    return 1
+
+
+def test_multi_var_function(arg1, arg2):
+    return arg1 + arg2
+
+
+class TestFunction(unittest.TestCase):
+    def test_constant_functions(self):
+        d = pybamm.Scalar(6)
+        funcd = pybamm.Function(test_const_function, d)
+        self.assertEqual(funcd.evaluate(), 1)
+
+    def test_function_of_one_variable(self):
+        a = pybamm.Symbol("a")
+        funca = pybamm.Function(test_function, a)
+        self.assertEqual(funca.name, "function (test_function)")
+        self.assertEqual(funca.children[0].name, a.name)
+
+        b = pybamm.Scalar(1)
+        sina = pybamm.Function(np.sin, b)
+        self.assertEqual(sina.evaluate(), np.sin(1))
+        self.assertEqual(sina.name, "function ({})".format(np.sin.__name__))
+
+        c = pybamm.Vector(np.linspace(0, 1))
+        cosb = pybamm.Function(np.cos, c)
+        np.testing.assert_array_equal(cosb.evaluate(), np.cos(c.evaluate()))
+
+        var = pybamm.StateVector(slice(0, 100))
+        y = np.linspace(0, 1, 100)[:, np.newaxis]
+        logvar = pybamm.Function(np.log1p, var)
+        np.testing.assert_array_equal(logvar.evaluate(y=y), np.log1p(y))
+
+    def test_with_autograd(self):
+        a = pybamm.StateVector(slice(0, 1))
+        y = np.array([5])
+        func = pybamm.Function(test_function, a)
+        self.assertEqual((func).diff(a).evaluate(y=y), 2)
+        self.assertEqual((func).diff(func).evaluate(), 1)
+        func = pybamm.Function(auto_np.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.Function(auto_np.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)))
+
+    def test_function_of_multiple_variables(self):
+        a = pybamm.Variable("a")
+        b = pybamm.Parameter("b")
+        func = pybamm.Function(test_multi_var_function, a, b)
+        self.assertEqual(func.name, "function (test_multi_var_function)")
+        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(test_multi_var_function, 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)
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index 7a5476e6f..3f58d9e0e 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -39,10 +39,21 @@ class TestSimplify(unittest.TestCase):
         self.assertIsInstance((f).simplify(), pybamm.Scalar)
         self.assertEqual((f).simplify().evaluate(), math.sin(1))
 
+        def myfunction(x, y):
+            return x * y
+        f = pybamm.Function(myfunction, a, b)
+        self.assertIsInstance((f).simplify(), pybamm.Scalar)
+        self.assertEqual((f).simplify().evaluate(), 0)
+
         # FunctionParameter
         f = pybamm.FunctionParameter("function", b)
         self.assertIsInstance((f).simplify(), pybamm.FunctionParameter)
-        self.assertEqual((f).simplify().child.id, b.id)
+        self.assertEqual((f).simplify().children[0].id, b.id)
+
+        f = pybamm.FunctionParameter("function", a, b)
+        self.assertIsInstance((f).simplify(), pybamm.FunctionParameter)
+        self.assertEqual((f).simplify().children[0].id, a.id)
+        self.assertEqual((f).simplify().children[1].id, b.id)
 
         # Gradient
         self.assertIsInstance((pybamm.grad(a)).simplify(), pybamm.Scalar)
diff --git a/tests/unit/test_expression_tree/test_unary_operators.py b/tests/unit/test_expression_tree/test_unary_operators.py
index deeb5914d..797be033a 100644
--- a/tests/unit/test_expression_tree/test_unary_operators.py
+++ b/tests/unit/test_expression_tree/test_unary_operators.py
@@ -5,15 +5,6 @@ import pybamm
 
 import unittest
 import numpy as np
-import autograd.numpy as auto_np
-
-
-def test_function(arg):
-    return arg + arg
-
-
-def test_const_function():
-    return 1
 
 
 class TestUnaryOperators(unittest.TestCase):
@@ -43,30 +34,6 @@ class TestUnaryOperators(unittest.TestCase):
         absb = pybamm.AbsoluteValue(b)
         self.assertEqual(absb.evaluate(), 4)
 
-    def test_function(self):
-        a = pybamm.Symbol("a")
-        funca = pybamm.Function(test_function, a)
-        self.assertEqual(funca.name, "function (test_function)")
-        self.assertEqual(funca.children[0].name, a.name)
-
-        b = pybamm.Scalar(1)
-        sina = pybamm.Function(np.sin, b)
-        self.assertEqual(sina.evaluate(), np.sin(1))
-        self.assertEqual(sina.name, "function ({})".format(np.sin.__name__))
-
-        c = pybamm.Vector(np.linspace(0, 1))
-        cosb = pybamm.Function(np.cos, c)
-        np.testing.assert_array_equal(cosb.evaluate(), np.cos(c.evaluate()))
-
-        var = pybamm.StateVector(slice(0, 100))
-        y = np.linspace(0, 1, 100)[:, np.newaxis]
-        logvar = pybamm.Function(np.log1p, var)
-        np.testing.assert_array_equal(logvar.evaluate(y=y), np.log1p(y))
-
-        d = pybamm.Scalar(6)
-        funcd = pybamm.Function(test_const_function, d)
-        self.assertEqual(funcd.evaluate(), 1)
-
     def test_gradient(self):
         a = pybamm.Symbol("a")
         grad = pybamm.Gradient(a)
@@ -145,17 +112,6 @@ class TestUnaryOperators(unittest.TestCase):
         with self.assertRaises(pybamm.UndefinedOperationError):
             absa.diff(a)
 
-        # function: use autograd
-        func = pybamm.Function(test_function, a)
-        self.assertEqual((func).diff(a).evaluate(y=y), 2)
-        self.assertEqual((func).diff(func).evaluate(), 1)
-        func = pybamm.Function(auto_np.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.Function(auto_np.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)))
-
         # spatial operator (not implemented)
         spatial_a = pybamm.SpatialOperator("name", a)
         with self.assertRaises(NotImplementedError):
diff --git a/tests/unit/test_parameters/test_parameter_values.py b/tests/unit/test_parameters/test_parameter_values.py
index 5121588ed..dc70d9037 100644
--- a/tests/unit/test_parameters/test_parameter_values.py
+++ b/tests/unit/test_parameters/test_parameter_values.py
@@ -216,6 +216,33 @@ class TestParameterValues(unittest.TestCase):
         processed_diff_func = parameter_values.process_symbol(diff_func)
         self.assertEqual(processed_diff_func.evaluate(), 6)
 
+    def test_multi_var_function_with_parameters(self):
+        def D(a, b):
+            return a * np.exp(b)
+
+        parameter_values = pybamm.ParameterValues({"a": 3, "b": 0})
+        a = pybamm.Parameter("a")
+        b = pybamm.Parameter("b")
+        func = pybamm.Function(D, a, b)
+
+        processed_func = parameter_values.process_symbol(func)
+        self.assertIsInstance(processed_func, pybamm.Function)
+        self.assertEqual(processed_func.evaluate(), 3)
+
+    def test_multi_var_function_parameter(self):
+        def D(a, b):
+            return a * np.exp(b)
+
+        parameter_values = pybamm.ParameterValues({"a": 3, "b": 0, "Diffusivity": D})
+
+        a = pybamm.Parameter("a")
+        b = pybamm.Parameter("b")
+        func = pybamm.FunctionParameter("Diffusivity", a, b)
+
+        processed_func = parameter_values.process_symbol(func)
+        self.assertIsInstance(processed_func, pybamm.Function)
+        self.assertEqual(processed_func.evaluate(), 3)
+
     def test_process_complex_expression(self):
         var1 = pybamm.Variable("var1")
         var2 = pybamm.Variable("var2")
diff --git a/tests/unit/test_parameters/test_standard_parameters_lead_acid.py b/tests/unit/test_parameters/test_standard_parameters_lead_acid.py
index e0f7578c6..9cd79eb54 100644
--- a/tests/unit/test_parameters/test_standard_parameters_lead_acid.py
+++ b/tests/unit/test_parameters/test_standard_parameters_lead_acid.py
@@ -51,6 +51,7 @@ class TestStandardParametersLeadAcid(unittest.TestCase):
                 ),
             },
         )
+
         param_eval = {
             name: parameter_values.process_symbol(parameter).evaluate()
             for name, parameter in parameters.items()