Network handles forward for all solvers

pina/solvers/garom.py

@@ -32,7 +32,6 @@ class GAROM(SolverInterface):
         problem,
         generator,
         discriminator,
-        extra_features=None,
         loss=None,
         optimizer_generator=torch.optim.Adam,
         optimizer_generator_kwargs={'lr': 0.001},
@@ -58,13 +57,6 @@ class GAROM(SolverInterface):
             for the generator.
         :param torch.nn.Module discriminator: The neural network model to use
             for the discriminator.
-        :param torch.nn.Module extra_features: The additional input
-            features to use as augmented input. It should either be a
-            list of torch.nn.Module, or a dictionary. If a list it is 
-            passed the extra features are passed to both network. If a 
-            dictionary is passed, the keys must be ``generator`` and
-            ``discriminator`` and the values a list of torch.nn.Module
-            extra features for each.
         :param torch.nn.Module loss: The loss function used as minimizer,
             default ``None``. If ``loss`` is ``None`` the defualt
             ``PowerLoss(p=1)`` is used, as in the original paper.
@@ -97,15 +89,9 @@ class GAROM(SolverInterface):
             parameters), and ``output_points``.
         """
 
-        if isinstance(extra_features, dict):
-            extra_features = [
-                extra_features['generator'], extra_features['discriminator']
-            ]
-
         super().__init__(
             models=[generator, discriminator],
             problem=problem,
-            extra_features=extra_features,
             optimizers=[optimizer_generator, optimizer_discriminator],
             optimizers_kwargs=[
                 optimizer_generator_kwargs, optimizer_discriminator_kwargs
@@ -200,7 +186,7 @@ class GAROM(SolverInterface):
 
         # generator loss
         r_loss = self._loss(snapshots, generated_snapshots)
-        d_fake = self.discriminator([generated_snapshots, parameters])
+        d_fake = self.discriminator.forward_map([generated_snapshots, parameters])
         g_loss = self._loss(d_fake, generated_snapshots) + self.regularizer * r_loss
 
         # backward step
@@ -220,8 +206,8 @@ class GAROM(SolverInterface):
         generated_snapshots = self.generator(parameters)
 
         # Discriminator pass
-        d_real = self.discriminator([snapshots, parameters])
-        d_fake = self.discriminator([generated_snapshots, parameters]) 
+        d_real = self.discriminator.forward_map([snapshots, parameters])
+        d_fake = self.discriminator.forward_map([generated_snapshots, parameters]) 
 
         # evaluate loss
         d_loss_real = self._loss(d_real, snapshots)

pina/solvers/pinn.py

@@ -83,13 +83,7 @@ class PINN(SolverInterface):
         :return: PINN solution.
         :rtype: torch.Tensor
         """
-        # extract torch.Tensor from corresponding label
-        x = x.extract(self.problem.input_variables).as_subclass(torch.Tensor)
-        # perform forward pass (using torch.Tensor) + converting to LabelTensor
-        output = self.neural_net(x).as_subclass(LabelTensor)
-        # set the labels for LabelTensor
-        output.labels = self.problem.output_variables
-        return output
+        return self.neural_net(x)
 
     def configure_optimizers(self):
         """

pina/solvers/solver.py

@@ -80,7 +80,7 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
                     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'
+                        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.')
@@ -91,6 +91,8 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
 
         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])

pina/solvers/supervised.py

@@ -72,13 +72,7 @@ class SupervisedSolver(SolverInterface):
         :return: Solver solution.
         :rtype: torch.Tensor
         """
-        # extract torch.Tensor from corresponding label
-        x = x.extract(self.problem.input_variables).as_subclass(torch.Tensor)
-        # perform forward pass (using torch.Tensor) + converting to LabelTensor
-        output = self.neural_net(x).as_subclass(LabelTensor)
-        # set the labels for LabelTensor
-        output.labels = self.problem.output_variables
-        return output
+        return self.neural_net(x)
 
     def configure_optimizers(self):
         """Optimizer configuration for the solver.
@@ -125,37 +119,6 @@ class SupervisedSolver(SolverInterface):
         self.log('mean_loss', float(loss), prog_bar=True, logger=True)
         return loss
 
-
-    def training_step_(self, batch, batch_idx):
-        """Solver training step.
-
-        :param batch: The batch element in the dataloader.
-        :type batch: tuple
-        :param batch_idx: The batch index.
-        :type batch_idx: int
-        :return: The sum of the loss functions.
-        :rtype: LabelTensor
-        """
-
-        for condition_name, samples in batch.items():
-
-            if condition_name not in self.problem.conditions:
-                raise RuntimeError('Something wrong happened.')
-
-            condition = self.problem.conditions[condition_name]
-
-            # data loss
-            if hasattr(condition, 'output_points'):
-                input_pts, output_pts = samples
-                loss = self.loss(self.forward(input_pts),
-                                 output_pts) * condition.data_weight
-            else:
-                raise RuntimeError(
-                    'Supervised solver works only in data-driven mode.')
-
-        self.log('mean_loss', float(loss), prog_bar=True, logger=True)
-        return loss
-
     @property
     def scheduler(self):
         """