Fix adaptive refinement (#571)

---------

Co-authored-by: Dario Coscia <[email protected]>

Author: FilippoOlivo

docs/source/_rst/_code.rst

@@ -238,7 +238,8 @@ Callbacks
 
     Processing callback <callback/processing_callback.rst>
     Optimizer callback <callback/optimizer_callback.rst>
-    Refinment callback <callback/adaptive_refinment_callback.rst>
+    R3 Refinment callback <callback/refinement/r3_refinement.rst>
+    Refinment Interface callback <callback/refinement/refinement_interface.rst>
     Weighting callback <callback/linear_weight_update_callback.rst>
 
 Losses and Weightings

…callback/adaptive_refinment_callback.rst …st/callback/refinement/r3_refinement.rstdocs/source/_rst/callback/adaptive_refinment_callback.rst renamed to docs/source/_rst/callback/refinement/r3_refinement.rst docs/source/_rst/callback/adaptive_refinment_callback.rst renamed to docs/source/_rst/callback/refinement/r3_refinement.rst

@@ -1,7 +1,7 @@
 Refinments callbacks
 =======================
 
-.. currentmodule:: pina.callback.adaptive_refinement_callback
+.. currentmodule:: pina.callback.refinement
 .. autoclass:: R3Refinement
    :members:
    :show-inheritance:

docs/source/_rst/callback/refinement/refinement_interface.rst

@@ -0,0 +1,7 @@
+Refinement Interface
+=======================
+
+.. currentmodule:: pina.callback.refinement
+.. autoclass:: RefinementInterface
+   :members:
+   :show-inheritance:

pina/callback/__init__.py

@@ -2,13 +2,13 @@
 
 __all__ = [
     "SwitchOptimizer",
-    "R3Refinement",
     "MetricTracker",
     "PINAProgressBar",
     "LinearWeightUpdate",
+    "R3Refinement",
 ]
 
 from .optimizer_callback import SwitchOptimizer
-from .adaptive_refinement_callback import R3Refinement
 from .processing_callback import MetricTracker, PINAProgressBar
 from .linear_weight_update_callback import LinearWeightUpdate
+from .refinement import R3Refinement

pina/callback/adaptive_refinement_callback.py

@@ -0,0 +1,11 @@
+"""
+Module for Pina Refinement callbacks.
+"""
+
+__all__ = [
+    "RefinementInterface",
+    "R3Refinement",
+]
+
+from .refinement_interface import RefinementInterface
+from .r3_refinement import R3Refinement

pina/callback/refinement/__init__.py

@@ -0,0 +1,88 @@
+"""Module for the R3Refinement callback."""
+
+import torch
+from torch import nn
+from torch.nn.modules.loss import _Loss
+from .refinement_interface import RefinementInterface
+from ...label_tensor import LabelTensor
+from ...utils import check_consistency
+from ...loss import LossInterface
+
+
+class R3Refinement(RefinementInterface):
+    """
+    PINA Implementation of an R3 Refinement Callback.
+    """
+
+    def __init__(
+        self, sample_every, residual_loss=nn.L1Loss, condition_to_update=None
+    ):
+        """
+        This callback implements the R3 (Retain-Resample-Release) routine for
+        sampling new points based on adaptive search.
+        The algorithm incrementally accumulates collocation points in regions
+        of high PDE residuals, and releases those with low residuals.
+        Points are sampled uniformly in all regions where sampling is needed.
+
+        .. seealso::
+
+            Original Reference: Daw, Arka, et al. *Mitigating Propagation
+            Failures in Physics-informed Neural Networks
+            using Retain-Resample-Release (R3) Sampling. (2023)*.
+            DOI: `10.48550/arXiv.2207.02338
+            <https://doi.org/10.48550/arXiv.2207.02338>`_
+
+        :param int sample_every: Frequency for sampling.
+        :param loss: Loss function
+        :type loss: LossInterface | ~torch.nn.modules.loss._Loss
+        :param condition_to_update: The conditions to update during the
+            refinement process. If None, all conditions with a conditions will
+            be updated. Default is None.
+        :type condition_to_update: list(str) | tuple(str) | str
+        :raises ValueError: If the condition_to_update is not a string or
+            iterable of strings.
+        :raises TypeError: If the residual_loss is not a subclass of
+            torch.nn.Module.
+
+
+        Example:
+            >>> r3_callback = R3Refinement(sample_every=5)
+        """
+        super().__init__(sample_every, condition_to_update)
+        # check consistency loss
+        check_consistency(residual_loss, (LossInterface, _Loss), subclass=True)
+        self.loss_fn = residual_loss(reduction="none")
+
+    def sample(self, current_points, condition_name, solver):
+        """
+        Sample new points based on the R3 refinement strategy.
+
+        :param current_points: Current points in the domain.
+        :param condition_name: Name of the condition to update.
+        :param PINNInterface solver: The solver object.
+        :return: New points sampled based on the R3 strategy.
+        :rtype: LabelTensor
+        """
+        # Compute residuals for the given condition (average over fields)
+        condition = solver.problem.conditions[condition_name]
+        target = solver.compute_residual(
+            current_points.requires_grad_(True), condition.equation
+        )
+        residuals = self.loss_fn(target, torch.zeros_like(target)).mean(
+            dim=tuple(range(1, target.ndim))
+        )
+
+        # Prepare new points
+        labels = current_points.labels
+        domain_name = solver.problem.conditions[condition_name].domain
+        domain = solver.problem.domains[domain_name]
+        num_old_points = self.initial_population_size[condition_name]
+        mask = (residuals > residuals.mean()).flatten()
+
+        if mask.any():  # Use high-residual points
+            pts = current_points[mask]
+            pts.labels = labels
+            retain_pts = len(pts)
+            samples = domain.sample(num_old_points - retain_pts, "random")
+            return LabelTensor.cat([pts, samples])
+        return domain.sample(num_old_points, "random")