162 lines
6.3 KiB
Python
162 lines
6.3 KiB
Python
import math
|
|
import torch
|
|
from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, UnsupervisedDataset, unsupervised_dataset
|
|
from pina.data import PinaDataLoader
|
|
from pina import LabelTensor, Condition
|
|
from pina.equation import Equation
|
|
from pina.domain import CartesianDomain
|
|
from pina.problem import SpatialProblem
|
|
from pina.operators import laplacian
|
|
from pina.equation.equation_factory import FixedValue
|
|
|
|
|
|
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.]]), ['x', 'y'])
|
|
out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
|
|
in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
|
|
out2_ = LabelTensor(torch.rand(60, 1), ['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(
|
|
input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
|
|
equation=my_laplace),
|
|
'data': Condition(
|
|
input_points=in_,
|
|
output_points=out_),
|
|
'data2': Condition(
|
|
input_points=in2_,
|
|
output_points=out2_),
|
|
'unsupervised': Condition(
|
|
input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
|
|
conditional_variables=LabelTensor(torch.ones(size=(45, 1)), ['alpha']),
|
|
),
|
|
'unsupervised2': Condition(
|
|
input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
|
|
conditional_variables=LabelTensor(torch.ones(size=(90, 1)), ['alpha']),
|
|
)
|
|
}
|
|
|
|
|
|
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
|
|
poisson = Poisson()
|
|
poisson.discretise_domain(10, 'grid', locations=boundaries)
|
|
|
|
|
|
def test_sample():
|
|
sample_dataset = SamplePointDataset(poisson, device='cpu')
|
|
assert len(sample_dataset) == 140
|
|
assert sample_dataset.input_points.shape == (140, 2)
|
|
assert sample_dataset.input_points.labels == ['x', 'y']
|
|
assert sample_dataset.condition_indices.dtype == torch.uint8
|
|
assert sample_dataset.condition_indices.max() == torch.tensor(4)
|
|
assert sample_dataset.condition_indices.min() == torch.tensor(0)
|
|
|
|
|
|
def test_data():
|
|
dataset = SupervisedDataset(poisson, device='cpu')
|
|
assert len(dataset) == 61
|
|
assert dataset['input_points'].shape == (61, 2)
|
|
assert dataset.input_points.shape == (61, 2)
|
|
assert dataset['input_points'].labels == ['x', 'y']
|
|
assert dataset.input_points.labels == ['x', 'y']
|
|
assert dataset['input_points', 3:].shape == (58, 2)
|
|
assert dataset[3:][1].labels == ['u']
|
|
assert dataset.output_points.shape == (61, 1)
|
|
assert dataset.output_points.labels == ['u']
|
|
assert dataset.condition_indices.dtype == torch.uint8
|
|
assert dataset.condition_indices.max() == torch.tensor(1)
|
|
assert dataset.condition_indices.min() == torch.tensor(0)
|
|
|
|
|
|
def test_unsupervised():
|
|
dataset = UnsupervisedDataset(poisson, device='cpu')
|
|
assert len(dataset) == 135
|
|
assert dataset.input_points.shape == (135, 2)
|
|
assert dataset.input_points.labels == ['x', 'y']
|
|
assert dataset.input_points[3:].shape == (132, 2)
|
|
|
|
assert dataset.conditional_variables.shape == (135, 1)
|
|
assert dataset.conditional_variables.labels == ['alpha']
|
|
assert dataset.condition_indices.dtype == torch.uint8
|
|
assert dataset.condition_indices.max() == torch.tensor(1)
|
|
assert dataset.condition_indices.min() == torch.tensor(0)
|
|
|
|
|
|
def test_data_module():
|
|
data_module = PinaDataModule(poisson, device='cpu')
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
assert isinstance(loader, PinaDataLoader)
|
|
assert isinstance(loader, PinaDataLoader)
|
|
|
|
data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False)
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
assert len(loader) == 24
|
|
for i in loader:
|
|
assert len(i) <= 10
|
|
len_ref = sum([math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
|
|
len_real = sum([len(dataset) for dataset in data_module.splits['train'].values()])
|
|
assert len_ref == len_real
|
|
|
|
supervised_dataset = SupervisedDataset(poisson, device='cpu')
|
|
data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[supervised_dataset])
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
for batch in loader:
|
|
assert len(batch) <= 10
|
|
|
|
physics_dataset = SamplePointDataset(poisson, device='cpu')
|
|
data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[physics_dataset])
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
for batch in loader:
|
|
assert len(batch) <= 10
|
|
|
|
unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu')
|
|
data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[unsupervised_dataset])
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
for batch in loader:
|
|
assert len(batch) <= 10
|
|
|
|
|
|
def test_loader():
|
|
data_module = PinaDataModule(poisson, device='cpu', batch_size=10)
|
|
data_module.setup()
|
|
loader = data_module.train_dataloader()
|
|
assert isinstance(loader, PinaDataLoader)
|
|
assert len(loader) == 24
|
|
for i in loader:
|
|
assert len(i) <= 10
|
|
assert i.supervised.input_points.labels == ['x', 'y']
|
|
assert i.physics.input_points.labels == ['x', 'y']
|
|
assert i.unsupervised.input_points.labels == ['x', 'y']
|
|
assert i.supervised.input_points.requires_grad == True
|
|
assert i.physics.input_points.requires_grad == True
|
|
assert i.unsupervised.input_points.requires_grad == True
|
|
test_loader() |