fix doc solver

Author: GiovanniCanali

pina/solver/__init__.py

@@ -1,6 +1,4 @@
-"""
-TODO
-"""
+"""Module for the solver classes."""
 
 __all__ = [
     "SolverInterface",

pina/solver/garom.py

@@ -1,4 +1,4 @@
-"""Module for GAROM"""
+"""Module for the GAROM solver."""
 
 import torch
 from torch.nn.modules.loss import _Loss
@@ -10,9 +10,9 @@
 
 class GAROM(MultiSolverInterface):
     """
-    GAROM solver class. This class implements Generative Adversarial
-    Reduced Order Model solver, using user specified ``models`` to solve
-    a specific order reduction``problem``.
+    GAROM solver class. This class implements Generative Adversarial Reduced
+    Order Model solver, using user specified ``models`` to solve a specific
+    order reduction ``problem``.
 
     .. seealso::
 
@@ -39,40 +39,28 @@ def __init__(
         regularizer=False,
     ):
         """
-        :param AbstractProblem problem: The formualation of the problem.
-        :param torch.nn.Module generator: The neural network model to use
-            for the generator.
-        :param torch.nn.Module discriminator: The neural network model to use
+        Initialization of the :class:`GAROM` class.
+        
+        :param AbstractProblem problem: The formulation of the problem.
+        :param torch.nn.Module generator: The generator model.
+        :param torch.nn.Module discriminator: The discriminator model.
+        :param torch.nn.Module loss: The loss function to be minimized.
+            If ``None``, ``PowerLoss(p=1)`` is used. Default is ``None``.
+        :param Optimizer optimizer_generator: The optimizer for the generator.
+            If `None`, the Adam optimizer is used. Default is ``None``.
+        :param Optimizer optimizer_discriminator: The optimizer for the 
+            discriminator. If `None`, the Adam optimizer is used.
+            Default is ``None``.
+        :param Scheduler scheduler_generator: The learning rate scheduler for
+            the generator.
+        :param Scheduler scheduler_discriminator: The learning rate scheduler
             for the discriminator.
-        :param torch.nn.Module loss: The loss function used as minimizer,
-            default ``None``. If ``loss`` is ``None`` the defualt
-            ``PowerLoss(p=1)`` is used, as in the original paper.
-        :param Optimizer optimizer_generator: The neural
-            network optimizer to use for the generator network
-            , default is `torch.optim.Adam`.
-        :param Optimizer optimizer_discriminator: The neural
-            network optimizer to use for the discriminator network
-            , default is `torch.optim.Adam`.
-        :param Scheduler scheduler_generator: Learning
-            rate scheduler for the generator.
-        :param Scheduler scheduler_discriminator: Learning
-            rate scheduler for the discriminator.
-        :param dict scheduler_discriminator_kwargs: LR scheduler constructor
-            keyword args.
-        :param gamma: Ratio of expected loss for generator and discriminator,
-            defaults to 0.3.
-        :type gamma: float
-        :param lambda_k: Learning rate for control theory optimization,
-            defaults to 0.001.
-        :type lambda_k: float
-        :param regularizer: Regularization term in the GAROM loss,
-            defaults to False.
-        :type regularizer: bool
-
-        .. warning::
-            The algorithm works only for data-driven model. Hence in the
-            ``problem`` definition the codition must only contain ``input``
-            (e.g. coefficient parameters, time parameters), and ``target``.
+        :param float gamma: Ratio of expected loss for generator and
+            discriminator. Default is ``0.3``.
+        :param float lambda_k: Learning rate for control theory optimization.
+            Default is ``0.001``.
+        :param bool regularizer: If ``True``, uses a regularization term in the
+            GAROM loss. Default is ``False``.
         """
 
         # set loss
@@ -112,19 +100,15 @@ def __init__(
 
     def forward(self, x, mc_steps=20, variance=False):
         """
-        Forward step for GAROM solver
+        Forward pass implementation.
 
-        :param x: The input tensor.
-        :type x: torch.Tensor
-        :param mc_steps: Number of montecarlo samples to approximate the
-            expected value, defaults to 20.
-        :type mc_steps: int
-        :param variance: Returining also the sample variance of the solution,
-            defaults to False.
-        :type variance: bool
+        :param torch.Tensor x: The input tensor.
+        :param int mc_steps: Number of Montecarlo samples to approximate the
+            expected value. Default is ``20``.
+        :param bool variance: If ``True``, the method returns also the variance
+            of the solution. Default is ``False``.
         :return: The expected value of the generator distribution. If
-            ``variance=True`` also the
-            sample variance is returned.
+            ``variance=True``, the method returns also the variance.
         :rtype: torch.Tensor | tuple(torch.Tensor, torch.Tensor)
         """
 
@@ -142,13 +126,24 @@ def forward(self, x, mc_steps=20, variance=False):
         return mean
 
     def sample(self, x):
-        """TODO"""
+        """
+        Sample from the generator distribution.
+
+        :param torch.Tensor x: The input tensor.
+        :return: The generated sample.
+        :rtype: torch.Tensor
+        """
         # sampling
         return self.generator(x)
 
     def _train_generator(self, parameters, snapshots):
         """
-        Private method to train the generator network.
+        Train the generator model.
+
+        :param torch.Tensor parameters: The input tensor.
+        :param torch.Tensor snapshots: The target tensor.
+        :return: The residual loss and the generator loss.
+        :rtype: tuple(torch.Tensor, torch.Tensor)
         """
         optimizer = self.optimizer_generator
         optimizer.zero_grad()
@@ -170,16 +165,13 @@ def _train_generator(self, parameters, snapshots):
 
     def on_train_batch_end(self, outputs, batch, batch_idx):
         """
-        This method is called at the end of each training batch, and ovverides
-        the PytorchLightining implementation for logging the checkpoints.
+        This method is called at the end of each training batch and overrides
+        the PyTorch Lightning implementation to log checkpoints.
 
-        :param torch.Tensor outputs: The output from the model for the
-            current batch.
-        :param tuple batch: The current batch of data.
+        :param torch.Tensor outputs: The ``model``'s output for the current
+            batch.
+        :param dict batch: The current batch of data.
         :param int batch_idx: The index of the current batch.
-        :return: Whatever is returned by the parent
-            method ``on_train_batch_end``.
-        :rtype: Any
         """
         # increase by one the counter of optimization to save loggers
         (
@@ -190,7 +182,12 @@ def on_train_batch_end(self, outputs, batch, batch_idx):
 
     def _train_discriminator(self, parameters, snapshots):
         """
-        Private method to train the discriminator network.
+        Train the discriminator model.
+
+        :param torch.Tensor parameters: The input tensor.
+        :param torch.Tensor snapshots: The target tensor.
+        :return: The residual loss and the generator loss.
+        :rtype: tuple(torch.Tensor, torch.Tensor)
         """
         optimizer = self.optimizer_discriminator
         optimizer.zero_grad()
@@ -215,8 +212,15 @@ def _train_discriminator(self, parameters, snapshots):
 
     def _update_weights(self, d_loss_real, d_loss_fake):
         """
-        Private method to Update the weights of the generator and discriminator
-        networks.
+        Update the weights of the generator and discriminator models.
+
+        :param torch.Tensor d_loss_real: The discriminator loss computed on
+            dataset samples.
+        :param torch.Tensor d_loss_fake: The discriminator loss computed on
+            generated samples.
+        :return: The difference between the loss computed on the dataset samples
+            and the loss computed on the generated samples.
+        :rtype: torch.Tensor
         """
 
         diff = torch.mean(self.gamma * d_loss_real - d_loss_fake)
@@ -227,11 +231,11 @@ def _update_weights(self, d_loss_real, d_loss_fake):
         return diff
 
     def optimization_cycle(self, batch):
-        """GAROM solver training step.
+        """
+        The optimization cycle for the GAROM solver.
 
-        :param batch: The batch element in the dataloader.
-        :type batch: tuple
-        :return: The sum of the loss functions.
+        :param tuple batch: The batch element in the dataloader.
+        :return: The loss of the optimization cycle.
         :rtype: LabelTensor
         """
         condition_loss = {}
@@ -258,6 +262,13 @@ def optimization_cycle(self, batch):
         return condition_loss
 
     def validation_step(self, batch):
+        """
+        The validation step for the PINN solver.
+
+        :param dict batch: The batch of data to use in the validation step.
+        :return: The loss of the validation step.
+        :rtype: torch.Tensor
+        """
         condition_loss = {}
         for condition_name, points in batch:
             parameters, snapshots = (
@@ -273,6 +284,13 @@ def validation_step(self, batch):
         return loss
 
     def test_step(self, batch):
+        """
+        The test step for the PINN solver.
+
+        :param dict batch: The batch of data to use in the test step.
+        :return: The loss of the test step.
+        :rtype: torch.Tensor
+        """
         condition_loss = {}
         for condition_name, points in batch:
             parameters, snapshots = (
@@ -289,30 +307,60 @@ def test_step(self, batch):
 
     @property
     def generator(self):
-        """TODO"""
+        """
+        The generator model.
+
+        :return: The generator model.
+        :rtype: torch.nn.Module
+        """
         return self.models[0]
 
     @property
     def discriminator(self):
-        """TODO"""
+        """
+        The discriminator model.
+
+        :return: The discriminator model.
+        :rtype: torch.nn.Module
+        """
         return self.models[1]
 
     @property
     def optimizer_generator(self):
-        """TODO"""
+        """
+        The optimizer for the generator.
+
+        :return: The optimizer for the generator.
+        :rtype: Optimizer
+        """
         return self.optimizers[0].instance
 
     @property
     def optimizer_discriminator(self):
-        """TODO"""
+        """
+        The optimizer for the discriminator.
+
+        :return: The optimizer for the discriminator.
+        :rtype: Optimizer
+        """
         return self.optimizers[1].instance
 
     @property
     def scheduler_generator(self):
-        """TODO"""
+        """
+        The scheduler for the generator.
+
+        :return: The scheduler for the generator.
+        :rtype: Scheduler
+        """
         return self.schedulers[0].instance
 
     @property
     def scheduler_discriminator(self):
-        """TODO"""
+        """
+        The scheduler for the discriminator.
+
+        :return: The scheduler for the discriminator.
+        :rtype: Scheduler
+        """
         return self.schedulers[1].instance

pina/solver/physic_informed_solver/pinn_interface.py

@@ -110,8 +110,8 @@ def test_step(self, batch):
     def loss_data(self, input_pts, output_pts):
         """
         Compute the data loss for the PINN solver by evaluating the loss
-        between the network's output and the true solution. This method
-        should only be overridden intentionally.
+        between the network's output and the true solution. This method should
+        not be overridden, if not intentionally.
 
         :param LabelTensor input_pts: The input points to the neural network.
         :param LabelTensor output_pts: The true solution to compare with the