Source code for matrix_toolkit.pde.equations.wave

"""
Wave equation
"""

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

from matrix_toolkit.pde.base import BasePDEEquation
from matrix_toolkit.pde.config import PDEConfig  




class WaveEquation(BasePDEEquation):
    """
    The wave equation:
    
    .. math::
    
        \frac{\partial^2 u}{\partial t^2} = c^2 \nabla^2 u
    
    Discretized in time using finite differences:
    
    .. math::
    
        \frac{u^{n+1} - 2u^n + u^{n-1}}{\Delta t^2} = c^2 \nabla^2 u^n
        
    Examples:
        >>> config = PDEConfig(
        ...     dimension=2,
        ...     mesh_size=64,
        ...     coefficients={'wave_speed': 1.0, 'dt': 0.01}
        ... )
        >>> eq = WaveEquation(config)
        >>> M = eq.generate_matrix()
    """
    


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

    


    def get_default_coefficients(self) -> Dict[str, Any]:
        return {
            'wave_speed': 1.0,   # c in c²∇²u
            'dt': 0.01,          # Time step
        }

    


    def generate_matrix(self) -> sparse.spmatrix:
        """
        Generate system matrix for wave equation
        
        Returns matrix M for: M u^{n+1} = RHS(u^n, u^{n-1})
        """
        coeffs = self.config.coefficients or self.get_default_coefficients()
        
        if self.config.random_params:
            coeffs = self.randomize_coefficients(self.config.random_seed)
        
        c = coeffs.get('wave_speed', 1.0)
        dt = coeffs.get('dt', 0.01)
        
        # Generate based on dimension
        if self.config.dimension == 1:
            return self._generate_1d(c, dt)
        elif self.config.dimension == 2:
            return self._generate_2d(c, dt)
        elif self.config.dimension == 3:
            return self._generate_3d(c, dt)

    
    def _generate_1d(self, c: float, dt: float) -> sparse.spmatrix:
        """Generate 1D wave equation matrix"""
        from matrix_toolkit.pde.dim1d.generators import generate_1d_wave
        
        n = self.config.mesh_size[0]
        h = self.config.get_mesh_spacing()[0]
        
        return generate_1d_wave(n, dt, c, h)
    
    def _generate_2d(self, c: float, dt: float) -> sparse.spmatrix:
        """Generate 2D wave equation matrix"""
        from matrix_toolkit.pde.dim2d.generators import generate_2d_wave
        
        nx, ny = self.config.mesh_size
        hx, hy = self.config.get_mesh_spacing()
        
        return generate_2d_wave(nx, ny, dt, c, hx, hy)
    
    def _generate_3d(self, c: float, dt: float) -> sparse.spmatrix:
        """Generate 3D wave equation matrix"""
        from matrix_toolkit.pde.dim3d.generators import generate_3d_wave
        
        nx, ny, nz = self.config.mesh_size
        hx, hy, hz = self.config.get_mesh_spacing()
        
        return generate_3d_wave(nx, ny, nz, dt, c, hx, hy, hz)