103 lines
3.5 KiB
Python
103 lines
3.5 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from pina.model.layers import PeriodicBoundaryEmbedding
|
|
from pina import LabelTensor
|
|
|
|
# test tolerance
|
|
tol = 1e-6
|
|
|
|
def check_same_columns(tensor):
|
|
# Get the first column and compute residual
|
|
residual = tensor - tensor[0]
|
|
zeros = torch.zeros_like(residual)
|
|
# Compare each column with the first column
|
|
all_same = torch.allclose(input=residual,other=zeros,atol=tol)
|
|
return all_same
|
|
|
|
def grad(u, x):
|
|
"""
|
|
Compute the first derivative of u with respect to x.
|
|
"""
|
|
return torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u),
|
|
create_graph=True, allow_unused=True,
|
|
retain_graph=True)[0]
|
|
|
|
def test_constructor():
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods=2)
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods={'x': 3, 'y' : 4})
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods={0: 3, 1 : 4})
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods=2, output_dimension=10)
|
|
with pytest.raises(TypeError):
|
|
PeriodicBoundaryEmbedding()
|
|
with pytest.raises(ValueError):
|
|
PeriodicBoundaryEmbedding(input_dimension=1., periods=1)
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods=1, output_dimension=1.)
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods={'x':'x'})
|
|
PeriodicBoundaryEmbedding(input_dimension=1, periods={0:'x'})
|
|
|
|
|
|
@pytest.mark.parametrize("period", [1, 4, 10])
|
|
@pytest.mark.parametrize("input_dimension", [1, 2, 3])
|
|
def test_forward_same_period(input_dimension, period):
|
|
func = torch.nn.Sequential(
|
|
PeriodicBoundaryEmbedding(input_dimension=input_dimension,
|
|
output_dimension=60, periods=period),
|
|
torch.nn.Tanh(),
|
|
torch.nn.Linear(60, 60),
|
|
torch.nn.Tanh(),
|
|
torch.nn.Linear(60, 1)
|
|
)
|
|
# coordinates
|
|
x = period * torch.tensor([[0.],[1.]])
|
|
if input_dimension == 2:
|
|
x = torch.cartesian_prod(x.flatten(),x.flatten())
|
|
elif input_dimension == 3:
|
|
x = torch.cartesian_prod(x.flatten(),x.flatten(),x.flatten())
|
|
x.requires_grad = True
|
|
# output
|
|
f = func(x)
|
|
assert check_same_columns(f)
|
|
|
|
|
|
|
|
# def test_forward_same_period_labels():
|
|
# func = torch.nn.Sequential(
|
|
# PeriodicBoundaryEmbedding(input_dimension=2,
|
|
# output_dimension=60, periods={'x':1, 'y':2}),
|
|
# torch.nn.Tanh(),
|
|
# torch.nn.Linear(60, 60),
|
|
# torch.nn.Tanh(),
|
|
# torch.nn.Linear(60, 1)
|
|
# )
|
|
# # coordinates
|
|
# tensor = torch.tensor([[0., 0.], [0., 2.], [1., 0.], [1., 2.]])
|
|
# with pytest.raises(RuntimeError):
|
|
# func(tensor)
|
|
# tensor = tensor.as_subclass(LabelTensor)
|
|
# tensor.labels = ['x', 'y']
|
|
# tensor.requires_grad = True
|
|
# # output
|
|
# f = func(tensor)
|
|
# assert check_same_columns(f)
|
|
|
|
# def test_forward_same_period_index():
|
|
# func = torch.nn.Sequential(
|
|
# PeriodicBoundaryEmbedding(input_dimension=2,
|
|
# output_dimension=60, periods={0:1, 1:2}),
|
|
# torch.nn.Tanh(),
|
|
# torch.nn.Linear(60, 60),
|
|
# torch.nn.Tanh(),
|
|
# torch.nn.Linear(60, 1)
|
|
# )
|
|
# # coordinates
|
|
# tensor = torch.tensor([[0., 0.], [0., 2.], [1., 0.], [1., 2.]])
|
|
# tensor.requires_grad = True
|
|
# # output
|
|
# f = func(tensor)
|
|
# assert check_same_columns(f)
|
|
# tensor = tensor.as_subclass(LabelTensor)
|
|
# tensor.labels = ['x', 'y']
|
|
# # output
|
|
# f = func(tensor)
|
|
# assert check_same_columns(f) |