Files
PINA/tests/test_problem.py
2025-03-19 17:46:35 +01:00

94 lines
3.3 KiB
Python

import torch
import pytest
from pina.problem.zoo import Poisson2DSquareProblem as Poisson
from pina import LabelTensor
from pina.domain import Union
from pina.domain import CartesianDomain
def test_discretise_domain():
n = 10
poisson_problem = Poisson()
boundaries = ['g1', 'g2', 'g3', 'g4']
poisson_problem.discretise_domain(n, 'grid', domains=boundaries)
for b in boundaries:
assert poisson_problem.discretised_domains[b].shape[0] == n
poisson_problem.discretise_domain(n, 'random', domains=boundaries)
for b in boundaries:
assert poisson_problem.discretised_domains[b].shape[0] == n
poisson_problem.discretise_domain(n, 'grid', domains=['D'])
assert poisson_problem.discretised_domains['D'].shape[0] == n ** 2
poisson_problem.discretise_domain(n, 'random', domains=['D'])
assert poisson_problem.discretised_domains['D'].shape[0] == n
poisson_problem.discretise_domain(n, 'latin', domains=['D'])
assert poisson_problem.discretised_domains['D'].shape[0] == n
poisson_problem.discretise_domain(n, 'lh', domains=['D'])
assert poisson_problem.discretised_domains['D'].shape[0] == n
poisson_problem.discretise_domain(n)
def test_variables_correct_order_sampling():
n = 10
poisson_problem = Poisson()
poisson_problem.discretise_domain(n,
'grid',
domains=['D'])
assert poisson_problem.discretised_domains['D'].labels == sorted(
poisson_problem.input_variables)
poisson_problem.discretise_domain(n, 'grid', domains=['D'])
assert poisson_problem.discretised_domains['D'].labels == sorted(
poisson_problem.input_variables)
def test_add_points():
poisson_problem = Poisson()
poisson_problem.discretise_domain(0,
'random',
domains=['D'])
new_pts = LabelTensor(torch.tensor([[0.5, -0.5]]), labels=['x', 'y'])
poisson_problem.add_points({'D': new_pts})
assert torch.isclose(poisson_problem.discretised_domains['D'].extract('x'),
new_pts.extract('x'))
assert torch.isclose(poisson_problem.discretised_domains['D'].extract('y'),
new_pts.extract('y'))
@pytest.mark.parametrize(
"mode",
[
'random',
'grid'
]
)
def test_custom_sampling_logic(mode):
poisson_problem = Poisson()
sampling_rules = {
'x': {'n': 100, 'mode': mode},
'y': {'n': 50, 'mode': mode}
}
poisson_problem.discretise_domain(sample_rules=sampling_rules)
for domain in ['g1', 'g2', 'g3', 'g4']:
assert poisson_problem.discretised_domains[domain].shape[0] == 100 * 50
assert poisson_problem.discretised_domains[domain].labels == ['x', 'y']
@pytest.mark.parametrize(
"mode",
[
'random',
'grid'
]
)
def test_wrong_custom_sampling_logic(mode):
d2 = CartesianDomain({'x': [1,2], 'y': [0,1] })
poisson_problem = Poisson()
poisson_problem.domains['D'] = Union([poisson_problem.domains['D'], d2])
sampling_rules = {
'x': {'n': 100, 'mode': mode},
'y': {'n': 50, 'mode': mode}
}
with pytest.raises(RuntimeError):
poisson_problem.discretise_domain(sample_rules=sampling_rules)