clean code

Author: Monthly Tag bot

pina/callback/refinement/r3_refinement.py

@@ -1,17 +1,22 @@
 """Module for the R3Refinement callback."""
 
 import torch
+import torch.nn as 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):
+    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.
@@ -33,47 +38,33 @@ def __init__(self, sample_every):
         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")
 
-        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])
+    def sample(self, current_points, condition_name, solver):
+        # 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))
         )
-        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.
+        # 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()
 
-        :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)
+        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])
+        else:
+            return domain.sample(num_old_points, "random")

pina/callback/refinement/refinement_interface.py

@@ -3,101 +3,121 @@
 network training process.
 """
 
-import torch
-from abc import ABCMeta
+from abc import ABCMeta, abstractmethod
 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
+from ...utils import check_consistency
+from ...solver.physics_informed_solver import PINNInterface
 
 
 class RefinementInterface(Callback, metaclass=ABCMeta):
     """
-    Interface class of Refinement
+    Interface class of Refinement approaches.
     """
 
-    def __init__(self, sample_every):
+    def __init__(self, sample_every, condition_to_update=None):
         """
         Initializes the RefinementInterface.
 
         :param int sample_every: The number of epochs between each refinement.
         """
+        # check consistency of the input
+        check_consistency(sample_every, int)
+        if condition_to_update is not None:
+            if not isinstance(condition_to_update, (list, tuple)):
+                raise ValueError(
+                    "'condition_to_update' must be iter of strings."
+                )
+            check_consistency(condition_to_update, str)
+        # store
         self.sample_every = sample_every
-        self.conditions = None
-        self.dataset = None
-        self.solver = None
+        self._condition_to_update = condition_to_update
+        self._dataset = None
+        self._initial_population_size = None
 
-    def on_train_start(self, trainer, _):
+    def on_train_start(self, trainer, solver):
         """
         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
+        # check we have valid conditions names
+        if self._condition_to_update is None:
+            self._condition_to_update = list(solver.problem.conditions.keys())
 
-    @property
-    def points(self):
-        """
-        Returns the points of the dataset.
-        """
-        return self.dataset.conditions_dict
+        for cond in self._condition_to_update:
+            if cond not in solver.problem.conditions:
+                raise RuntimeError(
+                    f"Condition '{cond}' not found in "
+                    f"{list(solver.problem.conditions.keys())}."
+                )
+            if not hasattr(solver.problem.conditions[cond], "domain"):
+                raise RuntimeError(
+                    f"Condition '{cond}' does not contain a domain to "
+                    "sample from."
+                )
+        # check solver
+        if not isinstance(solver, PINNInterface):
+            raise RuntimeError(
+                f"Refinment strategies are currently implemented only "
+                "for physics informed based solvers. Please use a Solver "
+                "inheriting from 'PINNInterface'."
+            )
+        # store dataset
+        self._dataset = trainer.datamodule.train_dataset
+        # compute initial population size
+        self._initial_population_size = self._compute_population_size(
+            self._condition_to_update
+        )
+        return super().on_train_epoch_start(trainer, solver)
 
-    def on_train_epoch_end(self, trainer, _):
+    def on_train_epoch_end(self, trainer, solver):
         """
         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()
+            self._update_points(solver)
+        return super().on_train_epoch_end(trainer, solver)
 
-    def update(self):
+    @abstractmethod
+    def sample(self, current_points, condition_name, solver):
         """
-        Performs the refinement of the points.
+        Samples new points based on the condition.
         """
-        new_points = {}
-        for name, condition in self.conditions.items():
-            new_points[name] = {"input": self.sample(name, condition)}
-        self.dataset.update_data(new_points)
+        pass
 
-    def per_point_residual(self, conditions_name=None):
+    @property
+    def dataset(self):
         """
-        Computes the residuals for a PINN object.
-
-        :return: the total loss, and pointwise loss.
-        :rtype: tuple
+        Returns the dataset for training.
         """
-        # 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
+        return self._dataset
 
-    def _compute_residual(self, pts, equation):
-        pts.requires_grad_(True)
-        pts.retain_grad()
-        return equation.residual(pts, self.solver.forward(pts))
+    @property
+    def initial_population_size(self):
+        """
+        Returns the dataset for training.
+        """
+        return self._initial_population_size
 
-    @abstractmethod
-    def sample(self, condition):
+    def _update_points(self, solver):
         """
-        Samples new points based on the condition.
+        Performs the refinement of the points.
         """
-        pass
+        new_points = {}
+        for name in self._condition_to_update:
+            current_points = self.dataset.conditions_dict[name]["input"]
+            new_points[name] = {
+                "input": self.sample(current_points, name, solver)
+            }
+        self.dataset.update_data(new_points)
+
+    def _compute_population_size(self, conditions):
+        return {
+            cond: len(self.dataset.conditions_dict[cond]["input"])
+            for cond in conditions
+        }

tests/test_callback/test_adaptive_refinement_callback.py

@@ -1,3 +1,7 @@
+import pytest
+
+from torch.nn import MSELoss
+
 from pina.solver import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
@@ -7,28 +11,27 @@
 
 # make the problem
 poisson_problem = Poisson()
-boundaries = ["g1", "g2", "g3", "g4"]
-n = 10
-poisson_problem.discretise_domain(n, "grid", domains=boundaries)
-poisson_problem.discretise_domain(n, "grid", domains="D")
+poisson_problem.discretise_domain(10, "grid", domains=["g1", "g2", "g3", "g4"])
+poisson_problem.discretise_domain(10, "grid", domains="D")
 model = FeedForward(
     len(poisson_problem.input_variables), len(poisson_problem.output_variables)
 )
-
-# make the solver
 solver = PINN(problem=poisson_problem, model=model)
 
 
-def test_r3constructor():
+def test_constructor():
+    # good constructor
     R3Refinement(sample_every=10)
+    R3Refinement(sample_every=10, residual_loss=MSELoss)
+    R3Refinement(sample_every=10, condition_to_update=["D"])
+    # wrong constructor
+    with pytest.raises(ValueError):
+        R3Refinement(sample_every="str")
+    with pytest.raises(ValueError):
+        R3Refinement(sample_every=10, condition_to_update=3)
 
 
-def test_r3refinement_routine():
-    model = FeedForward(
-        len(poisson_problem.input_variables),
-        len(poisson_problem.output_variables),
-    )
-    solver = PINN(problem=poisson_problem, model=model)
+def test_sample():
     trainer = Trainer(
         solver=solver,
         callbacks=[R3Refinement(sample_every=1)],
@@ -43,4 +46,5 @@ def test_r3refinement_routine():
         loc: len(pts)
         for loc, pts in trainer.data_module.train_dataset.input.items()
     }
+    assert before_n_points == trainer.callbacks[0].initial_population_size
     assert before_n_points == after_n_points