tests: Add a vector grad test and simplify existing tests

Author: EdCaunt

tests/test_derivatives.py

@@ -4,7 +4,7 @@
 
 from devito import (Grid, Function, TimeFunction, Eq, Operator, NODE, cos, sin,
                     ConditionalDimension, left, right, centered, div, grad,
-                    curl, laplace, VectorFunction)
+                    curl, laplace, VectorFunction, TensorFunction)
 from devito.finite_differences import Derivative, Differentiable, diffify
 from devito.finite_differences.differentiable import (Add, EvalDerivative, IndexSum,
                                                       IndexDerivative, Weights,
@@ -636,11 +636,22 @@ def test_grad_w_side(self, side):
         expr1 = VectorFunction(name=f"{f.name}_vec", space_order=f.space_order,
                                components=comps, grid=grid).evaluate
 
-        expr2 = f.grad(side=side).evaluate
-        expr3 = grad(f, side=side).evaluate
+        assert expr1 == f.grad(side=side).evaluate
+        assert expr1 == grad(f, side=side).evaluate
 
-        assert expr1 == expr2
-        assert expr1 == expr3
+    @pytest.mark.parametrize('side', [left, right, centered])
+    def test_vector_grad_w_side(self, side):
+        grid = Grid(shape=(11, 11))
+        f = VectorFunction(name='f', grid=grid, space_order=2, staggered=(None, None))
+
+        comps = ((f[0].dx(side=side), f[0].dy(side=side)),
+                 (f[1].dx(side=side), f[1].dy(side=side)))
+
+        expr1 = TensorFunction(name=f"{f.name}_tens", space_order=f.space_order,
+                               components=comps, grid=grid).evaluate
+
+        assert expr1 == f.grad(side=side).evaluate
+        assert expr1 == grad(f, side=side).evaluate
 
     @pytest.mark.parametrize('side', [left, right, centered])
     def test_div_w_side(self, side):
@@ -649,11 +660,8 @@ def test_div_w_side(self, side):
 
         expr1 = (f[0].dx(side=side) + f[1].dy(side=side)).evaluate
 
-        expr2 = f.div(side=side).evaluate
-        expr3 = div(f, side=side).evaluate
-
-        assert expr1 == expr2
-        assert expr1 == expr3
+        assert expr1 == f.div(side=side).evaluate
+        assert expr1 == div(f, side=side).evaluate
 
     @pytest.mark.parametrize('side', [left, right, centered])
     def test_curl_w_side(self, side):
@@ -668,11 +676,8 @@ def test_curl_w_side(self, side):
         expr1 = VectorFunction(name=f"{f.name}_vec", space_order=f.space_order,
                                components=comps, grid=grid).evaluate
 
-        expr2 = f.curl(side=side).evaluate
-        expr3 = curl(f, side=side).evaluate
-
-        assert expr1 == expr2
-        assert expr1 == expr3
+        assert expr1 == f.curl(side=side).evaluate
+        assert expr1 == curl(f, side=side).evaluate
 
     @pytest.mark.parametrize('side', [left, right, centered])
     def test_laplace_w_side(self, side):
@@ -681,11 +686,8 @@ def test_laplace_w_side(self, side):
 
         expr1 = (f.dx2(side=side) + f.dy2(side=side)).evaluate
 
-        expr2 = f.laplacian(side=side).evaluate
-        expr3 = laplace(f, side=side).evaluate
-
-        assert expr1 == expr2
-        assert expr1 == expr3
+        assert expr1 == f.laplacian(side=side).evaluate
+        assert expr1 == laplace(f, side=side).evaluate
 
     def test_substitution(self):
         grid = Grid((11, 11))