Merge pull request #9 from JuliaGeodynamics/Fix_3D

Add Zhang and Shun limiter

Author: Iddingsite

examples/2D/2D_linear_deformation.jl

@@ -0,0 +1,101 @@
+using FiniteDiffWENO5
+using GLMakie
+
+function main(; nx = 400, ny = 400)
+
+    Lx = 1.0
+    Δx = Lx / nx
+    Δy = Lx / ny
+
+    x = range(0, stop = Lx, length = nx)
+
+    # Courant number
+    CFL = 0.7
+    period = 1
+
+    # Grid x assumed defined
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
+    grid = (x .* ones(ny)', ones(nx) .* y')
+
+    # Staggered grids:
+    # vx at x-faces → (nx+1, ny)
+    x_vx = range(0, stop = Lx, length = nx + 1)
+    y_vx = range(Δy / 2, stop = Lx - Δy / 2, length = ny)
+
+    # vy at y-faces → (nx, ny+1)
+    x_vy = range(Δx / 2, stop = Lx - Δx / 2, length = nx)
+    y_vy = range(0, stop = Lx, length = ny + 1)
+
+    # Make 2D coordinate arrays with correct orientation
+    X_vx = repeat(x_vx, 1, ny)         # (nx+1, ny)
+    Y_vx = repeat(y_vx', nx + 1, 1)      # (nx+1, ny)
+
+    X_vy = repeat(x_vy, 1, ny + 1)       # (nx, ny+1)
+    Y_vy = repeat(y_vy', nx, 1)        # (nx, ny+1)
+
+    # Define velocity field (example: divergence-free vortex)
+    vx0 = .-2π .* sin.(π .* X_vx) .* cos.(π .* Y_vx)  # (nx+1, ny)
+    vy0 = 2π .* cos.(π .* X_vy) .* sin.(π .* Y_vy)   # (nx, ny+1)
+
+    v = (; x = vx0, y = vy0)
+
+    x0 = 1 / 4
+    c = 0.08
+
+    u0 = zeros(ny, nx)
+
+    for I in CartesianIndices((ny, nx))
+        u0[I] = sign(exp(-((grid[1][I] - x0)^2 + (grid[2][I]' - x0)^2) / c^2) - 0.5) * 0.5 + 0.5
+    end
+
+    u = copy(u0)
+    weno = WENOScheme(u; boundary = (2, 2, 2, 2), stag = true, multithreading = true)
+
+
+    # grid size
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
+
+    tmax = 1
+    t = 0
+    counter = 0
+
+    f = Figure(size = (800, 600))
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
+    u_obser = Observable(u0)
+    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
+    Colorbar(f[1, 2], label = "u", hm)
+    display(f)
+    counter_half = 0
+
+    while t < tmax
+        WENO_step!(u, v, weno, Δt, Δx, Δy)
+
+
+        if t > tmax / 2 && counter_half == 0
+            println("reversing velocity field...")
+            vx0 .= .- vx0
+            vy0 .= .- vy0
+            counter_half += 1
+        end
+
+        t += Δt
+
+        if t + Δt > tmax
+            Δt = tmax - t
+        end
+
+        if counter % 100 == 0
+            u_obser[] = u
+            ax.title = "t = $(round(t, digits = 2))"
+        end
+
+        counter += 1
+
+
+    end
+
+    return
+end
+
+main(nx = 400, ny = 400)

ext/ChmyExt.jl

@@ -1,5 +1,6 @@
 module ChmyExt
 using FiniteDiffWENO5
+using FiniteDiffWENO5: zhang_shu_limit, weno5_reconstruction_upwind, weno5_reconstruction_downwind
 using MuladdMacro
 using UnPack
 using Chmy
@@ -21,7 +22,7 @@ Create a WENO scheme structure for the given field `u` on the specified `grid` u
 - `boundary::NTuple{2N, Int}`: Tuple specifying the boundary conditions for each dimension (0: homogeneous Neumann, 1: homogeneous Dirichlet, 2: periodic). Default is periodic (2).
 - `stag::Bool`: Whether the grid is staggered (velocities on cell faces) or not (velocities on cell centers).
 """
-function WENOScheme(c0::AbstractField{T, N}, grid::StructuredGrid; boundary::NTuple = (2, 2), stag::Bool = true, kwargs...) where {T, N}
+function WENOScheme(c0::AbstractField{T, N}, grid::StructuredGrid; boundary::NTuple = (2, 2), stag::Bool = true, lim_ZS::Bool = false, kwargs...) where {T, N}
 
     # check that boundary conditions are correctly defined
     @assert length(boundary) == 2N "Boundary conditions must be a tuple of length $(2N) for $(N)D data."
@@ -43,7 +44,7 @@ function WENOScheme(c0::AbstractField{T, N}, grid::StructuredGrid; boundary::NTu
     TFlux = typeof(fl)
     TArray = typeof(du)
 
-    return WENOScheme{T, TArray, TFlux, N_boundary}(stag = stag, boundary = boundary, multithreading = multithreading, fl = fl, fr = fr, du = du, ut = ut)
+    return WENOScheme{T, TArray, TFlux, N_boundary}(stag = stag, boundary = boundary, multithreading = multithreading, lim_ZS = lim_ZS, fl = fl, fr = fr, du = du, ut = ut)
 end
 
 function WENOScheme(c0::AbstractField; kwargs...)
@@ -75,7 +76,8 @@ include("KAExt3D.jl")
                v::NamedTuple{(:x,), <:Tuple{<:AbstractField{<:Real, 1}}},
                weno::FiniteDiffWENO5.WENOScheme,
                Δt, Δx,
-               grid::StructuredGrid, arch) where {T_field <: AbstractField{<:Real, 1}}
+               grid::StructuredGrid, arch;
+               u_min = 0.0, u_max = 0.0) where {T_field <: AbstractField{<:Real, 1}}
 
 Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method with WENO5 spatial discretization using Chmy.jl fields in 1D.
 
@@ -86,8 +88,11 @@ Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method
 - `Δt`: Time step size.
 - `Δx`: Spatial grid size.
 - `grid::StructuredGrid`: Computational grid from Chmy.
+- `arch::Backend`: The KernelAbstractions backend in use (e.g., CPU(), CUDABackend(), etc.).
+- `u_min`: Minimum value of `u` for the Zhang-Shu positivity limiter.
+- `u_max`: Maximum value of `u` for the Zhang-Shu positivity limiter.
 """
-function WENO_step!(u::T_field, v::NamedTuple{(:x,), <:Tuple{<:AbstractField{<:Real, 1}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, grid::StructuredGrid, arch) where {T_field <: AbstractField{<:Real, 1}}
+function WENO_step!(u::T_field, v::NamedTuple{(:x,), <:Tuple{<:AbstractField{<:Real, 1}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, grid::StructuredGrid, arch; u_min = 0.0, u_max = 0.0) where {T_field <: AbstractField{<:Real, 1}}
 
     @assert get_backend(u) == get_backend(v.x)
 
@@ -98,19 +103,19 @@ function WENO_step!(u::T_field, v::NamedTuple{(:x,), <:Tuple{<:AbstractField{<:R
     nx = grid.axes[1].length
     Δx_ = inv(Δx)
 
-    @unpack ut, du, fl, fr, stag, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
 
-    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_1D! => (du, fl, fr, v, stag, Δx_, grid))
 
     interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_1D! => (du, fl, fr, v, stag, Δx_, grid))
 
     interior(ut) .= @muladd 0.75 .* interior(u) .+ 0.25 .* interior(ut) .- 0.25 .* Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_1D => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_1D! => (du, fl, fr, v, stag, Δx_, grid))
 
     interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
@@ -124,7 +129,8 @@ end
                v::NamedTuple{(:x, :y), <:Tuple{Vararg{AbstractField{<:Real}, 2}}},
                weno::FiniteDiffWENO5.WENOScheme,
                Δt, Δx,
-               grid::StructuredGrid, arch) where T_field <: AbstractField{<:Real} where names
+               grid::StructuredGrid, arch;
+               u_min = 0.0, u_max = 0.0) where T_field <: AbstractField{<:Real} where names
 
 Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method with WENO5 spatial discretization using Chmy.jl fields in 2D.
 
@@ -135,8 +141,11 @@ Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method
 - `Δt`: Time step size.
 - `Δx`: Spatial grid size.
 - `grid::StructuredGrid`: Computational grid from Chmy.
+- `arch::Backend`: The KernelAbstractions backend in use (e.g., CPU(), CUDABackend(), etc.).
+- `u_min`: Minimum value of `u` for the Zhang-Shu positivity limiter.
+- `u_max`: Maximum value of `u` for the Zhang-Shu positivity limiter.
 """
-function WENO_step!(u::T_field, v::NamedTuple{(:x, :y), <:Tuple{Vararg{AbstractField{<:Real}, 2}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, Δy, grid::StructuredGrid, arch) where {T_field <: AbstractField{<:Real, 2}}
+function WENO_step!(u::T_field, v::NamedTuple{(:x, :y), <:Tuple{Vararg{AbstractField{<:Real}, 2}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, Δy, grid::StructuredGrid, arch; u_min = 0.0, u_max = 0.0) where {T_field <: AbstractField{<:Real, 2}}
 
     @assert get_backend(u) == get_backend(v.x)
     @assert get_backend(u) == get_backend(v.y)
@@ -150,22 +159,22 @@ function WENO_step!(u::T_field, v::NamedTuple{(:x, :y), <:Tuple{Vararg{AbstractF
     Δx_ = inv(Δx)
     Δy_ = inv(Δy)
 
-    @unpack ut, du, fl, fr, stag, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
 
-    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, u, boundary, ny, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, u, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_2D! => (du, fl, fr, v, stag, Δx_, Δy_, grid))
 
     interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_2D! => (du, fl, fr, v, stag, Δx_, Δy_, grid))
 
     interior(ut) .= @muladd 0.75 .* interior(u) .+ 0.25 .* interior(ut) .- 0.25 .* Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_2D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_2D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_2D! => (du, fl, fr, v, stag, Δx_, Δy_, grid))
 
     interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
@@ -179,7 +188,8 @@ end
                v::NamedTuple{names, <:Tuple{Vararg{AbstractField{<:Real}, 2}}},
                weno::FiniteDiffWENO5.WENOScheme,
                Δt, Δx, Δy, Δz,
-               grid::StructuredGrid, arch) where T_field <: AbstractArray{<:Real, 3}
+               grid::StructuredGrid, arch;
+               u_min = 0.0, u_max = 0.0) where T_field <: AbstractArray{<:Real, 3}
 
 Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method with WENO5 spatial discretization using Chmy.jl fields in 3D.
 
@@ -192,8 +202,11 @@ Advance the solution `u` by one time step using the 3rd-order Runge-Kutta method
 - `Δy`: Spatial grid size.
 - `Δz`: Spatial grid size.
 - `grid::StructuredGrid`: Computational grid from Chmy.
+- `arch::Backend`: The KernelAbstractions backend in use (e.g., CPU(), CUDABackend(), etc.).
+- `u_min`: Minimum value of `u` for the Zhang-Shu positivity limiter.
+- `u_max`: Maximum value of `u` for the Zhang-Shu positivity limiter.
 """
-function WENO_step!(u::T_field, v::NamedTuple{(:x, :y, :z), <:Tuple{Vararg{AbstractField{<:Real}, 3}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, Δy, Δz, grid::StructuredGrid, arch) where {T_field <: AbstractArray{<:Real, 3}}
+function WENO_step!(u::T_field, v::NamedTuple{(:x, :y, :z), <:Tuple{Vararg{AbstractField{<:Real}, 3}}}, weno::FiniteDiffWENO5.WENOScheme, Δt, Δx, Δy, Δz, grid::StructuredGrid, arch; u_min = 0.0, u_max = 0.0) where {T_field <: AbstractArray{<:Real, 3}}
 
     @assert get_backend(u) == get_backend(v.x)
     @assert get_backend(u) == get_backend(v.y)
@@ -210,25 +223,25 @@ function WENO_step!(u::T_field, v::NamedTuple{(:x, :y, :z), <:Tuple{Vararg{Abstr
     Δy_ = inv(Δy)
     Δz_ = inv(Δz)
 
-    @unpack ut, du, fl, fr, stag, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
 
-    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, u, boundary, ny, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, u, boundary, nz, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, u, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, u, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, u, boundary, nz, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_3D! => (du, fl, fr, v, stag, Δx_, Δy_, Δz_, grid))
 
     interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, ut, boundary, nz, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, ut, boundary, nz, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_3D! => (du, fl, fr, v, stag, Δx_, Δy_, Δz_, grid))
 
     interior(ut) .= @muladd 0.75 .* interior(u) .+ 0.25 .* interior(ut) .- 0.25 .* Δt .* interior(du)
 
-    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, grid))
-    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, ut, boundary, nz, χ, γ, ζ, ϵ, grid))
+    launch(arch, grid, WENO_flux_KA_3D_x => (fl.x, fr.x, ut, boundary, nx, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_y => (fl.y, fr.y, ut, boundary, ny, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+    launch(arch, grid, WENO_flux_KA_3D_z => (fl.z, fr.z, ut, boundary, nz, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
     launch(arch, grid, WENO_semi_discretisation_weno5_KA_3D! => (du, fl, fr, v, stag, Δx_, Δy_, Δz_, grid))
 
     interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)