PINA/pina/model/layers/pod.py

"""Module for Base Continuous Convolution class."""

from abc import ABCMeta, abstractmethod
import torch
from .stride import Stride
from .utils_convolution import optimizing
import warnings


class PODBlock(torch.nn.Module):
    """
    POD layer: it projects the input field on the proper orthogonal
    decomposition basis.  It needs to be fitted to the data before being used
    with the method :meth:`fit`, which invokes the singular value decomposition.
    The layer is not trainable.

    .. note::
        All the POD modes are stored in memory, avoiding to recompute them when the rank changes but increasing the memory usage.
    """

    def __init__(self, rank, scale_coefficients=True):
        """
        Build the POD layer with the given rank.

        :param int rank: The rank of the POD layer.
        :param bool scale_coefficients: If True, the coefficients are scaled
            after the projection to have zero mean and unit variance.
        """
        super().__init__()
        self.__scale_coefficients = scale_coefficients
        self._basis = None
        self._scaler = None
        self._rank = rank

    @property
    def rank(self):
        """
        The rank of the POD layer.

        :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")

        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
        """
        if self._basis is None:
            return None

        return self._basis[: self.rank]

    @property
    def scaler(self):
        """
        The scaler. It is a dictionary with the keys `'mean'` and `'std'` that
        store the mean and the standard deviation of the coefficients.

        :rtype: dict
        """
        if self._scaler is None:
            return

        return {
            "mean": self._scaler["mean"][: self.rank],
            "std": self._scaler["std"][: self.rank],
        }

    @property
    def scale_coefficients(self):
        """
        If True, the coefficients are scaled after the projection to have zero
        mean and unit variance.

        :rtype: bool
        """
        return self.__scale_coefficients

    def fit(self, X, randomized=True):
        """
        Set the POD basis by performing the singular value decomposition of the
        given tensor. If `self.scale_coefficients` is True, the coefficients
        are scaled after the projection to have zero mean and unit variance.

        :param torch.Tensor X: The tensor to be reduced.
        """
        self._fit_pod(X, randomized)

        if self.__scale_coefficients:
            self._fit_scaler(torch.matmul(self._basis, X.T))

    def _fit_scaler(self, coeffs):
        """
        Private merhod that computes the mean and the standard deviation of the
        given coefficients, allowing to scale them to have zero mean and unit
        variance. Mean and standard deviation are stored in the private member
        `_scaler`.

        :param torch.Tensor coeffs: The coefficients to be scaled.
        """
        self._scaler = {
            "std": torch.std(coeffs, dim=1),
            "mean": torch.mean(coeffs, dim=1),
        }

    def _fit_pod(self, X, randomized):
        """
        Private method that computes the POD basis of the given tensor and stores it in the private member `_basis`.

        :param torch.Tensor X: The tensor to be reduced.
        """
        if X.device.type == "mps":  #  svd_lowrank not arailable for mps
            warnings.warn(
                "svd_lowrank not available for mps, using svd instead."
                "This may slow down computations.",
                ResourceWarning,
            )
            self._basis = torch.svd(X.T)[0].T
        else:
            if randomized:
                warnings.warn(
                    "Considering a randomized algorithm to compute the POD basis"
                )
                self._basis = torch.svd_lowrank(X.T, q=X.shape[0])[0].T
            else:
                self._basis = torch.svd(X.T)[0].T

    def forward(self, X):
        """
        The forward pass of the POD layer. By default it executes the
        :meth:`reduce` method, reducing the input tensor to its POD
        representation. The POD layer needs to be fitted before being used.

        :param torch.Tensor X: The input tensor to be reduced.
        :return: The reduced tensor.
        :rtype: torch.Tensor
        """
        return self.reduce(X)

    def reduce(self, X):
        """
        Reduce the input tensor to its POD representation. The POD layer needs
        to be fitted before being used.

        :param torch.Tensor X: The input tensor to be reduced.
        :return: The reduced tensor.
        :rtype: torch.Tensor
        """
        if self._basis is None:
            raise RuntimeError(
                "The POD layer needs to be fitted before being used."
            )

        coeff = torch.matmul(self.basis, X.T)
        if coeff.ndim == 1:
            coeff = coeff.unsqueeze(1)

        coeff = coeff.T
        if self.__scale_coefficients:
            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."
            )

        if self.__scale_coefficients:
            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