diff --git a/docs/Project.toml b/docs/Project.toml index 42f7b47da..7f842d7ca 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,6 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" HarwellRutherfordBoeing = "ce388394-9b3f-5993-a911-eb95552e4f2e" ILUZero = "88f59080-6952-5380-9ea5-54057fb9a43f" diff --git a/docs/src/matrix_free.md b/docs/src/matrix_free.md index c77ce346c..b1c639595 100644 --- a/docs/src/matrix_free.md +++ b/docs/src/matrix_free.md @@ -152,6 +152,129 @@ opJ = LinearOperator(Float64, 3, 2, false, false, (y, v) -> J(y, v), lsmr(opJ, -F(xk)) ``` +### Example 3: Solving the Poisson equation with FFT and IFFT + +In applications related to partial differential equations (PDEs), linear systems can arise from discretizing differential operators. +Storing such operators as explicit matrices is computationally expensive and unnecessary when matrix-free methods can be used, particularly with structured grids. + +The FFT is an algorithm that computes the discrete Fourier transform (DFT) of a sequence, transforming data from the spatial domain to the frequency domain. +In the context of solving PDEs, it simplifies the application of differential operators like the Laplacian by converting derivatives into algebraic operations. + +For a function $u(x)$ discretized on a periodic grid with $n$ points, the FFT of $u$ is: + +```math +\hat{u}_k = \sum_{j=0}^{n-1} u_j e^{-i k x_j}, +``` + +where $\hat{u}_k$ represents the Fourier coefficients for the frequency $k$, and $u_j$ is the value of $u$ at the grid point $x_j$ defined as $x_j = \frac{2 \pi j}{L}$ with period $L$. +The inverse FFT (IFFT) reconstructs $u$ from its Fourier coefficients: + +```math +u_j = \frac{1}{n} \sum_{k=0}^{n-1} \hat{u}_k e^{i k x_j}. +``` + +In Fourier space, the Laplacian operator $\frac{d^2}{dx^2}$ becomes a simple multiplication by $-k^2$, where $k$ is the wavenumber derived from the grid size. +This transforms the Poisson equation $\frac{d^2 u(x)}{dx^2} = f(x)$ into an algebraic equation in the frequency domain: + +```math +-k^2 \hat{u}_k = \hat{f}_k. +``` + +By solving for $\hat{u}_k$ and applying the inverse FFT, we can recover the solution $u(x)$ efficiently. + +The inverse FFT (IFFT) is used to convert data from the frequency domain back to the spatial domain. +Once the solution in frequency space is obtained by dividing the Fourier coefficients $\hat{f}_k$ by $-k^2$, the IFFT is applied to transform the result back to the original grid points in the spatial domain. + +This example consists of solving the 1D Poisson equation on a periodic domain $[0, 4\pi]$: + +```math +\frac{d^2 u(x)}{dx^2} = f(x), +``` + +where $u(x)$ is the unknown solution, and $f(x)$ is the given source term. +We solve this equation using [FFTW.jl](https://github.com/JuliaMath/FFTW.jl) to compute the matrix-free action of the Laplacian within the conjugate gradient solver. + +```@example fft_poisson +using FFTW, Krylov, LinearAlgebra + +# Define the problem size and domain +n = 32768 # Number of grid points (2^15) +L = 4π # Length of the domain +x = LinRange(0, L, n+1)[1:end-1] # Periodic grid (excluding the last point) + +# Define the source term f(x) +f = sin.(x) + +# Define a matrix-free operator using FFT and IFFT +struct FFTPoissonOperator + n::Int + L::Float64 + complex::Bool + k::Vector{Float64} # Store Fourier wave numbers +end + +function FFTPoissonOperator(n::Int, L::Float64, complex::Bool) + if complex + k = Vector{Float64}(undef, n) + else + k = Vector{Float64}(undef, n÷2 + 1) + end + k[1] = 0 # DC component -- f(x) = sin(x) has a mean of 0 + for j in 1:(n÷2) + k[j+1] = 2 * π * j / L # Positive wave numbers + end + if complex + for j in 1:(n÷2 - 1) + k[n-j+1] = -2 * π * j / L # Negative wave numbers + end + end + return FFTPoissonOperator(n, L, complex, k) +end + +Base.size(A::FFTPoissonOperator) = (n, n) + +function Base.eltype(A::FFTPoissonOperator) + type = A.complex ? ComplexF64 : Float64 + return type +end + +function LinearAlgebra.mul!(y::Vector, A::FFTPoissonOperator, u::Vector) + # Transform the input vector `u` to the frequency domain using `fft` or `rfft`. + # If the operator is complex, use the full FFT; otherwise, use the real FFT. + if A.complex + u_hat = fft(u) + else + u_hat = rfft(u) + end + + # In Fourier space, solve the system by multiplying with -k^2 (corresponding to the second derivative). + # This step applies the Laplacian operator in the frequency domain. + u_hat .= -u_hat .* (A.k .^ 2) + + # Transform the result back to the spatial domain using `ifft` or `irfft`. + # If the operator is complex, use the full inverse FFT; otherwise, use the inverse real FFT. + if A.complex + y .= ifft(u_hat) + else + y .= irfft(u_hat, A.n) + end + + return y +end + + +# Create the matrix-free operator for the Poisson equation +complex = false +A = FFTPoissonOperator(n, L, complex) + +# Solve the linear system using CR +u_sol, stats = cg(A, f, atol=1e-10, rtol=0.0, verbose=1) + +# The exact solution is u(x) = -sin(x) +u_star = -sin.(x) +u_star ≈ u_sol +``` + 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.