Low Rank Neural Operator (#270)

* add the Low Rank Neural Operator as Model * add the Low Rank Layer as Layer * adding tests * adding doc --------- Co-authored-by: Dario Coscia <dariocoscia@Dario-Coscia.local> Co-authored-by: Nicola Demo <demo.nicola@gmail.com>
2024-04-02 10:24:23 +02:00
parent e0aeb923f3
commit 766856494a
9 changed files with 499 additions and 1 deletions
--- a/tests/test_model/test_lno.py
+++ b/tests/test_model/test_lno.py
@@ -0,0 +1,141 @@
+import torch
+from pina.model import LowRankNeuralOperator
+from pina import LabelTensor
+import pytest
+
+
+batch_size = 15
+n_layers = 4
+embedding_dim = 24
+func = torch.nn.Tanh
+rank = 4
+n_kernel_layers = 3
+field_indices = ['u']
+coordinates_indices = ['x', 'y']
+
+def test_constructor():
+    # working constructor
+    lifting_net = torch.nn.Linear(len(coordinates_indices) + len(field_indices),
+                                  embedding_dim)
+    projecting_net = torch.nn.Linear(embedding_dim +  len(coordinates_indices),
+                                     len(field_indices))
+    LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+
+    # not working constructor
+    with pytest.raises(ValueError):
+        LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=3.2, # wrong
+        rank=rank)
+
+        LowRankNeuralOperator(
+        lifting_net=[0], # wrong
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank) 
+
+        LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=[0], # wront
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+
+        LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=[0], #wrong
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+
+        LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=[0], #wrong
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+
+        lifting_net = torch.nn.Linear(len(coordinates_indices),
+                                    embedding_dim)
+        LowRankNeuralOperator(
+            lifting_net=lifting_net,
+            projecting_net=projecting_net,
+            coordinates_indices=coordinates_indices,
+            field_indices=field_indices,
+            n_kernel_layers=n_kernel_layers,
+            rank=rank)
+        
+        lifting_net = torch.nn.Linear(len(coordinates_indices) + len(field_indices),
+                                  embedding_dim)
+        projecting_net = torch.nn.Linear(embedding_dim,
+                                     len(field_indices))
+        LowRankNeuralOperator(
+            lifting_net=lifting_net,
+            projecting_net=projecting_net,
+            coordinates_indices=coordinates_indices,
+            field_indices=field_indices,
+            n_kernel_layers=n_kernel_layers,
+            rank=rank)
+    
+
+def test_forward():
+    lifting_net = torch.nn.Linear(len(coordinates_indices) + len(field_indices),
+                                  embedding_dim)
+    projecting_net = torch.nn.Linear(embedding_dim +  len(coordinates_indices),
+                                     len(field_indices))
+    lno = LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+    
+    input_ = LabelTensor(
+        torch.rand(batch_size, 100,
+        len(coordinates_indices) + len(field_indices)),
+        coordinates_indices + field_indices)
+
+    out = lno(input_)
+    assert out.shape == torch.Size(
+        [batch_size, input_.shape[1], len(field_indices)])
+
+
+def test_backward():
+    lifting_net = torch.nn.Linear(len(coordinates_indices) + len(field_indices),
+                                  embedding_dim)
+    projecting_net = torch.nn.Linear(embedding_dim +  len(coordinates_indices),
+                                     len(field_indices))
+    lno=LowRankNeuralOperator(
+        lifting_net=lifting_net,
+        projecting_net=projecting_net,
+        coordinates_indices=coordinates_indices,
+        field_indices=field_indices,
+        n_kernel_layers=n_kernel_layers,
+        rank=rank)
+    input_ = LabelTensor(
+        torch.rand(batch_size, 100,
+                   len(coordinates_indices) + len(field_indices)),
+                   coordinates_indices + field_indices)
+    input_ = input_.requires_grad_()
+    out = lno(input_)
+    tmp = torch.linalg.norm(out)
+    tmp.backward()
+    grad = input_.grad
+    assert grad.shape == torch.Size(
+        [batch_size, input_.shape[1],
+         len(coordinates_indices) + len(field_indices)])