Implement integrals of matrix and vector functions (#297)

* add py.typed

* rename AnyFunction -> Function

* integrate components of vector function

* ruff

* DefElement branch

Author: mscroggs

.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

symfem/basis_functions.py

@@ -9,7 +9,7 @@
 
 import symfem
 from symfem.functions import (
-    AnyFunction,
+    Function,
     FunctionInput,
     ScalarFunction,
     SympyFormat,
@@ -22,7 +22,7 @@
 __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,15 +40,15 @@ def __init__(self, scalar=False, vector=False, matrix=False):
         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:
             The basis function
         """
         pass
 
-    def __add__(self, other: typing.Any) -> AnyFunction:
+    def __add__(self, other: typing.Any) -> Function:
         """Add.
 
         Args:
@@ -59,7 +59,7 @@ def __add__(self, other: typing.Any) -> 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 @@ def __radd__(self, other: typing.Any) -> 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 @@ def __sub__(self, other: typing.Any) -> AnyFunction:
         """
         return self.get_function().__sub__(other)
 
-    def __rsub__(self, other: typing.Any) -> AnyFunction:
+    def __rsub__(self, other: typing.Any) -> Function:
         """Subtract.
 
         Args:
@@ -92,15 +92,15 @@ def __rsub__(self, other: typing.Any) -> AnyFunction:
         """
         return self.get_function().__rsub__(other)
 
-    def __neg__(self) -> AnyFunction:
+    def __neg__(self) -> Function:
         """Negate.
 
         Returns:
             Negated function
         """
         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 @@ def __truediv__(self, other: typing.Any) -> 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 @@ def __rtruediv__(self, other: typing.Any) -> 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 @@ def __mul__(self, other: typing.Any) -> 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 @@ def __rmul__(self, other: typing.Any) -> 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 @@ def __matmul__(self, other: typing.Any) -> 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 @@ def __rmatmul__(self, other: typing.Any) -> 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 @@ def as_tex(self) -> str:
         """
         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 @@ def diff(self, variable: sympy.core.symbol.Symbol) -> 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 @@ def directional_derivative(self, direction: PointType) -> 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 @@ def jacobian_component(self, component: typing.Tuple[int, int]) -> 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 @@ def jacobian(self, dim: int) -> 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 @@ def dot(self, other_in: FunctionInput) -> 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,23 +259,23 @@ def cross(self, other_in: FunctionInput) -> AnyFunction:
         """
         return self.get_function().cross(other_in)
 
-    def div(self) -> AnyFunction:
+    def div(self) -> Function:
         """Compute the divergence of the function.
 
         Returns:
             The divergence
         """
         return self.get_function().div()
 
-    def grad(self, dim: int) -> AnyFunction:
+    def grad(self, dim: int) -> Function:
         """Compute the gradient of the function.
 
         Returns:
             The gradient
         """
         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 @@ def integral(
         domain: Reference,
         vars: AxisVariablesNotSingle = x,
         dummy_vars: AxisVariablesNotSingle = t,
-    ) -> ScalarFunction:
+    ) -> Function:
         """Compute the integral of the function.
 
         Args:
@@ -321,7 +321,7 @@ def subs(self, vars: AxisVariables, values: ValuesToSubstitute) -> BasisFunction
         """
         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,15 +353,15 @@ def transpose(self) -> ScalarFunction:
             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:
             The function with floats as coefficients
         """
         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 @@ def __init__(self, f: BasisFunction, vars: AxisVariables, values: ValuesToSubsti
         self._vars = vars
         self._values = values
 
-    def get_function(self) -> AnyFunction:
+    def get_function(self) -> Function:
         """Return the symbolic function.
 
         Returns:

symfem/elements/bernstein.py

@@ -11,7 +11,7 @@
 
 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 @@ def dof_point(self) -> PointType:
         """
         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:

symfem/elements/dual.py

@@ -9,7 +9,7 @@
 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 @@ def __init__(
         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 @@ def get_dual_matrix(self) -> sympy.matrices.dense.MutableDenseMatrix:
         """
         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 @@ def get_basis_functions(
         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 @@ def dof_entities(self) -> typing.List[typing.Tuple[int, int]]:
     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:

symfem/elements/guzman_neilan.py

@@ -13,7 +13,7 @@
 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 @@
 __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":