Implement grad of VectorFunction (#299)

* add py.typed

* rename AnyFunction -> Function

* integrate components of vector function

* ruff

* write test

* DefElement branch

* diff(VectorFunction)

* ruff

* no need for ._vec

Author: mscroggs

.github/workflows/defelement.yml

@@ -27,7 +27,7 @@ jobs:
         with:
           path: ./defelement
           repository: DefElement/DefElement
-          ref: mscroggs/AnyFunction->Function
+          ref: main
       - name: Install requirements
         run: |
           cd defelement

CHANGELOG_SINCE_LAST_VERSION.md

@@ -1 +1,3 @@
 - Allow "gl" variant of Lagrange element
+- Implement component-wise integrals of VectorFunction and MatrixFunction
+- Implement grad of VectorFunction

symfem/functions.py

@@ -275,7 +275,7 @@ def div(self):
         """
 
     @abstractmethod
-    def grad(self, dim: int):
+    def grad(self, dim: int) -> Function:
         """Compute the grad of the function.
 
         Returns:
@@ -757,7 +757,7 @@ def div(self):
         """
         raise ValueError("Cannot compute the div of a scalar-valued function.")
 
-    def grad(self, dim: int) -> VectorFunction:
+    def grad(self, dim: int) -> Function:
         """Compute the grad of the function.
 
         Returns:
@@ -1215,13 +1215,13 @@ def div(self) -> ScalarFunction:
             out += i.diff(j)
         return out
 
-    def grad(self):
+    def grad(self, dim: int) -> Function:
         """Compute the grad of the function.
 
         Returns:
             The gradient
         """
-        raise ValueError("Cannot compute the grad of a vector-valued function.")
+        return MatrixFunction([f.grad(dim) for f in self._vec])
 
     def curl(self) -> VectorFunction:
         """Compute the curl of the function.
@@ -1721,7 +1721,7 @@ def div(self):
         """
         raise ValueError("Cannot compute the div of a matrix-valued function.")
 
-    def grad(self):
+    def grad(self, dim: int) -> Function:
         """Compute the grad of the function.
 
         Returns:

symfem/piecewise_functions.py

@@ -413,7 +413,7 @@ def div(self) -> PiecewiseFunction:
         """
         return PiecewiseFunction({shape: f.div() for shape, f in self._pieces.items()}, self.tdim)
 
-    def grad(self, dim: int) -> PiecewiseFunction:
+    def grad(self, dim: int) -> Function:
         """Compute the grad of the function.
 
         Returns:

test/test_functions.py

@@ -192,6 +192,21 @@ def test_piecewise_scalar_function_subs():
     assert f2.subs(x[:2], (1, 1)) == 1
 
 
+def test_vector_grad():
+    f = VectorFunction([x[0], x[1], x[1] ** 2 - 3 * x[2]])
+    d = f.grad(3)
+
+    assert d[0, 0] == 1
+    assert d[0, 1] == 0
+    assert d[0, 2] == 0
+    assert d[1, 0] == 0
+    assert d[1, 1] == 1
+    assert d[1, 2] == 0
+    assert d[2, 0] == 0
+    assert d[2, 1] == 2 * x[1]
+    assert d[2, 2] == -3
+
+
 def test_vector_integral():
     f = VectorFunction([1, x[0], x[0] * x[1] ** 2])
     i = f.integral(create_reference("triangle"))