PINA/pina/solvers/solver.py

""" Solver module. """

import lightning
import torch
import sys

from abc import ABCMeta, abstractmethod
from ..problem import AbstractProblem
from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler
from ..loss import WeightingInterface
from ..loss.scalar_weighting import _NoWeighting
from ..utils import check_consistency, labelize_forward
from torch._dynamo.eval_frame import OptimizedModule


class SolverInterface(lightning.pytorch.LightningModule, metaclass=ABCMeta):
    """
    SolverInterface base class. This class is a wrapper of LightningModule.
    """

    def __init__(self,
                 problem,
                 weighting,
                 use_lt):
        """
        :param problem: A problem definition instance.
        :type problem: AbstractProblem
        :param weighting: The loss weighting to use.
        :type weighting: WeightingInterface
        :param use_lt: Using LabelTensors as input during training.
        :type use_lt: bool
        """
        super().__init__()

        # check consistency of the problem
        check_consistency(problem, AbstractProblem)
        self._check_solver_consistency(problem)
        self._pina_problem = problem

        # check consistency of the weighting and hook the condition names
        if weighting is None:
            weighting = _NoWeighting()
        check_consistency(weighting, WeightingInterface)
        self._pina_weighting = weighting
        weighting.condition_names = list(self._pina_problem.conditions.keys())

        # check consistency use_lt
        check_consistency(use_lt, bool)
        self._use_lt = use_lt

        # if use_lt is true add extract operation in input
        if use_lt is True:
            self.forward = labelize_forward(
                forward=self.forward,
                input_variables=problem.input_variables,
                output_variables=problem.output_variables,
            )

        # PINA private attributes (some are overridden by derived classes)
        self._pina_problem = problem
        self._pina_models = None
        self._pina_optimizers = None
        self._pina_schedulers = None

    def _check_solver_consistency(self, problem):
        for condition in problem.conditions.values():
            check_consistency(condition, self.accepted_conditions_types)

    def _optimization_cycle(self, batch):
        """
        Perform a private optimization cycle by computing the loss for each
        condition in the given batch. The loss are later aggregated using the
        specific weighting schema.

        :param batch: A batch of data, where each element is a tuple containing
            a condition name and a dictionary of points.
        :type batch: list of tuples (str, dict)
        :return: The computed loss for the all conditions in the batch,
            cast to a subclass of `torch.Tensor`. It should return a dict
            containing the condition name and the associated scalar loss.
        :rtype: dict(torch.Tensor)
        """
        losses = self.optimization_cycle(batch)
        for name, value in losses.items():
            self.store_log(f'{name}_loss', value.item(), self.get_batch_size(batch))
        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
        return loss

    def training_step(self, batch):
        """
        Solver training step.

        :param batch: The batch element in the dataloader.
        :type batch: tuple
        :return: The sum of the loss functions.
        :rtype: LabelTensor
        """
        loss = self._optimization_cycle(batch=batch)
        self.store_log('train_loss', loss, self.get_batch_size(batch))
        return loss

    def validation_step(self, batch):
        """
        Solver validation step.

        :param batch: The batch element in the dataloader.
        :type batch: tuple
        """
        loss = self._optimization_cycle(batch=batch)
        self.store_log('val_loss', loss, self.get_batch_size(batch))

    def test_step(self, batch):
        """
        Solver test step.

        :param batch: The batch element in the dataloader.
        :type batch: tuple
        """
        loss = self._optimization_cycle(batch=batch)
        self.store_log('test_loss', loss, self.get_batch_size(batch))

    def store_log(self, name, value, batch_size):
        self.log(name=name,
                 value=value,
                 batch_size=batch_size,
                 **self.trainer.logging_kwargs
                 )

    @abstractmethod
    def forward(self, *args, **kwargs):
        pass

    @abstractmethod
    def optimization_cycle(self, batch):
        """
        Perform an optimization cycle by computing the loss for each condition
        in the given batch.

        :param batch: A batch of data, where each element is a tuple containing
            a condition name and a dictionary of points.
        :type batch: list of tuples (str, dict)
        :return: The computed loss for the all conditions in the batch,
            cast to a subclass of `torch.Tensor`. It should return a dict
            containing the condition name and the associated scalar loss.
        :rtype: dict(torch.Tensor)
        """
        pass

    @property
    def problem(self):
        """
        The problem formulation.
        """
        return self._pina_problem

    @property
    def use_lt(self):
        """
        Using LabelTensor in training.
        """
        return self._use_lt

    @property
    def weighting(self):
        """
        The weighting mechanism.
        """
        return self._pina_weighting

    @staticmethod
    def get_batch_size(batch):
        # assuming batch is a custom Batch object
        batch_size = 0
        for data in batch:
            batch_size += len(data[1]['input_points'])
        return batch_size

    @staticmethod
    def default_torch_optimizer():
        return TorchOptimizer(torch.optim.Adam, lr=0.001)

    @staticmethod
    def default_torch_scheduler():
        return TorchScheduler(torch.optim.lr_scheduler.ConstantLR)

    def on_train_start(self):
        """
        Hook that is called before training begins.
        Used to compile the model if the trainer is set to compile.
        """
        super().on_train_start()
        if self.trainer.compile:
            self._compile_model()

    def on_test_start(self):
        """
        Hook that is called before training begins.
        Used to compile the model if the trainer is set to compile.
        """
        super().on_train_start()
        if self.trainer.compile and not self._check_already_compiled():
            self._compile_model()

    def _check_already_compiled(self):
        models = self._pina_models
        if len(models) == 1 and isinstance(self._pina_models[0],
                                           torch.nn.ModuleDict):
            models = list(self._pina_models.values())
        for model in models:
            if not isinstance(model, (OptimizedModule, torch.nn.ModuleDict)):
                return False
        return True

    @staticmethod
    def _perform_compilation(model):
        model_device = next(model.parameters()).device
        try:
            if model_device == torch.device("mps:0"):
                model = torch.compile(model, backend="eager")
            else:
                model = torch.compile(model, backend="inductor")
        except Exception as e:
            print("Compilation failed, running in normal mode.:\n", e)
        return model


class SingleSolverInterface(SolverInterface):
    def __init__(self,
                 problem,
                 model,
                 optimizer=None,
                 scheduler=None,
                 weighting=None,
                 use_lt=True):
        """
        :param problem: A problem definition instance.
        :type problem: AbstractProblem
        :param model: A torch nn.Module instances.
        :type model: torch.nn.Module
        :param Optimizer optimizers: A neural network optimizers to use.
        :param Scheduler optimizers: A neural network scheduler to use.
        :param WeightingInterface weighting: The loss weighting to use.
        :param bool use_lt: Using LabelTensors as input during training.
        """
        if optimizer is None:
            optimizer = self.default_torch_optimizer()

        if scheduler is None:
            scheduler = self.default_torch_scheduler()

        super().__init__(problem=problem,
                         use_lt=use_lt,
                         weighting=weighting)

        # check consistency of models argument and encapsulate in list
        check_consistency(model, torch.nn.Module)
        # check scheduler consistency and encapsulate in list
        check_consistency(scheduler, Scheduler)
        # check optimizer consistency and encapsulate in list
        check_consistency(optimizer, Optimizer)

        # initialize the model (needed by Lightining to go to different devices)
        self._pina_models = torch.nn.ModuleList([model])
        self._pina_optimizers = [optimizer]
        self._pina_schedulers = [scheduler]

    def forward(self, x):
        """
        Forward pass implementation for the solver.

        :param torch.Tensor x: Input tensor.
        :return: Solver solution.
        :rtype: torch.Tensor
        """
        x = self.model(x)
        return x

    def configure_optimizers(self):
        """
        Optimizer configuration for the solver.

        :return: The optimizers and the schedulers
        :rtype: tuple(list, list)
        """
        self.optimizer.hook(self.model.parameters())
        self.scheduler.hook(self.optimizer)
        return (
            [self.optimizer.instance],
            [self.scheduler.instance]
        )

    def _compile_model(self):
        if isinstance(self._pina_models[0], torch.nn.ModuleDict):
            self._compile_module_dict()
        else:
            self._compile_single_model()

    def _compile_module_dict(self):
        for name, model in self._pina_models[0].items():
            self._pina_models[0][name] = self._perform_compilation(model)

    def _compile_single_model(self):
        self._pina_models[0] = self._perform_compilation(self._pina_models[0])

    @property
    def model(self):
        """
        Model for training.
        """
        return self._pina_models[0]

    @property
    def scheduler(self):
        """
        Scheduler for training.
        """
        return self._pina_schedulers[0]

    @property
    def optimizer(self):
        """
        Optimizer for training.
        """
        return self._pina_optimizers[0]


class MultiSolverInterface(SolverInterface):
    """
    Multiple Solver base class. This class inherits is a wrapper of
    SolverInterface class
    """

    def __init__(self,
                 problem,
                 models,
                 optimizers=None,
                 schedulers=None,
                 weighting=None,
                 use_lt=True):
        """
        :param problem: A problem definition instance.
        :type problem: AbstractProblem
        :param models: Multiple torch nn.Module instances.
        :type model: list[torch.nn.Module] | tuple[torch.nn.Module]
        :param list(Optimizer) optimizers: A list of neural network
           optimizers to use.
        :param list(Scheduler) optimizers: A list of neural network
           schedulers to use.
        :param WeightingInterface weighting: The loss weighting to use.
        :param bool use_lt: Using LabelTensors as input during training.
        """
        if not isinstance(models, (list, tuple)) or len(models) < 2:
            raise ValueError(
                'models should be list[torch.nn.Module] or '
                'tuple[torch.nn.Module] with len greater than '
                'one.'
            )

        if any(opt is None for opt in optimizers):
            optimizers = [
                self.default_torch_optimizer() if opt is None else opt
                for opt in optimizers
            ]

        if any(sched is None for sched in schedulers):
            schedulers = [
                self.default_torch_scheduler() if sched is None else sched
                for sched in schedulers
            ]

        super().__init__(problem=problem,
                         use_lt=use_lt,
                         weighting=weighting)

        # check consistency of models argument and encapsulate in list
        check_consistency(models, torch.nn.Module)

        # check scheduler consistency and encapsulate in list
        check_consistency(schedulers, Scheduler)

        # check optimizer consistency and encapsulate in list
        check_consistency(optimizers, Optimizer)

        # check length consistency optimizers
        if len(models) != len(optimizers):
            raise ValueError(
                "You must define one optimizer for each model."
                f"Got {len(models)} models, and {len(optimizers)}"
                " optimizers."
            )

        # initialize the model
        self._pina_models = torch.nn.ModuleList(models)
        self._pina_optimizers = optimizers
        self._pina_schedulers = schedulers

    def configure_optimizers(self):
        """Optimizer configuration for the solver.

        :return: The optimizers and the schedulers
        :rtype: tuple(list, list)
        """
        for optimizer, scheduler, model in zip(self.optimizers,
                                               self.schedulers,
                                               self.models):
            optimizer.hook(model.parameters())
            scheduler.hook(optimizer)

        return (
            [optimizer.instance for optimizer in self.optimizers],
            [scheduler.instance for scheduler in self.schedulers]
        )

    def _compile_model(self):
        for i, model in enumerate(self._pina_models):
            if not isinstance(model, torch.nn.ModuleDict):
                self._pina_models[i] = self._perform_compilation(model)

    @property
    def models(self):
        """
        The torch model."""
        return self._pina_models

    @property
    def optimizers(self):
        """
        The torch model."""
        return self._pina_optimizers

    @property
    def schedulers(self):
        """
        The torch model."""
        return self._pina_schedulers