Documentation for v0.1 version (#199)

* Adding Equations, solving typos
* improve _code.rst
* the team rst and restuctore index.rst
* fixing errors

---------

Co-authored-by: Dario Coscia <dariocoscia@dhcp-015.eduroam.sissa.it>
This commit is contained in:
Dario Coscia
2023-11-08 14:39:00 +01:00
committed by Nicola Demo
parent 3f9305d475
commit 8b7b61b3bd
144 changed files with 2741 additions and 1766 deletions

View File

@@ -15,6 +15,7 @@ def func(x, mu1, mu2):
norm = x[:, 0]**2 + x[:, 1]**2
return torch.exp(-(x_m1 + x_m2))
class ParametricGaussian(AbstractProblem):
output_variables = [f'u_{i}' for i in range(900)]
@@ -24,7 +25,7 @@ class ParametricGaussian(AbstractProblem):
params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])
# define domain
x = torch.linspace(-1, 1, 30)
x = torch.linspace(-1, 1, 30)
domain = torch.cartesian_prod(x, x)
triang = tri.Triangulation(domain[:, 0], domain[:, 1])
sol = []
@@ -34,15 +35,18 @@ class ParametricGaussian(AbstractProblem):
# define conditions
conditions = {
'data': Condition(
input_points=params,
output_points=snapshots)
'data': Condition(input_points=params, output_points=snapshots)
}
# simple Generator Network
class Generator(nn.Module):
def __init__(self, input_dimension, parameters_dimension,
noise_dimension, activation=torch.nn.SiLU):
def __init__(self,
input_dimension,
parameters_dimension,
noise_dimension,
activation=torch.nn.SiLU):
super().__init__()
self._noise_dimension = noise_dimension
@@ -53,13 +57,12 @@ class Generator(nn.Module):
self._activation(),
torch.nn.Linear(input_dimension // 6, input_dimension // 3),
self._activation(),
torch.nn.Linear(input_dimension // 3, input_dimension)
)
torch.nn.Linear(input_dimension // 3, input_dimension))
self.condition = torch.nn.Sequential(
torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
self._activation(),
torch.nn.Linear(2 * self._noise_dimension, 5 * self._noise_dimension)
)
torch.nn.Linear(2 * self._noise_dimension,
5 * self._noise_dimension))
def forward(self, param):
# uniform sampling in [-1, 1]
@@ -78,8 +81,12 @@ class Generator(nn.Module):
# Simple Discriminator Network
class Discriminator(nn.Module):
def __init__(self, input_dimension, parameter_dimension,
hidden_dimension, activation=torch.nn.ReLU):
def __init__(self,
input_dimension,
parameter_dimension,
hidden_dimension,
activation=torch.nn.ReLU):
super().__init__()
self._activation = activation
@@ -88,10 +95,9 @@ class Discriminator(nn.Module):
self._activation(),
torch.nn.Linear(input_dimension // 3, input_dimension // 6),
self._activation(),
torch.nn.Linear(input_dimension // 6, hidden_dimension)
)
torch.nn.Linear(input_dimension // 6, hidden_dimension))
self.decoding = torch.nn.Sequential(
torch.nn.Linear(2*hidden_dimension, input_dimension // 6),
torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
self._activation(),
torch.nn.Linear(input_dimension // 6, input_dimension // 3),
self._activation(),
@@ -101,9 +107,8 @@ class Discriminator(nn.Module):
self.condition = torch.nn.Sequential(
torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
self._activation(),
torch.nn.Linear(hidden_dimension // 2, hidden_dimension)
)
torch.nn.Linear(hidden_dimension // 2, hidden_dimension))
def forward(self, data):
x, condition = data
encoding = self.encoding(x)
@@ -114,49 +119,49 @@ class Discriminator(nn.Module):
problem = ParametricGaussian()
def test_constructor():
GAROM(problem = problem,
generator = Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator = Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64)
)
GAROM(problem=problem,
generator=Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator=Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64))
def test_train_cpu():
solver = GAROM(problem = problem,
generator = Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator = Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64)
)
solver = GAROM(problem=problem,
generator=Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator=Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64))
trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
trainer.train()
def test_sample():
solver = GAROM(problem = problem,
generator = Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator = Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64)
)
solver = GAROM(problem=problem,
generator=Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator=Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64))
solver.sample(problem.params)
assert solver.sample(problem.params).shape == problem.snapshots.shape
def test_forward():
solver = GAROM(problem = problem,
generator = Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator = Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64)
)
solver = GAROM(problem=problem,
generator=Generator(input_dimension=900,
parameters_dimension=2,
noise_dimension=12),
discriminator=Discriminator(input_dimension=900,
parameter_dimension=2,
hidden_dimension=64))
solver(problem.params, mc_steps=100, variance=True)
assert solver(problem.params).shape == problem.snapshots.shape
assert solver(problem.params).shape == problem.snapshots.shape

View File

@@ -1,30 +1,31 @@
import torch
import pytest
from pina.problem import SpatialProblem
from pina.operators import laplacian
from pina.geometry import CartesianDomain
from pina import Condition, LabelTensor, PINN
from pina import Condition, LabelTensor
from pina.solvers import PINN
from pina.trainer import Trainer
from pina.model import FeedForward
from pina.equation.equation import Equation
from pina.equation.equation_factory import FixedValue
from pina.plotter import Plotter
from pina.loss import LpLoss
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]})
@@ -54,42 +55,48 @@ class Poisson(SpatialProblem):
}
def poisson_sol(self, pts):
return -(
torch.sin(pts.extract(['x'])*torch.pi) *
torch.sin(pts.extract(['y'])*torch.pi)
)/(2*torch.pi**2)
return -(torch.sin(pts.extract(['x']) * torch.pi) *
torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
truth_solution = poisson_sol
class myFeature(torch.nn.Module):
"""
Feature: sin(x)
"""
def __init__(self):
super(myFeature, self).__init__()
def forward(self, x):
t = (torch.sin(x.extract(['x'])*torch.pi) *
torch.sin(x.extract(['y'])*torch.pi))
t = (torch.sin(x.extract(['x']) * torch.pi) *
torch.sin(x.extract(['y']) * torch.pi))
return LabelTensor(t, ['sin(x)sin(y)'])
# make the problem
poisson_problem = Poisson()
model = FeedForward(len(poisson_problem.input_variables),len(poisson_problem.output_variables))
model_extra_feats = FeedForward(len(poisson_problem.input_variables)+1,len(poisson_problem.output_variables))
model = FeedForward(len(poisson_problem.input_variables),
len(poisson_problem.output_variables))
model_extra_feats = FeedForward(
len(poisson_problem.input_variables) + 1,
len(poisson_problem.output_variables))
extra_feats = [myFeature()]
def test_constructor():
PINN(problem = poisson_problem, model=model, extra_features=None)
PINN(problem=poisson_problem, model=model, extra_features=None)
def test_constructor_extra_feats():
model_extra_feats = FeedForward(len(poisson_problem.input_variables)+1,len(poisson_problem.output_variables))
PINN(problem = poisson_problem, model=model_extra_feats, extra_features=extra_feats)
model_extra_feats = FeedForward(
len(poisson_problem.input_variables) + 1,
len(poisson_problem.output_variables))
PINN(problem=poisson_problem,
model=model_extra_feats,
extra_features=extra_feats)
def test_train_cpu():
poisson_problem = Poisson()
@@ -100,14 +107,21 @@ def test_train_cpu():
trainer = Trainer(solver=pinn, max_epochs=1, accelerator='cpu', batch_size=20)
trainer.train()
def test_train_restore():
tmpdir = "tests/tmp_restore"
poisson_problem = Poisson()
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu', default_root_dir=tmpdir)
pinn = PINN(problem=poisson_problem,
model=model,
extra_features=None,
loss=LpLoss())
trainer = Trainer(solver=pinn,
max_epochs=5,
accelerator='cpu',
default_root_dir=tmpdir)
trainer.train()
ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
t = ntrainer.train(
@@ -115,31 +129,40 @@ def test_train_restore():
import shutil
shutil.rmtree(tmpdir)
def test_train_load():
tmpdir = "tests/tmp_load"
poisson_problem = Poisson()
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
trainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu',
default_root_dir=tmpdir)
pinn = PINN(problem=poisson_problem,
model=model,
extra_features=None,
loss=LpLoss())
trainer = Trainer(solver=pinn,
max_epochs=15,
accelerator='cpu',
default_root_dir=tmpdir)
trainer.train()
new_pinn = PINN.load_from_checkpoint(
f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt',
problem = poisson_problem, model=model)
test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
assert new_pinn.forward(test_pts).extract(['u']).shape == pinn.forward(test_pts).extract(['u']).shape
torch.testing.assert_close(new_pinn.forward(test_pts).extract(['u']), pinn.forward(test_pts).extract(['u']))
assert new_pinn.forward(test_pts).extract(
['u']).shape == pinn.forward(test_pts).extract(['u']).shape
torch.testing.assert_close(
new_pinn.forward(test_pts).extract(['u']),
pinn.forward(test_pts).extract(['u']))
import shutil
shutil.rmtree(tmpdir)
# # TODO fix asap. Basically sampling few variables
# # works only if both variables are in a range.
# # if one is fixed and the other not, this will
# # not work. This test also needs to be fixed and
# # not work. This test also needs to be fixed and
# # insert in test problem not in test pinn.
# def test_train_cpu_sampling_few_vars():
# poisson_problem = Poisson()
@@ -158,12 +181,15 @@ def test_train_extra_feats_cpu():
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
pinn = PINN(problem = poisson_problem, model=model_extra_feats, extra_features=extra_feats)
pinn = PINN(problem=poisson_problem,
model=model_extra_feats,
extra_features=extra_feats)
trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
trainer.train()
# TODO, fix GitHub actions to run also on GPU
# def test_train_gpu():
# def test_train_gpu():
# poisson_problem = Poisson()
# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
# n = 10
@@ -171,7 +197,6 @@ def test_train_extra_feats_cpu():
# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'})
# trainer.train()
"""
def test_train_gpu(): #TODO fix ASAP
poisson_problem = Poisson()