Cleaning Solver

pina/_src/solver/REFACTORING.md

@@ -0,0 +1,161 @@
+# Solver Submodule Refactoring
+
+## Motivation
+
+The solver submodule uses a mixin-based architecture for modular construction. While extensible,
+the growing number of mixins and override chains made it difficult to read and understand which
+method is overridden by which mixin. This refactoring compresses methods, clarifies override
+semantics, and makes it easier to add new mixins.
+
+## Design Decisions
+
+| Decision | Resolution |
+|---|---|
+| `_compute_condition_loss` has 4 variants | **Pro.** Different scientific methods coexist in one framework. |
+| SelfAdaptive solver | Keep as solver class, not a weighting class. Trainable-pointwise-weight + min-max loop is solver-level logic. |
+| Competitive solver | Fine as-is. |
+| GradientEnhanced | Keep as mixin. Moving to Condition changes the interface too much. |
+| RBA mixin | Keep as mixin. Operates at a different stage than GradientEnhanced. |
+| All mixins stay | ConditionAggregator, SingleModel, MultiModel, Ensemble, ManualOpt, PhysicsInformed, Autoregressive, GradientEnhanced, RBA — all retained for future extensibility. |
+| Forward override conflict | Resolved by removing stacking forward from MultiModelMixin — only EnsembleMixin owns it. |
+| GradientEnhanced + RBA | Composable, not exclusive. GradientEnhanced overrides `_regularize_condition_loss` (additive), RBA overrides `_weight_condition_loss` (multiplicative). |
+
+## Architecture
+
+### Layer 0: Interface
+`SolverInterface` — abstract contract (training_step, validation_step, test_step, properties).
+
+### Layer 1: Core
+`BaseSolver` — full pipeline: init, forward (default single-model), loss computation, hooks,
+properties. All mixins and concrete solvers build on this.
+
+### Layer 2: Structural Mixins
+| Mixin | Responsibility |
+|---|---|
+| `SingleModelMixin` | Single-model forward, optimizer config, convenience properties (model, optimizer, scheduler) |
+| `MultiModelMixin` | Multi-optimizer config, convenience properties (models, optimizers, schedulers) — no forward |
+| `EnsembleMixin` | Stacking forward (all models), per-model loss iteration |
+| `ManualOptimizationMixin` | Lightning manual optimization loop (zero-grad, backward, step) |
+| `ConditionAggregatorMixin` | Batch iteration, per-condition loss, logging, weighting aggregation |
+
+### Layer 3: Composed Base Solvers
+| Solver | Inherits |
+|---|---|
+| `SingleModelSolver` | SingleModel + CondAgg + BaseSolver |
+| `MultiModelSolver` | ManualOpt + MultiModel + CondAgg + BaseSolver |
+| `EnsembleSolver` | ManualOpt + Ensemble + CondAgg + BaseSolver |
+
+### Layer 4: Feature Mixins
+| Mixin | Hook overridden | Behavior |
+|---|---|---|
+| `PhysicsInformedMixin` | `validation_step`, `test_step` | Wraps with `@torch.enable_grad()` |
+| `AutoregressiveMixin` | `_loss_from_residual` | Adaptive temporal weighting |
+| `GradientEnhancedMixin` | `_regularize_condition_loss` | Additive gradient penalty |
+| `ResidualBasedAttentionMixin` | `_weight_condition_loss` | Multiplicative attention weighting |
+
+### Layer 5: Concrete Solvers
+All 14 concrete solver classes. Most are thin wrappers (own `__init__` + `accepted_conditions_types`).
+Complex ones (Causal, SelfAdaptive, Competitive) override `_compute_condition_loss` and/or
+`training_step`.
+
+## Method x Mixin Table
+
+### BaseSolver + Structural Mixins + Feature Mixins
+
+| Method | BaseSolver | SingleModelMixin | MultiModelMixin | EnsembleMixin | ManualOptMixin | CondAggMixin | PhysicsInfMixin | AutoregressiveMixin | GradientEnhancedMixin | RBAMixin |
+|---|---|---|---|---|---|---|---|---|---|---|
+| `__init__` | **Merged:** validates problem, stores state, labelizes forward, stores models/opts/scheds, sets up weighting + loss | | | | | | | | | |
+| `forward` | **Default:** `self._pina_models[0](x)` | **Overrides:** `self.model(x)` | | **Overrides:** if `_active_model_idx`, eval one model; else `torch.stack([m(x) for m in self.models])` | | | | | | |
+| `configure_optimizers` | | **Defines:** hook single model + inverse params, return single lists | **Defines:** hook all models + inverse params, return lists | *(inherits MultiModel)* | | | | | | |
+| `training_step` | Calls `batch_evaluation_step`, logs `train_loss` | | | | **Overrides:** zero-grad -> super() -> manual backward -> step all | | | | | |
+| `validation_step` | Calls `batch_evaluation_step`, logs `val_loss` | | | | | | **Overrides:** wraps `super()` with `@torch.enable_grad()` | | | |
+| `test_step` | Calls `batch_evaluation_step`, logs `test_loss` | | | | | | **Overrides:** wraps `super()` with `@torch.enable_grad()` | | | |
+| `on_train_batch_end` | | | | | **Defines:** syncs Lightning manual-opt counters | | | | | |
+| `batch_evaluation_step` | | | | | | **Defines:** iterates batch -> `_compute_condition_loss` per condition -> logs -> aggregates via weighting | | | | |
+| `_compute_condition_loss` | Clone -> prepare -> evaluate -> loss -> regularize -> weight -> reduce | | | | | | | | | |
+| `_prepare_condition_data` | Identity hook | | | | | | | | **Overrides:** sets `requires_grad_(True)` on input | |
+| `_regularize_condition_loss` | Identity hook (additive) | | | | | | | | **Overrides:** computes `grad(residual, input)`, adds gradient penalty | |
+| `_weight_condition_loss` | Identity hook (multiplicative) | | | | | | | | | **Overrides:** normalizes residual norms, updates attention weights, multiplies loss |
+| `_loss_from_residual` | `loss_fn(residual, zeros)` | | | | | | | **Overrides:** adds adaptive temporal weighting via `_get_weights` | | |
+| `_apply_reduction` | Applies none/mean/sum | | | | | | | | | |
+| `get_batch_size` | Static: sums input lengths | | | | | | | | | |
+| `on_train_epoch_start` | | | | | | | | **Defines:** optionally resets running averages | | |
+| `_init_autoregressive_components` | | | | | | | | **Defines:** stores eps, running averages, step counters | | |
+| `_get_weights` | | | | | | | | **Defines:** running average per condition -> `_compute_adaptive_weights` | | |
+| `_compute_adaptive_weights` | | | | | | | | **Defines:** `exp(-eps * cumsum(step_loss))` | | |
+| `preprocess_step` | | | | | | | | **Defines:** identity hook before rollout step | | |
+| `postprocess_step` | | | | | | | | **Defines:** identity hook after rollout step | | |
+| `predict` | | | | | | | | **Defines:** autoregressive rollout | | |
+| `_init_gradient_enhanced_components` | | | | | | | | | **Defines:** validates params, ensures SpatialProblem + LabelTensor | |
+| `_init_residual_attention_components` | | | | | | | | | | **Defines:** validates params, registers per-condition attention buffers |
+| `problem` | Property -> `_pina_problem` | | | | | | | | | |
+| `use_lt` | Property -> `_use_lt` | | | | | | | | | |
+| `weighting` | Property -> `_pina_weighting` | | | | | | | | | |
+| `loss` | Property -> `_loss_fn` | | | | | | | | | |
+| `model` | | Property -> `_pina_models[0]` | | | | | | | | |
+| `optimizer` | | Property -> `_pina_optimizers[0]` | | | | | | | | |
+| `scheduler` | | Property -> `_pina_schedulers[0]` | | | | | | | | |
+| `models` | | | Property -> `_pina_models` | *(inherits)* | | | | | | |
+| `optimizers` | | | Property -> `_pina_optimizers` | *(inherits)* | | | | | | |
+| `schedulers` | | | Property -> `_pina_schedulers` | *(inherits)* | | | | | | |
+| `num_models` | | | Property -> `len(self.models)` | *(inherits)* | | | | | | |
+
+### Concrete Solvers
+
+| Solver | Inherits from | Overrides `training_step` | Overrides `_compute_condition_loss` | Extra own methods |
+|---|---|---|---|---|
+| **SingleModelSolver** | Single + CondAgg + Base | -- | -- | `__init__`, `accepted_conditions_types` |
+| **MultiModelSolver** | ManualOpt + Multi + CondAgg + Base | -- | -- | `__init__` |
+| **EnsembleSolver** | ManualOpt + Ensemble + CondAgg + Base | -- | -- | `__init__` |
+| **SupervisedSingleModelSolver** | SingleModelSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **SupervisedEnsembleSolver** | EnsembleSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **PhysicsInformedSingleModelSolver** | PhysInf + SingleModelSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **PhysicsInformedEnsembleSolver** | PhysInf + EnsembleSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **AutoregressiveSingleModelSolver** | AutoReg + SingleModelSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **AutoregressiveEnsembleSolver** | AutoReg + EnsembleSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **GradientPhysicsInformedSingleModelSolver** | PhysInf + GradEnh + SingleModelSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **RBAPhysicsInformedSingleModelSolver** | PhysInf + RBA + SingleModelSolver | -- | -- | `__init__`, `accepted_conditions_types` |
+| **CausalPhysicsInformedSingleModelSolver** | PhysInf + SingleModelSolver | -- | Time-step loop + causal weighting | `__init__`, `accepted_conditions_types`, `temporal_variable`, `spatial_variables`, `_compute_weights` |
+| **SelfAdaptivePhysicsInformedSolver** | PhysInf + MultiModelSolver | Min-max: weights first, then model | Trainable weight application | `__init__`, `accepted_conditions_types`, `model`, `weights`, `optimizer_model`, `optimizer_weights`, `scheduler_model`, `scheduler_weights` |
+| **CompetitivePhysicsInformedSolver** | PhysInf + MultiModelSolver | Min-max: model first, then discriminator with `-loss` | Multiplies residual by discriminator output | `__init__`, `accepted_conditions_types`, `model`, `discriminator`, `optimizer_model`, `optimizer_discriminator`, `scheduler_model`, `scheduler_discriminator` |
+
+## Override Conflict Map
+
+| Method | Overridden by | Conflict? |
+|---|---|---|
+| `forward` | BaseSolver (default), SingleModelMixin, EnsembleMixin | Clean. SingleModel overrides for single-model; Ensemble overrides for stacking. |
+| `training_step` | BaseSolver, ManualOptimizationMixin, SelfAdaptive, Competitive | SelfAdaptive/Competitive completely replace ManualOpt's training_step with custom min-max loops. |
+| `_compute_condition_loss` | BaseSolver, EnsembleMixin, Causal, SelfAdaptive, Competitive | 4 genuinely different implementations. Pro: different scientific methods coexist. |
+| `_loss_from_residual` | BaseSolver, AutoregressiveMixin | Clean single override. |
+| `_prepare_condition_data` | BaseSolver, GradientEnhancedMixin | Clean single override. |
+| `_regularize_condition_loss` | BaseSolver, GradientEnhancedMixin | Clean. GradientEnhanced adds gradient penalty (additive). |
+| `_weight_condition_loss` | BaseSolver, ResidualBasedAttentionMixin | Clean. RBA multiplies by attention weights (multiplicative). |
+| `validation_step` | BaseSolver, PhysicsInformedMixin | Clean single override. |
+| `test_step` | BaseSolver, PhysicsInformedMixin | Clean single override. |
+
+## Changes Applied
+
+### Change 1: BaseSolver init merge
+Merged `_init_solver_components` + `_init_weighting_and_loss` into `BaseSolver.__init__`.
+Concrete solvers now call `BaseSolver.__init__` with all args at once.
+
+### Change 2: MultiModelMixin loses stacking forward
+Removed `forward` from `MultiModelMixin`. It now only manages optimizers/schedulers and provides
+list properties. Stacking forward is owned solely by `EnsembleMixin`.
+
+### Change 3: ManualOptimizationMixin loses `_init_manual_optimization`
+Removed `_init_manual_optimization`. `MultiModelSolver` and `EnsembleSolver` set
+`self.automatic_optimization = False` directly in their `__init__`.
+
+### Change 4: Split regularization into two hooks
+Added `_weight_condition_loss` (multiplicative hook) alongside `_regularize_condition_loss`
+(additive hook). `_compute_condition_loss` calls both in sequence. GradientEnhanced overrides the
+additive hook; RBA overrides the multiplicative hook. They are now composable.
+
+### Change 5: SelfAdaptive/Competitive lose forward override
+Both removed their `forward` override. With `MultiModelMixin` no longer stacking, the default
+single-model forward from `BaseSolver` is correct for both.
+
+### Default forward on BaseSolver
+Added `BaseSolver.forward(x)` returning `self._pina_models[0](x)`. Provides a sensible default
+for any solver that doesn't explicitly override forward.

pina/_src/solver/base_solver.py

@@ -31,11 +31,38 @@ class BaseSolver(SolverInterface, metaclass=ABCMeta):
         "sum": lambda x: x.sum(),
     }
 
-    def __init__(self, problem, use_lt=True):
+    def __init__(
+        self,
+        problem,
+        models,
+        optimizers=None,
+        schedulers=None,
+        weighting=None,
+        loss=None,
+        use_lt=True,
+    ):
         """
         Initialization of the :class:`BaseSolver` class.
 
         :param BaseProblem problem: The problem to be solved.
+        :param models: The model or list of models used by the solver.
+        :type models: torch.nn.Module | list[torch.nn.Module]
+        :param optimizers: The optimizer or list of optimizers 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``.
+        :type optimizers: TorchOptimizer | list[TorchOptimizer]
+        :param schedulers: The scheduler or list of schedulers 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``.
+        :type schedulers: TorchScheduler | list[TorchScheduler]
+        :param BaseWeighting weighting: The weighting strategy used to combine
+            condition losses. If ``None``, no weighting is applied. Default is
+            ``None``.
+        :param loss: The loss function used to compute residual losses.
+            If ``None``, :class:`torch.nn.MSELoss` is used. Default is ``None``.
+        :type loss: torch.nn.Module | BaseDualLoss
         :param bool use_lt: If ``True``, the solver uses LabelTensors as input.
             Default is ``True``.
         :raises ValueError: If ``use_lt`` is not a boolean.
@@ -56,7 +83,7 @@ def __init__(self, problem, use_lt=True):
         for condition in problem.conditions.values():
             check_consistency(condition, self.accepted_conditions_types)
 
-        # Initialize the solver components
+        # Store problem and use_lt
         self._pina_problem = problem
         self._use_lt = use_lt
 
@@ -76,6 +103,27 @@ def __init__(self, problem, use_lt=True):
                 output_variables=problem.output_variables,
             )
 
+        # Initialize solver components (models, optimizers, schedulers)
+        self._init_solver_components(
+            models=models,
+            optimizers=optimizers,
+            schedulers=schedulers,
+        )
+
+        # Initialize the weighting scheme and loss function
+        self._init_weighting_and_loss(weighting=weighting, loss=loss)
+
+    def forward(self, x):
+        """
+        Default forward pass: evaluates the first model on the input.
+
+        :param x: The input data.
+        :type x: torch.Tensor | LabelTensor | Data | Graph
+        :return: The output of the first model.
+        :rtype: torch.Tensor | LabelTensor | Data | Graph
+        """
+        return self._pina_models[0](x)
+
     def reset(self):
         """
         Reset the internal solver state, clearing the stored problem, models,
@@ -294,14 +342,22 @@ def _compute_condition_loss(self, condition, data, batch_idx):
         # Compute the tensor loss from the residual tensor
         condition_tensor_loss = self._loss_from_residual(condition_name)
 
-        # Optional regularization hook, e.g gradient-enhanced or residual-based
+        # Additive regularization hook, e.g. gradient-enhanced penalty
         condition_tensor_loss = self._regularize_condition_loss(
             condition_tensor_loss=condition_tensor_loss,
             condition_name=condition_name,
             data=data,
             batch_idx=batch_idx,
         )
 
+        # Multiplicative weighting hook, e.g. residual-based attention
+        condition_tensor_loss = self._weight_condition_loss(
+            condition_tensor_loss=condition_tensor_loss,
+            condition_name=condition_name,
+            data=data,
+            batch_idx=batch_idx,
+        )
+
         # Compute the scalar loss from the tensor loss and return it
         condition_scalar_loss = self._apply_reduction(condition_tensor_loss)
 
@@ -327,10 +383,9 @@ def _regularize_condition_loss(
         batch_idx,
     ):
         """
-        Regularize the condition loss if needed. This method can be overridden
-        by mixins to implement specific regularization strategies, such as
-        adding a gradient penalty in gradient-enhanced solvers or applying
-        residual-based attention.
+        Additive regularization hook for the condition loss. Override this to
+        implement additive regularization strategies such as gradient-enhanced
+        penalties.
 
         :param condition_tensor_loss: The original tensor loss for the
             condition.
@@ -343,6 +398,28 @@ def _regularize_condition_loss(
         """
         return condition_tensor_loss
 
+    def _weight_condition_loss(
+        self,
+        condition_tensor_loss,
+        condition_name,
+        data,
+        batch_idx,
+    ):
+        """
+        Multiplicative weighting hook for the condition loss. Override this to
+        implement multiplicative weighting strategies such as residual-based
+        attention.
+
+        :param condition_tensor_loss: The tensor loss for the condition.
+        :type condition_tensor_loss: torch.Tensor | LabelTensor
+        :param str condition_name: The name of the condition.
+        :param dict data: The data corresponding to the condition.
+        :param int batch_idx: The index of the current batch.
+        :return: The weighted tensor loss for the condition.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        return condition_tensor_loss
+
     def _loss_from_residual(self, condition_name=None):
         """
         Compute the tensor loss from the residual tensor.

pina/_src/solver/competitive_physics_informed_solver.py

@@ -195,17 +195,6 @@ def training_step(self, batch, batch_idx):
 
         return loss
 
-    def forward(self, x):
-        """
-        Forward pass through the model.
-
-        :param x: The input data.
-        :type x: torch.Tensor | LabelTensor | Data | Graph
-        :return: The output of the model.
-        :rtype: torch.Tensor | LabelTensor | Data | Graph
-        """
-        return self.model(x)
-
     def _compute_condition_loss(self, condition, data, batch_idx):
         """
         Compute the scalar loss for a given condition and its data.

pina/_src/solver/ensemble_solver.py

@@ -72,18 +72,16 @@ def __init__(
             Default is ``True``.
         """
 
-        # Initialize the base solver
-        BaseSolver.__init__(self, problem=problem, use_lt=use_lt)
-
-        # Initialize the components of the solver
-        self._init_solver_components(
+        BaseSolver.__init__(
+            self,
+            problem=problem,
             models=models,
             optimizers=optimizers,
             schedulers=schedulers,
+            weighting=weighting,
+            loss=loss,
+            use_lt=use_lt,
         )
 
-        # Initialize the weighting scheme for the conditions and the loss
-        self._init_weighting_and_loss(weighting=weighting, loss=loss)
-
         # Activate manual optimization
-        self._init_manual_optimization()
+        self.automatic_optimization = False

pina/_src/solver/mixin/manual_optimization_mixin.py

@@ -11,12 +11,6 @@ class ManualOptimizationMixin:
     :class:`~pina._src.solver.base_solver.BaseSolver`.
     """
 
-    def _init_manual_optimization(self):
-        """
-        Disable Lightning's automatic optimization.
-        """
-        self.automatic_optimization = False
-
     def training_step(self, batch, batch_idx):
         """
         Solver training step.