CliMA / Oceananigans.jl

🌊 Julia software for fast, friendly, flexible, ocean-flavored fluid dynamics on CPUs and GPUs
https://clima.github.io/OceananigansDocumentation/stable
MIT License
990 stars 194 forks source link

Spectral solver for the nonhydrostatic pressure must produce divergence-free velocity field. #8

Closed ali-ramadhan closed 5 years ago

ali-ramadhan commented 5 years ago

I guess this is not something I was thinking of but John pointed out that it's crucial that the Fourier-spectral solver returns a nonhydrostatic pressure that when used to update the velocity field, produces a velocity field that is non-divergent at every grid point. Otherwise mass is being unphysically accumulated and tracer quantities will also be accumulated due to nonzero Q(∇·u) terms in the flux divergence operators ∇·(uQ) = Q(∇·u) + u·∇Q, leading to divergences and blowups.

Right now the wavenumbers are computed as

kx = 2π/Lx  # DFT
ky = 2π/Ly  # DFT
kz = 1π/Ly  # DCT

which should lead to a solver whose solutions converge spectrally. While it may solve for the pressure at the center of the cells very accurately, if ∇·u is non-zero this will be a big problem.

This will require some testing on my part to see which solver best satisfies ∇·u. If we can satisfy it to machine precision, that would be amazing. If not, hopefully it can satisfy it better than the conjugate-gradient method and then we can use the continuity equation to enforce ∇·u=0.

An alternative (not sure if this would work) is to discretize the derivative operators using a second-order centered-difference scheme (which I believe I've done for the 1D solver, and previous 3D solver) which explicitly places the discretization points on the center of the cells. Then the wavenumbers are

kˣ² = (4 / Δx²) * sin(πl / Nˣ)²  # DFT
kʸ² = (4 / Δy²) * sin(πm / Nʸ)²  # DFT
kᶻ² = (2 / Δz²) * (cos(πn / Nᶻ) - 1)  # DCT

and of course you expect second-order convergence. But if it better satisfies ∇·u=0 then it might be the way to go. You can also derive wavenumbers for fourth-order discretization.

EDIT: Fixed second-order wavenumbers.

ali-ramadhan commented 5 years ago

Comment from @christophernhill

There are some worked through notes from Gil Strang here

https://ocw.mit.edu/courses/mathematics/18-086-mathematical-methods-for-engineers-ii-spring-2006/readings/am35.pdf

will try and find my notes too.

I think your code looks right for this, but you should be able to check by comparing discrete del2 on phi. It should match RHS for any G that has RHS mean == 0

ali-ramadhan commented 5 years ago

@christophernhill Here is the 1D spectral solver from spectral_solvers.jl

# Solve a 1D Poisson equation ∇²ϕ = d²ϕ/dx² = f(x) with periodic boundary
# conditions and domain length L using the Fourier-spectral method. Solutions to
# Poisson's equation with periodic boundary conditions are unique up to a
# constant so you may need to appropriately normalize the solution if you care
# about the numerical value of the solution itself and not just derivatives of
# the solution.
function solve_poisson_1d_pbc(f, L, wavenumbers)
    N = length(f)  # Number of grid points (excluding the periodic end point).
    n = 0:N        # Wavenumber indices.

    if wavenumbers == :second_order
        # Wavenumbers if Laplacian is discretized using second-order
        # centered-difference scheme. Gives second-order convergence and ensures
        # that ∇²ϕ == f(x) to machine precision.
        Δx = L / N
        k² = @. (4 / Δx^2) * sin(π*n / N)^2
    elseif wavenumbers == :analytic
        # Wavenumbers if the derivatives are not discretized, should give
        # spectral convergence so that ϕ is accurate but ∇²ϕ ≈ f(x) as the ∇²
        # operator must be discretized.
        k² = @. ((2*π / L) * n)^2
    end

    # Forward transform the real-valued source term.
    fh = FFTW.rfft(f)

    # Calculate the Fourier coefficients of the source term.
    # We only need to compute the first (N-1)/2 + 1 Fourier coefficients
    # as ϕh(N-i) = ϕh(i) for a real-valued f.
    ϕh = - fh ./ k²[1:Int(N/2 + 1)]

    # Setting the DC/zero Fourier component to zero.
    ϕh[1] = 0

    # Take the inverse transform of the . We need to specify that the input f
    # had a length of N as an rrft of length N/2 may have come from an array
    # of length N (if N is even) or N-1 (if N is odd).
    ϕ = FFTW.irfft(ϕh, N)
end

where I use the "analytic wavenumbers" if I want spectral convergence for the solution

# Testing with the cos(2πz/H) source term should give a spectral solution that
# is numerically accurate within ≈ machine epsilon, so in these tests we can
# use ≈ to check the spectral solution with the analytic solution.
#
# With the source term f(z) = cos(2πz/H), the analytic solution is given by
# ϕ(z) = -(H/2π)² cos(2πz/H) and is periodic as ϕ(0) = ϕ(H).

function test_solve_poisson_1d_pbc_cosine_source(N)
    H = 12       # Length of domain.
    Nz = N       # Number of grid points.
    Δz = H / Nz  # Grid spacing.
    z = Δz * (0:(Nz-1))  # Grid point locations.

    f = cos.(2*π*z ./ H)  # Source term.
    ϕa = @. -(H / (2π))^2 * cos((2π / H) * z)  # Analytic solution.

    ϕs = solve_poisson_1d_pbc(f, H, :analytic)

    ϕs ≈ ϕa
end

and the "second-order wavenumbers" if I want the discretized Laplacian of the solution to equal the source term to numerical precision

function test_solve_poisson_1d_pbc_divergence_free(N)
    laplacian1d(f) = circshift(f, 1) - 2 .* f + circshift(f, -1)

    A = rand(N)
    A .= A .- mean(A)
    B = solve_poisson_1d_pbc(A, N, :second_order)
    A′ = laplacian1d(B)
    A ≈ A′
end

Both tests are from test_spectral_solvers.jl and pass. But I know that test_solve_poisson_1d_pbc_divergence_free fails when solve_poisson_1d_pbc is called with the wavenumbers=:analytic argument, which is why I think the choice of wavenumber matters.

glwagner commented 5 years ago

Two small notes:

"""
    f(x)

Returns `x^2`.
"""
f(x) = x^2

Doc strings are automatically detected and compiled to documentation by Julia's Documenter.jl package.

Next, I'm having a little trouble with the logic of these functions:

What is happening under the flag wavenumbers=:second_order? If derivatives are calculated with second-order finite differences, then we have no need for Fourier wavenumbers --- correct? What does this flag mean?

Why are the square wavenumbers equal to k² = @. (4 / Δx^2) * sin(π*n / N)^2? These are not equispaced and not equal to the Fourier wavenumbers.

ali-ramadhan commented 5 years ago

Thanks for the suggestions!

I'm not sure what to call them exactly so I'll try to describe.

"Analytic" case: You want to solve ∇²ϕ = f(x,y,z) so you plug in the inverse discrete Fourier transform (or inverse discrete cosine transform) into both sides and solving the resulting equation in Fourier space for the Fourier coefficients you get ϕh_lmn = - fh_lmn / (kx² + ky² + kz²) where ϕh is ϕ^hat, fh is f^hat (the Fourier coefficients), l,m,n are the wavenumber indices (as i,j,k are the spatial indices), and kx = 2π/Lx, ky = 2π/Ly, and kz = 1π/Ly (DCT).

The problem with this "regular" spectral method is that taking the Laplacian of the solution ∇²ϕ does not produce f(x,y,z) to machine precision if the Laplacian is discretized using a second-order centered difference scheme, which is what the ocean model essentially does to calculate the pressure gradients to step forward the velocities. If ∇²ϕ does not equal f(x,y,z) to machine precision then ∇·u will be non-zero and mass will accumulate, eventually leading the model to blow up.

"Second-order" case: We want ∇·u=0 so what we'll do now is discretize Poisson's equation using second-order centered differences so that ∇²ϕ = f(x,y,z) becomes [ϕ(i+1, j, k) - 2ϕ(i, j, k) + ϕ(i-1, j, k)] / Δx² + ... similarly for y and z ... = f(i, j, k). Then you plug in the IDFT and IDCT expressions into all the ϕ(i+1, j, k), ϕ(i, j, k), ϕ(i-1, j, k), etc. terms and solve the equation in Fourier space.

The solution will look again look like ϕh_lmn = - fh_lmn / (kx² + ky² + kz²) except now

kx² = (4 / Δx²) * sin(πl / Nx)²  # DFT
ky² = (4 / Δy²) * sin(πm / Ny)²  # DFT
kz² = (2 / Δz²) * (cos(πn / Nz) - 1)  # DCT

If you use these wavenumbers then you lose the spectral convergence and instead only get second-order convergence as expected because we've discretized the derivative operator (I've tested this but only in a Jupyter notebook right now). However, because we've discretized the Laplacian in the same way we do so for calculating pressure gradients, calculating the discretized Laplacian ∇²ϕ now produces f(x,y,z) to machine precision, which should ensure that ∇·u=0.

I have this working in the 1D periodic case (DFT) and the 1D case with Neumann boundary conditions (DCT) as well as the 2D periodic case (DFTs in x,y) but I'm still trying to figure out the 2D mixed case (e.g. DFT in x, DCT in y).

I don't know if what I've written makes much sense (or if I'm even describing it correctly) but I know it works as we are chiefly interested in keeping ∇·u=0 which requires that ∇²ϕ = f(x,y,z) if ∇² is discretized using centered differences.

glwagner commented 5 years ago

I don't understand what's being done for the "second-order" case, but perhaps this is my own shortcoming. Why does this procedure better ensure ∇·u=0? This constraint can be enforced in Fourier space and does not require second-order finite differences.

EDIT: I just read your explanation again and I understand the issue (in a superficial sense).

Would it work to output the pressure correction derivatives ϕx, ϕy, and ϕz from the spectral solver? The pressure correction is found in Fourier space, so finding the derivatives is cheap.

I'd have to write down the discretized system to figure out if what you're doing makes sense.

ali-ramadhan commented 5 years ago

Sorry for the slow reply. Might be easier to discuss in person, I don't think I did a good job of explaining what I'm doing. I'd be very interested in learning how to enforce this through some sort of pressure correction.

christophernhill commented 5 years ago

Here are two relevant articles

https://doi.org/10.1016/0021-9991(88)90102-7 https://doi.org/10.1016/j.amc.2015.03.011

christophernhill commented 5 years ago

@ali-ramadhan and @glwagner I tried the code below in some simple tests while waiting for my car to be serviced.

using Pkg
Pkg.add("LinearAlgebra")
using LinearAlgebra

function solveLinearSystem(A,f)
 # Solve Aϕ=f
 tol=1.e-12
 E=eigen(A);
 L=E.values;
 V=E.vectors;
 # Get amplitudes, F, of eigenvectors that give f
 F=V'*f
 # Get inverse eigenvalues (zeroing inverse for v. small ones)
 rL=map(x -> if (abs(x)>tol) 1.0/x;  else 0. ; end , L);
 # Get amplitudes, Φ, of eigenvectors that give ϕ
 Φ=F.*rL
 # Solve for ϕ given Φ
 ϕ=V*Φ
 println(A*ϕ,f,ϕ)
 return ϕ
end

Acyc=[-2.  1   1
       1  -2.  1
       1   1  -2.];
Aneu=[-1   1   0
       1. -2.  1
       0.  1.  -1];
s=size(Acyc);
nx=s[1];
g=rand(nx+1,1);
divg=g[1:end-1]-g[2:end]
mdivg=sum(divg)./size(divg)[1]
divg=divg.-mdivg
solveLinearSystem(Acyc,divg)
solveLinearSystem(Aneu,divg)

This is algorithm that underlies the FFT approach. The FFT just optimizes (and makes it more complicated) by utilizing the fact that the eigenvector/eigenvalue coefficients for the simple, constant spacing Poisson problem, are the cos and sin terms in an FFT. Code appears to work so I am going to try and hack together a "plugin" for solve_poisson_3d_mbc.

The code won't be super high performance (or work for really big problems) but (fingers crossed) it should give something clean (and short) to get started and help with debugging/optimizing on GPU. Then we can work on various FFT approaches on CPU and GPU (3-d FFT, 2-d + cyclic reduction), Greg's thought on saving for gradients in continuous form.

In principle the eigenvectors and eigenvalues above should correspond with amplitudes that come out of FFTW - except that there are a bunch of 1/2 factors, N versus m numbers, complex versus split cos/sin notation bits that need to be carefully understood etc....

Just going to learn a little about sparse matrices in Julia - I assume they must exist! Hopefully the car will take a little longer to be finished.

Chris

christophernhill commented 5 years ago

sad

eigen(A) not supported for sparse matrices. Use for example eigs(A) from the Arpack package instead. .....

ali-ramadhan commented 5 years ago

Thanks for posting this @christophernhill ! Gotta run around this afternoon but will try playing around with solveLinearSystem(A,f) later today.

I don't know much about sparse matrices in Julia either. CUDA has cuSPARSE for GPUs but I think you're stuck with lower-level BLAS-like functionality.

christophernhill commented 5 years ago

@ali-ramadhan I added some code here https://github.com/ali-ramadhan/OceanDispatch.jl/blob/fixing-nonhydrostatic-algorithm/src/eig_solvers.jl . This is for maybe helping test/debug - it uses too much memory and CPU for any real work. The eigenvectors need N^2 storage (where N is Nx.Ny.Nz)!

Code runs and gives machine precision error for a Laplace like operator with cyclic in X and Y and Neumann in Z ( see - solve_poisson_3d_mbc_eig). I pushed to repo and branch to see if it would work, and to my surprise it did, hope that's OK!

Going to look at FFT bits some now.

ali-ramadhan commented 5 years ago

Thanks @christophernhill this looks awesome!

@mg547 and I were going through your 1D solver and we now understand what's going on. Will try out the 3D solver tomorrow!

christophernhill commented 5 years ago

Great - lets catch up tomorrow. I got as for as puzzling over the absence of finite difference form scaling in solve_poisson_3d_mbc - but its getting late!

ali-ramadhan commented 5 years ago

@glwagner Mukund (@mg547) and I are meeting tomorrow around 1pm to discuss this spectral solver stuff. Are you free then?

EDIT: Woops thought Greg replied!

christophernhill commented 5 years ago

@ali-ramadhan I haven't quite got to looking at weights yet, but have you thought about doing something like this

    fz  = FFTW.r2r(f,FFTW.REDFT10,3)
    fxyz= FFTW.fft(fz,[1,2])

and this

    ixyz=FFTW.ifft(fxyz,[1,2])
    iz  =FFTW.r2r(real.(ixyz),FFTW.REDFT01,3)/(2*Nz)

Not sure if that is completely correct syntax, but something do the Neumann forward(backward) transform first(last) and periodic ones second(first)? The current formulation looks like it might loose some cyclic info when you drop complex parts and also may get confusing around even/odd Nx, Ny, Nz. Usually working in complex FFT space generally simplifies juggling factors and sizes around odd v even dimensions.

christophernhill commented 5 years ago

@ali-ramadhan I think the

    fz  = FFTW.r2r(f,FFTW.REDFT10,3)
    fxyz= FFTW.fft(fz,[1,2])

plus

ixyz=FFTW.ifft(fxyz,[1,2])
iz  =FFTW.r2r(real.(ixyz),FFTW.REDFT01,3)/(2*Nz)

looks to work. There is a rough gist here

https://gist.github.com/christophernhill/5d9c9aafe3fab413ded728da034cb01e

where everything lines up for various sizes. I'll tidy up and add some dx, dy, dz foo later today (ish....)

ali-ramadhan commented 5 years ago

@christophernhill Thanks for posting the gist! @mg547 and I were able to get it working with our Laplacian operator. Going to try and inject it into the pressure solver now, fingers crossed!

christophernhill commented 5 years ago

Exciting On Mon, Dec 3, 2018 at 6:03 PM Ali Ramadhan notifications@github.com wrote:

@christophernhill Thanks for posting this! @mg547 and I were able to get it working with our Laplacian operator. Going to try and inject it into the pressure solver now, fingers crossed!

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub, or mute the thread.

ali-ramadhan commented 5 years ago

Closing as this was solved by Chris' solver which is now implemented in the following solvers:

  1. solve_poisson_3d_ppn(f, Nx, Ny, Nz, Δx, Δy, Δz)
  2. solve_poisson_3d_ppn!(g::RegularCartesianGrid, f::CellField, ϕ::CellField)
  3. solve_poisson_3d_ppn_planned!(ssp::SpectralSolverParameters, g::RegularCartesianGrid, f::CellField, ϕ::CellField)