Merge pull request #10 from JuliaGeodynamics/upwind_debug

add simple upwind mode

Author: Iddingsite

examples/1D/1D_linear_KA.jl

@@ -4,8 +4,6 @@ using KernelAbstractions
 
 function main(backend = CPU(), nx = 400)
 
-    backend = CPU()
-    nx = 400
 
     x_min = -1.0
     x_max = 1.0

examples/2D/2D_linear_KA.jl

@@ -4,9 +4,6 @@ using GLMakie
 
 function main(; backend = CPU(), nx = 400, ny = 400, stag = true)
 
-    nx = 400
-    ny = 400
-    stag = true
 
     Lx = 1.0
     Δx = Lx / nx
@@ -43,7 +40,7 @@ function main(; backend = CPU(), nx = 400, ny = 400, stag = true)
     u = KernelAbstractions.zeros(backend, Float64, nx, ny)
     copyto!(u, u0)
 
-    weno = WENOScheme(u, backend; boundary = (2, 2, 2, 2), stag = stag, multithreading = true)
+    weno = WENOScheme(u, backend; boundary = (2, 2, 2, 2), stag = stag)
 
     if stag
         v = (;

examples/2D/2D_linear_deformation_limiter.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, lim_ZS = 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; u_min = 0.0, u_max = 1.0)
+
+
+        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)

examples/2D/2D_linear_rotation_KA.jl

@@ -33,7 +33,7 @@ function main(; backend = CPU(), nx = 400, ny = 400)
     u = KernelAbstractions.zeros(backend, Float64, nx, ny)
     copyto!(u, u0)
 
-    weno = WENOScheme(u, backend; boundary = (2, 2, 2, 2), stag = false, multithreading = true)
+    weno = WENOScheme(u, backend; boundary = (2, 2, 2, 2), stag = false)
 
 
     v = (;

examples/3D/3D_linear_chmy.jl

@@ -54,7 +54,7 @@ function main(; backend = CPU(), nx = 50, ny = 50, nz = 50)
 
     u = Field(backend, grid, Center())
     set!(u, u0)
-    weno = WENOScheme(u, grid; boundary = (2, 2, 2, 2, 2, 2), stag = false, multithreading = true)
+    weno = WENOScheme(u, grid; boundary = (2, 2, 2, 2, 2, 2), stag = false)
 
     Δt = CFL * min(Δx, Δy, Δz)^(5 / 3)
     tmax = period * Lx / max(maximum(abs.(vx0)), maximum(abs.(vy0)), maximum(abs.(vz0)))
@@ -82,6 +82,7 @@ function main(; backend = CPU(), nx = 50, ny = 50, nz = 50)
                 KernelAbstractions.synchronize(backend)
                 u_obser[] = (interior(u) |> Array)[:, :, div(nz, 2)]
                 ax.title = "t = $(round(t, digits = 2))"
+                display(f)
             end
         end
 

ext/ChmyExt.jl

@@ -22,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, lim_ZS::Bool = false, kwargs...) where {T, N}
+function WENOScheme(c0::AbstractField{T, N}, grid::StructuredGrid; boundary::NTuple = (2, 2), stag::Bool = true, lim_ZS::Bool = false, upwind_mode = false) 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."
@@ -44,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, lim_ZS = lim_ZS, 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, upwind_mode = upwind_mode)
 end
 
 function WENOScheme(c0::AbstractField; kwargs...)
@@ -103,22 +103,27 @@ 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, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ, upwind_mode = weno
 
-    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))
+    if !upwind_mode
 
-    interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
+        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))
 
-    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 interior(u) .- Δt .* interior(du)
 
-    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, χ, γ, ζ, ϵ, lim_ZS, u_min, u_max, grid))
+        launch(arch, grid, WENO_semi_discretisation_weno5_KA_1D! => (du, fl, fr, v, stag, Δx_, 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)
 
-    interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
+        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)
+    else
+        launch(arch, grid, upwind_update_KA_1D! => (u, v, nx, Δx_, Δt, stag, boundary, grid))
+    end
 
     return nothing
 end
@@ -159,25 +164,29 @@ 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, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ, upwind_mode = weno
 
-    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))
+    if !upwind_mode
+        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)
+        interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
 
-    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))
+        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)
+        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, χ, γ, ζ, ϵ, 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))
+        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)
+        interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
+    else
+        launch(arch, grid, upwind_update_KA_2D! => (u, v, nx, ny, Δx_, Δy_, Δt, stag, boundary, grid))
+    end
 
     return nothing
 end
@@ -223,28 +232,32 @@ 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, lim_ZS, boundary, χ, γ, ζ, ϵ = weno
+    @unpack ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ, upwind_mode = weno
 
-    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))
+    if !upwind_mode
+        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)
+        interior(ut) .= @muladd interior(u) .- Δt .* interior(du)
 
-    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))
+        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)
+        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, χ, γ, ζ, ϵ, 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))
+        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)
+        interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
+    else
+        launch(arch, grid, upwind_update_KA_3D! => (u, v, nx, ny, nz, Δx_, Δy_, Δz_, Δt, stag, boundary, grid))
+    end
 
     return nothing
 end