evaluate for physics condition

Author: ndem0

pina/_src/condition/domain_equation_condition.py

@@ -1,11 +1,13 @@
 """Module for the DomainEquationCondition class."""
 
-from pina._src.condition.condition_base import ConditionBase
+from pina._src.condition.equation_condition_base import (
+    EquationConditionBase,
+)
 from pina._src.domain.domain_interface import DomainInterface
 from pina._src.equation.equation_interface import EquationInterface
 
 
-class DomainEquationCondition(ConditionBase):
+class DomainEquationCondition(EquationConditionBase):
     """
     The class :class:`DomainEquationCondition` defines a condition based on a
     ``domain`` and an ``equation``. This condition is typically used in
@@ -92,4 +94,29 @@ def store_data(self, **kwargs):
         :rtype: dict
         """
         setattr(self, "domain", kwargs.get("domain"))
-        setattr(self, "equation", kwargs.get("equation"))
+        setattr(self, "_equation", kwargs.get("equation"))
+
+    @property
+    def equation(self):
+        """
+        Return the equation associated with this condition.
+
+        :return: Equation associated with this condition.
+        :rtype: EquationInterface
+        """
+        return self._equation
+
+    @equation.setter
+    def equation(self, value):
+        """
+        Set the equation associated with this condition.
+
+        :param EquationInterface value: The equation to associate
+            with this condition.
+        """
+        if not isinstance(value, EquationInterface):
+            raise TypeError(
+                "The equation must be an instance of "
+                "EquationInterface."
+            )
+        self._equation = value

pina/_src/condition/equation_condition_base.py

@@ -0,0 +1,43 @@
+"""Module for the EquationConditionBase class."""
+
+from pina._src.condition.condition_base import ConditionBase
+
+
+class EquationConditionBase(ConditionBase):
+    """
+    Base class for conditions that involve an equation.
+
+    This class provides the :meth:`evaluate` method, which computes the
+    non-aggregated residual of the equation given the input samples and a
+    solver. It is intended to be subclassed by conditions that define an
+    ``equation`` attribute, such as
+    :class:`~pina.condition.DomainEquationCondition` and
+    :class:`~pina.condition.InputEquationCondition`.
+    """
+
+    def evaluate(self, samples, solver):
+        """
+        Evaluate the equation residual on the given samples using the solver.
+
+        This method computes the non-aggregated, element-wise residual of the
+        equation. It performs a forward pass of the solver's model on the
+        input samples and then evaluates the equation residual. The returned
+        tensor is **not** reduced (i.e., no mean, sum, etc.), preserving the
+        per-sample residual values.
+
+        :param samples: The input samples on which to evaluate the residual.
+        :type samples: ~pina.label_tensor.LabelTensor
+        :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.
+        :rtype: ~pina.label_tensor.LabelTensor
+
+        :Example:
+
+            >>> residuals = condition.evaluate(input_samples, solver)
+            >>> # residuals is a non-reduced tensor of shape (n_samples, ...)
+        """
+        return self.equation.residual(
+            samples, solver.forward(samples), solver._params
+        )

pina/_src/condition/input_equation_condition.py

@@ -1,13 +1,15 @@
 """Module for the InputEquationCondition class and its subclasses."""
 
-from pina._src.condition.condition_base import ConditionBase
+from pina._src.condition.equation_condition_base import (
+    EquationConditionBase,
+)
 from pina._src.core.label_tensor import LabelTensor
 from pina._src.core.graph import Graph
 from pina._src.equation.equation_interface import EquationInterface
 from pina._src.condition.data_manager import _DataManager
 
 
-class InputEquationCondition(ConditionBase):
+class InputEquationCondition(EquationConditionBase):
     """
     The class :class:`InputEquationCondition` defines a condition based on
     ``input`` data and an ``equation``. This condition is typically used in

pina/condition/__init__.py

@@ -10,6 +10,7 @@
     "Condition",
     "ConditionInterface",
     "ConditionBase",
+    "EquationConditionBase",
     "DomainEquationCondition",
     "InputTargetCondition",
     "InputEquationCondition",
@@ -18,6 +19,9 @@
 
 from pina._src.condition.condition_interface import ConditionInterface
 from pina._src.condition.condition_base import ConditionBase
+from pina._src.condition.equation_condition_base import (
+    EquationConditionBase,
+)
 from pina._src.condition.condition import Condition
 from pina._src.condition.domain_equation_condition import (
     DomainEquationCondition,

tests/test_condition/test_domain_equation_condition.py

@@ -1,9 +1,20 @@
 import pytest
+import torch
 from pina import Condition
+from pina import LabelTensor
 from pina.domain import CartesianDomain
 from pina._src.equation.equation_factory import FixedValue
+from pina.equation import Equation
 from pina.condition import DomainEquationCondition
 
+
+class DummySolver:
+    def __init__(self):
+        self._params = {"shift": torch.tensor(0.25)}
+
+    def forward(self, samples):
+        return samples.extract(["x"]) - samples.extract(["y"])
+
 example_domain = CartesianDomain({"x": [0, 1], "y": [0, 1]})
 example_equation = FixedValue(0.0)
 
@@ -27,3 +38,17 @@ def test_getitem_not_implemented():
     cond = Condition(domain=example_domain, equation=FixedValue(0.0))
     with pytest.raises(NotImplementedError):
         cond[0]
+
+
+def test_evaluate_domain_equation_condition():
+    def equation_func(input_, output_, params_):
+        return output_ + input_.extract(["y"]) - params_["shift"]
+
+    samples = LabelTensor(torch.randn(12, 2), labels=["x", "y"])
+    cond = Condition(domain=example_domain, equation=Equation(equation_func))
+    solver = DummySolver()
+
+    residual = cond.evaluate(samples, solver)
+    expected = samples.extract(["x"]) - solver._params["shift"]
+
+    torch.testing.assert_close(residual, expected)

tests/test_condition/test_input_equation_condition.py

@@ -8,6 +8,14 @@
 from pina._src.condition.data_manager import _DataManager
 
 
+class DummySolver:
+    def __init__(self):
+        self._params = {"shift": torch.tensor(1.5)}
+
+    def forward(self, samples):
+        return samples.extract(["x"]) + samples.extract(["y"])
+
+
 def _create_pts_and_equation():
     def dummy_equation(pts):
         return pts["x"] ** 2 + pts["y"] ** 2 - 1
@@ -77,3 +85,17 @@ def test_getitems_tensor_equation_condition():
     assert isinstance(item, _DataManager)
     assert hasattr(item, "input")
     assert item.input.shape == (3, 2)
+
+
+def test_evaluate_tensor_equation_condition():
+    def equation_func(input_, output_, params_):
+        return output_ - input_.extract(["x"]) - params_["shift"]
+
+    pts = LabelTensor(torch.randn(10, 2), labels=["x", "y"])
+    condition = Condition(input=pts, equation=Equation(equation_func))
+    solver = DummySolver()
+
+    residual = condition.evaluate(pts, solver)
+    expected = pts.extract(["y"]) - solver._params["shift"]
+
+    torch.testing.assert_close(residual, expected)