Inverse problem implementation (#177)
* inverse problem implementation * add tutorial7 for inverse Poisson problem * fix doc in equation, equation_interface, system_equation --------- Co-authored-by: Dario Coscia <dariocoscia@dhcp-015.eduroam.sissa.it>
This commit is contained in:
committed by
Nicola Demo
parent
a9f14ac323
commit
0b7a307cf1
@@ -37,7 +37,7 @@ class Condition:
|
||||
>>> example_input_pts = LabelTensor(
|
||||
>>> torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
|
||||
>>> example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ['a', 'b'])
|
||||
>>>
|
||||
>>>
|
||||
>>> Condition(
|
||||
>>> input_points=example_input_pts,
|
||||
>>> output_points=example_output_pts)
|
||||
|
||||
@@ -9,4 +9,4 @@ __all__ = [
|
||||
|
||||
from .equation import Equation
|
||||
from .equation_factory import FixedFlux, FixedGradient, Laplace, FixedValue
|
||||
from .system_equation import SystemEquation
|
||||
from .system_equation import SystemEquation
|
||||
|
||||
@@ -8,7 +8,7 @@ class Equation(EquationInterface):
|
||||
"""
|
||||
Equation class for specifing any equation in PINA.
|
||||
Each ``equation`` passed to a ``Condition`` object
|
||||
must be an ``Equation`` or ``SystemEquation``.
|
||||
must be an ``Equation`` or ``SystemEquation``.
|
||||
|
||||
:param equation: A ``torch`` callable equation to
|
||||
evaluate the residual.
|
||||
@@ -20,14 +20,26 @@ class Equation(EquationInterface):
|
||||
f'{equation}')
|
||||
self.__equation = equation
|
||||
|
||||
def residual(self, input_, output_):
|
||||
def residual(self, input_, output_, params_ = None):
|
||||
"""
|
||||
Residual computation of the equation.
|
||||
|
||||
:param LabelTensor input_: Input points to evaluate the equation.
|
||||
:param LabelTensor output_: Output vectors given my a model (e.g,
|
||||
:param LabelTensor output_: Output vectors given by a model (e.g,
|
||||
a ``FeedForward`` model).
|
||||
:param dict params_: Dictionary of parameters related to the inverse
|
||||
problem (if any).
|
||||
If the equation is not related to an ``InverseProblem``, the
|
||||
parameters are initialized to ``None`` and the residual is
|
||||
computed as ``equation(input_, output_)``.
|
||||
Otherwise, the parameters are automatically initialized in the
|
||||
ranges specified by the user.
|
||||
|
||||
:return: The residual evaluation of the specified equation.
|
||||
:rtype: LabelTensor
|
||||
"""
|
||||
return self.__equation(input_, output_)
|
||||
if params_ is None:
|
||||
result = self.__equation(input_, output_)
|
||||
else:
|
||||
result = self.__equation(input_, output_, params_)
|
||||
return result
|
||||
|
||||
@@ -11,3 +11,17 @@ class EquationInterface(metaclass=ABCMeta):
|
||||
the output variables, the condition(s), and the domain(s) where the
|
||||
conditions are applied.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def residual(self, input_, output_, params_):
|
||||
"""
|
||||
Residual computation of the equation.
|
||||
|
||||
:param LabelTensor input_: Input points to evaluate the equation.
|
||||
:param LabelTensor output_: Output vectors given by my model (e.g., a ``FeedForward`` model).
|
||||
:param dict params_: Dictionary of unknown parameters, eventually
|
||||
related to an ``InverseProblem``.
|
||||
:return: The residual evaluation of the specified equation.
|
||||
:rtype: LabelTensor
|
||||
"""
|
||||
pass
|
||||
|
||||
@@ -11,14 +11,14 @@ class SystemEquation(Equation):
|
||||
System of Equation class for specifing any system
|
||||
of equations in PINA.
|
||||
Each ``equation`` passed to a ``Condition`` object
|
||||
must be an ``Equation`` or ``SystemEquation``.
|
||||
A ``SystemEquation`` is specified by a list of
|
||||
must be an ``Equation`` or ``SystemEquation``.
|
||||
A ``SystemEquation`` is specified by a list of
|
||||
equations.
|
||||
|
||||
:param Callable equation: A ``torch`` callable equation to
|
||||
evaluate the residual
|
||||
:param str reduction: Specifies the reduction to apply to the output:
|
||||
``none`` | ``mean`` | ``sum`` | ``callable``. ``none``: no reduction
|
||||
``none`` | ``mean`` | ``sum`` | ``callable``. ``none``: no reduction
|
||||
will be applied, ``mean``: the sum of the output will be divided
|
||||
by the number of elements in the output, ``sum``: the output will
|
||||
be summed. ``callable`` a callable function to perform reduction,
|
||||
@@ -43,19 +43,28 @@ class SystemEquation(Equation):
|
||||
raise NotImplementedError(
|
||||
'Only mean and sum reductions implemented.')
|
||||
|
||||
def residual(self, input_, output_):
|
||||
def residual(self, input_, output_, params_=None):
|
||||
"""
|
||||
Residual computation of the equation.
|
||||
Residual computation for the equations of the system.
|
||||
|
||||
:param LabelTensor input_: Input points to evaluate the equation.
|
||||
:param LabelTensor output_: Output vectors given my a model (e.g,
|
||||
:param LabelTensor input_: Input points to evaluate the system of
|
||||
equations.
|
||||
:param LabelTensor output_: Output vectors given by a model (e.g,
|
||||
a ``FeedForward`` model).
|
||||
:return: The residual evaluation of the specified equation,
|
||||
:param dict params_: Dictionary of parameters related to the inverse
|
||||
problem (if any).
|
||||
If the equation is not related to an ``InverseProblem``, the
|
||||
parameters are initialized to ``None`` and the residual is
|
||||
computed as ``equation(input_, output_)``.
|
||||
Otherwise, the parameters are automatically initialized in the
|
||||
ranges specified by the user.
|
||||
|
||||
:return: The residual evaluation of the specified system of equations,
|
||||
aggregated by the ``reduction`` defined in the ``__init__``.
|
||||
:rtype: LabelTensor
|
||||
"""
|
||||
residual = torch.hstack(
|
||||
[equation.residual(input_, output_) for equation in self.equations])
|
||||
[equation.residual(input_, output_, params_) for equation in self.equations])
|
||||
|
||||
if self.reduction == 'none':
|
||||
return residual
|
||||
|
||||
@@ -205,6 +205,7 @@ class Plotter:
|
||||
plt.savefig(filename)
|
||||
else:
|
||||
plt.show()
|
||||
plt.close()
|
||||
|
||||
def plot_loss(self,
|
||||
trainer,
|
||||
|
||||
@@ -3,9 +3,11 @@ __all__ = [
|
||||
'SpatialProblem',
|
||||
'TimeDependentProblem',
|
||||
'ParametricProblem',
|
||||
'InverseProblem',
|
||||
]
|
||||
|
||||
from .abstract_problem import AbstractProblem
|
||||
from .spatial_problem import SpatialProblem
|
||||
from .timedep_problem import TimeDependentProblem
|
||||
from .parametric_problem import ParametricProblem
|
||||
from .inverse_problem import InverseProblem
|
||||
|
||||
@@ -109,6 +109,14 @@ class AbstractProblem(metaclass=ABCMeta):
|
||||
samples = condition.input_points
|
||||
self.input_pts[condition_name] = samples
|
||||
self._have_sampled_points[condition_name] = True
|
||||
if hasattr(self, 'unknown_parameter_domain'):
|
||||
# initialize the unknown parameters of the inverse problem given
|
||||
# the domain the user gives
|
||||
self.unknown_parameters = {}
|
||||
for i, var in enumerate(self.unknown_variables):
|
||||
range_var = self.unknown_parameter_domain.range_[var]
|
||||
tensor_var = torch.rand(1, requires_grad=True) * range_var[1] + range_var[0]
|
||||
self.unknown_parameters[var] = torch.nn.Parameter(tensor_var)
|
||||
|
||||
def discretise_domain(self,
|
||||
n,
|
||||
@@ -203,6 +211,7 @@ class AbstractProblem(metaclass=ABCMeta):
|
||||
self.input_variables):
|
||||
self._have_sampled_points[location] = True
|
||||
|
||||
|
||||
def add_points(self, new_points):
|
||||
"""
|
||||
Adding points to the already sampled points.
|
||||
@@ -237,7 +246,7 @@ class AbstractProblem(metaclass=ABCMeta):
|
||||
@property
|
||||
def have_sampled_points(self):
|
||||
"""
|
||||
Check if all points for
|
||||
Check if all points for
|
||||
``Location`` are sampled.
|
||||
"""
|
||||
return all(self._have_sampled_points.values())
|
||||
@@ -245,7 +254,7 @@ class AbstractProblem(metaclass=ABCMeta):
|
||||
@property
|
||||
def not_sampled_points(self):
|
||||
"""
|
||||
Check which points are
|
||||
Check which points are
|
||||
not sampled.
|
||||
"""
|
||||
# variables which are not sampled
|
||||
@@ -257,3 +266,4 @@ class AbstractProblem(metaclass=ABCMeta):
|
||||
if not is_sample:
|
||||
not_sampled.append(condition_name)
|
||||
return not_sampled
|
||||
|
||||
|
||||
71
pina/problem/inverse_problem.py
Normal file
71
pina/problem/inverse_problem.py
Normal file
@@ -0,0 +1,71 @@
|
||||
"""Module for the ParametricProblem class"""
|
||||
from abc import abstractmethod
|
||||
|
||||
from .abstract_problem import AbstractProblem
|
||||
|
||||
|
||||
class InverseProblem(AbstractProblem):
|
||||
"""
|
||||
The class for the definition of inverse problems, i.e., problems
|
||||
with unknown parameters that have to be learned during the training process
|
||||
from given data.
|
||||
|
||||
Here's an example of a spatial inverse ODE problem, i.e., a spatial
|
||||
ODE problem with an unknown parameter `alpha` as coefficient of the
|
||||
derivative term.
|
||||
|
||||
:Example:
|
||||
>>> from pina.problem import SpatialProblem, InverseProblem
|
||||
>>> from pina.operators import grad
|
||||
>>> from pina.equation import ParametricEquation, FixedValue
|
||||
>>> from pina import Condition
|
||||
>>> from pina.geometry import CartesianDomain
|
||||
>>> import torch
|
||||
>>>
|
||||
>>> class InverseODE(SpatialProblem, InverseProblem):
|
||||
>>>
|
||||
>>> output_variables = ['u']
|
||||
>>> spatial_domain = CartesianDomain({'x': [0, 1]})
|
||||
>>> unknown_parameter_domain = CartesianDomain({'alpha': [1, 10]})
|
||||
>>>
|
||||
>>> def ode_equation(input_, output_, params_):
|
||||
>>> u_x = grad(output_, input_, components=['u'], d=['x'])
|
||||
>>> u = output_.extract(['u'])
|
||||
>>> return params_.extract(['alpha']) * u_x - u
|
||||
>>>
|
||||
>>> def solution_data(input_, output_):
|
||||
>>> x = input_.extract(['x'])
|
||||
>>> solution = torch.exp(x)
|
||||
>>> return output_ - solution
|
||||
>>>
|
||||
>>> conditions = {
|
||||
>>> 'x0': Condition(CartesianDomain({'x': 0}), FixedValue(1.0)),
|
||||
>>> 'D': Condition(CartesianDomain({'x': [0, 1]}), ParametricEquation(ode_equation)),
|
||||
>>> 'data': Condition(CartesianDomain({'x': [0, 1]}), Equation(solution_data))
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def unknown_parameter_domain(self):
|
||||
"""
|
||||
The parameters' domain of the problem.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
def unknown_variables(self):
|
||||
"""
|
||||
The parameters of the problem.
|
||||
"""
|
||||
return self.unknown_parameter_domain.variables
|
||||
|
||||
@property
|
||||
def unknown_parameters(self):
|
||||
"""
|
||||
The parameters of the problem.
|
||||
"""
|
||||
return self.__unknown_parameters
|
||||
|
||||
@unknown_parameters.setter
|
||||
def unknown_parameters(self, value):
|
||||
self.__unknown_parameters = value
|
||||
|
||||
@@ -14,7 +14,7 @@ class SpatialProblem(AbstractProblem):
|
||||
:Example:
|
||||
>>> from pina.problem import SpatialProblem
|
||||
>>> from pina.operators import grad
|
||||
>>> from pina.equations import Equation, FixedValue
|
||||
>>> from pina.equation import Equation, FixedValue
|
||||
>>> from pina import Condition
|
||||
>>> from pina.geometry import CartesianDomain
|
||||
>>> import torch
|
||||
@@ -33,7 +33,6 @@ class SpatialProblem(AbstractProblem):
|
||||
>>> conditions = {
|
||||
>>> 'x0': Condition(CartesianDomain({'x': 0, 'alpha':[1, 10]}), FixedValue(1.)),
|
||||
>>> 'D': Condition(CartesianDomain({'x': [0, 1], 'alpha':[1, 10]}), Equation(ode_equation))}
|
||||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -14,7 +14,7 @@ class TimeDependentProblem(AbstractProblem):
|
||||
:Example:
|
||||
>>> from pina.problem import SpatialProblem, TimeDependentProblem
|
||||
>>> from pina.operators import grad, laplacian
|
||||
>>> from pina.equations import Equation, FixedValue
|
||||
>>> from pina.equation import Equation, FixedValue
|
||||
>>> from pina import Condition
|
||||
>>> from pina.geometry import CartesianDomain
|
||||
>>> import torch
|
||||
@@ -43,7 +43,6 @@ class TimeDependentProblem(AbstractProblem):
|
||||
>>> 'gamma1': Condition(CartesianDomain({'x':0, 't':[0, 1]}), FixedValue(0.)),
|
||||
>>> 'gamma2': Condition(CartesianDomain({'x':3, 't':[0, 1]}), FixedValue(0.)),
|
||||
>>> 'D': Condition(CartesianDomain({'x': [0, 3], 't':[0, 1]}), Equation(wave_equation))}
|
||||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -11,6 +11,7 @@ from .solver import SolverInterface
|
||||
from ..label_tensor import LabelTensor
|
||||
from ..utils import check_consistency
|
||||
from ..loss import LossInterface
|
||||
from ..problem import InverseProblem
|
||||
from torch.nn.modules.loss import _Loss
|
||||
|
||||
torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732
|
||||
@@ -18,14 +19,14 @@ torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732
|
||||
|
||||
class PINN(SolverInterface):
|
||||
"""
|
||||
PINN solver class. This class implements Physics Informed Neural
|
||||
PINN solver class. This class implements Physics Informed Neural
|
||||
Network solvers, using a user specified ``model`` to solve a specific
|
||||
``problem``.
|
||||
``problem``. It can be used for solving both forward and inverse problems.
|
||||
|
||||
.. seealso::
|
||||
|
||||
**Original reference**: Karniadakis, G. E., Kevrekidis, I. G., Lu, L.,
|
||||
Perdikaris, P., Wang, S., & Yang, L. (2021).
|
||||
**Original reference**: Karniadakis, G. E., Kevrekidis, I. G., Lu, L.,
|
||||
Perdikaris, P., Wang, S., & Yang, L. (2021).
|
||||
Physics-informed machine learning. Nature Reviews Physics, 3(6), 422-440.
|
||||
<https://doi.org/10.1038/s42254-021-00314-5>`_.
|
||||
"""
|
||||
@@ -45,7 +46,7 @@ class PINN(SolverInterface):
|
||||
},
|
||||
):
|
||||
'''
|
||||
:param AbstractProblem problem: The formualation of the problem.
|
||||
:param AbstractProblem problem: The formulation of the problem.
|
||||
:param torch.nn.Module model: The neural network model to use.
|
||||
:param torch.nn.Module loss: The loss function used as minimizer,
|
||||
default :class:`torch.nn.MSELoss`.
|
||||
@@ -74,12 +75,18 @@ class PINN(SolverInterface):
|
||||
self._loss = loss
|
||||
self._neural_net = self.models[0]
|
||||
|
||||
# inverse problem handling
|
||||
if isinstance(self.problem, InverseProblem):
|
||||
self._params = self.problem.unknown_parameters
|
||||
else:
|
||||
self._params = None
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass implementation for the PINN
|
||||
solver.
|
||||
|
||||
:param torch.Tensor x: Input tensor.
|
||||
:param torch.Tensor x: Input tensor.
|
||||
:return: PINN solution.
|
||||
:rtype: torch.Tensor
|
||||
"""
|
||||
@@ -93,17 +100,30 @@ class PINN(SolverInterface):
|
||||
:return: The optimizers and the schedulers
|
||||
:rtype: tuple(list, list)
|
||||
"""
|
||||
# if the problem is an InverseProblem, add the unknown parameters
|
||||
# to the parameters that the optimizer needs to optimize
|
||||
if isinstance(self.problem, InverseProblem):
|
||||
self.optimizers[0].add_param_group(
|
||||
{'params': [self._params[var] for var in self.problem.unknown_variables]}
|
||||
)
|
||||
return self.optimizers, [self.scheduler]
|
||||
|
||||
|
||||
def _clamp_inverse_problem_params(self):
|
||||
for v in self._params:
|
||||
self._params[v].data.clamp_(
|
||||
self.problem.unknown_parameter_domain.range_[v][0],
|
||||
self.problem.unknown_parameter_domain.range_[v][1])
|
||||
|
||||
def _loss_data(self, input, output):
|
||||
return self.loss(self.forward(input), output)
|
||||
|
||||
|
||||
def _loss_phys(self, samples, equation):
|
||||
residual = equation.residual(samples, self.forward(samples))
|
||||
try:
|
||||
residual = equation.residual(samples, self.forward(samples))
|
||||
except TypeError: # this occurs when the function has three inputs, i.e. inverse problem
|
||||
residual = equation.residual(samples, self.forward(samples), self._params)
|
||||
return self.loss(torch.zeros_like(residual, requires_grad=True), residual)
|
||||
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
"""
|
||||
PINN solver training step.
|
||||
@@ -137,15 +157,20 @@ class PINN(SolverInterface):
|
||||
else:
|
||||
raise ValueError("Batch size not supported")
|
||||
|
||||
# TODO for users this us hard to remebeber when creating a new solver, to fix in a smarter way
|
||||
# TODO for users this us hard to remember when creating a new solver, to fix in a smarter way
|
||||
loss = loss.as_subclass(torch.Tensor)
|
||||
|
||||
# add condition losses and accumulate logging for each epoch
|
||||
# # add condition losses and accumulate logging for each epoch
|
||||
condition_losses.append(loss * condition.data_weight)
|
||||
self.log(condition_name + '_loss', float(loss),
|
||||
prog_bar=True, logger=True, on_epoch=True, on_step=False)
|
||||
|
||||
# add to tot loss and accumulate logging for each epoch
|
||||
# clamp unknown parameters of the InverseProblem to their domain ranges (if needed)
|
||||
if isinstance(self.problem, InverseProblem):
|
||||
self._clamp_inverse_problem_params()
|
||||
|
||||
# TODO Fix the bug, tot_loss is a label tensor without labels
|
||||
# we need to pass it as a torch tensor to make everything work
|
||||
total_loss = sum(condition_losses)
|
||||
self.log('mean_loss', float(total_loss / len(condition_losses)),
|
||||
prog_bar=True, logger=True, on_epoch=True, on_step=False)
|
||||
|
||||
Reference in New Issue
Block a user