fix batch_idx + clamp_params + tensor_loss

Author: GiovanniCanali

pina/_src/problem/base_problem.py

@@ -248,6 +248,7 @@ def move_discretisation_into_conditions(self):
                 # Set the domain and problem attributes of the new condition
                 new_condition.domain = cond.domain
                 new_condition.problem = self
+                new_condition.name = name
 
                 # Replace the old condition in the conditions dictionary
                 self.conditions[name] = new_condition

pina/_src/solver/autoregressive_single_model_solver.py

@@ -67,11 +67,10 @@ def __init__(
         :param torch.nn.Module model: The model used by the solver.
         :param TorchOptimizer optimizer: The optimizer used by the solver.
             If ``None``, the ``torch.optim.Adam`` optimizer with a learning rate
-            of ``0.001`` is used for each model. Default is ``None``.
+            of ``0.001`` is used. Default is ``None``.
         :param TorchScheduler scheduler: The scheduler used by the solver.
             If ``None``, the ``torch.optim.lr_scheduler.ConstantLR`` scheduler
-            with a factor of ``1.0`` is used for each model.
-            Default is ``None``.
+            with a factor of ``1.0`` is used. Default is ``None``.
         :param BaseWeighting weighting: The weighting strategy used to combine
             condition losses. If ``None``, no weighting is applied. Default is
             ``None``.

pina/_src/solver/base_solver.py

@@ -214,16 +214,17 @@ def _init_solver_components(
         self._pina_optimizers = optimizers
         self._pina_schedulers = schedulers
 
-    def training_step(self, batch):
+    def training_step(self, batch, batch_idx):
         """
         Solver training step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
-        loss = self.batch_evaluation_step(batch=batch)
+        loss = self.batch_evaluation_step(batch=batch, batch_idx=batch_idx)
         self.log(
             name="train_loss",
             value=loss.item(),
@@ -232,16 +233,17 @@ def training_step(self, batch):
         )
         return loss
 
-    def validation_step(self, batch):
+    def validation_step(self, batch, batch_idx):
         """
         Solver validation step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
-        loss = self.batch_evaluation_step(batch=batch)
+        loss = self.batch_evaluation_step(batch=batch, batch_idx=batch_idx)
         self.log(
             name="val_loss",
             value=loss.item(),
@@ -250,16 +252,17 @@ def validation_step(self, batch):
         )
         return loss
 
-    def test_step(self, batch):
+    def test_step(self, batch, batch_idx):
         """
         Solver test step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
-        loss = self.batch_evaluation_step(batch=batch)
+        loss = self.batch_evaluation_step(batch=batch, batch_idx=batch_idx)
         self.log(
             name="test_loss",
             value=loss.item(),
@@ -268,13 +271,14 @@ def test_step(self, batch):
         )
         return loss
 
-    def _compute_condition_loss(self, condition, data):
+    def _compute_condition_loss(self, condition, data, batch_idx):
         """
         Compute the scalar loss for a given condition and its data.
 
         :param BaseCondition condition: The condition for which to compute the
             loss.
         :param dict data: The data corresponding to the condition.
+        :param int batch_idx: The index of the current batch.
         :return: The scalar loss for the condition.
         :rtype: torch.Tensor
         """
@@ -289,26 +293,27 @@ def _compute_condition_loss(self, condition, data):
         # Retrieve condition name for more complex weighting schemes
         condition_name = condition.name if hasattr(condition, "name") else None
 
-        # Compute the scalar loss from the residual tensor and return it
-        condition_loss = self._loss_from_residual(condition_name)
+        # Compute the tensor loss from the residual tensor
+        condition_tensor_loss = self._loss_from_residual(condition_name)
 
-        return condition_loss
+        # Compute the scalar loss from the tensor loss and return it
+        condition_scalar_loss = self._apply_reduction(condition_tensor_loss)
+
+        return condition_scalar_loss
 
     def _loss_from_residual(self, condition_name=None):
         """
-        Compute the scalar loss from the residual tensor.
+        Compute the tensor loss from the residual tensor.
 
         :param str condition_name: The name of the condition.
-        :return: The scalar loss computed from the residual tensor.
+        :return: The tensor loss computed from the residual tensor.
         :rtype: torch.Tensor | LabelTensor
         """
         # Compute the loss tensor and appply reduction
-        loss_tensor = self._loss_fn(
+        return self._loss_fn(
             self.residual_tensor, torch.zeros_like(self.residual_tensor)
         )
 
-        return self._apply_reduction(loss_tensor)
-
     def _apply_reduction(self, value):
         """
         Apply the specified reduction to the loss tensor.

pina/_src/solver/mixin/autoregressive_mixin.py

@@ -47,10 +47,10 @@ def _init_autoregressive_components(
 
     def _loss_from_residual(self, condition_name=None):
         """
-        Compute the scalar loss from the residual tensor.
+        Compute the tensor loss from the residual tensor.
 
         :param str condition_name: The name of the condition.
-        :return: The scalar loss computed from the residual tensor.
+        :return: The tensor loss computed from the residual tensor.
         :rtype: torch.Tensor | LabelTensor
         """
         # Compute the step losses from the residual tensor
@@ -62,10 +62,7 @@ def _loss_from_residual(self, condition_name=None):
         with torch.no_grad():
             weights = self._get_weights(condition_name or "default", step_loss)
 
-        # Compute the weighted step losses
-        weighted_step_loss = step_loss * weights
-
-        return self._apply_reduction(weighted_step_loss)
+        return step_loss * weights
 
     def _get_weights(self, condition_name, step_loss):
         """

pina/_src/solver/mixin/condition_aggregator_mixin.py

@@ -10,7 +10,7 @@ class _ConditionAggregatorMixin:
     :class:`~pina._src.solver.base_solver.BaseSolver`.
     """
 
-    def batch_evaluation_step(self, batch):
+    def batch_evaluation_step(self, batch, batch_idx):
         """
         Evaluate and aggregate the losses for all conditions in a batch.
 
@@ -21,6 +21,7 @@ def batch_evaluation_step(self, batch):
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The aggregated scalar loss for the batch.
         :rtype: torch.Tensor
         """
@@ -34,8 +35,12 @@ def batch_evaluation_step(self, batch):
             condition_losses[condition_name] = self._compute_condition_loss(
                 condition=self.problem.conditions[condition_name],
                 data=dict(data),
+                batch_idx=batch_idx,
             )
 
+        # Clamp parameters - null operation if problem is not InverseProblem
+        self._clamp_params()
+
         # Log the individual condition losses
         for name, value in condition_losses.items():
             self.log(

pina/_src/solver/mixin/manual_optimization_mixin.py

@@ -14,12 +14,13 @@ def _init_manual_optimization(self):
         """
         self.automatic_optimization = False
 
-    def training_step(self, batch):
+    def training_step(self, batch, batch_idx):
         """
         Solver training step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
@@ -28,7 +29,7 @@ def training_step(self, batch):
             opt.instance.zero_grad()
 
         # Perform the forward pass and compute the loss
-        loss = super().training_step(batch)
+        loss = super().training_step(batch, batch_idx)
 
         # Perform the backward pass
         self.manual_backward(loss)

pina/_src/solver/mixin/multi_model_mixin.py

@@ -31,7 +31,7 @@ def configure_optimizers(self):
         Configure the optimizers and schedulers for all models.
 
         :return: The optimizer and the scheduler
-        :rtype: tuple[list[Optimizer], list[Scheduler]]
+        :rtype: tuple[list[TorchOptimizer], list[TorchScheduler]]
         """
         # Iterate over models, optimizers, and schedulers to hook them together
         for optimizer, scheduler, model in zip(

pina/_src/solver/mixin/physics_informed_mixin.py

@@ -12,25 +12,27 @@ class _PhysicsInformedMixin:
     """
 
     @torch.enable_grad()
-    def validation_step(self, batch):
+    def validation_step(self, batch, batch_idx):
         """
         Solver validation step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
-        return super().validation_step(batch)
+        return super().validation_step(batch, batch_idx)
 
     @torch.enable_grad()
-    def test_step(self, batch):
+    def test_step(self, batch, batch_idx):
         """
         Solver test step.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
         :return: The loss of the training step.
         :rtype: torch.Tensor
         """
-        return super().test_step(batch)
+        return super().test_step(batch, batch_idx)

pina/_src/solver/mixin/single_model_mixin.py

@@ -28,7 +28,7 @@ def configure_optimizers(self):
         Configure the optimizer and scheduler for the single model.
 
         :return: The optimizer and the scheduler
-        :rtype: tuple[list[Optimizer], list[Scheduler]]
+        :rtype: tuple[list[TorchOptimizer], list[TorchScheduler]]
         """
         # Hook the optimizer to the model parameters
         self.optimizer.hook(self.model.parameters())

pina/_src/solver/physics_informed_single_model_solver.py

@@ -76,11 +76,10 @@ def __init__(
         :param torch.nn.Module model: The model used by the solver.
         :param TorchOptimizer optimizer: The optimizer used by the solver.
             If ``None``, the ``torch.optim.Adam`` optimizer with a learning rate
-            of ``0.001`` is used for each model. Default is ``None``.
+            of ``0.001`` is used. Default is ``None``.
         :param TorchScheduler scheduler: The scheduler used by the solver.
             If ``None``, the ``torch.optim.lr_scheduler.ConstantLR`` scheduler
-            with a factor of ``1.0`` is used for each model.
-            Default is ``None``.
+            with a factor of ``1.0`` is used. Default is ``None``.
         :param BaseWeighting weighting: The weighting strategy used to combine
             condition losses. If ``None``, no weighting is applied. Default is
             ``None``.