Source code for scimba_torch.physical_models.temporal_pde.heat_equation

"""Heat equation in strong form."""

from typing import Callable

import torch

from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace
from scimba_torch.physical_models.elliptic_pde.laplacians import (
    Laplacian1DDirichletStrongForm,
    Laplacian1DNeumannStrongForm,
    Laplacian2DDirichletStrongForm,
)
from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import (
    GenericFirstOrderTemporalPDE,
)
from scimba_torch.utils.scimba_tensors import LabelTensor
from scimba_torch.utils.typing_protocols import VarArgCallable




class HeatEquation1DStrongForm(GenericFirstOrderTemporalPDE):
    """Implementation of a 1D heat equation with Neumann BCs in strong form.

    Args:
        space: The approximation space for the problem
        init: Callable for the initial condition
        f: Source term function (default is zero)
        g: Neumann boundary condition function (default is zero)
        **kwargs: Additional keyword arguments
    """

    def __init__(
        self,
        space: AbstractApproxSpace,
        init: Callable,
        f: Callable | None = None,
        g: Callable | None = None,
        **kwargs,
    ):
        # Rémi says: below we initialize rhs and bc_rhs of the
        # space component to default values because it will not be used;
        # rhs and bc_rhs of the time-dependent PDE are used instead
        space_component = Laplacian1DNeumannStrongForm(space)
        super().__init__(space, space_component, init, f, g, **kwargs)





class HeatEquation1DDirichletStrongForm(GenericFirstOrderTemporalPDE):
    r"""Implementation of a 1D heat equation with Dirichlet BCs in strong form.

    Args:
        space: The approximation space for the problem
        init: Callable for the initial condition
        f: Source term function (default is zero)
        g: Dirichlet boundary condition function (default is zero)
        **kwargs: Additional keyword arguments
    """

    def __init__(
        self,
        space: AbstractApproxSpace,
        init: Callable,
        f: Callable | None = None,
        g: Callable | None = None,
        **kwargs,
    ):
        # Rémi says: below we initialize rhs and bc_rhs of the
        # space component to default values because it will not be used;
        # rhs and bc_rhs of the time-dependent PDE are used instead
        space_component = Laplacian1DDirichletStrongForm(space)
        super().__init__(space, space_component, init, f, g, **kwargs)





class HeatEquation2DStrongForm(GenericFirstOrderTemporalPDE):
    r"""Implementation of a 2D Laplacian problem with Dirichlet BCs in strong form.

    Args:
        space: The approximation space for the problem
        init: Callable for the initial condition
        f: Source term function (default is zero)
        g: Dirichlet boundary condition function (default is zero)
        **kwargs: Additional keyword arguments
    """

    def __init__(
        self,
        space: AbstractApproxSpace,
        init: Callable,
        f: Callable | None = None,
        g: Callable | None = None,
        **kwargs,
    ):
        # Rémi says: below we initialize rhs and bc_rhs of the
        # space component to default values because it will not be used;
        # rhs and bc_rhs of the time-dependent PDE are used instead
        space_component = Laplacian2DDirichletStrongForm(space)
        super().__init__(space, space_component, init, f, g, **kwargs)





class HeatEquation2DStrongFormImplicit(GenericFirstOrderTemporalPDE):
    r"""2D heat equation for implicit discrete_pinns.

    Args:
        space: the approx. space.
        init: the rhs of the initial condition.
        f: the rhs of the residual (default is zero).
        g: the rhs of the boundary condition (default is zero).
        **kwargs: Additional keyword arguments.
    """

    dt: float
    alpha: float

    def __init__(
        self,
        space: AbstractApproxSpace,
        init: Callable,
        f: Callable | None = None,
        g: Callable | None = None,
        **kwargs,
    ):
        # Rémi says: below we initialize rhs and bc_rhs of the
        # space component to default values because it will not be used;
        # rhs and bc_rhs of the time-dependent PDE are used instead
        space_component = Laplacian2DDirichletStrongForm(space)
        super().__init__(space, space_component, init, f, g, **kwargs)

        self.dt = kwargs.get("dt", 1e-3)
        self.alpha = kwargs.get("alpha", 1.0)



    def functional_operator(
        self,
        func: VarArgCallable,
        # t: LabelTensor,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the functional operator.

        Args:
            func: Callable representing the function
            x: Spatial coordinate tensor
            mu: Parameter tensor
            theta: Additional parameters tensor

        Returns:
            Functional operator tensor
        """
        # space operator
        # grad_func = torch.func.jacrev(func, 0)
        # space_op = -mu[0] * torch.func.jacrev(grad_func, 0)(x, mu, theta)
        # return func(x, mu, theta) - self.alpha * self.dt * space_op[0, 0]
        space_op = self.space_component.functional_operator(func, x, mu, theta)
        return func(x, mu, theta) - self.alpha * self.dt * space_op


    # Dirichlet conditions


    def functional_operator_bc(
        self,
        func: VarArgCallable,
        t: LabelTensor,
        x: torch.Tensor,
        n: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the functional operator for boundary conditions.

        Args:
            func: Callable representing the function
            t: Temporal coordinate tensor
            x: Spatial coordinate tensor
            n: Normal vector tensor
            mu: Parameter tensor
            theta: Additional parameters tensor

        Returns:
            Functional operator tensor for boundary conditions
        """
        return func(x, mu, theta)




    def functional_operator_ic(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the functional operator for initial conditions.

        Args:
            func: Callable representing the function
            x: Spatial coordinate tensor
            mu: Parameter tensor
            theta: Additional parameters tensor

        Returns:
            Functional operator tensor for initial conditions
        """
        return func(x, mu, theta)