Simplify Graph class (#459)

* Simplifying Graph class and adjust tests

---------

Co-authored-by: Dario Coscia <[email protected]>

Author: 2 people

pina/graph.py

@@ -15,12 +15,18 @@ def test_supervised_tensor_collector():
     class SupervisedProblem(AbstractProblem):
         output_variables = None
         conditions = {
-            'data1': Condition(input_points=torch.rand((10, 2)),
-                               output_points=torch.rand((10, 2))),
-            'data2': Condition(input_points=torch.rand((20, 2)),
-                               output_points=torch.rand((20, 2))),
-            'data3': Condition(input_points=torch.rand((30, 2)),
-                               output_points=torch.rand((30, 2))),
+            "data1": Condition(
+                input_points=torch.rand((10, 2)),
+                output_points=torch.rand((10, 2)),
+            ),
+            "data2": Condition(
+                input_points=torch.rand((20, 2)),
+                output_points=torch.rand((20, 2)),
+            ),
+            "data3": Condition(
+                input_points=torch.rand((30, 2)),
+                output_points=torch.rand((30, 2)),
+            ),
         }
 
     problem = SupervisedProblem()
@@ -31,65 +37,58 @@ class SupervisedProblem(AbstractProblem):
 
 def test_pinn_collector():
     def laplace_equation(input_, output_):
-        force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
-                      torch.sin(input_.extract(['y']) * torch.pi))
-        delta_u = laplacian(output_.extract(['u']), input_)
+        force_term = torch.sin(input_.extract(["x"]) * torch.pi) * torch.sin(
+            input_.extract(["y"]) * torch.pi
+        )
+        delta_u = laplacian(output_.extract(["u"]), input_)
         return delta_u - force_term
 
     my_laplace = Equation(laplace_equation)
-    in_ = LabelTensor(torch.tensor([[0., 1.]], requires_grad=True), ['x', 'y'])
-    out_ = LabelTensor(torch.tensor([[0.]], requires_grad=True), ['u'])
+    in_ = LabelTensor(
+        torch.tensor([[0.0, 1.0]], requires_grad=True), ["x", "y"]
+    )
+    out_ = LabelTensor(torch.tensor([[0.0]], requires_grad=True), ["u"])
 
     class Poisson(SpatialProblem):
-        output_variables = ['u']
-        spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+        output_variables = ["u"]
+        spatial_domain = CartesianDomain({"x": [0, 1], "y": [0, 1]})
 
         conditions = {
-            'gamma1':
-                Condition(domain=CartesianDomain({
-                    'x': [0, 1],
-                    'y': 1
-                }),
-                    equation=FixedValue(0.0)),
-            'gamma2':
-                Condition(domain=CartesianDomain({
-                    'x': [0, 1],
-                    'y': 0
-                }),
-                    equation=FixedValue(0.0)),
-            'gamma3':
-                Condition(domain=CartesianDomain({
-                    'x': 1,
-                    'y': [0, 1]
-                }),
-                    equation=FixedValue(0.0)),
-            'gamma4':
-                Condition(domain=CartesianDomain({
-                    'x': 0,
-                    'y': [0, 1]
-                }),
-                    equation=FixedValue(0.0)),
-            'D':
-                Condition(domain=CartesianDomain({
-                    'x': [0, 1],
-                    'y': [0, 1]
-                }),
-                    equation=my_laplace),
-            'data':
-                Condition(input_points=in_, output_points=out_)
+            "gamma1": Condition(
+                domain=CartesianDomain({"x": [0, 1], "y": 1}),
+                equation=FixedValue(0.0),
+            ),
+            "gamma2": Condition(
+                domain=CartesianDomain({"x": [0, 1], "y": 0}),
+                equation=FixedValue(0.0),
+            ),
+            "gamma3": Condition(
+                domain=CartesianDomain({"x": 1, "y": [0, 1]}),
+                equation=FixedValue(0.0),
+            ),
+            "gamma4": Condition(
+                domain=CartesianDomain({"x": 0, "y": [0, 1]}),
+                equation=FixedValue(0.0),
+            ),
+            "D": Condition(
+                domain=CartesianDomain({"x": [0, 1], "y": [0, 1]}),
+                equation=my_laplace,
+            ),
+            "data": Condition(input_points=in_, output_points=out_),
         }
 
         def poisson_sol(self, pts):
-            return -(torch.sin(pts.extract(['x']) * torch.pi) *
-                     torch.sin(pts.extract(['y']) * torch.pi)) / (
-                    2 * torch.pi ** 2)
+            return -(
+                torch.sin(pts.extract(["x"]) * torch.pi)
+                * torch.sin(pts.extract(["y"]) * torch.pi)
+            ) / (2 * torch.pi**2)
 
         truth_solution = poisson_sol
 
     problem = Poisson()
-    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
-    problem.discretise_domain(10, 'grid', domains=boundaries)
-    problem.discretise_domain(10, 'grid', domains='D')
+    boundaries = ["gamma1", "gamma2", "gamma3", "gamma4"]
+    problem.discretise_domain(10, "grid", domains=boundaries)
+    problem.discretise_domain(10, "grid", domains="D")
 
     collector = Collector(problem)
     collector.store_fixed_data()
@@ -98,31 +97,34 @@ def poisson_sol(self, pts):
     for k, v in problem.conditions.items():
         if isinstance(v, InputOutputPointsCondition):
             assert list(collector.data_collections[k].keys()) == [
-                'input_points', 'output_points']
+                "input_points",
+                "output_points",
+            ]
 
     for k, v in problem.conditions.items():
         if isinstance(v, DomainEquationCondition):
             assert list(collector.data_collections[k].keys()) == [
-                'input_points', 'equation']
+                "input_points",
+                "equation",
+            ]
 
 
 def test_supervised_graph_collector():
     pos = torch.rand((100, 3))
     x = [torch.rand((100, 3)) for _ in range(10)]
-    graph_list_1 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
+    graph_list_1 = [RadiusGraph(pos=pos, radius=0.4, x=x_) for x_ in x]
     out_1 = torch.rand((10, 100, 3))
+
     pos = torch.rand((50, 3))
     x = [torch.rand((50, 3)) for _ in range(10)]
-    graph_list_2 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
+    graph_list_2 = [RadiusGraph(pos=pos, radius=0.4, x=x_) for x_ in x]
     out_2 = torch.rand((10, 50, 3))
 
     class SupervisedProblem(AbstractProblem):
         output_variables = None
         conditions = {
-            'data1': Condition(input_points=graph_list_1,
-                               output_points=out_1),
-            'data2': Condition(input_points=graph_list_2,
-                               output_points=out_2),
+            "data1": Condition(input_points=graph_list_1, output_points=out_1),
+            "data2": Condition(input_points=graph_list_2, output_points=out_2),
         }
 
     problem = SupervisedProblem()