add multifield advection

Project.toml

@@ -1,23 +1,21 @@
 name = "FiniteDiffWENO5"
 uuid = "4b1b4241-afed-47c3-a843-c6e03821b375"
-version = "0.0.4"
+version = "0.0.5"
 authors = ["Hugo Dominguez  <hdomingu@uni-mainz.de>"]
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
-UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [weakdeps]
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Chmy = "33a72cf0-4690-46d7-b987-06506c2248b9"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 
 [extensions]
-KAExt = ["KernelAbstractions"]
 ChmyExt = ["KernelAbstractions", "Chmy"]
+KAExt = ["KernelAbstractions"]
 
 [compat]
 Chmy = "0.1"
 KernelAbstractions = "0.9"
 MuladdMacro = "0.2"
-UnPack = "1.0"
 julia = "1.10"

README.md

@@ -125,6 +125,47 @@ Which produces the following result:
 
 ![](/docs/src/assets/1D_linear_advection.png)
 
+## Multi-field advection
+
+If you have multiple scalar fields (e.g. different chemical components) that share the same velocity, you can advect them all in a single call by passing a tuple of arrays. The same `WENOScheme` buffers are reused for each field, so there is no extra memory overhead. Each field gets its own `u_min` / `u_max` bounds for the Zhang-Shu limiter.
+
+```julia
+using FiniteDiffWENO5
+
+nx, ny = 200, 200
+Lx = 1.0
+Δx, Δy = Lx / nx, Lx / ny
+
+# Three chemical components with different initial distributions
+c1 = rand(nx, ny)
+c2 = zeros(nx, ny)
+c3 = ones(nx, ny)
+
+# Shared velocity field
+v = (; x = ones(nx, ny), y = 0.5 .* ones(nx, ny))
+
+# Create the WENO scheme from any one of the fields (they must all have the same size and type)
+weno = WENOScheme(c1; boundary = (2, 2, 2, 2), stag = false)
+
+Δt = 0.7 * min(Δx, Δy)^(5 / 3)
+
+# Advect all three fields in one call with per-field limiter bounds
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy;
+    u_min = (0.0, 0.0, 0.0),
+    u_max = (1.0, 1.0, 1.0))
+```
+
+This also works with the KernelAbstractions and Chmy backends:
+
+```julia
+# KernelAbstractions
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy, backend;
+    u_min = (0.0, 0.0, 0.0), u_max = (1.0, 1.0, 1.0))
+
+# Chmy
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy, grid, arch;
+    u_min = (0.0, 0.0, 0.0), u_max = (1.0, 1.0, 1.0))
+```
 
 ## Funding & author
 

docs/src/GPU.md

@@ -0,0 +1,106 @@
+
+## Running on GPU
+
+To use this package on a GPU, you need to use [KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl) to define the arrays and run the computations.
+The KernelAbstractions.jl package introduces a macro-based programming model that simplifies vendor-specific GPU programming by abstracting away its complexities. This allows hardware-independent kernels to be written that can be compiled and executed on different device backends without altering the high-level code or compromising performance.
+
+
+This will load an extension from FiniteDiffWENO5GPU.jl:
+
+```julia
+using KernelAbstractions
+using FiniteDiffWENO5
+```
+
+That way, the two functions `WENOScheme()` and `WENO_step!()` will automatically dispatch to the array types if the argument `backend` is provided. See the KernelAbstractions.jl documentation for more details on how to set up the GPU backend.
+
+
+
+
+nx = 400
+
+x_min = -1.0
+x_max = 1.0
+Lx = x_max - x_min
+
+x = range(x_min, stop = x_max, length = nx)
+
+# Courant number
+CFL = 0.4
+period = 4
+
+# Parameters for Shu test
+z = -0.7
+δ = 0.005
+β = log(2) / (36 * δ^2)
+a = 0.5
+α = 10
+
+# Functions
+G(x, β, z) = exp.(-β .* (x .- z) .^ 2)
+F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a) .^ 2, 0.0))
+
+# Grid x assumed defined
+u0_vec = zeros(length(x))
+
+# Gaussian-like smooth bump at x in [-0.8, -0.6]
+idx = (x .>= -0.8) .& (x .<= -0.6)
+u0_vec[idx] .= (1 / 6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
+
+# Heaviside step at x in [-0.4, -0.2]
+idx = (x .>= -0.4) .& (x .<= -0.2)
+u0_vec[idx] .= 1.0
+
+# Piecewise linear ramp at x in [0, 0.2]
+# Triangular spike at x=0.1, base width 0.2
+idx = abs.(x .- 0.1) .<= 0.1
+u0_vec[idx] .= 1 .- 10 .* abs.(x[idx] .- 0.1)
+
+# Elliptic/smooth bell at x in [0.4, 0.6]
+idx = (x .>= 0.4) .& (x .<= 0.6)
+u0_vec[idx] .= (1 / 6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
+
+
+u = copy(u0_vec)
+
+# fix here boundary to periodic when equal to 2
+# staggering = true means that the advection velocity is defined on the sides of the cells and should be of size nx+1 compared to the scalar field u.
+# Multithreading to false means that the computations will be done on a single thread.
+# Set to true to enable multithreading if the Julia session was started with multiple threads (e.g. `julia -t 4`).
+weno = WENOScheme(u; boundary = (2, 2), stag = true, multithreading = false)
+
+# advection velocity with size nx+1 for staggered grid (here constant)
+a = (; x = ones(nx + 1))
+
+# grid size
+Δx = x[2] - x[1]
+Δt = CFL * Δx^(5 / 3)
+
+tmax = period * (Lx + Δx) / maximum(abs.(a.x))
+
+t = 0
+
+while t < tmax
+    # take as input the scalar field u, the advection velocity a as a NamedTuple,
+    # the WENO scheme struct weno, the time step Δt and the grid size Δx
+    WENO_step!(u, a, weno, Δt, Δx)
+
+    t += Δt
+
+    if t + Δt > tmax
+        Δt = tmax - t
+    end
+end
+
+f = Figure(size = (800, 600), dpi = 400)
+ax = Axis(f[1, 1], title = "1D linear advection after $period periods", xlabel = "x", ylabel = "u")
+lines!(ax, x, u0_vec, label = "Exact", linestyle = :dash, color = :red)
+scatter!(ax, x, u, label = "WENO5")
+xlims!(ax, x_min, x_max)
+axislegend(ax)
+display(f)
+```
+
+which outputs:
+
+![1D advection](./assets/1D_linear_advection.png)

docs/src/GettingStarted.md

@@ -96,4 +96,50 @@ display(f)
 
 which outputs:
 
-![1D advection](./assets/1D_linear_advection.png)
+![1D advection](./assets/1D_linear_advection.png)
+
+## Advecting multiple fields
+
+If you have multiple scalar fields (e.g. different chemical components) that share the same velocity, you can advect them all in a single call by passing a tuple of arrays. The same `WENOScheme` buffers are reused for each field, so there is no extra memory overhead.
+
+Each field gets its own `u_min` / `u_max` bounds for the Zhang-Shu limiter, passed as tuples matching the number of fields.
+
+```julia
+using FiniteDiffWENO5
+
+nx = 200
+ny = 200
+Lx = 1.0
+Δx = Lx / nx
+Δy = Lx / ny
+
+# Three chemical components with different initial distributions
+c1 = rand(nx, ny)    # component 1
+c2 = zeros(nx, ny)   # component 2
+c3 = ones(nx, ny)    # component 3
+
+# Shared velocity field
+v = (; x = ones(nx, ny), y = 0.5 .* ones(nx, ny))
+
+# Create the WENO scheme from any one of the fields (they must all have the same size/type)
+weno = WENOScheme(c1; boundary = (2, 2, 2, 2), stag = false)
+
+Δt = 0.7 * min(Δx, Δy)^(5 / 3)
+
+# Advect all three fields in one call
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy;
+    u_min = (0.0, 0.0, 0.0),
+    u_max = (1.0, 1.0, 1.0))
+```
+
+This also works with KernelAbstractions and Chmy backends — just pass the extra backend/grid arguments as usual:
+
+```julia
+# KernelAbstractions
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy, backend;
+    u_min = (0.0, 0.0, 0.0), u_max = (1.0, 1.0, 1.0))
+
+# Chmy
+WENO_step!((c1, c2, c3), v, weno, Δt, Δx, Δy, grid, arch;
+    u_min = (0.0, 0.0, 0.0), u_max = (1.0, 1.0, 1.0))
+```

examples/2D/2D_multi_field.jl

@@ -0,0 +1,92 @@
+using FiniteDiffWENO5
+using GLMakie
+
+function main(; nx = 400, ny = 400)
+
+    Lx = 1.0
+    Δx = Lx / nx
+    Δy = Lx / ny
+
+    # Courant number
+    CFL = 0.7
+    period = 1
+
+    # Grid
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
+    grid = (x .* ones(ny)', ones(nx) .* y')
+
+    # Shared velocity field
+    vx0 = ones(nx, ny)
+    vy0 = ones(nx, ny)
+    v = (; x = vy0, y = vx0)
+
+    # Three chemical components with different initial conditions
+    x0 = 1 / 4
+    c_width1 = 0.08
+    c_width2 = 0.06
+    c_width3 = 0.1
+
+    c1 = zeros(ny, nx)
+    c2 = zeros(ny, nx)
+    c3 = zeros(ny, nx)
+
+    for I in CartesianIndices((ny, nx))
+        c1[I] = sign(exp(-((grid[1][I] - x0)^2 + (grid[2][I]' - x0)^2) / c_width1^2) - 0.5) * 0.5 + 0.5
+        c2[I] = exp(-((grid[1][I] - 0.5)^2 + (grid[2][I]' - 0.5)^2) / c_width2^2)
+        c3[I] = exp(-((grid[1][I] - 0.75)^2 + (grid[2][I]' - 0.25)^2) / c_width3^2)
+    end
+
+    c1_0 = copy(c1)
+    c2_0 = copy(c2)
+    c3_0 = copy(c3)
+
+    # Create a single WENOScheme shared by all fields
+    weno = WENOScheme(c1; boundary = (2, 2, 2, 2), stag = false, multithreading = true)
+
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
+    tmax = period * Lx / max(maximum(abs.(vx0)), maximum(abs.(vy0)))
+
+    t = 0
+    counter = 0
+
+    # Plot setup
+    f = Figure(size = (1200, 400))
+    ax1 = Axis(f[1, 1], title = "Component 1")
+    ax2 = Axis(f[1, 2], title = "Component 2")
+    ax3 = Axis(f[1, 3], title = "Component 3")
+    c1_obs = Observable(c1_0)
+    c2_obs = Observable(c2_0)
+    c3_obs = Observable(c3_0)
+    heatmap!(ax1, x, y, c1_obs; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
+    heatmap!(ax2, x, y, c2_obs; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
+    heatmap!(ax3, x, y, c3_obs; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
+    display(f)
+
+    while t < tmax
+        # Advect all 3 components in a single call
+        WENO_step!(
+            (c1, c2, c3), v, weno, Δt, Δx, Δy;
+            u_min = (0.0, 0.0, 0.0),
+            u_max = (1.0, 1.0, 1.0)
+        )
+
+        t += Δt
+
+        if t + Δt > tmax
+            Δt = tmax - t
+        end
+
+        if counter % 100 == 0
+            c1_obs[] = c1
+            c2_obs[] = c2
+            c3_obs[] = c3
+        end
+
+        counter += 1
+    end
+
+    return nothing
+end
+
+main(nx = 400, ny = 400)