Source code for pina.equation.equation

""" Module for Equation. """

from .equation_interface import EquationInterface




class Equation(EquationInterface):

    def __init__(self, equation):
        """
        Equation class for specifing any equation in PINA.
        Each ``equation`` passed to a ``Condition`` object
        must be an ``Equation`` or ``SystemEquation``.

        :param equation: A ``torch`` callable equation to
            evaluate the residual.
        :type equation: Callable
        """
        if not callable(equation):
            raise ValueError(
                "equation must be a callable function."
                "Expected a callable function, got "
                f"{equation}"
            )
        self.__equation = equation



    def residual(self, input_, output_, params_=None):
        """
        Residual computation of the equation.

        :param LabelTensor input_: Input points to evaluate the equation.
        :param LabelTensor output_: Output vectors given by a model (e.g,
            a ``FeedForward`` model).
        :param dict params_: Dictionary of parameters related to the inverse
            problem (if any).
            If the equation is not related to an ``InverseProblem``, the
            parameters are initialized to ``None`` and the residual is
            computed as ``equation(input_, output_)``.
            Otherwise, the parameters are automatically initialized in the
            ranges specified by the user.

        :return: The residual evaluation of the specified equation.
        :rtype: LabelTensor
        """
        if params_ is None:
            result = self.__equation(input_, output_)
        else:
            result = self.__equation(input_, output_, params_)
        return result