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>

pina/solvers/pinn.py

@@ -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)