From cbb43a53926c41502d880afa51703d6669346c19 Mon Sep 17 00:00:00 2001 From: ndem0 Date: Fri, 9 Feb 2024 11:25:00 +0000 Subject: [PATCH] :art: Format Python code with psf/black --- pina/__init__.py | 9 +- pina/adaptive_functions/adaptive_cos.py | 20 +- pina/adaptive_functions/adaptive_exp.py | 24 +- pina/adaptive_functions/adaptive_linear.py | 13 +- pina/adaptive_functions/adaptive_relu.py | 14 +- pina/adaptive_functions/adaptive_sin.py | 22 +- pina/adaptive_functions/adaptive_softplus.py | 14 +- pina/adaptive_functions/adaptive_square.py | 14 +- pina/adaptive_functions/adaptive_tanh.py | 28 ++- pina/callbacks/__init__.py | 2 +- .../callbacks/adaptive_refinment_callbacks.py | 29 +-- pina/callbacks/optimizer_callbacks.py | 27 ++- pina/callbacks/processing_callbacks.py | 9 +- pina/condition.py | 38 +-- pina/dataset.py | 96 ++++---- pina/equation/__init__.py | 12 +- pina/equation/equation.py | 11 +- pina/equation/equation_factory.py | 1 + pina/equation/equation_interface.py | 1 + pina/equation/system_equation.py | 20 +- pina/geometry/__init__.py | 11 +- pina/geometry/cartesian.py | 44 ++-- pina/geometry/difference_domain.py | 11 +- pina/geometry/ellipsoid.py | 38 +-- pina/geometry/exclusion_domain.py | 11 +- pina/geometry/intersection_domain.py | 11 +- pina/geometry/operation_interface.py | 11 +- pina/geometry/simplex.py | 34 +-- pina/geometry/union_domain.py | 10 +- pina/label_tensor.py | 73 +++--- pina/loss.py | 18 +- pina/meta.py | 15 +- pina/model/__init__.py | 12 +- pina/model/deeponet.py | 90 ++++--- pina/model/feed_forward.py | 92 ++++---- pina/model/fno.py | 63 ++--- pina/model/layers/__init__.py | 26 ++- pina/model/layers/convolution.py | 46 ++-- pina/model/layers/convolution_2d.py | 135 ++++++----- pina/model/layers/fourier.py | 73 +++--- pina/model/layers/integral.py | 4 +- pina/model/layers/pod.py | 40 ++-- pina/model/layers/residual.py | 22 +- pina/model/layers/spectral.py | 220 +++++++++++------- pina/model/layers/stride.py | 3 +- pina/model/layers/utils_convolution.py | 9 +- pina/model/multi_feed_forward.py | 3 +- pina/model/network.py | 12 +- pina/operators.py | 65 +++--- pina/plotter.py | 183 ++++++++------- pina/problem/__init__.py | 10 +- pina/problem/abstract_problem.py | 78 ++++--- pina/problem/inverse_problem.py | 2 +- pina/problem/parametric_problem.py | 1 + pina/problem/spatial_problem.py | 1 + pina/problem/timedep_problem.py | 1 + pina/solvers/__init__.py | 8 +- pina/solvers/garom.py | 115 ++++++--- pina/solvers/pinn.py | 89 ++++--- pina/solvers/solver.py | 64 ++--- pina/solvers/supervised.py | 62 ++--- pina/trainer.py | 29 ++- pina/utils.py | 22 +- pina/writer.py | 7 +- 64 files changed, 1323 insertions(+), 955 deletions(-) diff --git a/pina/__init__.py b/pina/__init__.py index 711161b..730b2ea 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -1,6 +1,11 @@ __all__ = [ - 'PINN', 'Trainer', 'LabelTensor', 'Plotter', 'Condition', - 'SamplePointDataset', 'SamplePointLoader' + "PINN", + "Trainer", + "LabelTensor", + "Plotter", + "Condition", + "SamplePointDataset", + "SamplePointLoader", ] from .meta import * diff --git a/pina/adaptive_functions/adaptive_cos.py b/pina/adaptive_functions/adaptive_cos.py index 10106aa..860a9e9 100644 --- a/pina/adaptive_functions/adaptive_cos.py +++ b/pina/adaptive_functions/adaptive_cos.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveCos(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,26 +18,28 @@ class AdaptiveCos(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self, alpha=None): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveCos, self).__init__() - #self.in_features = in_features + # self.in_features = in_features # initialize alpha if alpha == None: self.alpha = Parameter( - torch.tensor(1.0)) # create a tensor out of alpha + torch.tensor(1.0) + ) # create a tensor out of alpha else: self.alpha = Parameter( - torch.tensor(alpha)) # create a tensor out of alpha + torch.tensor(alpha) + ) # create a tensor out of alpha self.alpha.requiresGrad = True # set requiresGrad to true! self.scale = Parameter(torch.tensor(1.0)) @@ -47,8 +49,8 @@ class AdaptiveCos(torch.nn.Module): self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' + """ return self.scale * (torch.cos(self.alpha * x + self.translate)) diff --git a/pina/adaptive_functions/adaptive_exp.py b/pina/adaptive_functions/adaptive_exp.py index c65406f..b6a1184 100644 --- a/pina/adaptive_functions/adaptive_exp.py +++ b/pina/adaptive_functions/adaptive_exp.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveExp(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,36 +18,36 @@ class AdaptiveExp(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveExp, self).__init__() self.scale = Parameter( - torch.normal(torch.tensor(1.0), - torch.tensor(0.1))) # create a tensor out of alpha + torch.normal(torch.tensor(1.0), torch.tensor(0.1)) + ) # create a tensor out of alpha self.scale.requiresGrad = True # set requiresGrad to true! self.alpha = Parameter( - torch.normal(torch.tensor(1.0), - torch.tensor(0.1))) # create a tensor out of alpha + torch.normal(torch.tensor(1.0), torch.tensor(0.1)) + ) # create a tensor out of alpha self.alpha.requiresGrad = True # set requiresGrad to true! self.translate = Parameter( - torch.normal(torch.tensor(0.0), - torch.tensor(0.1))) # create a tensor out of alpha + torch.normal(torch.tensor(0.0), torch.tensor(0.1)) + ) # create a tensor out of alpha self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' + """ return self.scale * (x + self.translate) diff --git a/pina/adaptive_functions/adaptive_linear.py b/pina/adaptive_functions/adaptive_linear.py index 42968c9..66688ef 100644 --- a/pina/adaptive_functions/adaptive_linear.py +++ b/pina/adaptive_functions/adaptive_linear.py @@ -1,10 +1,11 @@ """ Implementation of adaptive linear layer. """ + import torch from torch.nn.parameter import Parameter class AdaptiveLinear(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -19,16 +20,16 @@ class AdaptiveLinear(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveLinear, self).__init__() self.scale = Parameter(torch.tensor(1.0)) @@ -38,8 +39,8 @@ class AdaptiveLinear(torch.nn.Module): self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' + """ return self.scale * (x + self.translate) diff --git a/pina/adaptive_functions/adaptive_relu.py b/pina/adaptive_functions/adaptive_relu.py index 0061462..8885434 100644 --- a/pina/adaptive_functions/adaptive_relu.py +++ b/pina/adaptive_functions/adaptive_relu.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveReLU(torch.nn.Module, Parameter): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,16 +18,16 @@ class AdaptiveReLU(torch.nn.Module, Parameter): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveReLU, self).__init__() self.scale = Parameter(torch.rand(1)) @@ -37,9 +37,9 @@ class AdaptiveReLU(torch.nn.Module, Parameter): self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' - #x += self.translate + """ + # x += self.translate return torch.relu(x + self.translate) * self.scale diff --git a/pina/adaptive_functions/adaptive_sin.py b/pina/adaptive_functions/adaptive_sin.py index 26a6ef3..b5da89a 100644 --- a/pina/adaptive_functions/adaptive_sin.py +++ b/pina/adaptive_functions/adaptive_sin.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveSin(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,35 +18,37 @@ class AdaptiveSin(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self, alpha=None): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveSin, self).__init__() # initialize alpha self.alpha = Parameter( - torch.normal(torch.tensor(1.0), - torch.tensor(0.1))) # create a tensor out of alpha + torch.normal(torch.tensor(1.0), torch.tensor(0.1)) + ) # create a tensor out of alpha self.alpha.requiresGrad = True # set requiresGrad to true! self.scale = Parameter( - torch.normal(torch.tensor(1.0), torch.tensor(0.1))) + torch.normal(torch.tensor(1.0), torch.tensor(0.1)) + ) self.scale.requiresGrad = True # set requiresGrad to true! self.translate = Parameter( - torch.normal(torch.tensor(0.0), torch.tensor(0.1))) + torch.normal(torch.tensor(0.0), torch.tensor(0.1)) + ) self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' + """ return self.scale * (torch.sin(self.alpha * x + self.translate)) diff --git a/pina/adaptive_functions/adaptive_softplus.py b/pina/adaptive_functions/adaptive_softplus.py index d306832..d62e2e5 100644 --- a/pina/adaptive_functions/adaptive_softplus.py +++ b/pina/adaptive_functions/adaptive_softplus.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveSoftplus(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,16 +18,16 @@ class AdaptiveSoftplus(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super().__init__() self.soft = torch.nn.Softplus() @@ -36,9 +36,9 @@ class AdaptiveSoftplus(torch.nn.Module): self.scale.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' - #x += self.translate + """ + # x += self.translate return self.soft(x) * self.scale diff --git a/pina/adaptive_functions/adaptive_square.py b/pina/adaptive_functions/adaptive_square.py index 9b341a0..f84c6aa 100644 --- a/pina/adaptive_functions/adaptive_square.py +++ b/pina/adaptive_functions/adaptive_square.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveSquare(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,16 +18,16 @@ class AdaptiveSquare(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self, alpha=None): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveSquare, self).__init__() self.scale = Parameter(torch.tensor(1.0)) @@ -37,8 +37,8 @@ class AdaptiveSquare(torch.nn.Module): self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' - return self.scale * (x + self.translate)**2 + """ + return self.scale * (x + self.translate) ** 2 diff --git a/pina/adaptive_functions/adaptive_tanh.py b/pina/adaptive_functions/adaptive_tanh.py index 3a2c719..40cad2d 100644 --- a/pina/adaptive_functions/adaptive_tanh.py +++ b/pina/adaptive_functions/adaptive_tanh.py @@ -3,7 +3,7 @@ from torch.nn.parameter import Parameter class AdaptiveTanh(torch.nn.Module): - ''' + """ Implementation of soft exponential activation. Shape: - Input: (N, *) where * means, any number of additional @@ -18,26 +18,28 @@ class AdaptiveTanh(torch.nn.Module): >>> a1 = soft_exponential(256) >>> x = torch.randn(256) >>> x = a1(x) - ''' + """ def __init__(self, alpha=None): - ''' + """ Initialization. INPUT: - in_features: shape of the input - aplha: trainable parameter aplha is initialized with zero value by default - ''' + """ super(AdaptiveTanh, self).__init__() - #self.in_features = in_features + # self.in_features = in_features # initialize alpha if alpha == None: self.alpha = Parameter( - torch.tensor(1.0)) # create a tensor out of alpha + torch.tensor(1.0) + ) # create a tensor out of alpha else: self.alpha = Parameter( - torch.tensor(alpha)) # create a tensor out of alpha + torch.tensor(alpha) + ) # create a tensor out of alpha self.alpha.requiresGrad = True # set requiresGrad to true! @@ -48,11 +50,13 @@ class AdaptiveTanh(torch.nn.Module): self.translate.requiresGrad = True # set requiresGrad to true! def forward(self, x): - ''' + """ Forward pass of the function. Applies the function to the input elementwise. - ''' + """ x += self.translate - return self.scale * (torch.exp(self.alpha * x) - torch.exp( - -self.alpha * x)) / (torch.exp(self.alpha * x) + - torch.exp(-self.alpha * x)) + return ( + self.scale + * (torch.exp(self.alpha * x) - torch.exp(-self.alpha * x)) + / (torch.exp(self.alpha * x) + torch.exp(-self.alpha * x)) + ) diff --git a/pina/callbacks/__init__.py b/pina/callbacks/__init__.py index c9ba520..9698136 100644 --- a/pina/callbacks/__init__.py +++ b/pina/callbacks/__init__.py @@ -1,4 +1,4 @@ -__all__ = ['SwitchOptimizer', 'R3Refinement', 'MetricTracker'] +__all__ = ["SwitchOptimizer", "R3Refinement", "MetricTracker"] from .optimizer_callbacks import SwitchOptimizer from .adaptive_refinment_callbacks import R3Refinement diff --git a/pina/callbacks/adaptive_refinment_callbacks.py b/pina/callbacks/adaptive_refinment_callbacks.py index 5ec149b..b5e2b70 100644 --- a/pina/callbacks/adaptive_refinment_callbacks.py +++ b/pina/callbacks/adaptive_refinment_callbacks.py @@ -1,4 +1,4 @@ -'''PINA Callbacks Implementations''' +"""PINA Callbacks Implementations""" # from lightning.pytorch.callbacks import Callback from pytorch_lightning.callbacks import Callback @@ -8,18 +8,17 @@ from ..utils import check_consistency class R3Refinement(Callback): - def __init__(self, sample_every): """ PINA Implementation of an R3 Refinement Callback. 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 + 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 + 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 `_ @@ -79,7 +78,7 @@ class R3Refinement(Callback): # !!!!!! From now everything is performed on CPU !!!!!! # average loss - avg = (tot_loss.mean()).to('cpu') + avg = (tot_loss.mean()).to("cpu") # points to keep old_pts = {} @@ -90,25 +89,29 @@ class R3Refinement(Callback): pts = pts.cpu().detach() residuals = res_loss[location].cpu() mask = (residuals > avg).flatten() - if any(mask): # if there are residuals greater than averge we append them - pts = pts[mask] # TODO masking remove labels + if any( + mask + ): # if there are residuals greater than averge we append them + pts = pts[mask] # TODO masking remove labels pts.labels = labels old_pts[location] = pts tot_points += len(pts) # extract new points to sample uniformally for each location n_points = (self._tot_pop_numb - tot_points) // len( - self._sampling_locations) + self._sampling_locations + ) remainder = (self._tot_pop_numb - tot_points) % len( - self._sampling_locations) + self._sampling_locations + ) n_uniform_points = [n_points] * len(self._sampling_locations) n_uniform_points[-1] += remainder # sample new points for numb_pts, loc in zip(n_uniform_points, self._sampling_locations): - trainer._model.problem.discretise_domain(numb_pts, - 'random', - locations=[loc]) + trainer._model.problem.discretise_domain( + numb_pts, "random", locations=[loc] + ) # adding previous population points trainer._model.problem.add_points(old_pts) @@ -133,7 +136,7 @@ class R3Refinement(Callback): locations = [] for condition_name in problem.conditions: condition = problem.conditions[condition_name] - if hasattr(condition, 'location'): + if hasattr(condition, "location"): locations.append(condition_name) self._sampling_locations = locations diff --git a/pina/callbacks/optimizer_callbacks.py b/pina/callbacks/optimizer_callbacks.py index 276983e..c11db88 100644 --- a/pina/callbacks/optimizer_callbacks.py +++ b/pina/callbacks/optimizer_callbacks.py @@ -1,4 +1,4 @@ -'''PINA Callbacks Implementations''' +"""PINA Callbacks Implementations""" from pytorch_lightning.callbacks import Callback import torch @@ -14,7 +14,7 @@ class SwitchOptimizer(Callback): This callback allows for switching between different optimizers during training, enabling the exploration of multiple optimization strategies without the need to stop training. - :param new_optimizers: The model optimizers to switch to. Can be a single + :param new_optimizers: The model optimizers to switch to. Can be a single :class:`torch.optim.Optimizer` or a list of them for multiple model solvers. :type new_optimizers: torch.optim.Optimizer | list :param new_optimizers_kwargs: The keyword arguments for the new optimizers. Can be a single dictionary @@ -23,7 +23,7 @@ class SwitchOptimizer(Callback): :param epoch_switch: The epoch at which to switch to the new optimizer. :type epoch_switch: int - :raises ValueError: If `epoch_switch` is less than 1 or if there is a mismatch in the number of + :raises ValueError: If `epoch_switch` is less than 1 or if there is a mismatch in the number of optimizers and their corresponding keyword argument dictionaries. Example: @@ -39,7 +39,7 @@ class SwitchOptimizer(Callback): check_consistency(epoch_switch, int) if epoch_switch < 1: - raise ValueError('epoch_switch must be greater than one.') + raise ValueError("epoch_switch must be greater than one.") if not isinstance(new_optimizers, list): new_optimizers = [new_optimizers] @@ -48,10 +48,12 @@ class SwitchOptimizer(Callback): len_optimizer_kwargs = len(new_optimizers_kwargs) if len_optimizer_kwargs != len_optimizer: - raise ValueError('You must define one dictionary of keyword' - ' arguments for each optimizers.' - f' Got {len_optimizer} optimizers, and' - f' {len_optimizer_kwargs} dicitionaries') + raise ValueError( + "You must define one dictionary of keyword" + " arguments for each optimizers." + f" Got {len_optimizer} optimizers, and" + f" {len_optimizer_kwargs} dicitionaries" + ) # save new optimizers self._new_optimizers = new_optimizers @@ -72,9 +74,12 @@ class SwitchOptimizer(Callback): if trainer.current_epoch == self._epoch_switch: optims = [] for idx, (optim, optim_kwargs) in enumerate( - zip(self._new_optimizers, self._new_optimizers_kwargs)): + zip(self._new_optimizers, self._new_optimizers_kwargs) + ): optims.append( - optim(trainer._model.models[idx].parameters(), - **optim_kwargs)) + optim( + trainer._model.models[idx].parameters(), **optim_kwargs + ) + ) trainer.optimizers = optims diff --git a/pina/callbacks/processing_callbacks.py b/pina/callbacks/processing_callbacks.py index 791d540..0a7359c 100644 --- a/pina/callbacks/processing_callbacks.py +++ b/pina/callbacks/processing_callbacks.py @@ -1,4 +1,4 @@ -'''PINA Callbacks Implementations''' +"""PINA Callbacks Implementations""" from pytorch_lightning.callbacks import Callback import torch @@ -6,7 +6,7 @@ import copy class MetricTracker(Callback): - + def __init__(self): """ PINA Implementation of a Lightning Callback for Metric Tracking. @@ -39,8 +39,9 @@ class MetricTracker(Callback): :return: None :rtype: None """ - self._collection.append(copy.deepcopy( - trainer.logged_metrics)) # track them + self._collection.append( + copy.deepcopy(trainer.logged_metrics) + ) # track them @property def metrics(self): diff --git a/pina/condition.py b/pina/condition.py index c03678e..5125fe0 100644 --- a/pina/condition.py +++ b/pina/condition.py @@ -1,4 +1,5 @@ """ Condition module. """ + from .label_tensor import LabelTensor from .geometry import Location from .equation.equation import Equation @@ -51,7 +52,11 @@ class Condition: """ __slots__ = [ - 'input_points', 'output_points', 'location', 'equation', 'data_weight' + "input_points", + "output_points", + "location", + "equation", + "data_weight", ] def _dictvalue_isinstance(self, dict_, key_, class_): @@ -65,27 +70,28 @@ class Condition: """ Constructor for the `Condition` class. """ - self.data_weight = kwargs.pop('data_weight', 1.0) + self.data_weight = kwargs.pop("data_weight", 1.0) if len(args) != 0: raise ValueError( - f'Condition takes only the following keyword arguments: {Condition.__slots__}.' + f"Condition takes only the following keyword arguments: {Condition.__slots__}." ) - if (sorted(kwargs.keys()) != sorted(['input_points', 'output_points']) - and sorted(kwargs.keys()) != sorted(['location', 'equation']) - and sorted(kwargs.keys()) != sorted( - ['input_points', 'equation'])): - raise ValueError(f'Invalid keyword arguments {kwargs.keys()}.') + if ( + sorted(kwargs.keys()) != sorted(["input_points", "output_points"]) + and sorted(kwargs.keys()) != sorted(["location", "equation"]) + and sorted(kwargs.keys()) != sorted(["input_points", "equation"]) + ): + raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.") - if not self._dictvalue_isinstance(kwargs, 'input_points', LabelTensor): - raise TypeError('`input_points` must be a torch.Tensor.') - if not self._dictvalue_isinstance(kwargs, 'output_points', LabelTensor): - raise TypeError('`output_points` must be a torch.Tensor.') - if not self._dictvalue_isinstance(kwargs, 'location', Location): - raise TypeError('`location` must be a Location.') - if not self._dictvalue_isinstance(kwargs, 'equation', Equation): - raise TypeError('`equation` must be a Equation.') + if not self._dictvalue_isinstance(kwargs, "input_points", LabelTensor): + raise TypeError("`input_points` must be a torch.Tensor.") + if not self._dictvalue_isinstance(kwargs, "output_points", LabelTensor): + raise TypeError("`output_points` must be a torch.Tensor.") + if not self._dictvalue_isinstance(kwargs, "location", Location): + raise TypeError("`location` must be a Location.") + if not self._dictvalue_isinstance(kwargs, "equation", Equation): + raise TypeError("`equation` must be a Equation.") for key, value in kwargs.items(): setattr(self, key, value) diff --git a/pina/dataset.py b/pina/dataset.py index 38b2200..c6a8d29 100644 --- a/pina/dataset.py +++ b/pina/dataset.py @@ -17,38 +17,41 @@ class SamplePointDataset(Dataset): self.condition_names = [] for name, condition in problem.conditions.items(): - if not hasattr(condition, 'output_points'): + if not hasattr(condition, "output_points"): pts_list.append(problem.input_pts[name]) self.condition_names.append(name) self.pts = LabelTensor.vstack(pts_list) if self.pts != []: - self.condition_indeces = torch.cat([ - torch.tensor([i]*len(pts_list[i])) - for i in range(len(self.condition_names)) - ], dim=0) - else: # if there are no sample points + self.condition_indeces = torch.cat( + [ + torch.tensor([i] * len(pts_list[i])) + for i in range(len(self.condition_names)) + ], + dim=0, + ) + else: # if there are no sample points self.condition_indeces = torch.tensor([]) self.pts = torch.tensor([]) self.pts = self.pts.to(device) self.condition_indeces = self.condition_indeces.to(device) - + def __len__(self): return self.pts.shape[0] - + class DataPointDataset(Dataset): def __init__(self, problem, device) -> None: super().__init__() input_list = [] - output_list = [] + output_list = [] self.condition_names = [] for name, condition in problem.conditions.items(): - if hasattr(condition, 'output_points'): + if hasattr(condition, "output_points"): input_list.append(problem.conditions[name].input_points) output_list.append(problem.conditions[name].output_points) self.condition_names.append(name) @@ -57,11 +60,14 @@ class DataPointDataset(Dataset): self.output_pts = LabelTensor.vstack(output_list) if self.input_pts != []: - self.condition_indeces = torch.cat([ - torch.tensor([i]*len(input_list[i])) - for i in range(len(self.condition_names)) - ], dim=0) - else: # if there are no data points + self.condition_indeces = torch.cat( + [ + torch.tensor([i] * len(input_list[i])) + for i in range(len(self.condition_names)) + ], + dim=0, + ) + else: # if there are no data points self.condition_indeces = torch.tensor([]) self.input_pts = torch.tensor([]) self.output_pts = torch.tensor([]) @@ -83,7 +89,9 @@ class SamplePointLoader: :vartype condition_names: list[str] """ - def __init__(self, sample_dataset, data_dataset, batch_size=None, shuffle=True) -> None: + def __init__( + self, sample_dataset, data_dataset, batch_size=None, shuffle=True + ) -> None: """ Constructor. @@ -94,9 +102,13 @@ class SamplePointLoader: Default is ``True``. """ if not isinstance(sample_dataset, SamplePointDataset): - raise TypeError(f'Expected SamplePointDataset, got {type(sample_dataset)}') + raise TypeError( + f"Expected SamplePointDataset, got {type(sample_dataset)}" + ) if not isinstance(data_dataset, DataPointDataset): - raise TypeError(f'Expected DataPointDataset, got {type(data_dataset)}') + raise TypeError( + f"Expected DataPointDataset, got {type(data_dataset)}" + ) self.n_data_conditions = len(data_dataset.condition_names) self.n_phys_conditions = len(sample_dataset.condition_names) @@ -106,25 +118,21 @@ class SamplePointLoader: self._prepare_data_dataset(data_dataset, batch_size, shuffle) self.condition_names = ( - sample_dataset.condition_names + data_dataset.condition_names) + sample_dataset.condition_names + data_dataset.condition_names + ) self.batch_list = [] for i in range(len(self.batch_sample_pts)): - self.batch_list.append( - ('sample', i) - ) + self.batch_list.append(("sample", i)) for i in range(len(self.batch_input_pts)): - self.batch_list.append( - ('data', i) - ) + self.batch_list.append(("data", i)) if shuffle: - self.random_idx = torch.randperm(len(self.batch_list)) + self.random_idx = torch.randperm(len(self.batch_list)) else: self.random_idx = torch.arange(len(self.batch_list)) - def _prepare_data_dataset(self, dataset, batch_size, shuffle): """ Prepare the dataset for data points. @@ -157,17 +165,18 @@ class SamplePointLoader: self.output_pts = dataset.output_pts[idx] self.tensor_conditions = dataset.condition_indeces[idx] - self.batch_input_pts = torch.tensor_split( - dataset.input_pts, batch_num) + self.batch_input_pts = torch.tensor_split(dataset.input_pts, batch_num) self.batch_output_pts = torch.tensor_split( - dataset.output_pts, batch_num) + dataset.output_pts, batch_num + ) for i in range(len(self.batch_input_pts)): self.batch_input_pts[i].labels = input_labels self.batch_output_pts[i].labels = output_labels - + self.batch_data_conditions = torch.tensor_split( - self.tensor_conditions, batch_num) + self.tensor_conditions, batch_num + ) def _prepare_sample_dataset(self, dataset, batch_size, shuffle): """ @@ -190,7 +199,7 @@ class SamplePointLoader: batch_num = len(dataset) // batch_size if len(dataset) % batch_size != 0: batch_num += 1 - + self.tensor_pts = dataset.pts self.tensor_conditions = dataset.condition_indeces @@ -198,13 +207,14 @@ class SamplePointLoader: # idx = torch.randperm(self.tensor_pts.shape[0]) # self.tensor_pts = self.tensor_pts[idx] # self.tensor_conditions = self.tensor_conditions[idx] - + self.batch_sample_pts = torch.tensor_split(self.tensor_pts, batch_num) for i in range(len(self.batch_sample_pts)): self.batch_sample_pts[i].labels = dataset.pts.labels self.batch_sample_conditions = torch.tensor_split( - self.tensor_conditions, batch_num) + self.tensor_conditions, batch_num + ) def __iter__(self): """ @@ -222,20 +232,20 @@ class SamplePointLoader: :return: An iterator over the points. :rtype: iter """ - #for i in self.random_idx: + # for i in self.random_idx: for i in range(len(self.batch_list)): type_, idx_ = self.batch_list[i] - if type_ == 'sample': + if type_ == "sample": d = { - 'pts': self.batch_sample_pts[idx_].requires_grad_(True), - 'condition': self.batch_sample_conditions[idx_], + "pts": self.batch_sample_pts[idx_].requires_grad_(True), + "condition": self.batch_sample_conditions[idx_], } else: d = { - 'pts': self.batch_input_pts[idx_].requires_grad_(True), - 'output': self.batch_output_pts[idx_], - 'condition': self.batch_data_conditions[idx_], + "pts": self.batch_input_pts[idx_].requires_grad_(True), + "output": self.batch_output_pts[idx_], + "condition": self.batch_data_conditions[idx_], } yield d @@ -246,4 +256,4 @@ class SamplePointLoader: :return: The number of batches. :rtype: int """ - return len(self.batch_list) \ No newline at end of file + return len(self.batch_list) diff --git a/pina/equation/__init__.py b/pina/equation/__init__.py index 653aa18..d9961b4 100644 --- a/pina/equation/__init__.py +++ b/pina/equation/__init__.py @@ -1,10 +1,10 @@ __all__ = [ - 'SystemEquation', - 'Equation', - 'FixedValue', - 'FixedGradient', - 'FixedFlux', - 'Laplace', + "SystemEquation", + "Equation", + "FixedValue", + "FixedGradient", + "FixedFlux", + "Laplace", ] from .equation import Equation diff --git a/pina/equation/equation.py b/pina/equation/equation.py index 1e34ebc..3a8f4b1 100644 --- a/pina/equation/equation.py +++ b/pina/equation/equation.py @@ -1,4 +1,5 @@ """ Module for Equation. """ + from .equation_interface import EquationInterface @@ -15,12 +16,14 @@ class Equation(EquationInterface): :type equation: Callable """ if not callable(equation): - raise ValueError('equation must be a callable function.' - 'Expected a callable function, got ' - f'{equation}') + raise ValueError( + "equation must be a callable function." + "Expected a callable function, got " + f"{equation}" + ) self.__equation = equation - def residual(self, input_, output_, params_ = None): + def residual(self, input_, output_, params_=None): """ Residual computation of the equation. diff --git a/pina/equation/equation_factory.py b/pina/equation/equation_factory.py index 4edbf53..5921b1f 100644 --- a/pina/equation/equation_factory.py +++ b/pina/equation/equation_factory.py @@ -1,4 +1,5 @@ """ Module """ + from .equation import Equation from ..operators import grad, div, laplacian diff --git a/pina/equation/equation_interface.py b/pina/equation/equation_interface.py index 5e5ec90..c64c180 100644 --- a/pina/equation/equation_interface.py +++ b/pina/equation/equation_interface.py @@ -1,4 +1,5 @@ """ Module for EquationInterface class """ + from abc import ABCMeta, abstractmethod diff --git a/pina/equation/system_equation.py b/pina/equation/system_equation.py index 28861e6..16d8c46 100644 --- a/pina/equation/system_equation.py +++ b/pina/equation/system_equation.py @@ -1,4 +1,5 @@ """ Module for SystemEquation. """ + import torch from .equation import Equation from ..utils import check_consistency @@ -6,7 +7,7 @@ from ..utils import check_consistency class SystemEquation(Equation): - def __init__(self, list_equation, reduction='mean'): + def __init__(self, list_equation, reduction="mean"): """ System of Equation class for specifing any system of equations in PINA. @@ -33,15 +34,16 @@ class SystemEquation(Equation): self.equations.append(Equation(equation)) # possible reduction - if reduction == 'mean': + if reduction == "mean": self.reduction = torch.mean - elif reduction == 'sum': + elif reduction == "sum": self.reduction = torch.sum - elif (reduction == 'none') or callable(reduction): + elif (reduction == "none") or callable(reduction): self.reduction = reduction else: raise NotImplementedError( - 'Only mean and sum reductions implemented.') + "Only mean and sum reductions implemented." + ) def residual(self, input_, output_, params_=None): """ @@ -64,9 +66,13 @@ class SystemEquation(Equation): :rtype: LabelTensor """ residual = torch.hstack( - [equation.residual(input_, output_, params_) for equation in self.equations]) + [ + equation.residual(input_, output_, params_) + for equation in self.equations + ] + ) - if self.reduction == 'none': + if self.reduction == "none": return residual return self.reduction(residual, dim=-1) diff --git a/pina/geometry/__init__.py b/pina/geometry/__init__.py index a936069..963136a 100644 --- a/pina/geometry/__init__.py +++ b/pina/geometry/__init__.py @@ -1,6 +1,13 @@ __all__ = [ - 'Location', 'CartesianDomain', 'EllipsoidDomain', 'Union', 'Intersection', - 'Exclusion', 'Difference', 'OperationInterface', 'SimplexDomain' + "Location", + "CartesianDomain", + "EllipsoidDomain", + "Union", + "Intersection", + "Exclusion", + "Difference", + "OperationInterface", + "SimplexDomain", ] from .location import Location diff --git a/pina/geometry/cartesian.py b/pina/geometry/cartesian.py index 84f6554..6008589 100644 --- a/pina/geometry/cartesian.py +++ b/pina/geometry/cartesian.py @@ -72,17 +72,17 @@ class CartesianDomain(Location): :rtype: torch.Tensor """ dim = bounds.shape[0] - if mode in ['chebyshev', 'grid'] and dim != 1: - raise RuntimeError('Something wrong in Span...') + if mode in ["chebyshev", "grid"] and dim != 1: + raise RuntimeError("Something wrong in Span...") - if mode == 'random': + if mode == "random": pts = torch.rand(size=(n, dim)) - elif mode == 'chebyshev': - pts = chebyshev_roots(n).mul(.5).add(.5).reshape(-1, 1) - elif mode == 'grid': + elif mode == "chebyshev": + pts = chebyshev_roots(n).mul(0.5).add(0.5).reshape(-1, 1) + elif mode == "grid": pts = torch.linspace(0, 1, n).reshape(-1, 1) # elif mode == 'lh' or mode == 'latin': - elif mode in ['lh', 'latin']: + elif mode in ["lh", "latin"]: pts = torch_lhs(n, dim) pts *= bounds[:, 1] - bounds[:, 0] @@ -90,7 +90,7 @@ class CartesianDomain(Location): return pts - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """Sample routine. :param n: Number of points to sample, see Note below @@ -145,7 +145,7 @@ class CartesianDomain(Location): """ def _1d_sampler(n, mode, variables): - """ Sample independentely the variables and cross the results""" + """Sample independentely the variables and cross the results""" tmp = [] for variable in variables: if variable in self.range_.keys(): @@ -158,17 +158,18 @@ class CartesianDomain(Location): result = tmp[0] for i in tmp[1:]: - result = result.append(i, mode='cross') + result = result.append(i, mode="cross") for variable in variables: if variable in self.fixed_.keys(): value = self.fixed_[variable] - pts_variable = torch.tensor([[value] - ]).repeat(result.shape[0], 1) + pts_variable = torch.tensor([[value]]).repeat( + result.shape[0], 1 + ) pts_variable = pts_variable.as_subclass(LabelTensor) pts_variable.labels = [variable] - result = result.append(pts_variable, mode='std') + result = result.append(pts_variable, mode="std") return result @@ -197,12 +198,13 @@ class CartesianDomain(Location): for variable in variables: if variable in self.fixed_.keys(): value = self.fixed_[variable] - pts_variable = torch.tensor([[value] - ]).repeat(result.shape[0], 1) + pts_variable = torch.tensor([[value]]).repeat( + result.shape[0], 1 + ) pts_variable = pts_variable.as_subclass(LabelTensor) pts_variable.labels = [variable] - result = result.append(pts_variable, mode='std') + result = result.append(pts_variable, mode="std") return result def _single_points_sample(n, variables): @@ -226,22 +228,22 @@ class CartesianDomain(Location): result = tmp[0] for i in tmp[1:]: - result = result.append(i, mode='std') + result = result.append(i, mode="std") return result if self.fixed_ and (not self.range_): return _single_points_sample(n, variables) - if variables == 'all': + if variables == "all": variables = list(self.range_.keys()) + list(self.fixed_.keys()) - if mode in ['grid', 'chebyshev']: + if mode in ["grid", "chebyshev"]: return _1d_sampler(n, mode, variables) - elif mode in ['random', 'lh', 'latin']: + elif mode in ["random", "lh", "latin"]: return _Nd_sampler(n, mode, variables) else: - raise ValueError(f'mode={mode} is not valid.') + raise ValueError(f"mode={mode} is not valid.") def is_inside(self, point, check_border=False): """Check if a point is inside the ellipsoid. diff --git a/pina/geometry/difference_domain.py b/pina/geometry/difference_domain.py index efdad63..d2ba414 100644 --- a/pina/geometry/difference_domain.py +++ b/pina/geometry/difference_domain.py @@ -20,7 +20,7 @@ class Difference(OperationInterface): the dimension of the geometry space. :param list geometries: A list of geometries from ``pina.geometry`` - such as ``EllipsoidDomain`` or ``CartesianDomain``. The first + such as ``EllipsoidDomain`` or ``CartesianDomain``. The first geometry in the list is the geometry from which points are sampled. The rest of the geometries are the geometries that are excluded from the first geometry to find the difference. @@ -39,7 +39,7 @@ class Difference(OperationInterface): Check if a point is inside the ``Difference`` domain. :param point: Point to be checked. - :type point: torch.Tensor + :type point: torch.Tensor :param bool check_border: If ``True``, the border is considered inside. :return: ``True`` if the point is inside the Exclusion domain, ``False`` otherwise. :rtype: bool @@ -49,7 +49,7 @@ class Difference(OperationInterface): return False return self.geometries[0].is_inside(point, check_border) - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """ Sample routine for ``Difference`` domain. @@ -77,9 +77,10 @@ class Difference(OperationInterface): 5 """ - if mode != 'random': + if mode != "random": raise NotImplementedError( - f'{mode} is not a valid mode for sampling.') + f"{mode} is not a valid mode for sampling." + ) sampled = [] diff --git a/pina/geometry/ellipsoid.py b/pina/geometry/ellipsoid.py index e99425f..0764d81 100644 --- a/pina/geometry/ellipsoid.py +++ b/pina/geometry/ellipsoid.py @@ -85,7 +85,7 @@ class EllipsoidDomain(Location): .. note:: When ``sample_surface`` in the ``__init()__`` - is set to ``True``, then the method only checks + is set to ``True``, then the method only checks points on the surface, and not inside the domain. :param point: Point to be checked. @@ -103,7 +103,7 @@ class EllipsoidDomain(Location): # get axis ellipse as tensors list_dict_vals = list(self._axis.values()) tmp = torch.tensor(list_dict_vals, dtype=torch.float) - ax_sq = LabelTensor(tmp.reshape(1, -1)**2, self.variables) + ax_sq = LabelTensor(tmp.reshape(1, -1) ** 2, self.variables) # get centers ellipse as tensors list_dict_vals = list(self._centers.values()) @@ -111,16 +111,18 @@ class EllipsoidDomain(Location): centers = LabelTensor(tmp.reshape(1, -1), self.variables) if not all([i in ax_sq.labels for i in point.labels]): - raise ValueError('point labels different from constructor' - f' dictionary labels. Got {point.labels},' - f' expected {ax_sq.labels}.') + raise ValueError( + "point labels different from constructor" + f" dictionary labels. Got {point.labels}," + f" expected {ax_sq.labels}." + ) # point square + shift center point_sq = (point - centers).pow(2) point_sq.labels = point.labels # calculate ellispoid equation - eqn = torch.sum(point_sq.extract(ax_sq.labels) / ax_sq) - 1. + eqn = torch.sum(point_sq.extract(ax_sq.labels) / ax_sq) - 1.0 # if we have sampled only the surface, we check that the # point is inside the surface border only @@ -160,8 +162,9 @@ class EllipsoidDomain(Location): dim = len(variables) # get values center - pairs_center = [(k, v) for k, v in self._centers.items() - if k in variables] + pairs_center = [ + (k, v) for k, v in self._centers.items() if k in variables + ] _, values_center = map(list, zip(*pairs_center)) values_center = torch.tensor(values_center) @@ -171,7 +174,7 @@ class EllipsoidDomain(Location): values_axis = torch.tensor(values_axis) # Sample in the unit sphere - if mode == 'random': + if mode == "random": # 1. Sample n points from the surface of a unit sphere # 2. Scale each dimension using torch.rand() # (a random number between 0-1) so that it lies within @@ -192,7 +195,7 @@ class EllipsoidDomain(Location): return pts - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """Sample routine. :param int n: Number of points to sample in the shape. @@ -238,12 +241,13 @@ class EllipsoidDomain(Location): for variable in variables: if variable in self.fixed_.keys(): value = self.fixed_[variable] - pts_variable = torch.tensor([[value] - ]).repeat(result.shape[0], 1) + pts_variable = torch.tensor([[value]]).repeat( + result.shape[0], 1 + ) pts_variable = pts_variable.as_subclass(LabelTensor) pts_variable.labels = [variable] - result = result.append(pts_variable, mode='std') + result = result.append(pts_variable, mode="std") return result def _single_points_sample(n, variables): @@ -267,17 +271,17 @@ class EllipsoidDomain(Location): result = tmp[0] for i in tmp[1:]: - result = result.append(i, mode='std') + result = result.append(i, mode="std") return result if self.fixed_ and (not self.range_): return _single_points_sample(n, variables) - if variables == 'all': + if variables == "all": variables = list(self.range_.keys()) + list(self.fixed_.keys()) - if mode in ['random']: + if mode in ["random"]: return _Nd_sampler(n, mode, variables) else: - raise NotImplementedError(f'mode={mode} is not implemented.') + raise NotImplementedError(f"mode={mode} is not implemented.") diff --git a/pina/geometry/exclusion_domain.py b/pina/geometry/exclusion_domain.py index 457289b..ed63db3 100644 --- a/pina/geometry/exclusion_domain.py +++ b/pina/geometry/exclusion_domain.py @@ -37,7 +37,7 @@ class Exclusion(OperationInterface): Check if a point is inside the ``Exclusion`` domain. :param point: Point to be checked. - :type point: torch.Tensor + :type point: torch.Tensor :param bool check_border: If ``True``, the border is considered inside. :return: ``True`` if the point is inside the Exclusion domain, ``False`` otherwise. :rtype: bool @@ -48,7 +48,7 @@ class Exclusion(OperationInterface): flag += 1 return flag == 1 - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """ Sample routine for ``Exclusion`` domain. @@ -76,9 +76,10 @@ class Exclusion(OperationInterface): 5 """ - if mode != 'random': + if mode != "random": raise NotImplementedError( - f'{mode} is not a valid mode for sampling.') + f"{mode} is not a valid mode for sampling." + ) sampled = [] @@ -104,4 +105,4 @@ class Exclusion(OperationInterface): sampled_points.append(sample) sampled += sampled_points - return LabelTensor(torch.cat(sampled), labels=self.variables) \ No newline at end of file + return LabelTensor(torch.cat(sampled), labels=self.variables) diff --git a/pina/geometry/intersection_domain.py b/pina/geometry/intersection_domain.py index e5ecb1a..b40d369 100644 --- a/pina/geometry/intersection_domain.py +++ b/pina/geometry/intersection_domain.py @@ -20,7 +20,7 @@ class Intersection(OperationInterface): with :math:`x` a point in :math:`\mathbb{R}^N` and :math:`N` the dimension of the geometry space. - :param list geometries: A list of geometries from ``pina.geometry`` + :param list geometries: A list of geometries from ``pina.geometry`` such as ``EllipsoidDomain`` or ``CartesianDomain``. The intersection will be taken between all the geometries in the list. The resulting geometry will be the intersection of all the geometries in the list. @@ -39,7 +39,7 @@ class Intersection(OperationInterface): Check if a point is inside the ``Intersection`` domain. :param point: Point to be checked. - :type point: torch.Tensor + :type point: torch.Tensor :param bool check_border: If ``True``, the border is considered inside. :return: ``True`` if the point is inside the Intersection domain, ``False`` otherwise. :rtype: bool @@ -50,7 +50,7 @@ class Intersection(OperationInterface): flag += 1 return flag == len(self.geometries) - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """ Sample routine for ``Intersection`` domain. @@ -78,9 +78,10 @@ class Intersection(OperationInterface): 5 """ - if mode != 'random': + if mode != "random": raise NotImplementedError( - f'{mode} is not a valid mode for sampling.') + f"{mode} is not a valid mode for sampling." + ) sampled = [] diff --git a/pina/geometry/operation_interface.py b/pina/geometry/operation_interface.py index 670e924..4f7709b 100644 --- a/pina/geometry/operation_interface.py +++ b/pina/geometry/operation_interface.py @@ -26,7 +26,7 @@ class OperationInterface(Location, metaclass=ABCMeta): @property def geometries(self): - """ + """ The geometries to perform set operation. """ return self._geometries @@ -40,15 +40,15 @@ class OperationInterface(Location, metaclass=ABCMeta): :rtype: list[str] """ return self.geometries[0].variables - - @ abstractmethod + + @abstractmethod def is_inside(self, point, check_border=False): """ Check if a point is inside the resulting domain after a set operation is applied. :param point: Point to be checked. - :type point: torch.Tensor + :type point: torch.Tensor :param bool check_border: If ``True``, the border is considered inside. :return: ``True`` if the point is inside the Intersection domain, ``False`` otherwise. :rtype: bool @@ -64,4 +64,5 @@ class OperationInterface(Location, metaclass=ABCMeta): for geometry in geometries: if geometry.variables != geometries[0].variables: raise NotImplementedError( - f'The geometries need to have same dimensions and labels.') + f"The geometries need to have same dimensions and labels." + ) diff --git a/pina/geometry/simplex.py b/pina/geometry/simplex.py index 6cfcfc7..2e78872 100644 --- a/pina/geometry/simplex.py +++ b/pina/geometry/simplex.py @@ -94,7 +94,7 @@ class SimplexDomain(Location): # respective coord bounded by the lowest and highest values span_dict[coord] = [ float(sorted_vertices[0][i]), - float(sorted_vertices[-1][i]) + float(sorted_vertices[-1][i]), ] return CartesianDomain(span_dict) @@ -120,16 +120,19 @@ class SimplexDomain(Location): """ if not all(label in self.variables for label in point.labels): - raise ValueError("Point labels different from constructor" - f" dictionary labels. Got {point.labels}," - f" expected {self.variables}.") + raise ValueError( + "Point labels different from constructor" + f" dictionary labels. Got {point.labels}," + f" expected {self.variables}." + ) point_shift = point - self._vertices_matrix[-1] point_shift = point_shift.tensor.reshape(-1, 1) # compute barycentric coordinates - lambda_ = torch.linalg.solve(self._vectors_shifted * 1.0, - point_shift * 1.0) + lambda_ = torch.linalg.solve( + self._vectors_shifted * 1.0, point_shift * 1.0 + ) lambda_1 = 1.0 - torch.sum(lambda_) lambdas = torch.vstack([lambda_, lambda_1]) @@ -137,8 +140,9 @@ class SimplexDomain(Location): if not check_border: return all(torch.gt(lambdas, 0.0)) and all(torch.lt(lambdas, 1.0)) - return all(torch.ge(lambdas, 0)) and (any(torch.eq(lambdas, 0)) - or any(torch.eq(lambdas, 1))) + return all(torch.ge(lambdas, 0)) and ( + any(torch.eq(lambdas, 0)) or any(torch.eq(lambdas, 1)) + ) def _sample_interior_randomly(self, n, variables): """ @@ -163,9 +167,9 @@ class SimplexDomain(Location): sampled_points = [] while len(sampled_points) < n: - sampled_point = self._cartesian_bound.sample(n=1, - mode="random", - variables=variables) + sampled_point = self._cartesian_bound.sample( + n=1, mode="random", variables=variables + ) if self.is_inside(sampled_point, self._sample_surface): sampled_points.append(sampled_point) @@ -196,9 +200,9 @@ class SimplexDomain(Location): # extract number of vertices number_of_vertices = self._vertices_matrix.shape[0] # extract idx lambda to set to zero randomly - idx_lambda = torch.randint(low=0, - high=number_of_vertices, - size=(1, )) + idx_lambda = torch.randint( + low=0, high=number_of_vertices, size=(1,) + ) # build lambda vector # 1. sampling [1, 2) lambdas = torch.rand((number_of_vertices, 1)) @@ -236,4 +240,4 @@ class SimplexDomain(Location): else: raise NotImplementedError(f"mode={mode} is not implemented.") - return LabelTensor(sample_pts, labels=self.variables) \ No newline at end of file + return LabelTensor(sample_pts, labels=self.variables) diff --git a/pina/geometry/union_domain.py b/pina/geometry/union_domain.py index 1141236..da2ead9 100644 --- a/pina/geometry/union_domain.py +++ b/pina/geometry/union_domain.py @@ -20,7 +20,7 @@ class Union(OperationInterface): with :math:`x` a point in :math:`\mathbb{R}^N` and :math:`N` the dimension of the geometry space. - :param list geometries: A list of geometries from ``pina.geometry`` + :param list geometries: A list of geometries from ``pina.geometry`` such as ``EllipsoidDomain`` or ``CartesianDomain``. :Example: @@ -50,7 +50,7 @@ class Union(OperationInterface): return True return False - def sample(self, n, mode='random', variables='all'): + def sample(self, n, mode="random", variables="all"): """ Sample routine for ``Union`` domain. @@ -93,8 +93,10 @@ class Union(OperationInterface): # different than zero. Notice that len(geometries) is # always smaller than remaider. sampled_points.append( - geometry.sample(num_points + int(i < remainder), mode, - variables)) + geometry.sample( + num_points + int(i < remainder), mode, variables + ) + ) # in case number of sampled points is smaller than the number of geometries if len(sampled_points) >= n: break diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 502d31f..c92dda9 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -1,4 +1,5 @@ """ Module for LabelTensor """ + from typing import Any import torch from torch import Tensor @@ -12,7 +13,7 @@ class LabelTensor(torch.Tensor): return super().__new__(cls, x, *args, **kwargs) def __init__(self, x, labels): - ''' + """ Construct a `LabelTensor` by passing a tensor and a list of column labels. Such labels uniquely identify the columns of the tensor, allowing for an easier manipulation. @@ -64,7 +65,7 @@ class LabelTensor(torch.Tensor): [0.9427, 0.5819], [0.9518, 0.1025], [0.8066, 0.9615]]) - ''' + """ if x.ndim == 1: x = x.reshape(-1, 1) @@ -72,10 +73,12 @@ class LabelTensor(torch.Tensor): labels = [labels] if len(labels) != x.shape[-1]: - raise ValueError('the tensor has not the same number of columns of ' - 'the passed labels.') + raise ValueError( + "the tensor has not the same number of columns of " + "the passed labels." + ) self._labels = labels - + @property def labels(self): """Property decorator for labels @@ -88,8 +91,10 @@ class LabelTensor(torch.Tensor): @labels.setter def labels(self, labels): if len(labels) != self.shape[self.ndim - 1]: # small check - raise ValueError('The tensor has not the same number of columns of ' - 'the passed labels.') + raise ValueError( + "The tensor has not the same number of columns of " + "the passed labels." + ) self._labels = labels # assign the label @@ -109,7 +114,7 @@ class LabelTensor(torch.Tensor): all_labels = [label for lt in label_tensors for label in lt.labels] if set(all_labels) != set(label_tensors[0].labels): - raise RuntimeError('The tensors to stack have different labels') + raise RuntimeError("The tensors to stack have different labels") labels = label_tensors[0].labels tensors = [lt.extract(labels) for lt in label_tensors] @@ -123,7 +128,7 @@ class LabelTensor(torch.Tensor): :return: A copy of the tensor. :rtype: LabelTensor """ - # # used before merging + # # used before merging # try: # out = LabelTensor(super().clone(*args, **kwargs), self.labels) # except: @@ -184,14 +189,15 @@ class LabelTensor(torch.Tensor): pass else: raise TypeError( - '`label_to_extract` should be a str, or a str iterator') + "`label_to_extract` should be a str, or a str iterator" + ) indeces = [] for f in label_to_extract: try: indeces.append(self.labels.index(f)) except ValueError: - raise ValueError(f'`{f}` not in the labels list') + raise ValueError(f"`{f}` not in the labels list") new_data = super(Tensor, self.T).__getitem__(indeces).T new_labels = [self.labels[idx] for idx in indeces] @@ -203,17 +209,16 @@ class LabelTensor(torch.Tensor): def detach(self): detached = super().detach() - if hasattr(self, '_labels'): + if hasattr(self, "_labels"): detached._labels = self._labels return detached - - def requires_grad_(self, mode = True): + def requires_grad_(self, mode=True): lt = super().requires_grad_(mode) lt.labels = self.labels return lt - def append(self, lt, mode='std'): + def append(self, lt, mode="std"): """ Return a copy of the merged tensors. @@ -223,22 +228,23 @@ class LabelTensor(torch.Tensor): :rtype: LabelTensor """ if set(self.labels).intersection(lt.labels): - raise RuntimeError('The tensors to merge have common labels') + raise RuntimeError("The tensors to merge have common labels") new_labels = self.labels + lt.labels - if mode == 'std': + if mode == "std": new_tensor = torch.cat((self, lt), dim=1) - elif mode == 'first': + elif mode == "first": raise NotImplementedError - elif mode == 'cross': + elif mode == "cross": tensor1 = self tensor2 = lt n1 = tensor1.shape[0] n2 = tensor2.shape[0] tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels) - tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), - labels=tensor2.labels) + tensor2 = LabelTensor( + tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels + ) new_tensor = torch.cat((tensor1, tensor2), dim=1) new_tensor = new_tensor.as_subclass(LabelTensor) @@ -250,34 +256,37 @@ class LabelTensor(torch.Tensor): Return a copy of the selected tensor. """ - if isinstance(index, str) or (isinstance(index, (tuple, list))and all(isinstance(a, str) for a in index)): + if isinstance(index, str) or ( + isinstance(index, (tuple, list)) + and all(isinstance(a, str) for a in index) + ): return self.extract(index) selected_lt = super(Tensor, self).__getitem__(index) - + try: len_index = len(index) except TypeError: len_index = 1 - + if isinstance(index, int) or len_index == 1: if selected_lt.ndim == 1: selected_lt = selected_lt.reshape(1, -1) - if hasattr(self, 'labels'): + if hasattr(self, "labels"): selected_lt.labels = self.labels elif len_index == 2: if selected_lt.ndim == 1: selected_lt = selected_lt.reshape(-1, 1) - if hasattr(self, 'labels'): + if hasattr(self, "labels"): if isinstance(index[1], list): selected_lt.labels = [self.labels[i] for i in index[1]] else: selected_lt.labels = self.labels[index[1]] else: selected_lt.labels = self.labels - + return selected_lt - + @property def tensor(self): return self.as_subclass(Tensor) @@ -286,9 +295,9 @@ class LabelTensor(torch.Tensor): return super().__len__() def __str__(self): - if hasattr(self, 'labels'): - s = f'labels({str(self.labels)})\n' + if hasattr(self, "labels"): + s = f"labels({str(self.labels)})\n" else: - s = 'no labels\n' + s = "no labels\n" s += super().__str__() - return s \ No newline at end of file + return s diff --git a/pina/loss.py b/pina/loss.py index 8c0b67e..3cbf888 100644 --- a/pina/loss.py +++ b/pina/loss.py @@ -5,7 +5,7 @@ from torch.nn.modules.loss import _Loss import torch from .utils import check_consistency -__all__ = ['LossInterface', 'LpLoss', 'PowerLoss'] +__all__ = ["LossInterface", "LpLoss", "PowerLoss"] class LossInterface(_Loss, metaclass=ABCMeta): @@ -14,10 +14,10 @@ class LossInterface(_Loss, metaclass=ABCMeta): should be inheritied from this class. """ - def __init__(self, reduction='mean'): + def __init__(self, reduction="mean"): """ :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction will be applied, ``mean``: the sum of the output will be divided by the number of elements in the output, ``sum``: the output will be summed. Note: ``size_average`` and ``reduce`` are in the @@ -41,7 +41,7 @@ class LossInterface(_Loss, metaclass=ABCMeta): """Simple helper function to check reduction :param reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction will be applied, ``mean``: the sum of the output will be divided by the number of elements in the output, ``sum``: the output will be summed. Note: ``size_average`` and ``reduce`` are in the @@ -101,14 +101,14 @@ class LpLoss(LossInterface): The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. """ - def __init__(self, p=2, reduction='mean', relative=False): + def __init__(self, p=2, reduction="mean", relative=False): """ :param int p: Degree of Lp norm. It specifies the type of norm to be calculated. See `list of possible orders in torch linalg `_ to for possible degrees. Default 2 (euclidean norm). :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. ``none``: no reduction + ``none`` | ``mean`` | ``sum``. ``none``: no reduction will be applied, ``mean``: the sum of the output will be divided by the number of elements in the output, ``sum``: the output will be summed. @@ -117,7 +117,7 @@ class LpLoss(LossInterface): super().__init__(reduction=reduction) # check consistency - check_consistency(p, (str,int,float)) + check_consistency(p, (str, int, float)) check_consistency(relative, bool) self.p = p @@ -174,14 +174,14 @@ class PowerLoss(LossInterface): The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. """ - def __init__(self, p=2, reduction='mean', relative=False): + def __init__(self, p=2, reduction="mean", relative=False): """ :param int p: Degree of Lp norm. It specifies the type of norm to be calculated. See `list of possible orders in torch linalg `_ to see the possible degrees. Default 2 (euclidean norm). :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction will be applied, ``mean``: the sum of the output will be divided by the number of elements in the output, ``sum``: the output will be summed. diff --git a/pina/meta.py b/pina/meta.py index 581a866..efad679 100644 --- a/pina/meta.py +++ b/pina/meta.py @@ -1,15 +1,22 @@ __all__ = [ - '__project__', '__title__', '__author__', '__copyright__', '__license__', - '__version__', '__mail__', '__maintainer__', '__status__' + "__project__", + "__title__", + "__author__", + "__copyright__", + "__license__", + "__version__", + "__mail__", + "__maintainer__", + "__status__", ] -__project__ = 'PINA' +__project__ = "PINA" __title__ = "pina" __author__ = "PINA Contributors" __copyright__ = "Copyright 2021-2024, PINA Contributors" __license__ = "MIT" __version__ = "0.1.0" -__mail__ = 'demo.nicola@gmail.com, dario.coscia@sissa.it' # TODO +__mail__ = "demo.nicola@gmail.com, dario.coscia@sissa.it" # TODO __maintainer__ = __author__ __status__ = "Alpha" __packagename__ = "pina-mathlab" diff --git a/pina/model/__init__.py b/pina/model/__init__.py index d4c3acb..aab81bf 100644 --- a/pina/model/__init__.py +++ b/pina/model/__init__.py @@ -1,10 +1,10 @@ __all__ = [ - 'FeedForward', - 'ResidualFeedForward', - 'MultiFeedForward', - 'DeepONet', - 'MIONet', - 'FNO', + "FeedForward", + "ResidualFeedForward", + "MultiFeedForward", + "DeepONet", + "MIONet", + "FNO", ] from .feed_forward import FeedForward, ResidualFeedForward diff --git a/pina/model/deeponet.py b/pina/model/deeponet.py index 290c877..b39f532 100644 --- a/pina/model/deeponet.py +++ b/pina/model/deeponet.py @@ -1,4 +1,5 @@ """Module for DeepONet model""" + import torch import torch.nn as nn from ..utils import check_consistency, is_function @@ -24,12 +25,14 @@ class MIONet(torch.nn.Module): """ - def __init__(self, - networks, - aggregator="*", - reduction="+", - scale=True, - translation=True): + def __init__( + self, + networks, + aggregator="*", + reduction="+", + scale=True, + translation=True, + ): """ :param dict networks: The neural networks to use as models. The ``dict`` takes as key a neural network, and @@ -121,8 +124,9 @@ class MIONet(torch.nn.Module): shapes.append(key(input_).shape[-1]) if not all(map(lambda x: x == shapes[0], shapes)): - raise ValueError('The passed networks have not the same ' - 'output dimension.') + raise ValueError( + "The passed networks have not the same " "output dimension." + ) # assign trunk and branch net with their input indeces self.models = torch.nn.ModuleList(networks.keys()) @@ -133,10 +137,16 @@ class MIONet(torch.nn.Module): self._init_reduction(reduction=reduction) # scale and translation - self._scale = torch.nn.Parameter(torch.tensor( - [1.0])) if scale else torch.tensor([1.0]) - self._trasl = torch.nn.Parameter(torch.tensor( - [1.0])) if translation else torch.tensor([1.0]) + self._scale = ( + torch.nn.Parameter(torch.tensor([1.0])) + if scale + else torch.tensor([1.0]) + ) + self._trasl = ( + torch.nn.Parameter(torch.tensor([1.0])) + if translation + else torch.tensor([1.0]) + ) @staticmethod def _symbol_functions(**kwargs): @@ -180,16 +190,18 @@ class MIONet(torch.nn.Module): return x.extract(indeces) except AttributeError: raise RuntimeError( - 'Not possible to extract input variables from tensor.' - ' Ensure that the passed tensor is a LabelTensor or' - ' pass list of integers to extract variables. For' - ' more information refer to warning in the documentation.') + "Not possible to extract input variables from tensor." + " Ensure that the passed tensor is a LabelTensor or" + " pass list of integers to extract variables. For" + " more information refer to warning in the documentation." + ) elif isinstance(indeces[0], int): return x[..., indeces] else: raise RuntimeError( - 'Not able to extract right indeces for tensor.' - ' For more information refer to warning in the documentation.') + "Not able to extract right indeces for tensor." + " For more information refer to warning in the documentation." + ) def forward(self, x): """ @@ -197,7 +209,7 @@ class MIONet(torch.nn.Module): :param LabelTensor or torch.Tensor x: The input tensor for the forward call. :return: The output computed by the DeepONet model. - :rtype: LabelTensor or torch.Tensor + :rtype: LabelTensor or torch.Tensor """ # forward pass @@ -267,7 +279,7 @@ class DeepONet(MIONet): DeepONet is a general architecture for learning Operators. Unlike traditional machine learning methods DeepONet is designed to map - entire functions to other functions. It can be trained both with + entire functions to other functions. It can be trained both with Physics Informed or Supervised learning strategies. .. seealso:: @@ -280,15 +292,17 @@ class DeepONet(MIONet): """ - def __init__(self, - branch_net, - trunk_net, - input_indeces_branch_net, - input_indeces_trunk_net, - aggregator="*", - reduction="+", - scale=True, - translation=True): + def __init__( + self, + branch_net, + trunk_net, + input_indeces_branch_net, + input_indeces_trunk_net, + aggregator="*", + reduction="+", + scale=True, + translation=True, + ): """ :param torch.nn.Module branch_net: The neural network to use as branch model. It has to take as input a :py:obj:`pina.label_tensor.LabelTensor` @@ -363,14 +377,15 @@ class DeepONet(MIONet): """ networks = { branch_net: input_indeces_branch_net, - trunk_net: input_indeces_trunk_net + trunk_net: input_indeces_trunk_net, } - super().__init__(networks=networks, - aggregator=aggregator, - reduction=reduction, - scale=scale, - translation=translation) - + super().__init__( + networks=networks, + aggregator=aggregator, + reduction=reduction, + scale=scale, + translation=translation, + ) def forward(self, x): """ @@ -378,11 +393,10 @@ class DeepONet(MIONet): :param LabelTensor or torch.Tensor x: The input tensor for the forward call. :return: The output computed by the DeepONet model. - :rtype: LabelTensor or torch.Tensor + :rtype: LabelTensor or torch.Tensor """ return super().forward(x) - @property def branch_net(self): """ diff --git a/pina/model/feed_forward.py b/pina/model/feed_forward.py index 63d9bb2..ba98680 100644 --- a/pina/model/feed_forward.py +++ b/pina/model/feed_forward.py @@ -1,4 +1,5 @@ """Module for FeedForward model""" + import torch import torch.nn as nn from ..utils import check_consistency @@ -29,24 +30,25 @@ class FeedForward(torch.nn.Module): :param bool bias: If ``True`` the MLP will consider some bias. """ - def __init__(self, - input_dimensions, - output_dimensions, - inner_size=20, - n_layers=2, - func=nn.Tanh, - layers=None, - bias=True): - """ - """ + def __init__( + self, + input_dimensions, + output_dimensions, + inner_size=20, + n_layers=2, + func=nn.Tanh, + layers=None, + bias=True, + ): + """ """ super().__init__() if not isinstance(input_dimensions, int): - raise ValueError('input_dimensions expected to be int.') + raise ValueError("input_dimensions expected to be int.") self.input_dimension = input_dimensions if not isinstance(output_dimensions, int): - raise ValueError('output_dimensions expected to be int.') + raise ValueError("output_dimensions expected to be int.") self.output_dimension = output_dimensions if layers is None: layers = [inner_size] * n_layers @@ -58,7 +60,8 @@ class FeedForward(torch.nn.Module): self.layers = [] for i in range(len(tmp_layers) - 1): self.layers.append( - nn.Linear(tmp_layers[i], tmp_layers[i + 1], bias=bias)) + nn.Linear(tmp_layers[i], tmp_layers[i + 1], bias=bias) + ) if isinstance(func, list): self.functions = func @@ -66,7 +69,7 @@ class FeedForward(torch.nn.Module): self.functions = [func for _ in range(len(self.layers) - 1)] if len(self.layers) != len(self.functions) + 1: - raise RuntimeError('uncosistent number of layers and functions') + raise RuntimeError("uncosistent number of layers and functions") unique_list = [] for layer, func in zip(self.layers[:-1], self.functions): @@ -97,7 +100,7 @@ class ResidualFeedForward(torch.nn.Module): .. seealso:: - **Original reference**: Wang, Sifan, Yujun Teng, and Paris Perdikaris. + **Original reference**: Wang, Sifan, Yujun Teng, and Paris Perdikaris. *Understanding and mitigating gradient flow pathologies in physics-informed neural networks*. SIAM Journal on Scientific Computing 43.5 (2021): A3055-A3081. DOI: `10.1137/20M1318043 @@ -124,16 +127,17 @@ class ResidualFeedForward(torch.nn.Module): dimension must be the same as ``inner_size``. """ - def __init__(self, - input_dimensions, - output_dimensions, - inner_size=20, - n_layers=2, - func=nn.Tanh, - bias=True, - transformer_nets=None): - """ - """ + def __init__( + self, + input_dimensions, + output_dimensions, + inner_size=20, + n_layers=2, + func=nn.Tanh, + bias=True, + transformer_nets=None, + ): + """ """ super().__init__() # check type consistency @@ -148,35 +152,42 @@ class ResidualFeedForward(torch.nn.Module): if transformer_nets is None: transformer_nets = [ EnhancedLinear( - nn.Linear(in_features=input_dimensions, - out_features=inner_size), nn.Tanh()), + nn.Linear( + in_features=input_dimensions, out_features=inner_size + ), + nn.Tanh(), + ), EnhancedLinear( - nn.Linear(in_features=input_dimensions, - out_features=inner_size), nn.Tanh()) + nn.Linear( + in_features=input_dimensions, out_features=inner_size + ), + nn.Tanh(), + ), ] elif isinstance(transformer_nets, (list, tuple)): if len(transformer_nets) != 2: raise ValueError( - 'transformer_nets needs to be a list of len two.') + "transformer_nets needs to be a list of len two." + ) for net in transformer_nets: if not isinstance(net, nn.Module): raise ValueError( - 'transformer_nets needs to be a list of torch.nn.Module.' + "transformer_nets needs to be a list of torch.nn.Module." ) x = torch.rand(10, input_dimensions) try: out = net(x) except RuntimeError: raise ValueError( - 'transformer network input incompatible with input_dimensions.' + "transformer network input incompatible with input_dimensions." ) if out.shape[-1] != inner_size: raise ValueError( - 'transformer network output incompatible with inner_size.' + "transformer network output incompatible with inner_size." ) else: RuntimeError( - 'Runtime error for transformer nets, check official documentation.' + "Runtime error for transformer nets, check official documentation." ) # assign variables @@ -193,10 +204,11 @@ class ResidualFeedForward(torch.nn.Module): self.layers = [] for i in range(len(tmp_layers) - 1): self.layers.append( - nn.Linear(tmp_layers[i], tmp_layers[i + 1], bias=bias)) - self.last_layer = nn.Linear(tmp_layers[len(tmp_layers) - 1], - output_dimensions, - bias=bias) + nn.Linear(tmp_layers[i], tmp_layers[i + 1], bias=bias) + ) + self.last_layer = nn.Linear( + tmp_layers[len(tmp_layers) - 1], output_dimensions, bias=bias + ) if isinstance(func, list): self.functions = func() @@ -204,7 +216,7 @@ class ResidualFeedForward(torch.nn.Module): self.functions = [func() for _ in range(len(self.layers))] if len(self.layers) != len(self.functions): - raise RuntimeError('uncosistent number of layers and functions') + raise RuntimeError("uncosistent number of layers and functions") unique_list = [] for layer, func in zip(self.layers, self.functions): @@ -228,7 +240,7 @@ class ResidualFeedForward(torch.nn.Module): # skip connections pass for layer in self.inner_layers.children(): x = layer(x) - x = (1. - x) * input_[0] + x * input_[1] + x = (1.0 - x) * input_[0] + x * input_[1] # last layer return self.last_layer(x) diff --git a/pina/model/fno.py b/pina/model/fno.py index e756443..93168e8 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -12,7 +12,7 @@ class FNO(torch.nn.Module): Fourier Neural Operator (FNO) is a general architecture for learning Operators. Unlike traditional machine learning methods FNO is designed to map - entire functions to other functions. It can be trained both with + entire functions to other functions. It can be trained both with Supervised learning strategies. FNO does global convolution by performing the operation on the Fourier space. @@ -25,17 +25,19 @@ class FNO(torch.nn.Module): `_ """ - def __init__(self, - lifting_net, - projecting_net, - n_modes, - dimensions=3, - padding=8, - padding_type="constant", - inner_size=20, - n_layers=2, - func=nn.Tanh, - layers=None): + def __init__( + self, + lifting_net, + projecting_net, + n_modes, + dimensions=3, + padding=8, + padding_type="constant", + inner_size=20, + n_layers=2, + func=nn.Tanh, + layers=None, + ): super().__init__() # check type consistency @@ -51,11 +53,12 @@ class FNO(torch.nn.Module): if isinstance(layers, (tuple, list)): check_consistency(layers, int) else: - raise ValueError('layers must be tuple or list of int.') + raise ValueError("layers must be tuple or list of int.") if not isinstance(n_modes, (list, tuple, int)): raise ValueError( - 'n_modes must be a int or list or tuple of valid modes.' - ' More information on the official documentation.') + "n_modes must be a int or list or tuple of valid modes." + " More information on the official documentation." + ) # assign variables # TODO check input lifting net and input projecting net @@ -71,7 +74,7 @@ class FNO(torch.nn.Module): elif dimensions == 3: fourier_layer = FourierBlock3D else: - raise NotImplementedError('FNO implemented only for 1D/2D/3D data.') + raise NotImplementedError("FNO implemented only for 1D/2D/3D data.") # Here we build the FNO by stacking Fourier Blocks @@ -83,17 +86,20 @@ class FNO(torch.nn.Module): if isinstance(func, list): if len(layers) != len(func): raise RuntimeError( - 'Uncosistent number of layers and functions.') + "Uncosistent number of layers and functions." + ) self._functions = func else: self._functions = [func for _ in range(len(layers))] # 3. Assign modes functions for each FNO layer if isinstance(n_modes, list): - if all(isinstance(i, list) - for i in n_modes) and len(layers) != len(n_modes): + if all(isinstance(i, list) for i in n_modes) and len(layers) != len( + n_modes + ): raise RuntimeError( - 'Uncosistent number of layers and functions.') + "Uncosistent number of layers and functions." + ) elif all(isinstance(i, int) for i in n_modes): n_modes = [n_modes] * len(layers) else: @@ -109,10 +115,13 @@ class FNO(torch.nn.Module): self._layers = [] for i in range(len(tmp_layers) - 1): self._layers.append( - fourier_layer(input_numb_fields=tmp_layers[i], - output_numb_fields=tmp_layers[i + 1], - n_modes=n_modes[i], - activation=self._functions[i])) + fourier_layer( + input_numb_fields=tmp_layers[i], + output_numb_fields=tmp_layers[i + 1], + n_modes=n_modes[i], + activation=self._functions[i], + ) + ) self._layers = nn.Sequential(*self._layers) # 5. Padding values for spectral conv @@ -139,8 +148,10 @@ class FNO(torch.nn.Module): :return: The output tensor obtained from the FNO. :rtype: torch.Tensor """ - if isinstance(x, LabelTensor): #TODO remove when Network is fixed - warnings.warn('LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor') + if isinstance(x, LabelTensor): # TODO remove when Network is fixed + warnings.warn( + "LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor" + ) x = x.as_subclass(torch.Tensor) # lifting the input in higher dimensional space diff --git a/pina/model/layers/__init__.py b/pina/model/layers/__init__.py index 62a2372..a6e4e0b 100644 --- a/pina/model/layers/__init__.py +++ b/pina/model/layers/__init__.py @@ -1,18 +1,22 @@ __all__ = [ - 'ContinuousConvBlock', - 'ResidualBlock', - 'EnhancedLinear', - 'SpectralConvBlock1D', - 'SpectralConvBlock2D', - 'SpectralConvBlock3D', - 'FourierBlock1D', - 'FourierBlock2D', - 'FourierBlock3D', - 'PODLayer' + "ContinuousConvBlock", + "ResidualBlock", + "EnhancedLinear", + "SpectralConvBlock1D", + "SpectralConvBlock2D", + "SpectralConvBlock3D", + "FourierBlock1D", + "FourierBlock2D", + "FourierBlock3D", + "PODLayer", ] from .convolution_2d import ContinuousConvBlock from .residual import ResidualBlock, EnhancedLinear -from .spectral import SpectralConvBlock1D, SpectralConvBlock2D, SpectralConvBlock3D +from .spectral import ( + SpectralConvBlock1D, + SpectralConvBlock2D, + SpectralConvBlock3D, +) from .fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D from .pod import PODLayer diff --git a/pina/model/layers/convolution.py b/pina/model/layers/convolution.py index 3b9fdcd..c6ae4e2 100644 --- a/pina/model/layers/convolution.py +++ b/pina/model/layers/convolution.py @@ -1,4 +1,5 @@ """Module for Base Continuous Convolution class.""" + from abc import ABCMeta, abstractmethod import torch from .stride import Stride @@ -10,14 +11,16 @@ class BaseContinuousConv(torch.nn.Module, metaclass=ABCMeta): Abstract class """ - def __init__(self, - input_numb_field, - output_numb_field, - filter_dim, - stride, - model=None, - optimize=False, - no_overlap=False): + def __init__( + self, + input_numb_field, + output_numb_field, + filter_dim, + stride, + model=None, + optimize=False, + no_overlap=False, + ): """ Base Class for Continuous Convolution. @@ -75,43 +78,44 @@ class BaseContinuousConv(torch.nn.Module, metaclass=ABCMeta): if isinstance(input_numb_field, int): self._input_numb_field = input_numb_field else: - raise ValueError('input_numb_field must be int.') + raise ValueError("input_numb_field must be int.") if isinstance(output_numb_field, int): self._output_numb_field = output_numb_field else: - raise ValueError('input_numb_field must be int.') + raise ValueError("input_numb_field must be int.") if isinstance(filter_dim, (tuple, list)): vect = filter_dim else: - raise ValueError('filter_dim must be tuple or list.') + raise ValueError("filter_dim must be tuple or list.") vect = torch.tensor(vect) self.register_buffer("_dim", vect, persistent=False) if isinstance(stride, dict): self._stride = Stride(stride) else: - raise ValueError('stride must be dictionary.') + raise ValueError("stride must be dictionary.") self._net = model if isinstance(optimize, bool): self._optimize = optimize else: - raise ValueError('optimize must be bool.') + raise ValueError("optimize must be bool.") # choosing how to initialize based on optimization if self._optimize: # optimizing decorator ensure the function is called # just once self._choose_initialization = optimizing( - self._initialize_convolution) + self._initialize_convolution + ) else: self._choose_initialization = self._initialize_convolution if not isinstance(no_overlap, bool): - raise ValueError('no_overlap must be bool.') + raise ValueError("no_overlap must be bool.") if no_overlap: raise NotImplementedError @@ -125,11 +129,13 @@ class BaseContinuousConv(torch.nn.Module, metaclass=ABCMeta): super().__init__() assert isinstance(input_dim, int) assert isinstance(output_dim, int) - self._model = torch.nn.Sequential(torch.nn.Linear(input_dim, 20), - torch.nn.ReLU(), - torch.nn.Linear(20, 20), - torch.nn.ReLU(), - torch.nn.Linear(20, output_dim)) + self._model = torch.nn.Sequential( + torch.nn.Linear(input_dim, 20), + torch.nn.ReLU(), + torch.nn.Linear(20, 20), + torch.nn.ReLU(), + torch.nn.Linear(20, output_dim), + ) def forward(self, x): return self._model(x) diff --git a/pina/model/layers/convolution_2d.py b/pina/model/layers/convolution_2d.py index 4668796..665ddaf 100644 --- a/pina/model/layers/convolution_2d.py +++ b/pina/model/layers/convolution_2d.py @@ -1,4 +1,5 @@ """Module for Continuous Convolution class""" + from .convolution import BaseContinuousConv from .utils_convolution import check_point, map_points_ from .integral import Integral @@ -31,14 +32,16 @@ class ContinuousConvBlock(BaseContinuousConv): """ - def __init__(self, - input_numb_field, - output_numb_field, - filter_dim, - stride, - model=None, - optimize=False, - no_overlap=False): + def __init__( + self, + input_numb_field, + output_numb_field, + filter_dim, + stride, + model=None, + optimize=False, + no_overlap=False, + ): """ :param input_numb_field: Number of fields :math:`N_{in}` in the input. :type input_numb_field: int @@ -112,16 +115,18 @@ class ContinuousConvBlock(BaseContinuousConv): ) """ - super().__init__(input_numb_field=input_numb_field, - output_numb_field=output_numb_field, - filter_dim=filter_dim, - stride=stride, - model=model, - optimize=optimize, - no_overlap=no_overlap) + super().__init__( + input_numb_field=input_numb_field, + output_numb_field=output_numb_field, + filter_dim=filter_dim, + stride=stride, + model=model, + optimize=optimize, + no_overlap=no_overlap, + ) # integral routine - self._integral = Integral('discrete') + self._integral = Integral("discrete") # create the network self._net = self._spawn_networks(model) @@ -146,15 +151,18 @@ class ContinuousConvBlock(BaseContinuousConv): nets.append(tmp) else: if not isinstance(model, object): - raise ValueError("Expected a python class inheriting" - " from torch.nn.Module") + raise ValueError( + "Expected a python class inheriting" " from torch.nn.Module" + ) for _ in range(self._input_numb_field * self._output_numb_field): tmp = model() if not isinstance(tmp, torch.nn.Module): - raise ValueError("The python class must be inherited from" - " torch.nn.Module. See the docstring for" - " an example.") + raise ValueError( + "The python class must be inherited from" + " torch.nn.Module. See the docstring for" + " an example." + ) nets.append(tmp) return torch.nn.ModuleList(nets) @@ -232,11 +240,17 @@ class ContinuousConvBlock(BaseContinuousConv): number_points = len(self._stride) # initialize the grid - grid = torch.zeros(size=(X.shape[0], self._output_numb_field, - number_points, filter_dim + 1), - device=X.device, - dtype=X.dtype) - grid[..., :-1] = (self._stride + self._dim * 0.5) + grid = torch.zeros( + size=( + X.shape[0], + self._output_numb_field, + number_points, + filter_dim + 1, + ), + device=X.device, + dtype=X.dtype, + ) + grid[..., :-1] = self._stride + self._dim * 0.5 # saving the grid self._grid = grid.detach() @@ -269,14 +283,14 @@ class ContinuousConvBlock(BaseContinuousConv): """ # choose the type of convolution - if type == 'forward': + if type == "forward": return self._make_grid_forward(X) - elif type == 'inverse': + elif type == "inverse": self._make_grid_transpose(X) else: raise TypeError - def _initialize_convolution(self, X, type='forward'): + def _initialize_convolution(self, X, type="forward"): """ Private method to intialize the convolution. The convolution is initialized by setting a grid and @@ -307,10 +321,10 @@ class ContinuousConvBlock(BaseContinuousConv): # initialize convolution if self.training: # we choose what to do based on optimization - self._choose_initialization(X, type='forward') + self._choose_initialization(X, type="forward") else: # we always initialize on testing - self._initialize_convolution(X, 'forward') + self._initialize_convolution(X, "forward") # create convolutional array conv = self._grid.clone().detach() @@ -322,7 +336,8 @@ class ContinuousConvBlock(BaseContinuousConv): # extract mapped points stacked_input, indeces_channels = self._extract_mapped_points( - batch_idx, self._index, x) + batch_idx, self._index, x + ) # compute the convolution @@ -339,9 +354,11 @@ class ContinuousConvBlock(BaseContinuousConv): # calculate filter value staked_output = net(single_channel_input[..., :-1]) # perform integral for all strides in one field - integral = self._integral(staked_output, - single_channel_input[..., -1], - indeces_channels[idx]) + integral = self._integral( + staked_output, + single_channel_input[..., -1], + indeces_channels[idx], + ) res_tmp.append(integral) # stacking integral results @@ -349,9 +366,9 @@ class ContinuousConvBlock(BaseContinuousConv): # sum filters (for each input fields) in groups # for different ouput fields - conv[batch_idx, ..., - -1] = res_tmp.reshape(self._output_numb_field, - self._input_numb_field, -1).sum(1) + conv[batch_idx, ..., -1] = res_tmp.reshape( + self._output_numb_field, self._input_numb_field, -1 + ).sum(1) return conv def transpose_no_overlap(self, integrals, X): @@ -382,10 +399,10 @@ class ContinuousConvBlock(BaseContinuousConv): # initialize convolution if self.training: # we choose what to do based on optimization - self._choose_initialization(X, type='inverse') + self._choose_initialization(X, type="inverse") else: # we always initialize on testing - self._initialize_convolution(X, 'inverse') + self._initialize_convolution(X, "inverse") # initialize grid X = self._grid_transpose.clone().detach() @@ -398,7 +415,8 @@ class ContinuousConvBlock(BaseContinuousConv): # extract mapped points stacked_input, indeces_channels = self._extract_mapped_points( - batch_idx, self._index, x) + batch_idx, self._index, x + ) # compute the transpose convolution @@ -414,8 +432,9 @@ class ContinuousConvBlock(BaseContinuousConv): # extract input for each field single_channel_input = stacked_input[idx] rep_idx = torch.tensor(indeces_channels[idx]) - integral = integrals[batch_idx, - idx_in, :].repeat_interleave(rep_idx) + integral = integrals[batch_idx, idx_in, :].repeat_interleave( + rep_idx + ) # extract filter net = self._net[idx_conv] # perform transpose convolution for all strides in one field @@ -426,9 +445,11 @@ class ContinuousConvBlock(BaseContinuousConv): # stacking integral results and sum # filters (for each input fields) in groups # for different output fields - res_tmp = torch.stack(res_tmp).reshape(self._input_numb_field, - self._output_numb_field, - -1).sum(0) + res_tmp = ( + torch.stack(res_tmp) + .reshape(self._input_numb_field, self._output_numb_field, -1) + .sum(0) + ) conv_transposed[batch_idx, ..., -1] = res_tmp return conv_transposed @@ -460,10 +481,10 @@ class ContinuousConvBlock(BaseContinuousConv): # initialize convolution if self.training: # we choose what to do based on optimization - self._choose_initialization(X, type='inverse') + self._choose_initialization(X, type="inverse") else: # we always initialize on testing - self._initialize_convolution(X, 'inverse') + self._initialize_convolution(X, "inverse") # initialize grid X = self._grid_transpose.clone().detach() @@ -479,11 +500,14 @@ class ContinuousConvBlock(BaseContinuousConv): # accumulator for the convolution on different batches accumulator_batch = torch.zeros( - size=(self._grid_transpose.shape[1], - self._grid_transpose.shape[2]), + size=( + self._grid_transpose.shape[1], + self._grid_transpose.shape[2], + ), requires_grad=True, device=X.device, - dtype=X.dtype).clone() + dtype=X.dtype, + ).clone() for stride_idx, current_stride in enumerate(self._stride): # indeces of points falling into filter range @@ -522,9 +546,10 @@ class ContinuousConvBlock(BaseContinuousConv): staked_output = net(nn_input_pts[idx_channel_out]) # perform integral for all strides in one field - integral = staked_output * integrals[batch_idx, - idx_channel_in, - stride_idx] + integral = ( + staked_output + * integrals[batch_idx, idx_channel_in, stride_idx] + ) # append results res_tmp.append(integral.flatten()) @@ -532,7 +557,7 @@ class ContinuousConvBlock(BaseContinuousConv): channel_sum = [] start = 0 for _ in range(self._output_numb_field): - tmp = res_tmp[start:start + self._input_numb_field] + tmp = res_tmp[start : start + self._input_numb_field] tmp = torch.vstack(tmp).sum(dim=0) channel_sum.append(tmp) start += self._input_numb_field diff --git a/pina/model/layers/fourier.py b/pina/model/layers/fourier.py index ef9e76a..3b6078e 100644 --- a/pina/model/layers/fourier.py +++ b/pina/model/layers/fourier.py @@ -2,14 +2,18 @@ import torch import torch.nn as nn from ...utils import check_consistency -from pina.model.layers import SpectralConvBlock1D, SpectralConvBlock2D, SpectralConvBlock3D +from pina.model.layers import ( + SpectralConvBlock1D, + SpectralConvBlock2D, + SpectralConvBlock3D, +) class FourierBlock1D(nn.Module): """ Fourier block implementation for three dimensional input tensor. The combination of Fourier blocks - make up the Fourier Neural Operator + make up the Fourier Neural Operator .. seealso:: @@ -21,11 +25,13 @@ class FourierBlock1D(nn.Module): """ - def __init__(self, - input_numb_fields, - output_numb_fields, - n_modes, - activation=torch.nn.Tanh): + def __init__( + self, + input_numb_fields, + output_numb_fields, + n_modes, + activation=torch.nn.Tanh, + ): super().__init__() """ PINA implementation of Fourier block one dimension. The module computes @@ -51,17 +57,18 @@ class FourierBlock1D(nn.Module): self._spectral_conv = SpectralConvBlock1D( input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, - n_modes=n_modes) + n_modes=n_modes, + ) self._activation = activation() self._linear = nn.Conv1d(input_numb_fields, output_numb_fields, 1) def forward(self, x): """ - Forward computation for Fourier Block. It performs a spectral + Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results. - :param x: The input tensor for fourier block, expect of size + :param x: The input tensor for fourier block, expect of size ``[batch, input_numb_fields, x]``. :type x: torch.Tensor :return: The output tensor obtained from the @@ -75,7 +82,7 @@ class FourierBlock2D(nn.Module): """ Fourier block implementation for two dimensional input tensor. The combination of Fourier blocks - make up the Fourier Neural Operator + make up the Fourier Neural Operator .. seealso:: @@ -87,18 +94,20 @@ class FourierBlock2D(nn.Module): """ - def __init__(self, - input_numb_fields, - output_numb_fields, - n_modes, - activation=torch.nn.Tanh): + def __init__( + self, + input_numb_fields, + output_numb_fields, + n_modes, + activation=torch.nn.Tanh, + ): """ PINA implementation of Fourier block two dimensions. The module computes the spectral convolution of the input with a linear kernel in the fourier space, and then it maps the input back to the physical space. The output is then added to a Linear tranformation of the input in the physical space. Finally an activation function is - applied to the output. + applied to the output. The block expects an input of size ``[batch, input_numb_fields, Nx, Ny]`` and returns an output of size ``[batch, output_numb_fields, Nx, Ny]``. @@ -118,17 +127,18 @@ class FourierBlock2D(nn.Module): self._spectral_conv = SpectralConvBlock2D( input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, - n_modes=n_modes) + n_modes=n_modes, + ) self._activation = activation() self._linear = nn.Conv2d(input_numb_fields, output_numb_fields, 1) def forward(self, x): """ - Forward computation for Fourier Block. It performs a spectral + Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results. - :param x: The input tensor for fourier block, expect of size + :param x: The input tensor for fourier block, expect of size ``[batch, input_numb_fields, x, y]``. :type x: torch.Tensor :return: The output tensor obtained from the @@ -142,7 +152,7 @@ class FourierBlock3D(nn.Module): """ Fourier block implementation for three dimensional input tensor. The combination of Fourier blocks - make up the Fourier Neural Operator + make up the Fourier Neural Operator .. seealso:: @@ -154,18 +164,20 @@ class FourierBlock3D(nn.Module): """ - def __init__(self, - input_numb_fields, - output_numb_fields, - n_modes, - activation=torch.nn.Tanh): + def __init__( + self, + input_numb_fields, + output_numb_fields, + n_modes, + activation=torch.nn.Tanh, + ): """ PINA implementation of Fourier block three dimensions. The module computes the spectral convolution of the input with a linear kernel in the fourier space, and then it maps the input back to the physical space. The output is then added to a Linear tranformation of the input in the physical space. Finally an activation function is - applied to the output. + applied to the output. The block expects an input of size ``[batch, input_numb_fields, Nx, Ny, Nz]`` and returns an output of size ``[batch, output_numb_fields, Nx, Ny, Nz]``. @@ -186,17 +198,18 @@ class FourierBlock3D(nn.Module): self._spectral_conv = SpectralConvBlock3D( input_numb_fields=input_numb_fields, output_numb_fields=output_numb_fields, - n_modes=n_modes) + n_modes=n_modes, + ) self._activation = activation() self._linear = nn.Conv3d(input_numb_fields, output_numb_fields, 1) def forward(self, x): """ - Forward computation for Fourier Block. It performs a spectral + Forward computation for Fourier Block. It performs a spectral convolution and a linear transformation of the input and sum the results. - :param x: The input tensor for fourier block, expect of size + :param x: The input tensor for fourier block, expect of size ``[batch, input_numb_fields, x, y, z]``. :type x: torch.Tensor :return: The output tensor obtained from the diff --git a/pina/model/layers/integral.py b/pina/model/layers/integral.py index 3269134..565aec3 100644 --- a/pina/model/layers/integral.py +++ b/pina/model/layers/integral.py @@ -10,9 +10,9 @@ class Integral(object): :type param: string """ - if param == 'discrete': + if param == "discrete": self.make_integral = self.integral_param_disc - elif param == 'continuous': + elif param == "continuous": self.make_integral = self.integral_param_cont else: raise TypeError diff --git a/pina/model/layers/pod.py b/pina/model/layers/pod.py index 299b0ac..f696d03 100644 --- a/pina/model/layers/pod.py +++ b/pina/model/layers/pod.py @@ -1,4 +1,5 @@ """Module for Base Continuous Convolution class.""" + from abc import ABCMeta, abstractmethod import torch from .stride import Stride @@ -38,17 +39,17 @@ class PODLayer(torch.nn.Module): :rtype: int """ return self._rank - + @rank.setter def rank(self, value): if value < 1 or not isinstance(value, int): - raise ValueError('The rank must be positive integer') + raise ValueError("The rank must be positive integer") self._rank = value @property def basis(self): - """ + """ The POD basis. It is a matrix whose columns are the first `self.rank` POD modes. :rtype: torch.Tensor @@ -56,7 +57,7 @@ class PODLayer(torch.nn.Module): if self._basis is None: return None - return self._basis[:self.rank] + return self._basis[: self.rank] @property def scaler(self): @@ -67,10 +68,12 @@ class PODLayer(torch.nn.Module): :rtype: dict """ if self._scaler is None: - return + return - return {'mean': self._scaler['mean'][:self.rank], - 'std': self._scaler['std'][:self.rank]} + return { + "mean": self._scaler["mean"][: self.rank], + "std": self._scaler["std"][: self.rank], + } @property def scale_coefficients(self): @@ -105,8 +108,9 @@ class PODLayer(torch.nn.Module): :param torch.Tensor coeffs: The coefficients to be scaled. """ self._scaler = { - 'std': torch.std(coeffs, dim=1), - 'mean': torch.mean(coeffs, dim=1)} + "std": torch.std(coeffs, dim=1), + "mean": torch.mean(coeffs, dim=1), + } def _fit_pod(self, X): """ @@ -114,7 +118,7 @@ class PODLayer(torch.nn.Module): :param torch.Tensor X: The tensor to be reduced. """ - if X.device.type == 'mps': # svd_lowrank not arailable for mps + if X.device.type == "mps": # svd_lowrank not arailable for mps self._basis = torch.svd(X.T)[0].T else: self._basis = torch.svd_lowrank(X.T, q=X.shape[0])[0].T @@ -142,7 +146,8 @@ class PODLayer(torch.nn.Module): """ if self._basis is None: raise RuntimeError( - 'The POD layer needs to be fitted before being used.') + "The POD layer needs to be fitted before being used." + ) coeff = torch.matmul(self.basis, X.T) if coeff.ndim == 1: @@ -150,28 +155,29 @@ class PODLayer(torch.nn.Module): coeff = coeff.T if self.__scale_coefficients: - coeff = (coeff - self.scaler['mean']) / self.scaler['std'] + coeff = (coeff - self.scaler["mean"]) / self.scaler["std"] return coeff def expand(self, coeff): - """ + """ Expand the given coefficients to the original space. The POD layer needs to be fitted before being used. - + :param torch.Tensor coeff: The coefficients to be expanded. :return: The expanded tensor. :rtype: torch.Tensor """ if self._basis is None: raise RuntimeError( - 'The POD layer needs to be trained before being used.') + "The POD layer needs to be trained before being used." + ) if self.__scale_coefficients: - coeff = coeff * self.scaler['std'] + self.scaler['mean'] + coeff = coeff * self.scaler["std"] + self.scaler["mean"] predicted = torch.matmul(self.basis.T, coeff.T).T if predicted.ndim == 1: predicted = predicted.unsqueeze(0) - return predicted \ No newline at end of file + return predicted diff --git a/pina/model/layers/residual.py b/pina/model/layers/residual.py index 6b3734c..edd9b07 100644 --- a/pina/model/layers/residual.py +++ b/pina/model/layers/residual.py @@ -16,18 +16,20 @@ class ResidualBlock(nn.Module): """ - def __init__(self, - input_dim, - output_dim, - hidden_dim, - spectral_norm=False, - activation=torch.nn.ReLU()): + def __init__( + self, + input_dim, + output_dim, + hidden_dim, + spectral_norm=False, + activation=torch.nn.ReLU(), + ): """ Initializes the ResidualBlock module. :param int input_dim: Dimension of the input to pass to the feedforward linear layer. - :param int output_dim: Dimension of the output from the + :param int output_dim: Dimension of the output from the residual layer. :param int hidden_dim: Hidden dimension for mapping the input (first block). @@ -82,6 +84,7 @@ class ResidualBlock(nn.Module): import torch import torch.nn as nn + class EnhancedLinear(torch.nn.Module): """ A wrapper class for enhancing a linear layer with activation and/or dropout. @@ -132,8 +135,9 @@ class EnhancedLinear(torch.nn.Module): self._model = torch.nn.Sequential(layer, self._drop(dropout)) elif (dropout is not None) and (activation is not None): - self._model = torch.nn.Sequential(layer, activation, - self._drop(dropout)) + self._model = torch.nn.Sequential( + layer, activation, self._drop(dropout) + ) def forward(self, x): """ diff --git a/pina/model/layers/spectral.py b/pina/model/layers/spectral.py index d86aa09..674f3e0 100644 --- a/pina/model/layers/spectral.py +++ b/pina/model/layers/spectral.py @@ -37,18 +37,23 @@ class SpectralConvBlock1D(nn.Module): self._output_channels = output_numb_fields # scaling factor - scale = (1. / (self._input_channels * self._output_channels)) - self._weights = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes, - dtype=torch.cfloat)) + scale = 1.0 / (self._input_channels * self._output_channels) + self._weights = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes, + dtype=torch.cfloat, + ) + ) def _compute_mult1d(self, input, weights): """ Compute the matrix multiplication of the input with the linear kernel weights. - :param input: The input tensor, expect of size + :param input: The input tensor, expect of size ``[batch, input_numb_fields, x]``. :type input: torch.Tensor :param weights: The kernel weights, expect of @@ -64,7 +69,7 @@ class SpectralConvBlock1D(nn.Module): """ Forward computation for Spectral Convolution. - :param x: The input tensor, expect of size + :param x: The input tensor, expect of size ``[batch, input_numb_fields, x]``. :type x: torch.Tensor :return: The output tensor obtained from the @@ -77,13 +82,16 @@ class SpectralConvBlock1D(nn.Module): x_ft = torch.fft.rfft(x) # Multiply relevant Fourier modes - out_ft = torch.zeros(batch_size, - self._output_channels, - x.size(-1) // 2 + 1, - device=x.device, - dtype=torch.cfloat) - out_ft[:, :, :self._modes] = self._compute_mult1d( - x_ft[:, :, :self._modes], self._weights) + out_ft = torch.zeros( + batch_size, + self._output_channels, + x.size(-1) // 2 + 1, + device=x.device, + dtype=torch.cfloat, + ) + out_ft[:, :, : self._modes] = self._compute_mult1d( + x_ft[:, :, : self._modes], self._weights + ) # Return to physical space return torch.fft.irfft(out_ft, n=x.size(-1)) @@ -119,17 +127,19 @@ class SpectralConvBlock2D(nn.Module): if isinstance(n_modes, (tuple, list)): if len(n_modes) != 2: raise ValueError( - 'Expected n_modes to be a list or tuple of len two, ' - 'with each entry corresponding to the number of modes ' - 'for each dimension ') + "Expected n_modes to be a list or tuple of len two, " + "with each entry corresponding to the number of modes " + "for each dimension " + ) elif isinstance(n_modes, int): n_modes = [n_modes] * 2 else: raise ValueError( - 'Expected n_modes to be a list or tuple of len two, ' - 'with each entry corresponding to the number of modes ' - 'for each dimension; or an int value representing the ' - 'number of modes for all dimensions') + "Expected n_modes to be a list or tuple of len two, " + "with each entry corresponding to the number of modes " + "for each dimension; or an int value representing the " + "number of modes for all dimensions" + ) # assign variables self._modes = n_modes @@ -137,24 +147,34 @@ class SpectralConvBlock2D(nn.Module): self._output_channels = output_numb_fields # scaling factor - scale = (1. / (self._input_channels * self._output_channels)) - self._weights1 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - dtype=torch.cfloat)) - self._weights2 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - dtype=torch.cfloat)) + scale = 1.0 / (self._input_channels * self._output_channels) + self._weights1 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + dtype=torch.cfloat, + ) + ) + self._weights2 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + dtype=torch.cfloat, + ) + ) def _compute_mult2d(self, input, weights): """ Compute the matrix multiplication of the input with the linear kernel weights. - :param input: The input tensor, expect of size + :param input: The input tensor, expect of size ``[batch, input_numb_fields, x, y]``. :type input: torch.Tensor :param weights: The kernel weights, expect of @@ -170,7 +190,7 @@ class SpectralConvBlock2D(nn.Module): """ Forward computation for Spectral Convolution. - :param x: The input tensor, expect of size + :param x: The input tensor, expect of size ``[batch, input_numb_fields, x, y]``. :type x: torch.Tensor :return: The output tensor obtained from the @@ -184,16 +204,22 @@ class SpectralConvBlock2D(nn.Module): x_ft = torch.fft.rfft2(x) # Multiply relevant Fourier modes - out_ft = torch.zeros(batch_size, - self._output_channels, - x.size(-2), - x.size(-1) // 2 + 1, - device=x.device, - dtype=torch.cfloat) - out_ft[:, :, :self._modes[0], :self._modes[1]] = self._compute_mult2d( - x_ft[:, :, :self._modes[0], :self._modes[1]], self._weights1) - out_ft[:, :, -self._modes[0]:, :self._modes[1]:] = self._compute_mult2d( - x_ft[:, :, -self._modes[0]:, :self._modes[1]], self._weights2) + out_ft = torch.zeros( + batch_size, + self._output_channels, + x.size(-2), + x.size(-1) // 2 + 1, + device=x.device, + dtype=torch.cfloat, + ) + out_ft[:, :, : self._modes[0], : self._modes[1]] = self._compute_mult2d( + x_ft[:, :, : self._modes[0], : self._modes[1]], self._weights1 + ) + out_ft[:, :, -self._modes[0] :, : self._modes[1] :] = ( + self._compute_mult2d( + x_ft[:, :, -self._modes[0] :, : self._modes[1]], self._weights2 + ) + ) # Return to physical space return torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1))) @@ -230,17 +256,19 @@ class SpectralConvBlock3D(nn.Module): if isinstance(n_modes, (tuple, list)): if len(n_modes) != 3: raise ValueError( - 'Expected n_modes to be a list or tuple of len three, ' - 'with each entry corresponding to the number of modes ' - 'for each dimension ') + "Expected n_modes to be a list or tuple of len three, " + "with each entry corresponding to the number of modes " + "for each dimension " + ) elif isinstance(n_modes, int): n_modes = [n_modes] * 3 else: raise ValueError( - 'Expected n_modes to be a list or tuple of len three, ' - 'with each entry corresponding to the number of modes ' - 'for each dimension; or an int value representing the ' - 'number of modes for all dimensions') + "Expected n_modes to be a list or tuple of len three, " + "with each entry corresponding to the number of modes " + "for each dimension; or an int value representing the " + "number of modes for all dimensions" + ) # assign variables self._modes = n_modes @@ -248,38 +276,58 @@ class SpectralConvBlock3D(nn.Module): self._output_channels = output_numb_fields # scaling factor - scale = (1. / (self._input_channels * self._output_channels)) - self._weights1 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - self._modes[2], - dtype=torch.cfloat)) - self._weights2 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - self._modes[2], - dtype=torch.cfloat)) - self._weights3 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - self._modes[2], - dtype=torch.cfloat)) - self._weights4 = nn.Parameter(scale * torch.rand(self._input_channels, - self._output_channels, - self._modes[0], - self._modes[1], - self._modes[2], - dtype=torch.cfloat)) + scale = 1.0 / (self._input_channels * self._output_channels) + self._weights1 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + self._modes[2], + dtype=torch.cfloat, + ) + ) + self._weights2 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + self._modes[2], + dtype=torch.cfloat, + ) + ) + self._weights3 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + self._modes[2], + dtype=torch.cfloat, + ) + ) + self._weights4 = nn.Parameter( + scale + * torch.rand( + self._input_channels, + self._output_channels, + self._modes[0], + self._modes[1], + self._modes[2], + dtype=torch.cfloat, + ) + ) def _compute_mult3d(self, input, weights): """ Compute the matrix multiplication of the input with the linear kernel weights. - :param input: The input tensor, expect of size + :param input: The input tensor, expect of size ``[batch, input_numb_fields, x, y, z]``. :type input: torch.Tensor :param weights: The kernel weights, expect of @@ -295,7 +343,7 @@ class SpectralConvBlock3D(nn.Module): """ Forward computation for Spectral Convolution. - :param x: The input tensor, expect of size + :param x: The input tensor, expect of size ``[batch, input_numb_fields, x, y, z]``. :type x: torch.Tensor :return: The output tensor obtained from the @@ -309,13 +357,15 @@ class SpectralConvBlock3D(nn.Module): x_ft = torch.fft.rfftn(x, dim=[-3, -2, -1]) # Multiply relevant Fourier modes - out_ft = torch.zeros(batch_size, - self._output_channels, - x.size(-3), - x.size(-2), - x.size(-1) // 2 + 1, - device=x.device, - dtype=torch.cfloat) + out_ft = torch.zeros( + batch_size, + self._output_channels, + x.size(-3), + x.size(-2), + x.size(-1) // 2 + 1, + device=x.device, + dtype=torch.cfloat, + ) slice0 = ( slice(None), diff --git a/pina/model/layers/stride.py b/pina/model/layers/stride.py index 6facb01..7832ac4 100644 --- a/pina/model/layers/stride.py +++ b/pina/model/layers/stride.py @@ -60,7 +60,8 @@ class Stride(object): if seq_direction != seq_jumps: raise IndexError( - "direction and jumps must have zero in the same index") + "direction and jumps must have zero in the same index" + ) if seq_jumps: for i in seq_jumps: diff --git a/pina/model/layers/utils_convolution.py b/pina/model/layers/utils_convolution.py index a2df5fe..5442ff4 100644 --- a/pina/model/layers/utils_convolution.py +++ b/pina/model/layers/utils_convolution.py @@ -3,8 +3,9 @@ import torch def check_point(x, current_stride, dim): max_stride = current_stride + dim - indeces = torch.logical_and(x[..., :-1] < max_stride, x[..., :-1] - >= current_stride).all(dim=-1) + indeces = torch.logical_and( + x[..., :-1] < max_stride, x[..., :-1] >= current_stride + ).all(dim=-1) return indeces @@ -32,12 +33,12 @@ def optimizing(f): def wrapper(*args, **kwargs): - if kwargs['type'] == 'forward': + if kwargs["type"] == "forward": if not wrapper.has_run_inverse: wrapper.has_run_inverse = True return f(*args, **kwargs) - if kwargs['type'] == 'inverse': + if kwargs["type"] == "inverse": if not wrapper.has_run: wrapper.has_run = True return f(*args, **kwargs) diff --git a/pina/model/multi_feed_forward.py b/pina/model/multi_feed_forward.py index 20d0240..b04708d 100644 --- a/pina/model/multi_feed_forward.py +++ b/pina/model/multi_feed_forward.py @@ -1,4 +1,5 @@ """Module for Multi FeedForward model""" + import torch from .feed_forward import FeedForward @@ -6,7 +7,7 @@ from .feed_forward import FeedForward class MultiFeedForward(torch.nn.Module): """ - The PINA implementation of MultiFeedForward network. + The PINA implementation of MultiFeedForward network. This model allows to create a network with multiple FeedForward combined together. The user has to define the `forward` method choosing how to diff --git a/pina/model/network.py b/pina/model/network.py index 5c5c23d..6fde803 100644 --- a/pina/model/network.py +++ b/pina/model/network.py @@ -6,13 +6,15 @@ from ..label_tensor import LabelTensor class Network(torch.nn.Module): - def __init__(self, model, input_variables, output_variables, extra_features=None): + def __init__( + self, model, input_variables, output_variables, extra_features=None + ): """ Network class with standard forward method and possibility to pass extra features. This class is used internally in PINA to convert any :class:`torch.nn.Module` s in a PINA module. - + :param model: The torch model to convert in a PINA model. :type model: torch.nn.Module :param list(str) input_variables: The input variables of the :class:`AbstractProblem`, whose type depends on the @@ -57,7 +59,9 @@ class Network(torch.nn.Module): :return torch.Tensor: Output of the network. """ # only labeltensors as input - assert isinstance(x, LabelTensor), "Expected LabelTensor as input to the model." + assert isinstance( + x, LabelTensor + ), "Expected LabelTensor as input to the model." # extract torch.Tensor from corresponding label # in case `input_variables = []` all points are used @@ -75,7 +79,7 @@ class Network(torch.nn.Module): output.labels = self._output_variables return output - + # TODO to remove in next releases (only used in GAROM solver) def forward_map(self, x): """ diff --git a/pina/operators.py b/pina/operators.py index 563188c..17b45d8 100644 --- a/pina/operators.py +++ b/pina/operators.py @@ -5,6 +5,7 @@ All operators take as input a tensor onto which computing the operator, a tensor to which computing the operator, the name of the output variables to calculate the operator for (in case of multidimensional functions), and the variables name on which the operator is calculated. """ + import torch from pina.label_tensor import LabelTensor @@ -49,24 +50,25 @@ def grad(output_, input_, components=None, d=None): """ if len(output_.labels) != 1: - raise RuntimeError('only scalar function can be differentiated') + raise RuntimeError("only scalar function can be differentiated") if not all([di in input_.labels for di in d]): - raise RuntimeError('derivative labels missing from input tensor') + raise RuntimeError("derivative labels missing from input tensor") output_fieldname = output_.labels[0] - gradients = torch.autograd.grad(output_, - input_, - grad_outputs=torch.ones( - output_.size(), - dtype=output_.dtype, - device=output_.device), - create_graph=True, - retain_graph=True, - allow_unused=True)[0] + gradients = torch.autograd.grad( + output_, + input_, + grad_outputs=torch.ones( + output_.size(), dtype=output_.dtype, device=output_.device + ), + create_graph=True, + retain_graph=True, + allow_unused=True, + )[0] gradients.labels = input_.labels gradients = gradients.extract(d) - gradients.labels = [f'd{output_fieldname}d{i}' for i in d] + gradients.labels = [f"d{output_fieldname}d{i}" for i in d] return gradients @@ -93,7 +95,8 @@ def grad(output_, input_, components=None, d=None): gradients = grad_scalar_output(c_output, input_, d) else: gradients = gradients.append( - grad_scalar_output(c_output, input_, d)) + grad_scalar_output(c_output, input_, d) + ) else: raise NotImplementedError @@ -133,7 +136,7 @@ def div(output_, input_, components=None, d=None): components = output_.labels if output_.shape[1] < 2 or len(components) < 2: - raise ValueError('div supported only for vector fields') + raise ValueError("div supported only for vector fields") if len(components) != len(d): raise ValueError @@ -142,16 +145,16 @@ def div(output_, input_, components=None, d=None): div = torch.zeros(input_.shape[0], 1, device=output_.device) labels = [None] * len(components) for i, (c, d) in enumerate(zip(components, d)): - c_fields = f'd{c}d{d}' + c_fields = f"d{c}d{d}" div[:, 0] += grad_output.extract(c_fields).sum(axis=1) labels[i] = c_fields div = div.as_subclass(LabelTensor) - div.labels = ['+'.join(labels)] + div.labels = ["+".join(labels)] return div -def laplacian(output_, input_, components=None, d=None, method='std'): +def laplacian(output_, input_, components=None, d=None, method="std"): """ Compute Laplace operator. The operator works for vectorial and scalar functions, with multiple input coordinates. @@ -182,26 +185,27 @@ def laplacian(output_, input_, components=None, d=None, method='std'): if len(components) != len(d) and len(components) != 1: raise ValueError - if method == 'divgrad': - raise NotImplementedError('divgrad not implemented as method') + if method == "divgrad": + raise NotImplementedError("divgrad not implemented as method") # TODO fix # grad_output = grad(output_, input_, components, d) # result = div(grad_output, input_, d=d) - elif method == 'std': + elif method == "std": if len(components) == 1: grad_output = grad(output_, input_, components=components, d=d) result = torch.zeros(output_.shape[0], 1, device=output_.device) for i, label in enumerate(grad_output.labels): gg = grad(grad_output, input_, d=d, components=[label]) - result[:, 0] += super(torch.Tensor, - gg.T).__getitem__(i) # TODO improve - labels = [f'dd{components[0]}'] + result[:, 0] += super(torch.Tensor, gg.T).__getitem__( + i + ) # TODO improve + labels = [f"dd{components[0]}"] else: - result = torch.empty(input_.shape[0], - len(components), - device=output_.device) + result = torch.empty( + input_.shape[0], len(components), device=output_.device + ) labels = [None] * len(components) for idx, (ci, di) in enumerate(zip(components, d)): @@ -212,7 +216,7 @@ def laplacian(output_, input_, components=None, d=None, method='std'): grad_output = grad(output_, input_, components=ci, d=di) result[:, idx] = grad(grad_output, input_, d=di).flatten() - labels[idx] = f'dd{ci}dd{di}' + labels[idx] = f"dd{ci}dd{di}" result = result.as_subclass(LabelTensor) result.labels = labels @@ -245,8 +249,11 @@ def advection(output_, input_, velocity_field, components=None, d=None): if components is None: components = output_.labels - tmp = grad(output_, input_, components, d).reshape(-1, len(components), - len(d)).transpose(0, 1) + tmp = ( + grad(output_, input_, components, d) + .reshape(-1, len(components), len(d)) + .transpose(0, 1) + ) tmp *= output_.extract(velocity_field) return tmp.sum(dim=2).T diff --git a/pina/plotter.py b/pina/plotter.py index d00b44d..63d8646 100644 --- a/pina/plotter.py +++ b/pina/plotter.py @@ -34,30 +34,33 @@ class Plotter: if variables is None: variables = problem.domain.variables - elif variables == 'spatial': + elif variables == "spatial": variables = problem.spatial_domain.variables - elif variables == 'temporal': + elif variables == "temporal": variables = problem.temporal_domain.variables if len(variables) not in [1, 2, 3]: - raise ValueError('Samples can be plotted only in ' - 'dimensions 1, 2 and 3.') + raise ValueError( + "Samples can be plotted only in " "dimensions 1, 2 and 3." + ) fig = plt.figure() - proj = '3d' if len(variables) == 3 else None + proj = "3d" if len(variables) == 3 else None ax = fig.add_subplot(projection=proj) for location in problem.input_pts: coords = problem.input_pts[location].extract(variables).T.detach() - if len(variables)==1: # 1D samples - ax.plot(coords.flatten(), - torch.zeros(coords.flatten().shape), - '.', - label=location, - **kwargs) - elif len(variables)==2: - ax.plot(*coords, '.', label=location, **kwargs) - elif len(variables)==3: - ax.scatter(*coords, '.', label=location, **kwargs) + if len(variables) == 1: # 1D samples + ax.plot( + coords.flatten(), + torch.zeros(coords.flatten().shape), + ".", + label=location, + **kwargs, + ) + elif len(variables) == 2: + ax.plot(*coords, ".", label=location, **kwargs) + elif len(variables) == 3: + ax.scatter(*coords, ".", label=location, **kwargs) ax.set_xlabel(variables[0]) try: @@ -94,27 +97,23 @@ class Plotter: """ fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 8)) - ax.plot(pts.extract(v), pred, label='Neural Network solution', **kwargs) + ax.plot(pts.extract(v), pred, label="Neural Network solution", **kwargs) if truth_solution: truth_output = truth_solution(pts).detach() - ax.plot(pts.extract(v), truth_output, - label='True solution', **kwargs) + ax.plot( + pts.extract(v), truth_output, label="True solution", **kwargs + ) # TODO: pred is a torch.Tensor, so no labels is available # extra variable for labels should be # passed in the function arguments. - # plt.ylabel(pred.labels[0]) + # plt.ylabel(pred.labels[0]) plt.legend() - def _2d_plot(self, - pts, - pred, - v, - res, - method, - truth_solution=None, - **kwargs): + def _2d_plot( + self, pts, pred, v, res, method, truth_solution=None, **kwargs + ): """Plot solution for two dimensional function :param pts: Points to plot the solution. @@ -136,44 +135,47 @@ class Plotter: pred_output = pred.reshape(res, res) if truth_solution: - truth_output = truth_solution(pts).float().reshape(res, res).as_subclass(torch.Tensor) + truth_output = ( + truth_solution(pts) + .float() + .reshape(res, res) + .as_subclass(torch.Tensor) + ) fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(16, 6)) - cb = getattr(ax[0], method)(*grids, pred_output, - **kwargs) + cb = getattr(ax[0], method)(*grids, pred_output, **kwargs) fig.colorbar(cb, ax=ax[0]) - ax[0].title.set_text('Neural Network prediction') - cb = getattr(ax[1], method)(*grids, truth_output, - **kwargs) + ax[0].title.set_text("Neural Network prediction") + cb = getattr(ax[1], method)(*grids, truth_output, **kwargs) fig.colorbar(cb, ax=ax[1]) - ax[1].title.set_text('True solution') - cb = getattr(ax[2], - method)(*grids, - (truth_output - pred_output), - **kwargs) + ax[1].title.set_text("True solution") + cb = getattr(ax[2], method)( + *grids, (truth_output - pred_output), **kwargs + ) fig.colorbar(cb, ax=ax[2]) - ax[2].title.set_text('Residual') + ax[2].title.set_text("Residual") else: fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 6)) - cb = getattr(ax, method)(*grids, pred_output, - **kwargs) + cb = getattr(ax, method)(*grids, pred_output, **kwargs) fig.colorbar(cb, ax=ax) - ax.title.set_text('Neural Network prediction') + ax.title.set_text("Neural Network prediction") - def plot(self, - solver, - components=None, - fixed_variables={}, - method='contourf', - res=256, - filename=None, - **kwargs): + def plot( + self, + solver, + components=None, + fixed_variables={}, + method="contourf", + res=256, + filename=None, + **kwargs, + ): """ Plot sample of SolverInterface output. :param SolverInterface solver: The ``SolverInterface`` object instance. - :param str | list(str) components: The output variable(s) to plot. - If None, all the output variables of the problem are selected. + :param str | list(str) components: The output variable(s) to plot. + If None, all the output variables of the problem are selected. Default value is None. :param dict fixed_variables: A dictionary with all the variables that should be kept fixed during the plot. The keys of the dictionary @@ -190,23 +192,28 @@ class Plotter: if components is None: components = solver.problem.output_variables - + if isinstance(components, str): components = [components] if not isinstance(components, list): - raise NotImplementedError('Output variables must be passed' - 'as a string or a list of strings.') - + raise NotImplementedError( + "Output variables must be passed" + "as a string or a list of strings." + ) + if len(components) > 1: - raise NotImplementedError('Multidimensional plots are not implemented, ' - 'set components to an available components of' - ' the problem.') + raise NotImplementedError( + "Multidimensional plots are not implemented, " + "set components to an available components of" + " the problem." + ) v = [ - var for var in solver.problem.input_variables + var + for var in solver.problem.input_variables if var not in fixed_variables.keys() ] - pts = solver.problem.domain.sample(res, 'grid', variables=v) + pts = solver.problem.domain.sample(res, "grid", variables=v) fixed_pts = torch.ones(pts.shape[0], len(fixed_variables)) fixed_pts *= torch.tensor(list(fixed_variables.values())) @@ -218,16 +225,20 @@ class Plotter: # computing soluting and sending to cpu predicted_output = solver.forward(pts).extract(components) - predicted_output = predicted_output.as_subclass(torch.Tensor).cpu().detach() + predicted_output = ( + predicted_output.as_subclass(torch.Tensor).cpu().detach() + ) pts = pts.cpu() - truth_solution = getattr(solver.problem, 'truth_solution', None) + truth_solution = getattr(solver.problem, "truth_solution", None) if len(v) == 1: - self._1d_plot(pts, predicted_output, v, method, truth_solution, - **kwargs) + self._1d_plot( + pts, predicted_output, v, method, truth_solution, **kwargs + ) elif len(v) == 2: - self._2d_plot(pts, predicted_output, v, res, method, truth_solution, - **kwargs) + self._2d_plot( + pts, predicted_output, v, res, method, truth_solution, **kwargs + ) plt.tight_layout() if filename: @@ -236,13 +247,15 @@ class Plotter: else: plt.show() - def plot_loss(self, - trainer, - metrics=None, - logy=False, - logx=False, - filename=None, - **kwargs): + def plot_loss( + self, + trainer, + metrics=None, + logy=False, + logx=False, + filename=None, + **kwargs, + ): """ Plot the loss function values during traininig. @@ -260,41 +273,43 @@ class Plotter: # check that MetricTracker has been used list_ = [ - idx for idx, s in enumerate(trainer.callbacks) + idx + for idx, s in enumerate(trainer.callbacks) if isinstance(s, MetricTracker) ] if not bool(list_): raise FileNotFoundError( - 'MetricTracker should be used as a callback during training to' - ' use this method.') + "MetricTracker should be used as a callback during training to" + " use this method." + ) # extract trainer metrics trainer_metrics = trainer.callbacks[list_[0]].metrics if metrics is None: - metrics = ['mean_loss'] + metrics = ["mean_loss"] elif not isinstance(metrics, list): - raise ValueError('metrics must be class list.') + raise ValueError("metrics must be class list.") # loop over metrics to plot for metric in metrics: if metric not in trainer_metrics: raise ValueError( - f'{metric} not a valid metric. Available metrics are {list(trainer_metrics.keys())}.' + f"{metric} not a valid metric. Available metrics are {list(trainer_metrics.keys())}." ) loss = trainer_metrics[metric] epochs = range(len(loss)) plt.plot(epochs, loss.cpu(), **kwargs) # plotting - plt.xlabel('epoch') - plt.ylabel('loss') + plt.xlabel("epoch") + plt.ylabel("loss") plt.legend() # log axis if logy: - plt.yscale('log') + plt.yscale("log") if logx: - plt.xscale('log') + plt.xscale("log") # saving in file if filename: diff --git a/pina/problem/__init__.py b/pina/problem/__init__.py index 680451c..35251aa 100644 --- a/pina/problem/__init__.py +++ b/pina/problem/__init__.py @@ -1,9 +1,9 @@ __all__ = [ - 'AbstractProblem', - 'SpatialProblem', - 'TimeDependentProblem', - 'ParametricProblem', - 'InverseProblem', + "AbstractProblem", + "SpatialProblem", + "TimeDependentProblem", + "ParametricProblem", + "InverseProblem", ] from .abstract_problem import AbstractProblem diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py index 311dbce..a368b40 100644 --- a/pina/problem/abstract_problem.py +++ b/pina/problem/abstract_problem.py @@ -1,4 +1,5 @@ """ Module for AbstractProblem class """ + from abc import ABCMeta, abstractmethod from ..utils import merge_tensors, check_consistency import torch @@ -40,13 +41,13 @@ class AbstractProblem(metaclass=ABCMeta): """ variables = [] - if hasattr(self, 'spatial_variables'): + if hasattr(self, "spatial_variables"): variables += self.spatial_variables - if hasattr(self, 'temporal_variable'): + if hasattr(self, "temporal_variable"): variables += self.temporal_variable - if hasattr(self, 'parameters'): + if hasattr(self, "parameters"): variables += self.parameters - if hasattr(self, 'custom_variables'): + if hasattr(self, "custom_variables"): variables += self.custom_variables return variables @@ -62,9 +63,9 @@ class AbstractProblem(metaclass=ABCMeta): :rtype: list[Location] """ domains = [ - getattr(self, f'{t}_domain') - for t in ['spatial', 'temporal', 'parameter'] - if hasattr(self, f'{t}_domain') + getattr(self, f"{t}_domain") + for t in ["spatial", "temporal", "parameter"] + if hasattr(self, f"{t}_domain") ] if len(domains) == 1: @@ -77,7 +78,7 @@ class AbstractProblem(metaclass=ABCMeta): [domain.update(d) for d in domains] return domain else: - raise RuntimeError('different domains') + raise RuntimeError("different domains") @input_variables.setter def input_variables(self, variables): @@ -105,24 +106,27 @@ class AbstractProblem(metaclass=ABCMeta): """ for condition_name in self.conditions: condition = self.conditions[condition_name] - if hasattr(condition, 'input_points'): + if hasattr(condition, "input_points"): samples = condition.input_points self.input_pts[condition_name] = samples self._have_sampled_points[condition_name] = True - if hasattr(self, 'unknown_parameter_domain'): + if hasattr(self, "unknown_parameter_domain"): # initialize the unknown parameters of the inverse problem given # the domain the user gives self.unknown_parameters = {} for i, var in enumerate(self.unknown_variables): range_var = self.unknown_parameter_domain.range_[var] - tensor_var = torch.rand(1, requires_grad=True) * range_var[1] + range_var[0] - self.unknown_parameters[var] = torch.nn.Parameter(tensor_var) + tensor_var = ( + torch.rand(1, requires_grad=True) * range_var[1] + + range_var[0] + ) + self.unknown_parameters[var] = torch.nn.Parameter( + tensor_var + ) - def discretise_domain(self, - n, - mode='random', - variables='all', - locations='all'): + def discretise_domain( + self, n, mode="random", variables="all", locations="all" + ): """ Generate a set of points to span the `Location` of all the conditions of the problem. @@ -157,28 +161,32 @@ class AbstractProblem(metaclass=ABCMeta): # check consistency mode check_consistency(mode, str) - if mode not in ['random', 'grid', 'lh', 'chebyshev', 'latin']: - raise TypeError(f'mode {mode} not valid.') + if mode not in ["random", "grid", "lh", "chebyshev", "latin"]: + raise TypeError(f"mode {mode} not valid.") # check consistency variables - if variables == 'all': + if variables == "all": variables = self.input_variables else: check_consistency(variables, str) if sorted(variables) != sorted(self.input_variables): - TypeError(f'Wrong variables for sampling. Variables ', - f'should be in {self.input_variables}.') + TypeError( + f"Wrong variables for sampling. Variables ", + f"should be in {self.input_variables}.", + ) # check consistency location - if locations == 'all': + if locations == "all": locations = [condition for condition in self.conditions] else: check_consistency(locations, str) if sorted(locations) != sorted(self.conditions): - TypeError(f'Wrong locations for sampling. Location ', - f'should be in {self.conditions}.') + TypeError( + f"Wrong locations for sampling. Location ", + f"should be in {self.conditions}.", + ) # sampling for location in locations: @@ -208,10 +216,10 @@ class AbstractProblem(metaclass=ABCMeta): # the condition is sampled if input_pts contains all labels if sorted(self.input_pts[location].labels) == sorted( - self.input_variables): + self.input_variables + ): self._have_sampled_points[location] = True - def add_points(self, new_points): """ Adding points to the already sampled points. @@ -221,8 +229,10 @@ class AbstractProblem(metaclass=ABCMeta): """ if sorted(new_points.keys()) != sorted(self.conditions): - TypeError(f'Wrong locations for new points. Location ', - f'should be in {self.conditions}.') + TypeError( + f"Wrong locations for new points. Location ", + f"should be in {self.conditions}.", + ) for location in new_points.keys(): # extract old and new points @@ -231,11 +241,14 @@ class AbstractProblem(metaclass=ABCMeta): # if they don't have the same variables error if sorted(old_pts.labels) != sorted(new_pts.labels): - TypeError(f'Not matching variables for old and new points ' - f'in condition {location}.') + TypeError( + f"Not matching variables for old and new points " + f"in condition {location}." + ) if old_pts.labels != new_pts.labels: new_pts = torch.hstack( - [new_pts.extract([i]) for i in old_pts.labels]) + [new_pts.extract([i]) for i in old_pts.labels] + ) new_pts.labels = old_pts.labels # merging @@ -266,4 +279,3 @@ class AbstractProblem(metaclass=ABCMeta): if not is_sample: not_sampled.append(condition_name) return not_sampled - diff --git a/pina/problem/inverse_problem.py b/pina/problem/inverse_problem.py index b9efd6b..5a83566 100644 --- a/pina/problem/inverse_problem.py +++ b/pina/problem/inverse_problem.py @@ -1,4 +1,5 @@ """Module for the ParametricProblem class""" + from abc import abstractmethod from .abstract_problem import AbstractProblem @@ -68,4 +69,3 @@ class InverseProblem(AbstractProblem): @unknown_parameters.setter def unknown_parameters(self, value): self.__unknown_parameters = value - diff --git a/pina/problem/parametric_problem.py b/pina/problem/parametric_problem.py index 17b8ed6..600eab0 100644 --- a/pina/problem/parametric_problem.py +++ b/pina/problem/parametric_problem.py @@ -1,4 +1,5 @@ """Module for the ParametricProblem class""" + from abc import abstractmethod from .abstract_problem import AbstractProblem diff --git a/pina/problem/spatial_problem.py b/pina/problem/spatial_problem.py index 67a1507..e344142 100644 --- a/pina/problem/spatial_problem.py +++ b/pina/problem/spatial_problem.py @@ -1,4 +1,5 @@ """Module for the SpatialProblem class""" + from abc import abstractmethod from .abstract_problem import AbstractProblem diff --git a/pina/problem/timedep_problem.py b/pina/problem/timedep_problem.py index ee34383..cefdb54 100644 --- a/pina/problem/timedep_problem.py +++ b/pina/problem/timedep_problem.py @@ -1,4 +1,5 @@ """Module for the TimeDependentProblem class""" + from abc import abstractmethod from .abstract_problem import AbstractProblem diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py index fdb6219..0562dc2 100644 --- a/pina/solvers/__init__.py +++ b/pina/solvers/__init__.py @@ -1,10 +1,4 @@ -__all__ = [ - 'PINN', - 'GAROM', - 'SupervisedSolver', - 'SolverInterface' - -] +__all__ = ["PINN", "GAROM", "SupervisedSolver", "SolverInterface"] from .garom import GAROM from .pinn import PINN diff --git a/pina/solvers/garom.py b/pina/solvers/garom.py index d7781bd..0885670 100644 --- a/pina/solvers/garom.py +++ b/pina/solvers/garom.py @@ -2,10 +2,13 @@ import torch import sys + 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 ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 from torch.optim.lr_scheduler import ConstantLR from .solver import SolverInterface @@ -18,12 +21,12 @@ class GAROM(SolverInterface): """ GAROM solver class. This class implements Generative Adversarial Reduced Order Model solver, using user specified ``models`` to solve - a specific order reduction``problem``. + a specific order reduction``problem``. .. seealso:: **Original reference**: Coscia, D., Demo, N., & Rozza, G. (2023). - *Generative Adversarial Reduced Order Modelling*. + *Generative Adversarial Reduced Order Modelling*. DOI: `arXiv preprint arXiv:2305.15881. `_. """ @@ -35,19 +38,13 @@ class GAROM(SolverInterface): discriminator, loss=None, optimizer_generator=torch.optim.Adam, - optimizer_generator_kwargs={'lr': 0.001}, + optimizer_generator_kwargs={"lr": 0.001}, optimizer_discriminator=torch.optim.Adam, - optimizer_discriminator_kwargs={'lr': 0.001}, + optimizer_discriminator_kwargs={"lr": 0.001}, scheduler_generator=ConstantLR, - scheduler_generator_kwargs={ - "factor": 1, - "total_iters": 0 - }, + scheduler_generator_kwargs={"factor": 1, "total_iters": 0}, scheduler_discriminator=ConstantLR, - scheduler_discriminator_kwargs={ - "factor": 1, - "total_iters": 0 - }, + scheduler_discriminator_kwargs={"factor": 1, "total_iters": 0}, gamma=0.3, lambda_k=0.001, regularizer=False, @@ -95,8 +92,10 @@ class GAROM(SolverInterface): problem=problem, optimizers=[optimizer_generator, optimizer_discriminator], optimizers_kwargs=[ - optimizer_generator_kwargs, optimizer_discriminator_kwargs - ]) + optimizer_generator_kwargs, + optimizer_discriminator_kwargs, + ], + ) # set automatic optimization for GANs self.automatic_optimization = False @@ -118,13 +117,14 @@ class GAROM(SolverInterface): # assign schedulers self._schedulers = [ scheduler_generator( - self.optimizers[0], **scheduler_generator_kwargs), + self.optimizers[0], **scheduler_generator_kwargs + ), scheduler_discriminator( - self.optimizers[1], - **scheduler_discriminator_kwargs) + self.optimizers[1], **scheduler_discriminator_kwargs + ), ] - # loss and writer + # loss and writer self._loss = loss # began hyperparameters @@ -141,7 +141,7 @@ class GAROM(SolverInterface): :param x: The input tensor. :type x: torch.Tensor - :param mc_steps: Number of montecarlo samples to approximate the + :param mc_steps: Number of montecarlo samples to approximate the expected value, defaults to 20. :type mc_steps: int :param variance: Returining also the sample variance of the solution, defaults to False. @@ -189,8 +189,12 @@ class GAROM(SolverInterface): # generator loss r_loss = self._loss(snapshots, generated_snapshots) - d_fake = self.discriminator.forward_map([generated_snapshots, parameters]) - g_loss = self._loss(d_fake, generated_snapshots) + self.regularizer * r_loss + d_fake = self.discriminator.forward_map( + [generated_snapshots, parameters] + ) + g_loss = ( + self._loss(d_fake, generated_snapshots) + self.regularizer * r_loss + ) # backward step g_loss.backward() @@ -210,7 +214,9 @@ class GAROM(SolverInterface): # Discriminator pass d_real = self.discriminator.forward_map([snapshots, parameters]) - d_fake = self.discriminator.forward_map([generated_snapshots, parameters]) + d_fake = self.discriminator.forward_map( + [generated_snapshots, parameters] + ) # evaluate loss d_loss_real = self._loss(d_real, snapshots) @@ -235,7 +241,7 @@ class GAROM(SolverInterface): self.k += self.lambda_k * diff.item() self.k = min(max(self.k, 0), 1) # Constraint to interval [0, 1] return diff - + def training_step(self, batch, batch_idx): """GAROM solver training step. @@ -248,42 +254,75 @@ class GAROM(SolverInterface): """ dataloader = self.trainer.train_dataloader - condition_idx = batch['condition'] + condition_idx = batch["condition"] - for condition_id in range(condition_idx.min(), condition_idx.max()+1): + for condition_id in range(condition_idx.min(), condition_idx.max() + 1): if sys.version_info >= (3, 8): condition_name = dataloader.condition_names[condition_id] else: - condition_name = dataloader.loaders.condition_names[condition_id] + condition_name = dataloader.loaders.condition_names[ + condition_id + ] condition = self.problem.conditions[condition_name] - pts = batch['pts'].detach() - out = batch['output'] + pts = batch["pts"].detach() + out = batch["output"] if condition_name not in self.problem.conditions: - raise RuntimeError('Something wrong happened.') + raise RuntimeError("Something wrong happened.") # for data driven mode - if not hasattr(condition, 'output_points'): - raise NotImplementedError('GAROM works only in data-driven mode.') + if not hasattr(condition, "output_points"): + raise NotImplementedError( + "GAROM works only in data-driven mode." + ) # get data snapshots = out[condition_idx == condition_id] parameters = pts[condition_idx == condition_id] d_loss_real, d_loss_fake, d_loss = self._train_discriminator( - parameters, snapshots) + parameters, snapshots + ) r_loss, g_loss = self._train_generator(parameters, snapshots) - + diff = self._update_weights(d_loss_real, d_loss_fake) # logging - self.log('mean_loss', float(r_loss), prog_bar=True, logger=True, on_epoch=True, on_step=False) - self.log('d_loss', float(d_loss), prog_bar=True, logger=True, on_epoch=True, on_step=False) - self.log('g_loss', float(g_loss), prog_bar=True, logger=True, on_epoch=True, on_step=False) - self.log('stability_metric', float(d_loss_real + torch.abs(diff)), prog_bar=True, logger=True, on_epoch=True, on_step=False) + self.log( + "mean_loss", + float(r_loss), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) + self.log( + "d_loss", + float(d_loss), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) + self.log( + "g_loss", + float(g_loss), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) + self.log( + "stability_metric", + float(d_loss_real + torch.abs(diff)), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) return diff --git a/pina/solvers/pinn.py b/pina/solvers/pinn.py index 5cafbba..008034f 100644 --- a/pina/solvers/pinn.py +++ b/pina/solvers/pinn.py @@ -1,9 +1,13 @@ """ Module for PINN """ + import torch + 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 ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 import sys from torch.optim.lr_scheduler import ConstantLR @@ -39,14 +43,11 @@ class PINN(SolverInterface): extra_features=None, loss=torch.nn.MSELoss(), optimizer=torch.optim.Adam, - optimizer_kwargs={'lr': 0.001}, + optimizer_kwargs={"lr": 0.001}, scheduler=ConstantLR, - scheduler_kwargs={ - "factor": 1, - "total_iters": 0 - }, + scheduler_kwargs={"factor": 1, "total_iters": 0}, ): - ''' + """ :param AbstractProblem problem: The formulation of the problem. :param torch.nn.Module model: The neural network model to use. :param torch.nn.Module loss: The loss function used as minimizer, @@ -59,12 +60,14 @@ class PINN(SolverInterface): :param torch.optim.LRScheduler scheduler: Learning rate scheduler. :param dict scheduler_kwargs: LR scheduler constructor keyword args. - ''' - super().__init__(models=[model], - problem=problem, - optimizers=[optimizer], - optimizers_kwargs=[optimizer_kwargs], - extra_features=extra_features) + """ + super().__init__( + models=[model], + problem=problem, + optimizers=[optimizer], + optimizers_kwargs=[optimizer_kwargs], + extra_features=extra_features, + ) # check consistency check_consistency(scheduler, LRScheduler, subclass=True) @@ -105,15 +108,21 @@ class PINN(SolverInterface): # to the parameters that the optimizer needs to optimize if isinstance(self.problem, InverseProblem): self.optimizers[0].add_param_group( - {'params': [self._params[var] for var in self.problem.unknown_variables]} - ) + { + "params": [ + self._params[var] + for var in self.problem.unknown_variables + ] + } + ) return self.optimizers, [self.scheduler] def _clamp_inverse_problem_params(self): for v in self._params: self._params[v].data.clamp_( - self.problem.unknown_parameter_domain.range_[v][0], - self.problem.unknown_parameter_domain.range_[v][1]) + self.problem.unknown_parameter_domain.range_[v][0], + self.problem.unknown_parameter_domain.range_[v][1], + ) def _loss_data(self, input, output): return self.loss(self.forward(input), output) @@ -121,9 +130,15 @@ class PINN(SolverInterface): def _loss_phys(self, samples, equation): try: residual = equation.residual(samples, self.forward(samples)) - except TypeError: # this occurs when the function has three inputs, i.e. inverse problem - residual = equation.residual(samples, self.forward(samples), self._params) - return self.loss(torch.zeros_like(residual, requires_grad=True), residual) + except ( + TypeError + ): # this occurs when the function has three inputs, i.e. inverse problem + residual = equation.residual( + samples, self.forward(samples), self._params + ) + return self.loss( + torch.zeros_like(residual, requires_grad=True), residual + ) def training_step(self, batch, batch_idx): """ @@ -140,23 +155,25 @@ class PINN(SolverInterface): dataloader = self.trainer.train_dataloader condition_losses = [] - condition_idx = batch['condition'] + condition_idx = batch["condition"] - for condition_id in range(condition_idx.min(), condition_idx.max()+1): + for condition_id in range(condition_idx.min(), condition_idx.max() + 1): if sys.version_info >= (3, 8): condition_name = dataloader.condition_names[condition_id] else: - condition_name = dataloader.loaders.condition_names[condition_id] + condition_name = dataloader.loaders.condition_names[ + condition_id + ] condition = self.problem.conditions[condition_name] - pts = batch['pts'] + pts = batch["pts"] if len(batch) == 2: samples = pts[condition_idx == condition_id] loss = self._loss_phys(samples, condition.equation) elif len(batch) == 3: samples = pts[condition_idx == condition_id] - ground_truth = batch['output'][condition_idx == condition_id] + ground_truth = batch["output"][condition_idx == condition_id] loss = self._loss_data(samples, ground_truth) else: raise ValueError("Batch size not supported") @@ -164,10 +181,16 @@ class PINN(SolverInterface): # TODO for users this us hard to remember when creating a new solver, to fix in a smarter way loss = loss.as_subclass(torch.Tensor) -# # add condition losses and accumulate logging for each epoch + # # add condition losses and accumulate logging for each epoch condition_losses.append(loss * condition.data_weight) - self.log(condition_name + '_loss', float(loss), - prog_bar=True, logger=True, on_epoch=True, on_step=False) + self.log( + condition_name + "_loss", + float(loss), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) # clamp unknown parameters of the InverseProblem to their domain ranges (if needed) if isinstance(self.problem, InverseProblem): @@ -176,8 +199,14 @@ class PINN(SolverInterface): # TODO Fix the bug, tot_loss is a label tensor without labels # we need to pass it as a torch tensor to make everything work total_loss = sum(condition_losses) - self.log('mean_loss', float(total_loss / len(condition_losses)), - prog_bar=True, logger=True, on_epoch=True, on_step=False) + self.log( + "mean_loss", + float(total_loss / len(condition_losses)), + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) return total_loss diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index 93f1bbf..324a023 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -15,12 +15,14 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): LightningModule methods. """ - def __init__(self, - models, - problem, - optimizers, - optimizers_kwargs, - extra_features=None): + def __init__( + self, + models, + problem, + optimizers, + optimizers_kwargs, + extra_features=None, + ): """ :param models: A torch neural network model instance. :type models: torch.nn.Module @@ -30,7 +32,7 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): use. :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args. :param list(torch.nn.Module) extra_features: The additional input - features to use as augmented input. If ``None`` no extra features + features to use as augmented input. If ``None`` no extra features are passed. If it is a list of :class:`torch.nn.Module`, the extra feature list is passed to all models. If it is a list of extra features' lists, each single list of extra feature is passed to a model. @@ -57,19 +59,23 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): # check length consistency optimizers if len_model != len_optimizer: - raise ValueError('You must define one optimizer for each model.' - f'Got {len_model} models, and {len_optimizer}' - ' optimizers.') + raise ValueError( + "You must define one optimizer for each model." + f"Got {len_model} models, and {len_optimizer}" + " optimizers." + ) # check length consistency optimizers kwargs if len_optimizer_kwargs != len_optimizer: - raise ValueError('You must define one dictionary of keyword' - ' arguments for each optimizers.' - f'Got {len_optimizer} optimizers, and' - f' {len_optimizer_kwargs} dicitionaries') + raise ValueError( + "You must define one dictionary of keyword" + " arguments for each optimizers." + f"Got {len_optimizer} optimizers, and" + f" {len_optimizer_kwargs} dicitionaries" + ) # extra features handling - if (extra_features is None) or (len(extra_features)==0): + if (extra_features is None) or (len(extra_features) == 0): extra_features = [None] * len_model else: # if we only have a list of extra features @@ -78,24 +84,28 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): else: # if we have a list of list extra features if len(extra_features) != len_model: raise ValueError( - 'You passed a list of extrafeatures list with len' - f'different of models len. Expected {len_model} ' - f'got {len(extra_features)}. If you want to use ' - 'the same list of extra features for all models, ' - 'just pass a list of extrafeatures and not a list ' - 'of list of extra features.') + "You passed a list of extrafeatures list with len" + f"different of models len. Expected {len_model} " + f"got {len(extra_features)}. If you want to use " + "the same list of extra features for all models, " + "just pass a list of extrafeatures and not a list " + "of list of extra features." + ) # assigning model and optimizers self._pina_models = [] self._pina_optimizers = [] for idx in range(len_model): - model_ = Network(model=models[idx], - input_variables=problem.input_variables, - output_variables=problem.output_variables, - extra_features=extra_features[idx]) - optim_ = optimizers[idx](model_.parameters(), - **optimizers_kwargs[idx]) + model_ = Network( + model=models[idx], + input_variables=problem.input_variables, + output_variables=problem.output_variables, + extra_features=extra_features[idx], + ) + optim_ = optimizers[idx]( + model_.parameters(), **optimizers_kwargs[idx] + ) self._pina_models.append(model_) self._pina_optimizers.append(optim_) diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index 8abf8a6..c6a8a35 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -1,10 +1,14 @@ """ Module for SupervisedSolver """ + import torch import sys + 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 ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 from torch.optim.lr_scheduler import ConstantLR @@ -18,7 +22,7 @@ from torch.nn.modules.loss import _Loss class SupervisedSolver(SolverInterface): """ SupervisedSolver solver class. This class implements a SupervisedSolver, - using a user specified ``model`` to solve a specific ``problem``. + using a user specified ``model`` to solve a specific ``problem``. """ def __init__( @@ -28,14 +32,11 @@ class SupervisedSolver(SolverInterface): extra_features=None, loss=torch.nn.MSELoss(), optimizer=torch.optim.Adam, - optimizer_kwargs={'lr': 0.001}, + optimizer_kwargs={"lr": 0.001}, scheduler=ConstantLR, - scheduler_kwargs={ - "factor": 1, - "total_iters": 0 - }, + scheduler_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. :param torch.nn.Module loss: The loss function used as minimizer, @@ -49,12 +50,14 @@ class SupervisedSolver(SolverInterface): :param torch.optim.LRScheduler scheduler: Learning rate scheduler. :param dict scheduler_kwargs: LR scheduler constructor keyword args. - ''' - super().__init__(models=[model], - problem=problem, - optimizers=[optimizer], - optimizers_kwargs=[optimizer_kwargs], - extra_features=extra_features) + """ + super().__init__( + models=[model], + problem=problem, + optimizers=[optimizer], + optimizers_kwargs=[optimizer_kwargs], + extra_features=extra_features, + ) # check consistency check_consistency(scheduler, LRScheduler, subclass=True) @@ -69,7 +72,7 @@ class SupervisedSolver(SolverInterface): def forward(self, x): """Forward pass implementation for the solver. - :param torch.Tensor x: Input tensor. + :param torch.Tensor x: Input tensor. :return: Solver solution. :rtype: torch.Tensor """ @@ -95,32 +98,39 @@ class SupervisedSolver(SolverInterface): """ dataloader = self.trainer.train_dataloader - condition_idx = batch['condition'] + condition_idx = batch["condition"] - for condition_id in range(condition_idx.min(), condition_idx.max()+1): + for condition_id in range(condition_idx.min(), condition_idx.max() + 1): if sys.version_info >= (3, 8): condition_name = dataloader.condition_names[condition_id] else: - condition_name = dataloader.loaders.condition_names[condition_id] + condition_name = dataloader.loaders.condition_names[ + condition_id + ] condition = self.problem.conditions[condition_name] - pts = batch['pts'] - out = batch['output'] + pts = batch["pts"] + out = batch["output"] if condition_name not in self.problem.conditions: - raise RuntimeError('Something wrong happened.') + raise RuntimeError("Something wrong happened.") # for data driven mode - if not hasattr(condition, 'output_points'): - raise NotImplementedError('Supervised solver works only in data-driven mode.') - + if not hasattr(condition, "output_points"): + raise NotImplementedError( + "Supervised solver works only in data-driven mode." + ) + output_pts = out[condition_idx == condition_id] input_pts = pts[condition_idx == condition_id] - loss = self.loss(self.forward(input_pts), output_pts) * condition.data_weight + loss = ( + self.loss(self.forward(input_pts), output_pts) + * condition.data_weight + ) loss = loss.as_subclass(torch.Tensor) - self.log('mean_loss', float(loss), prog_bar=True, logger=True) + self.log("mean_loss", float(loss), prog_bar=True, logger=True) return loss @property diff --git a/pina/trainer.py b/pina/trainer.py index bef8b80..0acecaa 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -5,6 +5,7 @@ from .utils import check_consistency from .dataset import SamplePointDataset, SamplePointLoader, DataPointDataset from .solvers.solver import SolverInterface + class Trainer(pytorch_lightning.Trainer): def __init__(self, solver, batch_size=None, **kwargs): @@ -29,18 +30,20 @@ class Trainer(pytorch_lightning.Trainer): check_consistency(solver, SolverInterface) if batch_size is not None: check_consistency(batch_size, int) - + self._model = solver self.batch_size = batch_size # create dataloader if solver.problem.have_sampled_points is False: - raise RuntimeError(f'Input points in {solver.problem.not_sampled_points} ' - 'training are None. Please ' - 'sample points in your problem by calling ' - 'discretise_domain function before train ' - 'in the provided locations.') - + raise RuntimeError( + f"Input points in {solver.problem.not_sampled_points} " + "training are None. Please " + "sample points in your problem by calling " + "discretise_domain function before train " + "in the provided locations." + ) + self._create_or_update_loader() def _create_or_update_loader(self): @@ -52,21 +55,23 @@ class Trainer(pytorch_lightning.Trainer): devices = self._accelerator_connector._parallel_devices if len(devices) > 1: - raise RuntimeError('Parallel training is not supported yet.') + raise RuntimeError("Parallel training is not supported yet.") device = devices[0] dataset_phys = SamplePointDataset(self._model.problem, device) dataset_data = DataPointDataset(self._model.problem, device) self._loader = SamplePointLoader( - dataset_phys, dataset_data, batch_size=self.batch_size, - shuffle=True) + dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True + ) def train(self, **kwargs): """ Train the solver method. """ - return super().fit(self._model, train_dataloaders=self._loader, **kwargs) - + return super().fit( + self._model, train_dataloaders=self._loader, **kwargs + ) + @property def solver(self): """ diff --git a/pina/utils.py b/pina/utils.py index b21499b..282dd53 100644 --- a/pina/utils.py +++ b/pina/utils.py @@ -1,4 +1,5 @@ """Utils module""" + from torch.utils.data import Dataset, DataLoader from functools import reduce import types @@ -12,13 +13,13 @@ import torch def check_consistency(object, object_instance, subclass=False): - """Helper function to check object inheritance consistency. + """Helper function to check object inheritance consistency. Given a specific ``'object'`` we check if the object is instance of a specific ``'object_instance'``, or in case ``'subclass=True'`` we check if the object is subclass if the ``'object_instance'``. - :param (iterable or class object) object: The object to check the inheritance + :param (iterable or class object) object: The object to check the inheritance :param Object object_instance: The parent class from where the object is expected to inherit :param str object_name: The name of the object @@ -39,9 +40,9 @@ def check_consistency(object, object_instance, subclass=False): raise ValueError(f"{type(obj).__name__} must be {object_instance}.") -def number_parameters(model, - aggregate=True, - only_trainable=True): # TODO: check +def number_parameters( + model, aggregate=True, only_trainable=True +): # TODO: check """ Return the number of parameters of a given `model`. @@ -79,8 +80,9 @@ def merge_two_tensors(tensor1, tensor2): n2 = tensor2.shape[0] tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels) - tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), - labels=tensor2.labels) + tensor2 = LabelTensor( + tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels + ) return tensor1.append(tensor2) @@ -95,13 +97,13 @@ def torch_lhs(n, dim): """ if not isinstance(n, int): - raise TypeError('number of point n must be int') + raise TypeError("number of point n must be int") if not isinstance(dim, int): - raise TypeError('dim must be int') + raise TypeError("dim must be int") if dim < 1: - raise ValueError('dim must be greater than one') + raise ValueError("dim must be greater than one") samples = torch.rand(size=(n, dim)) diff --git a/pina/writer.py b/pina/writer.py index 36d1999..831c1cc 100644 --- a/pina/writer.py +++ b/pina/writer.py @@ -1,4 +1,5 @@ """ Module for plotting. """ + import matplotlib.pyplot as plt import numpy as np import torch @@ -11,7 +12,7 @@ class Writer: Implementation of a writer class, for textual output. """ - def __init__(self, frequency_print=10, header='any') -> None: + def __init__(self, frequency_print=10, header="any") -> None: """ The constructor of the class. @@ -28,7 +29,7 @@ class Writer: """ header = [] for condition_name in trainer.problem.conditions: - header.append(f'{condition_name}') + header.append(f"{condition_name}") return header @@ -46,4 +47,4 @@ class Writer: """ if trainer.trained_epoch % self._frequency_print == 0: - print(f'Epoch {trainer.trained_epoch:05d}: {loss.item():.5e}') + print(f"Epoch {trainer.trained_epoch:05d}: {loss.item():.5e}")