111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from pina.model.block 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.0, periods=1)
|
|
PeriodicBoundaryEmbedding(
|
|
input_dimension=1, periods=1, output_dimension=1.0
|
|
)
|
|
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.0], [1.0]])
|
|
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()
|