folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement Dataset, Dataloader and DataModule class and fix SupervisedSolver

Browse Source
This commit is contained in:
FilippoOlivo
2024-10-16 11:24:37 +02:00
committed by Nicola Demo
parent b9753c34b2
commit c9304fb9bb
30 changed files with 770 additions and 784 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
tests/test_dataset.py
View File

@@ -1,44 +1,45 @@
import math
import torch
import pytest
from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
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.model import FeedForward
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))
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(
location=CartesianDomain({'x': [0, 1], 'y': 1}),
domain=CartesianDomain({'x': [0, 1], 'y': 1}),
equation=FixedValue(0.0)),
'gamma2': Condition(
location=CartesianDomain({'x': [0, 1], 'y': 0}),
domain=CartesianDomain({'x': [0, 1], 'y': 0}),
equation=FixedValue(0.0)),
'gamma3': Condition(
location=CartesianDomain({'x': 1, 'y': [0, 1]}),
domain=CartesianDomain({'x': 1, 'y': [0, 1]}),
equation=FixedValue(0.0)),
'gamma4': Condition(
location=CartesianDomain({'x': 0, 'y': [0, 1]}),
domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
equation=FixedValue(0.0)),
'D': Condition(
input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
@@ -48,75 +49,114 @@ class Poisson(SpatialProblem):
output_points=out_),
'data2': Condition(
input_points=in2_,
output_points=out2_)
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.pts.shape == (140, 2)
assert sample_dataset.pts.labels == ['x', 'y']
assert sample_dataset.condition_indeces.dtype == torch.int64
assert sample_dataset.condition_indeces.max() == torch.tensor(4)
assert sample_dataset.condition_indeces.min() == torch.tensor(0)
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 = DataPointDataset(poisson, device='cpu')
dataset = SupervisedDataset(poisson, device='cpu')
assert len(dataset) == 61
assert dataset.input_pts.shape == (61, 2)
assert dataset.input_pts.labels == ['x', 'y']
assert dataset.output_pts.shape == (61, 1 )
assert dataset.output_pts.labels == ['u']
assert dataset.condition_indeces.dtype == torch.int64
assert dataset.condition_indeces.max() == torch.tensor(1)
assert dataset.condition_indeces.min() == torch.tensor(0)
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():
sample_dataset = SamplePointDataset(poisson, device='cpu')
data_dataset = DataPointDataset(poisson, device='cpu')
loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
for batch in loader:
assert len(batch) in [2, 3]
assert batch['pts'].shape[0] <= 10
assert batch['pts'].requires_grad == True
assert batch['pts'].labels == ['x', 'y']
loader2 = SamplePointLoader(sample_dataset, data_dataset, batch_size=None)
assert len(list(loader2)) == 2
def test_loader2():
poisson2 = Poisson()
del poisson.conditions['data2']
del poisson2.conditions['data']
poisson2.discretise_domain(10, 'grid', locations=boundaries)
sample_dataset = SamplePointDataset(poisson, device='cpu')
data_dataset = DataPointDataset(poisson, device='cpu')
loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
for batch in loader:
assert len(batch) == 2 # only phys condtions
assert batch['pts'].shape[0] <= 10
assert batch['pts'].requires_grad == True
assert batch['pts'].labels == ['x', 'y']
def test_loader3():
poisson2 = Poisson()
del poisson.conditions['gamma1']
del poisson.conditions['gamma2']
del poisson.conditions['gamma3']
del poisson.conditions['gamma4']
del poisson.conditions['D']
sample_dataset = SamplePointDataset(poisson, device='cpu')
data_dataset = DataPointDataset(poisson, device='cpu')
loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
for batch in loader:
assert len(batch) == 2 # only phys condtions
assert batch['pts'].shape[0] <= 10
assert batch['pts'].requires_grad == True
assert batch['pts'].labels == ['x', 'y']
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()

222
tests/test_solvers/test_supervised_solver.py
View File

@@ -1,50 +1,27 @@
import torch
from pina.problem import AbstractProblem
import pytest
from pina.problem import AbstractProblem, SpatialProblem
from pina import Condition, LabelTensor
from pina.solvers import SupervisedSolver
from pina.trainer import Trainer
from pina.model import FeedForward
from pina.loss import LpLoss
from pina.solvers import GraphSupervisedSolver
from pina.equation.equation import Equation
from pina.equation.equation_factory import FixedValue
from pina.operators import laplacian
from pina.domain import CartesianDomain
from pina.trainer import Trainer
in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1'])
out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
class NeuralOperatorProblem(AbstractProblem):
input_variables = ['u_0', 'u_1']
output_variables = ['u']
domains = {
'pts': LabelTensor(
torch.rand(100, 2),
labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}}
)
}
conditions = {
'data' : Condition(
domain='pts',
output_points=LabelTensor(
torch.rand(100, 1),
labels={1: {'name': 'output', 'dof': ['u']}}
)
)
'data': Condition(input_points=in_, output_points=out_),
}
class NeuralOperatorProblemGraph(AbstractProblem):
input_variables = ['x', 'y', 'u_0', 'u_1']
output_variables = ['u']
domains = {
'pts': LabelTensor(
torch.rand(100, 4),
labels={1: {'name': 'space', 'dof': ['x', 'y', 'u_0', 'u_1']}}
)
}
conditions = {
'data' : Condition(
domain='pts',
output_points=LabelTensor(
torch.rand(100, 1),
labels={1: {'name': 'output', 'dof': ['u']}}
)
)
}
class myFeature(torch.nn.Module):
"""
@@ -61,117 +38,106 @@ class myFeature(torch.nn.Module):
problem = NeuralOperatorProblem()
problem_graph = NeuralOperatorProblemGraph()
# make the problem + extra feats
extra_feats = [myFeature()]
model = FeedForward(len(problem.input_variables),
len(problem.output_variables))
model = FeedForward(len(problem.input_variables), len(problem.output_variables))
model_extra_feats = FeedForward(
len(problem.input_variables) + 1,
len(problem.output_variables))
len(problem.input_variables) + 1, len(problem.output_variables))
def test_constructor():
SupervisedSolver(problem=problem, model=model)
# def test_constructor_extra_feats():
# SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats)
test_constructor()
'''
class AutoSolver(SupervisedSolver):
def forward(self, input):
from pina.graph import Graph
print(Graph)
print(input)
if not isinstance(input, Graph):
input = Graph.build('radius', nodes_coordinates=input, nodes_data=torch.rand(input.shape), radius=0.2)
print(input)
print(input.data.edge_index)
print(input.data)
g = self._model(input.data, edge_index=input.data.edge_index)
g.labels = {1: {'name': 'output', 'dof': ['u']}}
return g
du_dt_new = LabelTensor(self.model(graph).reshape(-1,1), labels = ['du'])
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
return du_dt_new
'''
class GraphModel(torch.nn.Module):
def __init__(self, in_channels, out_channels):
from torch_geometric.nn import GCNConv, NNConv
super().__init__()
self.conv1 = GCNConv(in_channels, 16)
self.conv2 = GCNConv(16, out_channels)
my_laplace = Equation(laplace_equation)
def forward(self, data, edge_index):
print(data)
x = data.x
print(x)
x = self.conv1(x, edge_index)
x = x.relu()
x = self.conv2(x, edge_index)
return x
def test_graph():
solver = GraphSupervisedSolver(problem=problem_graph, model=GraphModel(2, 1), loss=LpLoss(),
nodes_coordinates=['x', 'y'], nodes_data=['u_0', 'u_1'])
trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20)
trainer.train()
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
def test_wrong_constructor():
poisson_problem = Poisson()
with pytest.raises(ValueError):
SupervisedSolver(problem=poisson_problem, model=model)
def test_train_cpu():
solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss())
trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20)
solver = SupervisedSolver(problem=problem, model=model)
trainer = Trainer(solver=solver,
max_epochs=200,
accelerator='gpu',
batch_size=5,
train_size=1,
test_size=0.,
eval_size=0.)
trainer.train()
test_train_cpu()
# def test_train_restore():
# tmpdir = "tests/tmp_restore"
# solver = SupervisedSolver(problem=problem,
# model=model,
# extra_features=None,
# loss=LpLoss())
# trainer = Trainer(solver=solver,
# max_epochs=5,
# accelerator='cpu',
# default_root_dir=tmpdir)
# trainer.train()
# ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
# t = ntrainer.train(
# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
# import shutil
# shutil.rmtree(tmpdir)
def test_extra_features_constructor():
SupervisedSolver(problem=problem,
model=model_extra_feats,
extra_features=extra_feats)
# def test_train_load():
# tmpdir = "tests/tmp_load"
# solver = SupervisedSolver(problem=problem,
# model=model,
# extra_features=None,
# loss=LpLoss())
# trainer = Trainer(solver=solver,
# max_epochs=15,
# accelerator='cpu',
# default_root_dir=tmpdir)
# trainer.train()
# new_solver = SupervisedSolver.load_from_checkpoint(
# f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
# problem = problem, model=model)
# test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
# assert new_solver.forward(test_pts).shape == (20, 1)
# assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
# torch.testing.assert_close(
# new_solver.forward(test_pts),
# solver.forward(test_pts))
# import shutil
# shutil.rmtree(tmpdir)
# def test_train_extra_feats_cpu():
# pinn = SupervisedSolver(problem=problem,
# model=model_extra_feats,
# extra_features=extra_feats)
# trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
# trainer.train()
test_graph()
def test_extra_features_train_cpu():
solver = SupervisedSolver(problem=problem,
model=model_extra_feats,
extra_features=extra_feats)
trainer = Trainer(solver=solver,
max_epochs=200,
accelerator='gpu',
batch_size=5)
trainer.train()