diff --git a/docs/make.jl b/docs/make.jl index e62aad428..d5f7e2cce 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -29,6 +29,7 @@ makedocs( "Warm-start" => "warm-start.md", "Matrix-free operators" => "matrix_free.md", "Callbacks" => "callbacks.md", + "Advanced features" => "internal.md", "Performance tips" => "tips.md", "Tutorials" => ["CG" => "examples/cg.md", "CAR" => "examples/car.md", diff --git a/docs/src/internal.md b/docs/src/internal.md new file mode 100644 index 000000000..d323d33d9 --- /dev/null +++ b/docs/src/internal.md @@ -0,0 +1,308 @@ +# Custom vector type for the 2D Poisson equation with halo regions + +## Introduction + +The 2D Poisson equation is a fundamental partial differential equation (PDE) widely used in physics and mathematics to model various phenomena, including temperature distribution and incompressible fluid flow. +It can be expressed as: + +```math +\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2} = f(x, y) +``` + +where: +- $u(x, y)$ is the unknown function we seek. +- $f(x, y)$ is a known function (the "source term") representing the distribution of sources or sinks in the domain. + +This equation is typically solved over a rectangular domain with boundary conditions specified at the edges. + +## Finite difference discretization + +To numerically solve the Poisson equation, we employ the finite difference method to approximate the second derivatives on a grid. +This approach involves dividing the domain into a grid of points and using differences between neighboring values to approximate derivatives. + +Assuming a square domain $[0, L] \times [0, L]$ discretized with $(N_x+2, N_y+2)$ points along each dimension (with grid spacing $h_x = \frac{L}{N_x+1}$ and $h_y = \frac{L}{N_y+1}$, let $u_{i,j}$ denote the approximation of $u(x_i, y_j)$ at the grid point $(x_i, y_j) = (ih, jh)$. + +### Discretized Laplacian + +The 2D Laplacian, $\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}$, can be approximated at each of the $N^2$ interior grid point $(i, j)$ using the central difference formula: + +```math +\frac{\partial^2 u}{\partial x^2} \approx \frac{u_{i+1,j} - 2u_{i,j} + u_{i-1,j}}{h^2} +``` + +```math +\frac{\partial^2 u}{\partial y^2} \approx \frac{u_{i,j+1} - 2u_{i,j} + u_{i,j-1}}{h^2} +``` + +Combining these yields the discrete form of the Poisson equation: + +```math +\frac{u_{i+1,j} - 2u_{i,j} + u_{i-1,j}}{h^2} + \frac{u_{i,j+1} - 2u_{i,j} + u_{i,j-1}}{h^2} = f_{i,j} +``` + +This represents a system of linear equations for the unknowns $u_{i,j}$ at each interior grid point. + +### Boundary conditions + +To complete the system, we need boundary conditions for $u$ along the domain edges. +Common options include: +- **Dirichlet boundary conditions**: specifying the value of $u$ on the boundary. +- **Neumann boundary conditions**: specifying the derivative (flux) of $u$ normal to the boundary. + +## Implementing halo regions with `MyVector` + +In practical applications, particularly in parallel computing, it is common to introduce **halo regions** (or ghost cells) around the grid. +These additional layers store boundary values from neighboring subdomains, allowing each subdomain to compute stencils near its boundaries independently without immediate communication with neighboring domains. +Halo regions simplify boundary condition management in distributed or multi-threaded environments. + +In Krylov.jl, while internal storage for each Krylov method typically expects an `AbstractVector`, specific applications can benefit from structured data layouts. +This is where a specialized vector type called **`MyVector`** comes into play. + +Using `MyVector` with halo regions enables the implementation of finite difference stencils without boundary condition checks, enhancing both readability and performance. +The **`OffsetArray`** type from the [OffsetArrays.jl](https://github.com/JuliaArrays/OffsetArrays.jl) package supports custom indexing, making it ideal for grids with halo regions. +By wrapping an `OffsetArray` within an `MyVector`, we can access elements using custom offsets that align with the grid's physical layout. +This configuration allows for ``if-less'' stencils, avoiding direct boundary condition checks within the core loop, resulting in cleaner and potentially faster code. + +The design of `MyVector` can be easily adapted for 1D, 2D, or 3D problems with minimal changes, providing flexibility in handling various grid configurations. + +## Definition of the `MyVector` + +The `MyVector` type is a specialized vector designed for efficient handling of grid-based computations, particularly in the context of finite difference methods with halo regions. +It is parameterized by: +- **`FC`**: The element type of the vector. +- **`D`**: The type of the data array, which utilizes `OffsetArray` to enable custom indexing. + +Here is the definition of the `MyVector`: + +```julia +using OffsetArrays + +struct MyVector{FC, D} <: AbstractVector{FC} + data::D + + function MyVector(data::D) where {D} + FC = eltype(data) + return new{FC, D}(data) + end +end + +# Constructor +function MyVector{FC,D}(::UndefInitializer, l::Int64) where {FC,D} + m = n = sqrt(l) |> Int + data = zeros(FC, m+2, n+2) + v = OffsetMatrix(data, 0:m+1, 0:n+1) + return MyVector(v) +end + +function Base.length(v::MyVector) + m, n = size(v.data) + l = (m - 2) * (n - 2) + return l +end + +function Base.size(v::MyVector) + l = length(v) + return (l,) +end + +function Base.getindex(v::MyVector, idx) + m, n = size(v.data) + row = div(idx-1, n-2) + 1 + col = mod(idx-1, n-2) + 1 + return v.data[row, col] +end +``` + +The functions `size` and `getindex` are defined to enable display in the REPL. + +## Using `MyVector` for the 2D Poisson equation + +By utilizing `MyVector`, we can implement the finite difference stencil for the Laplacian operator efficiently, eliminating the need for conditional checks for boundary elements. + +### Stencil implementation + +Assuming `data` is initialized as an `OffsetArray` with appropriate halo regions, we can define a matrix-free Laplacian operator and apply a typical 5-point stencil operation as follows: + +```julia +# Define a matrix-free Laplacian operator +struct LaplacianOperator + Nx::Int # Number of grid points in the x-direction + Ny::Int # Number of grid points in the y-direction + Δx::Float64 # Grid spacing in the x-direction + Δy::Float64 # Grid spacing in the y-direction +end + +# Define size and element type for the operator +Base.size(A::LaplacianOperator) = (A.Nx * A.Ny, A.Nx * A.Ny) +Base.eltype(A::LaplacianOperator) = Float64 + +function LinearAlgebra.mul!(y::MyVector{Float64}, A::LaplacianOperator, u::MyVector{Float64}) + # Apply the discrete Laplacian in 2D + for i in 1:A.Nx + for j in 1:A.Ny + # Calculate second derivatives using finite differences + dx2 = (u.data[i-1,j] - 2 * u.data[i,j] + u.data[i+1,j]) / (A.Δx)^2 + dy2 = (u.data[i,j-1] - 2 * u.data[i,j] + u.data[i,j+1]) / (A.Δy)^2 + + # Update the output vector with the Laplacian result + y.data[i,j] = dx2 + dy2 + end + end + + return y +end +``` + +### Benefits of Using `MyVector` + +Utilizing `MyVector` offers several significant benefits for solving the 2D Poisson equation: + +- **Simplified Code**: The custom indexing capabilities of `OffsetArray` allow for the incorporation of boundary data from halo regions. This eliminates the need for boundary checks within the core loop, resulting in clearer and more maintainable code. + +- **Performance**: By removing boundary checks, we reduce branching in the code, which enhances computational efficiency, particularly for large grids. +This leads to faster execution times and better overall performance. + +- **Flexibility**: `MyVector` can be easily extended to accommodate more complex stencils or additional dimensions (e.g., 3D grids) by simply adjusting the offsets. +This adaptability makes it a powerful tool for various numerical applications. + +By leveraging these advantages, we can efficiently solve the 2D Poisson equation while maintaining a clear and concise code structure. + +## Required methods for Krylov.jl compatibility + +To integrate `MyVector` with Krylov.jl, the following operations must be defined for compatibility: + +```julia +using Krylov +import Krylov.FloatOrComplex + +function Krylov.kdot(n::Integer, x::MyVector{T}, y::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + _y = y.data + res = zero(T) + for i = 1:mx-1 + for j = 1:nx-1 + res += _x[i,j] * _y[i,j] + end + end + return res +end + +function Krylov.knorm(n::Integer, x::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + res = zero(T) + for i = 1:mx-1 + for j = 1:nx-1 + res += _x[i,j]^2 + end + end + return sqrt(res) +end + +function Krylov.kscal!(n::Integer, s::T, x::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + for i = 1:mx-1 + for j = 1:nx-1 + _x[i,j] = s * _x[i,j] + end + end + return x +end + +function Krylov.kaxpy!(n::Integer, s::T, x::MyVector{T}, y::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + _y = y.data + for i = 1:mx-1 + for j = 1:nx-1 + _y[i,j] += s * _x[i,j] + end + end + return y +end + +function Krylov.kaxpby!(n::Integer, s::T, x::MyVector{T}, t::T, y::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + _y = y.data + for i = 1:mx-1 + for j = 1:nx-1 + _y[i,j] = s * _x[i,j] + t * _y[i,j] + end + end + return y +end + +function Krylov.kcopy!(n::Integer, y::MyVector{T}, x::MyVector{T}) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + _y = y.data + for i = 1:mx-1 + for j = 1:nx-1 + _y[i,j] = _x[i,j] + end + end + return y +end + +function Krylov.kfill!(x::MyVector{T}, val::T) where T <: FloatOrComplex + mx, nx = size(x.data) + _x = x.data + for i = 1:mx-1 + for j = 1:nx-1 + _x[i,j] = val + end + end + return x +end +``` + +These methods enable Krylov.jl to use custom vector types, allowing for operations like dot products, norms, scalar multiplication, and element-wise updates, which are essential for Krylov solvers. + +## Complete example + +```julia +using Krylov, LinearAlgebra, OffsetArrays + +# Parameters +L = 1.0 # Length of the square domain +Nx = 200 # Number of interior grid points in x +Ny = 200 # Number of interior grid points in y +Δx = L / (Nx + 1) # Grid spacing in x +Δy = L / (Ny + 1) # Grid spacing in y + +# Define the source term f(x,y) +f(x,y) = -2 * π * π * sin(π * x) * sin(π * y) + +# Create the matrix-free Laplacian operator +A = LaplacianOperator(Nx, Ny, Δx, Δy) + +# Create the right-hand side +rhs = zeros(Float64, Nx+2, Ny+2) +data = OffsetArray(rhs, 0:Nx+1, 0:Ny+1) +for i in 1:Nx + for j in 1:Ny + xi = i * Δx + yj = j * Δy + data[i,j] = f(xi, yj) + end +end +b = MyVector(data) + +# Solve the system with CG +u_sol, stats = Krylov.cg(A, b, atol=1e-12, rtol=0.0, verbose=1) + +# The exact solution is u(x,y) = sin(πx) * sin(πy) +u_star = [sin(π * i * Δx) * sin(π * j * Δy) for i=1:Nx, j=1:Ny] +norm(u_sol.data[1:Nx, 1:Ny] - u_star, Inf) +``` + +## Conclusion + +This overview illustrates the effective implementation of a 2D Poisson equation solver utilizing a specialized vector type storage with `MyVector` to accommodate halo regions. +By leveraging custom indexing, we significantly enhance both code readability and performance, enabling a flexible framework that can be adapted to a variety of applications. + +!!! info + The package [Oceanigans.jl](https://github.com/CliMA/Oceananigans.jl) utilizes a similar approach with its type `Field` to solve linear systems involving millions of variables efficiently with Krylov.jl.