Reimplement refinement

Author: FilippoOlivo

pina/callback/__init__.py

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

pina/callback/adaptive_refinement_callback.py

@@ -0,0 +1,3 @@
+__all__ = ["R3Refinement"]
+
+from .r3_refinement import R3Refinement

pina/callback/refinement/__init__.py

@@ -0,0 +1,79 @@
+"""Module for the R3Refinement callback."""
+
+import torch
+from .refinement_interface import RefinementInterface
+from ...label_tensor import LabelTensor
+from ...utils import check_consistency
+
+
+class R3Refinement(RefinementInterface):
+    """
+    PINA Implementation of an R3 Refinement Callback.
+    """
+
+    def __init__(self, sample_every):
+        """
+        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.
+        :raises ValueError: If `sample_every` is not an integer.
+
+        Example:
+            >>> r3_callback = R3Refinement(sample_every=5)
+        """
+
+        super().__init__(sample_every=sample_every)
+        self.const_pts = None
+
+    def sample(self, condition_name, condition):
+        avg_res, res = self.per_point_residual([condition_name])
+        pts = self.dataset.conditions_dict[condition_name]["input"]
+        domain = condition.domain
+        labels = pts.labels
+        pts = pts.cpu().detach().as_subclass(torch.Tensor)
+        residuals = res[condition_name]
+        mask = (residuals > avg_res).flatten()
+        if any(mask):  # append residuals greater than average
+            pts = (pts[mask]).as_subclass(LabelTensor)
+            pts.labels = labels
+            numb_pts = self.const_pts[condition_name] - len(pts)
+        else:
+            numb_pts = self.const_pts[condition_name]
+            pts = None
+        self.problem.discretise_domain(numb_pts, "random", domains=[domain])
+        sampled_points = self.problem.discretised_domains[domain]
+        tmp = (
+            sampled_points
+            if pts is None
+            else LabelTensor.cat([pts, sampled_points])
+        )
+        return tmp
+
+    def on_train_start(self, trainer, _):
+        """
+        Callback function called at the start of training.
+
+        This method extracts the locations for sampling from the problem
+        conditions and calculates the total population.
+
+        :param trainer: The trainer object managing the training process.
+        :type trainer: pytorch_lightning.Trainer
+        :param _: Placeholder argument (not used).
+        """
+        super().on_train_start(trainer, _)
+        self.const_pts = {}
+        for condition in self.conditions:
+            pts = self.dataset.conditions_dict[condition]["input"]
+            self.const_pts[condition] = len(pts)

pina/callback/refinement/r3_refinement.py

@@ -0,0 +1,103 @@
+"""
+RefinementInterface class for handling the refinement of points in a neural
+network training process.
+"""
+
+import torch
+from abc import ABCMeta
+from lightning.pytorch import Callback
+from torch_geometric.data.feature_store import abstractmethod
+from torch_geometric.nn.conv import point_transformer_conv
+from ...condition.domain_equation_condition import DomainEquationCondition
+
+
+class RefinementInterface(Callback):
+    """
+    Interface class of Refinement
+    """
+
+    def __init__(self, sample_every):
+        """
+        Initializes the RefinementInterface.
+
+        :param int sample_every: The number of epochs between each refinement.
+        """
+        self.sample_every = sample_every
+        self.conditions = None
+        self.dataset = None
+        self.solver = None
+
+    def on_train_start(self, trainer, _):
+        """
+        Called when the training begins. It initializes the conditions and
+        dataset.
+
+        :param lightning.pytorch.Trainer trainer: The trainer object.
+        :param _: Unused argument.
+        """
+        self.problem = trainer.solver.problem
+        self.solver = trainer.solver
+        self.conditions = {}
+        for name, cond in self.problem.conditions.items():
+            if isinstance(cond, DomainEquationCondition):
+                self.conditions[name] = cond
+        self.dataset = trainer.datamodule.train_dataset
+
+    @property
+    def points(self):
+        """
+        Returns the points of the dataset.
+        """
+        return self.dataset.conditions_dict
+
+    def on_train_epoch_end(self, trainer, _):
+        """
+        Performs the refinement at the end of each training epoch (if needed).
+
+        :param lightning.pytorch.Trainer trainer: The trainer object.
+        :param _: Unused argument.
+        """
+        if trainer.current_epoch % self.sample_every == 0:
+            self.update()
+
+    def update(self):
+        """
+        Performs the refinement of the points.
+        """
+        new_points = {}
+        for name, condition in self.conditions.items():
+            new_points[name] = {"input": self.sample(name, condition)}
+        self.dataset.update_data(new_points)
+
+    def per_point_residual(self, conditions_name=None):
+        """
+        Computes the residuals for a PINN object.
+
+        :return: the total loss, and pointwise loss.
+        :rtype: tuple
+        """
+        # compute residual
+        res_loss = {}
+        tot_loss = []
+        points = self.points
+        if conditions_name is None:
+            conditions_name = list(self.conditions.keys())
+        for name in conditions_name:
+            cond = self.conditions[name]
+            cond_points = points[name]["input"]
+            target = self._compute_residual(cond_points, cond.equation)
+            res_loss[name] = torch.abs(target).as_subclass(torch.Tensor)
+            tot_loss.append(torch.abs(target))
+        return torch.vstack(tot_loss).tensor.mean(), res_loss
+
+    def _compute_residual(self, pts, equation):
+        pts.requires_grad_(True)
+        pts.retain_grad()
+        return equation.residual(pts, self.solver.forward(pts))
+
+    @abstractmethod
+    def sample(self, condition):
+        """
+        Samples new points based on the condition.
+        """
+        pass