added GNO with test and tutorial

.github/workflows/testing_pr.yml

@@ -28,6 +28,7 @@ jobs:
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install .[test]
+        python3 -m pip install torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cpu.html
 
     - name: Test with pytest
       run: |

pina/model/__init__.py

@@ -5,9 +5,11 @@
     'DeepONet',
     'MIONet',
     'FNO',
+    'GNO'
 ]
 
 from .feed_forward import FeedForward, ResidualFeedForward
 from .multi_feed_forward import MultiFeedForward
 from .deeponet import DeepONet, MIONet
 from .fno import FNO
+from .gno import GNO

pina/model/gno.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import pool
+from . import FeedForward
+
+class GNO(nn.Module):
+    def __init__(self,
+                input_features,
+                output_features,
+                points,
+                radius=0.05,
+                inner_size=20,
+                n_layers=2,
+                func=nn.Tanh,
+                ):
+
+        super().__init__()
+        self.input_features=input_features
+        self.output_features=output_features
+        self.num_points=points.shape[0]
+        self.points_size=points.shape[1]
+        self.nn=FeedForward(2*self.points_size+2*self.input_features, output_features*output_features, inner_size, n_layers, func)
+        self.points=points
+        self.neigh=pool.radius(points, points, radius)
+        self.linear=nn.Linear(self.input_features,self.output_features,bias=False)
+
+    def forward(self, batch):
+        points_x=self.points[self.neigh[1]].unsqueeze(0).repeat(batch.shape[0],1,1)
+        points_y=self.points[self.neigh[0]].unsqueeze(0).repeat(batch.shape[0],1,1)
+        batch_x=batch[:,self.neigh[1],:]
+        batch_y=batch[:,self.neigh[0],:]
+        new_batch=torch.concatenate((points_x,points_y,batch_x,batch_y),dim=2)
+        new_batch=self.nn(new_batch).reshape(batch.shape[0],-1,self.output_features,self.output_features)
+        new_batch=torch.matmul(new_batch,batch_y.unsqueeze(-1)).squeeze(-1)
+
+        tmp_list=self.neigh[0].unsqueeze(0).unsqueeze(-1).repeat(batch.shape[0],1,self.output_features)
+        tmp_array=torch.zeros(batch.shape[0],batch.shape[1],self.output_features,requires_grad=True)
+        lin_part=self.linear(batch)
+        k_part=torch.scatter_reduce(tmp_array,1,tmp_list,new_batch,reduce='mean')
+        new_batch=lin_part+k_part
+        return new_batch
+
+
+
+
+
+    def forward_eval(self,batch,points):
+        neigh=pool.radius(points, points, 0.05)
+        points_x=points[neigh[1]].unsqueeze(0).repeat(batch.shape[0],1,1)
+        points_y=points[neigh[0]].unsqueeze(0).repeat(batch.shape[0],1,1)
+        batch_x=batch[:,neigh[1],:]
+        batch_y=batch[:,neigh[0],:]
+        new_batch=torch.concatenate((points_x,points_y,batch_x,batch_y),dim=2)
+        new_batch=self.nn(new_batch).reshape(batch.shape[0],-1,self.output_features,self.output_features)
+        new_batch=torch.matmul(new_batch,batch_y.unsqueeze(-1)).squeeze(-1)
+        tmp_list=neigh[0].unsqueeze(0).unsqueeze(-1).repeat(batch.shape[0],1,self.output_features)
+        tmp_array=torch.zeros(batch.shape[0],batch.shape[1],self.output_features,requires_grad=True)
+        lin_part=self.linear(batch)
+        k_part=torch.scatter_reduce(tmp_array,1,tmp_list,new_batch,reduce='mean')
+        new_batch=lin_part+k_part
+        return new_batch
+
+
+
+

setup.py

@@ -15,7 +15,7 @@
 KEYWORDS = 'physics-informed neural-network'
 
 REQUIRED = [
-    'numpy', 'matplotlib', 'torch', 'lightning', 'pytorch_lightning'
+    'numpy', 'matplotlib', 'torch', 'lightning', 'pytorch_lightning','torch_geometric'
 ]
 
 EXTRAS = {

tests/test_model/test_gno.py

@@ -0,0 +1,28 @@
+import torch
+from pina.model import GNO
+
+output_channels = 5
+batch_size = 15
+
+
+def test_constructor():
+    input_channels = 1
+    output_channels = 1
+    #minimuum constructor
+    GNO(input_channels, output_channels, torch.rand(10000, 2))
+
+    #all constructor
+    GNO(input_channels, output_channels, torch.rand(100, 2),radius=0.5,inner_size=5,n_layers=5,func=torch.nn.ReLU)
+
+
+
+def test_forward():
+    input_channels = 1
+    output_channels = 1
+    input_ = torch.rand(batch_size, 1000, input_channels)
+    points=torch.rand(1000,2)
+    gno = GNO(input_channels, output_channels, points)
+    out = gno(input_)
+    assert out.shape == torch.Size([batch_size, points.shape[0], output_channels])
+
+

tutorials/tutorial_dev_9/tutorial_dev_9.py

@@ -0,0 +1,85 @@
+import torch
+from time import time
+from pina.model import GNO
+from pina import Condition,LabelTensor
+from pina.problem import AbstractProblem
+from pina.solvers import SupervisedSolver
+from pina.trainer import Trainer
+from pina.loss import LpLoss
+
+#Data generation
+
+torch.manual_seed(0)
+
+def sample_unit_circle(num_points):
+    radius=torch.rand(num_points,1)
+    angle=torch.rand(num_points,1)*2*torch.pi
+    x=radius*torch.cos(angle)
+    y=radius*torch.sin(angle)
+    data=torch.cat((x,y),dim=1)
+    return data
+
+#sin(a*x+b*y)
+def compute_input(data,theta):
+    data=data.reshape(1,-1,2)
+    z=torch.sin(theta[:,:,0]*data[:,:,0]+theta[:,:,1]*data[:,:,1])
+    return z
+
+#1+convolution of sin(a*x+b*y) with sin(x) over [0,2pi]x[0,2pi    ]
+def compute_output(data,theta):
+    data=data.reshape(1,-1,2)
+    z=1-4*torch.sin(torch.pi*theta[:,:,0])*torch.sin(torch.pi*theta[:,:,1])*torch.cos(theta[:,:,0]*(torch.pi*data[:,:,0])+theta[:,:,1]*(torch.pi*data[:,:,1]))/((theta[:,:,0]**2-1)*theta[:,:,1])
+    return z
+
+
+
+theta=1+0.01*torch.rand(300,1,2)
+data_coarse=sample_unit_circle(1000)
+output_coarse=compute_output(data_coarse,theta).unsqueeze(-1) 
+input_coarse=compute_input(data_coarse,theta).unsqueeze(-1)
+data_dense=sample_unit_circle(1000)
+output_dense=compute_output(data_dense,theta).unsqueeze(-1) 
+input_dense=compute_input(data_dense,theta).unsqueeze(-1)
+
+
+model=GNO(1,1,data_coarse,inner_size=500,n_layers=4)
+class GNOSolver(AbstractProblem):
+    input_variables=['input']
+    input_points=LabelTensor(input_coarse,input_variables)
+    output_variables=['output']
+    output_points=LabelTensor(output_coarse,output_variables)
+    conditions={"data":Condition(input_points=input_points,output_points=output_points)}
+
+batch_size=1
+problem=GNOSolver()
+solver=SupervisedSolver(problem,model,optimizer_kwargs={'lr':1e-3},optimizer=torch.optim.AdamW)
+trainer=Trainer(solver=solver,max_epochs=5,accelerator='cpu',enable_model_summary=False,batch_size=batch_size)
+loss=LpLoss(2,relative=True)
+
+start_time=time()
+trainer.train()
+end_time=time()
+print(end_time-start_time) 
+solver.neural_net=solver.neural_net.eval()
+loss=torch.nn.MSELoss()
+num_batches=len(input_coarse)//batch_size
+num=0
+dem=0
+for i in range(num_batches):
+    input_variables=['input']
+    myinput=LabelTensor(input_coarse[i].unsqueeze(0),input_variables)
+    tmp=model(myinput).detach().squeeze(0)
+    num=num+torch.linalg.norm(tmp-output_coarse[i])**2
+    dem=dem+torch.linalg.norm(output_coarse[i])**2
+print("Training mse loss is", torch.sqrt(num/dem))
+
+
+num=0
+dem=0
+for i in range(num_batches):
+    input_variables=['input']
+    myinput=LabelTensor(input_dense[i].unsqueeze(0),input_variables)
+    tmp=model.forward_eval(myinput,data_dense).detach().squeeze(0)
+    num=num+torch.linalg.norm(tmp-output_dense[i])**2
+    dem=dem+torch.linalg.norm(output_dense[i])**2
+print("Super Resolution mse loss is", torch.sqrt(num/dem))