"""
Base classes for PDE matrix generation
"""
from abc import ABC, abstractmethod
from typing import Any, Optional, Dict, Union, List
import numpy as np
from matrix_toolkit.pde.config import PDEConfig
from matrix_toolkit.converters import ConverterFactory
class BasePDEEquation(ABC):
"""Base class for PDE equations"""
def __init__(self, config: PDEConfig):
"""
Initialize PDE equation
Args:
config: PDE configuration
"""
self.config = config
self.name = self.__class__.__name__
@abstractmethod
def get_operator_name(self) -> str:
"""Get the name of the differential operator"""
pass
@abstractmethod
def get_default_coefficients(self) -> Dict[str, Any]:
"""Get default PDE coefficients"""
pass
def randomize_coefficients(self, seed: Optional[int] = None) -> Dict[str, Any]:
"""
Generate random coefficients within specified ranges
Args:
seed: Random seed
Returns:
Dictionary of randomized coefficients
"""
if seed is not None:
np.random.seed(seed)
coeffs = self.get_default_coefficients()
if self.config.random_params and self.config.param_ranges:
for param, (min_val, max_val) in self.config.param_ranges.items():
if param in coeffs:
coeffs[param] = np.random.uniform(min_val, max_val)
return coeffs
@abstractmethod
def generate_matrix(self) -> Any:
"""Generate the discretized matrix"""
pass
def generate_rhs(self, source_function=None) -> np.ndarray:
"""
Generate right-hand side vector
Args:
source_function: Source term function f(x) or f(x,y) or f(x,y,z)
Returns:
Right-hand side vector
"""
# Default: zero RHS
return np.zeros(self.config.total_dofs())
[docs]
class PDEMatrixGenerator:
"""
Main interface for generating PDE matrices
Examples:
>>> from matrix_toolkit.pde import PDEMatrixGenerator, PDEConfig
>>>
>>> # 2D Poisson equation
>>> config = PDEConfig(
... dimension=2,
... mesh_size=64,
... backend='scipy',
... format='csr'
... )
>>>
>>> gen = PDEMatrixGenerator('poisson', config)
>>> A = gen.generate()
>>> print(A.shape) # (4096, 4096)
"""
[docs]
def __init__(self, equation_type: str, config: PDEConfig):
"""
Initialize PDE matrix generator
Args:
equation_type: Type of equation ('poisson', 'heat', 'wave', etc.)
config: PDE configuration
"""
self.equation_type = equation_type.lower()
self.config = config
self.equation = self._create_equation()
def _create_equation(self) -> BasePDEEquation:
"""Create appropriate equation object"""
from matrix_toolkit.pde.equations import (
PoissonEquation,
HeatEquation,
WaveEquation,
ConvectionDiffusionEquation,
HelmholtzEquation,
StokesEquation,
ElasticityEquation,
MaxwellEquation,
ReactionDiffusionEquation,
FisherKPPEquation,
GrayScottEquation,
NavierStokesEquation,
AdvectionDiffusionEquation,
BurgersEquation,
Burgers2DEquation,
SchrodingerEquation,
KleinGordonEquation,
)
equation_map = {
# Basic equations
'poisson': PoissonEquation,
'heat': HeatEquation,
'wave': WaveEquation,
'convection_diffusion': ConvectionDiffusionEquation,
'cd': ConvectionDiffusionEquation,
'helmholtz': HelmholtzEquation,
# Fluid mechanics
'stokes': StokesEquation,
'navier_stokes': NavierStokesEquation,
'ns': NavierStokesEquation,
'burgers': BurgersEquation,
'burgers2d': Burgers2DEquation,
'advection_diffusion': AdvectionDiffusionEquation,
'ad': AdvectionDiffusionEquation,
# Solid mechanics
'elasticity': ElasticityEquation,
'elastic': ElasticityEquation,
# Electromagnetics
'maxwell': MaxwellEquation,
# Reaction-diffusion
'reaction_diffusion': ReactionDiffusionEquation,
'rd': ReactionDiffusionEquation,
'fisher_kpp': FisherKPPEquation,
'fisher': FisherKPPEquation,
'gray_scott': GrayScottEquation,
# Quantum mechanics
'schrodinger': SchrodingerEquation,
'klein_gordon': KleinGordonEquation,
'kg': KleinGordonEquation,
}
if self.equation_type not in equation_map:
raise ValueError(
f"Unknown equation type: {self.equation_type}. "
f"Available: {list(equation_map.keys())}"
)
return equation_map[self.equation_type](self.config)
[docs]
def generate(
self,
n_matrices: int = 1,
randomize: bool = None
) -> Union[Any, List[Any]]:
"""
Generate PDE matrix/matrices
Args:
n_matrices: Number of matrices to generate
randomize: Whether to randomize parameters (overrides config)
Returns:
Single matrix if n_matrices=1, otherwise list of matrices
"""
if randomize is None:
randomize = self.config.random_params
matrices = []
for i in range(n_matrices):
# Update random seed for each iteration
if randomize and self.config.random_seed is not None:
# Create new config with updated seed
import copy
temp_config = copy.deepcopy(self.config)
temp_config.random_seed = self.config.random_seed + i
# Temporarily replace equation's config
old_config = self.equation.config
self.equation.config = temp_config
# Generate matrix
matrix = self.equation.generate_matrix()
# Restore original config
if randomize and self.config.random_seed is not None:
self.equation.config = old_config
# Convert to desired backend and format
matrix = self._convert_matrix(matrix)
matrices.append(matrix)
return matrices[0] if n_matrices == 1 else matrices
def _convert_matrix(self, matrix: Any) -> Any:
"""Convert matrix to desired backend and format"""
# Get converter for target backend
converter = ConverterFactory.get_converter(self.config.backend)
# Convert
converted = converter.convert(
matrix,
format=self.config.format,
dtype=self.config.dtype
)
return converted
[docs]
def generate_batch(
self,
n_matrices: int,
vary_params: List[str] = None,
**param_ranges
) -> List[Any]:
"""
Generate batch of matrices with varying parameters
Args:
n_matrices: Number of matrices to generate
vary_params: Which parameters to vary
**param_ranges: Parameter ranges for randomization
Returns:
List of matrices
Examples:
>>> gen.generate_batch(
... n_matrices=10,
... vary_params=['diffusion_coeff'],
... diffusion_coeff=(0.1, 10.0)
... )
"""
if param_ranges:
self.config.param_ranges.update(param_ranges)
if vary_params:
# Only randomize specified parameters
old_ranges = self.config.param_ranges.copy()
self.config.param_ranges = {
k: v for k, v in old_ranges.items() if k in vary_params
}
matrices = self.generate(n_matrices=n_matrices, randomize=True)
if vary_params:
self.config.param_ranges = old_ranges
return matrices