From a6df3b051d40ad3b9b07b888f938be70241eb1d6 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Wed, 17 Apr 2024 16:47:15 +0200 Subject: [PATCH 1/3] competitive pinns --- pina/solvers/__init__.py | 2 + pina/solvers/pinns/competitive_pinn.py | 310 +++++++++++++++ tests/test_solvers/test_competitive_pinn.py | 418 ++++++++++++++++++++ 3 files changed, 730 insertions(+) create mode 100644 pina/solvers/pinns/competitive_pinn.py create mode 100644 tests/test_solvers/test_competitive_pinn.py diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py index 512df64b2..ee9a2b6ab 100644 --- a/pina/solvers/__init__.py +++ b/pina/solvers/__init__.py @@ -2,6 +2,7 @@ "PINN", "GPINN", "CausalPINN", + "CompetitivePINN" "GAROM", "SupervisedSolver", "SolverInterface" @@ -10,6 +11,7 @@ from .garom import GAROM from .pinns.pinn import PINN from .pinns.gpinn import GPINN +from .pinns.competitive_pinn import CompetitivePINN from .pinns.causalpinn import CausalPINN from .supervised import SupervisedSolver from .solver import SolverInterface diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py new file mode 100644 index 000000000..8efbb012a --- /dev/null +++ b/pina/solvers/pinns/competitive_pinn.py @@ -0,0 +1,310 @@ +""" Module for PINN """ + +import torch +import copy + +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 + +from torch.optim.lr_scheduler import ConstantLR + +from .basepinn import PINNInterface +from pina.utils import check_consistency +from pina.problem import InverseProblem + + +class CompetitivePINN(PINNInterface): + """ + TODO + + .. warning Condition Data Weight Not Suppoerted + """ + + def __init__( + self, + problem, + model, + discriminator=None, + loss=torch.nn.MSELoss(), + optimizer_model=torch.optim.Adam, + optimizer_model_kwargs={"lr": 0.001}, + optimizer_discriminator=torch.optim.Adam, + optimizer_discriminator_kwargs={"lr": 0.001}, + scheduler_model=ConstantLR, + scheduler_model_kwargs={"factor": 1, "total_iters": 0}, + scheduler_discriminator=ConstantLR, + scheduler_discriminator_kwargs={"factor": 1, "total_iters": 0}, + ): + """ + :param AbstractProblem problem: The formualation of the problem. + :param torch.nn.Module model: The neural network model to use + for the model. + :param torch.nn.Module discriminator: The neural network model to use + for the discriminator. If ``None``, the discriminator network will + have the same architecture as the model network. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.optim.Optimizer optimizer_model: The neural + network optimizer to use for the model network + , default is `torch.optim.Adam`. + :param dict optimizer_model_kwargs: Optimizer constructor keyword + args. for the model. + :param torch.optim.Optimizer optimizer_discriminator: The neural + network optimizer to use for the discriminator network + , default is `torch.optim.Adam`. + :param dict optimizer_discriminator_kwargs: Optimizer constructor + keyword args. for the discriminator. + :param torch.optim.LRScheduler scheduler_model: Learning + rate scheduler for the model. + :param dict scheduler_model_kwargs: LR scheduler constructor + keyword args. + :param torch.optim.LRScheduler scheduler_discriminator: Learning + rate scheduler for the discriminator. + """ + if discriminator is None: + discriminator = copy.deepcopy(model) + + super().__init__( + models=[model, discriminator], + problem=problem, + optimizers=[optimizer_model, optimizer_discriminator], + optimizers_kwargs=[ + optimizer_model_kwargs, + optimizer_discriminator_kwargs, + ], + extra_features=None, # PIGAN doesn't take extra features + loss=loss + ) + + # set automatic optimization for GANs + self.automatic_optimization = False + + # check consistency + check_consistency(scheduler_model, LRScheduler, subclass=True) + check_consistency(scheduler_model_kwargs, dict) + check_consistency(scheduler_discriminator, LRScheduler, subclass=True) + check_consistency(scheduler_discriminator_kwargs, dict) + + # assign schedulers + self._schedulers = [ + scheduler_model( + self.optimizers[0], **scheduler_model_kwargs + ), + scheduler_discriminator( + self.optimizers[1], **scheduler_discriminator_kwargs + ), + ] + + self._model = self.models[0] + self._discriminator = self.models[1] + + + def forward(self, x): + """ + Forward pass implementation for the PINN solver. + + :param LabelTensor x: Input tensor for the PINN solver. It expects + a tensor :math:`N \times D`, where :math:`N` the number of points + in the mesh, :math:`D` the dimension of the problem, + :return: PINN solution evaluated at the input points. + :rtype: LabelTensor + """ + return self.neural_net(x) + + + def training_step(self, batch, _): + """ + PINN solver training step. + + :param batch: The batch element in the dataloader. + :type batch: tuple + :param batch_idx: The batch index. + :type batch_idx: int + :return: The sum of the loss functions. + :rtype: LabelTensor + """ + + condition_idx = batch["condition"] + + for condition_id in range(condition_idx.min(), condition_idx.max() + 1): + + condition_name = self._dataloader.condition_names[condition_id] + condition = self.problem.conditions[condition_name] + pts = batch["pts"] + + if len(batch) == 2: + samples = pts[condition_idx == condition_id] + self.loss_phys(samples=samples, + equation=condition.equation, + condition_name=condition_name) + elif len(batch) == 3: + samples = pts[condition_idx == condition_id] + ground_truth = batch["output"][condition_idx == condition_id] + self._loss_data(samples, ground_truth, condition_name) + else: + raise ValueError("Batch size not supported") + + # clamp unknown parameters in InverseProblem (if needed), otherwise + # it returns None + self._clamp_params() + return + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the PINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics residual calculated based on given + samples and equation. + :rtype: LabelTensor + """ + return self.compute_residual(samples=samples, equation=equation) + + + def configure_optimizers(self): + """ + Optimizer configuration for the Competitive PINN solver. + + :return: The optimizers and the schedulers + :rtype: tuple(list, list) + """ + # if the problem is an InverseProblem, add the unknown parameters + # to the parameters that the optimizer needs to optimize + if isinstance(self.problem, InverseProblem): + self.optimizer_model.add_param_group( + { + "params": [ + self._params[var] + for var in self.problem.unknown_variables + ] + } + ) + return self.optimizers, self._schedulers + + + def _loss_phys(self, samples, equation, condition_name): + # train one step of discriminator + discriminator_bets = self.discriminator(samples.clone()) + self._train_discriminator(samples, equation, discriminator_bets) + # detaching samples from the computational graph to erase it and setting + # the gradient to true to create a new computational graph. + # In alternative set `retain_graph=True`. + samples = samples.detach() + samples.requires_grad = True + # train one step of the model + self._train_model(samples, equation, condition_name, discriminator_bets) + + def _train_discriminator(self, samples, equation, discriminator_bets): + """ + Trains the discriminator network of the Competitive PINN. + + :param LabelTensor samples: Input samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation representing + the physics. + :param Tensor discriminator_bets: Predictions made by the discriminator + network. + """ + # manual optimization + self.optimizer_discriminator.zero_grad() + # compute residual, we detach because the weights of the generator + # model are fixed + residual = self.loss_phys(samples=samples, + equation=equation).detach() + # compute competitive residual, the minus is because we maximise + competitive_residual = residual * discriminator_bets + loss_val = - self.loss( + torch.zeros_like(competitive_residual, requires_grad=True), + competitive_residual + ).as_subclass(torch.Tensor) + # backprop + loss_val.backward() + self.optimizer_discriminator.step() + return + + def _train_model(self, samples, equation, + condition_name, discriminator_bets): + """ + Trains the model network of the Competitive PINN. + + :param LabelTensor samples: Input samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation representing + the physics. + :param str condition_name: The condition name for tracking purposes. + :param Tensor discriminator_bets: Predictions made by the discriminator + network. + """ + # manual optimization + self.optimizer_model.zero_grad() + # compute residual (detached for discriminator) and log + residual = self.loss_phys(samples=samples, equation=equation) + # store logging + with torch.no_grad(): + loss_residual = self.loss( + torch.zeros_like(residual), + residual + ).as_subclass(torch.Tensor) + self.store_log(name=condition_name+f'_loss', + loss_val=float(loss_residual)) + # compute competitive residual, discriminator_bets are detached becase + # we optimize only the generator model + competitive_residual = residual * discriminator_bets.detach() + loss_val = self.loss( + torch.zeros_like(competitive_residual, requires_grad=True), + competitive_residual + ).as_subclass(torch.Tensor) + # backprop + loss_val.backward() + self.optimizer_model.step() + return + + def _loss_data(self, input_tensor, output_tensor, condition_name): + """ + Computes the data loss for the PINN solver based on input, + output, and condition name. This function is a wrapper of the function + :meth:`loss_data` used internally in PINA to handle the logging step. + + :param LabelTensor input_tensor: The input to the neural networks. + :param LabelTensor output_tensor: The true solution to compare the + network solution. + :param str condition_name: The condition name for tracking purposes. + :return: The computed data loss. + :rtype: torch.Tensor + """ + self.optimizer_model.zero_grad() + loss_val = super()._loss_data(input_tensor, + output_tensor, condition_name) + loss_val.backward() + self.optimizer_model.step() + return + + + @property + def neural_net(self): + return self._model + + @property + def discriminator(self): + return self._discriminator + + @property + def optimizer_model(self): + return self.optimizers[0] + + @property + def optimizer_discriminator(self): + return self.optimizers[1] + + @property + def scheduler_model(self): + return self._schedulers[0] + + @property + def scheduler_discriminator(self): + return self._schedulers[1] \ No newline at end of file diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py new file mode 100644 index 000000000..9c6f6c8ec --- /dev/null +++ b/tests/test_solvers/test_competitive_pinn.py @@ -0,0 +1,418 @@ +import torch +import pytest + +from pina.problem import SpatialProblem, InverseProblem +from pina.operators import laplacian +from pina.geometry import CartesianDomain +from pina import Condition, LabelTensor +from pina.solvers import CompetitivePINN as PINN +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.loss import LpLoss + + +def laplace_equation(input_, output_): + force_term = (torch.sin(input_.extract(['x']) * torch.pi) * + torch.sin(input_.extract(['y']) * torch.pi)) + delta_u = laplacian(output_.extract(['u']), input_) + return delta_u - force_term + + +my_laplace = Equation(laplace_equation) +in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y']) +out_ = LabelTensor(torch.tensor([[0.]]), ['u']) +in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y']) +out2_ = LabelTensor(torch.rand(60, 1), ['u']) + + +class InversePoisson(SpatialProblem, InverseProblem): + ''' + Problem definition for the Poisson equation. + ''' + output_variables = ['u'] + x_min = -2 + x_max = 2 + y_min = -2 + y_max = 2 + data_input = LabelTensor(torch.rand(10, 2), ['x', 'y']) + data_output = LabelTensor(torch.rand(10, 1), ['u']) + spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]}) + # define the ranges for the parameters + unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]}) + + def laplace_equation(input_, output_, params_): + ''' + Laplace equation with a force term. + ''' + force_term = torch.exp( + - 2*(input_.extract(['x']) - params_['mu1'])**2 + - 2*(input_.extract(['y']) - params_['mu2'])**2) + delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y']) + + return delta_u - force_term + + # define the conditions for the loss (boundary conditions, equation, data) + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max], + 'y': y_max}), + equation=FixedValue(0.0, components=['u'])), + 'gamma2': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': y_min + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma3': Condition(location=CartesianDomain( + {'x': x_max, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma4': Condition(location=CartesianDomain( + {'x': x_min, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'D': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': [y_min, y_max] + }), + equation=Equation(laplace_equation)), + 'data': Condition(input_points=data_input.extract(['x', 'y']), + output_points=data_output) + } + + +class Poisson(SpatialProblem): + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]}) + + conditions = { + 'gamma1': Condition( + location=CartesianDomain({'x': [0, 1], 'y': 1}), + equation=FixedValue(0.0)), + 'gamma2': Condition( + location=CartesianDomain({'x': [0, 1], 'y': 0}), + equation=FixedValue(0.0)), + 'gamma3': Condition( + location=CartesianDomain({'x': 1, 'y': [0, 1]}), + equation=FixedValue(0.0)), + 'gamma4': Condition( + location=CartesianDomain({'x': 0, 'y': [0, 1]}), + equation=FixedValue(0.0)), + 'D': Condition( + input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']), + equation=my_laplace), + 'data': Condition( + input_points=in_, + output_points=out_), + 'data2': Condition( + input_points=in2_, + output_points=out2_) + } + + def poisson_sol(self, pts): + return -(torch.sin(pts.extract(['x']) * torch.pi) * + torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2) + + truth_solution = poisson_sol + + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['x']) * torch.pi) * + torch.sin(x.extract(['y']) * torch.pi)) + return LabelTensor(t, ['sin(x)sin(y)']) + + +# make the problem +poisson_problem = Poisson() +model = FeedForward(len(poisson_problem.input_variables), + len(poisson_problem.output_variables)) +model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) +extra_feats = [myFeature()] + + +def test_constructor(): + PINN(problem=poisson_problem, model=model) + PINN(problem=poisson_problem, model=model, discriminator = model) + + +def test_constructor_extra_feats(): + with pytest.raises(TypeError): + model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) + PINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_train_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu') + t = ntrainer.train( + ckpt_path=f'{tmpdir}/lightning_logs/version_0/' + 'checkpoints/epoch=4-step=10.ckpt') + import shutil + shutil.rmtree(tmpdir) + + +def test_train_load(): + tmpdir = "tests/tmp_load" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +def test_train_inverse_problem_cpu(): + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# poisson_problem = InversePoisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] +# n = 100 +# poisson_problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = PINN(problem=poisson_problem, +# model=model, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') +# t = ntrainer.train( +# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') +# import shutil +# shutil.rmtree(tmpdir) + + +def test_train_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +# # TODO fix asap. Basically sampling few variables +# # works only if both variables are in a range. +# # if one is fixed and the other not, this will +# # not work. This test also needs to be fixed and +# # insert in test problem not in test pinn. +# def test_train_cpu_sampling_few_vars(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x']) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y']) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'}) +# trainer.train() + + +# TODO, fix GitHub actions to run also on GPU +# def test_train_gpu(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_gpu(): #TODO fix ASAP +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_2(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_extra_feats(): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + + +# def test_train_2_extra_feats(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_optimizer_kwargs(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3}) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_lr_scheduler(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN( +# problem, +# model, +# lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, +# lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} +# ) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# # def test_train_batch(): +# # pinn = PINN(problem, model, batch_size=6) +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + + +# # def test_train_batch_2(): +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # expected_keys = [[], list(range(0, 50, 3))] +# # param = [0, 3] +# # for i, truth_key in zip(param, expected_keys): +# # pinn = PINN(problem, model, batch_size=6) +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(50, save_loss=i) +# # assert list(pinn.history_loss.keys()) == truth_key + + +# if torch.cuda.is_available(): + +# # def test_gpu_train(): +# # pinn = PINN(problem, model, batch_size=20, device='cuda') +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 100 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + +# def test_gpu_train_nobatch(): +# pinn = PINN(problem, model, batch_size=None, device='cuda') +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 100 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + From 2986efc335a771193cdeaaf296455ec935bf0caa Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 19 Apr 2024 19:33:54 +0200 Subject: [PATCH 2/3] fixing after rebase/ fix loss --- pina/solvers/pinns/competitive_pinn.py | 94 ++++++-------------------- 1 file changed, 20 insertions(+), 74 deletions(-) diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py index 8efbb012a..f225d8912 100644 --- a/pina/solvers/pinns/competitive_pinn.py +++ b/pina/solvers/pinns/competitive_pinn.py @@ -1,4 +1,4 @@ -""" Module for PINN """ +""" Module for CompetitivePINN """ import torch import copy @@ -21,7 +21,7 @@ class CompetitivePINN(PINNInterface): """ TODO - .. warning Condition Data Weight Not Suppoerted + .. warning Does Not Support Extra Features """ def __init__( @@ -76,7 +76,7 @@ def __init__( optimizer_model_kwargs, optimizer_discriminator_kwargs, ], - extra_features=None, # PIGAN doesn't take extra features + extra_features=None, # CompetitivePINN doesn't take extra features loss=loss ) @@ -106,7 +106,6 @@ def __init__( def forward(self, x): """ Forward pass implementation for the PINN solver. - :param LabelTensor x: Input tensor for the PINN solver. It expects a tensor :math:`N \times D`, where :math:`N` the number of points in the mesh, :math:`D` the dimension of the problem, @@ -116,58 +115,6 @@ def forward(self, x): return self.neural_net(x) - def training_step(self, batch, _): - """ - PINN solver training step. - - :param batch: The batch element in the dataloader. - :type batch: tuple - :param batch_idx: The batch index. - :type batch_idx: int - :return: The sum of the loss functions. - :rtype: LabelTensor - """ - - condition_idx = batch["condition"] - - for condition_id in range(condition_idx.min(), condition_idx.max() + 1): - - condition_name = self._dataloader.condition_names[condition_id] - condition = self.problem.conditions[condition_name] - pts = batch["pts"] - - if len(batch) == 2: - samples = pts[condition_idx == condition_id] - self.loss_phys(samples=samples, - equation=condition.equation, - condition_name=condition_name) - elif len(batch) == 3: - samples = pts[condition_idx == condition_id] - ground_truth = batch["output"][condition_idx == condition_id] - self._loss_data(samples, ground_truth, condition_name) - else: - raise ValueError("Batch size not supported") - - # clamp unknown parameters in InverseProblem (if needed), otherwise - # it returns None - self._clamp_params() - return - - def loss_phys(self, samples, equation): - """ - Computes the physics loss for the PINN solver based on given - samples and equation. - - :param LabelTensor samples: The samples to evaluate the physics loss. - :param EquationInterface equation: The governing equation - representing the physics. - :return: The physics residual calculated based on given - samples and equation. - :rtype: LabelTensor - """ - return self.compute_residual(samples=samples, equation=equation) - - def configure_optimizers(self): """ Optimizer configuration for the Competitive PINN solver. @@ -189,7 +136,7 @@ def configure_optimizers(self): return self.optimizers, self._schedulers - def _loss_phys(self, samples, equation, condition_name): + def loss_phys(self, samples, equation): # train one step of discriminator discriminator_bets = self.discriminator(samples.clone()) self._train_discriminator(samples, equation, discriminator_bets) @@ -199,7 +146,9 @@ def _loss_phys(self, samples, equation, condition_name): samples = samples.detach() samples.requires_grad = True # train one step of the model - self._train_model(samples, equation, condition_name, discriminator_bets) + loss_val = self._train_model(samples, equation, discriminator_bets) + self.store_log(loss_value=float(loss_val)) + return loss_val def _train_discriminator(self, samples, equation, discriminator_bets): """ @@ -215,8 +164,8 @@ def _train_discriminator(self, samples, equation, discriminator_bets): self.optimizer_discriminator.zero_grad() # compute residual, we detach because the weights of the generator # model are fixed - residual = self.loss_phys(samples=samples, - equation=equation).detach() + residual = self.compute_residual(samples=samples, + equation=equation).detach() # compute competitive residual, the minus is because we maximise competitive_residual = residual * discriminator_bets loss_val = - self.loss( @@ -228,30 +177,28 @@ def _train_discriminator(self, samples, equation, discriminator_bets): self.optimizer_discriminator.step() return - def _train_model(self, samples, equation, - condition_name, discriminator_bets): + def _train_model(self, samples, equation, discriminator_bets): """ Trains the model network of the Competitive PINN. :param LabelTensor samples: Input samples to evaluate the physics loss. :param EquationInterface equation: The governing equation representing the physics. - :param str condition_name: The condition name for tracking purposes. - :param Tensor discriminator_bets: Predictions made by the discriminator + :param Tensor discriminator_bets: Predictions made by the discriminator. network. + :return: The computed data loss. + :rtype: torch.Tensor """ # manual optimization self.optimizer_model.zero_grad() # compute residual (detached for discriminator) and log - residual = self.loss_phys(samples=samples, equation=equation) + residual = self.compute_residual(samples=samples, equation=equation) # store logging with torch.no_grad(): loss_residual = self.loss( torch.zeros_like(residual), residual - ).as_subclass(torch.Tensor) - self.store_log(name=condition_name+f'_loss', - loss_val=float(loss_residual)) + ) # compute competitive residual, discriminator_bets are detached becase # we optimize only the generator model competitive_residual = residual * discriminator_bets.detach() @@ -262,9 +209,9 @@ def _train_model(self, samples, equation, # backprop loss_val.backward() self.optimizer_model.step() - return + return loss_residual - def _loss_data(self, input_tensor, output_tensor, condition_name): + def loss_data(self, input_tensor, output_tensor): """ Computes the data loss for the PINN solver based on input, output, and condition name. This function is a wrapper of the function @@ -273,16 +220,15 @@ def _loss_data(self, input_tensor, output_tensor, condition_name): :param LabelTensor input_tensor: The input to the neural networks. :param LabelTensor output_tensor: The true solution to compare the network solution. - :param str condition_name: The condition name for tracking purposes. :return: The computed data loss. :rtype: torch.Tensor """ self.optimizer_model.zero_grad() - loss_val = super()._loss_data(input_tensor, - output_tensor, condition_name) + loss_val = super().loss_data( + input_tensor, output_tensor).as_subclass(torch.Tensor) loss_val.backward() self.optimizer_model.step() - return + return loss_val @property From 640d6438ebac55b690aa198c61d301185d3cc170 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Mon, 22 Apr 2024 17:41:09 +0200 Subject: [PATCH 3/3] fixing final issues --- pina/solvers/pinns/competitive_pinn.py | 86 +++++++++++++++++++++++--- 1 file changed, 79 insertions(+), 7 deletions(-) diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py index f225d8912..df8319823 100644 --- a/pina/solvers/pinns/competitive_pinn.py +++ b/pina/solvers/pinns/competitive_pinn.py @@ -19,9 +19,20 @@ class CompetitivePINN(PINNInterface): """ - TODO + CompetitivePINN solver class. This class implements Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. - .. warning Does Not Support Extra Features + .. seealso:: + + **Original reference**: Zeng, Qi, et al. + "Competitive physics informed networks." International Conference on + Learning Representations, ICLR 2022 + OpenReview Preprint `_. + + .. warning:: + This solver does not currently support the possibility to pass + ``extra_feature``. """ def __init__( @@ -102,7 +113,25 @@ def __init__( self._model = self.models[0] self._discriminator = self.models[1] + def on_train_batch_end(self,outputs, batch, batch_idx): + """ + This method is called at the end of each training batch, and ovverides + the PytorchLightining implementation for logging the checkpoints. + :param outputs: The output from the model for the current batch. + :type outputs: Any + :param batch: The current batch of data. + :type batch: Any + :param batch_idx: The index of the current batch. + :type batch_idx: int + :return: Whatever is returned by the parent + method ``on_train_batch_end``. + :rtype: Any + """ + # increase by one the counter of optimization to save loggers + self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += 1 + return super().on_train_batch_end(outputs, batch, batch_idx) + def forward(self, x): """ Forward pass implementation for the PINN solver. @@ -114,7 +143,6 @@ def forward(self, x): """ return self.neural_net(x) - def configure_optimizers(self): """ Optimizer configuration for the Competitive PINN solver. @@ -135,8 +163,17 @@ def configure_optimizers(self): ) return self.optimizers, self._schedulers - def loss_phys(self, samples, equation): + """ + Computes the physics loss for the PINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor""" # train one step of discriminator discriminator_bets = self.discriminator(samples.clone()) self._train_discriminator(samples, equation, discriminator_bets) @@ -173,7 +210,7 @@ def _train_discriminator(self, samples, equation, discriminator_bets): competitive_residual ).as_subclass(torch.Tensor) # backprop - loss_val.backward() + self.manual_backward(loss_val) self.optimizer_discriminator.step() return @@ -207,7 +244,7 @@ def _train_model(self, samples, equation, discriminator_bets): competitive_residual ).as_subclass(torch.Tensor) # backprop - loss_val.backward() + self.manual_backward(loss_val) self.optimizer_model.step() return loss_residual @@ -230,27 +267,62 @@ def loss_data(self, input_tensor, output_tensor): self.optimizer_model.step() return loss_val - @property def neural_net(self): + """ + Returns the neural network model. + + :return: The neural network model. + :rtype: torch.nn.Module + """ return self._model @property def discriminator(self): + """ + Returns the discriminator model (if applicable). + + :return: The discriminator model. + :rtype: torch.nn.Module + """ return self._discriminator @property def optimizer_model(self): + """ + Returns the optimizer associated with the neural network model. + + :return: The optimizer for the neural network model. + :rtype: torch.optim.Optimizer + """ return self.optimizers[0] @property def optimizer_discriminator(self): + """ + Returns the optimizer associated with the discriminator (if applicable). + + :return: The optimizer for the discriminator. + :rtype: torch.optim.Optimizer + """ return self.optimizers[1] @property def scheduler_model(self): + """ + Returns the scheduler associated with the neural network model. + + :return: The scheduler for the neural network model. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ return self._schedulers[0] @property def scheduler_discriminator(self): + """ + Returns the scheduler associated with the discriminator (if applicable). + + :return: The scheduler for the discriminator. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ return self._schedulers[1] \ No newline at end of file