58 lines
1.5 KiB
Python
58 lines
1.5 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
|