"""
1D PDE matrix generators using finite differences
"""
import numpy as np
from scipy import sparse
from typing import Optional, Tuple
from matrix_toolkit.pde.config import PDEConfig, BoundaryCondition, StencilType
[docs]
def generate_1d_laplacian(
n: int,
h: float = None,
boundary: BoundaryCondition = BoundaryCondition.DIRICHLET,
stencil_order: StencilType = StencilType.SECOND_ORDER
) -> sparse.spmatrix:
"""
Generate 1D Laplacian matrix: -d²u/dx²
Args:
n: Number of interior grid points
h: Grid spacing (default: 1/(n+1))
boundary: Boundary condition type
stencil_order: Finite difference accuracy order
Returns:
Sparse matrix in CSR format
Examples:
>>> A = generate_1d_laplacian(100)
>>> A.shape
(100, 100)
"""
if h is None:
h = 1.0 / (n + 1)
if stencil_order == StencilType.SECOND_ORDER:
# Second order: [-1, 2, -1] / h²
diagonals = [
-np.ones(n-1), # Lower diagonal
2*np.ones(n), # Main diagonal
-np.ones(n-1), # Upper diagonal
]
offsets = [-1, 0, 1]
elif stencil_order == StencilType.FOURTH_ORDER:
# Fourth order: [1, -16, 30, -16, 1] / (12h²)
if n < 3:
raise ValueError("Need at least 3 points for 4th order stencil")
diagonals = [
np.ones(n-2), # -2 offset
-16*np.ones(n-1), # -1 offset
30*np.ones(n), # 0 offset
-16*np.ones(n-1), # +1 offset
np.ones(n-2), # +2 offset
]
offsets = [-2, -1, 0, 1, 2]
# Modify boundaries to use lower order
diagonals[2][0] = 2 # First point
diagonals[2][-1] = 2 # Last point
factor = 1.0 / 12.0
else:
raise ValueError(f"Unsupported stencil order: {stencil_order}")
A = sparse.diags(diagonals, offsets, shape=(n, n), format='csr')
# Scale by 1/h²
factor = 12.0 if stencil_order == StencilType.FOURTH_ORDER else 1.0
A = A / (h**2 * factor)
# Apply boundary conditions
if boundary == BoundaryCondition.NEUMANN:
# Modify first and last rows for Neumann BC
A[0, 0] = 1.0 / h**2
A[-1, -1] = 1.0 / h**2
elif boundary == BoundaryCondition.PERIODIC:
# Add wrap-around connections
A = A.tolil()
A[0, -1] = -1.0 / h**2
A[-1, 0] = -1.0 / h**2
A = A.tocsr()
return A
[docs]
def generate_1d_convection_diffusion(
n: int,
velocity: float = 1.0,
diffusion: float = 1.0,
h: float = None,
boundary: BoundaryCondition = BoundaryCondition.DIRICHLET
) -> sparse.spmatrix:
"""
Generate 1D convection-diffusion matrix: -ε d²u/dx² + v du/dx
Args:
n: Number of grid points
velocity: Convection velocity (v)
diffusion: Diffusion coefficient (ε)
h: Grid spacing
boundary: Boundary condition
Returns:
Sparse CSR matrix
"""
if h is None:
h = 1.0 / (n + 1)
# Diffusion part: -ε d²u/dx²
A_diff = diffusion * generate_1d_laplacian(n, h, boundary)
# Convection part: v du/dx (central differences)
diagonals = [
-velocity * np.ones(n-1), # Lower diagonal
np.zeros(n), # Main diagonal (zero for central diff)
velocity * np.ones(n-1), # Upper diagonal
]
A_conv = sparse.diags(diagonals, [-1, 0, 1], shape=(n, n), format='csr')
A_conv = A_conv / (2 * h)
# Combined operator
A = A_diff + A_conv
return A
[docs]
def generate_1d_heat(
n: int,
dt: float,
diffusion: float = 1.0,
h: float = None,
theta: float = 0.5
) -> sparse.spmatrix:
"""
Generate 1D heat equation matrix: du/dt = ε d²u/dx²
Uses theta-method: implicit (θ=1), explicit (θ=0), Crank-Nicolson (θ=0.5)
Args:
n: Number of grid points
dt: Time step
diffusion: Thermal diffusivity
h: Spatial grid spacing
theta: Implicitness parameter
Returns:
System matrix M for: M u^{n+1} = ... u^n
"""
if h is None:
h = 1.0 / (n + 1)
# Laplacian
L = generate_1d_laplacian(n, h)
# Identity
I = sparse.eye(n, format='csr')
# Theta-method: (I - θ dt ε L) u^{n+1} = (I + (1-θ) dt ε L) u^n
M = I - theta * dt * diffusion * L
return M
[docs]
def generate_1d_wave(
n: int,
dt: float,
wave_speed: float = 1.0,
h: float = None
) -> sparse.spmatrix:
"""
Generate 1D wave equation matrix: d²u/dt² = c² d²u/dx²
Args:
n: Number of grid points
dt: Time step
wave_speed: Wave speed c
h: Grid spacing
Returns:
System matrix M
"""
if h is None:
h = 1.0 / (n + 1)
# Laplacian
L = generate_1d_laplacian(n, h)
# Identity
I = sparse.eye(n, format='csr')
# Wave equation discretization
M = I - (wave_speed * dt)**2 * L
return M
[docs]
def generate_1d_helmholtz(
n: int,
wavenumber: float,
h: float = None
) -> sparse.spmatrix:
"""
Generate 1D Helmholtz matrix: -d²u/dx² - k²u
Args:
n: Number of grid points
wavenumber: Wave number k
h: Grid spacing
Returns:
Sparse CSR matrix
"""
if h is None:
h = 1.0 / (n + 1)
# Laplacian
L = generate_1d_laplacian(n, h)
# Add -k² term
I = sparse.eye(n, format='csr')
A = L - (wavenumber**2) * I
return A
[docs]
def generate_1d_biharmonic(
n: int,
h: float = None
) -> sparse.spmatrix:
"""
Generate 1D biharmonic matrix: d⁴u/dx⁴
Args:
n: Number of grid points
h: Grid spacing
Returns:
Sparse CSR matrix
"""
if h is None:
h = 1.0 / (n + 1)
# Laplacian
L = generate_1d_laplacian(n, h)
# Biharmonic = Laplacian²
A = L @ L
return A