diff --git a/pina/callback/__init__.py b/pina/callback/__init__.py index 3bb12ad..4cf5147 100644 --- a/pina/callback/__init__.py +++ b/pina/callback/__init__.py @@ -3,8 +3,10 @@ __all__ = [ "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 diff --git a/pina/callback/linear_weight_update_callback.py b/pina/callback/linear_weight_update_callback.py new file mode 100644 index 0000000..02a8878 --- /dev/null +++ b/pina/callback/linear_weight_update_callback.py @@ -0,0 +1,85 @@ +"""PINA Callbacks Implementations""" + +import warnings +from lightning.pytorch.callbacks import Callback +from ..utils import check_consistency +from ..loss import ScalarWeighting + + +class LinearWeightUpdate(Callback): + """ + Callback to linearly adjust the weight of a condition from an + initial value to a target value over a specified number of epochs. + """ + + def __init__( + self, target_epoch, condition_name, initial_value, target_value + ): + """ + Callback initialization. + + :param int target_epoch: The epoch at which the weight of the condition + should reach the target value. + :param str condition_name: The name of the condition whose weight + should be adjusted. + :param float initial_value: The initial value of the weight. + :param float target_value: The target value of the weight. + """ + super().__init__() + self.target_epoch = target_epoch + self.condition_name = condition_name + self.initial_value = initial_value + self.target_value = target_value + + # Check consistency + check_consistency(self.target_epoch, int, subclass=False) + check_consistency(self.condition_name, str, subclass=False) + check_consistency(self.initial_value, (float, int), subclass=False) + check_consistency(self.target_value, (float, int), subclass=False) + + def on_train_start(self, trainer, solver): + """ + Initialize the weight of the condition to the specified `initial_value`. + + :param Trainer trainer: a pina:class:`Trainer` instance. + :param SolverInterface solver: a pina:class:`SolverInterface` instance. + """ + # Check that the target epoch is valid + if not 0 < self.target_epoch <= trainer.max_epochs: + raise ValueError( + "`target_epoch` must be greater than 0" + " and less than or equal to `max_epochs`." + ) + + # Check that the condition is a problem condition + if self.condition_name not in solver.problem.conditions: + raise ValueError( + f"`{self.condition_name}` must be a problem condition." + ) + + # Check that the initial value is not equal to the target value + if self.initial_value == self.target_value: + warnings.warn( + "`initial_value` is equal to `target_value`. " + "No effective adjustment will be performed.", + UserWarning, + ) + + # Check that the weighting schema is ScalarWeighting + if not isinstance(solver.weighting, ScalarWeighting): + raise ValueError("The weighting schema must be ScalarWeighting.") + + # Initialize the weight of the condition + solver.weighting.weights[self.condition_name] = self.initial_value + + def on_train_epoch_start(self, trainer, solver): + """ + Adjust at each epoch the weight of the condition. + + :param Trainer trainer: a pina:class:`Trainer` instance. + :param SolverInterface solver: a pina:class:`SolverInterface` instance. + """ + if 0 < trainer.current_epoch <= self.target_epoch: + solver.weighting.weights[self.condition_name] += ( + self.target_value - self.initial_value + ) / (self.target_epoch - 1) diff --git a/tests/test_callback/test_linear_weight_update_callback.py b/tests/test_callback/test_linear_weight_update_callback.py new file mode 100644 index 0000000..c1f4cf3 --- /dev/null +++ b/tests/test_callback/test_linear_weight_update_callback.py @@ -0,0 +1,164 @@ +import pytest +import math +from pina.solver import PINN +from pina.loss import ScalarWeighting +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.problem.zoo import Poisson2DSquareProblem as Poisson +from pina.callback import LinearWeightUpdate + + +# Define the problem +poisson_problem = Poisson() +poisson_problem.discretise_domain(50, "grid") +cond_name = list(poisson_problem.conditions.keys())[0] + +# Define the model +model = FeedForward( + input_dimensions=len(poisson_problem.input_variables), + output_dimensions=len(poisson_problem.output_variables), + layers=[32, 32], +) + +# Define the weighting schema +weights_dict = {key: 1 for key in poisson_problem.conditions.keys()} +weighting = ScalarWeighting(weights=weights_dict) + +# Define the solver +solver = PINN(problem=poisson_problem, model=model, weighting=weighting) + +# Value used for testing +epochs = 10 + + +@pytest.mark.parametrize("initial_value", [1, 5.5]) +@pytest.mark.parametrize("target_value", [10, 25.5]) +def test_constructor(initial_value, target_value): + LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value=initial_value, + target_value=target_value, + ) + + # Target_epoch must be int + with pytest.raises(ValueError): + LinearWeightUpdate( + target_epoch=10.0, + condition_name=cond_name, + initial_value=0, + target_value=1, + ) + + # Condition_name must be str + with pytest.raises(ValueError): + LinearWeightUpdate( + target_epoch=epochs, + condition_name=100, + initial_value=0, + target_value=1, + ) + + # Initial_value must be float or int + with pytest.raises(ValueError): + LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value="0", + target_value=1, + ) + + # Target_value must be float or int + with pytest.raises(ValueError): + LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value=0, + target_value="1", + ) + + +@pytest.mark.parametrize("initial_value, target_value", [(1, 10), (10, 1)]) +def test_training(initial_value, target_value): + callback = LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value=initial_value, + target_value=target_value, + ) + trainer = Trainer( + solver=solver, + callbacks=[callback], + accelerator="cpu", + max_epochs=epochs, + ) + trainer.train() + + # Check that the final weight value matches the target value + final_value = solver.weighting.weights[cond_name] + assert math.isclose(final_value, target_value) + + # Target_epoch must be greater than 0 + with pytest.raises(ValueError): + callback = LinearWeightUpdate( + target_epoch=0, + condition_name=cond_name, + initial_value=0, + target_value=1, + ) + trainer = Trainer( + solver=solver, + callbacks=[callback], + accelerator="cpu", + max_epochs=5, + ) + trainer.train() + + # Target_epoch must be less than or equal to max_epochs + with pytest.raises(ValueError): + callback = LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value=0, + target_value=1, + ) + trainer = Trainer( + solver=solver, + callbacks=[callback], + accelerator="cpu", + max_epochs=epochs - 1, + ) + trainer.train() + + # Condition_name must be a problem condition + with pytest.raises(ValueError): + callback = LinearWeightUpdate( + target_epoch=epochs, + condition_name="not_a_condition", + initial_value=0, + target_value=1, + ) + trainer = Trainer( + solver=solver, + callbacks=[callback], + accelerator="cpu", + max_epochs=epochs, + ) + trainer.train() + + # Weighting schema must be ScalarWeighting + with pytest.raises(ValueError): + callback = LinearWeightUpdate( + target_epoch=epochs, + condition_name=cond_name, + initial_value=0, + target_value=1, + ) + unweighted_solver = PINN(problem=poisson_problem, model=model) + trainer = Trainer( + solver=unweighted_solver, + callbacks=[callback], + accelerator="cpu", + max_epochs=epochs, + ) + trainer.train()