Source code for matrix_toolkit.pde.equations.heat

"""
Heat equation: ∂u/∂t = α∇²u
"""

import numpy as np
from scipy import sparse
from typing import Dict, Any

from matrix_toolkit.pde.base import BasePDEEquation
from matrix_toolkit.pde.config import PDEConfig  # 添加这行




class HeatEquation(BasePDEEquation):
    """
    The heat/diffusion equation:
    
    .. math::
    
        \frac{\partial u}{\partial t} = \alpha \nabla^2 u
    
    Time discretization using the theta-method:
    
    - :math:`\theta = 0`: Forward Euler (explicit)
    - :math:`\theta = 0.5`: Crank-Nicolson
    - :math:`\theta = 1`: Backward Euler (implicit)
    
    Examples:
        >>> config = PDEConfig(
        ...     dimension=2,
        ...     mesh_size=32,
        ...     coefficients={'dt': 0.01, 'theta': 0.5}
        ... )
        >>> eq = HeatEquation(config)
        >>> M = eq.generate_matrix()  # System matrix
    """
    


    def get_operator_name(self) -> str:
        return "Heat"

    


    def get_default_coefficients(self) -> Dict[str, Any]:
        return {
            'diffusion': 1.0,    # Thermal diffusivity α
            'dt': 0.01,          # Time step
            'theta': 0.5,        # Theta parameter (0.5 = Crank-Nicolson)
        }

    


    def generate_matrix(self) -> sparse.spmatrix:
        """
        Generate system matrix for heat equation
        
        Returns matrix M such that: M u^{n+1} = RHS(u^n)
        """
        coeffs = self.config.coefficients or self.get_default_coefficients()
        
        if self.config.random_params:
            coeffs = self.randomize_coefficients(self.config.random_seed)
        
        alpha = coeffs.get('diffusion', 1.0)
        dt = coeffs.get('dt', 0.01)
        theta = coeffs.get('theta', 0.5)
        
        # Get spatial discretization (Laplacian)
        from matrix_toolkit.pde.equations.poisson import PoissonEquation
        
        # Create temporary config for Laplacian
        laplacian_config = PDEConfig(
            dimension=self.config.dimension,
            domain=self.config.domain,
            mesh_size=self.config.mesh_size,
            boundary_condition=self.config.boundary_condition
        )
        
        poisson = PoissonEquation(laplacian_config)
        L = poisson.generate_matrix()
        
        # Identity matrix
        n = self.config.total_dofs()
        I = sparse.eye(n, format='csr')
        
        # Theta-method: (I - θ dt α L) u^{n+1} = (I + (1-θ) dt α L) u^n
        M = I - theta * dt * alpha * L
        
        return M

    


    def generate_rhs_matrix(self) -> sparse.spmatrix:
        """Generate RHS matrix for explicit part"""
        coeffs = self.config.coefficients or self.get_default_coefficients()
        
        alpha = coeffs.get('diffusion', 1.0)
        dt = coeffs.get('dt', 0.01)
        theta = coeffs.get('theta', 0.5)
        
        from matrix_toolkit.pde.equations.poisson import PoissonEquation
        
        laplacian_config = PDEConfig(
            dimension=self.config.dimension,
            domain=self.config.domain,
            mesh_size=self.config.mesh_size,
            boundary_condition=self.config.boundary_condition
        )
        
        poisson = PoissonEquation(laplacian_config)
        L = poisson.generate_matrix()
        
        n = self.config.total_dofs()
        I = sparse.eye(n, format='csr')
        
        R = I + (1 - theta) * dt * alpha * L
        
        return R