"""
2D PDE matrix generators
"""
import numpy as np
from scipy import sparse
from typing import Optional, Tuple
from matrix_toolkit.pde.config import PDEConfig, BoundaryCondition
[docs]
def generate_2d_laplacian(
nx: int,
ny: int,
hx: float = None,
hy: float = None,
boundary: BoundaryCondition = BoundaryCondition.DIRICHLET
) -> sparse.spmatrix:
"""
Generate 2D Laplacian matrix: -(∂²u/∂x² + ∂²u/∂y²)
Uses standard 5-point stencil:
1/(h²) * [-1 ]
[-1 4 -1]
[ -1 ]
Args:
nx: Number of grid points in x-direction
ny: Number of grid points in y-direction
hx: Grid spacing in x (default: 1/(nx+1))
hy: Grid spacing in y (default: 1/(ny+1))
boundary: Boundary condition
Returns:
Sparse CSR matrix of size (nx*ny, nx*ny)
Examples:
>>> A = generate_2d_laplacian(32, 32)
>>> A.shape
(1024, 1024)
"""
if hx is None:
hx = 1.0 / (nx + 1)
if hy is None:
hy = 1.0 / (ny + 1)
n = nx * ny
# Create 1D Laplacians
from matrix_toolkit.pde.dim1d.generators import generate_1d_laplacian
Lx = generate_1d_laplacian(nx, hx, boundary)
Ly = generate_1d_laplacian(ny, hy, boundary)
# Kronecker sum: Lx ⊗ I + I ⊗ Ly
Ix = sparse.eye(nx, format='csr')
Iy = sparse.eye(ny, format='csr')
# L = Lx ⊗ Iy + Ix ⊗ Ly
A = sparse.kron(Lx, Iy, format='csr') + sparse.kron(Ix, Ly, format='csr')
return A
[docs]
def generate_2d_convection_diffusion(
nx: int,
ny: int,
velocity: Tuple[float, float] = (1.0, 1.0),
diffusion: float = 1.0,
hx: float = None,
hy: float = None,
boundary: BoundaryCondition = BoundaryCondition.DIRICHLET
) -> sparse.spmatrix:
"""
Generate 2D convection-diffusion matrix:
-ε(∂²u/∂x² + ∂²u/∂y²) + vx ∂u/∂x + vy ∂u/∂y
Args:
nx, ny: Grid points in each direction
velocity: Convection velocity (vx, vy)
diffusion: Diffusion coefficient ε
hx, hy: Grid spacings
boundary: Boundary condition
Returns:
Sparse CSR matrix
"""
if hx is None:
hx = 1.0 / (nx + 1)
if hy is None:
hy = 1.0 / (ny + 1)
vx, vy = velocity
# Diffusion part
A_diff = diffusion * generate_2d_laplacian(nx, ny, hx, hy, boundary)
# Convection part: vx ∂/∂x + vy ∂/∂y
# ∂/∂x term
Dx = sparse.diags(
[-np.ones(nx-1), np.zeros(nx), np.ones(nx-1)],
[-1, 0, 1],
shape=(nx, nx),
format='csr'
) / (2 * hx)
# ∂/∂y term
Dy = sparse.diags(
[-np.ones(ny-1), np.zeros(ny), np.ones(ny-1)],
[-1, 0, 1],
shape=(ny, ny),
format='csr'
) / (2 * hy)
Ix = sparse.eye(nx, format='csr')
Iy = sparse.eye(ny, format='csr')
# vx ∂/∂x ⊗ Iy + Ix ⊗ vy ∂/∂y
A_conv = vx * sparse.kron(Dx, Iy, format='csr') + \
vy * sparse.kron(Ix, Dy, format='csr')
A = A_diff + A_conv
return A
[docs]
def generate_2d_heat(
nx: int,
ny: int,
dt: float,
diffusion: float = 1.0,
hx: float = None,
hy: float = None,
theta: float = 0.5
) -> sparse.spmatrix:
"""
Generate 2D heat equation matrix: ∂u/∂t = ε(∂²u/∂x² + ∂²u/∂y²)
Args:
nx, ny: Grid points
dt: Time step
diffusion: Thermal diffusivity
hx, hy: Grid spacings
theta: Implicitness parameter
Returns:
System matrix M
"""
if hx is None:
hx = 1.0 / (nx + 1)
if hy is None:
hy = 1.0 / (ny + 1)
# Laplacian
L = generate_2d_laplacian(nx, ny, hx, hy)
# Identity
n = nx * ny
I = sparse.eye(n, format='csr')
# Theta-method
M = I - theta * dt * diffusion * L
return M
[docs]
def generate_2d_wave(
nx: int,
ny: int,
dt: float,
wave_speed: float = 1.0,
hx: float = None,
hy: float = None
) -> sparse.spmatrix:
"""
Generate 2D wave equation matrix: ∂²u/∂t² = c²(∂²u/∂x² + ∂²u/∂y²)
Args:
nx, ny: Grid points
dt: Time step
wave_speed: Wave speed c
hx, hy: Grid spacings
Returns:
System matrix M
"""
if hx is None:
hx = 1.0 / (nx + 1)
if hy is None:
hy = 1.0 / (ny + 1)
# Laplacian
L = generate_2d_laplacian(nx, ny, hx, hy)
# Identity
n = nx * ny
I = sparse.eye(n, format='csr')
# Wave equation
M = I - (wave_speed * dt)**2 * L
return M
[docs]
def generate_2d_helmholtz(
nx: int,
ny: int,
wavenumber: float,
hx: float = None,
hy: float = None
) -> sparse.spmatrix:
"""
Generate 2D Helmholtz matrix: -Δu - k²u
Args:
nx, ny: Grid points
wavenumber: Wave number k
hx, hy: Grid spacings
Returns:
Sparse CSR matrix
"""
# Laplacian
A = generate_2d_laplacian(nx, ny, hx, hy)
# Add -k² term
n = nx * ny
A = A - (wavenumber**2) * sparse.eye(n, format='csr')
return A
[docs]
def generate_2d_biharmonic(
nx: int,
ny: int,
hx: float = None,
hy: float = None
) -> sparse.spmatrix:
"""
Generate 2D biharmonic matrix: Δ²u = ΔΔu
Args:
nx, ny: Grid points
hx, hy: Grid spacings
Returns:
Sparse CSR matrix
"""
# Compute Laplacian
L = generate_2d_laplacian(nx, ny, hx, hy)
# Biharmonic = Laplacian²
A = L @ L
return A