109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
import torch
|
|
import pytest
|
|
import math
|
|
|
|
from pina.model.block.rbf_block import RBFBlock
|
|
|
|
x = torch.linspace(-1, 1, 100)
|
|
toy_params = torch.linspace(0, 1, 10).unsqueeze(1)
|
|
toy_snapshots = torch.vstack([torch.exp(-(x**2)) * c for c in toy_params])
|
|
toy_params_test = torch.linspace(0, 1, 3).unsqueeze(1)
|
|
toy_snapshots_test = torch.vstack(
|
|
[torch.exp(-(x**2)) * c for c in toy_params_test]
|
|
)
|
|
|
|
kernels = [
|
|
"linear",
|
|
"thin_plate_spline",
|
|
"cubic",
|
|
"quintic",
|
|
"multiquadric",
|
|
"inverse_multiquadric",
|
|
"inverse_quadratic",
|
|
"gaussian",
|
|
]
|
|
|
|
noscale_invariant_kernels = [
|
|
"multiquadric",
|
|
"inverse_multiquadric",
|
|
"inverse_quadratic",
|
|
"gaussian",
|
|
]
|
|
|
|
scale_invariant_kernels = ["linear", "thin_plate_spline", "cubic", "quintic"]
|
|
|
|
|
|
def test_constructor_default():
|
|
rbf = RBFBlock()
|
|
assert rbf.kernel == "thin_plate_spline"
|
|
assert rbf.epsilon == 1
|
|
assert rbf.smoothing == 0.0
|
|
|
|
|
|
@pytest.mark.parametrize("kernel", kernels)
|
|
@pytest.mark.parametrize("epsilon", [0.1, 1.0, 10.0])
|
|
def test_constructor_epsilon(kernel, epsilon):
|
|
if kernel in scale_invariant_kernels:
|
|
rbf = RBFBlock(kernel=kernel)
|
|
assert rbf.kernel == kernel
|
|
assert rbf.epsilon == 1
|
|
elif kernel in noscale_invariant_kernels:
|
|
with pytest.raises(ValueError):
|
|
rbf = RBFBlock(kernel=kernel)
|
|
rbf = RBFBlock(kernel=kernel, epsilon=epsilon)
|
|
assert rbf.kernel == kernel
|
|
assert rbf.epsilon == epsilon
|
|
|
|
assert rbf.smoothing == 0.0
|
|
|
|
|
|
@pytest.mark.parametrize("kernel", kernels)
|
|
@pytest.mark.parametrize("epsilon", [0.1, 1.0, 10.0])
|
|
@pytest.mark.parametrize("degree", [2, 3, 4])
|
|
@pytest.mark.parametrize("smoothing", [1e-5, 1e-3, 1e-1])
|
|
def test_constructor_all(kernel, epsilon, degree, smoothing):
|
|
rbf = RBFBlock(
|
|
kernel=kernel, epsilon=epsilon, degree=degree, smoothing=smoothing
|
|
)
|
|
assert rbf.kernel == kernel
|
|
assert rbf.epsilon == epsilon
|
|
assert rbf.degree == degree
|
|
assert rbf.smoothing == smoothing
|
|
assert rbf.y == None
|
|
assert rbf.d == None
|
|
assert rbf.powers == None
|
|
assert rbf._shift == None
|
|
assert rbf._scale == None
|
|
assert rbf._coeffs == None
|
|
|
|
|
|
def test_fit():
|
|
rbf = RBFBlock()
|
|
rbf.fit(toy_params, toy_snapshots)
|
|
ndim = toy_params.shape[1]
|
|
torch.testing.assert_close(rbf.y, toy_params)
|
|
torch.testing.assert_close(rbf.d, toy_snapshots)
|
|
assert rbf.powers.shape == (math.comb(rbf.degree + ndim, ndim), ndim)
|
|
assert rbf._shift.shape == (ndim,)
|
|
assert rbf._scale.shape == (ndim,)
|
|
assert rbf._coeffs.shape == (
|
|
rbf.powers.shape[0] + toy_snapshots.shape[0],
|
|
toy_snapshots.shape[1],
|
|
)
|
|
|
|
|
|
def test_forward():
|
|
rbf = RBFBlock()
|
|
rbf.fit(toy_params, toy_snapshots)
|
|
c = rbf(toy_params)
|
|
assert c.shape == toy_snapshots.shape
|
|
torch.testing.assert_close(c, toy_snapshots)
|
|
|
|
|
|
def test_forward_unseen_parameters():
|
|
rbf = RBFBlock()
|
|
rbf.fit(toy_params, toy_snapshots)
|
|
c = rbf(toy_params_test)
|
|
assert c.shape == toy_snapshots_test.shape
|
|
torch.testing.assert_close(c, toy_snapshots_test)
|