133 lines
4.9 KiB
Python
133 lines
4.9 KiB
Python
import torch
|
|
import pytest
|
|
from pina import Condition, LabelTensor, Graph
|
|
from pina.condition import InputOutputPointsCondition, DomainEquationCondition
|
|
from pina.graph import RadiusGraph
|
|
from pina.problem import AbstractProblem, SpatialProblem
|
|
from pina.domain import CartesianDomain
|
|
from pina.equation.equation import Equation
|
|
from pina.equation.equation_factory import FixedValue
|
|
from pina.operators import laplacian
|
|
from pina.collector import Collector
|
|
|
|
|
|
# def test_supervised_tensor_collector():
|
|
# class SupervisedProblem(AbstractProblem):
|
|
# output_variables = None
|
|
# conditions = {
|
|
# 'data1' : Condition(input_points=torch.rand((10,2)),
|
|
# output_points=torch.rand((10,2))),
|
|
# 'data2' : Condition(input_points=torch.rand((20,2)),
|
|
# output_points=torch.rand((20,2))),
|
|
# 'data3' : Condition(input_points=torch.rand((30,2)),
|
|
# output_points=torch.rand((30,2))),
|
|
# }
|
|
# problem = SupervisedProblem()
|
|
# collector = Collector(problem)
|
|
# for v in collector.conditions_name.values():
|
|
# assert v in problem.conditions.keys()
|
|
# assert all(collector._is_conditions_ready.values())
|
|
|
|
def test_pinn_collector():
|
|
def laplace_equation(input_, output_):
|
|
force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
|
|
torch.sin(input_.extract(['y']) * torch.pi))
|
|
delta_u = laplacian(output_.extract(['u']), input_)
|
|
return delta_u - force_term
|
|
|
|
my_laplace = Equation(laplace_equation)
|
|
in_ = LabelTensor(torch.tensor([[0., 1.]], requires_grad=True), ['x', 'y'])
|
|
out_ = LabelTensor(torch.tensor([[0.]], requires_grad=True), ['u'])
|
|
|
|
class Poisson(SpatialProblem):
|
|
output_variables = ['u']
|
|
spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
|
|
|
|
conditions = {
|
|
'gamma1':
|
|
Condition(domain=CartesianDomain({
|
|
'x': [0, 1],
|
|
'y': 1
|
|
}),
|
|
equation=FixedValue(0.0)),
|
|
'gamma2':
|
|
Condition(domain=CartesianDomain({
|
|
'x': [0, 1],
|
|
'y': 0
|
|
}),
|
|
equation=FixedValue(0.0)),
|
|
'gamma3':
|
|
Condition(domain=CartesianDomain({
|
|
'x': 1,
|
|
'y': [0, 1]
|
|
}),
|
|
equation=FixedValue(0.0)),
|
|
'gamma4':
|
|
Condition(domain=CartesianDomain({
|
|
'x': 0,
|
|
'y': [0, 1]
|
|
}),
|
|
equation=FixedValue(0.0)),
|
|
'D':
|
|
Condition(domain=CartesianDomain({
|
|
'x': [0, 1],
|
|
'y': [0, 1]
|
|
}),
|
|
equation=my_laplace),
|
|
'data':
|
|
Condition(input_points=in_, output_points=out_)
|
|
}
|
|
|
|
def poisson_sol(self, pts):
|
|
return -(torch.sin(pts.extract(['x']) * torch.pi) *
|
|
torch.sin(pts.extract(['y']) * torch.pi)) / (
|
|
2 * torch.pi ** 2)
|
|
|
|
truth_solution = poisson_sol
|
|
|
|
problem = Poisson()
|
|
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
|
|
problem.discretise_domain(10, 'grid', domains=boundaries)
|
|
problem.discretise_domain(10, 'grid', domains='D')
|
|
|
|
collector = Collector(problem)
|
|
collector.store_fixed_data()
|
|
collector.store_sample_domains()
|
|
|
|
for k, v in problem.conditions.items():
|
|
if isinstance(v, InputOutputPointsCondition):
|
|
assert list(collector.data_collections[k].keys()) == [
|
|
'input_points', 'output_points']
|
|
|
|
for k, v in problem.conditions.items():
|
|
if isinstance(v, DomainEquationCondition):
|
|
assert list(collector.data_collections[k].keys()) == [
|
|
'input_points', 'equation']
|
|
|
|
|
|
def test_supervised_graph_collector():
|
|
pos = torch.rand((100, 3))
|
|
x = [torch.rand((100, 3)) for _ in range(10)]
|
|
graph_list_1 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
|
|
out_1 = torch.rand((10, 100, 3))
|
|
pos = torch.rand((50, 3))
|
|
x = [torch.rand((50, 3)) for _ in range(10)]
|
|
graph_list_2 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
|
|
out_2 = torch.rand((10, 50, 3))
|
|
|
|
class SupervisedProblem(AbstractProblem):
|
|
output_variables = None
|
|
conditions = {
|
|
'data1': Condition(input_points=graph_list_1,
|
|
output_points=out_1),
|
|
'data2': Condition(input_points=graph_list_2,
|
|
output_points=out_2),
|
|
}
|
|
|
|
problem = SupervisedProblem()
|
|
collector = Collector(problem)
|
|
collector.store_fixed_data()
|
|
# assert all(collector._is_conditions_ready.values())
|
|
for v in collector.conditions_name.values():
|
|
assert v in problem.conditions.keys()
|