Source code for matrix_toolkit.pde.config

"""
Configuration classes for PDE matrix generation
"""

from dataclasses import dataclass, field
from typing import Optional, Tuple, Union, List, Dict, Any
from enum import Enum
import numpy as np




class BoundaryCondition(Enum):
    """Boundary condition types"""
    DIRICHLET = "dirichlet"      # u = 0 on boundary
    NEUMANN = "neumann"          # du/dn = 0 on boundary
    PERIODIC = "periodic"        # periodic BC
    ROBIN = "robin"              # alpha*u + beta*du/dn = 0
    MIXED = "mixed"              # different BC on different sides





class DiscretizationType(Enum):
    """Discretization method types"""
    FINITE_DIFFERENCE = "fd"     # Finite differences
    FINITE_ELEMENT = "fe"        # Finite elements (basic)
    FINITE_VOLUME = "fv"         # Finite volumes





class StencilType(Enum):
    """Finite difference stencil types"""
    SECOND_ORDER = 2             # Second order accurate
    FOURTH_ORDER = 4             # Fourth order accurate
    SIXTH_ORDER = 6              # Sixth order accurate





@dataclass
class PDEConfig:
    """Configuration for PDE matrix generation"""
    
    # Domain configuration
    dimension: int = 2                          # 1D, 2D, or 3D
    domain: Tuple[float, ...] = (0.0, 1.0)     # Domain bounds per dimension
    
    # Mesh configuration
    mesh_size: Union[int, Tuple[int, ...]] = 32  # Grid points per dimension
    mesh_type: str = "uniform"                    # uniform, stretched, random
    
    # Discretization
    discretization: DiscretizationType = DiscretizationType.FINITE_DIFFERENCE
    stencil_order: StencilType = StencilType.SECOND_ORDER
    
    # Boundary conditions
    boundary_condition: Union[BoundaryCondition, List[BoundaryCondition]] = BoundaryCondition.DIRICHLET
    
    # PDE parameters
    coefficients: Dict[str, Any] = field(default_factory=dict)
    
    # Random parameter generation
    random_params: bool = False
    random_seed: Optional[int] = None
    param_ranges: Dict[str, Tuple[float, float]] = field(default_factory=dict)
    
    # Output format
    backend: str = "scipy"           # scipy, cupy, jax, torch, numpy
    format: str = "csr"              # csr, csc, coo, dense, etc.
    dtype: str = "float64"           # float32, float64, complex64, complex128
    
    # Additional options
    include_rhs: bool = False        # Include right-hand side vector
    include_exact_solution: bool = False  # Include exact solution if known
    normalize: bool = False          # Normalize matrix



    def __post_init__(self):
        """Validate and process configuration"""
        # Ensure domain is tuple
        if isinstance(self.domain, (int, float)):
            self.domain = (0.0, float(self.domain))
        
        # Ensure mesh_size matches dimension
        if isinstance(self.mesh_size, int):
            self.mesh_size = tuple([self.mesh_size] * self.dimension)
        elif len(self.mesh_size) != self.dimension:
            raise ValueError(
                f"mesh_size length ({len(self.mesh_size)}) "
                f"must match dimension ({self.dimension})"
            )
        
        # Reshape domain to match dimension
        if len(self.domain) == 2 and self.dimension > 1:
            # Assume same bounds for all dimensions
            self.domain = tuple(self.domain * self.dimension)
        elif len(self.domain) != 2 * self.dimension:
            raise ValueError(
                f"domain must have {2 * self.dimension} values "
                f"for {self.dimension}D problem"
            )

    


    def get_mesh_spacing(self) -> Tuple[float, ...]:
        """Calculate mesh spacing for each dimension"""
        spacings = []
        for i in range(self.dimension):
            x_min = self.domain[2*i]
            x_max = self.domain[2*i + 1]
            n = self.mesh_size[i]
            h = (x_max - x_min) / (n - 1)
            spacings.append(h)
        return tuple(spacings)

    


    def total_dofs(self) -> int:
        """Calculate total degrees of freedom"""
        return int(np.prod(self.mesh_size))