ensemble solver

Author: ndem0

pina/_src/callback/refinement/r3_refinement.py

@@ -6,7 +6,7 @@
 )
 from pina._src.core.label_tensor import LabelTensor
 from pina._src.core.utils import check_consistency
-from pina._src.loss.loss_interface import LossInterface
+from pina._src.loss.loss_interface import DualLossInterface as LossInterface
 
 
 class R3Refinement(RefinementInterface):

pina/_src/callback/refinement/refinement_interface.py

@@ -6,9 +6,7 @@
 from abc import ABCMeta, abstractmethod
 from lightning.pytorch import Callback
 from pina._src.core.utils import check_consistency
-from pina._src.solver.physics_informed_solver.pinn_interface import (
-    PINNInterface,
-)
+from pina._src.solver.pinn import PINN as PINNInterface
 
 
 class RefinementInterface(Callback, metaclass=ABCMeta):

pina/_src/core/trainer.py

@@ -6,12 +6,12 @@
 import lightning
 from pina._src.core.utils import check_consistency, custom_warning_format
 from pina._src.data.data_module import PinaDataModule
-from pina._src.solver.supervised_solver.supervised_solver_interface import (
+from pina._src.solver.solver_interface import (
     SolverInterface,
 )
-from pina._src.solver.physics_informed_solver.pinn_interface import (
-    PINNInterface,
-)
+# from pina._src.solver.physics_informed_solver.pinn_interface import (
+#     PINNInterface,
+# )
 
 # set the warning for compile options
 warnings.formatwarning = custom_warning_format

pina/_src/loss/loss_interface.py

@@ -5,15 +5,15 @@
 import torch
 
 
-class LossInterface(_Loss, metaclass=ABCMeta):
+class DualLossInterface(_Loss, metaclass=ABCMeta):
     """
     Abstract base class for all losses. All classes defining a loss function
     should inherit from this interface.
     """
 
     def __init__(self, reduction="mean"):
         """
-        Initialization of the :class:`LossInterface` class.
+        Initialization of the :class:`DualLossInterface` class.
 
         :param str reduction: The reduction method for the loss.
             Available options: ``none``, ``mean``, ``sum``.

pina/_src/solver/ensemble_simple_solver.py

@@ -0,0 +1,103 @@
+"""Module for the DeepEnsemble simple solver."""
+
+from pina._src.solver.multi_model_simple_solver import MultiModelSimpleSolver
+
+
+class DeepEnsembleSimpleSolver(MultiModelSimpleSolver):
+    r"""
+    Deep Ensemble Simple Solver class. This class implements a Deep Ensemble
+    solver for generic conditions (data, equations, or domain residuals) using
+    user-specified ``models`` to solve a specific ``problem``.
+
+    It is the ensemble counterpart of
+    :class:`~pina.solver.SingleModelSimpleSolver`: each model in the ensemble
+    evaluates every condition independently, and the per-model scalar losses
+    are averaged to produce the final condition loss.
+
+    An ensemble model is constructed by combining multiple models that solve
+    the same type of problem. Mathematically, this creates an implicit
+    distribution :math:`p(\mathbf{u} \mid \mathbf{s})` over the possible
+    outputs :math:`\mathbf{u}`, given the original input :math:`\mathbf{s}`.
+    The models :math:`\mathcal{M}_{i\in (1,\dots,r)}` in
+    the ensemble work collaboratively to capture different
+    aspects of the data or task, with each model contributing a distinct
+    prediction
+    :math:`\mathbf{y}_{i}=\mathcal{M}_i(\mathbf{u} \mid \mathbf{s})`.
+    By aggregating these predictions, the ensemble
+    model can achieve greater robustness and accuracy compared to individual
+    models, leveraging the diversity of the models to reduce overfitting and
+    improve generalization. Furthemore, statistical metrics can
+    be computed, e.g. the ensemble mean and variance:
+
+    .. math::
+        \mathbf{\mu} = \frac{1}{N}\sum_{i=1}^r \mathbf{y}_{i}
+
+    .. math::
+        \mathbf{\sigma^2} = \frac{1}{N}\sum_{i=1}^r
+        (\mathbf{y}_{i} - \mathbf{\mu})^2
+
+    During training the condition loss is minimised by each ensemble model
+    independently and then averaged:
+
+    .. math::
+        \mathcal{L}_{\rm{condition}} = \frac{1}{N_{\rm{ensemble}}}
+        \sum_{i=1}^{N_{\rm{ensemble}}}
+        \mathcal{L}_i(\mathcal{M}_i, \mathbf{s})
+
+    where :math:`\mathcal{L}` is a specific loss function, typically the MSE:
+
+    .. math::
+        \mathcal{L}(v) = \| v \|^2_2.
+
+    .. seealso::
+
+        **Original reference**: Lakshminarayanan, B., Pritzel, A., & Blundell,
+        C. (2017). *Simple and scalable predictive uncertainty estimation
+        using deep ensembles*. Advances in neural information
+        processing systems, 30.
+        DOI: `arXiv:1612.01474 <https://arxiv.org/abs/1612.01474>`_.
+    """
+
+    def __init__(
+        self,
+        problem,
+        models,
+        optimizers=None,
+        schedulers=None,
+        weighting=None,
+        loss=None,
+        use_lt=True,
+        ensemble_dim=0,
+    ):
+        """
+        Initialization of the :class:`DeepEnsembleSimpleSolver` class.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param list[torch.nn.Module] models: The neural network models to be
+            used. Must be a list or tuple with at least two models.
+        :param list[Optimizer] optimizers: The optimizers to be used.
+            If ``None``, the :class:`torch.optim.Adam` optimizer is used for
+            each model. Default is ``None``.
+        :param list[Scheduler] schedulers: The learning rate schedulers.
+            If ``None``, :class:`torch.optim.lr_scheduler.ConstantLR` is used
+            for each model. 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 element-wise loss module.
+            If ``None``, :class:`torch.nn.MSELoss` is used. Default is
+            ``None``.
+        :param bool use_lt: If ``True``, the solver uses LabelTensors as
+            input. Default is ``True``.
+        :param int ensemble_dim: The dimension along which the per-model
+            outputs are stacked in :meth:`forward`. Default is ``0``.
+        """
+        super().__init__(
+            problem=problem,
+            models=models,
+            optimizers=optimizers,
+            schedulers=schedulers,
+            weighting=weighting,
+            loss=loss,
+            use_lt=use_lt,
+            ensemble_dim=ensemble_dim,
+        )