Update matrix_free.md

docs/src/matrix_free.md

@@ -297,14 +297,15 @@ u_star ≈ u_sol
 The Helmholtz equation in 3D is a fundamental equation used in various fields like acoustics, electromagnetism, and quantum mechanics to model stationary wave phenomena.
 It is given by:
 ```math
-\nabla^2 u + k^2 u = 0
+\nabla^2 u(x,y,z) + k^2 u(x,y,z) = f(x,y,z)
 ```
 where:
 - $u(x, y, z)$ is the unknown function representing, for example, a pressure field in acoustics, a scalar potential in electromagnetism, or a wave function in quantum mechanics,
 - $\nabla^2$ is the Laplacian operator in three dimensions,
-- $k$ is the wave number, related to the frequency of the wave by $k = \frac{2\pi}{\lambda}$, where $\lambda$ is the wavelength.
+- $k$ is the wave number, related to the frequency of the wave by $k = \frac{2\pi}{\lambda}$, where $\lambda$ is the wavelength,
+- $f(x,y,z)$ is a source term that drives the wave phenomena, such as a forcing function or external influence.
 
-We can discretize the Helmholtz equation using finite differences on a uniform 3D grid with grid spacings $\Delta x$, $\Delta y$, and $\Delta z$.
+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:
@@ -321,19 +322,17 @@ For a grid point $(i, j, k)$, the second derivatives are approximated as follows
 ```
 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} = 0
+\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 leads to an equation at each grid point, resulting in a large and sparse linear system when assembled across the entire 3D grid.
+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.
 
 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 need to vectorize both the solution and the computational domain.
+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 us to efficiently apply the operator in 3D while leveraging the vectorized framework for linear algebra operations.
-
-We will solve the Helmholtz equation in a cubic domain using Dirichlet boundary conditions.
+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
@@ -344,9 +343,14 @@ Nx, Ny, Nz = 40, 40, 40  # Number of grid points
 Δ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.1         # Wavelength of the wave
+wavelength = 0.5         # Wavelength of the wave
 k = 2 * π / wavelength   # Wave number
 
+# Create the grid points
+x = 0:Δx:L  # Points in x dimension
+y = 0:Δy:L  # Points in y dimension
+z = 0:Δz:L  # Points in z dimension
+
 # Define a matrix-free Helmholtz operator
 struct HelmholtzOperator
     Nx::Int
@@ -363,20 +367,49 @@ 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)
-    # We reshape the vectors y and u into 3D arrays
+    # 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 and add k^2 * u
-    for i in 2:A.Nx-1
-        for j in 2:A.Ny-1
-            for k in 2:A.Nz-1
-                Y[i,j,k] = (
-                    (U[i+1,j,k] - 2 * U[i,j,k] + U[i-1,j,k]) / (A.Δx)^2 +
-                    (U[i,j+1,k] - 2 * U[i,j,k] + U[i,j-1,k]) / (A.Δy)^2 +
-                    (U[i,j,k+1] - 2 * U[i,j,k] + U[i,j,k-1]) / (A.Δz)^2 +
-                    (A.k)^2 * U[i,j,k]
-                )
+    # 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
+                    # Forward difference on x boundary
+                    dx2 = (U[i+2,j,k] - 2 * U[i+1,j,k] + U[i,j,k]) / (A.Δx)^2
+                elseif i == A.Nx
+                    # Backward difference on x boundary
+                    dx2 = (U[i,j,k] - 2 * U[i-1,j,k] + U[i-2,j,k]) / (A.Δx)^2
+                else
+                    # Centered difference for interior points
+                    dx2 = (U[i+1,j,k] - 2 * U[i,j,k] + U[i-1,j,k]) / (A.Δx)^2
+                end
+
+                if j == 1
+                    # Forward difference on y boundary
+                    dy2 = (U[i,j+2,k] - 2 * U[i,j+1,k] + U[i,j,k]) / (A.Δy)^2
+                elseif j == A.Ny
+                    # Backward difference on y boundary
+                    dy2 = (U[i,j,k] - 2 * U[i,j-1,k] + U[i,j-2,k]) / (A.Δy)^2
+                else
+                    # Centered difference for interior points
+                    dy2 = (U[i,j+1,k] - 2 * U[i,j,k] + U[i,j-1,k]) / (A.Δy)^2
+                end
+
+                if k == 1
+                    # Forward difference on z boundary
+                    dz2 = (U[i,j,k+2] - 2 * U[i,j,k+1] + U[i,j,k]) / (A.Δz)^2
+                elseif k == A.Nz
+                    # Backward difference on z boundary
+                    dz2 = (U[i,j,k] - 2 * U[i,j,k-1] + U[i,j,k-2]) / (A.Δz)^2
+                else
+                    # Centered difference for interior points
+                    dz2 = (U[i,j,k+1] - 2 * U[i,j,k] + U[i,j,k-1]) / (A.Δz)^2
+                end
+
+                # Helmholtz operator application
+                Y[i,j,k] = dx2 + dy2 + dz2 + (A.k)^2 * U[i,j,k]
             end
         end
     end
@@ -387,43 +420,37 @@ end
 # Create the matrix-free operator for the Helmholtz equation
 A = HelmholtzOperator(Nx, Ny, Nz, Δx, Δy, Δz, k)
 
-# Create the right-hand side
-B = zeros(Nx, Ny, Nz)
-B[1, :, :] .= 1   # Set boundary at x = 0
-B[Nx, :, :] .= 1  # Set boundary at x = 1
-B[:, 1, :] .= 1   # Set boundary at y = 0
-B[:, Ny, :] .= 1  # Set boundary at y = 1
-B[:, :, 1] .= 1   # Set boundary at z = 0
-B[:, :, Nz] .= 1  # Set boundary at z = 1
+# Create the right-hand side with the source term
+# f(x, y, z) = -2k² * sin(kx) * sin(ky) * sin(kz)
+F = zeros(Nx, Ny, Nz)
+for ii in 1:Nx
+    for jj in 1:Ny
+        for kk in 1:Nz
+            F[ii,jj,kk] = -2 * k^2 * sin(k * x[ii]) * sin(k * y[jj]) * sin(k * z[kk])
+        end
+    end
+end
 
 # Vectorize the right-hand side
-b = vec(B)
-
-# u(x,y,z) = − sin(πx) * sin(πy) * sin(πz) + 1
+f = vec(F)
 
 # Solve the linear system using GMRES
-u_sol, stats = gmres(A, b, atol=1e-10, rtol=0.0, verbose=1)
+u_sol, stats = gmres(A, f, atol=1e-10, rtol=0.0, verbose=1)
 
 # Reshape the solution back to a 3D array for visualization or further processing
 U_sol = reshape(u_sol, Nx, Ny, Nz)
 
-# The exact solution is
-# u(x,y,z) = - sin(πx) * sin(πy) * sin(πz) + 1
-
-# Create the grid points
-x = 0:Δx:L  # Points in x dimension
-y = 0:Δy:L  # Points in y dimension
-z = 0:Δz:L  # Points in z dimension
-
-# Compute the exact solution
+# Compute the exact solution u(x,y,z) = sin(kx) * sin(ky) * sin(kz)
 U_star = zeros(Nx, Ny, Nz)
-for i in 1:Nx
-    for j in 1:Ny
-        for k in 1:Nz
-            U_star[i,j,k] = -sin(π * x[i]) * sin(π * y[j]) * sin(π * z[k]) + 1
+for ii in 1:Nx
+    for jj in 1:Ny
+        for kk in 1:Nz
+            U_star[ii,jj,kk] = sin(k * x[ii]) * sin(k * y[jj]) * sin(k * z[kk])
         end
     end
 end
+
+norm(U_sol - U_star)
 ```
 
 Note that preconditioners can be also implemented as abstract operators.