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>
This commit is contained in:
Dario Coscia
2024-04-02 10:24:23 +02:00
committed by GitHub
parent e0aeb923f3
commit 766856494a
9 changed files with 499 additions and 1 deletions

View File

@@ -0,0 +1,58 @@
import torch
import pytest
from pina.model.layers import LowRankBlock
from pina import LabelTensor
input_dimensions=2
embedding_dimenion=1
rank=4
inner_size=20
n_layers=2
func=torch.nn.Tanh
bias=True
def test_constructor():
LowRankBlock(input_dimensions=input_dimensions,
embedding_dimenion=embedding_dimenion,
rank=rank,
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias)
def test_constructor_wrong():
with pytest.raises(ValueError):
LowRankBlock(input_dimensions=input_dimensions,
embedding_dimenion=embedding_dimenion,
rank=0.5,
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias)
def test_forward():
block = LowRankBlock(input_dimensions=input_dimensions,
embedding_dimenion=embedding_dimenion,
rank=rank,
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias)
data = LabelTensor(torch.rand(10, 30, 3), labels=['x', 'y', 'u'])
block(data.extract('u'), data.extract(['x', 'y']))
def test_backward():
block = LowRankBlock(input_dimensions=input_dimensions,
embedding_dimenion=embedding_dimenion,
rank=rank,
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias)
data = LabelTensor(torch.rand(10, 30, 3), labels=['x', 'y', 'u'])
data.requires_grad_(True)
out = block(data.extract('u'), data.extract(['x', 'y']))
loss = out.mean()
loss.backward()

View File

@@ -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)])