refact

tests/test_solvers/test_supervised_solver.py

@@ -5,14 +5,17 @@ from pina import Condition, LabelTensor
 from pina.solvers import SupervisedSolver
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.loss import LpLoss
+from pina.loss.loss_interface import LpLoss
 
 
 class NeuralOperatorProblem(AbstractProblem):
     input_variables = ['u_0', 'u_1']
     output_variables = ['u']
     domains = {
-        'pts': LabelTensor(torch.rand(100, 2), labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}})
+        'pts': LabelTensor(
+            torch.rand(100, 2),
+            labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}}
+        )
     }
     conditions = {
         'data' : Condition(
@@ -56,9 +59,51 @@ def test_constructor():
 #     SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
 
 
+class AutoSolver(SupervisedSolver):
+
+    def forward(self, input):
+        from pina.graph import Graph
+        print(Graph)
+        print(input)
+        if not isinstance(input, Graph):
+            input = Graph.build('radius', nodes_coordinates=input, nodes_data=torch.rand(input.shape), radius=0.2)
+        print(input)
+        print(input.data.edge_index)
+        print(input.data)
+        g = self.model[0](input.data, edge_index=input.data.edge_index)
+        g.labels = {1: {'name': 'output', 'dof': ['u']}}
+        return g
+        du_dt_new = LabelTensor(self.model[0](graph).reshape(-1,1), labels = ['du'])
+
+        return du_dt_new
+
+class GraphModel(torch.nn.Module):
+    def __init__(self, in_channels, out_channels):
+        from torch_geometric.nn import GCNConv, NNConv
+        super().__init__()
+        self.conv1 = GCNConv(in_channels, 16)
+        self.conv2 = GCNConv(16, out_channels)
+
+    def forward(self, data, edge_index):
+        print(data)
+        x = data.x
+        print(x)
+        x = self.conv1(x, edge_index)
+        x = x.relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+def test_graph():
+
+    solver = AutoSolver(problem = problem, model=GraphModel(2, 1), loss=LpLoss())
+    trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
+    trainer.train()
+    assert False
+
+
 def test_train_cpu():
     solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss())
-    trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
+    trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20)
     trainer.train()