Enforce gauge condition in preconditioners for ConjugateGradientPoissonSolver

[glwagner]

Also does a bit of clean up.

In the future we need to move ConjugateGradientPoissonSolver to Solvers module.

Also I would like to rename PreconditionedConjugateGradientSolver to just ConjugateGradientSolver cause repetitive stress injury from too much typing is a real problem.

[glwagner]

@xkykai I think this is the crucial lines. Without this the pressure does not have a zero mean and I suspect that can throw off the CG solver. But I'm checking.

[glwagner]

@jm-c might be good to have your input

[navidcy]

I have found that I needed to enforce the zero mean (e.g., when solving Laplace's or Poisson's equation) when I was using the conjugate gradient solver with @elise-palethorpe to compare with MultiGrid

[glwagner]

I think this was a key ingredient missing from previous implementation. The FFT-based preconditioner zeros out the mean over the underlying grid, but does not zero out the mean on the immersed grid.

[xkykai]

Checked the derivation, looks good to me!

[navidcy]

I like what's suggested in this PR!

[xkykai]

Checked the derivation, looks good to me!

[xkykai]

@Yixiao-Zhang here's the implementation of the gauge fixing that you proposed!

[glwagner]

@Yixiao-Zhang here's the implementation of the gauge fixing that you proposed!

Isn't this a bit different from what @Yixiao-Zhang proposed? This PR uses the same criteria that we use for the FFT-based solver without immersed boundaries, ie set the mean pressure to 0:

Oceananigans.jl/src/Solvers/fft_based_poisson_solver.jl

Lines 112 to 115 in 45838a5

	# If m === 0, the "zeroth mode" at `i, j, k = 1, 1, 1` is undetermined;
	# we set this to zero by default. Another slant on this "problem" is that
	# λx[1, 1, 1] + λy[1, 1, 1] + λz[1, 1, 1] = 0, which yields ϕ[1, 1, 1] = Inf or NaN.
	m === 0 && CUDA.@allowscalar ϕc[1, 1, 1] = 0

In what I saw, the mean pressure was not set to 0.

[Yixiao-Zhang]

It is different. I think the issue is that the conjugate gradient (CG) method can be numerically unstable for positive semi-definite matrices. I have a Python script (shown at the end of this comment) that demonstrates this. The residual keeps oscillating (see the attached figure).

Besides, if we add a small random perturbation to b so that Ax = b does not have a solution (this line is commented out for the figure shown below), the residual can grow over time and become larger than 1e4.

There are two potential solutions:

1. Modify the CG method to make it numerically stably for positive semi-definite matrices
2. Apply the CG method to another matrix that is positive definite and gives the same solution

My method belongs to the second category. I create a new matrix

$$\tilde{A} = A + cv_0v_0^\dagger$$

where $c$ is a positive real number and $v_0$ is the basis of the null space of $A$.

In contrast, enforcing the zero mean belongs to the first category. It definitely helps, since the CG method will likely blow up if the mean is not zero (which means there is no solution). However, dealing with round-off errors is always tricky, especially in this case where fast convergence leads to a rapid decrease in the "signal-to-error ratio".

import numpy as np
import matplotlib.pyplot as plt


def gram_schmidt(X, row_vecs=True, norm = True):
        if not row_vecs:
            X = X.T
        Y = X[0:1,:].copy()
        for i in range(1, X.shape[0]):
            proj = np.diag((X[i,:].dot(Y.T)/np.linalg.norm(Y,axis=1)**2).flat).dot(Y)
            Y = np.vstack((Y, X[i,:] - proj.sum(0)))
        if norm:
            Y = np.diag(1/np.linalg.norm(Y,axis=1)).dot(Y)
        if row_vecs:
            return Y
        else:
            return Y.T


def cg_solve(A,b,N):

    x = np.zeros_like(b)

    r = b - A.dot(x)

    p = r.copy()

    rs = []

    for i in range(N):

        rs.append(np.linalg.norm(r))

        Ap = A.dot(p)

        alpha = np.dot(p,r)/np.dot(p,Ap)
        x = x + alpha*p

        # r = b - A.dot(x)  # a numerically more stable version
        r -= alpha * Ap

        beta = -np.dot(r,Ap)/np.dot(p,Ap)
        p = r + beta*p

    return np.array(rs)


def main():
    N = 256

    # create a positive definite diagonal matrix
    matrix_D = np.diag(np.random.rand(N) + 1.0)

    # create a vector
    b0 = np.random.normal(size=N)
    b0 /= np.linalg.norm(b0)

    # set one eigenvalue value to zero and update the vector
    matrix_D[0, 0] = 0.0
    b0[0] = 0.0

    # create an orthogonal matrix
    matrix_q = gram_schmidt(np.random.normal(size=(N, N)))

    # create a positive semi-definite symmetric matrix
    matrix_A = np.dot(matrix_q, np.dot(matrix_D, matrix_q.T))
    b = np.dot(matrix_q, b0)

    # add a random perturbation
    # b += 1e-5 * np.random.normal(size=N)

    # return the residuals at every iteration
    rs = cg_solve(matrix_A, b, N=64)

    # plot residuals as a function of iterations
    fig, ax = plt.subplots()

    ax.plot(rs, marker='o')

    ax.set_yscale('log')

    ax.set_xlabel('Number of Iterations')
    ax.set_ylabel('Residual')

    plt.savefig('cg_residual.png')


if __name__ == '__main__':
    main()

[glwagner]

1. Modify the CG method to make it numerically stably for positive semi-definite matrices

Right, although I am not sure that this is done correctly in this PR. There may be more work to do...

My method belongs to the second category. I create a new matrix

$$\tilde{A} = A + cv_0v_0^\dagger$$

where $c$ is a positive real number and $v_0$ is the basis of the null space of $A$.

I was a bit confused about the motivation for defining a new matrix, but with some help from @amontoison I feel I understand this better. The basic idea applies to any semi-definite system; the idea is to "shift" an operator by the identity:

$$ \tilde A = A + \epsilon I $$

Where $\epsilon$ is a constant.

The key to this method is the smallness of $\epsilon$. Note that we are solving a different system: if the original system is $A x = b$, then the new system is

$$ \left (A + \epsilon I \right ) y = b $$

We do not automatically have $y=x$ or even $y \approx x$. However, we may have $y \approx x$ when $\epsilon$ is smaller than the tolerance of the underlying CG algorithm, for example.

In addition to making $\epsilon$ small, another strategy to recover $x$ from $y$ is to solve a new system:

$$ A z = b - A y $$

which, we can show, allows us to recover $x$ via $x = y + z$.

For the purpose of implementing shifted operators, some methods in Krylov.jl support provision of a "shift" as an argument to the solver, eg the argument λ here: https://jso.dev/Krylov.jl/dev/solvers/sid/#Krylov.symmlq.

It'd be easy to support a shift of the Poisson operator, the caveat being that the problem being solved is no longer exact. We could additionally support finding the exact solution through a second CG iteration via the method described above, but this would be a bit more work.

[amontoison]

@glwagner The goal of shifting the system is to recover a positive definite matrix by adding \lambda to all eigenvalue of the problem.
The proposition of @Yixiao-Zhang is to do that by adding v_0 * v_0^T. v_0 * v_0^T is ONLY semi-positive definite so not all eigenvalues will be modified. But it should lead to a positive definite A + c * v_0 * v_0^T (rank-one updating) by only updating one eigenvalue (the 0 eigenvalue).

I just don't understand why we will find the same solution because CG should return one solution among an infinity of solution (before the shift).

[glwagner]

Okay thanks to you both, I think I understand this better. I was a bit confused by notation. Let's call $e = [1, 1, ..., 1]$ the one vector and $A$ the discrete form of the Laplacian operator $\nabla^2$. Then --- on a uniform grid --- we have the equivalencies

$$ \nabla^2 \mapsto A \qquad \text{and} \qquad e^T \phi \mapsto \frac{N}{V} \int \phi \mathrm{d} V \equiv N \bar \phi $$

where $e^T$ denotes the transpose of $e$ (so $e^T \phi$ is a dot product), $V$ is the volume of the domain, $N = N_x N_y N_z$ is the total number of elements in the 3D discretization of $\phi$, and $\bar \phi$ denotes the mean over the volume. In other words, the product $e^T \phi$ is proportional to the volume-mean of $\phi$.

We write the Poisson equation

$$ \nabla^2 \phi = R $$

where $R$ is some right hand side. As we well understand, $\bar \phi$ is undetermined from this equation. Physically, pressure in an incompressible system is indeterminant up to a constant; only the pressure gradient enters the equations of motion. We thus require an additional condition to close the system which determines $\bar \phi$, and usually we pick $\bar \phi = 0$. This constraint is implemented in this PR.

Another way to close the system is to change the Poisson equation. Written discretely the proposal is to use

$$ \big ( \underbrace{A + \delta e e^T}_{\equiv \tilde A} \big ) \phi = R $$

where $\delta$ is a constant. I was a bit confused about this form until I understood that $e^T \phi = N \bar \phi$, so that in continuous from this would read

$$ \Bigg ( \underbrace{\nabla^2 + \frac{\delta N}{V} \int \mathrm{d} V}_{\tilde L} \Bigg ) \phi = \nabla^2 \phi + \delta N \bar \phi = R $$

where $\tilde L$ is the "regularized Laplacian operator". $\tilde L$ is regularized because taking the mean is well-defined:

$$ \overline{\tilde L \phi} = \delta N \bar \phi = \bar R $$

unlike the unregularized Poisson equation when $\bar R \ne 0$.

In terms of implementation, this is simply implemented by defining a new linear operator that adds the additional term, eg

Oceananigans.jl/src/Solvers/conjugate_gradient_poisson_solver.jl

Lines 47 to 51 in 4f6ffd7

	if !isnothing(L.δ)
	# Add regularization
	ϕ̄ = mean(ϕ)
	parent(Lϕ) .+= L.δ * ϕ̄
	end

Here are a few more loose ends:

1. I am wondering whether this perturbation or shifting of the Laplacian is independent of the preconditioner. For example, the FFT-based preconditioner also applies a gauge condition:

Oceananigans.jl/src/Solvers/fft_based_poisson_solver.jl

Lines 109 to 115 in efb8b71

	# Solve the discrete screened Poisson equation (∇² + m) ϕ = b.
	@. ϕc = - b / (λx + λy + λz - m)
	# If m === 0, the "zeroth mode" at `i, j, k = 1, 1, 1` is undetermined;
	# we set this to zero by default. Another slant on this "problem" is that
	# λx[1, 1, 1] + λy[1, 1, 1] + λz[1, 1, 1] = 0, which yields ϕ[1, 1, 1] = Inf or NaN.
	m === 0 && CUDA.@allowscalar ϕc[1, 1, 1] = 0

Therefore it is unclear to me whether additionally solving the regularized Poisson equation is necessary or valid when we already have a linear constraint embedded in the preconditioner. Also, note that the constrained implemented in this PR is simply a shifting of the constraint embedded in the FFT-based preconditioner. While the FFT preconditioner zeros out the mean pressure over the whole domain, the constraint in this PR takes into account the immersed boundary when zeroing the mean pressure.

I just don't understand why we will find the same solution because CG should return one solution among an infinity of solution (before the shift).

The "infinity of solutions" differ only in their mean value right? Or in other words, in the arbitrary value $e^T \phi$. I am maybe being too loose with language.

[glwagner]

I'm experimenting with this regularization on #3848

[Yixiao-Zhang]

In my implementation (Yixiao-Zhang@c7983b8002b91cd5939018a7c999eae77e2105ac\), the preconditioner is also perturbed. The goal is to solve the following equation in an immersed boundary grid.

$$\nabla^2 \phi + \delta N \bar \phi = R $$

The perturbed preconditioner solves this equation in a regular grid:

$$ \phi = L_\mathrm{FFT}R + \frac{1}{\delta N}\overline{R}$$

where $L_\mathrm{FFT}$ is the FFT-based Poisson solver. Due to the gauge condition of $L_\mathrm{FFT}$,

$$\nabla^2L_\mathrm{FFT}R = R - \overline{R}$$

$$\overline{L_\mathrm{FFT}R} = 0$$

for a random $R$ whose mean is not necessarily zero. The latter is simply the gauge condition; the former means that $L_\mathrm{FFT}$ ignores the mean component of $R$ (which can be understood as $\frac{1}{N}ee^TR$).

It is easy to verify that the preconditioner gives the exact solution in a regular grid.

Besides, in my implementation (Yixiao-Zhang@c7983b8002b91cd5939018a7c999eae77e2105ac\), I perturb the preconditioner directly by adding the mean of the input vector. However, a more efficient solution is to a new parameter to the FFT-based Poisson solver, so that it solves

$$\nabla^2 \phi + m \phi - \mu \overline{\phi} = r $$

where $\mu$ represents the role of $\delta$ or $c$ in previous comments.

[Yixiao-Zhang]

One way to motivate "defining a new matrix" is to view the conjugate gradient iteration method as a process of minimization of the quadratic function

$$f(x; A, b) = \frac{1}{2}x^\top A x - x^\top b$$

We can see

$$\mathrm{d}f = \mathrm{d}x^\top(Ax - b)$$

In our case, $A$ is positive semi-definite, and thus a gauge condition is required to ensure there is only one solution. The gauge condition is $\overline{x} = 0$.

The problem becomes minimization of $f(x)$ under a constraint of $\overline{x} = 0$. The constraint can be absorbed into the function to minimize:

$$\tilde{f}(x; A, b) = \frac{1}{2}x^\top A x + \frac{1}{2}cN\overline{x}^2 - x^\top b$$

where $c > 0$. Now, the problem becomes minimization of $\tilde{f}(x)$ for any $x\in \mathbb{R}^N$.

Similarly,

$$\mathrm{d}\tilde{f} = \mathrm{d}x^\top(Ax + \frac{c}{N}ee^\top x - b)$$

We can see $A + \frac{c}{N}ee^\top$ is actually $\tilde{A}$, the perturbed matrix defined in my previous comment.

Actually, we can verify

$$f(x; \tilde{A}, b) = \tilde{f}(x; A, b) $$

[glwagner]

However, a more efficient solution is to a new parameter to the FFT-based Poisson solver

Isn't this simply setting $\bar \phi = \bar r / \mu$ in the FFT solver, rather than $\bar \phi = 0$ as we do currently? Why does this help? (I'm not implying that it doesn't help, I just don't understand why trading one constraint for another is advantageous). Also given the ease of implementing such a change, did you test it?

Also if we set $\bar r = 0$ after computing the RHS here:

Oceananigans.jl/src/Models/NonhydrostaticModels/solve_for_pressure.jl

Lines 88 to 89 in efb8b71

	launch!(arch, grid, :xyz, _compute_source_term!, rhs, grid, Δt, Ũ)
	return solve!(pressure, solver.conjugate_gradient_solver, rhs)

then I suppose we would have $\bar \phi = \bar r / \mu = 0$ (I tested this briefly though and it didn't seem to improve the problems observed on #3831 / #3848, but we could keep trying).