From 11899b516bf0fe49fa53f3daba050608dfa79e46 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Wed, 1 Nov 2023 15:35:20 +0100 Subject: [PATCH] WENO 5 advection scheme (#65) * interpolate velocities from P nodes * weno5 advection scheme * Weno5 convection miniapp * add comment referring to Hugo's paper * improve comment * remove dead function from script * format * add MuladdMacro and Adapt --- Project.toml | 2 + miniapps/convection/WENO5/Layered_rheology.jl | 160 ++++++++ .../convection/WENO5/WENO_convection2D.jl | 342 ++++++++++++++++++ src/Interpolations.jl | 58 ++- src/MetaJustRelax.jl | 3 + src/advection/weno5.jl | 287 +++++++++++++++ 6 files changed, 846 insertions(+), 6 deletions(-) create mode 100644 miniapps/convection/WENO5/Layered_rheology.jl create mode 100644 miniapps/convection/WENO5/WENO_convection2D.jl create mode 100644 src/advection/weno5.jl diff --git a/Project.toml b/Project.toml index de4addb0..8a49ed43 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "0.1.0" [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" @@ -12,6 +13,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" +MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" ParallelStencil = "94395366-693c-11ea-3b26-d9b7aac5d958" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" diff --git a/miniapps/convection/WENO5/Layered_rheology.jl b/miniapps/convection/WENO5/Layered_rheology.jl new file mode 100644 index 00000000..2c935162 --- /dev/null +++ b/miniapps/convection/WENO5/Layered_rheology.jl @@ -0,0 +1,160 @@ +# from "Fingerprinting secondary mantle plumes", Cloetingh et al. 2022 + +function init_rheologies(; is_plastic = true) + + # Dislocation and Diffusion creep + disl_upper_crust = DislocationCreep(A=5.07e-18, n=2.3, E=154e3, V=6e-6, r=0.0, R=8.3145) + disl_lower_crust = DislocationCreep(A=2.08e-23, n=3.2, E=238e3, V=6e-6, r=0.0, R=8.3145) + disl_lithospheric_mantle = DislocationCreep(A=2.51e-17, n=3.5, E=530e3, V=6e-6, r=0.0, R=8.3145) + disl_sublithospheric_mantle = DislocationCreep(A=2.51e-17, n=3.5, E=530e3, V=6e-6, r=0.0, R=8.3145) + diff_lithospheric_mantle = DislocationCreep(A=2.51e-17, n=1.0, E=530e3, V=6e-6, r=0.0, R=8.3145) + diff_sublithospheric_mantle = DislocationCreep(A=2.51e-17, n=1.0, E=530e3, V=6e-6, r=0.0, R=8.3145) + + # Elasticity + el_upper_crust = SetConstantElasticity(; G=25e9, ν=0.5) + el_lower_crust = SetConstantElasticity(; G=25e9, ν=0.5) + el_lithospheric_mantle = SetConstantElasticity(; G=67e9, ν=0.5) + el_sublithospheric_mantle = SetConstantElasticity(; G=67e9, ν=0.5) + β_upper_crust = inv(get_Kb(el_upper_crust)) + β_lower_crust = inv(get_Kb(el_lower_crust)) + β_lithospheric_mantle = inv(get_Kb(el_lithospheric_mantle)) + β_sublithospheric_mantle = inv(get_Kb(el_sublithospheric_mantle)) + + # Physical properties using GeoParams ---------------- + η_reg = 1e16 + cohesion = 3e6 + friction = asind(0.2) + pl_crust = if is_plastic + DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + friction = asind(0.3) + pl = if is_plastic + DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + pl_wz = if is_plastic + DruckerPrager_regularised(; C = 2e6, ϕ=2.0, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + + # crust + K_crust = TP_Conductivity(; + a = 0.64, + b = 807e0, + c = 0.77, + d = 0.00004*1e-6, + ) + + K_mantle = TP_Conductivity(; + a = 0.73, + b = 1293e0, + c = 0.77, + d = 0.00004*1e-6, + ) + + # Define rheolgy struct + rheology = ( + # Name = "UpperCrust", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=2.75e3, β=β_upper_crust, T0=0.0, α = 3.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=7.5e2), + Conductivity = K_crust, + CompositeRheology = CompositeRheology((disl_upper_crust, el_upper_crust, pl_crust)), + Elasticity = el_upper_crust, + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "LowerCrust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3e3, β=β_lower_crust, T0=0.0, α = 3.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=7.5e2), + Conductivity = K_crust, + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + CompositeRheology = CompositeRheology((disl_lower_crust, el_lower_crust, pl_crust)), + Elasticity = el_lower_crust, + ), + # Name = "LithosphericMantle", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3.3e3, β=β_lithospheric_mantle, T0=0.0, α = 3e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + Conductivity = K_mantle, + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + CompositeRheology = CompositeRheology((disl_lithospheric_mantle, diff_lithospheric_mantle, el_lithospheric_mantle, pl)), + Elasticity = el_lithospheric_mantle, + ), + # Name = "SubLithosphericMantle", + # SetMaterialParams(; + # Phase = 4, + # Density = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), + # HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + # Conductivity = K_mantle, + # RadioactiveHeat = ConstantRadioactiveHeat(0.0), + # CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), + # Elasticity = el_sublithospheric_mantle, + # ), + # Name = "Plume", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=3.3e3-50, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + Conductivity = K_mantle, + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), + Elasticity = el_sublithospheric_mantle, + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 5, + Density = ConstantDensity(; ρ=1e3), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + Conductivity = ConstantConductivity(; k=15.0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + ), + ) +end + +function init_phases!(phases, particles, Lx; d=650e3, r=50e3) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) + @inbounds for ip in JustRelax.cellaxes(phases) + # quick escape + JustRelax.@cell(index[ip, i, j]) == 0 && continue + + x = JustRelax.@cell px[ip, i, j] + depth = -(JustRelax.@cell py[ip, i, j]) + if 0e0 ≤ depth ≤ 21e3 + @cell phases[ip, i, j] = 1.0 + + elseif 35e3 ≥ depth > 21e3 + @cell phases[ip, i, j] = 2.0 + + elseif 90e3 ≥ depth > 35e3 + @cell phases[ip, i, j] = 3.0 + + elseif depth > 90e3 + @cell phases[ip, i, j] = 3.0 + + elseif depth < 0e0 + @cell phases[ip, i, j] = 5.0 + + end + + # plume - rectangular + if ((x - Lx * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2) + JustRelax.@cell phases[ip, i, j] = 4.0 + end + end + return nothing + end + + @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) +end diff --git a/miniapps/convection/WENO5/WENO_convection2D.jl b/miniapps/convection/WENO5/WENO_convection2D.jl new file mode 100644 index 00000000..a742e483 --- /dev/null +++ b/miniapps/convection/WENO5/WENO_convection2D.jl @@ -0,0 +1,342 @@ +using CUDA +CUDA.allowscalar(false) # for safety + +using JustRelax, JustRelax.DataIO, JustPIC +import JustRelax.@cell + +## NOTE: need to run one of the lines below if one wishes to switch from one backend to another +# set_backend("Threads_Float64_2D") +# set_backend("CUDA_Float64_2D") + +# setup ParallelStencil.jl environment +model = PS_Setup(:CUDA, Float64, 2) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions, CellArrays + +# Load file with all the rheology configurations +include("Layered_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +@inline init_particle_fields(particles) = @zeros(size(particles.coords[1])...) +@inline init_particle_fields(particles, nfields) = tuple([zeros(particles.coords[1]) for i in 1:nfields]...) +@inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N)) +@inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N)) + +function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, nx, ny) + ni = nx, ny + ncells = nx * ny + np = max_xcell * ncells + px, py = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(2)) + inject = @fill(false, nx, ny, eltype=Bool) + index = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool) + + @parallel_indices (i, j) function fill_coords_index(px, py, index) + # lower-left corner of the cell + x0, y0 = x[i], y[j] + # fill index array + for l in 1:nxcell + JustRelax.@cell px[l, i, j] = x0 + dx * rand(0.05:1e-5:0.95) + JustRelax.@cell py[l, i, j] = y0 + dy * rand(0.05:1e-5:0.95) + JustRelax.@cell index[l, i, j] = true + end + return nothing + end + + @parallel (1:nx, 1:ny) fill_coords_index(px, py, index) + + return Particles( + (px, py), index, inject, nxcell, max_xcell, min_xcell, np, (nx, ny) + ) +end + +# Velocity helper grids for the particle advection +function velocity_grids(xci, xvi, di) + dx, dy = di + yVx = LinRange(xci[2][1] - dy, xci[2][end] + dy, length(xci[2])+2) + xVy = LinRange(xci[1][1] - dx, xci[1][end] + dx, length(xci[1])+2) + grid_vx = xvi[1], yVx + grid_vy = xVy, xvi[2] + + return grid_vx, grid_vy +end + +import ParallelStencil.INDICES +const idx_j = INDICES[2] +macro all_j(A) + esc(:($A[$idx_j])) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) + return nothing +end + +# Initial thermal profile +@parallel_indices (i, j) function init_T!(T, z) + depth = -z[j] + + # (depth - 15e3) because we have 15km of sticky air + if depth < 0e0 + T[i + 1, j] = 273e0 + + elseif 0e0 ≤ (depth) < 35e3 + dTdZ = (923-273)/35e3 + offset = 273e0 + T[i + 1, j] = (depth) * dTdZ + offset + + elseif 110e3 > (depth) ≥ 35e3 + dTdZ = (1492-923)/75e3 + offset = 923 + T[i + 1, j] = (depth - 35e3) * dTdZ + offset + + elseif (depth) ≥ 110e3 + dTdZ = (1837 - 1492)/590e3 + offset = 1492e0 + T[i + 1, j] = (depth - 110e3) * dTdZ + offset + + end + + return nothing +end + +# Thermal rectangular perturbation +function rectangular_perturbation!(T, xc, yc, r, xvi) + + @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + depth = abs(y[j]) + dTdZ = (2047 - 2017) / 50e3 + offset = 2017 + T[i + 1, j] = (depth - 585e3) * dTdZ + offset + end + return nothing + end + ni = length.(xvi) + @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) + + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false) + + # Physical domain ------------------------------------ + ly = 700e3 # domain length in y + lx = ly * ar # domain length in x + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / ni # grid step in x- and -y + origin = 0.0, -ly # origin coordinates (15km f sticky air layer) + xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies(; is_plastic = true) + κ = (10 / (rheology[1].HeatCapacity[1].cp * rheology[1].Density[1].ρ0)) + dt = dt_diff = 0.5 * min(di...)^2 / κ / 2.01 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Weno model ----------------------------------------- + weno = WENO5(ni=(nx,ny).+1, method=Val{2}()) # ni.+1 for Temp + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 20, 40, 1 + particles = init_particles_cellarrays( + nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_particle_fields_cellarrays(particles, Val(3)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + xc_anomaly = lx/2 # origin of thermal anomaly + yc_anomaly = -610e3 # origin of thermal anomaly + r_anomaly = 25e3 # radius of perturbation + init_phases!(pPhases, particles, lx; d=abs(yc_anomaly), r=r_anomaly) + phase_ratios = PhaseRatio(ni, length(rheology)) + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.1 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + periodicity = (left = false, right = false, top = false, bot = false), + ) + # initialize thermal profile - Half space cooling + @parallel (@idx ni .+ 1) init_T!(thermal.T, xvi[2]) + thermal_bcs!(thermal.T, thermal_bc) + + rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) + @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + for _ in 1:1 + @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + end + # Rheology + η = @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e16, 1e24) + ) + η_vep = copy(η) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), + periodicity = (left = false, right = false, top = false, bot = false), + ) + + # IO ----- ------------------------------------------- + # if it does not exist, make folder where figures are stored + if save_vtk + vtk_dir = figdir*"\\vtk" + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + # Plot initial T and η profiles + let + Yv = [y for x in xvi[1], y in xvi[2]][:] + Y = [y for x in xci[1], y in xci[2]][:] + fig = Figure(resolution = (1200, 900)) + ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") + ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") + scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + scatter!(ax2, Array(log10.(η[:])), Y./1e3) + ylims!(ax1, minimum(xvi[2])./1e3, 0) + ylims!(ax2, minimum(xvi[2])./1e3, 0) + hideydecorations!(ax2) + save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # WENO arrays + T_WENO = @zeros(ni.+1) + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + # Time loop + t, it = 0.0, 0 + while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) + ) + @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + # ------------------------------ + + # Stokes solver ---------------- + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax=50e3, + nout=1e3, + viscosity_cutoff=(1e18, 1e24) + ) + @parallel (JustRelax.@idx ni) tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di, dt_diff) + # ------------------------------ + + # Weno advection + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + T_WENO .= thermal.T[2:end-1, :] + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) + thermal.T[2:end-1, :] .= T_WENO + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + # advect particles in memory + shuffle_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject = check_injection(particles) + inject && inject_particles_phase!(particles, pPhases, tuple(), tuple(), xvi) + # update phase ratios + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 10) == 0 + # Make Makie figure + fig = Figure(resolution = (900, 900), title = "t = $t") + ax1 = Axis(fig[1, 1]) + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(T_WENO) , colormap=:batlow) + Colorbar(fig[1,2], h1) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end +## END OF MAIN SCRIPT ---------------------------------------------------------------- + + +# (Path)/folder where output data and figures are stored +figdir = "Weno2D" +ar = 1 # aspect ratio +n = 256 +nx = n*ar - 2 +ny = n - 2 +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) +else + igg +end + +# run main script +main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny); \ No newline at end of file diff --git a/src/Interpolations.jl b/src/Interpolations.jl index 779cee06..f61affaa 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -78,18 +78,18 @@ end ## 2D -function velocity2vertex!(Vx_v, Vy_v, Vx, Vy; ghost_nodes=false) +function velocity2vertex!(Vx_v, Vy_v, Vx, Vy; ghost_nodes=true) ni = size(Vx_v) if !ghost_nodes @parallel (@idx ni) Vx2vertex_noghost!(Vx_v, Vx) @parallel (@idx ni) Vy2vertex_noghost!(Vy_v, Vy) else - @parallel (@idx ni) Vx2vertex_ghost!(Vx_v, Vx) - @parallel (@idx ni) Vy2vertex_ghost!(Vy_v, Vy) + @parallel (@idx ni) Vx2vertex_LinP!(Vx_v, Vx) + @parallel (@idx ni) Vy2vertex_LinP!(Vy_v, Vy) end end -function velocity2vertex(Vx, Vy, nv_x, nv_y; ghost_nodes=false) +function velocity2vertex(Vx, Vy, nv_x, nv_y; ghost_nodes=true) Vx_v = @allocate(nv_x, nv_y) Vy_v = @allocate(nv_x, nv_y) @@ -97,8 +97,8 @@ function velocity2vertex(Vx, Vy, nv_x, nv_y; ghost_nodes=false) @parallel (@idx ni) Vx2vertex_noghost!(Vx_v, Vx) @parallel (@idx ni) Vy2vertex_noghost!(Vy_v, Vy) else - @parallel (@idx ni) Vx2vertex_ghost!(Vx_v, Vx) - @parallel (@idx ni) Vy2vertex_ghost!(Vy_v, Vy) + @parallel (@idx ni) Vx2vertex_LinP!(Vx_v, Vx) + @parallel (@idx ni) Vy2vertex_LinP!(Vy_v, Vy) end end @@ -143,6 +143,52 @@ end return nothing end +@parallel_indices (i, j) function Vx2vertex_LinP!(V, Vx) + @inline av(A, B) = (A + B) * 0.5 + + nx, ny = size(Vx) + + iSW, jSW = clamp(i - 1, 1, nx), clamp(j, 1, ny) + iS, jS = clamp(i, 1, nx), clamp(j, 1, ny) + iSE, jSE = clamp(i + 1, 1, nx), clamp(j, 1, ny) + + iNE, jNE = clamp(i + 1, 1, nx), clamp(j + 1, 1, ny) + iN, jN = clamp(i, 1, nx), clamp(j + 1, 1, ny) + iNW, jNW = clamp(i - 1, 1, nx), clamp(j + 1, 1, ny) + + V_SW = av(Vx[iSW, jSW], Vx[iS, jS]) + V_SE = av(Vx[iS, jS], Vx[iSE, jSE]) + V_NW = av(Vx[iNW, jNW], Vx[iN, jN]) + V_NE = av(Vx[iN, jN], Vx[iNE, jNE]) + + V[i, j] = 0.25 * (V_SW + V_SE + V_NW + V_NE) + + return nothing +end + +@parallel_indices (i, j) function Vy2vertex_LinP!(V, Vy) + @inline av(A, B) = (A + B) * 0.5 + + nx, ny = size(Vy) + + iSW, jSW = clamp(i, 1, nx), clamp(j - 1, 1, ny) + iW, jW = clamp(i, 1, nx), clamp(j, 1, ny) + iSE, jSE = clamp(i, 1, nx), clamp(j + 1, 1, ny) + + iNE, jNE = clamp(i + 1, 1, nx), clamp(j - 1, 1, ny) + iE, jE = clamp(i + 1, 1, nx), clamp(j, 1, ny) + iNW, jNW = clamp(i + 1, 1, nx), clamp(j + 1, 1, ny) + + V_SW = av(Vy[iSW, jSW], Vy[iW, jW]) + V_SE = av(Vy[iW, jW], Vy[iSE, jSE]) + V_NW = av(Vy[iNW, jNW], Vy[iE, jE]) + V_NE = av(Vy[iE, jE], Vy[iNE, jNE]) + + V[i, j] = 0.25 * (V_SW + V_SE + V_NW + V_NE) + + return nothing +end + ## 3D """ diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl index b0d0472e..0c1d7f4c 100644 --- a/src/MetaJustRelax.jl +++ b/src/MetaJustRelax.jl @@ -130,6 +130,9 @@ function environment!(model::PS_Setup{T,N}) where {T,N} include(joinpath(@__DIR__, "Interpolations.jl")) export vertex2center!, center2vertex!, temperature2center! + + include(joinpath(@__DIR__, "advection/weno5.jl")) + export WENO5, WENO_advection! end # conditional submodule load diff --git a/src/advection/weno5.jl b/src/advection/weno5.jl new file mode 100644 index 00000000..0b4d0b9f --- /dev/null +++ b/src/advection/weno5.jl @@ -0,0 +1,287 @@ +using MuladdMacro, Adapt + +## Weno5 advection scheme. Implemantation based on the repository from +# https://gmd.copernicus.org/preprints/gmd-2023-189/ + +abstract type AbstractWENO end + +""" + WENO5{T, N, A, M} <: AbstractWENO + +The `WENO5` is a structure representing the Weighted Essentially Non-Oscillatory (WENO) scheme of order 5 for the solution of hyperbolic partial differential equations. + +# Fields +- `d0L`, `d1L`, `d2L`: Upwind constants. +- `d0R`, `d1R`, `d2R`: Downwind constants. +- `c1`, `c2`: Constants for betas. +- `sc1`, `sc2`, `sc3`, `sc4`, `sc5`: Stencil weights. +- `ϵ`: Tolerance. +- `ni`: Grid size. +- `ut`: Intermediate buffer array. +- `fL`, `fR`, `fB`, `fT`: Fluxes. +- `method`: Method (1:JS, 2:Z). + +# Description +The `WENO5` structure contains the parameters and temporary variables used in the WENO scheme. These include the upwind and downwind constants, the constants for betas, the stencil candidate weights, the tolerance, the grid size, the fluxes, and the method. +""" +@kwdef struct WENO5{T,N,A,M} <: AbstractWENO + # upwind constants + d0L::T = 1 / 10 + d1L::T = 3 / 5 + d2L::T = 3 / 10 + # downwind constants + d0R::T = 3 / 10 + d1R::T = 3 / 5 + d2R::T = 1 / 10 + # for betas + c1::T = 13 / 12 + c2::T = 1 / 4 + # stencil weights + sc1::T = 1 / 3 + sc2::T = 7 / 6 + sc3::T = 11 / 6 + sc4::T = 1 / 6 + sc5::T = 5 / 6 + # tolerance + ϵ::T = 1e-6 + # grid size + ni::NTuple{N,Int64} + # fluxes + ut::A = @zeros(ni...) + fL::A = @zeros(ni...) + fR::A = @zeros(ni...) + fB::A = @zeros(ni...) + fT::A = @zeros(ni...) + # method + method::M = Val{1} # 1:JS, 2:Z +end + +Adapt.@adapt_structure WENO5 + +# check if index is on the boundary, if yes take value on the opposite for periodic, if not, don't change the value +# @inline limit_periodic(a, n) = a > n ? n : (a < 1 ? 1 : a) +@inline function limit_periodic(a, n) + a > n && return n + a < 1 && return 1 + return a +end + +## Betas +@inline function weno_betas(u1, u2, u3, u4, u5, weno) + β0 = @muladd weno.c1 * (u1 - 2 * u2 + u3)^2 + weno.c2 * (u1 - 4 * u2 + 3 * u3)^2 + β1 = @muladd weno.c1 * (u2 - 2 * u3 + u4)^2 + weno.c2 * (u2 - u4)^2 + β2 = @muladd weno.c1 * (u3 - 2 * u4 + u5)^2 + weno.c2 * (3 * u3 - 4 * u4 + u5)^2 + return β0, β1, β2 +end + +## Upwind alphas +@inline function weno_alphas_upwind(::WENO5, ::Type{Any}, β0, β1, β2) + return error("Unknown method for the WENO Scheme") +end + +@inline function weno_alphas_upwind(weno::WENO5, ::Val{1}, β0, β1, β2) + α0L = weno.d0L * inv(β0 + weno.ϵ)^2 + α1L = weno.d1L * inv(β1 + weno.ϵ)^2 + α2L = weno.d2L * inv(β2 + weno.ϵ)^2 + return α0L, α1L, α2L +end + +@inline function weno_alphas_upwind(weno::WENO5, ::Val{2}, β0, β1, β2) + τ = abs(β0 - β2) + α0L = weno.d0L * (1 + (τ * inv(β0 + weno.ϵ))^2) + α1L = weno.d1L * (1 + (τ * inv(β1 + weno.ϵ))^2) + α2L = weno.d2L * (1 + (τ * inv(β2 + weno.ϵ))^2) + return α0L, α1L, α2L +end + +## Downwind alphas +@inline function weno_alphas_downwind(::WENO5, ::Any, β0, β1, β2) + return error("Unknown method for the WENO Scheme") +end + +@inline function weno_alphas_downwind(weno::WENO5, ::Val{1}, β0, β1, β2) + α0R = weno.d0R * inv(β0 + weno.ϵ)^2 + α1R = weno.d1R * inv(β1 + weno.ϵ)^2 + α2R = weno.d2R * inv(β2 + weno.ϵ)^2 + return α0R, α1R, α2R +end + +@inline function weno_alphas_downwind(weno::WENO5, ::Val{2}, β0, β1, β2) + τ = abs(β0 - β2) + α0R = weno.d0R * (1 + (τ * inv(β0 + weno.ϵ))^2) + α1R = weno.d1R * (1 + (τ * inv(β1 + weno.ϵ))^2) + α2R = weno.d2R * (1 + (τ * inv(β2 + weno.ϵ))^2) + return α0R, α1R, α2R +end + +## Stencil candidates +@inline function stencil_candidate_upwind(u1, u2, u3, u4, u5, weno) + s0 = @muladd weno.sc1 * u1 - weno.sc2 * u2 + weno.sc3 * u3 + s1 = @muladd -weno.sc4 * u2 + weno.sc5 * u3 + weno.sc1 * u4 + s2 = @muladd weno.sc1 * u3 + weno.sc5 * u4 - weno.sc4 * u5 + return s0, s1, s2 +end + +@inline function stencil_candidate_downwind(u1, u2, u3, u4, u5, weno) + s0 = @muladd -weno.sc4 * u1 + weno.sc5 * u2 + weno.sc1 * u3 + s1 = @muladd weno.sc1 * u2 + weno.sc5 * u3 - weno.sc4 * u4 + s2 = @muladd weno.sc3 * u3 - weno.sc2 * u4 + weno.sc1 * u5 + return s0, s1, s2 +end + +# UP/DOWN-WIND FLUXES +@inline function WENO_u_downwind(u1, u2, u3, u4, u5, weno) + return _WENO_u( + u1, u2, u3, u4, u5, weno, weno_alphas_downwind, stencil_candidate_downwind + ) +end + +@inline function WENO_u_upwind(u1, u2, u3, u4, u5, weno) + return _WENO_u(u1, u2, u3, u4, u5, weno, weno_alphas_upwind, stencil_candidate_upwind) +end + +@inline function _WENO_u( + u1, u2, u3, u4, u5, weno, fun_alphas::F1, fun_stencil::F2 +) where {F1,F2} + + # Smoothness indicators + β = weno_betas(u1, u2, u3, u4, u5, weno) + + # classical approach + α0, α1, α2 = fun_alphas(weno, weno.method, β...) + + _α = inv(α0 + α1 + α2) + w0 = α0 * _α + w1 = α1 * _α + w2 = α2 * _α + + # Candidate stencils + s0, s1, s2 = fun_stencil(u1, u2, u3, u4, u5, weno) + + # flux down/up-wind + f = @muladd w0 * s0 + w1 * s1 + w2 * s2 + + return f +end + +# FLUXES + +## x-axis +@inline function WENO_flux_downwind_x(u, nx, weno, i, j) + return _WENO_flux_x(u, nx, weno, i, j, WENO_u_downwind) +end +@inline function WENO_flux_upwind_x(u, nx, weno, i, j) + return _WENO_flux_x(u, nx, weno, i, j, WENO_u_upwind) +end + +@inline function _WENO_flux_x(u, nx, weno, i, j, fun::F) where {F} + jw, jww = clamp(j - 1, 1, nx), clamp(j - 2, 1, nx) + je, jee = clamp(j + 1, 1, nx), clamp(j + 2, 1, nx) + + u1 = @inbounds u[i, jww] + u2 = @inbounds u[i, jw] + u3 = @inbounds u[i, j] + u4 = @inbounds u[i, je] + u5 = @inbounds u[i, jee] + return f = fun(u1, u2, u3, u4, u5, weno) +end + +## y-axis +@inline function WENO_flux_downwind_y(u, ny, weno, i, j) + return _WENO_flux_y(u, ny, weno, i, j, WENO_u_downwind) +end +@inline function WENO_flux_upwind_y(u, ny, weno, i, j) + return _WENO_flux_y(u, ny, weno, i, j, WENO_u_upwind) +end + +@inline function _WENO_flux_y(u, ny, weno, i, j, fun::F) where {F} + iw, iww = clamp(i - 1, 1, ny), clamp(i - 2, 1, ny) + ie, iee = clamp(i + 1, 1, ny), clamp(i + 2, 1, ny) + + u1 = @inbounds u[iww, j] + u2 = @inbounds u[iw, j] + u3 = @inbounds u[i, j] + u4 = @inbounds u[ie, j] + u5 = @inbounds u[iee, j] + + return f = fun(u1, u2, u3, u4, u5, weno) +end + +@inline function weno_rhs(vy, vx, weno, _dx, _dy, nx, ny, i, j) + jW, jE = clamp(j - 1, 1, ny), clamp(j + 1, 1, ny) + iS, iN = clamp(i - 1, 1, nx), clamp(i + 1, 1, nx) + + vx_ij = vx[i, j] + vy_ij = vy[i, j] + + return r = @inbounds begin + @muladd begin + max(vx_ij, 0) * (weno.fL[i, j] - weno.fL[i, jW]) * _dx + + min(vx_ij, 0) * (weno.fR[i, jE] - weno.fR[i, j]) * _dx + + max(vy_ij, 0) * (weno.fB[i, j] - weno.fB[iS, j]) * _dy + + min(vy_ij, 0) * (weno.fT[iN, j] - weno.fT[i, j]) * _dy + end + end +end + +@parallel_indices (i, j) function weno_f!(u, weno, nx, ny) + weno.fL[i, j] = WENO_flux_upwind_x(u, nx, weno, i, j) + weno.fR[i, j] = WENO_flux_downwind_x(u, nx, weno, i, j) + weno.fB[i, j] = WENO_flux_upwind_y(u, ny, weno, i, j) + weno.fT[i, j] = WENO_flux_downwind_y(u, ny, weno, i, j) + return nothing +end + +## WENO-5 ADVECTION +""" + WENO_advection!(u, Vxi, weno, di, ni, dt) + +Perform the advection step of the Weighted Essentially Non-Oscillatory (WENO) scheme for the solution of hyperbolic partial differential equations. + +# Arguments +- `u`: field to be advected. +- `Vxi`: velocity field. +- `weno`: structure containing the WENO scheme parameters and temporary variables. +- `di`: grid spacing. +- `ni`: number of grid points. +- `dt`: time step. + +# Description +The function first calculates the fluxes using the WENO scheme. Then it performs three steps of the WENO scheme. Each step involves calculating the right-hand side of the WENO equation and updating the solution `u`. The updating of the solution `u` is done using different combinations of the original solution and the temporary solution `weno.ut`. +""" +function WENO_advection!(u, Vxi, weno, di, dt) + _di = inv.(di) + ni = nx, ny = size(u) + one_third = inv(3) + two_thirds = 2 * one_third + + @parallel (1:nx, 1:ny) weno_f!(u, weno, nx, ny) + @parallel (1:nx, 1:ny) weno_step1!(weno, u, Vxi, _di, ni, dt) + + @parallel (1:nx, 1:ny) weno_f!(weno.ut, weno, nx, ny) + @parallel (1:nx, 1:ny) weno_step2!(weno, u, Vxi, _di, ni, dt) + + @parallel (1:nx, 1:ny) weno_f!(weno.ut, weno, nx, ny) + @parallel (1:nx, 1:ny) weno_step3!(u, weno, Vxi, _di, ni, dt, one_third, two_thirds) +end + +@parallel_indices (i, j) function weno_step1!(weno, u, Vxi, _di, ni, dt) + rᵢ = weno_rhs(Vxi..., weno, _di..., ni..., i, j) + @inbounds weno.ut[i, j] = muladd(-dt, rᵢ, u[i, j]) + return nothing +end + +@parallel_indices (i, j) function weno_step2!(weno, u, Vxi, _di, ni, dt) + rᵢ = weno_rhs(Vxi..., weno, _di..., ni..., i, j) + @inbounds weno.ut[i, j] = @muladd 0.75 * u[i, j] + 0.25 * weno.ut[i, j] - 0.25 * dt * rᵢ + return nothing +end + +@parallel_indices (i, j) function weno_step3!( + u, weno, Vxi, _di, ni, dt, one_third, two_thirds +) + rᵢ = weno_rhs(Vxi..., weno, _di..., ni..., i, j) + @inbounds u[i, j] = @muladd one_third * u[i, j] + two_thirds * weno.ut[i, j] - + two_thirds * dt * rᵢ + return nothing +end