equation class, fix minor bugs, diff domain (#89)

* equation class
* difference domain
* dummy dataloader
* writer class
* refactoring and minor fix

tests/test_cartesian.py

@@ -0,0 +1,21 @@
+import torch
+import pytest
+
+from pina import LabelTensor, Condition, CartesianDomain, PINN
+from pina.problem import SpatialProblem
+from pina.model import FeedForward
+from pina.operators import nabla
+
+
+
+def test_constructor():
+    CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+
+def test_is_inside():
+    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5]]), ['x', 'y'])
+    pt_2 = LabelTensor(torch.tensor([[1.0, 0.5]]), ['x', 'y'])
+    pt_3 = LabelTensor(torch.tensor([[1.5, 0.5]]), ['x', 'y'])
+    domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+    for pt, exp_result in zip([pt_1, pt_2, pt_3], [True, True, False]):
+        assert domain.is_inside(pt) == exp_result

tests/test_condition.py

@@ -5,12 +5,10 @@ from pina import LabelTensor, Condition, CartesianDomain, PINN
 from pina.problem import SpatialProblem
 from pina.model import FeedForward
 from pina.operators import nabla
+from pina.equation.equation_factory import FixedValue
 
 
 example_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-def example_dirichlet(input_, output_):
-    value = 0.0
-    return output_.extract(['u']) - value
 example_input_pts = LabelTensor(torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
 example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ['a', 'b'])
 
@@ -21,22 +19,22 @@ def test_init_inputoutput():
     with pytest.raises(TypeError):
         Condition(input_points=3., output_points='example')
     with pytest.raises(TypeError):
-        Condition(input_points=example_domain, output_points=example_dirichlet)
+        Condition(input_points=example_domain, output_points=example_domain)
 
 def test_init_locfunc():
-    Condition(location=example_domain, function=example_dirichlet)
+    Condition(location=example_domain, equation=FixedValue(0.0))    
     with pytest.raises(ValueError):
-        Condition(example_domain, example_dirichlet)
+        Condition(example_domain, FixedValue(0.0))
     with pytest.raises(TypeError):
-        Condition(location=3., function='example')
+        Condition(location=3., equation='example')
     with pytest.raises(TypeError):
-        Condition(location=example_input_pts, function=example_output_pts)
+        Condition(location=example_input_pts, equation=example_output_pts)
 
 def test_init_inputfunc():
-    Condition(input_points=example_input_pts, function=example_dirichlet)
+    Condition(input_points=example_input_pts, equation=FixedValue(0.0))
     with pytest.raises(ValueError):
-        Condition(example_domain, example_dirichlet)
+        Condition(example_domain, FixedValue(0.0))
     with pytest.raises(TypeError):
-        Condition(input_points=3., function='example')
+        Condition(input_points=3., equation='example')
     with pytest.raises(TypeError):
-        Condition(input_points=example_domain, function=example_output_pts)
+        Condition(input_points=example_domain, equation=example_output_pts)

tests/test_label_tensor.py

@@ -27,6 +27,7 @@ def test_labels():
 def test_extract():
     label_to_extract = ['a', 'c']
     tensor = LabelTensor(data, labels)
+    print(tensor)
     new = tensor.extract(label_to_extract)
     assert new.labels == label_to_extract
     assert new.shape[1] == len(label_to_extract)
@@ -79,3 +80,11 @@ def test_merge():
 
     tensor_bc = tensor_b.append(tensor_c)
     assert torch.allclose(tensor_bc, tensor.extract(['b', 'c']))
+
+
+def test_getitem():
+    tensor = LabelTensor(data, labels)
+    tensor_view = tensor[:5]
+
+    assert tensor_view.labels == labels
+    assert torch.allclose(tensor_view, data[:5])

tests/test_operators.py

@@ -8,7 +8,11 @@ def func_vec(x):
     return x**2
 
 def func_scalar(x):
-    return x[:, 0]**2 + x[:, 1]**2 + x[:, 2]**3
+    print('X')
+    x_ = x.extract(['x'])
+    y_ = x.extract(['y'])
+    mu_ = x.extract(['mu'])
+    return x_**2 + y_**2 + mu_**3
 
 data = torch.rand((20, 3), requires_grad=True)
 inp = LabelTensor(data, ['x', 'y', 'mu'])

tests/test_pinn.py

@@ -5,40 +5,41 @@ from pina import LabelTensor, Condition, CartesianDomain, PINN
 from pina.problem import SpatialProblem
 from pina.model import FeedForward
 from pina.operators import nabla
+from pina.equation.equation import Equation
+from pina.equation.equation_factory import FixedValue
+
 
 in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
 out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
 
+def laplace_equation(input_, output_):
+    force_term = (torch.sin(input_.extract(['x'])*torch.pi) *
+                    torch.sin(input_.extract(['y'])*torch.pi))
+    nabla_u = nabla(output_.extract(['u']), input_)
+    return nabla_u - force_term
+
+my_laplace = Equation(laplace_equation)
+
 class Poisson(SpatialProblem):
     output_variables = ['u']
     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 
-    def laplace_equation(input_, output_):
-        force_term = (torch.sin(input_.extract(['x'])*torch.pi) *
-                      torch.sin(input_.extract(['y'])*torch.pi))
-        nabla_u = nabla(output_, input_, components=['u'], d=['x', 'y'])
-        return nabla_u - force_term
-
-    def nil_dirichlet(input_, output_):
-        value = 0.0
-        return output_.extract(['u']) - value
-
     conditions = {
         'gamma1': Condition(
             location=CartesianDomain({'x': [0, 1], 'y':  1}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.0)),
         'gamma2': Condition(
             location=CartesianDomain({'x': [0, 1], 'y': 0}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.0)),
         'gamma3': Condition(
             location=CartesianDomain({'x':  1, 'y': [0, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.0)),
         'gamma4': Condition(
             location=CartesianDomain({'x': 0, 'y': [0, 1]}),
-            function=nil_dirichlet),
+            equation=FixedValue(0.0)),
         'D': Condition(
             location=CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
-            function=laplace_equation),
+            equation=my_laplace),
         'data': Condition(
             input_points=in_,
             output_points=out_)
@@ -137,7 +138,7 @@ def test_train():
     pinn.span_pts(n, 'grid', locations=['D'])
     pinn.train(5)
 
-
+"""
 def test_train_2():
     boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
     n = 10
@@ -243,3 +244,4 @@ if torch.cuda.is_available():
         pinn.span_pts(n, 'grid', locations=boundaries)
         pinn.span_pts(n, 'grid', locations=['D'])
         pinn.train(5)
+"""