Skip to content

Symmetrizing Laplacian in order to use conjugate gradient in nonuniform grids #4563

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

Draft
wants to merge 22 commits into
base: main
Choose a base branch
from

Conversation

xkykai
Copy link
Collaborator

@xkykai xkykai commented May 29, 2025

As noted by @Yixiao-Zhang, the Laplacian operator in its current form is not symmetric when grids are non-uniform. Mathematically, conjugate gradient method should be used for symmetric (semi-) positive definite operators.
This is still very much work in progress, but I will outline the mathematical details of why and how below.

The Laplacian operator is non-symmetric on nonuniform grids.

We first note that the Laplacian operator $\nabla^2$ on the staggered grid is given by

@inline function ∇²ᶜᶜᶜ(i, j, k, grid, c)
    return V⁻¹ᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_∂xᶠᶜᶜ, c) +
                                    δyᵃᶜᵃ(i, j, k, grid, Ay_∂yᶜᶠᶜ, c) +
                                    δzᵃᵃᶜ(i, j, k, grid, Az_∂zᶜᶜᶠ, c))
end

To see why it is nonsymmetric, consider the second derivative in x $\partial_x^2p$:

$$(\partial_x^2p)^{ccc} = \frac{1}{V_i^{ccc}} \delta x_i^{caa} \left[\frac{A x_i^{fcc}}{\Delta x_i^{fcc}} \delta x_i^{faa} \left[ p \right] \right]$$

where

$$\begin{align} \delta x_i ^{faa} \left[ f \right] &= f[i] - f[i-1], \\\ \delta x_i ^{caa} \left[ f \right] &= f[i+1] - f[i]. \end{align}$$

Plugging everything into $(\partial_x^2p)^{caa}$, and grouping terms by their locations we get

$$(\partial_x^2p)^{ccc} = \frac{1}{V_i^{ccc}} \left( p[i+1] \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} - p[i] \left( \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} + \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} \right) + p[i-1] \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} \right)$$

Writing this out as a matrix, we get entries of the second derivative operator as

$$(\partial_x^2)^{ccc}_{ij} = \begin{cases} -\frac{1}{V_i^{ccc}} \left( \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} + \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} \right), & \text{for } i = j,\\\ \frac{1}{V_i^{ccc}} \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} & \text{for } j = i+1,\\\ \frac{1}{V_i^{ccc}} \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} & \text{for } j = i-1,\\\ 0 & \text{otherwise}. \end{cases}$$

From this we can see that $(\partial_x^2)^{ccc}_{ij}$ is nonsymmetric when grids are nonuniform. A concrete example can be shown if one tries to write out $(\partial_x^2)^{ccc}_{12} \neq (\partial_x^2)^{ccc}_{21}$.

Symmetrizing the Laplacian operator for nonuniform grids

However, in nonuniform grids, we can perform a symmetrization to create a "symmetrized Laplacian" so that it is still amenable to conjugate gradient methods,
Let our original system be
$$Lp = b$$
where $L$ is nonsymmetric. To symmetrize the system, we use the cell volume operator $V$ and solve the transformed system

$$\begin{align} \tilde{L} \tilde{x} &= \tilde{b}, \\\ \tilde{L} &= V^{\frac{1}{2}} L V^{-\frac{1}{2}}, \\\ \tilde{x} &= V^{\frac{1}{2}} x, \\\ \tilde{b} &= V^{\frac{1}{2}} b, \\\ V_{ij} &= V_{i}^{ccc} \delta_{ij}, \\\ \end{align}$$

where $\delta_{ij}$ is the Kronecker delta. In words, $V$ is the operator that gives the volume of the cell at specific index $i$.
Now we will show that $\tilde{L}$ is symmetric.

Consider again only the $x$-direction, as it is analagous in $y$ and $z$. In matrix form,

$$\tilde{L}_{ij} = \left( V^{\frac{1}{2}} (\partial_x^2) V^{-\frac{1}{2}} \right)^{ccc}_{ij} = \begin{cases} -\frac{1}{V_i^{ccc}} \left( \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} + \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} \right), & \text{for } i = j,\\\ (V_i^{ccc} V_{i+1}^{ccc})^{-\frac{1}{2}} \, \frac{A x_{i+1}^{fcc}}{\Delta x_{i+1}^{fcc}} & \text{for } j = i+1,\\\ (V_i^{ccc} V_{i-1}^{ccc})^{-\frac{1}{2}} \, \frac{1}{V_i^{ccc}} \frac{A x_{i}^{fcc}}{\Delta x_{i}^{fcc}} & \text{for } j = i-1,\\\ 0 & \text{otherwise}. \end{cases}$$

A substitution to show $\tilde{L}_{12} = \tilde{L}_{21}$ will illustrate that $\tilde{L}_{ij}$ is indeed symmetric.

Implementation details

Where to do the volume transform?

There are 2 ways to implement this, and they differ in what goes where.

  1. At the beginning of each solve, we can transform $b$ to $\tilde{b}$, then precondition (or not) and solve the system using conjugate gradient to obtain $\tilde{p}$. After solve! for CG is completed we then transform back by doing $p = V^{-\frac{1}{2}}\tilde{p}$. In this approach, the linear operator is given $\tilde{p}$, which it then has to first transform back to $p$ to compute the symmetrized laplacian, since in practice the symmetrized linear operator is really doing $V^{\frac{1}{2}} L p$. This is what is currently implemented in my code sketch.
  2. Alternatively, the linear operator is given $p$, and it computed $\tilde{L} \tilde{p}$ in return. After the CG iteration, the $\tilde{p} \rightarrow p$ operation is performed before being passed into the linear operation. To me this seemed more clunky, which is why I have opted for option 1 for now.

What should the API be?

In my mind there are 2 ways to expose the symmetrized Laplacian to the user:

  1. We keep it under ConjugateGradientPoissonSolver,
  2. or we create a new solver (something like ConjugateGradientSymmetrizedPoissonSolver.

I am inclined to use option 1 since using the asymmetric Laplacian when grids are nonuniform is probably a no go. If we are going with option 1, I have attempted to sketch out how it might look, dispatching the appropriate form of laplacian depending on the grid type (uniform or not)

Enforcing gauge condition.

This is also another subtlety that came up when I am testing the code, but I haven't really looked closely and will investigate further. Currently the gauge condition is proposed to be enforced after every CG iteration. In the nonsymmetric case, this amounts to enforcing the gauge condition on $\tilde{x}$ rather than $x$. This also means that $x$ will not have zero mean once scaled. Thus we will likely have to think about enforcing gauge condition at different locations depending on which Laplacian is used.

Finally, here's a video of a test case that I was testing using a nonuniform grid in z, with the symmetrized Laplacian in CG with no preconditioner. Unfortunately at this stage using FourierTridiagonalPoissonSolver causes numerical instability so it still remains a work in progress.

symmetrized_nonuniform_staircase_2D_convection_noprec_cg.jld2.mp4

I have also tried the symmetrized version on a uniform grid (with and without FFTBasedPoissonSolver as preconditioner) which theoretically should perform identically as the old Laplacian we have. The solutions look fine, so perhaps this idea and implementation is still not that crazy.

Sorry this is a long explanation @glwagner @simone-silvestri happy to hear your thoughts, also I'm more than happy to explain over a zoom call since this is honestly too much text.

@simone-silvestri
Copy link
Collaborator

simone-silvestri commented May 30, 2025

The Laplacian operator should be symmetric independently of the uniformity of the grid. If the Laplacian is discretized correctly there should be no need for any transformation.

You can test it out very easily. In the testing suite, we have a file called test_matrix_poisson_solver.jl. You can extract a couple of methods from there to compute the matrix associated with a Laplacian operator. I can copy them here:

@kernel function _compute_poisson_weights(Ax, Ay, Az, grid)
    i, j, k = @index(Global, NTuple)
    Ax[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δyᶠᶜᵃ(i, j, k, grid) / Δxᶠᶜᵃ(i, j, k, grid)
    Ay[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δxᶜᶠᵃ(i, j, k, grid) / Δyᶜᶠᵃ(i, j, k, grid)
    Az[i, j, k] = Δxᶜᶜᵃ(i, j, k, grid) * Δyᶜᶜᵃ(i, j, k, grid) / Δzᵃᵃᶠ(i, j, k, grid)
end

function compute_poisson_weights(grid)
    N = size(grid)
    Ax = on_architecture(architecture(grid), zeros(N...))
    Ay = on_architecture(architecture(grid), zeros(N...))
    Az = on_architecture(architecture(grid), zeros(N...))
    C  = on_architecture(architecture(grid), zeros(grid, N...))
    D  = on_architecture(architecture(grid), zeros(grid, N...))

    launch!(architecture(grid), grid, :xyz, _compute_poisson_weights, Ax, Ay, Az, grid)

    return (Ax, Ay, Az, C, D)
end

You'll see that these weights really correspond to the entries of the matrix associated with the Laplacian operator.
You can easily compute the associated matrix for a non-unifrom grid:

julia> z = [-i + rand() / 2 for i in 10:-1:0]
11-element Vector{Float64}:
 -9.546382726500523
 -8.710469429065327
 -7.782210612640547
 -6.621562097299453
 -5.724305498989565
 -4.994706647474553
 -3.7018693176721933
 -2.8099576284091086
 -1.8047376493589797
 -0.7983261867480465
  0.1841839876367073

julia> grid = RectilinearGrid(size = (10, 10), x = (0, 1), z = z, topology=(Periodic, Flat, Bounded))
10×1×10 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── Periodic x  [0.0, 1.0)           regularly spaced with Δx=0.1
├── Flat y
└── Bounded  z  [-9.54638, 0.184184] variably spaced with min(Δz)=0.729599, max(Δz)=1.29284

Then if you do

weights = compute_poisson_weights(grid)
solver = HeptadiagonalIterativeSolver(weights; grid, preconditioner_method = nothing)

and test the symmetricity of the resulting matrix you'll see that is satisfied:

julia> A = deepcopy(solver.matrix);

julia> B = deepcopy(solver.matrix);

julia> all(A .== (B + B'I) ./ 2)
true

You can also add an immersed boundary and the matrix still remains symmetric. Note that the volume needs to be multipled on the RHS, and it should not divide the operator, exactly how we compute the RHS for the tridiagonal fourier pressure solve and how it is also shown in test_matrix_poisson_solver.jl

@simone-silvestri
Copy link
Collaborator

We can also inspect the matrix visually (for small enough problems)

julia> x = z = [-i + rand() / 2 for i in 3:-1:0]
4-element Vector{Float64}:
 -2.522319420238921
 -1.7595545526990484
 -0.9419593158192162
  0.2004281087138336

julia> grid = RectilinearGrid(; size = (3, 3), x, z , topology=(Periodic, Flat, Bounded))
3×1×3 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── Periodic x  [-2.52232, 0.200428) variably spaced with min(Δx)=0.762765, max(Δx)=1.14239
├── Flat y
└── Bounded  z  [-2.52232, 0.200428] variably spaced with min(Δz)=0.762765, max(Δz)=1.14239

julia> weights = compute_poisson_weights(grid)
([0.8007389967863805; 0.9653051420453994; 0.778338382840614;;; 0.8582990874606024; 1.0346948579546007; 0.8342882342545385;;; 1.1992610032136195; 1.4457305285525734; 1.1657117657454614], [0.5818102431531194; 0.6236329225598759; 0.871372992553168;;; 0.6236329225598759; 0.6684619713685889; 0.9340105169696403;;; 0.871372992553168; 0.9340105169696403; 1.3050490277312548], [1.0; 1.0718837110534505; 1.497692766340386;;; 0.9653051420453994; 1.0346948579546007; 1.4457305285525734;;; 0.778338382840614; 0.8342882342545385; 1.1657117657454614], [0.0; 0.0; 0.0;;; 0.0; 0.0; 0.0;;; 0.0; 0.0; 0.0], [0.0; 0.0; 0.0;;; 0.0; 0.0; 0.0;;; 0.0; 0.0; 0.0])

julia> solver  = HeptadiagonalIterativeSolver(weights, grid = grid, preconditioner_method = nothing)
Matrix-based iterative solver with:
├── Problem size: (3, 1, 3)
├── Grid: 3×1×3 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── Periodic x  [-2.52232, 0.200428) variably spaced with min(Δx)=0.762765, max(Δx)=1.14239
├── Flat y
└── Bounded  z  [-2.52232, 0.200428] variably spaced with min(Δz)=0.762765, max(Δz)=1.14239
├── Solution method: cg!
└── Preconditioner: nothing

julia> Array(solver.matrix)
9×9 Matrix{Float64}:
 -2.73135    0.965305   0.800739   0.965305   0.0        0.0        0.0        0.0        0.0
  0.965305  -2.77834    0.778338   0.0        1.03469    0.0        0.0        0.0        0.0
  0.800739   0.778338  -3.02481    0.0        0.0        1.44573    0.0        0.0        0.0
  0.965305   0.0        0.0       -3.63664    1.03469    0.858299   0.778338   0.0        0.0
  0.0        1.03469    0.0        1.03469   -3.73797    0.834288   0.0        0.834288   0.0
  0.0        0.0        1.44573    0.858299   0.834288  -4.30403    0.0        0.0        1.16571
  0.0        0.0        0.0        0.778338   0.0        0.0       -3.42333    1.44573    1.19926
  0.0        0.0        0.0        0.0        0.834288   0.0        1.44573   -3.44573    1.16571
  0.0        0.0        0.0        0.0        0.0        1.16571    1.19926    1.16571   -3.53068

@simone-silvestri
Copy link
Collaborator

I suggest to build the matrix from the operator and checking that, indeed, the matrix is nonsymmetric in your case.
You can do it by executing the linear operation on a field on 1s and reshaping the result into columns

function initialize_matrix(template_field, linear_operator!, args...)
    Nx, Ny, Nz = size(template_field)
    A = spzeros(eltype(template_field.grid), Nx*Ny*Nz, Nx*Ny*Nz)
    make_column(f) = reshape(interior(f), Nx*Ny*Nz)

    eᵢⱼₖ = similar(template_field)
    ∇²eᵢⱼₖ = similar(template_field)

    for k = 1:Nz, j in 1:Ny, i in 1:Nx
        parent(eᵢⱼₖ) .= 0
        parent(∇²eᵢⱼₖ) .= 0
        eᵢⱼₖ[i, j, k] = 1
        fill_halo_regions!(eᵢⱼₖ)
        linear_operator!(∇²eᵢⱼₖ, eᵢⱼₖ, args...)

        A[:, Ny*Nx*(k-1) + Nx*(j-1) + i] .= make_column(∇²eᵢⱼₖ)
    end
    
    return A
end

If this is the case, the operator is wrong which implies that the Laplacian is incorrectly discretized and needs to be corrected. We should have an operator that works in all cases (uniform, non-uniform, and with immersed boundaries) without specialization.

@xkykai
Copy link
Collaborator Author

xkykai commented May 30, 2025

I suggest to build the matrix from the operator and checking that, indeed, the matrix is nonsymmetric in your case.

Indeed, here's a MWE to test compute_laplacian! which is found here

function compute_laplacian!(∇²ϕ, ϕ)

using Oceananigans
using Oceananigans.Solvers: compute_laplacian!
using Oceananigans.BoundaryConditions: fill_halo_regions!
using SparseArrays

xs = collect(range(0, 1, length=6))
xs[2:end-1] .+= rand(length(xs[2:end-1])) * (1 / 6) / 5
@info xs

grid = RectilinearGrid(CPU(), Float64,
                        size = (5), 
                        halo = 4,
                        x = xs,
                        topology = (Bounded, Flat, Flat))

function initialize_matrix(template_field, linear_operator!, args...)
    Nx, Ny, Nz = size(template_field)
    A = spzeros(eltype(template_field.grid), Nx*Ny*Nz, Nx*Ny*Nz)
    make_column(f) = reshape(interior(f), Nx*Ny*Nz)

    eᵢⱼₖ = similar(template_field)
    ∇²eᵢⱼₖ = similar(template_field)

    for k = 1:Nz, j in 1:Ny, i in 1:Nx
        parent(eᵢⱼₖ) .= 0
        parent(∇²eᵢⱼₖ) .= 0
        eᵢⱼₖ[i, j, k] = 1
        fill_halo_regions!(eᵢⱼₖ)
        linear_operator!(∇²eᵢⱼₖ, eᵢⱼₖ, args...)

        A[:, Ny*Nx*(k-1) + Nx*(j-1) + i] .= make_column(∇²eᵢⱼₖ)
    end
    
    return A
end

test_field = CenterField(grid)
A = initialize_matrix(test_field, compute_laplacian!)
display(A)

and the output is

[ Info: [0.0, 0.22960116301095507, 0.4189282613194224, 0.6149861209945792, 0.804943345574523, 1.0]

5×5 SparseMatrixCSC{Float64, Int64} with 13 stored entries:
 -20.793    20.793                        
  25.2161  -52.6269   27.4108              
           26.4698  -52.8964   26.4266     
                    27.2753  -54.6216   27.3463
                             26.6313  -26.6313

This shows that compute_laplacian! and hence

@inline function ∇²ᶜᶜᶜ(i, j, k, grid, c)
is nonsymmetric indeed.
Should we fix this? How to fix it? In my mind it seemed to make sense that the Laplacian operator could be nonsymmetric on nonuniform grids, but perhaps I am missing something...

@xkykai
Copy link
Collaborator Author

xkykai commented May 30, 2025

just to add I also tested the compute_symmetric_laplacian! that I wrote:

function compute_symmetric_laplacian!(∇²ϕ, ϕ)

A_symmetric = initialize_matrix(test_field, compute_symmetric_laplacian!)
display(A_symmetric)

and it does indeed produce a symmetrized Laplacian

5×5 SparseMatrixCSC{Float64, Int64} with 13 stored entries:
 -20.793   22.898                        
  22.898  -52.6269   26.9362              
          26.9362  -52.8964   26.8476     
                   26.8476  -54.6216   26.9864
                            26.9864  -26.6313

@simone-silvestri
Copy link
Collaborator

simone-silvestri commented May 30, 2025

Should we fix this? How to fix it? In my mind it seemed to make sense that the Laplacian operator could be nonsymmetric on nonuniform grids, but perhaps I am missing something...

The Laplacian should be a symmetric operator. We can use the weights that I posted above to infer the correct discretization of the Laplacian. The discrete Laplacian should satisfy

    Axᵢ₊₁ pᵢ₊₁ + Axᵢ pᵢ₋₁ + Ayⱼ₊₁ pⱼ₊₁ + Ayⱼ pⱼ₋₁ + Azₖ₊₁ pₖ₊₁ + Azₖ pₖ₋₁
    - 2 ( Axᵢ₊₁ + Axᵢ + Ayⱼ₊₁ + Ayⱼ + Azₖ₊₁ + Azₖ ) pᵢⱼₖ  = Vᶜᶜᶜ b

where

@kernel function _compute_poisson_weights(Ax, Ay, Az, grid)
    i, j, k = @index(Global, NTuple)
    Ax[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δyᶠᶜᵃ(i, j, k, grid) / Δxᶠᶜᵃ(i, j, k, grid)
    Ay[i, j, k] = Δzᵃᵃᶜ(i, j, k, grid) * Δxᶜᶠᵃ(i, j, k, grid) / Δyᶜᶠᵃ(i, j, k, grid)
    Az[i, j, k] = Δxᶜᶜᵃ(i, j, k, grid) * Δyᶜᶜᵃ(i, j, k, grid) / Δzᵃᵃᶠ(i, j, k, grid)
end

If you write it like this, you are guaranteed to end up with a symmetric operator

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants