Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[documentation] Add a matrix-free example with FFTW.jl #883

Merged
merged 5 commits into from
Oct 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/Project.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
123 changes: 123 additions & 0 deletions docs/src/matrix_free.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
Loading