From eed6afb6f2ef9e9c8da44c566f2ed0fc79b34caf Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sat, 26 Oct 2024 21:50:30 -0700 Subject: [PATCH] Matrix-free example example with a discretized PDE --- docs/src/matrix_free.md | 139 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) diff --git a/docs/src/matrix_free.md b/docs/src/matrix_free.md index e5519bf0b..95b87bc86 100644 --- a/docs/src/matrix_free.md +++ b/docs/src/matrix_free.md @@ -290,6 +290,145 @@ u_star = -sin.(x) u_star ≈ u_sol ``` +## Example with discretized PDE + +### Example 4: Solving the 3D Helmholtz equation + +The Helmholtz equation in 3D is a fundamental equation used in various fields like acoustics, electromagnetism, and quantum mechanics to model stationary wave phenomena. + +The equation is given by: +```math +\nabla^2 u(x,y,z) + k^2 u(x,y,z) = f(x,y,z) +``` +In this equation, $u(x, y, z)$ represents the unknown function, which could describe a pressure field in acoustics, a scalar potential in electromagnetism, or a wave function in quantum mechanics. +The operator $\nabla^2$ denotes the Laplacian in three dimensions. The wave number $k$ is related to the frequency of the wave through the equation $k = \frac{2\pi}{\lambda}$, where $\lambda$ is the wavelength. +Finally, $f(x,y,z)$ is a source term that drives the wave phenomena, acting as a forcing function or external influence. + +To discretize the Helmholtz equation, we use finite differences on a uniform 3D grid with grid spacings $\Delta x$, $\Delta y$, and $\Delta z$. +For a grid point $(i, j, k)$, the second derivatives are approximated as follows: + +- In the $x$-direction: +```math +\frac{\partial^2 u}{\partial x^2} \approx \frac{u_{i+1,j,k} - 2u_{i,j,k} + u_{i-1,j,k}}{\Delta x^2} +``` +- In the $y$-direction: +```math +\frac{\partial^2 u}{\partial y^2} \approx \frac{u_{i,j+1,k} - 2u_{i,j,k} + u_{i,j-1,k}}{\Delta y^2} +``` +- In the $z$-direction: +```math +\frac{\partial^2 u}{\partial z^2} \approx \frac{u_{i,j,k+1} - 2u_{i,j,k} + u_{i,j,k-1}}{\Delta z^2} +``` +Combining these, the discretized Helmholtz equation becomes: +```math +\frac{u_{i+1,j,k} - 2u_{i,j,k} + u_{i-1,j,k}}{\Delta x^2} + \frac{u_{i,j+1,k} - 2u_{i,j,k} + u_{i,j-1,k}}{\Delta y^2} + \frac{u_{i,j,k+1} - 2u_{i,j,k} + u_{i,j,k-1}}{\Delta z^2} + k^2 u_{i,j,k} = f_{i,j,k} +``` + +This discretization results in an equation at each grid point, resulting in a large and sparse linear system when assembled across the entire 3D grid. +To simplify the example, we impose Dirichlet boundary conditions with the solution $u(x, y, z) = 0$ on the boundary of the cubic domain. + +Explicitly constructing this large sparse matrix is often impractical and unnecessary. +Instead, we can define a function that directly applies the Helmholtz operator to the 3D grid, avoiding the need to form the matrix explicitly. + +Krylov.jl operates on vectors, so we must vectorize both the solution and the computational domain. +However, we can still maintain the structure of the original 3D operator by using `reshape` and `vec`. +This approach enables a simpler and efficient application of the operator in 3D while leveraging the vectorized framework for linear algebra operations. + +```@example helmholtz +using Krylov, LinearAlgebra + +# Parameters +L = 1.0 # Length of the cubic domain +Nx = 200 # Number of interior grid points in x +Ny = 200 # Number of interior grid points in y +Nz = 200 # Number of interior grid points in z +Δx = L / (Nx + 1) # Grid spacing in x +Δy = L / (Ny + 1) # Grid spacing in y +Δz = L / (Nz + 1) # Grid spacing in z +wavelength = 0.5 # Wavelength of the wave +k = 2 * π / wavelength # Wave number + +# Create the grid points +x = 0:Δx:L # Points in x dimension (Nx + 2) +y = 0:Δy:L # Points in y dimension (Ny + 2) +z = 0:Δz:L # Points in z dimension (Nz + 2) + +# Define a matrix-free Helmholtz operator +struct HelmholtzOperator + Nx::Int + Ny::Int + Nz::Int + Δx::Float64 + Δy::Float64 + Δz::Float64 + k::Float64 +end + +Base.size(A::HelmholtzOperator) = (A.Nx * A.Ny * A.Nz, A.Nx * A.Ny * A.Nz) + +Base.eltype(A::HelmholtzOperator) = Float64 + +function LinearAlgebra.mul!(y::Vector, A::HelmholtzOperator, u::Vector) + # Reshape vectors y and u into 3D arrays + U = reshape(u, A.Nx, A.Ny, A.Nz) + Y = reshape(y, A.Nx, A.Ny, A.Nz) + + # Apply the discrete Laplacian in 3D with k^2 * u + for i in 1:A.Nx + for j in 1:A.Ny + for k in 1:A.Nz + if i == 1 + dx2 = (U[i+1,j,k] -2 * U[i,j,k]) / (A.Δx)^2 + elseif i == A.Nx + dx2 = (-2 * U[i,j,k] + U[i-1,j,k]) / (A.Δx)^2 + else + dx2 = (U[i+1,j,k] -2 * U[i,j,k] + U[i-1,j,k]) / (A.Δx)^2 + end + + if j == 1 + dy2 = (U[i,j+1,k] -2 * U[i,j,k]) / (A.Δy)^2 + elseif j == A.Ny + dy2 = (-2 * U[i,j,k] + U[i,j-1,k]) / (A.Δy)^2 + else + dy2 = (U[i,j+1,k] -2 * U[i,j,k] + U[i,j-1,k]) / (A.Δy)^2 + end + + if k == 1 + dz2 = (U[i,j,k+1] -2 * U[i,j,k]) / (A.Δz)^2 + elseif k == A.Nz + dz2 = (-2 * U[i,j,k] + U[i,j,k-1]) / (A.Δz)^2 + else + dz2 = (U[i,j,k+1] -2 * U[i,j,k] + U[i,j,k-1]) / (A.Δz)^2 + end + + Y[i,j,k] = dx2 + dy2 + dz2 + (A.k)^2 * U[i,j,k] + end + end + end + + return y +end + +# Create the matrix-free operator for the Helmholtz equation +A = HelmholtzOperator(Nx, Ny, Nz, Δx, Δy, Δz, k) + +# Source term f(x, y, z) = -2k² * sin(kx) * sin(ky) * sin(kz) +F = [-2 * k^2 * sin(k * x[ii+1]) * sin(k * y[jj+1]) * sin(k * z[kk+1]) for ii in 1:Nx, jj in 1:Ny, kk in 1:Nz] +f = vec(F) + +# Solve the linear system using MinAres +u_sol, stats = minares(A, f, atol=1e-10, rtol=0.0, verbose=1) + +# Solution as 3D array +U_sol = reshape(u_sol, Nx, Ny, Nz) + +# Exact solution u(x,y,z) = sin(kx) * sin(ky) * sin(kz) +U_star = [sin(k * x[ii+1]) * sin(k * y[jj+1]) * sin(k * z[kk+1]) for ii in 1:Nx, jj in 1:Ny, kk in 1:Nz] + +# Compute the maximum error between the numerical solution U_sol and the exact solution U_star +norm(U_sol - U_star, Inf) +``` + Note that preconditioners can be also implemented as abstract operators. For instance, we could compute the Cholesky factorization of $M$ and $N$ and create linear operators that perform the forward and backsolves.