fix rendering part 2

Author: GiovanniCanali

pina/problem/zoo/advection.py

@@ -63,6 +63,9 @@ class AdvectionProblem(SpatialProblem, TimeDependentProblem):
         training physics-informed neural networks*.
         arXiv preprint arXiv:2308.08468 (2023).
         DOI: `arXiv:2308.08468  <https://arxiv.org/abs/2308.08468>`_.
+
+    :Example:
+        >>> problem = AdvectionProblem(c=1.0)
     """
 
     output_variables = ["u"]

pina/problem/zoo/allen_cahn.py

@@ -49,6 +49,9 @@ class AllenCahnProblem(TimeDependentProblem, SpatialProblem):
         Computer Methods in Applied Mechanics and Engineering 421 (2024): 116805
         DOI: `10.1016/
         j.cma.2024.116805 <https://doi.org/10.1016/j.cma.2024.116805>`_.
+    
+    :Example:
+        >>> problem = AllenCahnProblem()
     """
 
     output_variables = ["u"]

pina/problem/zoo/diffusion_reaction.py

@@ -59,6 +59,9 @@ class DiffusionReactionProblem(TimeDependentProblem, SpatialProblem):
         **Original reference**: Si, Chenhao, et al. *Complex Physics-Informed
         Neural Network.* arXiv preprint arXiv:2502.04917 (2025).
         DOI: `arXiv:2502.04917 <https://arxiv.org/abs/2502.04917>`_.
+
+    :Example:
+        >>> problem = DiffusionReactionProblem()
     """
 
     output_variables = ["u"]

pina/problem/zoo/helmholtz.py

@@ -53,6 +53,9 @@ class HelmholtzProblem(SpatialProblem):
         **Original reference**: Si, Chenhao, et al. *Complex Physics-Informed
         Neural Network.* arXiv preprint arXiv:2502.04917 (2025).
         DOI: `arXiv:2502.04917 <https://arxiv.org/abs/2502.04917>`_.
+    
+    :Example:
+        >>> problem = HelmholtzProblem()
     """
 
     output_variables = ["u"]

pina/problem/zoo/inverse_poisson_2d_square.py

@@ -50,6 +50,9 @@ class InversePoisson2DSquareProblem(SpatialProblem, InverseProblem):
     Implementation of the inverse 2-dimensional Poisson problem in the square
     domain :math:`[0, 1] \times [0, 1]`,
     with unknown parameter domain :math:`[-1, 1] \times [-1, 1]`.
+
+    :Example:
+        >>> problem = InversePoisson2DSquareProblem()
     """
 
     output_variables = ["u"]

pina/problem/zoo/poisson_2d_square.py

@@ -30,6 +30,9 @@ class Poisson2DSquareProblem(SpatialProblem):
     r"""
     Implementation of the 2-dimensional Poisson problem in the square domain
     :math:`[0, 1] \times [0, 1]`.
+
+    :Example:
+        >>> problem = Poisson2DSquareProblem()
     """
 
     output_variables = ["u"]

pina/solver/garom.py

@@ -45,16 +45,22 @@ def __init__(
         :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``.
+            If ``None``, :class:`~pina.loss.power_loss.PowerLoss` with ``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.
+            If `None`, the :class:`torch.optim.Adam` optimizer is used.
             Default is ``None``.
+        :param Optimizer optimizer_discriminator: The optimizer for the 
+            discriminator. If `None`, the :class:`torch.optim.Adam` optimizer is
+            used. Default is ``None``.
         :param Scheduler scheduler_generator: The learning rate scheduler for
             the generator.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
         :param Scheduler scheduler_discriminator: The learning rate scheduler
             for the discriminator.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
         :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.
@@ -109,7 +115,7 @@ def forward(self, x, mc_steps=20, variance=False):
             of the solution. Default is ``False``.
         :return: The expected value of the generator distribution. If
             ``variance=True``, the method returns also the variance.
-        :rtype: torch.Tensor | tuple(torch.Tensor, torch.Tensor)
+        :rtype: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
         """
 
         # sampling
@@ -143,7 +149,7 @@ def _train_generator(self, parameters, snapshots):
         :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)
+        :rtype: tuple[torch.Tensor, torch.Tensor]
         """
         optimizer = self.optimizer_generator
         optimizer.zero_grad()
@@ -170,7 +176,8 @@ def on_train_batch_end(self, outputs, batch, batch_idx):
 
         :param torch.Tensor outputs: The ``model``'s output for the current
             batch.
-        :param dict batch: The current batch of data.
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
         :param int batch_idx: The index of the current batch.
         """
         # increase by one the counter of optimization to save loggers
@@ -187,7 +194,7 @@ def _train_discriminator(self, parameters, snapshots):
         :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)
+        :rtype: tuple[torch.Tensor, torch.Tensor]
         """
         optimizer = self.optimizer_discriminator
         optimizer.zero_grad()
@@ -234,9 +241,12 @@ def optimization_cycle(self, batch):
         """
         The optimization cycle for the GAROM solver.
 
-        :param tuple batch: The batch element in the dataloader.
-        :return: The loss of the optimization cycle.
-        :rtype: LabelTensor
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :return: The losses computed for all conditions in the batch, casted
+            to a subclass of :class:`torch.Tensor`. It should return a dict
+            containing the condition name and the associated scalar loss.
+        :rtype: dict
         """
         condition_loss = {}
         for condition_name, points in batch:
@@ -265,7 +275,8 @@ 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.
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
         :return: The loss of the validation step.
         :rtype: torch.Tensor
         """
@@ -287,7 +298,8 @@ 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.
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
         :return: The loss of the test step.
         :rtype: torch.Tensor
         """

pina/solver/physic_informed_solver/causal_pinn.py

@@ -53,14 +53,16 @@ class CausalPINN(PINN):
     .. seealso::
 
         **Original reference**: Wang, Sifan, Shyam Sankaran, and Paris
-        Perdikaris. "Respecting causality for training physics-informed
-        neural networks." Computer Methods in Applied Mechanics
-        and Engineering 421 (2024): 116813.
-        DOI `10.1016 <https://doi.org/10.1016/j.cma.2024.116813>`_.
+        Perdikaris. 
+        *Respecting causality for training physics-informed
+        neural networks.*
+        Computer Methods in Applied Mechanics and Engineering 421 (2024):116813.
+        DOI: `10.1016 <https://doi.org/10.1016/j.cma.2024.116813>`_.
 
     .. note::
         This class is only compatible with problems that inherit from  the
-        :class:`~pina.problem.TimeDependentProblem` class.
+        :class:`~pina.problem.time_dependent_problem.TimeDependentProblem`
+        class.
     """
 
     def __init__(
@@ -77,17 +79,19 @@ def __init__(
         Initialization of the :class:`CausalPINN` class.
 
         :param AbstractProblem problem: The problem to be solved. It must
-            inherit from at least :class:`~pina.problem.TimeDependentProblem`.
+            inherit from at least
+            :class:`~pina.problem.time_dependent_problem.TimeDependentProblem`.
         :param torch.nn.Module model: The neural network model to be used.
-        :param torch.optim.Optimizer optimizer: The optimizer to be used
-            If `None`, the Adam optimizer is used. Default is ``None``.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
-            If `None`, the constant learning rate scheduler is used.
+        :param Optimizer optimizer: The optimizer to be used.
+            If `None`, the :class:`torch.optim.Adam` optimizer is used.
             Default is ``None``.
+        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
         :param WeightingInterface weighting: The weighting schema to be used.
             If `None`, no weighting schema is used. Default is ``None``.
         :param torch.nn.Module loss: The loss function to be minimized.
-            If `None`, the Mean Squared Error (MSE) loss is used.
+            If `None`, the :class:`torch.nn.MSELoss` loss is used.
             Default is `None`.
         :param float eps: The exponential decay parameter. Default is ``100``.
         :raises ValueError: If the problem is not a TimeDependentProblem.

pina/solver/physic_informed_solver/competitive_pinn.py

@@ -46,8 +46,8 @@ class CompetitivePINN(PINNInterface, MultiSolverInterface):
     .. seealso::
 
         **Original reference**: Zeng, Qi, et al.
-        "Competitive physics informed networks." International Conference on
-        Learning Representations, ICLR 2022
+        *Competitive physics informed networks.*
+        International Conference on Learning Representations, ICLR 2022
         `OpenReview Preprint <https://openreview.net/forum?id=z9SIj-IM7tn>`_.
     """
 
@@ -72,21 +72,23 @@ def __init__(
             If `None`, the discriminator is a deepcopy of the ``model``.
             Default is ``None``.
         :param torch.optim.Optimizer optimizer_model: The optimizer of the
-            ``model``. If `None`, the Adam optimizer is used.
-            Default is ``None``.
+            ``model``. If `None`, the :class:`torch.optim.Adam` optimizer is
+            used. Default is ``None``.
         :param torch.optim.Optimizer optimizer_discriminator: The optimizer of
-            the ``discriminator``. If `None`, the Adam optimizer is used.
-            Default is ``None``.
-        :param torch.optim.LRScheduler scheduler_model: Learning rate scheduler
-            for the ``model``. If `None`, the constant learning rate scheduler
-            is used. Default is ``None``.
-        :param torch.optim.LRScheduler scheduler_discriminator: Learning rate
-            scheduler for the ``discriminator``. If `None`, the constant
-            learning rate scheduler is used. Default is ``None``.
+            the ``discriminator``. If `None`, the :class:`torch.optim.Adam`
+            optimizer is used. Default is ``None``.
+        :param Scheduler scheduler_model: Learning rate scheduler for the
+            ``model``.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
+        :param Scheduler scheduler_discriminator: Learning rate scheduler for
+            the ``discriminator``.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
         :param WeightingInterface weighting: The weighting schema to be used.
             If `None`, no weighting schema is used. Default is ``None``.
         :param torch.nn.Module loss: The loss function to be minimized.
-            If `None`, the Mean Squared Error (MSE) loss is used.
+            If `None`, the :class:`torch.nn.MSELoss` loss is used.
             Default is `None`.
         """
         if discriminator is None:
@@ -118,7 +120,8 @@ def training_step(self, batch):
         """
         Solver training step, overridden to perform manual optimization.
 
-        :param dict batch: The batch element in the dataloader.
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
         :return: The aggregated loss.
         :rtype: LabelTensor
         """
@@ -163,7 +166,7 @@ def configure_optimizers(self):
         Optimizer configuration.
 
         :return: The optimizers and the schedulers
-        :rtype: tuple(list, list)
+        :rtype: tuple[list[Optimizer], list[Scheduler]]
         """
         # If the problem is an InverseProblem, add the unknown parameters
         # to the parameters to be optimized
@@ -198,7 +201,8 @@ def on_train_batch_end(self, outputs, batch, batch_idx):
 
         :param torch.Tensor outputs: The ``model``'s output for the current
             batch.
-        :param dict batch: The current batch of data.
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
         :param int batch_idx: The index of the current batch.
         """
         # increase by one the counter of optimization to save loggers
@@ -234,7 +238,7 @@ def optimizer_model(self):
         The optimizer associated to the model.
 
         :return: The optimizer for the model.
-        :rtype: torch.optim.Optimizer
+        :rtype: Optimizer
         """
         return self.optimizers[0]
 
@@ -244,7 +248,7 @@ def optimizer_discriminator(self):
         The optimizer associated to the discriminator.
 
         :return: The optimizer for the discriminator.
-        :rtype: torch.optim.Optimizer
+        :rtype: Optimizer
         """
         return self.optimizers[1]
 
@@ -254,7 +258,7 @@ def scheduler_model(self):
         The scheduler associated to the model.
 
         :return: The scheduler for the model.
-        :rtype: torch.optim.lr_scheduler._LRScheduler
+        :rtype: Scheduler
         """
         return self.schedulers[0]
 
@@ -264,6 +268,6 @@ def scheduler_discriminator(self):
         The scheduler associated to the discriminator.
 
         :return: The scheduler for the discriminator.
-        :rtype: torch.optim.lr_scheduler._LRScheduler
+        :rtype: Scheduler
         """
         return self.schedulers[1]

pina/solver/physic_informed_solver/gradient_pinn.py

@@ -46,15 +46,15 @@ class GradientPINN(PINN):
 
     .. seealso::
 
-        **Original reference**: Yu, Jeremy, et al. "Gradient-enhanced
-        physics-informed neural networks for forward and inverse
-        PDE problems." Computer Methods in Applied Mechanics
-        and Engineering 393 (2022): 114823.
+        **Original reference**: Yu, Jeremy, et al.
+        *Gradient-enhanced physics-informed neural networks for forward and
+        inverse PDE problems.*
+        Computer Methods in Applied Mechanics and Engineering 393 (2022):114823.
         DOI: `10.1016 <https://doi.org/10.1016/j.cma.2022.114823>`_.
 
     .. note::
         This class is only compatible with problems that inherit from  the
-        :class:`~pina.problem.SpatialProblem` class. 
+        :class:`~pina.problem.spatial_problem.SpatialProblem` class. 
     """
 
     def __init__(
@@ -70,18 +70,20 @@ def __init__(
         Initialization of the :class:`GradientPINN` class.
 
         :param AbstractProblem problem: The problem to be solved.
-            It must inherit from at least :class:`~pina.problem.SpatialProblem`
-            to compute the gradient of the loss.
+            It must inherit from at least
+            :class:`~pina.problem.spatial_problem.SpatialProblem` to compute the
+            gradient of the loss.
         :param torch.nn.Module model: The neural network model to be used.
-        :param torch.optim.Optimizer optimizer: The optimizer to be used.
-            If `None`, the Adam optimizer is used. Default is ``None``.
-        :param torch.optim.LRScheduler scheduler: Learning rate scheduler.
-            If `None`, the constant learning rate scheduler is used.
+        :param Optimizer optimizer: The optimizer to be used.
+            If `None`, the :class:`torch.optim.Adam` optimizer is used.
             Default is ``None``.
+        :param Scheduler scheduler: Learning rate scheduler.
+            If `None`, the :class:`torch.optim.lr_scheduler.ConstantLR`
+            scheduler is used. Default is ``None``.
         :param WeightingInterface weighting: The weighting schema to be used.
             If `None`, no weighting schema is used. Default is ``None``.
         :param torch.nn.Module loss: The loss function to be minimized.
-            If `None`, the Mean Squared Error (MSE) loss is used.
+            If `None`, the :class:`torch.nn.MSELoss` loss is used.
             Default is `None`.
         :raises ValueError: If the problem is not a SpatialProblem.
         """