Repository: mscroggs/symfem

Issue:
utility functions for matrix multiplication etc.
Hi everyone,

first of all thanks for this nice package!. It saved me considerable time already solving integrals, so thank you again. 

Problem

I generalized the example given in the demos

https://symfem.readthedocs.io/en/latest/demos/stiffness_matrix.html 

to derive the stiffness matrix for anisotropic heat conduction where the heat conduction becomes a matrix. I tried it with the provided MatrixFunction by symfem, but could not figure out how to combine this with the VectorFunction. My current solution uses lists to mimic matrices and then self-coded "matrix-multiplication" to get the code done in few comprehensible lines of code. It now looks like this

[symfem_quadrilateral_heat.txt](https://github.com/user-attachments/files/19074080/symfem_quadrilateral_heat.txt)

which in my opinion is still a bit unnecessarily complicated. In the end I basically just compute 

$\int_{\Omega} (\nabla N)^{T} \boldsymbol{k} \nabla N d \Omega$

where N is the vecotr of basis function, k the heat conductivity tensor and Omega the domain. So it would be neat, if one could formulate it like that. For linear elasticity this becomes more even annoying as one has 

$\int_{\Omega} \boldsymbol{B}^{T} \boldsymbol{C} \boldsymbol{B} d \Omega$

and the Bmatrix has then to be constructed by the user himself.

1. Utility functions for matrices
2. Util. functions for recurring matrices like the B-matrix for extracting the infinitesimal strains from the displacement vector.

If I overlooked any sympy/symfem features, sorry and thanks for your help in advance. If I can help with anything, I am happy to, but my experience with sympy and your library is limited.

Regards
Stefan

Hints from issue comments:
@mscroggs: The `@` operator should perform matrix-vector multiplication of a MatrixFunction and a VectorFunction, returning a VectorFunction. Updating your code to the following, it looks like it gives the same result:

```python
from sympy import symbols
from symfem.functions import VectorFunction, MatrixFunction
from symfem import create_element,create_reference
from symfem.symbols import x

def generate_constMatrix(ncol,nrow,name):
    """
    Generate matrix full of symbolic constants as a list of lists

    Parameters
    ----------
    ncol : int
        number of rows.
    nrow : int
        number of cols.
    name : str
        name to put in front of indices.

    Returns
    -------
    M : list
        list of list e. g. for a 2x2 [[M11,M12],[M21,M22]]
    """
    M = []
    for i in range(1,nrow+1):
        M.append(list(symbols(
            " ".join([name+str(i)+str(j) for j in range(1,ncol+1)]) ) ) )
    return MatrixFunction(M)  # I added MatrixFunction() here

# Define the vertices and triangles of the mesh
vertices = [(-1, -1), (1, -1), (1, 1), (-1, 1)]
rectangle = [0, 1, 3, 2]
# Create a Lagrange element
element = create_element("quadrilateral", "Lagrange", 1)
# Create a matrix of zeros with the correct shape
matrix = [[0 for i in range(4)] for j in range(4)]
# anisotropic heat conductivity or equivalent
K = generate_constMatrix(2,2,"k")
# Get the vertices of the triangle
vs = tuple(vertices[i] for i in rectangle)
# Create a reference cell with these vertices: this will be used
# to compute the integral over the triangle
ref = create_reference("quadrilateral", vertices=vs)
# Map the basis functions to the cell
basis = element.map_to_cell(vs)
#
for test_i, test_f in zip(rectangle, basis):
    #
    for trial_i, trial_f in zip(rectangle, basis):
        # Compute the integral of grad(u)-dot-grad(v) for each pair of basis
        # functions u and v. The second input (x) into `ref.integral` tells
        # symfem which variables to use in the integral.
        #integrand = test_f.grad(2).dot(trial_f.grad(2))
        f = test_f.grad(2)
        g = trial_f.grad(2)
        #f = VectorFunction([f.dot(VectorFunction([k11,k21])),
        #                    f.dot(VectorFunction([k12,k22]))])
        f = f @ K  # Here, I replaced the function call with @
        integrand = f.dot(g)
        matrix[test_i][trial_i] += integrand.integral(ref, x)

print(matrix)
```

I believe that `MatrixFunction @ MatrixFunction` should also work, but if you have any issues with what you're doing with elasticity, let me know.@stefanhiemer: Thanks for the quick help. I edited the example further 

[symfem_quadrilateral_heat.txt](https://github.com/user-attachments/files/19090304/symfem_quadrilateral_heat.txt)

to come closer to the goal of basically just computing the integral in matrix form:

$\int_{\Omega} (\nabla N)^{T} \boldsymbol{k} \nabla N d \Omega$

It works, but two things are still a little awkward. 
1. I try to get the gradients of the basis functions via `gradN = VectorFunction(basis).jacobian(ndim)` which just raises a `NotImplementedError`. I guess as the name indicates, you did not have the time yet to implement it. I fix it via this line 
`gradN = MatrixFunction([VectorFunction([basis[i] for i in rectangle]).diff(var) for var in ["x","y","z"][:ndim]])`
The question: Have I again missed something, or is this for the moment the way to go?
2. More or less same story: I calculate the final matrix here 
ideally via `integrand.integral(ref,x)` which again raises a `NotImplementedError`
I now do it this way: 
`for i,j in product(range(integrand.shape[0]),range(integrand.shape[1])):`
`        matrix[i][j] += integrand[i,j].integral(ref, x)`
Same question: is this the method to go on for now?


Thanks and Regards
Stefan@mscroggs: Currently, your method is the best you can do, but I've opened PRs to make these easier.

#297 will help with (2): it makes `VectorFunction.integral` and `MatrixFunction.integral` do componentwise integrals.

#299 will help with (1): it makes `VectorFunction.grad` return a matrix function, where each row is the grad of one component of the vector function.

I'll get both those merged once the tests pass. When the second is merged, it'll auto-close this issue. Feel free to reopen or make a new issue if you have further issues

Solution patch:
diff --git a/.github/workflows/defelement.yml b/.github/workflows/defelement.yml
index 8bf45ee..67a07d8 100644
--- a/.github/workflows/defelement.yml
+++ b/.github/workflows/defelement.yml
@@ -27,7 +27,7 @@ jobs:
         with:
           path: ./defelement
           repository: DefElement/DefElement
-          ref: main
+          ref: mscroggs/AnyFunction->Function
       - name: Install requirements
         run: |
           cd defelement
diff --git a/symfem/basis_functions.py b/symfem/basis_functions.py
index 447aa53..39af19b 100644
--- a/symfem/basis_functions.py
+++ b/symfem/basis_functions.py
@@ -9,7 +9,7 @@ import sympy
 
 import symfem
 from symfem.functions import (
-    AnyFunction,
+    Function,
     FunctionInput,
     ScalarFunction,
     SympyFormat,
@@ -22,7 +22,7 @@ from symfem.symbols import AxisVariables, AxisVariablesNotSingle, t, x
 __all__ = ["BasisFunction", "SubbedBasisFunction"]
 
 
-class BasisFunction(AnyFunction):
+class BasisFunction(Function):
     """A basis function of a finite element.
 
     This basis function can be used before the element's basis functions have been computed. When
@@ -40,7 +40,7 @@ class BasisFunction(AnyFunction):
         super().__init__(scalar=scalar, vector=vector, matrix=matrix)
 
     @abstractmethod
-    def get_function(self) -> AnyFunction:
+    def get_function(self) -> Function:
         """Get the actual basis function.
 
         Returns:
@@ -48,7 +48,7 @@ class BasisFunction(AnyFunction):
         """
         pass
 
-    def __add__(self, other: typing.Any) -> AnyFunction:
+    def __add__(self, other: typing.Any) -> Function:
         """Add.
 
         Args:
@@ -59,7 +59,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__add__(other)
 
-    def __radd__(self, other: typing.Any) -> AnyFunction:
+    def __radd__(self, other: typing.Any) -> Function:
         """Add.
 
         Args:
@@ -70,7 +70,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__radd__(other)
 
-    def __sub__(self, other: typing.Any) -> AnyFunction:
+    def __sub__(self, other: typing.Any) -> Function:
         """Subtract.
 
         Args:
@@ -81,7 +81,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__sub__(other)
 
-    def __rsub__(self, other: typing.Any) -> AnyFunction:
+    def __rsub__(self, other: typing.Any) -> Function:
         """Subtract.
 
         Args:
@@ -92,7 +92,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__rsub__(other)
 
-    def __neg__(self) -> AnyFunction:
+    def __neg__(self) -> Function:
         """Negate.
 
         Returns:
@@ -100,7 +100,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__neg__()
 
-    def __truediv__(self, other: typing.Any) -> AnyFunction:
+    def __truediv__(self, other: typing.Any) -> Function:
         """Divide.
 
         Args:
@@ -111,7 +111,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__truediv__(other)
 
-    def __rtruediv__(self, other: typing.Any) -> AnyFunction:
+    def __rtruediv__(self, other: typing.Any) -> Function:
         """Divide.
 
         Args:
@@ -122,7 +122,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__rtruediv__(other)
 
-    def __mul__(self, other: typing.Any) -> AnyFunction:
+    def __mul__(self, other: typing.Any) -> Function:
         """Multiply.
 
         Args:
@@ -133,7 +133,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__mul__(other)
 
-    def __rmul__(self, other: typing.Any) -> AnyFunction:
+    def __rmul__(self, other: typing.Any) -> Function:
         """Multiply.
 
         Args:
@@ -144,7 +144,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__rmul__(other)
 
-    def __matmul__(self, other: typing.Any) -> AnyFunction:
+    def __matmul__(self, other: typing.Any) -> Function:
         """Multiply.
 
         Args:
@@ -155,7 +155,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__matmul__(other)
 
-    def __rmatmul__(self, other: typing.Any) -> AnyFunction:
+    def __rmatmul__(self, other: typing.Any) -> Function:
         """Multiply.
 
         Args:
@@ -166,7 +166,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().__rmatmul__(other)
 
-    def __pow__(self, other: typing.Any) -> AnyFunction:
+    def __pow__(self, other: typing.Any) -> Function:
         """Raise to a power.
 
         Args:
@@ -193,7 +193,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().as_tex()
 
-    def diff(self, variable: sympy.core.symbol.Symbol) -> AnyFunction:
+    def diff(self, variable: sympy.core.symbol.Symbol) -> Function:
         """Differentiate the function.
 
         Args:
@@ -204,7 +204,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().diff(variable)
 
-    def directional_derivative(self, direction: PointType) -> AnyFunction:
+    def directional_derivative(self, direction: PointType) -> Function:
         """Compute a directional derivative.
 
         Args:
@@ -215,7 +215,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().directional_derivative(direction)
 
-    def jacobian_component(self, component: typing.Tuple[int, int]) -> AnyFunction:
+    def jacobian_component(self, component: typing.Tuple[int, int]) -> Function:
         """Compute a component of the jacobian.
 
         Args:
@@ -226,7 +226,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().jacobian_component(component)
 
-    def jacobian(self, dim: int) -> AnyFunction:
+    def jacobian(self, dim: int) -> Function:
         """Compute the jacobian.
 
         Args:
@@ -237,7 +237,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().jacobian(dim)
 
-    def dot(self, other_in: FunctionInput) -> AnyFunction:
+    def dot(self, other_in: FunctionInput) -> Function:
         """Compute the dot product with another function.
 
         Args:
@@ -248,7 +248,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().dot(other_in)
 
-    def cross(self, other_in: FunctionInput) -> AnyFunction:
+    def cross(self, other_in: FunctionInput) -> Function:
         """Compute the cross product with another function.
 
         Args:
@@ -259,7 +259,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().cross(other_in)
 
-    def div(self) -> AnyFunction:
+    def div(self) -> Function:
         """Compute the divergence of the function.
 
         Returns:
@@ -267,7 +267,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().div()
 
-    def grad(self, dim: int) -> AnyFunction:
+    def grad(self, dim: int) -> Function:
         """Compute the gradient of the function.
 
         Returns:
@@ -275,7 +275,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().grad(dim)
 
-    def curl(self) -> AnyFunction:
+    def curl(self) -> Function:
         """Compute the curl of the function.
 
         Returns:
@@ -296,7 +296,7 @@ class BasisFunction(AnyFunction):
         domain: Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -321,7 +321,7 @@ class BasisFunction(AnyFunction):
         """
         return SubbedBasisFunction(self, vars, values)
 
-    def __getitem__(self, key) -> AnyFunction:
+    def __getitem__(self, key) -> Function:
         """Forward all other function calls to symbolic function."""
         return self.get_function().__getitem__(key)
 
@@ -353,7 +353,7 @@ class BasisFunction(AnyFunction):
             raise AttributeError(f"'{self.__class__.__name__}' object has no attribute 'transpose'")
         return self.get_function().transpose()
 
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
@@ -361,7 +361,7 @@ class BasisFunction(AnyFunction):
         """
         return self.get_function().with_floats()
 
-    def __iter__(self) -> typing.Iterator[AnyFunction]:
+    def __iter__(self) -> typing.Iterator[Function]:
         """Iterate through components of vector function."""
         f = self.get_function()
         return f.__iter__()
@@ -397,7 +397,7 @@ class SubbedBasisFunction(BasisFunction):
         self._vars = vars
         self._values = values
 
-    def get_function(self) -> AnyFunction:
+    def get_function(self) -> Function:
         """Return the symbolic function.
 
         Returns:
diff --git a/symfem/elements/bernstein.py b/symfem/elements/bernstein.py
index fbe9d57..ef78ead 100644
--- a/symfem/elements/bernstein.py
+++ b/symfem/elements/bernstein.py
@@ -11,7 +11,7 @@ import sympy
 
 from symfem.finite_element import CiarletElement
 from symfem.functionals import BaseFunctional, ListOfFunctionals, PointEvaluation
-from symfem.functions import AnyFunction, FunctionInput
+from symfem.functions import Function, FunctionInput
 from symfem.geometry import PointType
 from symfem.polynomials import orthogonal_basis, polynomial_set_1d
 from symfem.references import Reference
@@ -137,7 +137,7 @@ class BernsteinFunctional(BaseFunctional):
         """
         return self.ref.sub_entity(*self.entity).midpoint()
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply the functional to a function.
 
         Args:
diff --git a/symfem/elements/dual.py b/symfem/elements/dual.py
index ecb5be9..8291fc2 100644
--- a/symfem/elements/dual.py
+++ b/symfem/elements/dual.py
@@ -9,7 +9,7 @@ import typing
 import sympy
 
 from symfem.finite_element import FiniteElement
-from symfem.functions import AnyFunction, FunctionInput, VectorFunction
+from symfem.functions import Function, FunctionInput, VectorFunction
 from symfem.geometry import PointType, SetOfPoints, SetOfPointsInput
 from symfem.piecewise_functions import PiecewiseFunction
 from symfem.references import DualPolygon, NonDefaultReferenceError
@@ -56,11 +56,11 @@ class DualCiarletElement(FiniteElement):
         super().__init__(
             reference, order, len(dual_coefficients), domain_dim, range_dim, range_shape=range_shape
         )
-        self._basis_functions: typing.Union[typing.List[AnyFunction], None] = None
+        self._basis_functions: typing.Union[typing.List[Function], None] = None
         self._dof_entities = dof_entities
         self._dof_directions = dof_directions
 
-    def get_polynomial_basis(self, reshape: bool = True) -> typing.List[AnyFunction]:
+    def get_polynomial_basis(self, reshape: bool = True) -> typing.List[Function]:
         """Get the symbolic polynomial basis for the element.
 
         Returns:
@@ -81,9 +81,7 @@ class DualCiarletElement(FiniteElement):
         """
         raise ValueError("Dual matrix not supported for barycentric dual elements.")
 
-    def get_basis_functions(
-        self, use_tensor_factorisation: bool = False
-    ) -> typing.List[AnyFunction]:
+    def get_basis_functions(self, use_tensor_factorisation: bool = False) -> typing.List[Function]:
         """Get the basis functions of the element.
 
         Args:
@@ -97,7 +95,7 @@ class DualCiarletElement(FiniteElement):
         if self._basis_functions is None:
             from symfem import create_element
 
-            bfs: typing.List[AnyFunction] = []
+            bfs: typing.List[Function] = []
             sub_e = create_element("triangle", self.fine_space, self.order)
             for coeff_list in self.dual_coefficients:
                 v0 = self.reference.origin
@@ -194,10 +192,10 @@ class DualCiarletElement(FiniteElement):
     def map_to_cell(
         self,
         vertices_in: SetOfPointsInput,
-        basis: typing.Optional[typing.List[AnyFunction]] = None,
+        basis: typing.Optional[typing.List[Function]] = None,
         forward_map: typing.Optional[PointType] = None,
         inverse_map: typing.Optional[PointType] = None,
-    ) -> typing.List[AnyFunction]:
+    ) -> typing.List[Function]:
         """Map the basis onto a cell using the appropriate mapping for the element.
 
         Args:
diff --git a/symfem/elements/guzman_neilan.py b/symfem/elements/guzman_neilan.py
index a7060e7..4704f38 100644
--- a/symfem/elements/guzman_neilan.py
+++ b/symfem/elements/guzman_neilan.py
@@ -13,7 +13,7 @@ from symfem.elements.bernardi_raugel import BernardiRaugel
 from symfem.elements.lagrange import Lagrange, VectorLagrange
 from symfem.finite_element import CiarletElement
 from symfem.functionals import DotPointEvaluation, ListOfFunctionals, NormalIntegralMoment
-from symfem.functions import AnyFunction, FunctionInput, ScalarFunction, VectorFunction
+from symfem.functions import Function, FunctionInput, ScalarFunction, VectorFunction
 from symfem.geometry import SetOfPoints, SetOfPointsInput
 from symfem.moments import make_integral_moment_dofs
 from symfem.piecewise_functions import PiecewiseFunction
@@ -23,7 +23,7 @@ from symfem.symbols import t, x
 __all__ = ["GuzmanNeilanFirstKind", "GuzmanNeilanSecondKind", "make_piecewise_lagrange"]
 
 
-def poly(reference: Reference, k: int) -> typing.List[AnyFunction]:
+def poly(reference: Reference, k: int) -> typing.List[Function]:
     """Generate the P^perp polynomial set."""
     if k < 2:
         if reference.name == "triangle":
diff --git a/symfem/finite_element.py b/symfem/finite_element.py
index 96d5710..0eab258 100644
--- a/symfem/finite_element.py
+++ b/symfem/finite_element.py
@@ -13,7 +13,7 @@ from symfem.basis_functions import BasisFunction
 from symfem.caching import load_cached_matrix, save_cached_matrix
 from symfem.functionals import ListOfFunctionals
 from symfem.functions import (
-    AnyFunction,
+    Function,
     FunctionInput,
     ScalarFunction,
     VectorFunction,
@@ -57,7 +57,7 @@ class NoTensorProduct(Exception):
 class FiniteElement(ABC):
     """Abstract finite element."""
 
-    _float_basis_functions: typing.Union[None, typing.List[AnyFunction]]
+    _float_basis_functions: typing.Union[None, typing.List[Function]]
     _value_scale: typing.Union[None, sympy.core.expr.Expr]
 
     def __init__(
@@ -212,9 +212,7 @@ class FiniteElement(ABC):
         """
 
     @abstractmethod
-    def get_basis_functions(
-        self, use_tensor_factorisation: bool = False
-    ) -> typing.List[AnyFunction]:
+    def get_basis_functions(self, use_tensor_factorisation: bool = False) -> typing.List[Function]:
         """Get the basis functions of the element.
 
         Args:
@@ -329,10 +327,10 @@ class FiniteElement(ABC):
     def map_to_cell(
         self,
         vertices_in: SetOfPointsInput,
-        basis: typing.Optional[typing.List[AnyFunction]] = None,
+        basis: typing.Optional[typing.List[Function]] = None,
         forward_map: typing.Optional[PointType] = None,
         inverse_map: typing.Optional[PointType] = None,
-    ) -> typing.List[AnyFunction]:
+    ) -> typing.List[Function]:
         """Map the basis onto a cell using the appropriate mapping for the element.
 
         Args:
@@ -346,7 +344,7 @@ class FiniteElement(ABC):
         """
 
     @abstractmethod
-    def get_polynomial_basis(self) -> typing.List[AnyFunction]:
+    def get_polynomial_basis(self) -> typing.List[Function]:
         """Get the symbolic polynomial basis for the element.
 
         Returns:
@@ -541,7 +539,7 @@ class FiniteElement(ABC):
         """
         raise NoTensorProduct()
 
-    def _get_basis_functions_tensor(self) -> typing.List[AnyFunction]:
+    def _get_basis_functions_tensor(self) -> typing.List[Function]:
         """Compute the basis functions using the space's tensor product factorisation.
 
         Returns:
@@ -617,9 +615,9 @@ class CiarletElement(FiniteElement):
         assert len(basis) == len(dofs)
         self._basis = [parse_function_input(b) for b in basis]
         for b in self._basis:
-            assert isinstance(b, AnyFunction)
+            assert isinstance(b, Function)
         self.dofs = dofs
-        self._basis_functions: typing.Union[typing.List[AnyFunction], None] = None
+        self._basis_functions: typing.Union[typing.List[Function], None] = None
         self._maximum_degree = None
         if reference.name == "pyramid":
             self._maximum_degree = order
@@ -660,7 +658,7 @@ class CiarletElement(FiniteElement):
         """
         return [d.entity for d in self.dofs]
 
-    def get_polynomial_basis(self) -> typing.List[AnyFunction]:
+    def get_polynomial_basis(self) -> typing.List[Function]:
         """Get the symbolic polynomial basis for the element.
 
         Returns:
@@ -710,9 +708,7 @@ class CiarletElement(FiniteElement):
             else:
                 return mat
 
-    def get_basis_functions(
-        self, use_tensor_factorisation: bool = False
-    ) -> typing.List[AnyFunction]:
+    def get_basis_functions(self, use_tensor_factorisation: bool = False) -> typing.List[Function]:
         """Get the basis functions of the element.
 
         Args:
@@ -727,7 +723,7 @@ class CiarletElement(FiniteElement):
             else:
                 minv = self.get_dual_matrix(inverse=True)
 
-                sfs: typing.List[AnyFunction] = []
+                sfs: typing.List[Function] = []
                 pb = self.get_polynomial_basis()
                 for i, dof in enumerate(self.dofs):
                     sf = ScalarFunction(minv[i, 0]) * pb[0]
@@ -785,10 +781,10 @@ class CiarletElement(FiniteElement):
     def map_to_cell(
         self,
         vertices_in: SetOfPointsInput,
-        basis: typing.Optional[typing.List[AnyFunction]] = None,
+        basis: typing.Optional[typing.List[Function]] = None,
         forward_map: typing.Optional[PointType] = None,
         inverse_map: typing.Optional[PointType] = None,
-    ) -> typing.List[AnyFunction]:
+    ) -> typing.List[Function]:
         """Map the basis onto a cell using the appropriate mapping for the element.
 
         Args:
@@ -809,7 +805,7 @@ class CiarletElement(FiniteElement):
             inverse_map = self.reference.get_inverse_map_to(vertices)
 
         try:
-            functions: typing.List[AnyFunction] = [ScalarFunction(0) for f in basis]
+            functions: typing.List[Function] = [ScalarFunction(0) for f in basis]
             for dim in range(self.reference.tdim + 1):
                 for e in range(self.reference.sub_entity_count(dim)):
                     entity_dofs = self.entity_dofs(dim, e)
@@ -882,7 +878,7 @@ class CiarletElement(FiniteElement):
 class DirectElement(FiniteElement):
     """Finite element defined directly."""
 
-    _basis_functions: typing.List[AnyFunction]
+    _basis_functions: typing.List[Function]
 
     def __init__(
         self,
@@ -994,9 +990,7 @@ class DirectElement(FiniteElement):
         """
         return self._basis_entities
 
-    def get_basis_functions(
-        self, use_tensor_factorisation: bool = False
-    ) -> typing.List[AnyFunction]:
+    def get_basis_functions(self, use_tensor_factorisation: bool = False) -> typing.List[Function]:
         """Get the basis functions of the element.
 
         Args:
@@ -1013,10 +1007,10 @@ class DirectElement(FiniteElement):
     def map_to_cell(
         self,
         vertices_in: SetOfPointsInput,
-        basis: typing.Optional[typing.List[AnyFunction]] = None,
+        basis: typing.Optional[typing.List[Function]] = None,
         forward_map: typing.Optional[PointType] = None,
         inverse_map: typing.Optional[PointType] = None,
-    ) -> typing.List[AnyFunction]:
+    ) -> typing.List[Function]:
         """Map the basis onto a cell using the appropriate mapping for the element.
 
         Args:
@@ -1035,7 +1029,7 @@ class DirectElement(FiniteElement):
         #                   **self.init_kwargs())
         # return e.get_basis_functions()
 
-    def get_polynomial_basis(self) -> typing.List[AnyFunction]:
+    def get_polynomial_basis(self) -> typing.List[Function]:
         """Get the symbolic polynomial basis for the element.
 
         Returns:
@@ -1097,7 +1091,7 @@ class DirectElement(FiniteElement):
 class EnrichedElement(FiniteElement):
     """Finite element defined directly."""
 
-    _basis_functions: typing.Optional[typing.List[AnyFunction]]
+    _basis_functions: typing.Optional[typing.List[Function]]
 
     def __init__(
         self,
@@ -1180,9 +1174,7 @@ class EnrichedElement(FiniteElement):
             positions += e.dof_entities()
         return positions
 
-    def get_basis_functions(
-        self, use_tensor_factorisation: bool = False
-    ) -> typing.List[AnyFunction]:
+    def get_basis_functions(self, use_tensor_factorisation: bool = False) -> typing.List[Function]:
         """Get the basis functions of the element.
 
         Args:
@@ -1204,10 +1196,10 @@ class EnrichedElement(FiniteElement):
     def map_to_cell(
         self,
         vertices_in: SetOfPointsInput,
-        basis: typing.Optional[typing.List[AnyFunction]] = None,
+        basis: typing.Optional[typing.List[Function]] = None,
         forward_map: typing.Optional[PointType] = None,
         inverse_map: typing.Optional[PointType] = None,
-    ) -> typing.List[AnyFunction]:
+    ) -> typing.List[Function]:
         """Map the basis onto a cell using the appropriate mapping for the element.
 
         Args:
@@ -1224,7 +1216,7 @@ class EnrichedElement(FiniteElement):
             out += e.map_to_cell(vertices_in, basis, forward_map, inverse_map)
         return out
 
-    def get_polynomial_basis(self) -> typing.List[AnyFunction]:
+    def get_polynomial_basis(self) -> typing.List[Function]:
         """Get the symbolic polynomial basis for the element.
 
         Returns:
@@ -1263,7 +1255,7 @@ class ElementBasisFunction(BasisFunction):
         self.element = element
         self.n = n
 
-    def get_function(self) -> AnyFunction:
+    def get_function(self) -> Function:
         """Get the actual basis function.
 
         Returns:
diff --git a/symfem/functionals.py b/symfem/functionals.py
index 932b9fe..9556495 100644
--- a/symfem/functionals.py
+++ b/symfem/functionals.py
@@ -7,7 +7,7 @@ import sympy
 
 from symfem import mappings
 from symfem.functions import (
-    AnyFunction,
+    Function,
     FunctionInput,
     ScalarFunction,
     VectorFunction,
@@ -115,8 +115,8 @@ class BaseFunctional(ABC):
         return self.entity[1]
 
     def perform_mapping(
-        self, fs: typing.List[AnyFunction], map: PointType, inverse_map: PointType
-    ) -> typing.List[AnyFunction]:
+        self, fs: typing.List[Function], map: PointType, inverse_map: PointType
+    ) -> typing.List[Function]:
         """Map functions to a cell.
 
         Args:
@@ -163,7 +163,7 @@ class BaseFunctional(ABC):
         return None
 
     def eval(
-        self, function: AnyFunction, symbolic: bool = True
+        self, function: Function, symbolic: bool = True
     ) -> typing.Union[ScalarFunction, float]:
         """Apply to the functional to a function.
 
@@ -202,7 +202,7 @@ class BaseFunctional(ABC):
         midpoint = self.reference.sub_entity(*self.entity).midpoint()
         return tuple(m + sympy.Rational(7, 10) * (p - m) for m, p in zip(midpoint, point))
 
-    def eval_symbolic(self, function: AnyFunction) -> ScalarFunction:
+    def eval_symbolic(self, function: Function) -> ScalarFunction:
         """Symbolically apply the functional to a function.
 
         Args:
@@ -219,7 +219,7 @@ class BaseFunctional(ABC):
         return e
 
     @abstractmethod
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -264,7 +264,7 @@ class PointEvaluation(BaseFunctional):
         self.point = parse_function_input(point_in)
         assert self.point.is_vector
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -330,7 +330,7 @@ class WeightedPointEvaluation(BaseFunctional):
         assert self.point.is_vector
         self.weight = weight
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -399,7 +399,7 @@ class DerivativePointEvaluation(BaseFunctional):
         assert self.point.is_vector
         self.derivative = derivative
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -432,8 +432,8 @@ class DerivativePointEvaluation(BaseFunctional):
         return self.dof_point()
 
     def perform_mapping(
-        self, fs: typing.List[AnyFunction], map: PointType, inverse_map: PointType
-    ) -> typing.List[AnyFunction]:
+        self, fs: typing.List[Function], map: PointType, inverse_map: PointType
+    ) -> typing.List[Function]:
         """Map functions to a cell.
 
         Args:
@@ -499,7 +499,7 @@ class PointDirectionalDerivativeEvaluation(BaseFunctional):
         assert self.point.is_vector
         self.dir = parse_function_input(direction_in)
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -628,7 +628,7 @@ class PointComponentSecondDerivativeEvaluation(BaseFunctional):
         assert self.point.is_vector
         self.component = component
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -703,7 +703,7 @@ class PointInnerProduct(BaseFunctional):
         assert self.lvec.is_vector
         assert self.rvec.is_vector
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -792,7 +792,7 @@ class DotPointEvaluation(BaseFunctional):
         assert self.point.is_vector
         self.vector = parse_function_input(vector_in)
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -876,7 +876,7 @@ class PointDivergenceEvaluation(BaseFunctional):
         self.point = parse_function_input(point_in)
         assert self.point.is_vector
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -938,7 +938,7 @@ class PointDivergenceEvaluation(BaseFunctional):
 class IntegralAgainst(BaseFunctional):
     """An integral against a function."""
 
-    f: AnyFunction
+    f: Function
 
     def __init__(
         self,
@@ -987,7 +987,7 @@ class IntegralAgainst(BaseFunctional):
         """
         return tuple(sympy.Rational(sum(i), len(i)) for i in zip(*self.integral_domain.vertices))
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1004,7 +1004,7 @@ class IntegralAgainst(BaseFunctional):
         integrand = self.dot(v1)
         return integrand.integral(self.integral_domain)
 
-    def dot(self, function: AnyFunction) -> ScalarFunction:
+    def dot(self, function: Function) -> ScalarFunction:
         """Dot a function with the moment function.
 
         Args:
@@ -1071,7 +1071,7 @@ class IntegralOfDivergenceAgainst(BaseFunctional):
         """
         return tuple(sympy.Rational(sum(i), len(i)) for i in zip(*self.integral_domain.vertices))
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1154,7 +1154,7 @@ class IntegralOfDirectionalMultiderivative(BaseFunctional):
         """
         return tuple(sympy.Rational(sum(i), len(i)) for i in zip(*self.integral_domain.vertices))
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1174,8 +1174,8 @@ class IntegralOfDirectionalMultiderivative(BaseFunctional):
         return integrand.integral(self.integral_domain)
 
     def perform_mapping(
-        self, fs: typing.List[AnyFunction], map: PointType, inverse_map: PointType
-    ) -> typing.List[AnyFunction]:
+        self, fs: typing.List[Function], map: PointType, inverse_map: PointType
+    ) -> typing.List[Function]:
         """Map functions to a cell.
 
         Args:
@@ -1220,7 +1220,7 @@ class IntegralOfDirectionalMultiderivative(BaseFunctional):
 class IntegralMoment(BaseFunctional):
     """An integral moment."""
 
-    f: AnyFunction
+    f: Function
 
     def __init__(
         self,
@@ -1277,7 +1277,7 @@ class IntegralMoment(BaseFunctional):
             )
             self.f *= id_def.volume() / self.integral_domain.volume()
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1295,7 +1295,7 @@ class IntegralMoment(BaseFunctional):
         integrand = self.dot(v1)
         return integrand.integral(self.integral_domain)
 
-    def dot(self, function: AnyFunction) -> ScalarFunction:
+    def dot(self, function: Function) -> ScalarFunction:
         """Dot a function with the moment function.
 
         Args:
@@ -1410,7 +1410,7 @@ class DerivativeIntegralMoment(IntegralMoment):
         super().__init__(reference, f, dof, entity=entity, mapping=mapping)
         self.dot_with = parse_function_input(dot_with_in)
 
-    def dot(self, function: AnyFunction) -> ScalarFunction:
+    def dot(self, function: Function) -> ScalarFunction:
         """Dot a function with the moment function.
 
         Args:
@@ -1433,7 +1433,7 @@ class DerivativeIntegralMoment(IntegralMoment):
         assert isinstance(dw, tuple)
         return dw
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1477,7 +1477,7 @@ class DivergenceIntegralMoment(IntegralMoment):
         assert f.is_scalar
         super().__init__(reference, f, dof, entity=entity, mapping=mapping)
 
-    def _eval_symbolic(self, function: AnyFunction) -> AnyFunction:
+    def _eval_symbolic(self, function: Function) -> Function:
         """Apply to the functional to a function.
 
         Args:
@@ -1723,7 +1723,7 @@ class InnerProductIntegralMoment(IntegralMoment):
         self.inner_with_left = inner_with_left
         self.inner_with_right = inner_with_right
 
-    def dot(self, function: AnyFunction) -> ScalarFunction:
+    def dot(self, function: Function) -> ScalarFunction:
         """Take the inner product of a function with the moment direction.
 
         Args:
diff --git a/symfem/functions.py b/symfem/functions.py
index b808440..41f70c7 100644
--- a/symfem/functions.py
+++ b/symfem/functions.py
@@ -15,7 +15,7 @@ __all__ = [
     "SingleSympyFormat",
     "SympyFormat",
     "ValuesToSubstitute",
-    "AnyFunction",
+    "Function",
     "ScalarFunction",
     "VectorFunction",
     "MatrixFunction",
@@ -47,7 +47,7 @@ def _to_sympy_format(item: typing.Any) -> SympyFormat:
     Returns:
         The item in Sympy format expected by functions
     """
-    if isinstance(item, AnyFunction):
+    if isinstance(item, Function):
         return item.as_sympy()
 
     if isinstance(item, int):
@@ -108,7 +108,7 @@ def _check_equal(first: SympyFormat, second: SympyFormat) -> bool:
     return False
 
 
-class AnyFunction(ABC):
+class Function(ABC):
     """A function."""
 
     def __init__(self, scalar: bool = False, vector: bool = False, matrix: bool = False):
@@ -189,7 +189,7 @@ class AnyFunction(ABC):
         """
 
     @abstractmethod
-    def subs(self, vars: AxisVariables, values: typing.Union[AnyFunction, _ValuesToSubstitute]):
+    def subs(self, vars: AxisVariables, values: typing.Union[Function, _ValuesToSubstitute]):
         """Substitute values into the function.
 
         Args:
@@ -304,7 +304,7 @@ class AnyFunction(ABC):
         domain: symfem.references.Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -317,7 +317,7 @@ class AnyFunction(ABC):
         """
 
     @abstractmethod
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
@@ -338,7 +338,7 @@ class AnyFunction(ABC):
             A version the function with floats as coefficients
         """
 
-    def __iter__(self) -> typing.Iterator[AnyFunction]:
+    def __iter__(self) -> typing.Iterator[Function]:
         """Iterate through components of vector function."""
         raise TypeError(f"'{self.__class__.__name__}' object is not iterable")
 
@@ -477,7 +477,7 @@ class AnyFunction(ABC):
         """Compute the determinant."""
         raise TypeError(f"object of type '{self.__class__.__name__}' has no len()")
 
-    def __getitem__(self, key) -> AnyFunction:
+    def __getitem__(self, key) -> Function:
         """Get a component or slice of the function."""
         raise ValueError(f"'{self.__class__.__name__}' object is not subscriptable")
 
@@ -522,10 +522,10 @@ class AnyFunction(ABC):
         return not self.__eq__(other)
 
 
-ValuesToSubstitute = typing.Union[AnyFunction, _ValuesToSubstitute]
+ValuesToSubstitute = typing.Union[Function, _ValuesToSubstitute]
 
 
-class ScalarFunction(AnyFunction):
+class ScalarFunction(Function):
     """A scalar-valued function."""
 
     _f: sympy.core.expr.Expr
@@ -657,7 +657,7 @@ class ScalarFunction(AnyFunction):
             The substituted function
         """
         subbed = self._f
-        if isinstance(values, AnyFunction):
+        if isinstance(values, Function):
             values = values.as_sympy()
         if isinstance(vars, (list, tuple)):
             assert isinstance(values, (list, tuple))
@@ -733,7 +733,7 @@ class ScalarFunction(AnyFunction):
         if isinstance(other, ScalarFunction):
             return self * other
 
-        if isinstance(other, AnyFunction) and other.is_scalar:
+        if isinstance(other, Function) and other.is_scalar:
             return other.dot(self)
 
         raise NotImplementedError()
@@ -786,7 +786,7 @@ class ScalarFunction(AnyFunction):
         domain: symfem.references.Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -889,7 +889,7 @@ class ScalarFunction(AnyFunction):
         for s, m1, m2, e in self._plot_beziers[(reference, n)]:
             img.add_bezier(s, m1, m2, e, colors.ORANGE)
 
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
@@ -932,7 +932,7 @@ class ScalarFunction(AnyFunction):
             raise ValueError(f"Unrecognised cell: {cell.name}")
 
 
-class VectorFunction(AnyFunction):
+class VectorFunction(Function):
     """A vector-valued function."""
 
     _vec: tuple[ScalarFunction, ...]
@@ -940,8 +940,8 @@ class VectorFunction(AnyFunction):
     def __init__(
         self,
         vec: typing.Union[
-            typing.Tuple[typing.Union[AnyFunction, int, sympy.core.expr.Expr], ...],
-            typing.List[typing.Union[AnyFunction, int, sympy.core.expr.Expr]],
+            typing.Tuple[typing.Union[Function, int, sympy.core.expr.Expr], ...],
+            typing.List[typing.Union[Function, int, sympy.core.expr.Expr]],
         ],
     ):
         """Create a vector-valued function.
@@ -954,7 +954,7 @@ class VectorFunction(AnyFunction):
         super().__init__(vector=True)
         vec_l = []
         for i in vec:
-            if isinstance(i, AnyFunction):
+            if isinstance(i, Function):
                 if isinstance(i, BasisFunction):
                     i = i.get_function()
                 assert isinstance(i, ScalarFunction)
@@ -1176,7 +1176,7 @@ class VectorFunction(AnyFunction):
                 out += i._f * j._f
             return ScalarFunction(out)
 
-        if isinstance(other, AnyFunction) and other.is_vector:
+        if isinstance(other, Function) and other.is_vector:
             return other.dot(self)
 
         raise NotImplementedError()
@@ -1254,7 +1254,7 @@ class VectorFunction(AnyFunction):
         domain: symfem.references.Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -1265,7 +1265,7 @@ class VectorFunction(AnyFunction):
         Returns:
             The integral
         """
-        raise NotImplementedError()
+        return VectorFunction([f.integral(domain, vars, dummy_vars) for f in self._vec])
 
     def __iter__(self):
         """Get iterable."""
@@ -1315,7 +1315,7 @@ class VectorFunction(AnyFunction):
         for p, q, size in self._plot_arrows[(reference, n)]:
             img.add_arrow(p, q, colors.ORANGE, size)
 
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
@@ -1338,7 +1338,7 @@ class VectorFunction(AnyFunction):
         return max(f.maximum_degree(cell) for f in self._vec)
 
 
-class MatrixFunction(AnyFunction):
+class MatrixFunction(Function):
     """A matrix-valued function."""
 
     _mat: typing.Tuple[typing.Tuple[ScalarFunction, ...], ...]
@@ -1346,12 +1346,10 @@ class MatrixFunction(AnyFunction):
     def __init__(
         self,
         mat: typing.Union[
-            typing.Tuple[
-                typing.Tuple[typing.Union[AnyFunction, int, sympy.core.expr.Expr], ...], ...
-            ],
-            typing.Tuple[typing.List[typing.Union[AnyFunction, int, sympy.core.expr.Expr]], ...],
-            typing.List[typing.Tuple[typing.Union[AnyFunction, int, sympy.core.expr.Expr], ...]],
-            typing.List[typing.List[typing.Union[AnyFunction, int, sympy.core.expr.Expr]]],
+            typing.Tuple[typing.Tuple[typing.Union[Function, int, sympy.core.expr.Expr], ...], ...],
+            typing.Tuple[typing.List[typing.Union[Function, int, sympy.core.expr.Expr]], ...],
+            typing.List[typing.Tuple[typing.Union[Function, int, sympy.core.expr.Expr], ...]],
+            typing.List[typing.List[typing.Union[Function, int, sympy.core.expr.Expr]]],
             sympy.matrices.dense.MutableDenseMatrix,
         ],
     ):
@@ -1370,7 +1368,7 @@ class MatrixFunction(AnyFunction):
         for i in mat:
             row = []
             for j in i:
-                if isinstance(j, AnyFunction):
+                if isinstance(j, Function):
                     if isinstance(j, BasisFunction):
                         j = j.get_function()
                     assert isinstance(j, ScalarFunction)
@@ -1699,7 +1697,7 @@ class MatrixFunction(AnyFunction):
                     out += self._mat[i][j] * other._mat[i][j]
             return out
 
-        if isinstance(other, AnyFunction) and other.is_matrix:
+        if isinstance(other, Function) and other.is_matrix:
             return other.dot(self)
 
         raise NotImplementedError()
@@ -1752,7 +1750,7 @@ class MatrixFunction(AnyFunction):
         domain: symfem.references.Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -1763,7 +1761,9 @@ class MatrixFunction(AnyFunction):
         Returns:
             The integral
         """
-        raise NotImplementedError()
+        return MatrixFunction(
+            [[f.integral(domain, vars, dummy_vars) for f in row] for row in self._mat]
+        )
 
     def det(self) -> ScalarFunction:
         """Compute the determinant.
@@ -1789,7 +1789,7 @@ class MatrixFunction(AnyFunction):
         assert isinstance(mat, sympy.matrices.dense.MutableDenseMatrix)
         return MatrixFunction(mat.transpose())
 
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
@@ -1813,20 +1813,20 @@ class MatrixFunction(AnyFunction):
 
 
 FunctionInput = typing.Union[
-    AnyFunction,
+    Function,
     sympy.core.expr.Expr,
     int,
-    typing.Tuple[typing.Union[sympy.core.expr.Expr, int, AnyFunction], ...],
-    typing.List[typing.Union[sympy.core.expr.Expr, int, AnyFunction]],
-    typing.Tuple[typing.Tuple[typing.Union[sympy.core.expr.Expr, int, AnyFunction], ...], ...],
-    typing.Tuple[typing.List[typing.Union[sympy.core.expr.Expr, int, AnyFunction]], ...],
-    typing.List[typing.Tuple[typing.Union[sympy.core.expr.Expr, int, AnyFunction], ...]],
-    typing.List[typing.List[typing.Union[sympy.core.expr.Expr, int, AnyFunction]]],
+    typing.Tuple[typing.Union[sympy.core.expr.Expr, int, Function], ...],
+    typing.List[typing.Union[sympy.core.expr.Expr, int, Function]],
+    typing.Tuple[typing.Tuple[typing.Union[sympy.core.expr.Expr, int, Function], ...], ...],
+    typing.Tuple[typing.List[typing.Union[sympy.core.expr.Expr, int, Function]], ...],
+    typing.List[typing.Tuple[typing.Union[sympy.core.expr.Expr, int, Function], ...]],
+    typing.List[typing.List[typing.Union[sympy.core.expr.Expr, int, Function]]],
     sympy.matrices.dense.MutableDenseMatrix,
 ]
 
 
-def parse_function_input(f: FunctionInput) -> AnyFunction:
+def parse_function_input(f: FunctionInput) -> Function:
     """Parse a function.
 
     Args:
@@ -1835,7 +1835,7 @@ def parse_function_input(f: FunctionInput) -> AnyFunction:
     Returns:
         The function as a Symfem function
     """
-    if isinstance(f, AnyFunction):
+    if isinstance(f, Function):
         return f
     if isinstance(f, (sympy.core.expr.Expr, int)):
         return ScalarFunction(f)
@@ -1857,7 +1857,7 @@ def parse_function_input(f: FunctionInput) -> AnyFunction:
 
 def parse_function_list_input(
     functions: typing.Union[typing.List[FunctionInput], typing.Tuple[FunctionInput, ...]],
-) -> typing.List[AnyFunction]:
+) -> typing.List[Function]:
     """Parse a list of functions.
 
     Args:
diff --git a/symfem/mappings.py b/symfem/mappings.py
index 3be9c39..76cd84f 100644
--- a/symfem/mappings.py
+++ b/symfem/mappings.py
@@ -6,7 +6,7 @@ from typing import Callable
 
 import sympy
 
-from symfem.functions import AnyFunction, FunctionInput, MatrixFunction, parse_function_input
+from symfem.functions import Function, FunctionInput, MatrixFunction, parse_function_input
 from symfem.geometry import PointType
 from symfem.symbols import x
 
@@ -41,7 +41,7 @@ def identity(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Map functions.
 
     Args:
@@ -64,7 +64,7 @@ def l2(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Map functions, scaling by the determinant of the jacobian.
 
     Args:
@@ -90,7 +90,7 @@ def covariant(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Map H(curl) functions.
 
     Args:
@@ -115,7 +115,7 @@ def covariant(
 
 def contravariant(
     f_in: FunctionInput, map: PointType, inverse_map: PointType, substitute: bool = True
-) -> AnyFunction:
+) -> Function:
     """Map H(div) functions.
 
     Args:
@@ -199,7 +199,7 @@ def identity_inverse_transpose(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse transpose of identity().
 
     Args:
@@ -222,7 +222,7 @@ def l2_inverse_transpose(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse transpose of l2().
 
     Args:
@@ -248,7 +248,7 @@ def covariant_inverse_transpose(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse transpose of covariant().
 
     Args:
@@ -276,7 +276,7 @@ def contravariant_inverse_transpose(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse transpose of contravariant().
 
     Args:
@@ -305,7 +305,7 @@ def identity_inverse(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse of identity().
 
     Args:
@@ -325,7 +325,7 @@ def l2_inverse(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse of l2().
 
     Args:
@@ -345,7 +345,7 @@ def covariant_inverse(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse of covariant().
 
     Args:
@@ -362,7 +362,7 @@ def covariant_inverse(
 
 def contravariant_inverse(
     f_in: FunctionInput, map: PointType, inverse_map: PointType, substitute: bool = True
-) -> AnyFunction:
+) -> Function:
     """Inverse of contravariant().
 
     Args:
@@ -382,7 +382,7 @@ def double_covariant_inverse(
     map: PointType,
     inverse_map: PointType,
     substitute: bool = True,
-) -> AnyFunction:
+) -> Function:
     """Inverse of double_covariant().
 
     Args:
@@ -399,7 +399,7 @@ def double_covariant_inverse(
 
 def double_contravariant_inverse(
     f_in: FunctionInput, map: PointType, inverse_map: PointType, substitute: bool = True
-) -> AnyFunction:
+) -> Function:
     """Inverse of double_contravariant().
 
     Args:
@@ -416,7 +416,7 @@ def double_contravariant_inverse(
 
 def get_mapping(
     mapname: str, inverse: bool = False, transpose: bool = False
-) -> Callable[[FunctionInput, PointType, PointType, bool], AnyFunction]:
+) -> Callable[[FunctionInput, PointType, PointType, bool], Function]:
     """Get a mapping.
 
     Args:
diff --git a/symfem/piecewise_functions.py b/symfem/piecewise_functions.py
index 974f1b1..92fa4fb 100644
--- a/symfem/piecewise_functions.py
+++ b/symfem/piecewise_functions.py
@@ -8,7 +8,7 @@ import sympy
 
 import symfem
 from symfem.functions import (
-    AnyFunction,
+    Function,
     FunctionInput,
     ScalarFunction,
     SympyFormat,
@@ -33,10 +33,10 @@ from symfem.symbols import AxisVariables, AxisVariablesNotSingle, t, x
 __all__ = ["PiecewiseFunction"]
 
 
-class PiecewiseFunction(AnyFunction):
+class PiecewiseFunction(Function):
     """A piecewise function."""
 
-    _pieces: typing.Dict[SetOfPoints, AnyFunction]
+    _pieces: typing.Dict[SetOfPoints, Function]
 
     def __init__(self, pieces: typing.Dict[SetOfPointsInput, FunctionInput], tdim: int):
         """Create a piecewise function.
@@ -125,7 +125,7 @@ class PiecewiseFunction(AnyFunction):
         out += "\\end{cases}"
         return out
 
-    def get_piece(self, point: PointType) -> AnyFunction:
+    def get_piece(self, point: PointType) -> Function:
         """Get a piece of the function.
 
         Args:
@@ -178,7 +178,7 @@ class PiecewiseFunction(AnyFunction):
         return True
 
     @property
-    def pieces(self) -> typing.Dict[SetOfPoints, AnyFunction]:
+    def pieces(self) -> typing.Dict[SetOfPoints, Function]:
         """Get the pieces of the function.
 
         Returns:
@@ -310,7 +310,7 @@ class PiecewiseFunction(AnyFunction):
         Returns:
             The substituted function
         """
-        if isinstance(values, AnyFunction):
+        if isinstance(values, Function):
             values = values.as_sympy()
 
         if isinstance(values, (tuple, list)) and len(values) == self.tdim:
@@ -539,7 +539,7 @@ class PiecewiseFunction(AnyFunction):
             ref = _piece_reference(self.tdim, shape)
             f.plot_values(ref, img, value_scale, n // 2)
 
-    def with_floats(self) -> AnyFunction:
+    def with_floats(self) -> Function:
         """Return a version the function with floats as coefficients.
 
         Returns:
diff --git a/symfem/plotting.py b/symfem/plotting.py
index d46976f..78d575c 100644
--- a/symfem/plotting.py
+++ b/symfem/plotting.py
@@ -5,7 +5,7 @@ from abc import ABC, abstractmethod
 
 import sympy
 
-from symfem.functions import AnyFunction, VectorFunction
+from symfem.functions import Function, VectorFunction
 from symfem.geometry import (
     PointType,
     PointTypeInput,
@@ -30,7 +30,7 @@ __all__ = [
     "Picture",
 ]
 
-PointOrFunction = typing.Union[PointTypeInput, AnyFunction]
+PointOrFunction = typing.Union[PointTypeInput, Function]
 SetOfPointsOrFunctions = typing.Union[
     typing.List[PointOrFunction], typing.Tuple[PointOrFunction, ...]
 ]
@@ -800,7 +800,7 @@ class Picture:
         Returns:
             The point as a tuple of Sympy expressions
         """
-        if isinstance(p, AnyFunction):
+        if isinstance(p, Function):
             p_s = p.as_sympy()
             assert isinstance(p_s, tuple)
             p = p_s
diff --git a/symfem/py.typed b/symfem/py.typed
new file mode 100644
index 0000000..e69de29


Test patch:
diff --git a/test/test_functions.py b/test/test_functions.py
index 55ebfbb..ed469bd 100644
--- a/test/test_functions.py
+++ b/test/test_functions.py
@@ -5,6 +5,7 @@ import sympy
 from symfem.functions import MatrixFunction, ScalarFunction, VectorFunction
 from symfem.piecewise_functions import PiecewiseFunction
 from symfem.symbols import x
+from symfem import create_reference
 
 
 def test_scalar_function_add_sub():
@@ -189,3 +190,22 @@ def test_piecewise_scalar_function_subs():
 
     assert f2.subs(x[0], 2) == f7
     assert f2.subs(x[:2], (1, 1)) == 1
+
+
+def test_vector_integral():
+    f = VectorFunction([1, x[0], x[0] * x[1] ** 2])
+    i = f.integral(create_reference("triangle"))
+
+    assert i[0] == sympy.Rational(1, 2)
+    assert i[1] == sympy.Rational(1, 6)
+    assert i[2] == sympy.Rational(1, 60)
+
+
+def test_matrix_integral():
+    f = MatrixFunction([[1, x[0]], [x[0] * x[1] ** 2, 1]])
+    i = f.integral(create_reference("triangle"))
+
+    assert i[0, 0] == sympy.Rational(1, 2)
+    assert i[0, 1] == sympy.Rational(1, 6)
+    assert i[1, 0] == sympy.Rational(1, 60)
+    assert i[1, 1] == sympy.Rational(1, 2)