Additional tests for operators (#384)

* Adding tests for operators

Author: GiovanniCanali

tests/test_operators.py

@@ -5,73 +5,120 @@
 from pina.operators import grad, div, laplacian
 
 
-def func_vec(x):
+def func_vector(x):
     return x**2
 
 
 def func_scalar(x):
-    print('X')
     x_ = x.extract(['x'])
     y_ = x.extract(['y'])
-    mu_ = x.extract(['mu'])
-    return x_**2 + y_**2 + mu_**3
+    z_ = x.extract(['z'])
+    return x_**2 + y_**2 + z_**2
 
 
-data = torch.rand((20, 3), requires_grad=True)
-inp = LabelTensor(data, ['x', 'y', 'mu'])
-labels = ['a', 'b', 'c']
-tensor_v = LabelTensor(func_vec(inp), labels)
-tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), labels[0])
-
+inp = LabelTensor(torch.rand((20, 3), requires_grad=True), ['x', 'y', 'z'])
+tensor_v = LabelTensor(func_vector(inp), ['a', 'b', 'c'])
+tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), ['a'])
 
 def test_grad_scalar_output():
     grad_tensor_s = grad(tensor_s, inp)
+    true_val = 2*inp
     assert grad_tensor_s.shape == inp.shape
     assert grad_tensor_s.labels == [
         f'd{tensor_s.labels[0]}d{i}' for i in inp.labels
     ]
+    assert torch.allclose(grad_tensor_s, true_val)
+
     grad_tensor_s = grad(tensor_s, inp, d=['x', 'y'])
+    true_val = 2*inp.extract(['x', 'y'])
     assert grad_tensor_s.shape == (inp.shape[0], 2)
     assert grad_tensor_s.labels == [
         f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y']
     ]
+    assert torch.allclose(grad_tensor_s, true_val)
 
 
 def test_grad_vector_output():
     grad_tensor_v = grad(tensor_v, inp)
+    true_val = torch.cat(
+        (2*inp.extract(['x']),
+         torch.zeros_like(inp.extract(['y'])),
+         torch.zeros_like(inp.extract(['z'])),
+         torch.zeros_like(inp.extract(['x'])),
+         2*inp.extract(['y']),
+         torch.zeros_like(inp.extract(['z'])),
+         torch.zeros_like(inp.extract(['x'])),
+         torch.zeros_like(inp.extract(['y'])),
+         2*inp.extract(['z'])
+        ), dim=1
+    )
     assert grad_tensor_v.shape == (20, 9)
-    grad_tensor_v = grad(tensor_v, inp, d=['x', 'mu'])
+    assert grad_tensor_v.labels == [
+        f'd{j}d{i}' for j in tensor_v.labels for i in inp.labels
+    ]
+    assert torch.allclose(grad_tensor_v, true_val)
+
+    grad_tensor_v = grad(tensor_v, inp, d=['x', 'y'])
+    true_val = torch.cat(
+        (2*inp.extract(['x']),
+         torch.zeros_like(inp.extract(['y'])),
+         torch.zeros_like(inp.extract(['x'])),
+         2*inp.extract(['y']),
+         torch.zeros_like(inp.extract(['x'])),
+         torch.zeros_like(inp.extract(['y']))
+        ), dim=1
+    )
     assert grad_tensor_v.shape == (inp.shape[0], 6)
+    assert grad_tensor_v.labels == [
+        f'd{j}d{i}' for j in tensor_v.labels for i in ['x', 'y']
+    ]
+    assert torch.allclose(grad_tensor_v, true_val)
 
 
 def test_div_vector_output():
-    grad_tensor_v = div(tensor_v, inp)
-    assert grad_tensor_v.shape == (20, 1)
-    grad_tensor_v = div(tensor_v, inp, components=['a', 'b'], d=['x', 'mu'])
-    assert grad_tensor_v.shape == (inp.shape[0], 1)
+    div_tensor_v = div(tensor_v, inp)
+    true_val = 2*torch.sum(inp, dim=1).reshape(-1,1)
+    assert div_tensor_v.shape == (20, 1)
+    assert div_tensor_v.labels == [f'dadx+dbdy+dcdz']
+    assert torch.allclose(div_tensor_v, true_val)
+
+    div_tensor_v = div(tensor_v, inp, components=['a', 'b'], d=['x', 'y'])
+    true_val = 2*torch.sum(inp.extract(['x', 'y']), dim=1).reshape(-1,1)
+    assert div_tensor_v.shape == (inp.shape[0], 1)
+    assert div_tensor_v.labels == [f'dadx+dbdy']
+    assert torch.allclose(div_tensor_v, true_val)
 
 
 def test_laplacian_scalar_output():
-    laplace_tensor_s = laplacian(tensor_s, inp, components=['a'], d=['x', 'y'])
+    laplace_tensor_s = laplacian(tensor_s, inp)
+    true_val = 6*torch.ones_like(laplace_tensor_s)
     assert laplace_tensor_s.shape == tensor_s.shape
     assert laplace_tensor_s.labels == [f"dd{tensor_s.labels[0]}"]
+    assert torch.allclose(laplace_tensor_s, true_val)
+
+    laplace_tensor_s = laplacian(tensor_s, inp, components=['a'], d=['x', 'y'])
     true_val = 4*torch.ones_like(laplace_tensor_s)
-    assert all((laplace_tensor_s - true_val == 0).flatten())
+    assert laplace_tensor_s.shape == tensor_s.shape
+    assert laplace_tensor_s.labels == [f"dd{tensor_s.labels[0]}"]
+    assert torch.allclose(laplace_tensor_s, true_val)
 
 
 def test_laplacian_vector_output():
     laplace_tensor_v = laplacian(tensor_v, inp)
+    true_val = 2*torch.ones_like(tensor_v)
     assert laplace_tensor_v.shape == tensor_v.shape
     assert laplace_tensor_v.labels == [
         f'dd{i}' for i in tensor_v.labels
     ]
+    assert torch.allclose(laplace_tensor_v, true_val)
+
     laplace_tensor_v = laplacian(tensor_v,
                                  inp,
                                  components=['a', 'b'],
                                  d=['x', 'y'])
+    true_val = 2*torch.ones_like(tensor_v.extract(['a', 'b']))
     assert laplace_tensor_v.shape == tensor_v.extract(['a', 'b']).shape
     assert laplace_tensor_v.labels == [
         f'dd{i}' for i in ['a', 'b']
     ]
-    true_val = 2*torch.ones_like(tensor_v.extract(['a', 'b']))
-    assert all((laplace_tensor_v - true_val == 0).flatten())
+    assert torch.allclose(laplace_tensor_v, true_val)