96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from pina.model.layers import PeriodicBoundaryEmbedding, FourierFeatureEmbedding
|
|
|
|
# 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():
|
|
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_backward_same_period_PeriodicBoundaryEmbedding(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)
|
|
# compute backward
|
|
loss = f.mean()
|
|
loss.backward()
|
|
|
|
def test_constructor_FourierFeatureEmbedding():
|
|
FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
|
|
sigma=1)
|
|
with pytest.raises(TypeError):
|
|
FourierFeatureEmbedding()
|
|
with pytest.raises(RuntimeError):
|
|
FourierFeatureEmbedding(input_dimension=1, output_dimension=3, sigma=1)
|
|
with pytest.raises(ValueError):
|
|
FourierFeatureEmbedding(input_dimension='x', output_dimension=20,
|
|
sigma=1)
|
|
FourierFeatureEmbedding(input_dimension=1, output_dimension='x',
|
|
sigma=1)
|
|
FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
|
|
sigma='x')
|
|
|
|
@pytest.mark.parametrize("output_dimension", [2, 4, 6])
|
|
@pytest.mark.parametrize("input_dimension", [1, 2, 3])
|
|
@pytest.mark.parametrize("sigma", [10, 1, 0.1])
|
|
def test_forward_backward_FourierFeatureEmbedding(input_dimension,
|
|
output_dimension,
|
|
sigma):
|
|
func = FourierFeatureEmbedding(input_dimension, output_dimension,
|
|
sigma)
|
|
# coordinates
|
|
x = torch.rand((10, input_dimension), requires_grad=True)
|
|
# output
|
|
f = func(x)
|
|
assert f.shape[-1] == output_dimension
|
|
# compute backward
|
|
loss = f.mean()
|
|
loss.backward() |