e516e779f9
* Unify integral kernel NO architectures with NeuralKernelOperator * Implement FNO based on NeuralKernelOperator * modify doc for FNO and add for FourierIntegralKernel, NeuralKernelOperator * adding tests --------- Co-authored-by: Dario Coscia <dariocoscia@Dario-Coscia.local> Co-authored-by: Dario Coscia <dariocoscia@dhcp-015.eduroam.sissa.it>
41 lines
1.7 KiB
Python
41 lines
1.7 KiB
Python
import torch
|
|
from pina.model import KernelNeuralOperator, FeedForward
|
|
|
|
input_dim = 2
|
|
output_dim = 4
|
|
embedding_dim = 24
|
|
batch_size = 10
|
|
numb = 256
|
|
data = torch.rand(size=(batch_size, numb, input_dim), requires_grad=True)
|
|
output_shape = torch.Size([batch_size, numb, output_dim])
|
|
|
|
|
|
lifting_operator = FeedForward(input_dimensions=input_dim, output_dimensions=embedding_dim)
|
|
projection_operator = FeedForward(input_dimensions=embedding_dim, output_dimensions=output_dim)
|
|
integral_kernels = torch.nn.Sequential(FeedForward(input_dimensions=embedding_dim,
|
|
output_dimensions=embedding_dim),
|
|
FeedForward(input_dimensions=embedding_dim,
|
|
output_dimensions=embedding_dim),)
|
|
|
|
def test_constructor():
|
|
KernelNeuralOperator(lifting_operator=lifting_operator,
|
|
integral_kernels=integral_kernels,
|
|
projection_operator=projection_operator)
|
|
|
|
def test_forward():
|
|
operator = KernelNeuralOperator(lifting_operator=lifting_operator,
|
|
integral_kernels=integral_kernels,
|
|
projection_operator=projection_operator)
|
|
out = operator(data)
|
|
assert out.shape == output_shape
|
|
|
|
def test_backward():
|
|
operator = KernelNeuralOperator(lifting_operator=lifting_operator,
|
|
integral_kernels=integral_kernels,
|
|
projection_operator=projection_operator)
|
|
out = operator(data)
|
|
loss = torch.nn.functional.mse_loss(out, torch.zeros_like(out))
|
|
loss.backward()
|
|
grad = data.grad
|
|
assert grad.shape == data.shape
|