input target evaluate

Author: ndem0

pina/_src/condition/equation_condition_base.py

@@ -15,18 +15,18 @@ class EquationConditionBase(ConditionBase):
     :class:`~pina.condition.InputEquationCondition`.
     """
 
-    def evaluate(self, samples, solver):
+    def evaluate(self, batch, solver):
         """
-        Evaluate the equation residual on the given samples using the solver.
+        Evaluate the equation residual on the given batch 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 batch: The batch containing the ``input`` entry.
+        :type batch: dict | _DataManager
         :param solver: The solver containing the model and any additional
             parameters (e.g., unknown parameters for inverse problems).
         :type solver: ~pina.solver.solver.SolverInterface
@@ -35,9 +35,12 @@ def evaluate(self, samples, solver):
 
         :Example:
 
-            >>> residuals = condition.evaluate(input_samples, solver)
+            >>> residuals = condition.evaluate(
+            ...     {"input": input_samples}, solver
+            ... )
             >>> # residuals is a non-reduced tensor of shape (n_samples, ...)
         """
+        samples = batch["input"]
         return self.equation.residual(
             samples, solver.forward(samples), solver._params
         )

pina/_src/condition/input_target_condition.py

@@ -112,3 +112,19 @@ def target(self):
             list[Data] | tuple[Graph] | tuple[Data]
         """
         return self.data.target
+
+    def evaluate(self, batch, solver):
+        """
+        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.
+
+        :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.
+        :rtype: LabelTensor | torch.Tensor | Graph | Data
+        """
+        return solver.forward(batch["input"]) - batch["target"]

tests/test_condition/test_domain_equation_condition.py

@@ -47,8 +47,9 @@ def equation_func(input_, output_, params_):
     samples = LabelTensor(torch.randn(12, 2), labels=["x", "y"])
     cond = Condition(domain=example_domain, equation=Equation(equation_func))
     solver = DummySolver()
+    batch = {"input": samples}
 
-    residual = cond.evaluate(samples, solver)
+    residual = cond.evaluate(batch, solver)
     expected = samples.extract(["x"]) - solver._params["shift"]
 
     torch.testing.assert_close(residual, expected)

tests/test_condition/test_input_equation_condition.py

@@ -94,8 +94,9 @@ def equation_func(input_, output_, params_):
     pts = LabelTensor(torch.randn(10, 2), labels=["x", "y"])
     condition = Condition(input=pts, equation=Equation(equation_func))
     solver = DummySolver()
+    batch = {"input": pts}
 
-    residual = condition.evaluate(pts, solver)
+    residual = condition.evaluate(batch, solver)
     expected = pts.extract(["y"]) - solver._params["shift"]
 
     torch.testing.assert_close(residual, expected)

tests/test_condition/test_input_target_condition.py

@@ -5,6 +5,11 @@
 from pina._src.condition.batch_manager import _BatchManager
 
 
+class DummySolver:
+    def forward(self, samples):
+        return 2 * samples
+
+
 def _create_tensor_data(use_lt=False):
     if use_lt:
         input_tensor = LabelTensor(torch.rand((10, 3)), ["x", "y", "z"])
@@ -210,6 +215,19 @@ def test_getitem_tensor_input_tensor_target_condition_label_tensor():
     assert torch.allclose(item.target, target_tensor[index])
 
 
+def test_evaluate_tensor_input_target_condition():
+    input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    target_tensor = torch.tensor([[1.5, 3.5], [5.5, 7.5]])
+    condition = Condition(input=input_tensor, target=target_tensor)
+    solver = DummySolver()
+
+    batch = {"input": condition.input, "target": condition.target}
+    loss = condition.evaluate(batch, solver)
+    expected = solver.forward(input_tensor) - target_tensor
+
+    torch.testing.assert_close(loss, expected)
+
+
 @pytest.mark.parametrize("use_lt", [True, False])
 def test_getitem_graph_input_tensor_target_condition(use_lt):
     input_graph, target_tensor = _create_graph_data(False, use_lt=use_lt)