first new generation solver

Author: ndem0

pina/_src/condition/equation_condition_base.py

@@ -15,7 +15,7 @@ class EquationConditionBase(ConditionBase):
     :class:`~pina.condition.InputEquationCondition`.
     """
 
-    def evaluate(self, batch, solver):
+    def evaluate(self, batch, solver, loss):
         """
         Evaluate the equation residual on the given batch using the solver.
 
@@ -30,17 +30,21 @@ def evaluate(self, batch, solver):
         :param solver: The solver containing the model and any additional
             parameters (e.g., unknown parameters for inverse problems).
         :type solver: ~pina.solver.solver.SolverInterface
-        :return: The non-aggregated residual tensor.
+        :param loss: The non-aggregating loss function to apply to the
+            computed residual against zero.
+        :type loss: torch.nn.Module
+        :return: The non-aggregated loss tensor.
         :rtype: ~pina.label_tensor.LabelTensor
 
         :Example:
 
             >>> residuals = condition.evaluate(
-            ...     {"input": input_samples}, solver
+            ...     {"input": input_samples}, solver, loss
             ... )
             >>> # residuals is a non-reduced tensor of shape (n_samples, ...)
         """
         samples = batch["input"]
-        return self.equation.residual(
+        residual = self.equation.residual(
             samples, solver.forward(samples), solver._params
         )
+        return residual**2

pina/_src/condition/input_target_condition.py

@@ -113,18 +113,20 @@ def target(self):
         """
         return self.data.target
 
-    def evaluate(self, batch, solver):
+    def evaluate(self, batch, solver, loss):
         """
         Evaluate the supervised condition on the given batch using the solver.
 
-        This method computes the element-wise prediction error associated with
-        the condition using the input and target stored in the provided batch.
+        This method computes the element-wise loss associated with the
+        condition using the input and target stored in the provided batch.
 
         :param batch: The batch containing ``input`` and ``target`` entries.
         :type batch: dict | _DataManager
         :param solver: The solver containing the model.
         :type solver: ~pina.solver.solver.SolverInterface
-        :return: The non-aggregated prediction error.
+        :param loss: The non-aggregating loss function to apply.
+        :type loss: torch.nn.Module
+        :return: The non-aggregated loss tensor.
         :rtype: LabelTensor | torch.Tensor | Graph | Data
         """
-        return solver.forward(batch["input"]) - batch["target"]
+        return loss(solver.forward(batch["input"]), batch["target"])

pina/_src/solver/single_model_simple_solver.py

@@ -0,0 +1,135 @@
+"""Module for the SingleModelSimpleSolver."""
+
+import torch
+from torch.nn.modules.loss import _Loss
+
+from pina._src.condition.domain_equation_condition import (
+    DomainEquationCondition,
+)
+from pina._src.condition.input_equation_condition import (
+    InputEquationCondition,
+)
+from pina._src.condition.input_target_condition import InputTargetCondition
+from pina._src.core.utils import check_consistency
+from pina._src.loss.loss_interface import LossInterface
+from pina._src.solver.solver import SingleSolverInterface
+
+
+class SingleModelSimpleSolver(SingleSolverInterface):
+    """
+    Minimal single-model solver with explicit residual evaluation, reduction,
+    and loss aggregation across conditions.
+
+    The solver orchestrates a uniform workflow for all conditions in the batch:
+
+         1. evaluate the condition and obtain a non-aggregated loss tensor;
+         2. apply a reduction to obtain a scalar loss for that condition;
+     4. return the per-condition losses, which are aggregated by the inherited
+       solver machinery through the configured weighting.
+    """
+
+    accepted_conditions_types = (
+        InputTargetCondition,
+        InputEquationCondition,
+        DomainEquationCondition,
+    )
+
+    def __init__(
+        self,
+        problem,
+        model,
+        optimizer=None,
+        scheduler=None,
+        weighting=None,
+        loss=None,
+        use_lt=True,
+    ):
+        """
+        Initialize the single-model simple solver.
+
+        :param AbstractProblem problem: The problem to be solved.
+        :param torch.nn.Module model: The neural network model to be used.
+        :param Optimizer optimizer: The optimizer to be used.
+        :param Scheduler scheduler: Learning rate scheduler.
+        :param WeightingInterface weighting: The weighting schema to be used.
+        :param torch.nn.Module loss: The element-wise loss module whose
+            reduction strategy is reused by the solver. If ``None``,
+            :class:`torch.nn.MSELoss` is used.
+        :param bool use_lt: If ``True``, the solver uses LabelTensors as input.
+        """
+        if loss is None:
+            loss = torch.nn.MSELoss()
+
+        check_consistency(loss, (LossInterface, _Loss), subclass=False)
+
+        super().__init__(
+            model=model,
+            problem=problem,
+            optimizer=optimizer,
+            scheduler=scheduler,
+            weighting=weighting,
+            use_lt=use_lt,
+        )
+
+        self._loss_fn = loss
+        self._reduction = getattr(loss, "reduction", "mean")
+
+        if hasattr(self._loss_fn, "reduction"):
+            self._loss_fn.reduction = "none"
+
+    def optimization_cycle(self, batch):
+        """
+        Compute one reduced loss per condition in the batch.
+
+        :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 reduced losses for all conditions.
+        :rtype: dict[str, torch.Tensor]
+        """
+        condition_losses = {}
+
+        for condition_name, data in batch:
+            condition = self.problem.conditions[condition_name]
+            condition_data = dict(data)
+
+            if hasattr(condition_data.get("input"), "requires_grad_"):
+                condition_data["input"] = condition_data[
+                    "input"
+                ].requires_grad_()
+
+            condition_loss_tensor = condition.evaluate(
+                condition_data, self, self._loss_fn
+            )
+            condition_losses[condition_name] = self._apply_reduction(
+                condition_loss_tensor
+            )
+
+        return condition_losses
+
+    def _apply_reduction(self, value):
+        """
+        Apply the configured reduction to a non-aggregated condition tensor.
+
+        :param value: The non-aggregated tensor returned by a condition.
+        :type value: torch.Tensor
+        :return: The reduced scalar tensor.
+        :rtype: torch.Tensor
+        :raises ValueError: If the reduction is not supported.
+        """
+        if self._reduction == "none":
+            return value
+        if self._reduction == "mean":
+            return value.mean()
+        if self._reduction == "sum":
+            return value.sum()
+        raise ValueError(f"Unsupported reduction '{self._reduction}'.")
+
+    @property
+    def loss(self):
+        """
+        The underlying element-wise loss module.
+
+        :return: The stored loss module.
+        :rtype: torch.nn.Module
+        """
+        return self._loss_fn

pina/solver/__init__.py

@@ -13,6 +13,7 @@
     "SolverInterface",
     "SingleSolverInterface",
     "MultiSolverInterface",
+    "SingleModelSimpleSolver",
     "PINNInterface",
     "PINN",
     "GradientPINN",
@@ -36,6 +37,9 @@
     SingleSolverInterface,
     MultiSolverInterface,
 )
+from pina._src.solver.single_model_simple_solver import (
+    SingleModelSimpleSolver,
+)
 from pina._src.solver.physics_informed_solver.pinn import PINNInterface, PINN
 from pina._src.solver.physics_informed_solver.gradient_pinn import GradientPINN
 from pina._src.solver.physics_informed_solver.causal_pinn import CausalPINN

tests/test_condition/test_domain_equation_condition.py

@@ -48,8 +48,12 @@ def equation_func(input_, output_, params_):
     cond = Condition(domain=example_domain, equation=Equation(equation_func))
     solver = DummySolver()
     batch = {"input": samples}
+    loss = torch.nn.MSELoss(reduction="none")
 
-    residual = cond.evaluate(batch, solver)
-    expected = samples.extract(["x"]) - solver._params["shift"]
+    residual = cond.evaluate(batch, solver, loss)
+    expected = loss(
+        samples.extract(["x"]) - solver._params["shift"],
+        torch.zeros_like(samples.extract(["x"]) - solver._params["shift"]),
+    )
 
     torch.testing.assert_close(residual, expected)

tests/test_condition/test_input_equation_condition.py

@@ -95,8 +95,12 @@ def equation_func(input_, output_, params_):
     condition = Condition(input=pts, equation=Equation(equation_func))
     solver = DummySolver()
     batch = {"input": pts}
+    loss = torch.nn.MSELoss(reduction="none")
 
-    residual = condition.evaluate(batch, solver)
-    expected = pts.extract(["y"]) - solver._params["shift"]
+    residual = condition.evaluate(batch, solver, loss)
+    expected = loss(
+        pts.extract(["y"]) - solver._params["shift"],
+        torch.zeros_like(pts.extract(["y"]) - solver._params["shift"]),
+    )
 
     torch.testing.assert_close(residual, expected)

tests/test_condition/test_input_target_condition.py

@@ -220,10 +220,11 @@ def test_evaluate_tensor_input_target_condition():
     target_tensor = torch.tensor([[1.5, 3.5], [5.5, 7.5]])
     condition = Condition(input=input_tensor, target=target_tensor)
     solver = DummySolver()
+    loss_fn = torch.nn.MSELoss(reduction="none")
 
     batch = {"input": condition.input, "target": condition.target}
-    loss = condition.evaluate(batch, solver)
-    expected = solver.forward(input_tensor) - target_tensor
+    loss = condition.evaluate(batch, solver, loss_fn)
+    expected = loss_fn(solver.forward(input_tensor), target_tensor)
 
     torch.testing.assert_close(loss, expected)
 

tests/test_solver/test_single_model_simple_solver.py

@@ -0,0 +1,100 @@
+import pytest
+import torch
+
+from pina import LabelTensor, Condition
+from pina.model import FeedForward
+from pina.trainer import Trainer
+from pina.solver import SingleModelSimpleSolver
+from pina.condition import (
+    InputTargetCondition,
+    InputEquationCondition,
+    DomainEquationCondition,
+)
+from pina.problem.zoo import (
+    Poisson2DSquareProblem as Poisson,
+    InversePoisson2DSquareProblem as InversePoisson,
+)
+from torch._dynamo.eval_frame import OptimizedModule
+
+
+problem = Poisson()
+problem.discretise_domain(10)
+inverse_problem = InversePoisson(load=True, data_size=0.01)
+inverse_problem.discretise_domain(10)
+
+input_pts = torch.rand(10, len(problem.input_variables))
+input_pts = LabelTensor(input_pts, problem.input_variables)
+output_pts = torch.rand(10, len(problem.output_variables))
+output_pts = LabelTensor(output_pts, problem.output_variables)
+problem.conditions["data"] = Condition(input=input_pts, target=output_pts)
+
+model = FeedForward(len(problem.input_variables), len(problem.output_variables))
+
+
+@pytest.mark.parametrize("problem", [problem, inverse_problem])
+def test_constructor(problem):
+    solver = SingleModelSimpleSolver(problem=problem, model=model)
+
+    assert solver.accepted_conditions_types == (
+        InputTargetCondition,
+        InputEquationCondition,
+        DomainEquationCondition,
+    )
+
+
+@pytest.mark.parametrize("problem", [problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("compile", [True, False])
+def test_solver_train(problem, batch_size, compile):
+    solver = SingleModelSimpleSolver(model=model, problem=problem)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=1.0,
+        val_size=0.0,
+        test_size=0.0,
+        #compile=compile,
+    )
+    trainer.train()
+    if trainer.compile:
+        assert isinstance(solver.model, OptimizedModule)
+
+
+@pytest.mark.parametrize("problem", [problem, inverse_problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("compile", [True, False])
+def test_solver_validation(problem, batch_size, compile):
+    solver = SingleModelSimpleSolver(model=model, problem=problem)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=0.9,
+        val_size=0.1,
+        test_size=0.0,
+        compile=compile,
+    )
+    trainer.train()
+    if trainer.compile:
+        assert isinstance(solver.model, OptimizedModule)
+
+
+@pytest.mark.parametrize("problem", [problem, inverse_problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("compile", [True, False])
+def test_solver_test(problem, batch_size, compile):
+    solver = SingleModelSimpleSolver(model=model, problem=problem)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=0.7,
+        val_size=0.2,
+        test_size=0.1,
+        #compile=compile,
+    )
+    trainer.test()