From 24983c49dfe4463cea90de0b0246f78c538ae44a Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 20 Oct 2024 20:47:57 +0200 Subject: [PATCH 01/77] initial prototyping --- MiniKernels.jl | 159 ++++++++ PlumeFreeSurface_2D.jl | 285 +++++++++++++ mask.jl | 380 ++++++++++++++++++ .../stokes2D/shear_band/ShearBand2D.jl | 7 +- src/stokes/StressKernels.jl | 2 +- src/stokes/VelocityKernels.jl | 69 ++-- 6 files changed, 861 insertions(+), 41 deletions(-) create mode 100644 MiniKernels.jl create mode 100644 PlumeFreeSurface_2D.jl create mode 100644 mask.jl diff --git a/MiniKernels.jl b/MiniKernels.jl new file mode 100644 index 000000000..cbd0b2a32 --- /dev/null +++ b/MiniKernels.jl @@ -0,0 +1,159 @@ +# unmasked versions +Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds...] +Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds.+1...] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where I<:Integer = A[i-1, j] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where I<:Integer = A[i+1, j] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j-1] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j+1] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i-1, j, k] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i+1, j, k] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j-1, k] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j+1, k] +Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k-1] +Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k+1] + +# masked versions +Base.@propagate_inbounds @inline center(A::T,ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = A[inds...] * ϕ[inds...] +Base.@propagate_inbounds @inline next(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = A[inds.+1...]* ϕ[inds.+1...] +Base.@propagate_inbounds @inline left(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i-1, j] * ϕ[i-1, j] +Base.@propagate_inbounds @inline right(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i+1, j] * ϕ[i+1, j] +Base.@propagate_inbounds @inline back(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i, j-1] * ϕ[i, j-1] +Base.@propagate_inbounds @inline front(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i, j+1] * ϕ[i, j+1] +Base.@propagate_inbounds @inline left(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i-1, j, k] * ϕ[i-1, j, k] +Base.@propagate_inbounds @inline right(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i+1, j, k] * ϕ[i+1, j, k] +Base.@propagate_inbounds @inline back(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j-1, k] * ϕ[i, j-1, k] +Base.@propagate_inbounds @inline front(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j+1, k] * ϕ[i, j+1, k] +Base.@propagate_inbounds @inline bot(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j, k-1] * ϕ[i, j, k-1] +Base.@propagate_inbounds @inline top(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j, k+1] * ϕ[i, j, k+1] + +## 2D mini kernels +const T2 = AbstractArray{T,2} where {T} + +# finite differences +Base.@propagate_inbounds @inline _d_xa(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + right(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + front(A, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::AbstractArray, _dz, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + top(A, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-front(A, I...) + next(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-right(A, I...) + next(A, I...)) * _dy +Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy + +# averages +Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) +Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) +Base.@propagate_inbounds @inline _av_xa(A::T, i, j) where {T<:T2} = (center(A, i, j) + right(A, i, j)) * 0.5 +Base.@propagate_inbounds @inline _av_ya(A::T, i, j) where {T<:T2} = (center(A, i, j) + front(A, i, j)) * 0.5 +Base.@propagate_inbounds @inline _av_xi(A::T, i, j) where {T<:T2} = (front(A, i, j), next(A, i, j)) * 0.5 +Base.@propagate_inbounds @inline _av_yi(A::T, i, j) where {T<:T2} = (right(A, i, j), next(A, i, j)) * 0.5 +# harmonic averages +Base.@propagate_inbounds @inline function _harm(A::T, i, j) where {T<:T2} + return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2)) +end +Base.@propagate_inbounds @inline function _harm_a(A::T, i, j) where {T<:T2} + return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1)) +end +Base.@propagate_inbounds @inline function _harm_xa(A::T, i, j) where {T<:T2} + return eltype(A)(2) * (inv(right(A, i, j)) + inv(center(A, i, j))) +end +Base.@propagate_inbounds @inline function _harm_ya(A::T, i, j) where {T<:T2} + return eltype(A)(2) * (inv(front(A, i, j)) + inv(center(A, i, j))) +end +#others +Base.@propagate_inbounds @inline function _gather(A::T, i, j) where {T<:T2} + return center(A, i, j), right(A, i, j), front(A, i, j), next(A, i, j) +end + +## 3D mini kernels +const T3 = AbstractArray{T,3} where {T} + +# finite differences +Base.@propagate_inbounds @inline function _d_xi(A::T, i, j, k, _dx) where {T<:T3} + return (-A[i, j + 1, k + 1] + next(A, i, j, k)) * _dx +end +Base.@propagate_inbounds @inline function _d_yi(A::T, i, j, k, _dy) where {T<:T3} + return (-A[i + 1, j, k + 1] + next(A, i, j, k)) * _dy +end +Base.@propagate_inbounds @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3} + return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz +end +# averages +Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) +Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + right(A, i, j, k)) +Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + front(A, i, j, k)) +Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + top(A, i, j, k)) +Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) +Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) +Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) +Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) +Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) +Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) +# harmonic averages +@inline function _harm_x(A::T, i, j, k) where {T<:T3} + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(right(A, i, j, k))) +end +@inline function _harm_y(A::T, i, j, k) where {T<:T3} + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(front(A, i, j, k))) +end +@inline function _harm_z(A::T, i, j, k) where {T<:T3} + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(top(A, i, j, k))) +end +@inline function _harm_xy(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k)) +end +@inline function _harm_xz(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:j, k:(k + 1))) +end +@inline function _harm_yz(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, i:i, j:(j + 1), k:(k + 1))) +end +@inline function _harm_xyi(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, (i - 1):i, (j - 1):j, k:k)) +end +@inline function _harm_xzi(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, (i - 1):i, j:j, (k - 1):k)) +end +@inline function _harm_yzi(A::T, i, j, k) where {T<:T3} + return eltype(A)(4) * inv(mysum(A, i:i, (j - 1):j, (k - 1):k)) +end + +# others +Base.@propagate_inbounds @inline function _gather_yz(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), front(A, i, j, k), top(A, i, j, k), A[i, j + 1, k + 1] +end +Base.@propagate_inbounds @inline function _gather_xz(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), right(A, i, j, k), top(A, i, j, k), A[i + 1, j, k + 1] +end +Base.@propagate_inbounds @inline function _gather_xy(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), right(A, i, j, k), front(A, i, j, k), A[i + 1, j + 1, k] +end +@inline _current(A::T, i, j, k) where {T<:T3} = center(A, i, j, k) + +## Because mysum(::generator) does not work inside CUDA kernels... +@inline mysum(A, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ranges...) + +@inline function mysum(f::F, A, ranges_i) where {F<:Function} + s = 0.0 + for i in ranges_i + s += f(A[i]) + end + return s +end + +@inline function mysum(f::F, A, ranges_i, ranges_j) where {F<:Function} + s = 0.0 + for i in ranges_i, j in ranges_j + s += f(center(A, i, j)) + end + return s +end + +@inline function mysum(f::F, A, ranges_i, ranges_j, ranges_k) where {F<:Function} + s = 0.0 + for i in ranges_i, j in ranges_j, k in ranges_k + s += f(center(A, i, j, k)) + end + return s +end diff --git a/PlumeFreeSurface_2D.jl b/PlumeFreeSurface_2D.jl new file mode 100644 index 000000000..bed3b9d19 --- /dev/null +++ b/PlumeFreeSurface_2D.jl @@ -0,0 +1,285 @@ +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) + +# Load script dependencies +using LinearAlgebra, GeoParams, GLMakie + +include("mask.jl") +include("MiniKernels.jl") + +# Velocity helper grids for the particle advection +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + @parallel f_x(A, B) +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 + +function init_phases!(phases, particles) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index) + r=100e3 + f(x, A, λ) = A * sin(π*x/λ) + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) + @index phases[ip, i, j] = 2.0 + + if 0e0 ≤ depth ≤ 100e3 + @index phases[ip, i, j] = 1.0 + + else + @index phases[ip, i, j] = 2.0 + + if ((x - 250e3)^2 + (depth - 250e3)^2 ≤ r^2) + @index phases[ip, i, j] = 3.0 + end + end + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +# (Path)/folder where output data and figures are stored +n = 100 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +# function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 100e3 # thickness of sticky air layer + ly = 400e3 + thick_air # domain length in y + lx = 500e3 # 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) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + # Name = "Air", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=1e1), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Mantle", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Plume", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ) + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 30, 40, 15 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + init_phases!(pPhases, particles) + phase_ratios = PhaseRatios(backend, length(rheology), ni) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # RockRatios + air_phase = 1 + ϕ = RockRatio(ni...) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + free_surface = true + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "FreeSurfacePlume" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + while it < 15 + + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 50e3, + nout = 1e3, + viscosity_cutoff = (-Inf, Inf), + free_surface = true + ) + ) + dt = compute_dt(stokes, di) / 2 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 1) == 0 + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) + arrows!( + ax, + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + color = :red, + ) + fig + save(joinpath(figdir, "$(it).png"), fig) + + end + end + return nothing +# end +## END OF MAIN SCRIPT ---------------------------------------------------------------- + +# # (Path)/folder where output data and figures are stored +# n = 100 +# nx = n +# ny = n +# igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid +# IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +# else +# igg +# end + +# main(igg, nx, ny) + +_di = inv.(di) +@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., ϕ, _di...) + + +@parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + copy(stokes.P), + stokes.P, + stokes.viscosity.η, + @zeros(ni...), + @zeros(ni.+1...), + stokes.τ.II, + stokes.viscosity.η_vep, + 0.2, + dt, + pt_stokes.θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ +) + +@parallel compute_V!( + @velocity(stokes)..., + @zeros(size(stokes.V.Vy)), + stokes.P, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + stokes.viscosity.η, + ϕ.Vx, + ϕ.Vy, + _di..., + 0, +) \ No newline at end of file diff --git a/mask.jl b/mask.jl new file mode 100644 index 000000000..0062033e7 --- /dev/null +++ b/mask.jl @@ -0,0 +1,380 @@ +# using JustRelax, JustRelax.JustRelax2D +using CellArrays +# using Test +# using ParallelStencil +# @init_parallel_stencil(Threads, Float64, 2) +abstract type AbstractMask end + +struct RockRatio{T, N} <: AbstractMask + center::T + vertex::T + Vx::T + Vy::T + Vz::Union{Nothing, T} + yz::Union{Nothing, T} + xz::Union{Nothing, T} + xy::Union{Nothing, T} + + function RockRatio(center::AbstractArray{F, N}, vertex::T, Vx::T, Vy::T, Vz::Union{Nothing, T}, yz::Union{Nothing, T}, xz::Union{Nothing, T}, xy::Union{Nothing, T}) where {F, N, T} + new{T, N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) + end +end + +RockRatio(ni::NTuple{N, Integer}) where N = RockRatio(ni...) + +function RockRatio(nx, ny) + ni = nx, ny + center = @zeros(ni...) + vertex = @zeros(ni.+1...) + Vx = @zeros(nx+1, ny) # no ghost nodes! + Vy = @zeros(nx, ny+1) # no ghost nodes! + + return RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) +end + +function RockRatio(nx, ny, nz) + ni = nx, ny, nz + center = @zeros(ni...) + vertex = @zeros(ni.+1...) + Vx = @zeros(nx+1, ny, nz) # no ghost nodes! + Vy = @zeros(nx, ny+1, nz) # no ghost nodes! + Vz = @zeros(nx, ny, nz+1) # no ghost nodes! + yz = @zeros(nx, ny + 1, nz + 1) + xz = @zeros(nx + 1, ny, nz + 1) + xy = @zeros(nx + 1, ny + 1, nz) + + return RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) +end + +@inline size_c(x::AbstractMask) = size(x.center) +@inline size_v(x::AbstractMask) = size(x.vertex) +@inline size_vx(x::AbstractMask) = size(x.Vx) +@inline size_vy(x::AbstractMask) = size(x.Vy) +@inline size_vz(x::AbstractMask) = size(x.Vz) +@inline size_yz(x::AbstractMask) = size(x.yz) +@inline size_xz(x::AbstractMask) = size(x.xz) +@inline size_xy(x::AbstractMask) = size(x.xy) + +@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = 1 - @index phase_ratio[air_phase, I...] +@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = @index phase_ratio[air_phase, I...] + +@parallel_indices (I...) function update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase) + if all(I .≤ size(ratio_center)) + ϕ.center[I...] = compute_rock_ratio(ratio_center, air_phase, I...) + end + ϕ.vertex[I...] = compute_rock_ratio(ratio_vertex, air_phase, I...) + return nothing +end + +@parallel_indices (I...) function update_rock_ratio_vel!(ϕ::RockRatio{T, N}) where {T, N} + # 2D + @inline av_x(A::AbstractArray{T, 2}) where T = _av_xa(A, I...) + @inline av_y(A::AbstractArray{T, 2}) where T = _av_ya(A, I...) + # 3D + @inline av_x(A::AbstractArray{T, 3}) where T = _av_yz(A, I...) + @inline av_y(A::AbstractArray{T, 3}) where T = _av_xz(A, I...) + @inline av_z(A::AbstractArray{T, 3}) where T = _av_xy(A, I...) + + all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + end + return nothing +end + +function update_rock_ratio!(ϕ::RockRatio, phase_ratios, air_phase) + nvi = size_v(ϕ) + + @parallel (@idx nvi) update_rock_ratio_cv!(ϕ, phase_ratios.center, phase_ratios.vertex, air_phase) + @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) + + return nothing +end + +# true means that is a nullspace and needs to be removed +# @inline Base.@propagate_inbounds isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} = iszero(A[I...]) + +# nullspace: check whether any point of the stencil is 0, if so, eliminate equation +function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} + v = true + Base.@nexprs 2 j -> begin + Base.@nexprs 2 i -> begin + @inline + ii = clamp(I[1] + 2 * i - 3, 1, size(A, 1)) + jj = clamp(I[2] + 2 * j - 3, 1, size(A, 2)) + v *= if N == 3 + A[ii, jj, I[3]] > 0 + else + A[ii, jj] > 0 + end + end + end + if N === 3 + Base.@nexprs 2 k -> begin + kk = clamp(I[3] + 2 * k - 3, 1, size(A, 3)) + v *= A[ii, jj, kk] > 0 + end + end + return v * (A[I...] > 0) +end + +@parallel_indices (I...) function compute_∇V!( + ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::RockRatio, _di::NTuple{N} +) where {T, N} + @inline d_xi(A) = _d_xi(A, _di[1], I...) + @inline d_yi(A) = _d_yi(A, _di[2], I...) + @inline d_zi(A) = _d_zi(A, _di[3], I...) + + f = d_xi, d_yi, d_zi + + if isvalid(ϕ.center, I...) + @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) + else + @inbounds ∇V[I...] = zero(T) + end + return nothing +end + +@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + +@parallel_indices (I...) function compute_P!( + P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, ϕ::RockRatio, dt, r, θ_dτ +) where {N} + if isvalid(ϕ.center, I...) + K = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, @cell(phase_ratio[I...])) + RP[I...], P[I...] = JustRelax2D._compute_P!(P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ) + else + RP[I...] = P[I...] = zero(eltype(P)) + end + return nothing +end + +@parallel (@idx ni) compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + stokes.viscosity.η, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ, +) + +# 2D kernel +@parallel_indices (I...) function update_stresses_center_vertex_ps!( + ε::NTuple{3, T}, # normal components @ centers; shear components @ vertices + ε_pl::NTuple{3}, # whole Voigt tensor @ centers + EII, # accumulated plastic strain rate @ centers + τ::NTuple{3}, # whole Voigt tensor @ centers + τshear_v::NTuple{1}, # shear tensor components @ vertices + τ_o::NTuple{3}, + τshear_ov::NTuple{1}, # shear tensor components @ vertices + Pr, + Pr_c, + η, + λ, + λv, + τII, + η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_center, + phase_vertex, + ϕ +) where T + τxyv = τshear_v[1] + τxyv_old = τshear_ov[1] + ni = size(Pr) + Ic = JustRelax2D.clamped_indices(ni, I...) + + if isvalid(ϕ.vertex, I...) + # interpolate to ith vertex + Pv_ij = JustRelax2D.av_clamped(Pr, Ic...) + εxxv_ij = JustRelax2D.av_clamped(ε[1], Ic...) + εyyv_ij = JustRelax2D.av_clamped(ε[2], Ic...) + τxxv_ij = JustRelax2D.av_clamped(τ[1], Ic...) + τyyv_ij = JustRelax2D.av_clamped(τ[2], Ic...) + τxxv_old_ij = JustRelax2D.av_clamped(τ_o[1], Ic...) + τyyv_old_ij = JustRelax2D.av_clamped(τ_o[2], Ic...) + EIIv_ij = JustRelax2D.av_clamped(EII, Ic...) + + ## vertex + phase = @inbounds phase_vertex[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = JustRelax2D.plastic_params_phase(rheology, EIIv_ij, phase) + _Gvdt = inv(JustRelax2D.fn_ratio(JustRelax2D.get_shear_modulus, rheology, phase) * dt) + Kv = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + ηv_ij = JustRelax2D.av_clamped(η, Ic...) + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = + (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * dτ_rv + dτyyv = + (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * dτ_rv + dτxyv = + ( + -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + + 2.0 * ηv_ij * ε[3][I...] + ) * dτ_rv + τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) + + # yield function @ center + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) + # stress correction @ vertex + λv[I...] = + (1.0 - relλ) * λv[I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + dQdτxy = 0.5 * (τxyv[I...] + dτxyv) / τIIv_ij + τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[I...] * dQdτxy * dτ_rv + else + # stress correction @ vertex + τxyv[I...] += dτxyv + end + else + τxyv[I...] = zero(eltype(T)) + end + + ## center + if all(I .≤ ni) + if isvalid(ϕ.center, I...) + # Material properties + phase = @inbounds phase_center[I...] + _Gdt = inv(JustRelax2D.fn_ratio(JustRelax2D.get_shear_modulus, rheology, phase) * dt) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = JustRelax2D.plastic_params_phase(rheology, EII[I...], phase) + K = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, phase) + volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ + ηij = η[I...] + dτ_r = 1.0 / (θ_dτ + ηij * _Gdt + 1.0) + + # cache strain rates for center calculations + τij, τij_o, εij = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) + + # visco-elastic strain rates @ center + εij_ve = @. εij + 0.5 * τij_o * _Gdt + εII_ve = GeoParams.second_invariant(εij_ve) + # stress increments @ center + dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r + τII_ij = GeoParams.second_invariant(dτij .+ τij) + # yield function @ center + F = τII_ij - C - Pr[I...] * sinϕ + + if is_pl && !iszero(τII_ij) + # stress correction @ center + λ[I...] = + (1.0 - relλ) * λ[I...] + + relλ .* (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) + dQdτij = @. 0.5 * (τij + dτij) / τII_ij + εij_pl = λ[I...] .* dQdτij + dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r + τij = dτij .+ τij + setindex!.(τ, τij, I...) + setindex!.(ε_pl, εij_pl, I...) + τII[I...] = GeoParams.second_invariant(τij) + Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ + η_vep[I...] = 0.5 * τII_ij / εII_ve + else + # stress correction @ center + setindex!.(τ, dτij .+ τij, I...) + η_vep[I...] = ηij + τII[I...] = τII_ij + end + + Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) + else + Pr_c[I...] = zero(eltype(T)) + τij, = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) + dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) + setindex!.(τ, dτij .+ τij, I...) + end + end + + return nothing +end + + +# with free surface stabilization +@parallel_indices (i, j) function compute_V!( + Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy, dt +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + nx, ny = size(ρgy) + + if all((i, j) .< size(Vx) .- 1) + dVx = if iszero(ϕ_Vx[i + 1, j]) + zero(T) + else + (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) * ηdτ / av_xa(ητ) + end + Vx[i + 1, j + 1] += dVx + end + + if all((i, j) .< size(Vy) .- 1) + dVy = if iszero(ϕ_Vy[i, j + 1]) + zero(T) + else + ρg_correction = if iszero(dt) + zero(dt) + else + θ = 1 + # Interpolated Vx into Vy node (includes density gradient) + Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] + # Vertical velocity + Vyᵢⱼ = Vy[i + 1, j + 1] + # Get necessary buoyancy forces + j_N = min(j + 1, ny) + ρg_S = ρgy[i, j] + ρg_N = ρgy[i, j_N] + # Spatial derivatives + ∂ρg∂y = (ρg_N - ρg_S) * _dy + # correction term + (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + end + (-d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction) * ηdτ / av_ya(ητ) + end + Vy[i + 1, j + 1] += dVy + end + + return nothing +end + + + +# #### TESTING GROUNDS +# N = 4 +# ni = n, n = (N,N) +# ϕ = RockRatio(n, n) +# @test RockRatio(n, n) isa RockRatio + +# @test RockRatio(ni) isa RockRatio +# @test size_c(ϕ) === ni +# @test size_v(ϕ) === ni.+1 +# @test size_vx(ϕ) === (N+1, N) +# @test size_vy(ϕ) === (N, N+1) + +# A = [ +# 0 1 1 +# 1 0 1 +# ] + +# @test isvalid(A, 1, 1) || isvalid(A, 2, 2) === false +# @test isvalid(A, 1, 2) && isvalid(A, 2, 1) === true +# @test isvalid_x(A, 1, 1) === false +# @test isvalid_x(A, 1, 3) === true +# @test isvalid_y(A, 1, 1) === false +# @test isvalid_y(A, 1, 2) === true diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index cc7ab9afb..24d575f63 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -65,7 +65,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") C = C, ϕ = ϕ, η_vp = η_reg, - Ψ = 0 + Ψ = 15 ) rheology = ( @@ -101,7 +101,8 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.95 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.95 / √2.1) + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, Re=3e0, r=0.7, CFL = 0.95 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) @@ -191,7 +192,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") return nothing end -n = 32 +n = 128 nx = n ny = n figdir = "ShearBands2D" diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 91de6cf33..c4b53685c 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -1,6 +1,6 @@ # Viscous function compute_stress_increment(τij::Real, τij_o::Real, ηij, εij::Real, _Gdt, dτ_r) - dτij = dτ_r * fma(2.0 * ηij, εij, fma(-((τij - τij_o)) * ηij, _Gdt, -τij)) + dτij = dτ_r * fma(2.0 * ηij, εij, fma(-(τij - τij_o) * ηij, _Gdt, -τij)) return dτij end diff --git a/src/stokes/VelocityKernels.jl b/src/stokes/VelocityKernels.jl index 713b6ecdc..d860483e5 100644 --- a/src/stokes/VelocityKernels.jl +++ b/src/stokes/VelocityKernels.jl @@ -11,16 +11,37 @@ return nothing end -@parallel_indices (i, j, k) function compute_∇V!( - ∇V::AbstractArray{T,3}, Vx, Vy, Vz, _dx, _dy, _dz -) where {T} - d_xi(A) = _d_xi(A, i, j, k, _dx) - d_yi(A) = _d_yi(A, i, j, k, _dy) - d_zi(A) = _d_zi(A, i, j, k, _dz) - - @inbounds ∇V[i, j, k] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) - return nothing -end +# @parallel_indices (I...) function compute_∇V!( +# ∇V::AbstractArray{T,3}, Vx, Vy, Vz, ϕ::RockRatio, _dx, _dy, _dz +# ) where {T} +# d_xi(A) = _d_xi(A, I..., _dx) +# d_yi(A) = _d_yi(A, I..., _dy) +# d_zi(A) = _d_zi(A, I..., _dz) + +# if isvalid(ϕ, I...) +# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) +# else +# @inbounds ∇V[I...] = zero(T) +# end +# return nothing +# end + +# @parallel_indices (I...) function compute_∇V!( +# ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::RockRatio, _di::NTuple{N} +# ) where {T, N} +# @inline d_xi(A) = _d_xi(A, I..., _di[1]) +# @inline d_yi(A) = _d_yi(A, I..., _di[2]) +# @inline d_zi(A) = _d_zi(A, I..., _di[3]) + +# f = d_xi, d_yi, d_zi + +# if isvalid(ϕ, I...) +# @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) +# else +# @inbounds ∇V[I...] = zero(T) +# end +# return nothing +# end ## DEVIATORIC STRAIN RATE TENSOR @@ -289,34 +310,8 @@ end if all((i, j) .≤ size(Rx)) Rx[i, j] = d_xa(τxx) + d_yi(τxy) - d_xa(P) - av_xa(ρgx) end + if all((i, j) .≤ size(Ry)) - # θ = 1.0 - # Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] - # # Vertical velocity - # Vyᵢⱼ = Vy[i + 1, j + 1] - # # Get necessary buoyancy forces - # i_W, i_E = max(i - 1, 1), min(i + 1, nx) - # j_N = min(j + 1, ny) - # ρg_stencil = ( - # ρgy[i_W, j], - # ρgy[i, j], - # ρgy[i_E, j], - # ρgy[i_W, j_N], - # ρgy[i, j_N], - # ρgy[i_E, j_N], - # ) - # ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25 - # ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25 - # ρg_S = ρg_stencil[2] - # ρg_N = ρg_stencil[5] - # # Spatial derivatives - # ∂ρg∂x = (ρg_E - ρg_W) * _dx - # ∂ρg∂y = (ρg_N - ρg_S) * _dy - # # correction term - # ρg_correction = (Vxᵢⱼ * ∂ρg∂x + Vyᵢⱼ * ∂ρg∂y) * θ * dt - # Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy) + ρg_correction - # # Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy) - θ = 1.0 # Interpolated Vx into Vy node (includes density gradient) Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] From 5d0f3a1bdfeadac830a524839a93036d6fb4824b Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 21 Oct 2024 18:20:22 +0200 Subject: [PATCH 02/77] up --- Duretz2016.jl | 396 ++++++++++++++++++++++++++++++++++++++++ MiniKernels.jl | 23 ++- PlumeFreeSurface_2D.jl | 295 +++++++++++++++++++++--------- RayleightTaylor_2D.jl | 403 +++++++++++++++++++++++++++++++++++++++++ VariationalSolver.jl | 312 +++++++++++++++++++++++++++++++ mask.jl | 263 ++++++++++++++++++++++----- 6 files changed, 1549 insertions(+), 143 deletions(-) create mode 100644 Duretz2016.jl create mode 100644 RayleightTaylor_2D.jl create mode 100644 VariationalSolver.jl diff --git a/Duretz2016.jl b/Duretz2016.jl new file mode 100644 index 000000000..9f7fab9a1 --- /dev/null +++ b/Duretz2016.jl @@ -0,0 +1,396 @@ +using Printf +using JustRelax, JustRelax.JustRelax2D +import JustRelax.JustRelax2D as JR +const backend_JR = CPUBackend + +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) + +# Load script dependencies +using LinearAlgebra, GeoParams, GLMakie + +include("mask.jl") +include("MiniKernels.jl") + +# Velocity helper grids for the particle advection +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + @parallel f_x(A, B) +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 + +topo_fun(x, r) = -√(r^2 - x^2) + +function init_phases!(phases, particles) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index) + r=0.5 + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = (@index py[ip, i, j]) + + h = -√(r^2 - x^2) + if depth ≤ h + @index phases[ip, i, j] = 2.0 + else + @index phases[ip, i, j] = 1.0 + end + + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +# (Path)/folder where output data and figures are stored +n = 101 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +# function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 0 # thickness of sticky air layer + ly = 0.5 # domain length in y + lx = 0.5 # 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.25, -0.75 # origin coordinates (15km f sticky air layer) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + L = 0.5 + H = 0.5 + r = 0.5 + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=1), + CompositeRheology = CompositeRheology((LinearViscous(; η=1),)), + Gravity = ConstantGravity(; g=1), + ), + # Name = "Mantle", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=1), + CompositeRheology = CompositeRheology((LinearViscous(; η=1),)), + Gravity = ConstantGravity(; g=1), + ), + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 30, 40, 15 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + init_phases!(pPhases, particles) + phase_ratios = PhaseRatios(backend, length(rheology), ni) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # RockRatios + air_phase = 1 + ϕ = RockRatio(ni...) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + free_surface = true + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "Duretz2016" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + iterMax = 15e3 + nout = 1e3 + viscosity_cutoff = (-Inf, Inf) + free_surface = false + ητ = @zeros(ni...) + # while it < 1 + + ## variational solver + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + # Vx_on_Vy = @zeros(size(stokes.V.Vy)) + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + while iter ≤ iterMax + err < ϵ && break + # for _ in 1:100 + JR.compute_maxloc!(ητ, η; window=(1, 1)) + # update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + @parallel (@idx ni) compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ + ) + + JR.update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) JR.compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + # if rem(iter, nout) == 0 + # @copy η0 η + # end + # if do_visc + # update_viscosity!( + # stokes, + # phase_ratios, + # args, + # rheology, + # viscosity_cutoff; + # relaxation=viscosity_relaxation, + # ) + # end + + @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + # update_halo!(stokes.τ.xy) + + # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( + # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] + # ) + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni.+1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + # ϕ.Vx, + # ϕ.Vy, + _di..., + ) + # apply boundary conditions + # velocity2displacement!(stokes, dt) + # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + # end + # f,ax,h=heatmap(stokes.V.Vy) + # # f,ax,h=heatmap(stokes.V.Vx) + # Colorbar(f[1,2], h, label="Vy"); f + # update_halo!(@velocity(stokes)...) + # end + + iter += 1 + + if iter % nout == 0 && iter > 1 + # er_η = norm_mpi(@.(log10(η) - log10(η0))) + # er_η < 1e-3 && (do_visc = false) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), + norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), + norm(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + # end + isnan(err) && error("NaN(s)") + isinf(err) && error("Inf(s)") + end + end + + dt = compute_dt(stokes, di) / 2 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 1) == 0 + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + # arrows!( + # ax, + # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + # color = :red, + # ) + fig + save(joinpath(figdir, "$(it).png"), fig) + + end + # end + return nothing +# end +# ## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) + +# # # (Path)/folder where output data and figures are stored +# # n = 100 +# # nx = n +# # ny = n +# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid +# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +# # else +# # igg +# # end + diff --git a/MiniKernels.jl b/MiniKernels.jl index cbd0b2a32..63963a852 100644 --- a/MiniKernels.jl +++ b/MiniKernels.jl @@ -13,18 +13,17 @@ Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k+1] # masked versions -Base.@propagate_inbounds @inline center(A::T,ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = A[inds...] * ϕ[inds...] -Base.@propagate_inbounds @inline next(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = A[inds.+1...]* ϕ[inds.+1...] -Base.@propagate_inbounds @inline left(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i-1, j] * ϕ[i-1, j] -Base.@propagate_inbounds @inline right(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i+1, j] * ϕ[i+1, j] -Base.@propagate_inbounds @inline back(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i, j-1] * ϕ[i, j-1] -Base.@propagate_inbounds @inline front(A::T, ϕ::T, i::I, j::I) where {T<:AbstractArray, I<:Integer} = A[i, j+1] * ϕ[i, j+1] -Base.@propagate_inbounds @inline left(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i-1, j, k] * ϕ[i-1, j, k] -Base.@propagate_inbounds @inline right(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i+1, j, k] * ϕ[i+1, j, k] -Base.@propagate_inbounds @inline back(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j-1, k] * ϕ[i, j-1, k] -Base.@propagate_inbounds @inline front(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j+1, k] * ϕ[i, j+1, k] -Base.@propagate_inbounds @inline bot(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j, k-1] * ϕ[i, j, k-1] -Base.@propagate_inbounds @inline top(A::T, ϕ::T, i::I, j::I, k::I) where {T<:AbstractArray, I<:Integer} = A[i, j, k+1] * ϕ[i, j, k+1] +for fn in (:center, :next, :left, :right, :back, :front) + @eval begin + Base.@propagate_inbounds @inline ($fn)(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) + end +end +# Base.@propagate_inbounds @inline center(A::T,ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = center(A, inds...) * center(ϕ, inds...) +# Base.@propagate_inbounds @inline next(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = next(A, inds...) * next(ϕ, inds...) +# Base.@propagate_inbounds @inline left(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = left(A, inds...) * left(ϕ, inds...) +# Base.@propagate_inbounds @inline right(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = right(A, inds...) * right(ϕ, inds...) +# Base.@propagate_inbounds @inline back(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = back(A, inds...) * back(ϕ, inds...) +# Base.@propagate_inbounds @inline front(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = front(A, inds...) * front(ϕ, inds...) ## 2D mini kernels const T2 = AbstractArray{T,2} where {T} diff --git a/PlumeFreeSurface_2D.jl b/PlumeFreeSurface_2D.jl index bed3b9d19..55a87880b 100644 --- a/PlumeFreeSurface_2D.jl +++ b/PlumeFreeSurface_2D.jl @@ -1,4 +1,6 @@ +using Printf using JustRelax, JustRelax.JustRelax2D +import JustRelax.JustRelax2D as JR const backend_JR = CPUBackend using JustPIC, JustPIC._2D @@ -39,7 +41,6 @@ function init_phases!(phases, particles) @parallel_indices (i, j) function init_phases!(phases, px, py, index) r=100e3 - f(x, A, λ) = A * sin(π*x/λ) @inbounds for ip in cellaxes(phases) # quick escape @@ -69,7 +70,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ # (Path)/folder where output data and figures are stored -n = 100 +n = 101 nx = n ny = n igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -79,10 +80,10 @@ else end ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -# function main(igg, nx, ny) +function main(igg, nx, ny) # Physical domain ------------------------------------ - thick_air = 100e3 # thickness of sticky air layer + thick_air = 15e3 # thickness of sticky air layer ly = 400e3 + thick_air # domain length in y lx = 500e3 # domain length in x ni = nx, ny # number of cells @@ -96,10 +97,16 @@ end # Physical properties using GeoParams ---------------- rheology = rheology = ( # Name = "Air", + # SetMaterialParams(; + # Phase = 1, + # Density = ConstantDensity(; ρ=1e1), + # CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), + # Gravity = ConstantGravity(; g=9.81), + # ), SetMaterialParams(; Phase = 1, - Density = ConstantDensity(; ρ=1e1), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), Gravity = ConstantGravity(; g=9.81), ), # Name = "Mantle", @@ -120,7 +127,7 @@ end # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 30, 40, 15 + nxcell, max_xcell, min_xcell = 60, 80, 40 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) @@ -144,7 +151,7 @@ end # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.95 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- @@ -155,7 +162,7 @@ end ρg = @zeros(ni...), @zeros(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) + # @parallel init_P!(stokes.P, ρg[2], xci[2]) compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions @@ -173,26 +180,183 @@ end # Time loop t, it = 0.0, 0 dt = 1e3 * (3600 * 24 * 365.25) - while it < 15 - - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - phase_ratios, - rheology, - args, - dt, - igg; - kwargs = (; - iterMax = 50e3, - nout = 1e3, - viscosity_cutoff = (-Inf, Inf), - free_surface = true - ) - ) + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + iterMax = 15e3 + nout = 1e3 + viscosity_cutoff = (-Inf, Inf) + free_surface = false + ητ = @zeros(ni...) + while it < 20 + + ## variational solver + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + # Vx_on_Vy = @zeros(size(stokes.V.Vy)) + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + while iter ≤ iterMax + err < ϵ && break + # for _ in 1:100 + JR.compute_maxloc!(ητ, η; window=(1, 1)) + # update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + @parallel (@idx ni) compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ + ) + + JR.update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) JR.compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + # if rem(iter, nout) == 0 + # @copy η0 η + # end + # if do_visc + # update_viscosity!( + # stokes, + # phase_ratios, + # args, + # rheology, + # viscosity_cutoff; + # relaxation=viscosity_relaxation, + # ) + # end + + @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + # update_halo!(stokes.τ.xy) + + # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( + # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] + # ) + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni.+1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + # ϕ.Vx, + # ϕ.Vy, + _di..., + ) + # apply boundary conditions + # velocity2displacement!(stokes, dt) + # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + # end + # f,ax,h=heatmap(stokes.V.Vy) + # # f,ax,h=heatmap(stokes.V.Vx) + # Colorbar(f[1,2], h, label="Vy"); f + # update_halo!(@velocity(stokes)...) + # end + + iter += 1 + + if iter % nout == 0 && iter > 1 + # er_η = norm_mpi(@.(log10(η) - log10(η0))) + # er_η < 1e-3 && (do_visc = false) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), + norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), + norm(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + # end + isnan(err) && error("NaN(s)") + isinf(err) && error("Inf(s)") + end + end + dt = compute_dt(stokes, di) / 2 # ------------------------------ @@ -205,6 +369,7 @@ end inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) @show it += 1 t += dt @@ -214,7 +379,7 @@ end nt = 5 fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) arrows!( ax, xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., @@ -227,59 +392,17 @@ end end end return nothing -# end -## END OF MAIN SCRIPT ---------------------------------------------------------------- - -# # (Path)/folder where output data and figures are stored -# n = 100 -# nx = n -# ny = n -# igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid -# IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -# else -# igg -# end - -# main(igg, nx, ny) - -_di = inv.(di) -@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., ϕ, _di...) - - -@parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - copy(stokes.P), - stokes.P, - stokes.viscosity.η, - @zeros(ni...), - @zeros(ni.+1...), - stokes.τ.II, - stokes.viscosity.η_vep, - 0.2, - dt, - pt_stokes.θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ -) - -@parallel compute_V!( - @velocity(stokes)..., - @zeros(size(stokes.V.Vy)), - stokes.P, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - stokes.viscosity.η, - ϕ.Vx, - ϕ.Vy, - _di..., - 0, -) \ No newline at end of file +end +# ## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) + +# # # (Path)/folder where output data and figures are stored +# # n = 100 +# # nx = n +# # ny = n +# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid +# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +# # else +# # igg +# # end + diff --git a/RayleightTaylor_2D.jl b/RayleightTaylor_2D.jl new file mode 100644 index 000000000..d12ef8921 --- /dev/null +++ b/RayleightTaylor_2D.jl @@ -0,0 +1,403 @@ +using Printf +using JustRelax, JustRelax.JustRelax2D +import JustRelax.JustRelax2D as JR +const backend_JR = CPUBackend + +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) + +# Load script dependencies +using LinearAlgebra, GeoParams, GLMakie + +include("mask.jl") +include("MiniKernels.jl") + +# Velocity helper grids for the particle advection +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + @parallel f_x(A, B) +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 + +function init_phases!(phases, particles, A) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) + + f(x, A, λ) = A * sin(π * x / λ) + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) + @index phases[ip, i, j] = 2.0 + + if 0e0 ≤ depth ≤ 100e3 + @index phases[ip, i, j] = 1.0 + + elseif depth > (-f(x, A, 500e3) + (200e3 - A)) + @index phases[ip, i, j] = 3.0 + + end + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, A) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +# (Path)/folder where output data and figures are stored +n = 101 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 100e3 # thickness of sticky air layer + ly = 500e3 + thick_air # domain length in y + lx = 500e3 # 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) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + # Name = "Air", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=1e0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e16),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Crust", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Mantle", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ) + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 60, 80, 40 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + A = 5e3 # Amplitude of the anomaly + phase_ratios = PhaseRatios(backend, length(rheology), ni) + init_phases!(pPhases, particles, A) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # RockRatios + air_phase = 1 + ϕ = RockRatio(ni...) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.7 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = false), + no_slip = (left = false, right = false, top = false, bot = true), + free_surface = true, + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "RayleighTaylor2D" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 10e3 * (3600 * 24 * 365.25) + dt_max = 50e3 * (3600 * 24 * 365.25) + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + iterMax = 15e3 + nout = 1e3 + viscosity_cutoff = (-Inf, Inf) + free_surface = false + ητ = @zeros(ni...) + while it < 20 + + ## variational solver + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + # Vx_on_Vy = @zeros(size(stokes.V.Vy)) + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + while iter ≤ iterMax + err < ϵ && break + # for _ in 1:100 + JR.compute_maxloc!(ητ, η; window=(1, 1)) + # update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + @parallel (@idx ni) compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ + ) + + JR.update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) JR.compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + # if rem(iter, nout) == 0 + # @copy η0 η + # end + # if do_visc + # update_viscosity!( + # stokes, + # phase_ratios, + # args, + # rheology, + # viscosity_cutoff; + # relaxation=viscosity_relaxation, + # ) + # end + + @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + # update_halo!(stokes.τ.xy) + + # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( + # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] + # ) + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni.+1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + # ϕ.Vx, + # ϕ.Vy, + _di..., + ) + # apply boundary conditions + # velocity2displacement!(stokes, dt) + # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + # end + # f,ax,h=heatmap(stokes.V.Vy) + # # f,ax,h=heatmap(stokes.V.Vx) + # Colorbar(f[1,2], h, label="Vy"); f + # update_halo!(@velocity(stokes)...) + # end + + iter += 1 + + if iter % nout == 0 && iter > 1 + # er_η = norm_mpi(@.(log10(η) - log10(η0))) + # er_η < 1e-3 && (do_visc = false) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), + norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), + norm(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + # end + isnan(err) && error("NaN(s)") + isinf(err) && error("Inf(s)") + end + end + + dt = compute_dt(stokes, di) / 2 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 1) == 0 + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + arrows!( + ax, + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + color = :red, + ) + fig + save(joinpath(figdir, "$(it).png"), fig) + + end + end + return nothing +end +# ## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) + +# # # (Path)/folder where output data and figures are stored +# # n = 100 +# # nx = n +# # ny = n +# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid +# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +# # else +# # igg +# # end + diff --git a/VariationalSolver.jl b/VariationalSolver.jl new file mode 100644 index 000000000..f6c70ad13 --- /dev/null +++ b/VariationalSolver.jl @@ -0,0 +1,312 @@ +# unpack + +_di = inv.(di) +(; ϵ, r, θ_dτ, ηdτ) = pt_stokes +(; η, η_vep) = stokes.viscosity +ni = size(stokes.P) + +# ~preconditioner +ητ = deepcopy(η) +# @hide_communication b_width begin # communication/computation overlap +JR.compute_maxloc!(ητ, η; window=(1, 1)) +# update_halo!(ητ) +# end + +iterMax = 50e3 +nout = 1e3 +viscosity_cutoff = (-Inf, Inf) +free_surface = false + +# errors +err = 2 * ϵ +iter = 0 +err_evo1 = Float64[] +err_evo2 = Float64[] +norm_Rx = Float64[] +norm_Ry = Float64[] +norm_∇V = Float64[] +sizehint!(norm_Rx, Int(iterMax)) +sizehint!(norm_Ry, Int(iterMax)) +sizehint!(norm_∇V, Int(iterMax)) +sizehint!(err_evo1, Int(iterMax)) +sizehint!(err_evo2, Int(iterMax)) + +# solver loop +@copy stokes.P0 stokes.P +wtime0 = 0.0 +relλ = 0.2 +θ = deepcopy(stokes.P) +λ = @zeros(ni...) +λv = @zeros(ni .+ 1...) +η0 = deepcopy(η) +do_visc = true + +for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 +end +Vx_on_Vy = @zeros(size(stokes.V.Vy)) + +# compute buoyancy forces and viscosity +compute_ρg!(ρg[end], phase_ratios, rheology, args) +compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + +while iter ≤ iterMax + iterMin < iter && err < ϵ && break + + wtime0 += @elapsed begin + JR.compute_maxloc!(ητ, η; window=(1, 1)) + # update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + @parallel (@idx ni) compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ + ) + + JR.update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) JR.compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + # if rem(iter, nout) == 0 + # @copy η0 η + # end + # if do_visc + # update_viscosity!( + # stokes, + # phase_ratios, + # args, + # rheology, + # viscosity_cutoff; + # relaxation=viscosity_relaxation, + # ) + # end + + @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + # update_halo!(stokes.τ.xy) + + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( + Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] + ) + + # @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + Vx_on_Vy, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ.Vx, + ϕ.Vy, + _di..., + dt * free_surface, + ) + # apply boundary conditions + # velocity2displacement!(stokes, dt) + free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + # update_halo!(@velocity(stokes)...) + # end + end + + iter += 1 + + if iter % nout == 0 && iter > 1 + # er_η = norm_mpi(@.(log10(η) - log10(η0))) + # er_η < 1e-3 && (do_visc = false) + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, + stokes.R.Ry, + @velocity(stokes)..., + Vx_on_Vy, + stokes.P, + @stress(stokes)..., + ρg..., + ϕ.Vx, + ϕ.Vy, + _di..., + dt * free_surface, + ) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), + norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), + norm(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximumi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end +end + +# # compute vorticity +# @parallel (@idx ni .+ 1) compute_vorticity!( +# stokes.ω.xy, @velocity(stokes)..., inv.(di)... +# ) + +# # accumulate plastic strain tensor +# @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + +update_rock_ratio!(ϕ, phase_ratios, air_phase) + +heatmap(stokes.V.Vx) +heatmap(stokes.V.Vy) +heatmap(stokes.P) +heatmap(stokes.τ.xx) +heatmap(stokes.τ.yy) +heatmap(stokes.τ.xy) + + +heatmap(log10.(abs.(stokes.R.Rx))) +heatmap(log10.(abs.(stokes.R.Ry))) +heatmap(log10.(abs.(stokes.V.Vy))) +heatmap(ϕ.Vy) +heatmap(ϕ.Vx) +heatmap(ϕ.vertex) +heatmap(ϕ.center) + + +f,ax,h = heatmap(ϕ.vertex) +Colorbar(f[1,2], h);f + +f,ax,h = heatmap(ϕ.center) +Colorbar(f[1,2], h);f + +@parallel compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ.Vx, + ϕ.Vy, + _di..., +) + +extrema(stokes.R.Ry) +extrema(stokes.V.Vy) + +f,ax,h = heatmap(stokes.V.Vy) +Colorbar(f[1,2], h, label="Vy") + +-d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + +d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) +d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) +_dx, _dy = inv.(di) + +a = [_d_yi(stokes.τ.xy, ϕ.vertex, _dy, i, j) for i in 1:nx, j in 1:ny-1] +heatmap(a) + +b = [_d_ya(stokes.P, ϕ.center, _dy, i, j) for i in 1:nx, j in 1:ny-1] +f,ax,h = heatmap(b) +Colorbar(f[1,2], h, label="Vy") + +c = [_d_ya(stokes.τ.yy, ϕ.center, _dy, i, j) for i in 1:nx, j in 1:ny-1] +heatmap(c) +d = @. a-b+c + +f,ax,h=heatmap(d) +Colorbar(f[1,2], h, label="Vy") + +lines(stokes.P[1,:]) +lines(stokes.V.Vy[1,:]) + +v = [isvalid_c(ϕ, i, j) for i in 1:nx, j in 1:ny] +heatmap(xci..., v) + +v = [isvalid_v(ϕ, i, j) for i in 1:nx, j in 1:ny] +heatmap(xvi..., v) + + +vx = [isvalid_vx(ϕ, i, j) for i in 1:nx+1, j in 1:ny] +heatmap(vx) + +vy = [isvalid_vy(ϕ, i, j) for i in 1:nx, j in 1:ny+1] +heatmap(vy) + +heatmap( + xci[1].*1e-3, xci[2].*1e-3, + Array([argmax(p) for p in phase_ratios.center]), + colormap = :grayC) +px, py = particles + +heatmap( + xvi[1].*1e-3, xvi[2].*1e-3, + Array([argmax(p) for p in phase_ratios.vertex]), + colormap = :grayC) + + +heatmap(xvi..., ϕ.vertex) +# heatmap(xci..., ϕ.center) +heatmap(xci..., P) +scatter!(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle, colormap=:grayC) +scatter(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle, colormap=:grayC) + +P = copy(stokes.P) +P = copy(ϕ.vertex) +P[iszero.(P)] .= NaN +heatmap(xci..., P) + +id = particles.index.data[:] +px, py = particles.coords +scatter!(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle) +scatter(px.data[:][id], py.data[:][id], color=pPhases.data[:][id], markersize=5.0, marker=:circle) diff --git a/mask.jl b/mask.jl index 0062033e7..cf7d1a3e7 100644 --- a/mask.jl +++ b/mask.jl @@ -60,9 +60,9 @@ end @parallel_indices (I...) function update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase) if all(I .≤ size(ratio_center)) - ϕ.center[I...] = compute_rock_ratio(ratio_center, air_phase, I...) + ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) end - ϕ.vertex[I...] = compute_rock_ratio(ratio_vertex, air_phase, I...) + ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) return nothing end @@ -119,6 +119,81 @@ function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} return v * (A[I...] > 0) end +function isvalid_c(ϕ::RockRatio, i, j) + vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) + vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + div = vx * vy + return div * (ϕ.center[i, j] > 0) +end + +function isvalid_v(ϕ::RockRatio, i, j) + nx, ny = size(ϕ.Vx) + il = clamp(i - 1, 1, nx) + i0 = clamp(i, 1, nx) + j0 = clamp(j, 1, ny) + vx = (ϕ.Vx[i0, j0] > 0) * (ϕ.Vx[il, j0] > 0) + + nx, ny = size(ϕ.Vy) + jl = max(j - 1, 1) + i0 = clamp(i, 1, nx) + j0 = clamp(j, 1, ny) + vy = (ϕ.Vy[i0, j0] > 0) * (ϕ.Vy[i0, jl] > 0) + div = vx * vy + return div * (ϕ.vertex[i, j] > 0) +end + +function isvalid_vx(ϕ::RockRatio, i, j) + nx, ny = size(ϕ.center) + il = clamp(i - 1, 1, nx) + i0 = clamp(i, 1, nx) + j0 = clamp(j, 1, ny) + c = (ϕ.center[i0, j0] > 0) * (ϕ.center[il, j0] > 0) + + _, ny = size(ϕ.vertex) + jt = clamp(j - 1, 1, ny) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, jt] > 0) + div = c * v + return div * (ϕ.Vx[i, j] > 0) +end + +function isvalid_vy(ϕ::RockRatio, i, j) + nx, ny = size(ϕ.center) + jb = clamp(j - 1, 1, ny) + i0 = clamp(i, 1, nx) + j0 = clamp(j, 1, ny) + c = (ϕ.center[i0, j0] > 0) * (ϕ.center[i0, jb] > 0) + + nx, ny = size(ϕ.vertex) + ir = clamp(i - 1, 1, ny) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[ir, j] > 0) + div = c * v + return div * (ϕ.Vy[i, j] > 0) +end + + +# function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} +# v = true +# Base.@nexprs 2 j -> begin +# Base.@nexprs 2 i -> begin +# @inline +# ii = clamp(I[1] + 2 * i - 3, 1, size(A, 1)) +# jj = clamp(I[2] + 2 * j - 3, 1, size(A, 2)) +# v *= if N == 3 +# A[ii, jj, I[3]] > 0 +# else +# A[ii, jj] > 0 +# end +# end +# end +# if N === 3 +# Base.@nexprs 2 k -> begin +# kk = clamp(I[3] + 2 * k - 3, 1, size(A, 3)) +# v *= A[ii, jj, kk] > 0 +# end +# end +# return v * (A[I...] > 0) +# end + @parallel_indices (I...) function compute_∇V!( ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::RockRatio, _di::NTuple{N} ) where {T, N} @@ -128,7 +203,7 @@ end f = d_xi, d_yi, d_zi - if isvalid(ϕ.center, I...) + if isvalid_c(ϕ, I...) @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) else @inbounds ∇V[I...] = zero(T) @@ -136,12 +211,10 @@ end return nothing end -@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - @parallel_indices (I...) function compute_P!( P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, ϕ::RockRatio, dt, r, θ_dτ ) where {N} - if isvalid(ϕ.center, I...) + if isvalid_c(ϕ, I...) K = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, @cell(phase_ratio[I...])) RP[I...], P[I...] = JustRelax2D._compute_P!(P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ) else @@ -149,20 +222,6 @@ end end return nothing end - -@parallel (@idx ni) compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - stokes.viscosity.η, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ, -) # 2D kernel @parallel_indices (I...) function update_stresses_center_vertex_ps!( @@ -193,7 +252,7 @@ end ni = size(Pr) Ic = JustRelax2D.clamped_indices(ni, I...) - if isvalid(ϕ.vertex, I...) + if isvalid_v(ϕ, I...) # interpolate to ith vertex Pv_ij = JustRelax2D.av_clamped(Pr, Ic...) εxxv_ij = JustRelax2D.av_clamped(ε[1], Ic...) @@ -244,7 +303,7 @@ end ## center if all(I .≤ ni) - if isvalid(ϕ.center, I...) + if isvalid_c(ϕ, I...) # Material properties phase = @inbounds phase_center[I...] _Gdt = inv(JustRelax2D.fn_ratio(JustRelax2D.get_shear_modulus, rheology, phase) * dt) @@ -290,9 +349,10 @@ end Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) else Pr_c[I...] = zero(eltype(T)) - τij, = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) + # τij, = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) - setindex!.(τ, dτij .+ τij, I...) + # setindex!.(τ, dτij .+ τij, I...) + setindex!.(τ, dτij, I...) end end @@ -300,38 +360,154 @@ end end -# with free surface stabilization +# # with free surface stabilization @parallel_indices (i, j) function compute_V!( - Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy, dt + Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ::RockRatio, _dx, _dy ) where {T} d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) harm_ya(A) = _av_ya(A, i, j) - nx, ny = size(ρgy) - if all((i, j) .< size(Vx) .- 1) - dVx = if iszero(ϕ_Vx[i + 1, j]) - zero(T) + if isvalid_vx(ϕ, i, j) + Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) else - (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) * ηdτ / av_xa(ητ) + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) end - Vx[i + 1, j + 1] += dVx end - if all((i, j) .< size(Vy) .- 1) - dVy = if iszero(ϕ_Vy[i, j + 1]) - zero(T) + if all((i, j) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j) + Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) else - ρg_correction = if iszero(dt) - zero(dt) + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + end + end + + return nothing +end + +# @parallel_indices (i, j) function compute_V!( +# Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy +# ) where {T} +# d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) +# d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) +# d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) +# d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) +# av_xa(A) = _av_xa(A, i, j) +# av_ya(A) = _av_ya(A, i, j) +# harm_xa(A) = _av_xa(A, i, j) +# harm_ya(A) = _av_ya(A, i, j) + +# if all((i, j) .< size(Vx) .- 1) +# if iszero(ϕ_Vx[i + 1, j]) +# Rx[i, j] = zero(T) +# Vx[i + 1, j + 1] = zero(T) +# else +# Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) +# Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) +# end +# end + +# if all((i, j) .< size(Vy) .- 1) +# if iszero(ϕ_Vy[i, j + 1]) +# Ry[i, j] = zero(T) +# Vy[i + 1, j + 1] = zero(T) +# else +# Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) +# Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) +# end +# end + +# return nothing +# end + +# @parallel_indices (i, j) function compute_V!( +# Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy, dt +# ) where {T} +# d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) +# d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) +# d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) +# d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) +# av_xa(A) = _av_xa(A, i, j) +# av_ya(A) = _av_ya(A, i, j) +# harm_xa(A) = _av_xa(A, i, j) +# harm_ya(A) = _av_ya(A, i, j) + +# nx, ny = size(ρgy) + +# if all((i, j) .< size(Vx) .- 1) +# dVx = if iszero(ϕ_Vx[i + 1, j]) +# zero(T) +# else +# (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) * ηdτ / av_xa(ητ) +# end +# Vx[i + 1, j + 1] += dVx +# end + +# if all((i, j) .< size(Vy) .- 1) +# dVy = if iszero(ϕ_Vy[i, j + 1]) +# zero(T) +# else +# ρg_correction = if iszero(dt) +# zero(dt) +# else +# θ = 1 +# # Interpolated Vx into Vy node (includes density gradient) +# Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] +# # Vertical velocity +# Vyᵢⱼ = Vy[i + 1, j + 1] +# # Get necessary buoyancy forces +# j_N = min(j + 1, ny) +# ρg_S = ρgy[i, j] +# ρg_N = ρgy[i, j_N] +# # Spatial derivatives +# ∂ρg∂y = (ρg_N - ρg_S) * _dy +# # correction term +# (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt +# end +# (-d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction) * ηdτ / av_ya(ητ) +# end +# Vy[i + 1, j + 1] += dVy +# end + +# return nothing +# end + +@parallel_indices (i, j) function compute_Res!( + Rx::AbstractArray{T,2}, Ry, Vx, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ρgx, ρgy, ϕ_Vx, ϕ_Vy, _dx, _dy, dt +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + @inline av_xa(A) = _av_xa(A, i, j) + @inline av_ya(A) = _av_ya(A, i, j) + + nx, ny = size(ρgy) + @inbounds begin + if all((i, j) .≤ size(Rx)) + Rx[i, j] = if iszero(ϕ_Vx[i + 1, j]) + zero(T) else - θ = 1 + -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx) + end + end + + if all((i, j) .≤ size(Ry)) + Ry[i, j] = if iszero(ϕ_Vy[i, j + 1]) + zero(T) + else + θ = 1.0 # Interpolated Vx into Vy node (includes density gradient) Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] # Vertical velocity @@ -341,20 +517,17 @@ end ρg_S = ρgy[i, j] ρg_N = ρgy[i, j_N] # Spatial derivatives - ∂ρg∂y = (ρg_N - ρg_S) * _dy + ∂ρg∂y = (ρg_N - ρg_S) * _dy # correction term - (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction end - (-d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction) * ηdτ / av_ya(ητ) end - Vy[i + 1, j + 1] += dVy end return nothing end - - # #### TESTING GROUNDS # N = 4 # ni = n, n = (N,N) From 672f3e193a8fd00f4163b5ac01f8b5fab8661875 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Tue, 22 Oct 2024 00:22:46 +0200 Subject: [PATCH 03/77] fixes --- Cavity2D.jl | 396 +++++++++++++++++++++++++++++++++++++++++++ Duretz2016.jl | 134 ++++++++++----- VariationalSolver.jl | 18 +- mask.jl | 131 +++++++------- 4 files changed, 561 insertions(+), 118 deletions(-) create mode 100644 Cavity2D.jl diff --git a/Cavity2D.jl b/Cavity2D.jl new file mode 100644 index 000000000..d7e1d57a4 --- /dev/null +++ b/Cavity2D.jl @@ -0,0 +1,396 @@ +using Printf +using JustRelax, JustRelax.JustRelax2D +import JustRelax.JustRelax2D as JR +const backend_JR = CPUBackend + +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) + +# Load script dependencies +using LinearAlgebra, GeoParams, GLMakie + +include("mask.jl") +include("MiniKernels.jl") + +# Velocity helper grids for the particle advection +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + @parallel f_x(A, B) +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 + +function init_phases!(phases, particles) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index) + r = 5 + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = abs(@index py[ip, i, j]) + @index phases[ip, i, j] = 2.0 + + if 0e0 ≤ depth ≤ 10e0 || ((x - 25)^2 + (depth - 55)^2 ≤ r^2) + @index phases[ip, i, j] = 1.0 + end + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +# (Path)/folder where output data and figures are stored +n = 101 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 5e0 # thickness of sticky air layer + ly = 50e0 + thick_air # domain length in y + lx = 50e0 # 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) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + # Name = "Air", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=2.7e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Rock", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=2.7e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 120, 120, 100 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + init_phases!(pPhases, particles) + phase_ratios = PhaseRatios(backend, length(rheology), ni) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # RockRatios + air_phase = 1 + ϕ = RockRatio(ni...) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0*π, r=0.7, CFL = 0.95 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = false), + no_slip = (left = false, right = false, top = false, bot = true), + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "FreeSurfacePlume" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + iterMax = 5e3 + nout = 1e3 + viscosity_cutoff = (-Inf, Inf) + free_surface = false + ητ = @zeros(ni...) + while it < 1 + + ## variational solver + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + # Vx_on_Vy = @zeros(size(stokes.V.Vy)) + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + # for a in 1:100 + while iter ≤ iterMax + err < ϵ && break + # for _ in 1:100 + JR.compute_maxloc!(ητ, η; window=(1, 1)) + # update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + @parallel (@idx ni) compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + ϕ, + dt, + pt_stokes.r, + pt_stokes.θ_dτ + ) + + JR.update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) JR.compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + # if rem(iter, nout) == 0 + # @copy η0 η + # end + # if do_visc + # update_viscosity!( + # stokes, + # phase_ratios, + # args, + # rheology, + # viscosity_cutoff; + # relaxation=viscosity_relaxation, + # ) + # end + + @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + # update_halo!(stokes.τ.xy) + + # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( + # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] + # ) + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni.+1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + # ϕ.Vx, + # ϕ.Vy, + _di..., + ) + # apply boundary conditions + # velocity2displacement!(stokes, dt) + # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + # end + # f,ax,h=heatmap(stokes.V.Vy) + # # f,ax,h=heatmap(stokes.V.Vx) + # Colorbar(f[1,2], h, label="Vy"); f + # update_halo!(@velocity(stokes)...) + # end + + iter += 1 + + if iter % nout == 0 && iter > 1 + # er_η = norm_mpi(@.(log10(η) - log10(η0))) + # er_η < 1e-3 && (do_visc = false) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), + norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), + norm(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + # end + isnan(err) && error("NaN(s)") + isinf(err) && error("Inf(s)") + end + end + heatmap(stokes.V.Vy) + + dt = compute_dt(stokes, di) / 2 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 1) == 0 + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + # arrows!( + # ax, + # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + # color = :red, + # ) + # heatmap!(ax, stokes.V.Vy) + # heatmap!(ax, stokes.τ.xx) + + # ind = iszero.(stokes.V.Vy) + # stokes.V.Vy[ind] .= NaN + heatmap!(ax, stokes.V.Vy) + display(fig) + # save(joinpath(figdir, "$(it).png"), fig) + + end + end + return stokes, ϕ +end +# ## END OF MAIN SCRIPT ---------------------------------------------------------------- +stoke, ϕ = main(igg, nx, ny); + +# # # (Path)/folder where output data and figures are stored +# # n = 100 +# # nx = n +# # ny = n +# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid +# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +# # else +# # igg +# # end + diff --git a/Duretz2016.jl b/Duretz2016.jl index 9f7fab9a1..674a7dcb9 100644 --- a/Duretz2016.jl +++ b/Duretz2016.jl @@ -43,6 +43,7 @@ function init_phases!(phases, particles) @parallel_indices (i, j) function init_phases!(phases, px, py, index) r=0.5 + rinc=0.1 @inbounds for ip in cellaxes(phases) # quick escape @@ -51,6 +52,7 @@ function init_phases!(phases, particles) x = @index px[ip, i, j] depth = (@index py[ip, i, j]) + # h = √(r^2 - x^2) - 0.4330127018922193*2 h = -√(r^2 - x^2) if depth ≤ h @index phases[ip, i, j] = 2.0 @@ -58,6 +60,15 @@ function init_phases!(phases, particles) @index phases[ip, i, j] = 1.0 end + # h = √(r^2 - x^2) - 0.55*2 + # h = -√(r^2 - x^2) - 0.2 + # if depth ≤ h + # @index phases[ip, i, j] = 1.0 + # end + + # if (x^2 + (depth + 0.75)^2 ≤ rinc^2) + # @index phases[ip, i, j] = 1.0 + # end end return nothing @@ -65,6 +76,14 @@ function init_phases!(phases, particles) @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) end + +# init_phases!(pPhases, particles) +# phase_ratios = PhaseRatios(backend, length(rheology), ni) +# update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) +# p=[argmax(p) for p in phase_ratios.center] +# heatmap(p) + + ## END OF HELPER FUNCTION ------------------------------------------------------------ # (Path)/folder where output data and figures are stored @@ -78,7 +97,7 @@ else end ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -# function main(igg, nx, ny) +function main(igg, nx, ny) # Physical domain ------------------------------------ thick_air = 0 # thickness of sticky air layer @@ -115,7 +134,7 @@ end # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 30, 40, 15 + nxcell, max_xcell, min_xcell = 40, 70, 15 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) @@ -173,7 +192,7 @@ end (; ϵ, r, θ_dτ, ηdτ) = pt_stokes (; η, η_vep) = stokes.viscosity ni = size(stokes.P) - iterMax = 15e3 + iterMax = 5e3 nout = 1e3 viscosity_cutoff = (-Inf, Inf) free_surface = false @@ -347,50 +366,73 @@ end dt = compute_dt(stokes, di) / 2 # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - @show it += 1 - t += dt - - if it == 1 || rem(it, 1) == 0 - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 5 - fig = Figure(size = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - # arrows!( - # ax, - # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - # color = :red, - # ) - fig - save(joinpath(figdir, "$(it).png"), fig) - - end - # end - return nothing -# end + heatmap(stokes.τ.xy) + + # # Advection -------------------- + # # advect particles in space + # advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # # advect particles in memory + # move_particles!(particles, xvi, particle_args) + # # check if we need to inject particles + # inject_particles_phase!(particles, pPhases, (), (), xvi) + # # update phase ratios + # update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # @show it += 1 + # t += dt + + # if it == 1 || rem(it, 1) == 0 + # velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + # nt = 5 + # fig = Figure(size = (900, 900), title = "t = $t") + # ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + # # arrows!( + # # ax, + # # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + # # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + # # color = :red, + # # ) + # fig + # save(joinpath(figdir, "$(it).png"), fig) + + # end + # # end + # return nothing + # heatmap(stokes.τ.xy) + ind = iszero.(stokes.P) + stokes.P[ind] .= NaN + f,ax,=heatmap(xci...,stokes.P) + ind = iszero.(stokes.V.Vy) + stokes.V.Vy[ind] .= NaN + f,ax,=heatmap(xci[1],xvi[2], stokes.V.Vy[2:end-1,:]) + + px, py = particles.coords + scatter!( + px.data[:], py.data[:],color=pPhases.data[:], + markersize=3) + h = @. -√(0.5^2 - xci[1]^2) + lines!(ax,xci[1], h, color=:red) + xlims!(ax, -0.25, -0.15) + ylims!(ax, -0.5, -0.4) + f + # ind = iszero.(stokes.V.Vx) + # stokes.V.Vx[ind] .= NaN + # heatmap(stokes.V.Vx) + # stokes +end # ## END OF MAIN SCRIPT ---------------------------------------------------------------- main(igg, nx, ny) -# # # (Path)/folder where output data and figures are stored -# # n = 100 -# # nx = n -# # ny = n -# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid -# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -# # else -# # igg -# # end +# heatmap(stokes.τ.xy) +# # ind = iszero.(stokes.P) +# # stokes.P[ind] .= NaN +# heatmap(stokes.P) +# ind = iszero.(stokes.V.Vy) +# stokes.V.Vy[ind] .= NaN +# heatmap(stokes.V.Vy) +# # ind = iszero.(stokes.V.Vx) +# # stokes.V.Vx[ind] .= NaN +# # heatmap(stokes.V.Vx) \ No newline at end of file diff --git a/VariationalSolver.jl b/VariationalSolver.jl index f6c70ad13..b8441c291 100644 --- a/VariationalSolver.jl +++ b/VariationalSolver.jl @@ -270,17 +270,29 @@ Colorbar(f[1,2], h, label="Vy") lines(stokes.P[1,:]) lines(stokes.V.Vy[1,:]) +xc = [x for x in xci[1], y in xci[2]] +yc = [y for x in xci[1], y in xci[2]] + v = [isvalid_c(ϕ, i, j) for i in 1:nx, j in 1:ny] heatmap(xci..., v) +heatmap(xci..., ϕ.center) +h = @. -√(0.5^2 - xci[1]^2) +lines!(xci[1], h, color=:red) + v = [isvalid_v(ϕ, i, j) for i in 1:nx, j in 1:ny] heatmap(xvi..., v) +heatmap(xvi..., ϕ.vertex) +h = @. -√(0.5^2 - xvi[1]^2) +lines!(xvi[1], h, color=:red) -vx = [isvalid_vx(ϕ, i, j) for i in 1:nx+1, j in 1:ny] -heatmap(vx) +vx = [isvalid_vx(ϕ, i+1, j) for i in 1:nx-1, j in 1:ny] +heatmap(xvi[1], xci[2], vx) +h = @. -√(0.5^2 - xci[1]^2) +lines!(xci[1], h, color=:red) -vy = [isvalid_vy(ϕ, i, j) for i in 1:nx, j in 1:ny+1] +vy = [isvalid_vy(ϕ, i, j+1) for i in 1:nx, j in 1:ny-1] heatmap(vy) heatmap( diff --git a/mask.jl b/mask.jl index cf7d1a3e7..edc040ac2 100644 --- a/mask.jl +++ b/mask.jl @@ -122,55 +122,48 @@ end function isvalid_c(ϕ::RockRatio, i, j) vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) - div = vx * vy - return div * (ϕ.center[i, j] > 0) + v = vx * vy + # return v * (ϕ.center[i, j] > 0) + return true end function isvalid_v(ϕ::RockRatio, i, j) nx, ny = size(ϕ.Vx) - il = clamp(i - 1, 1, nx) - i0 = clamp(i, 1, nx) - j0 = clamp(j, 1, ny) - vx = (ϕ.Vx[i0, j0] > 0) * (ϕ.Vx[il, j0] > 0) + j_bot = max(j - 1, 1) + j0 = min(j, ny) + vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - nx, ny = size(ϕ.Vy) - jl = max(j - 1, 1) - i0 = clamp(i, 1, nx) - j0 = clamp(j, 1, ny) - vy = (ϕ.Vy[i0, j0] > 0) * (ϕ.Vy[i0, jl] > 0) - div = vx * vy - return div * (ϕ.vertex[i, j] > 0) + nx, ny = size(ϕ.Vy) + i_left = max(i - 1, 1) + i0 = min(i, nx) + vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + v = vx * vy + # return v * (ϕ.vertex[i, j] > 0) + return true end function isvalid_vx(ϕ::RockRatio, i, j) - nx, ny = size(ϕ.center) - il = clamp(i - 1, 1, nx) - i0 = clamp(i, 1, nx) - j0 = clamp(j, 1, ny) - c = (ϕ.center[i0, j0] > 0) * (ϕ.center[il, j0] > 0) - - _, ny = size(ϕ.vertex) - jt = clamp(j - 1, 1, ny) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, jt] > 0) - div = c * v - return div * (ϕ.Vx[i, j] > 0) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i-1, j] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j+1] > 0) + cv = c * v + return cv * (ϕ.Vx[i, j] > 0) + # c = (ϕ.center[i, j] > 0) || (ϕ.center[i-1, j] > 0) + # v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j+1] > 0) + # cv = c || v + # return cv || (ϕ.Vx[i, j] > 0) end function isvalid_vy(ϕ::RockRatio, i, j) - nx, ny = size(ϕ.center) - jb = clamp(j - 1, 1, ny) - i0 = clamp(i, 1, nx) - j0 = clamp(j, 1, ny) - c = (ϕ.center[i0, j0] > 0) * (ϕ.center[i0, jb] > 0) - - nx, ny = size(ϕ.vertex) - ir = clamp(i - 1, 1, ny) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[ir, j] > 0) - div = c * v - return div * (ϕ.Vy[i, j] > 0) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) + cv = c * v + return cv * (ϕ.Vy[i, j] > 0) + # c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) + # v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) + # cv = c || v + # return cv || (ϕ.Vy[i, j] > 0) end - # function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} # v = true # Base.@nexprs 2 j -> begin @@ -374,7 +367,7 @@ end harm_ya(A) = _av_ya(A, i, j) if all((i, j) .< size(Vx) .- 1) - if isvalid_vx(ϕ, i, j) + if isvalid_vx(ϕ, i + 1, j) Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) else @@ -384,7 +377,7 @@ end end if all((i, j) .< size(Vy) .- 1) - if isvalid_vy(ϕ, i, j) + if isvalid_vy(ϕ, i, j + 1) Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) else @@ -396,40 +389,40 @@ end return nothing end -# @parallel_indices (i, j) function compute_V!( -# Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy -# ) where {T} -# d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) -# d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) -# d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) -# d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) -# av_xa(A) = _av_xa(A, i, j) -# av_ya(A) = _av_ya(A, i, j) -# harm_xa(A) = _av_xa(A, i, j) -# harm_ya(A) = _av_ya(A, i, j) +@parallel_indices (i, j) function compute_V!( + Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) -# if all((i, j) .< size(Vx) .- 1) -# if iszero(ϕ_Vx[i + 1, j]) -# Rx[i, j] = zero(T) -# Vx[i + 1, j + 1] = zero(T) -# else -# Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) -# Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) -# end -# end + if all((i, j) .< size(Vx) .- 1) + if iszero(ϕ_Vx[i + 1, j]) + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) + else + Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + end + end -# if all((i, j) .< size(Vy) .- 1) -# if iszero(ϕ_Vy[i, j + 1]) -# Ry[i, j] = zero(T) -# Vy[i + 1, j + 1] = zero(T) -# else -# Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) -# Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) -# end -# end + if all((i, j) .< size(Vy) .- 1) + if iszero(ϕ_Vy[i, j + 1]) + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + else + Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + end + end -# return nothing -# end + return nothing +end # @parallel_indices (i, j) function compute_V!( # Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy, dt From 12a659bf91bb069b9782b2330a712a53846daf36 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 13:50:47 +0200 Subject: [PATCH 04/77] up miniapp --- .../RayleighTaylor2D.jl | 107 +++++++++--------- 1 file changed, 53 insertions(+), 54 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index d0efb3310..203d05eb2 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -1,5 +1,4 @@ -# using CUDA -using JustRelax, JustRelax.JustRelax2D +using JustRelax, JustRelax.JustRelax2D const backend_JR = CPUBackend using JustPIC, JustPIC._2D @@ -9,9 +8,9 @@ using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) # Load script dependencies -using LinearAlgebra, GeoParams, CairoMakie#GLMakie +using LinearAlgebra, GeoParams, GLMakie -## START OF HELPER FUNCTION ---------------------------------------------------------- +# Velocity helper grids for the particle advection function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -26,8 +25,9 @@ macro all_j(A) esc(:($A[$idx_j])) end -@parallel_indices (i, j) function init_P!(P, ρg, z) - P[i, j] = sum(abs(ρg[i, jj] * z[jj]) for jj in j:size(P, 2)) +# 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 @@ -62,8 +62,18 @@ function init_phases!(phases, particles, A) end ## END OF HELPER FUNCTION ------------------------------------------------------------ +# (Path)/folder where output data and figures are stored +n = 101 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function RT_2D(igg, nx, ny) +function main(igg, nx, ny) # Physical domain ------------------------------------ thick_air = 100e3 # thickness of sticky air layer @@ -104,7 +114,7 @@ function RT_2D(igg, nx, ny) # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 30, 40, 10 + nxcell, max_xcell, min_xcell = 60, 80, 40 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) @@ -121,10 +131,15 @@ function RT_2D(igg, nx, ny) update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) # ---------------------------------------------------- + # RockRatios + air_phase = 1 + ϕ = RockRatio(ni...) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.98 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- @@ -153,18 +168,30 @@ function RT_2D(igg, nx, ny) # Time loop t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) + dt = 10e3 * (3600 * 24 * 365.25) dt_max = 50e3 * (3600 * 24 * 365.25) - while it < 500 - + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + iterMax = 15e3 + nout = 1e3 + viscosity_cutoff = (-Inf, Inf) + free_surface = false + ητ = @zeros(ni...) + while it < 1 + + ## variational solver # Stokes solver ---------------- - solve!( + solve_VariationalStokes!( stokes, pt_stokes, di, flow_bcs, ρg, phase_ratios, + ϕ, rheology, args, dt, @@ -173,20 +200,11 @@ function RT_2D(igg, nx, ny) iterMax = 50e3, iterMin = 1e3, viscosity_relaxation = 1e-2, - nout = 5e3, - free_surface = true, + nout = 2e3, viscosity_cutoff = (-Inf, Inf) ) ) - dt = if it ≤ 10 - min(compute_dt(stokes, di), 1e3 * (3600 * 24 * 365.25)) - elseif 10 < it ≤ 20 - min(compute_dt(stokes, di), 10e3 * (3600 * 24 * 365.25)) - elseif 20 < it ≤ 30 - min(compute_dt(stokes, di), 25e3 * (3600 * 24 * 365.25)) - else - min(compute_dt(stokes, di), dt_max) - end + dt = compute_dt(stokes, di, dt_max) # ------------------------------ # Advection -------------------- @@ -198,52 +216,33 @@ function RT_2D(igg, nx, ny) inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) @show it += 1 t += dt if it == 1 || rem(it, 5) == 0 - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 2 - - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:]./1e3 - pyv = ppy.data[:]./1e3 - clr = pPhases.data[:] + px, py = particles.coords + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - scatter!( - ax, - pxv, pyv, - color=clr, - colormap = :lajolla, - markersize = 3 - ) + # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + scatter!(ax, px.data[:].*1e-3, py.data[:].*1e-3, color = pPhases.data[:], colormap = :grayC) arrows!( ax, xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :darkblue, + color = :red, ) fig save(joinpath(figdir, "$(it).png"), fig) - end + end end - return -end -## END OF MAIN SCRIPT ---------------------------------------------------------------- - -# (Path)/folder where output data and figures are stored -n = 100 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg + return nothing end +# ## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) -RT_2D(igg, nx, ny) From f97170a6f0a99774b14dd3ffc8ec3fb56c44be55 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 13:51:18 +0200 Subject: [PATCH 05/77] move variational Stokes into JR proper --- src/JustRelax.jl | 5 +- src/MiniKernels.jl | 103 +++++---- src/advection/weno5.jl | 2 +- src/common.jl | 15 ++ src/stokes/PressureKernels.jl | 4 +- src/stokes/StressRotation.jl | 18 +- src/stokes/VelocityKernels.jl | 75 ++----- src/thermal_diffusion/DiffusionPT_kernels.jl | 83 ++++--- src/variational_stokes/MiniKernels.jl | 14 ++ src/variational_stokes/PressureKernels.jl | 61 +++++ src/variational_stokes/Stokes2D.jl | 220 +++++++++++++++++++ src/variational_stokes/StressKernels.jl | 135 ++++++++++++ src/variational_stokes/VelocityKernels.jl | 51 +++++ src/variational_stokes/mask.jl | 149 +++++++++++++ src/variational_stokes/types.jl | 20 ++ 15 files changed, 798 insertions(+), 157 deletions(-) create mode 100644 src/variational_stokes/MiniKernels.jl create mode 100644 src/variational_stokes/PressureKernels.jl create mode 100644 src/variational_stokes/Stokes2D.jl create mode 100644 src/variational_stokes/StressKernels.jl create mode 100644 src/variational_stokes/VelocityKernels.jl create mode 100644 src/variational_stokes/mask.jl create mode 100644 src/variational_stokes/types.jl diff --git a/src/JustRelax.jl b/src/JustRelax.jl index d7bbbb1d4..8d80ebff6 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -1,5 +1,6 @@ module JustRelax +using Adapt using Reexport @reexport using ImplicitGlobalGrid using LinearAlgebra @@ -25,10 +26,10 @@ PTArray(::T) where {T} = error(ArgumentError("Unknown backend $T")) export PTArray, CPUBackend, CUDABackend, AMDGPUBackend include("types/stokes.jl") -# export StokesArrays, PTStokesCoeffs include("types/heat_diffusion.jl") -# export ThermalArrays, PTThermalCoeffs + +include("variational_stokes/types.jl") include("types/weno.jl") diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl index 0695eb56f..ce73cb1be 100644 --- a/src/MiniKernels.jl +++ b/src/MiniKernels.jl @@ -1,72 +1,79 @@ +# unmasked versions +Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds...] +Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds.+1...] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where I<:Integer = A[i-1, j] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where I<:Integer = A[i+1, j] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j-1] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j+1] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i-1, j, k] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i+1, j, k] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j-1, k] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j+1, k] +Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k-1] +Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k+1] + ## 2D mini kernels const T2 = AbstractArray{T,2} where {T} # finite differences -@inline _d_xa(A::T, i, j, _dx) where {T<:T2} = (-A[i, j] + A[i + 1, j]) * _dx -@inline _d_ya(A::T, i, j, _dy) where {T<:T2} = (-A[i, j] + A[i, j + 1]) * _dy -@inline _d_xi(A::T, i, j, _dx) where {T<:T2} = (-A[i, j + 1] + A[i + 1, j + 1]) * _dx -@inline _d_yi(A::T, i, j, _dy) where {T<:T2} = (-A[i + 1, j] + A[i + 1, j + 1]) * _dy +Base.@propagate_inbounds @inline _d_xa(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + right(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + front(A, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::AbstractArray, _dz, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + top(A, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-front(A, I...) + next(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-right(A, I...) + next(A, I...)) * _dy + # averages -@inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) -@inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) -@inline _av_xa(A::T, i, j) where {T<:T2} = (A[i, j] + A[i + 1, j]) * 0.5 -@inline _av_ya(A::T, i, j) where {T<:T2} = (A[i, j] + A[i, j + 1]) * 0.5 -@inline _av_xi(A::T, i, j) where {T<:T2} = (A[i, j + 1], A[i + 1, j + 1]) * 0.5 -@inline _av_yi(A::T, i, j) where {T<:T2} = (A[i + 1, j], A[i + 1, j + 1]) * 0.5 +Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) +Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) +Base.@propagate_inbounds @inline _av_xa(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (center(A, I...) + right(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_ya(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (center(A, I...) + front(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_xi(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (front(A, I...), next(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_yi(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (right(A, I...), next(A, I...)) * 0.5 # harmonic averages -@inline function _harm(A::T, i, j) where {T<:T2} +Base.@propagate_inbounds @inline function _harm(A::T, i, j) where {T<:T2} return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2)) end -@inline function _harm_a(A::T, i, j) where {T<:T2} +Base.@propagate_inbounds @inline function _harm_a(A::T, i, j) where {T<:T2} return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1)) end -@inline function _harm_xa(A::T, i, j) where {T<:T2} - return eltype(A)(2) * (inv(A[i + 1, j]) + inv(A[i, j])) +Base.@propagate_inbounds @inline function _harm_xa(A::T, I::Vararg{Integer, 2}) where {T<:T2} + return eltype(A)(2) * (inv(right(A, I...)) + inv(center(A, I...))) end -@inline function _harm_ya(A::T, i, j) where {T<:T2} - return eltype(A)(2) * (inv(A[i, j + 1]) + inv(A[i, j])) +Base.@propagate_inbounds @inline function _harm_ya(A::T, I::Vararg{Integer, 2}) where {T<:T2} + return eltype(A)(2) * (inv(front(A, I...)) + inv(center(A, I...))) end #others -@inline function _gather(A::T, i, j) where {T<:T2} - return A[i, j], A[i + 1, j], A[i, j + 1], A[i + 1, j + 1] +Base.@propagate_inbounds @inline function _gather(A::T, I::Vararg{Integer, 2}) where {T<:T2} + return center(A, I...), right(A, I...), front(A, I...), next(A, I...) end ## 3D mini kernels const T3 = AbstractArray{T,3} where {T} # finite differences -@inline _d_xa(A::T, i, j, k, _dx) where {T<:T3} = (-A[i, j, k] + A[i + 1, j, k]) * _dx -@inline _d_ya(A::T, i, j, k, _dy) where {T<:T3} = (-A[i, j, k] + A[i, j + 1, k]) * _dy -@inline _d_za(A::T, i, j, k, _dz) where {T<:T3} = (-A[i, j, k] + A[i, j, k + 1]) * _dz -@inline function _d_xi(A::T, i, j, k, _dx) where {T<:T3} - return (-A[i, j + 1, k + 1] + A[i + 1, j + 1, k + 1]) * _dx -end -@inline function _d_yi(A::T, i, j, k, _dy) where {T<:T3} - return (-A[i + 1, j, k + 1] + A[i + 1, j + 1, k + 1]) * _dy -end @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3} - return (-A[i + 1, j + 1, k] + A[i + 1, j + 1, k + 1]) * _dz + return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz end # averages -@inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) -@inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i + 1, j, k]) -@inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i, j + 1, k]) -@inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i, j, k + 1]) -@inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) -@inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) -@inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) -@inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) -@inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) -@inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) +Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) +Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + right(A, i, j, k)) +Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + front(A, i, j, k)) +Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + top(A, i, j, k)) +Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) +Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) +Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) +Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) +Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) +Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) # harmonic averages @inline function _harm_x(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(A[i, j, k]) + inv(A[i + 1, j, k])) + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(right(A, i, j, k))) end @inline function _harm_y(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(A[i, j, k]) + inv(A[i, j + 1, k])) + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(front(A, i, j, k))) end @inline function _harm_z(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(A[i, j, k]) + inv(A[i, j, k + 1])) + return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(top(A, i, j, k))) end @inline function _harm_xy(A::T, i, j, k) where {T<:T3} return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k)) @@ -88,16 +95,16 @@ end end # others -@inline function _gather_yz(A::T, i, j, k) where {T<:T3} - return A[i, j, k], A[i, j + 1, k], A[i, j, k + 1], A[i, j + 1, k + 1] +Base.@propagate_inbounds @inline function _gather_yz(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), front(A, i, j, k), top(A, i, j, k), A[i, j + 1, k + 1] end -@inline function _gather_xz(A::T, i, j, k) where {T<:T3} - return A[i, j, k], A[i + 1, j, k], A[i, j, k + 1], A[i + 1, j, k + 1] +Base.@propagate_inbounds @inline function _gather_xz(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), right(A, i, j, k), top(A, i, j, k), A[i + 1, j, k + 1] end -@inline function _gather_xy(A::T, i, j, k) where {T<:T3} - return A[i, j, k], A[i + 1, j, k], A[i, j + 1, k], A[i + 1, j + 1, k] +Base.@propagate_inbounds @inline function _gather_xy(A::T, i, j, k) where {T<:T3} + return center(A, i, j, k), right(A, i, j, k), front(A, i, j, k), A[i + 1, j + 1, k] end -@inline _current(A::T, i, j, k) where {T<:T3} = A[i, j, k] +@inline _current(A::T, i, j, k) where {T<:T3} = center(A, i, j, k) ## Because mysum(::generator) does not work inside CUDA kernels... @inline mysum(A, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ranges...) diff --git a/src/advection/weno5.jl b/src/advection/weno5.jl index fb2e98753..b97c7c53b 100644 --- a/src/advection/weno5.jl +++ b/src/advection/weno5.jl @@ -1,4 +1,4 @@ -using MuladdMacro, Adapt +using MuladdMacro # 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) diff --git a/src/common.jl b/src/common.jl index c88da13ba..c88600955 100644 --- a/src/common.jl +++ b/src/common.jl @@ -82,6 +82,21 @@ export rotate_stress_particles! include("stokes/VelocityKernels.jl") +# variational Stokes +include("variational_stokes/mask.jl") +export RockRatio, update_rock_ratio! + +include("variational_stokes/PressureKernels.jl") + +include("variational_stokes/MiniKernels.jl") + +include("variational_stokes/StressKernels.jl") + +include("variational_stokes/VelocityKernels.jl") + +include("variational_stokes/Stokes2D.jl") +export solve_VariationalStokes! + # thermal diffusion include("thermal_diffusion/DiffusionPT.jl") diff --git a/src/stokes/PressureKernels.jl b/src/stokes/PressureKernels.jl index 03bfd7c4c..468fa9c7c 100644 --- a/src/stokes/PressureKernels.jl +++ b/src/stokes/PressureKernels.jl @@ -59,12 +59,12 @@ function compute_P!( r, θ_dτ; ΔTc=nothing, - ϕ=nothing, + melt_fraction=nothing, kwargs..., ) where {N} ni = size(P) @parallel (@idx ni) compute_P_kernel!( - P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ, ΔTc, ϕ + P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ, ΔTc, melt_fraction ) return nothing end diff --git a/src/stokes/StressRotation.jl b/src/stokes/StressRotation.jl index 860ca3111..c7119422e 100644 --- a/src/stokes/StressRotation.jl +++ b/src/stokes/StressRotation.jl @@ -3,10 +3,10 @@ using StaticArrays # Vorticity tensor @parallel_indices (I...) function compute_vorticity!(ωxy, Vx, Vy, _dx, _dy) - dx(A) = _d_xa(A, I..., _dx) - dy(A) = _d_ya(A, I..., _dy) + Base.@propagate_inbounds @inline dx(A) = _d_xa(A, _dx, I...) + Base.@propagate_inbounds @inline dy(A) = _d_ya(A, _dy, I...) - ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) + @inbounds ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) return nothing end @@ -14,13 +14,13 @@ end @parallel_indices (I...) function compute_vorticity!( ωyz, ωxz, ωxy, Vx, Vy, Vz, _dx, _dy, _dz ) - dx(A) = _d_xa(A, I..., _dx) - dy(A) = _d_ya(A, I..., _dy) - dz(A) = _d_za(A, I..., _dz) + Base.@propagate_inbounds @inline dx(A) = _d_xa(A, _dx, I...) + Base.@propagate_inbounds @inline dy(A) = _d_ya(A, _dy, I...) + Base.@propagate_inbounds @inline dz(A) = _d_za(A, _dz, I...) - ωyz[I...] = 0.5 * (dy(Vz) - dz(Vy)) - ωxz[I...] = 0.5 * (dz(Vx) - dx(Vz)) - ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) + @inbounds ωyz[I...] = 0.5 * (dy(Vz) - dz(Vy)) + @inbounds ωxz[I...] = 0.5 * (dz(Vx) - dx(Vz)) + @inbounds ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) return nothing end diff --git a/src/stokes/VelocityKernels.jl b/src/stokes/VelocityKernels.jl index d860483e5..6dca036b5 100644 --- a/src/stokes/VelocityKernels.jl +++ b/src/stokes/VelocityKernels.jl @@ -3,55 +3,24 @@ @parallel_indices (i, j) function compute_∇V!( ∇V::AbstractArray{T,2}, Vx, Vy, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) + d_xi(A) = _d_xi(A,_dx, i, j) + d_yi(A) = _d_yi(A,_dy, i, j) ∇V[i, j] = d_xi(Vx) + d_yi(Vy) return nothing end -# @parallel_indices (I...) function compute_∇V!( -# ∇V::AbstractArray{T,3}, Vx, Vy, Vz, ϕ::RockRatio, _dx, _dy, _dz -# ) where {T} -# d_xi(A) = _d_xi(A, I..., _dx) -# d_yi(A) = _d_yi(A, I..., _dy) -# d_zi(A) = _d_zi(A, I..., _dz) - -# if isvalid(ϕ, I...) -# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) -# else -# @inbounds ∇V[I...] = zero(T) -# end -# return nothing -# end - -# @parallel_indices (I...) function compute_∇V!( -# ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::RockRatio, _di::NTuple{N} -# ) where {T, N} -# @inline d_xi(A) = _d_xi(A, I..., _di[1]) -# @inline d_yi(A) = _d_yi(A, I..., _di[2]) -# @inline d_zi(A) = _d_zi(A, I..., _di[3]) - -# f = d_xi, d_yi, d_zi - -# if isvalid(ϕ, I...) -# @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) -# else -# @inbounds ∇V[I...] = zero(T) -# end -# return nothing -# end ## DEVIATORIC STRAIN RATE TENSOR @parallel_indices (i, j) function compute_strain_rate!( εxx::AbstractArray{T,2}, εyy, εxy, ∇V, Vx, Vy, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xi(A) = _d_xi(A,_dx, i, j) + d_yi(A) = _d_yi(A,_dy, i, j) + d_xa(A) = _d_xa(A,_dx, i, j) + d_ya(A) = _d_ya(A,_dy, i, j) if all((i, j) .≤ size(εxx)) ∇V_ij = ∇V[i, j] / 3.0 @@ -114,10 +83,10 @@ end @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xi(A) = _d_xi(A,_dx, i, j) + d_yi(A) = _d_yi(A,_dy, i, j) + d_xa(A) = _d_xa(A,_dx, i, j) + d_ya(A) = _d_ya(A,_dy, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -138,10 +107,10 @@ end @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, _dx, _dy, dt ) where {T} - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xi(A) = _d_xi(A,_dx, i, j) + d_yi(A) = _d_yi(A,_dy, i, j) + d_xa(A) = _d_xa(A,_dx, i, j) + d_ya(A) = _d_ya(A,_dy, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -277,10 +246,10 @@ end @parallel_indices (i, j) function compute_Res!( Rx::AbstractArray{T,2}, Ry, P, τxx, τyy, τxy, ρgx, ρgy, _dx, _dy ) where {T} - @inline d_xa(A) = _d_xa(A, i, j, _dx) - @inline d_ya(A) = _d_ya(A, i, j, _dy) - @inline d_xi(A) = _d_xi(A, i, j, _dx) - @inline d_yi(A) = _d_yi(A, i, j, _dy) + @inline d_xa(A) = _d_xa(A,_dx, i, j) + @inline d_ya(A) = _d_ya(A,_dy, i, j) + @inline d_xi(A) = _d_xi(A,_dx, i, j) + @inline d_yi(A) = _d_yi(A,_dy, i, j) @inline av_xa(A) = _av_xa(A, i, j) @inline av_ya(A) = _av_ya(A, i, j) @@ -298,10 +267,10 @@ end @parallel_indices (i, j) function compute_Res!( Rx::AbstractArray{T,2}, Ry, Vx, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ρgx, ρgy, _dx, _dy, dt ) where {T} - @inline d_xa(A) = _d_xa(A, i, j, _dx) - @inline d_ya(A) = _d_ya(A, i, j, _dy) - @inline d_xi(A) = _d_xi(A, i, j, _dx) - @inline d_yi(A) = _d_yi(A, i, j, _dy) + @inline d_xa(A) = _d_xa(A,_dx, i, j) + @inline d_ya(A) = _d_ya(A,_dy, i, j) + @inline d_xi(A) = _d_xi(A,_dx, i, j) + @inline d_yi(A) = _d_yi(A,_dy, i, j) @inline av_xa(A) = _av_xa(A, i, j) @inline av_ya(A) = _av_ya(A, i, j) diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl index 31a58b3fc..190ec4b81 100644 --- a/src/thermal_diffusion/DiffusionPT_kernels.jl +++ b/src/thermal_diffusion/DiffusionPT_kernels.jl @@ -1,17 +1,15 @@ ## 3D KERNELS -@parallel_indices (i, j, k) function compute_flux!( +@parallel_indices (I...) function compute_flux!( qTx::AbstractArray{_T,3}, qTy, qTz, qTx2, qTy2, qTz2, T, K, θr_dτ, _dx, _dy, _dz ) where {_T} - d_xi(A) = _d_xi(A, i, j, k, _dx) - d_yi(A) = _d_yi(A, i, j, k, _dy) - d_zi(A) = _d_zi(A, i, j, k, _dz) - av_xy(A) = _av_xy(A, i, j, k) - av_xz(A) = _av_xz(A, i, j, k) - av_yz(A) = _av_yz(A, i, j, k) - - I = i, j, k + d_xi(A) = _d_xi(A, _dx, I...) + d_yi(A) = _d_yi(A, _dy, I...) + d_zi(A) = _d_zi(A, _dz, I...) + av_xy(A) = _av_xy(A, I...) + av_xz(A) = _av_xz(A, I...) + av_yz(A) = _av_yz(A, I...) if all(I .≤ size(qTx)) qx = qTx2[I...] = -av_yz(K) * d_xi(T) @@ -47,16 +45,17 @@ end _dz, args, ) where {_T} - d_xi(A) = _d_xi(A, i, j, k, _dx) - d_yi(A) = _d_yi(A, i, j, k, _dy) - d_zi(A) = _d_zi(A, i, j, k, _dz) - av_xy(A) = _av_xy(A, i, j, k) - av_xz(A) = _av_xz(A, i, j, k) - av_yz(A) = _av_yz(A, i, j, k) + I = i, j, k + + @inline d_xi(A) = _d_xi(A, _dx, I...) + @inline d_yi(A) = _d_yi(A, _dy, I...) + @inline d_zi(A) = _d_zi(A, _dz, I...) + @inline av_xy(A) = _av_xy(A, I...) + @inline av_xz(A) = _av_xz(A, I...) + @inline av_yz(A) = _av_yz(A, I...) get_K(idx, args) = compute_phase(compute_conductivity, rheology, idx, args) - I = i, j, k @inbounds if all(I .≤ size(qTx)) T_ijk = (T[(I .+ 1)...] + T[i, j + 1, k + 1]) * 0.5 @@ -122,9 +121,9 @@ end _dz, ) where {_T} av(A) = _av(A, I...) - d_xa(A) = _d_xa(A, I..., _dx) - d_ya(A) = _d_ya(A, I..., _dy) - d_za(A) = _d_za(A, I..., _dz) + d_xa(A) = _d_xa(A, _dx, I...) + d_ya(A) = _d_ya(A, _dy, I...) + d_za(A) = _d_za(A, _dz, I...) I1 = I .+ 1 T[I1...] = @@ -158,9 +157,9 @@ end args, ) where {_T} av(A) = _av(A, i, j, k) - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) + d_xa(A) = _d_xa(A, _dx, i, j, k) + d_ya(A) = _d_ya(A, _dy, i, j, k) + d_za(A) = _d_za(A, _dz, i, j, k) I = i + 1, j + 1, k + 1 @@ -197,9 +196,9 @@ end _dy, _dz, ) where {_T} - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) + d_xa(A) = _d_xa(A, _dx, i, j, k) + d_ya(A) = _d_ya(A, _dy, i, j, k) + d_za(A) = _d_za(A, _dz, i, j, k) av(A) = _av(A, i, j, k) I = i + 1, j + 1, k + 1 @@ -230,9 +229,9 @@ end _dz, args, ) where {_T} - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) + d_xa(A) = _d_xa(A, _dx, i, j, k) + d_ya(A) = _d_ya(A, _dy, i, j, k) + d_za(A) = _d_za(A, _dz, i, j, k) av(A) = _av(A, i, j, k) I = i + 1, j + 1, k + 1 @@ -257,8 +256,8 @@ end ) where {_T} nx = size(θr_dτ, 1) - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 @@ -279,8 +278,8 @@ end ) where {_T} nx = size(θr_dτ, 1) - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 @@ -334,8 +333,8 @@ end ) where {_T,N,C1,C2,C3,C4} nx = size(θr_dτ, 1) - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 compute_K(phase, args) = fn_ratio(compute_conductivity, rheology, phase, args) @@ -379,8 +378,8 @@ end ) where {_T} nx, ny = size(ρCp) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) #! format: off function av(A) ( @@ -428,8 +427,8 @@ end j0 = clamp(j - 1, 1, ny) j1 = clamp(j, 1, ny) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 T_ij = T[i + 1, j + 1] @@ -462,8 +461,8 @@ end ) where {_T} nx, ny = size(ρCp) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) #! format: off function av(A) ( @@ -506,8 +505,8 @@ end j0 = clamp(j - 1, 1, ny) j1 = clamp(j, 1, ny) - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 T_ij = T[i + 1, j + 1] diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl new file mode 100644 index 000000000..1d2561a6f --- /dev/null +++ b/src/variational_stokes/MiniKernels.jl @@ -0,0 +1,14 @@ +# masked versions +for fn in (:center, :next, :left, :right, :back, :front) + @eval begin + Base.@propagate_inbounds @inline ($fn)(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) + end +end + + +# finite differences +Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy \ No newline at end of file diff --git a/src/variational_stokes/PressureKernels.jl b/src/variational_stokes/PressureKernels.jl new file mode 100644 index 000000000..dc67c3a36 --- /dev/null +++ b/src/variational_stokes/PressureKernels.jl @@ -0,0 +1,61 @@ +function compute_P!( + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio::JustPIC.PhaseRatios, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ; + ΔTc=nothing, + kwargs..., +) where {N} + ni = size(P) + @parallel (@idx ni) compute_P_kernel!( + P, P0, RP, ∇V, η, rheology, phase_ratio.center, ϕ, dt, r, θ_dτ, ΔTc + ) + return nothing +end + +@parallel_indices (I...) function compute_P!( + P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, ϕ::JustRelax.RockRatio, dt, r, θ_dτ, ::Nothing +) where {N} + if isvalid_c(ϕ, I...) + K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...])) + RP[I...], P[I...] = _compute_P!(P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ) + else + RP[I...] = P[I...] = zero(eltype(P)) + end + return nothing +end + +@parallel_indices (I...) function compute_P_kernel!( + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio::C, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ, + ΔTc, + ::Nothing, +) where {N,C<:JustRelax.CellArray} + if isvalid_c(ϕ, I...) + phase_ratio_I = phase_ratio[I...] + K = fn_ratio(get_bulk_modulus, rheology, phase_ratio_I) + α = fn_ratio(get_thermal_expansion, rheology, phase_ratio_I) + RP[I...], P[I...] = _compute_P!( + P[I...], P0[I...], ∇V[I...], ΔTc[I...], α, η[I...], K, dt, r, θ_dτ + ) + else + RP[I...] = P[I...] = zero(eltype(P)) + end + return nothing +end \ No newline at end of file diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl new file mode 100644 index 000000000..8ce08c9c4 --- /dev/null +++ b/src/variational_stokes/Stokes2D.jl @@ -0,0 +1,220 @@ +## 2D VISCO-ELASTIC STOKES SOLVER + +# backend trait +function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs) + solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) + return nothing +end + +# entry point for extensions +solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) = _solve_VS!(stokes, args...; kwargs...) + +function _solve_VS!( + stokes::JustRelax.StokesArrays, + pt_stokes, + di::NTuple{2,T}, + flow_bcs::AbstractFlowBoundaryConditions, + ρg, + phase_ratios::JustPIC.PhaseRatios, + ϕ::JustRelax.RockRatio, + rheology, + args, + dt, + igg::IGG; + viscosity_cutoff=(-Inf, Inf), + viscosity_relaxation=1e-2, + iterMax=50e3, + iterMin=1e2, + nout=500, + b_width=(4, 4, 0), + verbose=true, + kwargs..., +) where {T} + + # unpack + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + displacement2velocity!(stokes, dt, flow_bcs) + + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + + while iter ≤ iterMax + iterMin < iter && err < ϵ && break + + wtime0 += @elapsed begin + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + dt, + r, + θ_dτ, + args, + ) + + update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + update_viscosity!( + stokes, + phase_ratios, + args, + rheology, + viscosity_cutoff; + relaxation=viscosity_relaxation, + ) + + @parallel (@idx ni .+ 1) update_stresses_center_vertex!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + update_halo!(stokes.τ.xy) + + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni.+1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + _di..., + ) + # apply boundary conditions + velocity2displacement!(stokes, dt) + free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) + # end + end + + iter += 1 + + if iter % nout == 0 && iter > 1 + errs = ( + norm_mpi(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / + length(stokes.R.Rx), + norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / + length(stokes.R.Ry), + norm_mpi(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum_mpi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) + + # accumulate plastic strain tensor + @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end \ No newline at end of file diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl new file mode 100644 index 000000000..2b5610fed --- /dev/null +++ b/src/variational_stokes/StressKernels.jl @@ -0,0 +1,135 @@ +# 2D kernel +@parallel_indices (I...) function update_stresses_center_vertex!( + ε::NTuple{3, T}, # normal components @ centers; shear components @ vertices + ε_pl::NTuple{3}, # whole Voigt tensor @ centers + EII, # accumulated plastic strain rate @ centers + τ::NTuple{3}, # whole Voigt tensor @ centers + τshear_v::NTuple{1}, # shear tensor components @ vertices + τ_o::NTuple{3}, + τshear_ov::NTuple{1}, # shear tensor components @ vertices + Pr, + Pr_c, + η, + λ, + λv, + τII, + η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_center, + phase_vertex, + ϕ::JustRelax.RockRatio +) where T + τxyv = τshear_v[1] + τxyv_old = τshear_ov[1] + ni = size(Pr) + Ic = clamped_indices(ni, I...) + + if isvalid_v(ϕ, I...) + # interpolate to ith vertex + Pv_ij = av_clamped(Pr, Ic...) + εxxv_ij = av_clamped(ε[1], Ic...) + εyyv_ij = av_clamped(ε[2], Ic...) + τxxv_ij = av_clamped(τ[1], Ic...) + τyyv_ij = av_clamped(τ[2], Ic...) + τxxv_old_ij = av_clamped(τ_o[1], Ic...) + τyyv_old_ij = av_clamped(τ_o[2], Ic...) + EIIv_ij = av_clamped(EII, Ic...) + + ## vertex + phase = @inbounds phase_vertex[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(rheology, EIIv_ij, phase) + _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + Kv = fn_ratio(get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + ηv_ij = av_clamped(η, Ic...) + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = + (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * dτ_rv + dτyyv = + (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * dτ_rv + dτxyv = + ( + -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + + 2.0 * ηv_ij * ε[3][I...] + ) * dτ_rv + τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) + + # yield function @ center + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) + # stress correction @ vertex + λv[I...] = + (1.0 - relλ) * λv[I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + dQdτxy = 0.5 * (τxyv[I...] + dτxyv) / τIIv_ij + τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[I...] * dQdτxy * dτ_rv + else + # stress correction @ vertex + τxyv[I...] += dτxyv + end + else + τxyv[I...] = zero(eltype(T)) + end + + ## center + if all(I .≤ ni) + if isvalid_c(ϕ, I...) + # Material properties + phase = @inbounds phase_center[I...] + _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase(rheology, EII[I...], phase) + K = fn_ratio(get_bulk_modulus, rheology, phase) + volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ + ηij = η[I...] + dτ_r = 1.0 / (θ_dτ + ηij * _Gdt + 1.0) + + # cache strain rates for center calculations + τij, τij_o, εij = cache_tensors(τ, τ_o, ε, I...) + + # visco-elastic strain rates @ center + εij_ve = @. εij + 0.5 * τij_o * _Gdt + εII_ve = GeoParams.second_invariant(εij_ve) + # stress increments @ center + dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r + τII_ij = GeoParams.second_invariant(dτij .+ τij) + # yield function @ center + F = τII_ij - C - Pr[I...] * sinϕ + + if is_pl && !iszero(τII_ij) + # stress correction @ center + λ[I...] = + (1.0 - relλ) * λ[I...] + + relλ .* (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) + dQdτij = @. 0.5 * (τij + dτij) / τII_ij + εij_pl = λ[I...] .* dQdτij + dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r + τij = dτij .+ τij + setindex!.(τ, τij, I...) + setindex!.(ε_pl, εij_pl, I...) + τII[I...] = GeoParams.second_invariant(τij) + Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ + η_vep[I...] = 0.5 * τII_ij / εII_ve + else + # stress correction @ center + setindex!.(τ, dτij .+ τij, I...) + η_vep[I...] = ηij + τII[I...] = τII_ij + end + + Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) + else + Pr_c[I...] = zero(eltype(T)) + # τij, = cache_tensors(τ, τ_o, ε, I...) + dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) + # setindex!.(τ, dτij .+ τij, I...) + setindex!.(τ, dτij, I...) + end + end + + return nothing +end \ No newline at end of file diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl new file mode 100644 index 000000000..eda4c3474 --- /dev/null +++ b/src/variational_stokes/VelocityKernels.jl @@ -0,0 +1,51 @@ +@parallel_indices (I...) function compute_∇V!( + ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} +) where {T, N} + @inline d_xi(A) = _d_xi(A, _di[1], I...) + @inline d_yi(A) = _d_yi(A, _di[2], I...) + @inline d_zi(A) = _d_zi(A, _di[3], I...) + + f = d_xi, d_yi, d_zi + + if isvalid_c(ϕ, I...) + @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) + else + @inbounds ∇V[I...] = zero(T) + end + return nothing +end + +@parallel_indices (i, j) function compute_V!( + Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ::JustRelax.RockRatio, _dx, _dy +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + if all((i, j) .< size(Vx) .- 1) + if isvalid_vx(ϕ, i + 1, j) + Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + else + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) + end + end + + if all((i, j) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j + 1) + Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + else + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + end + end + + return nothing +end \ No newline at end of file diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl new file mode 100644 index 000000000..0e5954f65 --- /dev/null +++ b/src/variational_stokes/mask.jl @@ -0,0 +1,149 @@ + +function RockRatio(nx, ny) + ni = nx, ny + center = @zeros(ni...) + vertex = @zeros(ni.+1...) + Vx = @zeros(nx+1, ny) # no ghost nodes! + Vy = @zeros(nx, ny+1) # no ghost nodes! + + return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) +end + +function RockRatio(nx, ny, nz) + ni = nx, ny, nz + center = @zeros(ni...) + vertex = @zeros(ni.+1...) + Vx = @zeros(nx+1, ny, nz) # no ghost nodes! + Vy = @zeros(nx, ny+1, nz) # no ghost nodes! + Vz = @zeros(nx, ny, nz+1) # no ghost nodes! + yz = @zeros(nx, ny + 1, nz + 1) + xz = @zeros(nx + 1, ny, nz + 1) + xy = @zeros(nx + 1, ny + 1, nz) + + return JustRelax.RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) +end + +@inline size_c(x::JustRelax.AbstractMask) = size(x.center) +@inline size_v(x::JustRelax.AbstractMask) = size(x.vertex) +@inline size_vx(x::JustRelax.AbstractMask) = size(x.Vx) +@inline size_vy(x::JustRelax.AbstractMask) = size(x.Vy) +@inline size_vz(x::JustRelax.AbstractMask) = size(x.Vz) +@inline size_yz(x::JustRelax.AbstractMask) = size(x.yz) +@inline size_xz(x::JustRelax.AbstractMask) = size(x.xz) +@inline size_xy(x::JustRelax.AbstractMask) = size(x.xy) + +""" + update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) + +Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The rock ratio object to be updated. +- `phase_ratios`: The ratios of different phases present. +- `air_phase`: The phase representing air. +""" +function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) + nvi = size_v(ϕ) + @parallel (@idx nvi) update_rock_ratio_cv!(ϕ, phase_ratios.center, phase_ratios.vertex, air_phase) + @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) + return nothing +end + +@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = 1 - @index phase_ratio[air_phase, I...] +@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = @index phase_ratio[air_phase, I...] + +@parallel_indices (I...) function update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase) + if all(I .≤ size(ratio_center)) + ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) + end + ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) + return nothing +end + +@parallel_indices (I...) function update_rock_ratio_vel!(ϕ::JustRelax.RockRatio{T, N}) where {T, N} + # 2D + @inline av_x(A::AbstractArray{T, 2}) where T = _av_xa(A, I...) + @inline av_y(A::AbstractArray{T, 2}) where T = _av_ya(A, I...) + # 3D + @inline av_x(A::AbstractArray{T, 3}) where T = _av_yz(A, I...) + @inline av_y(A::AbstractArray{T, 3}) where T = _av_xz(A, I...) + @inline av_z(A::AbstractArray{T, 3}) where T = _av_xy(A, I...) + + all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + end + return nothing +end + +""" + isvalid_c(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.center[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_c(ϕ::JustRelax.RockRatio, i, j) + vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) + vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + v = vx * vy + return v * (ϕ.center[i, j] > 0) +end + +""" + isvalid_v(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.vertex[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_v(ϕ::JustRelax.RockRatio, i, j) + nx, ny = size(ϕ.Vx) + j_bot = max(j - 1, 1) + j0 = min(j, ny) + vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + + nx, ny = size(ϕ.Vy) + i_left = max(i - 1, 1) + i0 = min(i, nx) + vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + v = vx * vy + return v * (ϕ.vertex[i, j] > 0) +end + +""" + isvalid_vx(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vx[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i-1, j] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j+1] > 0) + cv = c * v + return cv * (ϕ.Vx[i, j] > 0) +end + +""" + isvalid_vy(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vy[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) + cv = c * v + return cv * (ϕ.Vy[i, j] > 0) +end \ No newline at end of file diff --git a/src/variational_stokes/types.jl b/src/variational_stokes/types.jl new file mode 100644 index 000000000..1b3c85372 --- /dev/null +++ b/src/variational_stokes/types.jl @@ -0,0 +1,20 @@ +abstract type AbstractMask end + +struct RockRatio{T, N} <: AbstractMask + center::T + vertex::T + Vx::T + Vy::T + Vz::Union{Nothing, T} + yz::Union{Nothing, T} + xz::Union{Nothing, T} + xy::Union{Nothing, T} + + function RockRatio(center::AbstractArray{F, N}, vertex::T, Vx::T, Vy::T, Vz::Union{Nothing, T}, yz::Union{Nothing, T}, xz::Union{Nothing, T}, xy::Union{Nothing, T}) where {F, N, T} + new{T, N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) + end +end + +RockRatio(ni::NTuple{N, Integer}) where N = RockRatio(ni...) + +Adapt.@adapt_structure RockRatio \ No newline at end of file From 1a650b85ea92fdce93a93a066f6473557d41dfad Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 13:51:54 +0200 Subject: [PATCH 06/77] remove prototype files --- MiniKernels.jl | 158 ------------ PlumeFreeSurface_2D.jl | 408 ------------------------------ RayleightTaylor_2D.jl | 403 ------------------------------ VariationalSolver.jl | 324 ------------------------ mask.jl | 546 ----------------------------------------- 5 files changed, 1839 deletions(-) delete mode 100644 MiniKernels.jl delete mode 100644 PlumeFreeSurface_2D.jl delete mode 100644 RayleightTaylor_2D.jl delete mode 100644 VariationalSolver.jl delete mode 100644 mask.jl diff --git a/MiniKernels.jl b/MiniKernels.jl deleted file mode 100644 index 63963a852..000000000 --- a/MiniKernels.jl +++ /dev/null @@ -1,158 +0,0 @@ -# unmasked versions -Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds...] -Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds.+1...] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where I<:Integer = A[i-1, j] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where I<:Integer = A[i+1, j] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j-1] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j+1] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i-1, j, k] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i+1, j, k] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j-1, k] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j+1, k] -Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k-1] -Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k+1] - -# masked versions -for fn in (:center, :next, :left, :right, :back, :front) - @eval begin - Base.@propagate_inbounds @inline ($fn)(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) - end -end -# Base.@propagate_inbounds @inline center(A::T,ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = center(A, inds...) * center(ϕ, inds...) -# Base.@propagate_inbounds @inline next(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = next(A, inds...) * next(ϕ, inds...) -# Base.@propagate_inbounds @inline left(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = left(A, inds...) * left(ϕ, inds...) -# Base.@propagate_inbounds @inline right(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = right(A, inds...) * right(ϕ, inds...) -# Base.@propagate_inbounds @inline back(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = back(A, inds...) * back(ϕ, inds...) -# Base.@propagate_inbounds @inline front(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = front(A, inds...) * front(ϕ, inds...) - -## 2D mini kernels -const T2 = AbstractArray{T,2} where {T} - -# finite differences -Base.@propagate_inbounds @inline _d_xa(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + right(A, I...)) * _dx -Base.@propagate_inbounds @inline _d_ya(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + front(A, I...)) * _dy -Base.@propagate_inbounds @inline _d_za(A::AbstractArray, _dz, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + top(A, I...)) * _dz -Base.@propagate_inbounds @inline _d_xi(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-front(A, I...) + next(A, I...)) * _dx -Base.@propagate_inbounds @inline _d_yi(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-right(A, I...) + next(A, I...)) * _dy -Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy -Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz -Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy - -# averages -Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) -Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) -Base.@propagate_inbounds @inline _av_xa(A::T, i, j) where {T<:T2} = (center(A, i, j) + right(A, i, j)) * 0.5 -Base.@propagate_inbounds @inline _av_ya(A::T, i, j) where {T<:T2} = (center(A, i, j) + front(A, i, j)) * 0.5 -Base.@propagate_inbounds @inline _av_xi(A::T, i, j) where {T<:T2} = (front(A, i, j), next(A, i, j)) * 0.5 -Base.@propagate_inbounds @inline _av_yi(A::T, i, j) where {T<:T2} = (right(A, i, j), next(A, i, j)) * 0.5 -# harmonic averages -Base.@propagate_inbounds @inline function _harm(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2)) -end -Base.@propagate_inbounds @inline function _harm_a(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1)) -end -Base.@propagate_inbounds @inline function _harm_xa(A::T, i, j) where {T<:T2} - return eltype(A)(2) * (inv(right(A, i, j)) + inv(center(A, i, j))) -end -Base.@propagate_inbounds @inline function _harm_ya(A::T, i, j) where {T<:T2} - return eltype(A)(2) * (inv(front(A, i, j)) + inv(center(A, i, j))) -end -#others -Base.@propagate_inbounds @inline function _gather(A::T, i, j) where {T<:T2} - return center(A, i, j), right(A, i, j), front(A, i, j), next(A, i, j) -end - -## 3D mini kernels -const T3 = AbstractArray{T,3} where {T} - -# finite differences -Base.@propagate_inbounds @inline function _d_xi(A::T, i, j, k, _dx) where {T<:T3} - return (-A[i, j + 1, k + 1] + next(A, i, j, k)) * _dx -end -Base.@propagate_inbounds @inline function _d_yi(A::T, i, j, k, _dy) where {T<:T3} - return (-A[i + 1, j, k + 1] + next(A, i, j, k)) * _dy -end -Base.@propagate_inbounds @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3} - return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz -end -# averages -Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) -Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + right(A, i, j, k)) -Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + front(A, i, j, k)) -Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + top(A, i, j, k)) -Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) -Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) -Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) -Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) -Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) -Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) -# harmonic averages -@inline function _harm_x(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(right(A, i, j, k))) -end -@inline function _harm_y(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(front(A, i, j, k))) -end -@inline function _harm_z(A::T, i, j, k) where {T<:T3} - return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(top(A, i, j, k))) -end -@inline function _harm_xy(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k)) -end -@inline function _harm_xz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:j, k:(k + 1))) -end -@inline function _harm_yz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, j:(j + 1), k:(k + 1))) -end -@inline function _harm_xyi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i - 1):i, (j - 1):j, k:k)) -end -@inline function _harm_xzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i - 1):i, j:j, (k - 1):k)) -end -@inline function _harm_yzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, (j - 1):j, (k - 1):k)) -end - -# others -Base.@propagate_inbounds @inline function _gather_yz(A::T, i, j, k) where {T<:T3} - return center(A, i, j, k), front(A, i, j, k), top(A, i, j, k), A[i, j + 1, k + 1] -end -Base.@propagate_inbounds @inline function _gather_xz(A::T, i, j, k) where {T<:T3} - return center(A, i, j, k), right(A, i, j, k), top(A, i, j, k), A[i + 1, j, k + 1] -end -Base.@propagate_inbounds @inline function _gather_xy(A::T, i, j, k) where {T<:T3} - return center(A, i, j, k), right(A, i, j, k), front(A, i, j, k), A[i + 1, j + 1, k] -end -@inline _current(A::T, i, j, k) where {T<:T3} = center(A, i, j, k) - -## Because mysum(::generator) does not work inside CUDA kernels... -@inline mysum(A, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ranges...) - -@inline function mysum(f::F, A, ranges_i) where {F<:Function} - s = 0.0 - for i in ranges_i - s += f(A[i]) - end - return s -end - -@inline function mysum(f::F, A, ranges_i, ranges_j) where {F<:Function} - s = 0.0 - for i in ranges_i, j in ranges_j - s += f(center(A, i, j)) - end - return s -end - -@inline function mysum(f::F, A, ranges_i, ranges_j, ranges_k) where {F<:Function} - s = 0.0 - for i in ranges_i, j in ranges_j, k in ranges_k - s += f(center(A, i, j, k)) - end - return s -end diff --git a/PlumeFreeSurface_2D.jl b/PlumeFreeSurface_2D.jl deleted file mode 100644 index 55a87880b..000000000 --- a/PlumeFreeSurface_2D.jl +++ /dev/null @@ -1,408 +0,0 @@ -using Printf -using JustRelax, JustRelax.JustRelax2D -import JustRelax.JustRelax2D as JR -const backend_JR = CPUBackend - -using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) - -# Load script dependencies -using LinearAlgebra, GeoParams, GLMakie - -include("mask.jl") -include("MiniKernels.jl") - -# Velocity helper grids for the particle advection -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - @parallel f_x(A, B) -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 - -function init_phases!(phases, particles) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index) - r=100e3 - - @inbounds for ip in cellaxes(phases) - # quick escape - @index(index[ip, i, j]) == 0 && continue - - x = @index px[ip, i, j] - depth = -(@index py[ip, i, j]) - @index phases[ip, i, j] = 2.0 - - if 0e0 ≤ depth ≤ 100e3 - @index phases[ip, i, j] = 1.0 - - else - @index phases[ip, i, j] = 2.0 - - if ((x - 250e3)^2 + (depth - 250e3)^2 ≤ r^2) - @index phases[ip, i, j] = 3.0 - end - end - - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) -end -## END OF HELPER FUNCTION ------------------------------------------------------------ - -# (Path)/folder where output data and figures are stored -n = 101 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 15e3 # thickness of sticky air layer - ly = 400e3 + thick_air # domain length in y - lx = 500e3 # 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) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - # Name = "Air", - # SetMaterialParams(; - # Phase = 1, - # Density = ConstantDensity(; ρ=1e1), - # CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), - # Gravity = ConstantGravity(; g=9.81), - # ), - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Mantle", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Plume", - SetMaterialParams(; - Phase = 3, - Density = ConstantDensity(; ρ=3.2e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), - Gravity = ConstantGravity(; g=9.81), - ) - ) - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 60, 80, 40 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi, di, ni - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - init_phases!(pPhases, particles) - phase_ratios = PhaseRatios(backend, length(rheology), ni) - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # ---------------------------------------------------- - - # RockRatios - air_phase = 1 - ϕ = RockRatio(ni...) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.95 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) - - # Boundary conditions - flow_bcs = VelocityBoundaryConditions(; - free_slip = (left = true, right = true, top = true, bot = true), - free_surface = true - ) - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "FreeSurfacePlume" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - iterMax = 15e3 - nout = 1e3 - viscosity_cutoff = (-Inf, Inf) - free_surface = false - ητ = @zeros(ni...) - while it < 20 - - ## variational solver - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - relλ = 0.2 - θ = deepcopy(stokes.P) - λ = @zeros(ni...) - λv = @zeros(ni .+ 1...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - # Vx_on_Vy = @zeros(size(stokes.V.Vy)) - - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - - while iter ≤ iterMax - err < ϵ && break - # for _ in 1:100 - JR.compute_maxloc!(ητ, η; window=(1, 1)) - # update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - @parallel (@idx ni) compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ - ) - - JR.update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) JR.compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - # if rem(iter, nout) == 0 - # @copy η0 η - # end - # if do_visc - # update_viscosity!( - # stokes, - # phase_ratios, - # args, - # rheology, - # viscosity_cutoff; - # relaxation=viscosity_relaxation, - # ) - # end - - @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - # update_halo!(stokes.τ.xy) - - # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( - # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] - # ) - - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni.+1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - # ϕ.Vx, - # ϕ.Vy, - _di..., - ) - # apply boundary conditions - # velocity2displacement!(stokes, dt) - # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - # end - # f,ax,h=heatmap(stokes.V.Vy) - # # f,ax,h=heatmap(stokes.V.Vx) - # Colorbar(f[1,2], h, label="Vy"); f - # update_halo!(@velocity(stokes)...) - # end - - iter += 1 - - if iter % nout == 0 && iter > 1 - # er_η = norm_mpi(@.(log10(η) - log10(η0))) - # er_η < 1e-3 && (do_visc = false) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), - norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - # end - isnan(err) && error("NaN(s)") - isinf(err) && error("Inf(s)") - end - end - - dt = compute_dt(stokes, di) / 2 - # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - @show it += 1 - t += dt - - if it == 1 || rem(it, 1) == 0 - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 5 - fig = Figure(size = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - arrows!( - ax, - xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :red, - ) - fig - save(joinpath(figdir, "$(it).png"), fig) - - end - end - return nothing -end -# ## END OF MAIN SCRIPT ---------------------------------------------------------------- -main(igg, nx, ny) - -# # # (Path)/folder where output data and figures are stored -# # n = 100 -# # nx = n -# # ny = n -# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid -# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -# # else -# # igg -# # end - diff --git a/RayleightTaylor_2D.jl b/RayleightTaylor_2D.jl deleted file mode 100644 index d12ef8921..000000000 --- a/RayleightTaylor_2D.jl +++ /dev/null @@ -1,403 +0,0 @@ -using Printf -using JustRelax, JustRelax.JustRelax2D -import JustRelax.JustRelax2D as JR -const backend_JR = CPUBackend - -using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) - -# Load script dependencies -using LinearAlgebra, GeoParams, GLMakie - -include("mask.jl") -include("MiniKernels.jl") - -# Velocity helper grids for the particle advection -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - @parallel f_x(A, B) -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 - -function init_phases!(phases, particles, A) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) - - f(x, A, λ) = A * sin(π * x / λ) - - @inbounds for ip in cellaxes(phases) - # quick escape - @index(index[ip, i, j]) == 0 && continue - - x = @index px[ip, i, j] - depth = -(@index py[ip, i, j]) - @index phases[ip, i, j] = 2.0 - - if 0e0 ≤ depth ≤ 100e3 - @index phases[ip, i, j] = 1.0 - - elseif depth > (-f(x, A, 500e3) + (200e3 - A)) - @index phases[ip, i, j] = 3.0 - - end - - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, A) -end -## END OF HELPER FUNCTION ------------------------------------------------------------ - -# (Path)/folder where output data and figures are stored -n = 101 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 100e3 # thickness of sticky air layer - ly = 500e3 + thick_air # domain length in y - lx = 500e3 # 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) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - # Name = "Air", - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=1e0), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e16),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Crust", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Mantle", - SetMaterialParams(; - Phase = 3, - Density = ConstantDensity(; ρ=3.2e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), - Gravity = ConstantGravity(; g=9.81), - ) - ) - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 60, 80, 40 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi, di, ni - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - A = 5e3 # Amplitude of the anomaly - phase_ratios = PhaseRatios(backend, length(rheology), ni) - init_phases!(pPhases, particles, A) - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # ---------------------------------------------------- - - # RockRatios - air_phase = 1 - ϕ = RockRatio(ni...) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.7 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_ρg!(ρg[2], phase_ratios, rheology, args) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) - - # Boundary conditions - flow_bcs = VelocityBoundaryConditions(; - free_slip = (left = true, right = true, top = true, bot = false), - no_slip = (left = false, right = false, top = false, bot = true), - free_surface = true, - ) - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "RayleighTaylor2D" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = 10e3 * (3600 * 24 * 365.25) - dt_max = 50e3 * (3600 * 24 * 365.25) - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - iterMax = 15e3 - nout = 1e3 - viscosity_cutoff = (-Inf, Inf) - free_surface = false - ητ = @zeros(ni...) - while it < 20 - - ## variational solver - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - relλ = 0.2 - θ = deepcopy(stokes.P) - λ = @zeros(ni...) - λv = @zeros(ni .+ 1...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - # Vx_on_Vy = @zeros(size(stokes.V.Vy)) - - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - - while iter ≤ iterMax - err < ϵ && break - # for _ in 1:100 - JR.compute_maxloc!(ητ, η; window=(1, 1)) - # update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - @parallel (@idx ni) compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ - ) - - JR.update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) JR.compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - # if rem(iter, nout) == 0 - # @copy η0 η - # end - # if do_visc - # update_viscosity!( - # stokes, - # phase_ratios, - # args, - # rheology, - # viscosity_cutoff; - # relaxation=viscosity_relaxation, - # ) - # end - - @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - # update_halo!(stokes.τ.xy) - - # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( - # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] - # ) - - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni.+1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - # ϕ.Vx, - # ϕ.Vy, - _di..., - ) - # apply boundary conditions - # velocity2displacement!(stokes, dt) - # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - # end - # f,ax,h=heatmap(stokes.V.Vy) - # # f,ax,h=heatmap(stokes.V.Vx) - # Colorbar(f[1,2], h, label="Vy"); f - # update_halo!(@velocity(stokes)...) - # end - - iter += 1 - - if iter % nout == 0 && iter > 1 - # er_η = norm_mpi(@.(log10(η) - log10(η0))) - # er_η < 1e-3 && (do_visc = false) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), - norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - # end - isnan(err) && error("NaN(s)") - isinf(err) && error("Inf(s)") - end - end - - dt = compute_dt(stokes, di) / 2 - # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - @show it += 1 - t += dt - - if it == 1 || rem(it, 1) == 0 - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 5 - fig = Figure(size = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - arrows!( - ax, - xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :red, - ) - fig - save(joinpath(figdir, "$(it).png"), fig) - - end - end - return nothing -end -# ## END OF MAIN SCRIPT ---------------------------------------------------------------- -main(igg, nx, ny) - -# # # (Path)/folder where output data and figures are stored -# # n = 100 -# # nx = n -# # ny = n -# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid -# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -# # else -# # igg -# # end - diff --git a/VariationalSolver.jl b/VariationalSolver.jl deleted file mode 100644 index b8441c291..000000000 --- a/VariationalSolver.jl +++ /dev/null @@ -1,324 +0,0 @@ -# unpack - -_di = inv.(di) -(; ϵ, r, θ_dτ, ηdτ) = pt_stokes -(; η, η_vep) = stokes.viscosity -ni = size(stokes.P) - -# ~preconditioner -ητ = deepcopy(η) -# @hide_communication b_width begin # communication/computation overlap -JR.compute_maxloc!(ητ, η; window=(1, 1)) -# update_halo!(ητ) -# end - -iterMax = 50e3 -nout = 1e3 -viscosity_cutoff = (-Inf, Inf) -free_surface = false - -# errors -err = 2 * ϵ -iter = 0 -err_evo1 = Float64[] -err_evo2 = Float64[] -norm_Rx = Float64[] -norm_Ry = Float64[] -norm_∇V = Float64[] -sizehint!(norm_Rx, Int(iterMax)) -sizehint!(norm_Ry, Int(iterMax)) -sizehint!(norm_∇V, Int(iterMax)) -sizehint!(err_evo1, Int(iterMax)) -sizehint!(err_evo2, Int(iterMax)) - -# solver loop -@copy stokes.P0 stokes.P -wtime0 = 0.0 -relλ = 0.2 -θ = deepcopy(stokes.P) -λ = @zeros(ni...) -λv = @zeros(ni .+ 1...) -η0 = deepcopy(η) -do_visc = true - -for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 -end -Vx_on_Vy = @zeros(size(stokes.V.Vy)) - -# compute buoyancy forces and viscosity -compute_ρg!(ρg[end], phase_ratios, rheology, args) -compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - -while iter ≤ iterMax - iterMin < iter && err < ϵ && break - - wtime0 += @elapsed begin - JR.compute_maxloc!(ητ, η; window=(1, 1)) - # update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - @parallel (@idx ni) compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ - ) - - JR.update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) JR.compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - # if rem(iter, nout) == 0 - # @copy η0 η - # end - # if do_visc - # update_viscosity!( - # stokes, - # phase_ratios, - # args, - # rheology, - # viscosity_cutoff; - # relaxation=viscosity_relaxation, - # ) - # end - - @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - # update_halo!(stokes.τ.xy) - - @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( - Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] - ) - - # @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - Vx_on_Vy, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ.Vx, - ϕ.Vy, - _di..., - dt * free_surface, - ) - # apply boundary conditions - # velocity2displacement!(stokes, dt) - free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - # update_halo!(@velocity(stokes)...) - # end - end - - iter += 1 - - if iter % nout == 0 && iter > 1 - # er_η = norm_mpi(@.(log10(η) - log10(η0))) - # er_η < 1e-3 && (do_visc = false) - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, - stokes.R.Ry, - @velocity(stokes)..., - Vx_on_Vy, - stokes.P, - @stress(stokes)..., - ρg..., - ϕ.Vx, - ϕ.Vy, - _di..., - dt * free_surface, - ) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), - norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximumi(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end -end - -# # compute vorticity -# @parallel (@idx ni .+ 1) compute_vorticity!( -# stokes.ω.xy, @velocity(stokes)..., inv.(di)... -# ) - -# # accumulate plastic strain tensor -# @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) - - -update_rock_ratio!(ϕ, phase_ratios, air_phase) - -heatmap(stokes.V.Vx) -heatmap(stokes.V.Vy) -heatmap(stokes.P) -heatmap(stokes.τ.xx) -heatmap(stokes.τ.yy) -heatmap(stokes.τ.xy) - - -heatmap(log10.(abs.(stokes.R.Rx))) -heatmap(log10.(abs.(stokes.R.Ry))) -heatmap(log10.(abs.(stokes.V.Vy))) -heatmap(ϕ.Vy) -heatmap(ϕ.Vx) -heatmap(ϕ.vertex) -heatmap(ϕ.center) - - -f,ax,h = heatmap(ϕ.vertex) -Colorbar(f[1,2], h);f - -f,ax,h = heatmap(ϕ.center) -Colorbar(f[1,2], h);f - -@parallel compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ.Vx, - ϕ.Vy, - _di..., -) - -extrema(stokes.R.Ry) -extrema(stokes.V.Vy) - -f,ax,h = heatmap(stokes.V.Vy) -Colorbar(f[1,2], h, label="Vy") - --d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) - -d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) -d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) -_dx, _dy = inv.(di) - -a = [_d_yi(stokes.τ.xy, ϕ.vertex, _dy, i, j) for i in 1:nx, j in 1:ny-1] -heatmap(a) - -b = [_d_ya(stokes.P, ϕ.center, _dy, i, j) for i in 1:nx, j in 1:ny-1] -f,ax,h = heatmap(b) -Colorbar(f[1,2], h, label="Vy") - -c = [_d_ya(stokes.τ.yy, ϕ.center, _dy, i, j) for i in 1:nx, j in 1:ny-1] -heatmap(c) -d = @. a-b+c - -f,ax,h=heatmap(d) -Colorbar(f[1,2], h, label="Vy") - -lines(stokes.P[1,:]) -lines(stokes.V.Vy[1,:]) - -xc = [x for x in xci[1], y in xci[2]] -yc = [y for x in xci[1], y in xci[2]] - -v = [isvalid_c(ϕ, i, j) for i in 1:nx, j in 1:ny] -heatmap(xci..., v) - -heatmap(xci..., ϕ.center) -h = @. -√(0.5^2 - xci[1]^2) -lines!(xci[1], h, color=:red) - -v = [isvalid_v(ϕ, i, j) for i in 1:nx, j in 1:ny] -heatmap(xvi..., v) - -heatmap(xvi..., ϕ.vertex) -h = @. -√(0.5^2 - xvi[1]^2) -lines!(xvi[1], h, color=:red) - -vx = [isvalid_vx(ϕ, i+1, j) for i in 1:nx-1, j in 1:ny] -heatmap(xvi[1], xci[2], vx) -h = @. -√(0.5^2 - xci[1]^2) -lines!(xci[1], h, color=:red) - -vy = [isvalid_vy(ϕ, i, j+1) for i in 1:nx, j in 1:ny-1] -heatmap(vy) - -heatmap( - xci[1].*1e-3, xci[2].*1e-3, - Array([argmax(p) for p in phase_ratios.center]), - colormap = :grayC) -px, py = particles - -heatmap( - xvi[1].*1e-3, xvi[2].*1e-3, - Array([argmax(p) for p in phase_ratios.vertex]), - colormap = :grayC) - - -heatmap(xvi..., ϕ.vertex) -# heatmap(xci..., ϕ.center) -heatmap(xci..., P) -scatter!(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle, colormap=:grayC) -scatter(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle, colormap=:grayC) - -P = copy(stokes.P) -P = copy(ϕ.vertex) -P[iszero.(P)] .= NaN -heatmap(xci..., P) - -id = particles.index.data[:] -px, py = particles.coords -scatter!(px.data[:], py.data[:], color=pPhases.data[:], markersize=5.0, marker=:circle) -scatter(px.data[:][id], py.data[:][id], color=pPhases.data[:][id], markersize=5.0, marker=:circle) diff --git a/mask.jl b/mask.jl deleted file mode 100644 index edc040ac2..000000000 --- a/mask.jl +++ /dev/null @@ -1,546 +0,0 @@ -# using JustRelax, JustRelax.JustRelax2D -using CellArrays -# using Test -# using ParallelStencil -# @init_parallel_stencil(Threads, Float64, 2) -abstract type AbstractMask end - -struct RockRatio{T, N} <: AbstractMask - center::T - vertex::T - Vx::T - Vy::T - Vz::Union{Nothing, T} - yz::Union{Nothing, T} - xz::Union{Nothing, T} - xy::Union{Nothing, T} - - function RockRatio(center::AbstractArray{F, N}, vertex::T, Vx::T, Vy::T, Vz::Union{Nothing, T}, yz::Union{Nothing, T}, xz::Union{Nothing, T}, xy::Union{Nothing, T}) where {F, N, T} - new{T, N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) - end -end - -RockRatio(ni::NTuple{N, Integer}) where N = RockRatio(ni...) - -function RockRatio(nx, ny) - ni = nx, ny - center = @zeros(ni...) - vertex = @zeros(ni.+1...) - Vx = @zeros(nx+1, ny) # no ghost nodes! - Vy = @zeros(nx, ny+1) # no ghost nodes! - - return RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) -end - -function RockRatio(nx, ny, nz) - ni = nx, ny, nz - center = @zeros(ni...) - vertex = @zeros(ni.+1...) - Vx = @zeros(nx+1, ny, nz) # no ghost nodes! - Vy = @zeros(nx, ny+1, nz) # no ghost nodes! - Vz = @zeros(nx, ny, nz+1) # no ghost nodes! - yz = @zeros(nx, ny + 1, nz + 1) - xz = @zeros(nx + 1, ny, nz + 1) - xy = @zeros(nx + 1, ny + 1, nz) - - return RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) -end - -@inline size_c(x::AbstractMask) = size(x.center) -@inline size_v(x::AbstractMask) = size(x.vertex) -@inline size_vx(x::AbstractMask) = size(x.Vx) -@inline size_vy(x::AbstractMask) = size(x.Vy) -@inline size_vz(x::AbstractMask) = size(x.Vz) -@inline size_yz(x::AbstractMask) = size(x.yz) -@inline size_xz(x::AbstractMask) = size(x.xz) -@inline size_xy(x::AbstractMask) = size(x.xy) - -@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = 1 - @index phase_ratio[air_phase, I...] -@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = @index phase_ratio[air_phase, I...] - -@parallel_indices (I...) function update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase) - if all(I .≤ size(ratio_center)) - ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) - end - ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) - return nothing -end - -@parallel_indices (I...) function update_rock_ratio_vel!(ϕ::RockRatio{T, N}) where {T, N} - # 2D - @inline av_x(A::AbstractArray{T, 2}) where T = _av_xa(A, I...) - @inline av_y(A::AbstractArray{T, 2}) where T = _av_ya(A, I...) - # 3D - @inline av_x(A::AbstractArray{T, 3}) where T = _av_yz(A, I...) - @inline av_y(A::AbstractArray{T, 3}) where T = _av_xz(A, I...) - @inline av_z(A::AbstractArray{T, 3}) where T = _av_xy(A, I...) - - all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - end - return nothing -end - -function update_rock_ratio!(ϕ::RockRatio, phase_ratios, air_phase) - nvi = size_v(ϕ) - - @parallel (@idx nvi) update_rock_ratio_cv!(ϕ, phase_ratios.center, phase_ratios.vertex, air_phase) - @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) - - return nothing -end - -# true means that is a nullspace and needs to be removed -# @inline Base.@propagate_inbounds isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} = iszero(A[I...]) - -# nullspace: check whether any point of the stencil is 0, if so, eliminate equation -function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} - v = true - Base.@nexprs 2 j -> begin - Base.@nexprs 2 i -> begin - @inline - ii = clamp(I[1] + 2 * i - 3, 1, size(A, 1)) - jj = clamp(I[2] + 2 * j - 3, 1, size(A, 2)) - v *= if N == 3 - A[ii, jj, I[3]] > 0 - else - A[ii, jj] > 0 - end - end - end - if N === 3 - Base.@nexprs 2 k -> begin - kk = clamp(I[3] + 2 * k - 3, 1, size(A, 3)) - v *= A[ii, jj, kk] > 0 - end - end - return v * (A[I...] > 0) -end - -function isvalid_c(ϕ::RockRatio, i, j) - vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) - vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) - v = vx * vy - # return v * (ϕ.center[i, j] > 0) - return true -end - -function isvalid_v(ϕ::RockRatio, i, j) - nx, ny = size(ϕ.Vx) - j_bot = max(j - 1, 1) - j0 = min(j, ny) - vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - - nx, ny = size(ϕ.Vy) - i_left = max(i - 1, 1) - i0 = min(i, nx) - vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) - v = vx * vy - # return v * (ϕ.vertex[i, j] > 0) - return true -end - -function isvalid_vx(ϕ::RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i-1, j] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j+1] > 0) - cv = c * v - return cv * (ϕ.Vx[i, j] > 0) - # c = (ϕ.center[i, j] > 0) || (ϕ.center[i-1, j] > 0) - # v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j+1] > 0) - # cv = c || v - # return cv || (ϕ.Vx[i, j] > 0) -end - -function isvalid_vy(ϕ::RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) - cv = c * v - return cv * (ϕ.Vy[i, j] > 0) - # c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) - # v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) - # cv = c || v - # return cv || (ϕ.Vy[i, j] > 0) -end - -# function isvalid(A::T, I::Vararg{Integer, N}) where {N, T<:AbstractArray} -# v = true -# Base.@nexprs 2 j -> begin -# Base.@nexprs 2 i -> begin -# @inline -# ii = clamp(I[1] + 2 * i - 3, 1, size(A, 1)) -# jj = clamp(I[2] + 2 * j - 3, 1, size(A, 2)) -# v *= if N == 3 -# A[ii, jj, I[3]] > 0 -# else -# A[ii, jj] > 0 -# end -# end -# end -# if N === 3 -# Base.@nexprs 2 k -> begin -# kk = clamp(I[3] + 2 * k - 3, 1, size(A, 3)) -# v *= A[ii, jj, kk] > 0 -# end -# end -# return v * (A[I...] > 0) -# end - -@parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::RockRatio, _di::NTuple{N} -) where {T, N} - @inline d_xi(A) = _d_xi(A, _di[1], I...) - @inline d_yi(A) = _d_yi(A, _di[2], I...) - @inline d_zi(A) = _d_zi(A, _di[3], I...) - - f = d_xi, d_yi, d_zi - - if isvalid_c(ϕ, I...) - @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) - else - @inbounds ∇V[I...] = zero(T) - end - return nothing -end - -@parallel_indices (I...) function compute_P!( - P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, ϕ::RockRatio, dt, r, θ_dτ -) where {N} - if isvalid_c(ϕ, I...) - K = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, @cell(phase_ratio[I...])) - RP[I...], P[I...] = JustRelax2D._compute_P!(P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ) - else - RP[I...] = P[I...] = zero(eltype(P)) - end - return nothing -end - -# 2D kernel -@parallel_indices (I...) function update_stresses_center_vertex_ps!( - ε::NTuple{3, T}, # normal components @ centers; shear components @ vertices - ε_pl::NTuple{3}, # whole Voigt tensor @ centers - EII, # accumulated plastic strain rate @ centers - τ::NTuple{3}, # whole Voigt tensor @ centers - τshear_v::NTuple{1}, # shear tensor components @ vertices - τ_o::NTuple{3}, - τshear_ov::NTuple{1}, # shear tensor components @ vertices - Pr, - Pr_c, - η, - λ, - λv, - τII, - η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_center, - phase_vertex, - ϕ -) where T - τxyv = τshear_v[1] - τxyv_old = τshear_ov[1] - ni = size(Pr) - Ic = JustRelax2D.clamped_indices(ni, I...) - - if isvalid_v(ϕ, I...) - # interpolate to ith vertex - Pv_ij = JustRelax2D.av_clamped(Pr, Ic...) - εxxv_ij = JustRelax2D.av_clamped(ε[1], Ic...) - εyyv_ij = JustRelax2D.av_clamped(ε[2], Ic...) - τxxv_ij = JustRelax2D.av_clamped(τ[1], Ic...) - τyyv_ij = JustRelax2D.av_clamped(τ[2], Ic...) - τxxv_old_ij = JustRelax2D.av_clamped(τ_o[1], Ic...) - τyyv_old_ij = JustRelax2D.av_clamped(τ_o[2], Ic...) - EIIv_ij = JustRelax2D.av_clamped(EII, Ic...) - - ## vertex - phase = @inbounds phase_vertex[I...] - is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = JustRelax2D.plastic_params_phase(rheology, EIIv_ij, phase) - _Gvdt = inv(JustRelax2D.fn_ratio(JustRelax2D.get_shear_modulus, rheology, phase) * dt) - Kv = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, phase) - volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ - ηv_ij = JustRelax2D.av_clamped(η, Ic...) - dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) - - # stress increments @ vertex - dτxxv = - (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * dτ_rv - dτyyv = - (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * dτ_rv - dτxyv = - ( - -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + - 2.0 * ηv_ij * ε[3][I...] - ) * dτ_rv - τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) - - # yield function @ center - Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) - # stress correction @ vertex - λv[I...] = - (1.0 - relλ) * λv[I...] + - relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) - dQdτxy = 0.5 * (τxyv[I...] + dτxyv) / τIIv_ij - τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[I...] * dQdτxy * dτ_rv - else - # stress correction @ vertex - τxyv[I...] += dτxyv - end - else - τxyv[I...] = zero(eltype(T)) - end - - ## center - if all(I .≤ ni) - if isvalid_c(ϕ, I...) - # Material properties - phase = @inbounds phase_center[I...] - _Gdt = inv(JustRelax2D.fn_ratio(JustRelax2D.get_shear_modulus, rheology, phase) * dt) - is_pl, C, sinϕ, cosϕ, sinψ, η_reg = JustRelax2D.plastic_params_phase(rheology, EII[I...], phase) - K = JustRelax2D.fn_ratio(JustRelax2D.get_bulk_modulus, rheology, phase) - volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ - ηij = η[I...] - dτ_r = 1.0 / (θ_dτ + ηij * _Gdt + 1.0) - - # cache strain rates for center calculations - τij, τij_o, εij = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) - - # visco-elastic strain rates @ center - εij_ve = @. εij + 0.5 * τij_o * _Gdt - εII_ve = GeoParams.second_invariant(εij_ve) - # stress increments @ center - dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r - τII_ij = GeoParams.second_invariant(dτij .+ τij) - # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ - - if is_pl && !iszero(τII_ij) - # stress correction @ center - λ[I...] = - (1.0 - relλ) * λ[I...] + - relλ .* (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) - dQdτij = @. 0.5 * (τij + dτij) / τII_ij - εij_pl = λ[I...] .* dQdτij - dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r - τij = dτij .+ τij - setindex!.(τ, τij, I...) - setindex!.(ε_pl, εij_pl, I...) - τII[I...] = GeoParams.second_invariant(τij) - Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ - η_vep[I...] = 0.5 * τII_ij / εII_ve - else - # stress correction @ center - setindex!.(τ, dτij .+ τij, I...) - η_vep[I...] = ηij - τII[I...] = τII_ij - end - - Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) - else - Pr_c[I...] = zero(eltype(T)) - # τij, = JustRelax2D.cache_tensors(τ, τ_o, ε, I...) - dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) - # setindex!.(τ, dτij .+ τij, I...) - setindex!.(τ, dτij, I...) - end - end - - return nothing -end - - -# # with free surface stabilization -@parallel_indices (i, j) function compute_V!( - Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ::RockRatio, _dx, _dy -) where {T} - d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) - d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) - d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) - av_xa(A) = _av_xa(A, i, j) - av_ya(A) = _av_ya(A, i, j) - harm_xa(A) = _av_xa(A, i, j) - harm_ya(A) = _av_ya(A, i, j) - - if all((i, j) .< size(Vx) .- 1) - if isvalid_vx(ϕ, i + 1, j) - Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) - Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) - else - Rx[i, j] = zero(T) - Vx[i + 1, j + 1] = zero(T) - end - end - - if all((i, j) .< size(Vy) .- 1) - if isvalid_vy(ϕ, i, j + 1) - Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) - Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) - else - Ry[i, j] = zero(T) - Vy[i + 1, j + 1] = zero(T) - end - end - - return nothing -end - -@parallel_indices (i, j) function compute_V!( - Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy -) where {T} - d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) - d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) - d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) - av_xa(A) = _av_xa(A, i, j) - av_ya(A) = _av_ya(A, i, j) - harm_xa(A) = _av_xa(A, i, j) - harm_ya(A) = _av_ya(A, i, j) - - if all((i, j) .< size(Vx) .- 1) - if iszero(ϕ_Vx[i + 1, j]) - Rx[i, j] = zero(T) - Vx[i + 1, j + 1] = zero(T) - else - Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) - Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) - end - end - - if all((i, j) .< size(Vy) .- 1) - if iszero(ϕ_Vy[i, j + 1]) - Ry[i, j] = zero(T) - Vy[i + 1, j + 1] = zero(T) - else - Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) - Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) - end - end - - return nothing -end - -# @parallel_indices (i, j) function compute_V!( -# Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ_Vx, ϕ_Vy, _dx, _dy, dt -# ) where {T} -# d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) -# d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) -# d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) -# d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) -# av_xa(A) = _av_xa(A, i, j) -# av_ya(A) = _av_ya(A, i, j) -# harm_xa(A) = _av_xa(A, i, j) -# harm_ya(A) = _av_ya(A, i, j) - -# nx, ny = size(ρgy) - -# if all((i, j) .< size(Vx) .- 1) -# dVx = if iszero(ϕ_Vx[i + 1, j]) -# zero(T) -# else -# (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) * ηdτ / av_xa(ητ) -# end -# Vx[i + 1, j + 1] += dVx -# end - -# if all((i, j) .< size(Vy) .- 1) -# dVy = if iszero(ϕ_Vy[i, j + 1]) -# zero(T) -# else -# ρg_correction = if iszero(dt) -# zero(dt) -# else -# θ = 1 -# # Interpolated Vx into Vy node (includes density gradient) -# Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] -# # Vertical velocity -# Vyᵢⱼ = Vy[i + 1, j + 1] -# # Get necessary buoyancy forces -# j_N = min(j + 1, ny) -# ρg_S = ρgy[i, j] -# ρg_N = ρgy[i, j_N] -# # Spatial derivatives -# ∂ρg∂y = (ρg_N - ρg_S) * _dy -# # correction term -# (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt -# end -# (-d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction) * ηdτ / av_ya(ητ) -# end -# Vy[i + 1, j + 1] += dVy -# end - -# return nothing -# end - -@parallel_indices (i, j) function compute_Res!( - Rx::AbstractArray{T,2}, Ry, Vx, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ρgx, ρgy, ϕ_Vx, ϕ_Vy, _dx, _dy, dt -) where {T} - d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) - d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) - d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) - @inline av_xa(A) = _av_xa(A, i, j) - @inline av_ya(A) = _av_ya(A, i, j) - - nx, ny = size(ρgy) - @inbounds begin - if all((i, j) .≤ size(Rx)) - Rx[i, j] = if iszero(ϕ_Vx[i + 1, j]) - zero(T) - else - -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx) - end - end - - if all((i, j) .≤ size(Ry)) - Ry[i, j] = if iszero(ϕ_Vy[i, j + 1]) - zero(T) - else - θ = 1.0 - # Interpolated Vx into Vy node (includes density gradient) - Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] - # Vertical velocity - Vyᵢⱼ = Vy[i + 1, j + 1] - # Get necessary buoyancy forces - j_N = min(j + 1, ny) - ρg_S = ρgy[i, j] - ρg_N = ρgy[i, j_N] - # Spatial derivatives - ∂ρg∂y = (ρg_N - ρg_S) * _dy - # correction term - ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt - -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + ρg_correction - end - end - end - - return nothing -end - -# #### TESTING GROUNDS -# N = 4 -# ni = n, n = (N,N) -# ϕ = RockRatio(n, n) -# @test RockRatio(n, n) isa RockRatio - -# @test RockRatio(ni) isa RockRatio -# @test size_c(ϕ) === ni -# @test size_v(ϕ) === ni.+1 -# @test size_vx(ϕ) === (N+1, N) -# @test size_vy(ϕ) === (N, N+1) - -# A = [ -# 0 1 1 -# 1 0 1 -# ] - -# @test isvalid(A, 1, 1) || isvalid(A, 2, 2) === false -# @test isvalid(A, 1, 2) && isvalid(A, 2, 1) === true -# @test isvalid_x(A, 1, 1) === false -# @test isvalid_x(A, 1, 3) === true -# @test isvalid_y(A, 1, 1) === false -# @test isvalid_y(A, 1, 2) === true From d238119cf90af8fc8983ef3178f3c1d0ab566c83 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 14:01:33 +0200 Subject: [PATCH 07/77] format --- src/MiniKernels.jl | 134 +++++++---- src/Utils.jl | 58 ++--- src/VariationalStokes/MiniKernels.jl | 20 ++ src/VariationalStokes/PressureKernels.jl | 74 +++++++ src/VariationalStokes/Stokes2D.jl | 221 +++++++++++++++++++ src/VariationalStokes/StressKernels.jl | 142 ++++++++++++ src/VariationalStokes/VelocityKernels.jl | 72 ++++++ src/VariationalStokes/mask.jl | 159 +++++++++++++ src/VariationalStokes/types.jl | 29 +++ src/rheology/BuoyancyForces.jl | 21 +- src/stokes/Stokes2D.jl | 4 +- src/stokes/VelocityKernels.jl | 47 ++-- src/thermal_diffusion/DiffusionExplicit.jl | 8 +- src/thermal_diffusion/DiffusionPT_kernels.jl | 1 - 14 files changed, 881 insertions(+), 109 deletions(-) create mode 100644 src/VariationalStokes/MiniKernels.jl create mode 100644 src/VariationalStokes/PressureKernels.jl create mode 100644 src/VariationalStokes/Stokes2D.jl create mode 100644 src/VariationalStokes/StressKernels.jl create mode 100644 src/VariationalStokes/VelocityKernels.jl create mode 100644 src/VariationalStokes/mask.jl create mode 100644 src/VariationalStokes/types.jl diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl index ce73cb1be..22f3e05aa 100644 --- a/src/MiniKernels.jl +++ b/src/MiniKernels.jl @@ -1,49 +1,85 @@ # unmasked versions -Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds...] -Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer, N}) where N = A[inds.+1...] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where I<:Integer = A[i-1, j] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where I<:Integer = A[i+1, j] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j-1] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where I<:Integer = A[i, j+1] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i-1, j, k] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i+1, j, k] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j-1, k] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j+1, k] -Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k-1] -Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where I<:Integer = A[i, j, k+1] +Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer,N}) where {N} = A[inds...] +Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer,N}) where {N} = A[inds .+ 1...] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ + i - 1, j +] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ + i + 1, j +] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ + i, j - 1 +] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ + i, j + 1 +] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i - 1, j, k +] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i + 1, j, k +] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i, j - 1, k +] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i, j + 1, k +] +Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i, j, k - 1 +] +Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ + i, j, k + 1 +] ## 2D mini kernels const T2 = AbstractArray{T,2} where {T} # finite differences -Base.@propagate_inbounds @inline _d_xa(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + right(A, I...)) * _dx -Base.@propagate_inbounds @inline _d_ya(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + front(A, I...)) * _dy -Base.@propagate_inbounds @inline _d_za(A::AbstractArray, _dz, I::Vararg{Integer, N}) where {N} = (-center(A, I...) + top(A, I...)) * _dz -Base.@propagate_inbounds @inline _d_xi(A::AbstractArray, _dx, I::Vararg{Integer, N}) where {N} = (-front(A, I...) + next(A, I...)) * _dx -Base.@propagate_inbounds @inline _d_yi(A::AbstractArray, _dy, I::Vararg{Integer, N}) where {N} = (-right(A, I...) + next(A, I...)) * _dy +Base.@propagate_inbounds @inline _d_xa( + A::AbstractArray, _dx, I::Vararg{Integer,N} +) where {N} = (-center(A, I...) + right(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya( + A::AbstractArray, _dy, I::Vararg{Integer,N} +) where {N} = (-center(A, I...) + front(A, I...)) * _dy +Base.@propagate_inbounds @inline _d_za( + A::AbstractArray, _dz, I::Vararg{Integer,N} +) where {N} = (-center(A, I...) + top(A, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi( + A::AbstractArray, _dx, I::Vararg{Integer,N} +) where {N} = (-front(A, I...) + next(A, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi( + A::AbstractArray, _dy, I::Vararg{Integer,N} +) where {N} = (-right(A, I...) + next(A, I...)) * _dy # averages -Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) -Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) -Base.@propagate_inbounds @inline _av_xa(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (center(A, I...) + right(A, I...)) * 0.5 -Base.@propagate_inbounds @inline _av_ya(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (center(A, I...) + front(A, I...)) * 0.5 -Base.@propagate_inbounds @inline _av_xi(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (front(A, I...), next(A, I...)) * 0.5 -Base.@propagate_inbounds @inline _av_yi(A::T, I::Vararg{Integer, 2}) where {T<:T2} = (right(A, I...), next(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = + 0.25 * mysum(A, (i+1):(i+2), (j+1):(j+2)) +Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = + 0.25 * mysum(A, (i):(i+1), (j):(j+1)) +Base.@propagate_inbounds @inline _av_xa(A::T, I::Vararg{Integer,2}) where {T<:T2} = + (center(A, I...) + right(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_ya(A::T, I::Vararg{Integer,2}) where {T<:T2} = + (center(A, I...) + front(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_xi(A::T, I::Vararg{Integer,2}) where {T<:T2} = + (front(A, I...), next(A, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_yi(A::T, I::Vararg{Integer,2}) where {T<:T2} = + (right(A, I...), next(A, I...)) * 0.5 # harmonic averages Base.@propagate_inbounds @inline function _harm(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2)) + return eltype(A)(4) * mysum(inv, A, (i+1):(i+2), (j+1):(j+2)) end Base.@propagate_inbounds @inline function _harm_a(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1)) + return eltype(A)(4) * mysum(inv, A, (i):(i+1), (j):(j+1)) end -Base.@propagate_inbounds @inline function _harm_xa(A::T, I::Vararg{Integer, 2}) where {T<:T2} +Base.@propagate_inbounds @inline function _harm_xa(A::T, I::Vararg{Integer,2}) where {T<:T2} return eltype(A)(2) * (inv(right(A, I...)) + inv(center(A, I...))) end -Base.@propagate_inbounds @inline function _harm_ya(A::T, I::Vararg{Integer, 2}) where {T<:T2} +Base.@propagate_inbounds @inline function _harm_ya(A::T, I::Vararg{Integer,2}) where {T<:T2} return eltype(A)(2) * (inv(front(A, I...)) + inv(center(A, I...))) end #others -Base.@propagate_inbounds @inline function _gather(A::T, I::Vararg{Integer, 2}) where {T<:T2} +Base.@propagate_inbounds @inline function _gather(A::T, I::Vararg{Integer,2}) where {T<:T2} return center(A, I...), right(A, I...), front(A, I...), next(A, I...) end @@ -55,16 +91,26 @@ const T3 = AbstractArray{T,3} where {T} return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz end # averages -Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) -Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + right(A, i, j, k)) -Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + front(A, i, j, k)) -Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + top(A, i, j, k)) -Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) -Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) -Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) -Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) -Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) -Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) +Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = + 0.125 * mysum(A, i:(i+1), j:(j+1), k:(k+1)) +Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = + 0.5 * (center(A, i, j, k) + right(A, i, j, k)) +Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = + 0.5 * (center(A, i, j, k) + front(A, i, j, k)) +Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = + 0.5 * (center(A, i, j, k) + top(A, i, j, k)) +Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, i:(i+1), j:(j+1), k:k) +Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, i:(i+1), j:j, k:(k+1)) +Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, i:i, j:(j+1), k:(k+1)) +Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, (i-1):i, (j-1):j, k:k) +Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, (i-1):i, j:j, (k-1):k) +Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = + 0.25 * mysum(A, i:i, (j-1):j, (k-1):k) # harmonic averages @inline function _harm_x(A::T, i, j, k) where {T<:T3} return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(right(A, i, j, k))) @@ -76,22 +122,22 @@ end return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(top(A, i, j, k))) end @inline function _harm_xy(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k)) + return eltype(A)(4) * inv(mysum(A, i:(i+1), j:(j+1), k:k)) end @inline function _harm_xz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:j, k:(k + 1))) + return eltype(A)(4) * inv(mysum(A, i:(i+1), j:j, k:(k+1))) end @inline function _harm_yz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, j:(j + 1), k:(k + 1))) + return eltype(A)(4) * inv(mysum(A, i:i, j:(j+1), k:(k+1))) end @inline function _harm_xyi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i - 1):i, (j - 1):j, k:k)) + return eltype(A)(4) * inv(mysum(A, (i-1):i, (j-1):j, k:k)) end @inline function _harm_xzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i - 1):i, j:j, (k - 1):k)) + return eltype(A)(4) * inv(mysum(A, (i-1):i, j:j, (k-1):k)) end @inline function _harm_yzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, (j - 1):j, (k - 1):k)) + return eltype(A)(4) * inv(mysum(A, i:i, (j-1):j, (k-1):k)) end # others diff --git a/src/Utils.jl b/src/Utils.jl index f9cce928e..3ec3c1380 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -83,8 +83,9 @@ end @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz -@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = - A.xx, A.yy, A.xy_c +@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = A.xx, +A.yy, +A.xy_c @inline function _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}}) where {T} return A.xx, A.yy, A.zz, A.yz_c, A.xz_c, A.xy_c end @@ -101,8 +102,9 @@ macro velocity(A) end @inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy -@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = - V.Vx, V.Vy, V.Vz +@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, +V.Vy, +V.Vz """ @displacement(U) @@ -115,10 +117,11 @@ macro displacement(A) end end -@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} = - U.Ux, U.Uy -@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,3}}) where {T} = - U.Ux, U.Uy, U.Uz +@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} = U.Ux, +U.Uy +@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,3}}) where {T} = U.Ux, +U.Uy, +U.Uz """ @qT(V) @@ -132,8 +135,9 @@ macro qT(A) end @inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx, A.qTy -@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = - A.qTx, A.qTy, A.qTz +@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = A.qTx, +A.qTy, +A.qTz """ @qT2(V) @@ -146,8 +150,8 @@ macro qT2(A) end end -@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = - A.qTx2, A.qTy2 +@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx2, +A.qTy2 @inline function unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} return A.qTx2, A.qTy2, A.qTz2 end @@ -360,8 +364,8 @@ end @inline function _maxloc_window_clamped(A, I, J, width_x, width_y) nx, ny = size(A) - I_range = (I - width_x):(I + width_x) - J_range = (J - width_y):(J + width_y) + I_range = (I-width_x):(I+width_x) + J_range = (J-width_y):(J+width_y) x = -Inf for j in J_range @@ -379,9 +383,9 @@ end @inline function _maxloc_window_clamped(A, I, J, K, width_x, width_y, width_z) nx, ny, nz = size(A) - I_range = (I - width_x):(I + width_x) - J_range = (J - width_y):(J + width_y) - K_range = (K - width_z):(K + width_z) + I_range = (I-width_x):(I+width_x) + J_range = (J-width_y):(J+width_y) + K_range = (K-width_z):(K+width_z) x = -Inf for k in K_range @@ -428,17 +432,21 @@ function compute_dt(::CPUBackendTrait, S::JustRelax.StokesArrays, args...) return _compute_dt(S, args...) end -@inline _compute_dt(S::JustRelax.StokesArrays, di) = - _compute_dt(@velocity(S), di, Inf, maximum) +@inline _compute_dt(S::JustRelax.StokesArrays, di) = _compute_dt( + @velocity(S), di, Inf, maximum +) -@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = - _compute_dt(@velocity(S), di, dt_diff, maximum) +@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = _compute_dt( + @velocity(S), di, dt_diff, maximum +) -@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) = - _compute_dt(@velocity(S), di, dt_diff, maximum_mpi) +@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) = _compute_dt( + @velocity(S), di, dt_diff, maximum_mpi +) -@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) = - _compute_dt(@velocity(S), di, Inf, maximum_mpi) +@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) = _compute_dt( + @velocity(S), di, Inf, maximum_mpi +) @inline function _compute_dt(V::NTuple, di, dt_diff, max_fun::F) where {F<:Function} n = inv(length(V) + 0.1) diff --git a/src/VariationalStokes/MiniKernels.jl b/src/VariationalStokes/MiniKernels.jl new file mode 100644 index 000000000..5584994cc --- /dev/null +++ b/src/VariationalStokes/MiniKernels.jl @@ -0,0 +1,20 @@ +# masked versions +for fn in (:center, :next, :left, :right, :back, :front) + @eval begin + Base.@propagate_inbounds @inline ($fn)( + A::T, ϕ::T, inds::Vararg{Integer,N} + ) where {T<:AbstractArray,N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) + end +end + +# finite differences +Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = + (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = + (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy diff --git a/src/VariationalStokes/PressureKernels.jl b/src/VariationalStokes/PressureKernels.jl new file mode 100644 index 000000000..86c49f6cf --- /dev/null +++ b/src/VariationalStokes/PressureKernels.jl @@ -0,0 +1,74 @@ +function compute_P!( + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio::JustPIC.PhaseRatios, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ; + ΔTc=nothing, + kwargs..., +) where {N} + ni = size(P) + @parallel (@idx ni) compute_P_kernel!( + P, P0, RP, ∇V, η, rheology, phase_ratio.center, ϕ, dt, r, θ_dτ, ΔTc + ) + return nothing +end + +@parallel_indices (I...) function compute_P!( + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ, + ::Nothing, +) where {N} + if isvalid_c(ϕ, I...) + K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...])) + RP[I...], P[I...] = _compute_P!( + P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ + ) + else + RP[I...] = P[I...] = zero(eltype(P)) + end + return nothing +end + +@parallel_indices (I...) function compute_P_kernel!( + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio::C, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ, + ΔTc, + ::Nothing, +) where {N,C<:JustRelax.CellArray} + if isvalid_c(ϕ, I...) + phase_ratio_I = phase_ratio[I...] + K = fn_ratio(get_bulk_modulus, rheology, phase_ratio_I) + α = fn_ratio(get_thermal_expansion, rheology, phase_ratio_I) + RP[I...], P[I...] = _compute_P!( + P[I...], P0[I...], ∇V[I...], ΔTc[I...], α, η[I...], K, dt, r, θ_dτ + ) + else + RP[I...] = P[I...] = zero(eltype(P)) + end + return nothing +end diff --git a/src/VariationalStokes/Stokes2D.jl b/src/VariationalStokes/Stokes2D.jl new file mode 100644 index 000000000..6121b698b --- /dev/null +++ b/src/VariationalStokes/Stokes2D.jl @@ -0,0 +1,221 @@ +## 2D VISCO-ELASTIC STOKES SOLVER + +# backend trait +function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs) + solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) + return nothing +end + +# entry point for extensions +function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + +function _solve_VS!( + stokes::JustRelax.StokesArrays, + pt_stokes, + di::NTuple{2,T}, + flow_bcs::AbstractFlowBoundaryConditions, + ρg, + phase_ratios::JustPIC.PhaseRatios, + ϕ::JustRelax.RockRatio, + rheology, + args, + dt, + igg::IGG; + viscosity_cutoff=(-Inf, Inf), + viscosity_relaxation=1e-2, + iterMax=50e3, + iterMin=1e2, + nout=500, + b_width=(4, 4, 0), + verbose=true, + kwargs..., +) where {T} + + # unpack + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + relλ = 0.2 + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv = @zeros(ni .+ 1...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + + # compute buoyancy forces and viscosity + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + displacement2velocity!(stokes, dt, flow_bcs) + + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + + while iter ≤ iterMax + iterMin < iter && err < ϵ && break + + wtime0 += @elapsed begin + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + + compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + dt, + r, + θ_dτ, + args, + ) + + update_ρg!(ρg[2], phase_ratios, rheology, args) + + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + update_viscosity!( + stokes, + phase_ratios, + args, + rheology, + viscosity_cutoff; + relaxation=viscosity_relaxation, + ) + + @parallel (@idx ni .+ 1) update_stresses_center_vertex!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.xy,), + @tensor_center(stokes.τ_o), + (stokes.τ_o.xy,), + θ, + stokes.P, + stokes.viscosity.η, + λ, + λv, + stokes.τ.II, + stokes.viscosity.η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + ϕ, + ) + update_halo!(stokes.τ.xy) + + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni .+ 1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + _di..., + ) + # apply boundary conditions + velocity2displacement!(stokes, dt) + free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) + # end + end + + iter += 1 + + if iter % nout == 0 && iter > 1 + errs = ( + norm_mpi(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / + length(stokes.R.Rx), + norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / + length(stokes.R.Ry), + norm_mpi(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum_mpi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) + + # accumulate plastic strain tensor + @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end diff --git a/src/VariationalStokes/StressKernels.jl b/src/VariationalStokes/StressKernels.jl new file mode 100644 index 000000000..6b1d848a3 --- /dev/null +++ b/src/VariationalStokes/StressKernels.jl @@ -0,0 +1,142 @@ +# 2D kernel +@parallel_indices (I...) function update_stresses_center_vertex!( + ε::NTuple{3,T}, # normal components @ centers; shear components @ vertices + ε_pl::NTuple{3}, # whole Voigt tensor @ centers + EII, # accumulated plastic strain rate @ centers + τ::NTuple{3}, # whole Voigt tensor @ centers + τshear_v::NTuple{1}, # shear tensor components @ vertices + τ_o::NTuple{3}, + τshear_ov::NTuple{1}, # shear tensor components @ vertices + Pr, + Pr_c, + η, + λ, + λv, + τII, + η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_center, + phase_vertex, + ϕ::JustRelax.RockRatio, +) where {T} + τxyv = τshear_v[1] + τxyv_old = τshear_ov[1] + ni = size(Pr) + Ic = clamped_indices(ni, I...) + + if isvalid_v(ϕ, I...) + # interpolate to ith vertex + Pv_ij = av_clamped(Pr, Ic...) + εxxv_ij = av_clamped(ε[1], Ic...) + εyyv_ij = av_clamped(ε[2], Ic...) + τxxv_ij = av_clamped(τ[1], Ic...) + τyyv_ij = av_clamped(τ[2], Ic...) + τxxv_old_ij = av_clamped(τ_o[1], Ic...) + τyyv_old_ij = av_clamped(τ_o[2], Ic...) + EIIv_ij = av_clamped(EII, Ic...) + + ## vertex + phase = @inbounds phase_vertex[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( + rheology, EIIv_ij, phase + ) + _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + Kv = fn_ratio(get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + ηv_ij = av_clamped(η, Ic...) + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = + (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * + dτ_rv + dτyyv = + (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * + dτ_rv + dτxyv = + ( + -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + + 2.0 * ηv_ij * ε[3][I...] + ) * dτ_rv + τIIv_ij = + √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) + + # yield function @ center + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) + # stress correction @ vertex + λv[I...] = + (1.0 - relλ) * λv[I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + dQdτxy = 0.5 * (τxyv[I...] + dτxyv) / τIIv_ij + τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[I...] * dQdτxy * dτ_rv + else + # stress correction @ vertex + τxyv[I...] += dτxyv + end + else + τxyv[I...] = zero(eltype(T)) + end + + ## center + if all(I .≤ ni) + if isvalid_c(ϕ, I...) + # Material properties + phase = @inbounds phase_center[I...] + _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase( + rheology, EII[I...], phase + ) + K = fn_ratio(get_bulk_modulus, rheology, phase) + volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ + ηij = η[I...] + dτ_r = 1.0 / (θ_dτ + ηij * _Gdt + 1.0) + + # cache strain rates for center calculations + τij, τij_o, εij = cache_tensors(τ, τ_o, ε, I...) + + # visco-elastic strain rates @ center + εij_ve = @. εij + 0.5 * τij_o * _Gdt + εII_ve = GeoParams.second_invariant(εij_ve) + # stress increments @ center + dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r + τII_ij = GeoParams.second_invariant(dτij .+ τij) + # yield function @ center + F = τII_ij - C - Pr[I...] * sinϕ + + if is_pl && !iszero(τII_ij) + # stress correction @ center + λ[I...] = + (1.0 - relλ) * λ[I...] + + relλ .* (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) + dQdτij = @. 0.5 * (τij + dτij) / τII_ij + εij_pl = λ[I...] .* dQdτij + dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r + τij = dτij .+ τij + setindex!.(τ, τij, I...) + setindex!.(ε_pl, εij_pl, I...) + τII[I...] = GeoParams.second_invariant(τij) + Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ + η_vep[I...] = 0.5 * τII_ij / εII_ve + else + # stress correction @ center + setindex!.(τ, dτij .+ τij, I...) + η_vep[I...] = ηij + τII[I...] = τII_ij + end + + Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) + else + Pr_c[I...] = zero(eltype(T)) + # τij, = cache_tensors(τ, τ_o, ε, I...) + dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) + # setindex!.(τ, dτij .+ τij, I...) + setindex!.(τ, dτij, I...) + end + end + + return nothing +end diff --git a/src/VariationalStokes/VelocityKernels.jl b/src/VariationalStokes/VelocityKernels.jl new file mode 100644 index 000000000..6ff77e206 --- /dev/null +++ b/src/VariationalStokes/VelocityKernels.jl @@ -0,0 +1,72 @@ +@parallel_indices (I...) function compute_∇V!( + ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} +) where {T,N} + @inline d_xi(A) = _d_xi(A, _di[1], I...) + @inline d_yi(A) = _d_yi(A, _di[2], I...) + @inline d_zi(A) = _d_zi(A, _di[3], I...) + + f = d_xi, d_yi, d_zi + + if isvalid_c(ϕ, I...) + @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) + else + @inbounds ∇V[I...] = zero(T) + end + return nothing +end + +@parallel_indices (i, j) function compute_V!( + Vx::AbstractArray{T,2}, + Vy, + Rx, + Ry, + P, + τxx, + τyy, + τxy, + ηdτ, + ρgx, + ρgy, + ητ, + ϕ::JustRelax.RockRatio, + _dx, + _dy, +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + if all((i, j) .< size(Vx) .- 1) + if isvalid_vx(ϕ, i + 1, j) + Rx[i, j] = + R_Vx = ( + -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - + av_xa(ρgx) + ) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + else + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) + end + end + + if all((i, j) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j + 1) + Ry[i, j] = + R_Vy = + -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - + av_ya(ρgy) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + else + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + end + end + + return nothing +end diff --git a/src/VariationalStokes/mask.jl b/src/VariationalStokes/mask.jl new file mode 100644 index 000000000..27ba90c86 --- /dev/null +++ b/src/VariationalStokes/mask.jl @@ -0,0 +1,159 @@ + +function RockRatio(nx, ny) + ni = nx, ny + center = @zeros(ni...) + vertex = @zeros(ni .+ 1...) + Vx = @zeros(nx+1, ny) # no ghost nodes! + Vy = @zeros(nx, ny+1) # no ghost nodes! + + return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) +end + +function RockRatio(nx, ny, nz) + ni = nx, ny, nz + center = @zeros(ni...) + vertex = @zeros(ni .+ 1...) + Vx = @zeros(nx+1, ny, nz) # no ghost nodes! + Vy = @zeros(nx, ny+1, nz) # no ghost nodes! + Vz = @zeros(nx, ny, nz+1) # no ghost nodes! + yz = @zeros(nx, ny + 1, nz + 1) + xz = @zeros(nx + 1, ny, nz + 1) + xy = @zeros(nx + 1, ny + 1, nz) + + return JustRelax.RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) +end + +@inline size_c(x::JustRelax.AbstractMask) = size(x.center) +@inline size_v(x::JustRelax.AbstractMask) = size(x.vertex) +@inline size_vx(x::JustRelax.AbstractMask) = size(x.Vx) +@inline size_vy(x::JustRelax.AbstractMask) = size(x.Vy) +@inline size_vz(x::JustRelax.AbstractMask) = size(x.Vz) +@inline size_yz(x::JustRelax.AbstractMask) = size(x.yz) +@inline size_xz(x::JustRelax.AbstractMask) = size(x.xz) +@inline size_xy(x::JustRelax.AbstractMask) = size(x.xy) + +""" + update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) + +Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The rock ratio object to be updated. +- `phase_ratios`: The ratios of different phases present. +- `air_phase`: The phase representing air. +""" +function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) + nvi = size_v(ϕ) + @parallel (@idx nvi) update_rock_ratio_cv!( + ϕ, phase_ratios.center, phase_ratios.vertex, air_phase + ) + @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) + return nothing +end + +@inline compute_rock_ratio( + phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} +) where {N} = 1 - @index phase_ratio[air_phase, I...] +@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[ + air_phase, I... +] + +@parallel_indices (I...) function update_rock_ratio_cv!( + ϕ, ratio_center, ratio_vertex, air_phase +) + if all(I .≤ size(ratio_center)) + ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) + end + ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) + return nothing +end + +@parallel_indices (I...) function update_rock_ratio_vel!( + ϕ::JustRelax.RockRatio{T,N} +) where {T,N} + # 2D + @inline av_x(A::AbstractArray{T,2}) where {T} = _av_xa(A, I...) + @inline av_y(A::AbstractArray{T,2}) where {T} = _av_ya(A, I...) + # 3D + @inline av_x(A::AbstractArray{T,3}) where {T} = _av_yz(A, I...) + @inline av_y(A::AbstractArray{T,3}) where {T} = _av_xz(A, I...) + @inline av_z(A::AbstractArray{T,3}) where {T} = _av_xy(A, I...) + + all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case + all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) + end + return nothing +end + +""" + isvalid_c(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.center[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_c(ϕ::JustRelax.RockRatio, i, j) + vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) + vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + v = vx * vy + return v * (ϕ.center[i, j] > 0) +end + +""" + isvalid_v(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.vertex[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_v(ϕ::JustRelax.RockRatio, i, j) + nx, ny = size(ϕ.Vx) + j_bot = max(j - 1, 1) + j0 = min(j, ny) + vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + + nx, ny = size(ϕ.Vy) + i_left = max(i - 1, 1) + i0 = min(i, nx) + vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + v = vx * vy + return v * (ϕ.vertex[i, j] > 0) +end + +""" + isvalid_vx(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vx[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) + cv = c * v + return cv * (ϕ.Vx[i, j] > 0) +end + +""" + isvalid_vy(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vy[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) + cv = c * v + return cv * (ϕ.Vy[i, j] > 0) +end diff --git a/src/VariationalStokes/types.jl b/src/VariationalStokes/types.jl new file mode 100644 index 000000000..aa8ef3642 --- /dev/null +++ b/src/VariationalStokes/types.jl @@ -0,0 +1,29 @@ +abstract type AbstractMask end + +struct RockRatio{T,N} <: AbstractMask + center::T + vertex::T + Vx::T + Vy::T + Vz::Union{Nothing,T} + yz::Union{Nothing,T} + xz::Union{Nothing,T} + xy::Union{Nothing,T} + + function RockRatio( + center::AbstractArray{F,N}, + vertex::T, + Vx::T, + Vy::T, + Vz::Union{Nothing,T}, + yz::Union{Nothing,T}, + xz::Union{Nothing,T}, + xy::Union{Nothing,T}, + ) where {F,N,T} + return new{T,N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) + end +end + +RockRatio(ni::NTuple{N,Integer}) where {N} = RockRatio(ni...) + +Adapt.@adapt_structure RockRatio diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 946cab3b8..8dcb4282e 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -91,17 +91,20 @@ Compute the buoyancy forces based on the given rheology, arguments, and phase ra end # without phase ratios -@inline update_ρg!(ρg::AbstractArray, rheology, args) = - update_ρg!(isconstant(rheology), ρg, rheology, args) +@inline update_ρg!(ρg::AbstractArray, rheology, args) = update_ρg!( + isconstant(rheology), ρg, rheology, args +) @inline update_ρg!(::ConstantDensityTrait, ρg, rheology, args) = nothing -@inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = - compute_ρg!(ρg, rheology, args) +@inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = compute_ρg!( + ρg, rheology, args +) # with phase ratios -@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = - update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args) +@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = update_ρg!( + isconstant(rheology), ρg, phase_ratios, rheology, args +) @inline update_ρg!( ::ConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args ) = nothing -@inline update_ρg!( - ::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args -) = compute_ρg!(ρg, phase_ratios, rheology, args) +@inline update_ρg!(::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args) = compute_ρg!( + ρg, phase_ratios, rheology, args +) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index ad2eb38ae..ce65dd0b7 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -377,7 +377,7 @@ function _solve!( center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) - @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( + @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( Vx_on_Vy, stokes.V.Vx ) @@ -601,7 +601,7 @@ function _solve!( ) update_halo!(stokes.τ.xy) - @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!( + @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!( Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] ) diff --git a/src/stokes/VelocityKernels.jl b/src/stokes/VelocityKernels.jl index 6dca036b5..da07e49a3 100644 --- a/src/stokes/VelocityKernels.jl +++ b/src/stokes/VelocityKernels.jl @@ -3,24 +3,23 @@ @parallel_indices (i, j) function compute_∇V!( ∇V::AbstractArray{T,2}, Vx, Vy, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A,_dx, i, j) - d_yi(A) = _d_yi(A,_dy, i, j) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) ∇V[i, j] = d_xi(Vx) + d_yi(Vy) return nothing end - ## DEVIATORIC STRAIN RATE TENSOR @parallel_indices (i, j) function compute_strain_rate!( εxx::AbstractArray{T,2}, εyy, εxy, ∇V, Vx, Vy, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A,_dx, i, j) - d_yi(A) = _d_yi(A,_dy, i, j) - d_xa(A) = _d_xa(A,_dx, i, j) - d_ya(A) = _d_ya(A,_dy, i, j) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) if all((i, j) .≤ size(εxx)) ∇V_ij = ∇V[i, j] / 3.0 @@ -83,10 +82,10 @@ end @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, _dx, _dy ) where {T} - d_xi(A) = _d_xi(A,_dx, i, j) - d_yi(A) = _d_yi(A,_dy, i, j) - d_xa(A) = _d_xa(A,_dx, i, j) - d_ya(A) = _d_ya(A,_dy, i, j) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -107,10 +106,10 @@ end @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, _dx, _dy, dt ) where {T} - d_xi(A) = _d_xi(A,_dx, i, j) - d_yi(A) = _d_yi(A,_dy, i, j) - d_xa(A) = _d_xa(A,_dx, i, j) - d_ya(A) = _d_ya(A,_dy, i, j) + d_xi(A) = _d_xi(A, _dx, i, j) + d_yi(A) = _d_yi(A, _dy, i, j) + d_xa(A) = _d_xa(A, _dx, i, j) + d_ya(A) = _d_ya(A, _dy, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -246,10 +245,10 @@ end @parallel_indices (i, j) function compute_Res!( Rx::AbstractArray{T,2}, Ry, P, τxx, τyy, τxy, ρgx, ρgy, _dx, _dy ) where {T} - @inline d_xa(A) = _d_xa(A,_dx, i, j) - @inline d_ya(A) = _d_ya(A,_dy, i, j) - @inline d_xi(A) = _d_xi(A,_dx, i, j) - @inline d_yi(A) = _d_yi(A,_dy, i, j) + @inline d_xa(A) = _d_xa(A, _dx, i, j) + @inline d_ya(A) = _d_ya(A, _dy, i, j) + @inline d_xi(A) = _d_xi(A, _dx, i, j) + @inline d_yi(A) = _d_yi(A, _dy, i, j) @inline av_xa(A) = _av_xa(A, i, j) @inline av_ya(A) = _av_ya(A, i, j) @@ -267,10 +266,10 @@ end @parallel_indices (i, j) function compute_Res!( Rx::AbstractArray{T,2}, Ry, Vx, Vy, Vx_on_Vy, P, τxx, τyy, τxy, ρgx, ρgy, _dx, _dy, dt ) where {T} - @inline d_xa(A) = _d_xa(A,_dx, i, j) - @inline d_ya(A) = _d_ya(A,_dy, i, j) - @inline d_xi(A) = _d_xi(A,_dx, i, j) - @inline d_yi(A) = _d_yi(A,_dy, i, j) + @inline d_xa(A) = _d_xa(A, _dx, i, j) + @inline d_ya(A) = _d_ya(A, _dy, i, j) + @inline d_xi(A) = _d_xi(A, _dx, i, j) + @inline d_yi(A) = _d_yi(A, _dy, i, j) @inline av_xa(A) = _av_xa(A, i, j) @inline av_ya(A) = _av_ya(A, i, j) @@ -279,7 +278,7 @@ end if all((i, j) .≤ size(Rx)) Rx[i, j] = d_xa(τxx) + d_yi(τxy) - d_xa(P) - av_xa(ρgx) end - + if all((i, j) .≤ size(Ry)) θ = 1.0 # Interpolated Vx into Vy node (includes density gradient) diff --git a/src/thermal_diffusion/DiffusionExplicit.jl b/src/thermal_diffusion/DiffusionExplicit.jl index 2fa055fb8..97a0fee4c 100644 --- a/src/thermal_diffusion/DiffusionExplicit.jl +++ b/src/thermal_diffusion/DiffusionExplicit.jl @@ -381,7 +381,7 @@ function JustRelax.solve!( # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( + @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy ) @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _dx, _dy) @@ -413,7 +413,7 @@ function JustRelax.solve!( # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( + @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, phase_ratios.center, args, _di... ) @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _di...) @@ -442,7 +442,7 @@ function JustRelax.solve!( nx, ny = size(thermal.T) # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( + @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy ) @parallel advect_T!( @@ -481,7 +481,7 @@ function JustRelax.solve!( nx, ny = size(thermal.T) # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( + @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, phases, rheology, args, _dx, _dy ) @parallel advect_T!( diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl index 190ec4b81..4acb29f11 100644 --- a/src/thermal_diffusion/DiffusionPT_kernels.jl +++ b/src/thermal_diffusion/DiffusionPT_kernels.jl @@ -56,7 +56,6 @@ end get_K(idx, args) = compute_phase(compute_conductivity, rheology, idx, args) - @inbounds if all(I .≤ size(qTx)) T_ijk = (T[(I .+ 1)...] + T[i, j + 1, k + 1]) * 0.5 args_ijk = (; T=T_ijk, P=av_yz(args.P)) From 32f0024c8232f42bdea38aace9e66c996071e17d Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 14:06:36 +0200 Subject: [PATCH 08/77] format --- src/variational_stokes/MiniKernels.jl | 20 +++-- src/variational_stokes/PressureKernels.jl | 19 ++++- src/variational_stokes/Stokes2D.jl | 43 ++++++----- src/variational_stokes/StressKernels.jl | 37 +++++---- src/variational_stokes/VelocityKernels.jl | 41 +++++++--- src/variational_stokes/mask.jl | 94 +++++++++++++---------- src/variational_stokes/types.jl | 29 ++++--- 7 files changed, 175 insertions(+), 108 deletions(-) diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 1d2561a6f..5584994cc 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -1,14 +1,20 @@ # masked versions for fn in (:center, :next, :left, :right, :back, :front) @eval begin - Base.@propagate_inbounds @inline ($fn)(A::T, ϕ::T, inds::Vararg{Integer, N}) where {T<:AbstractArray, N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) + Base.@propagate_inbounds @inline ($fn)( + A::T, ϕ::T, inds::Vararg{Integer,N} + ) where {T<:AbstractArray,N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) end end - # finite differences -Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy -Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer, N}) where {N, T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz -Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer, N}) where {N, T} = (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer, N}) where {N, T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy \ No newline at end of file +Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy +Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer,N}) where {N,T} = + (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz +Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = + (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx +Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = + (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy diff --git a/src/variational_stokes/PressureKernels.jl b/src/variational_stokes/PressureKernels.jl index dc67c3a36..86c49f6cf 100644 --- a/src/variational_stokes/PressureKernels.jl +++ b/src/variational_stokes/PressureKernels.jl @@ -21,11 +21,24 @@ function compute_P!( end @parallel_indices (I...) function compute_P!( - P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, ϕ::JustRelax.RockRatio, dt, r, θ_dτ, ::Nothing + P, + P0, + RP, + ∇V, + η, + rheology::NTuple{N,MaterialParams}, + phase_ratio, + ϕ::JustRelax.RockRatio, + dt, + r, + θ_dτ, + ::Nothing, ) where {N} if isvalid_c(ϕ, I...) K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...])) - RP[I...], P[I...] = _compute_P!(P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ) + RP[I...], P[I...] = _compute_P!( + P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ + ) else RP[I...] = P[I...] = zero(eltype(P)) end @@ -58,4 +71,4 @@ end RP[I...] = P[I...] = zero(eltype(P)) end return nothing -end \ No newline at end of file +end diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 8ce08c9c4..6121b698b 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -7,7 +7,9 @@ function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwarg end # entry point for extensions -solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) = _solve_VS!(stokes, args...; kwargs...) +function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end function _solve_VS!( stokes::JustRelax.StokesArrays, @@ -109,7 +111,7 @@ function _solve_VS!( @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... ) - + update_viscosity!( stokes, phase_ratios, @@ -144,25 +146,24 @@ function _solve_VS!( ) update_halo!(stokes.τ.xy) - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni.+1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - _di..., - ) - # apply boundary conditions - velocity2displacement!(stokes, dt) - free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - update_halo!(@velocity(stokes)...) + @parallel (@idx ni .+ 1) compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + _di..., + ) + # apply boundary conditions + velocity2displacement!(stokes, dt) + free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) # end end @@ -217,4 +218,4 @@ function _solve_VS!( norm_Ry=norm_Ry, norm_∇V=norm_∇V, ) -end \ No newline at end of file +end diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl index 2b5610fed..6b1d848a3 100644 --- a/src/variational_stokes/StressKernels.jl +++ b/src/variational_stokes/StressKernels.jl @@ -1,6 +1,6 @@ # 2D kernel @parallel_indices (I...) function update_stresses_center_vertex!( - ε::NTuple{3, T}, # normal components @ centers; shear components @ vertices + ε::NTuple{3,T}, # normal components @ centers; shear components @ vertices ε_pl::NTuple{3}, # whole Voigt tensor @ centers EII, # accumulated plastic strain rate @ centers τ::NTuple{3}, # whole Voigt tensor @ centers @@ -20,8 +20,8 @@ rheology, phase_center, phase_vertex, - ϕ::JustRelax.RockRatio -) where T + ϕ::JustRelax.RockRatio, +) where {T} τxyv = τshear_v[1] τxyv_old = τshear_ov[1] ni = size(Pr) @@ -29,18 +29,20 @@ if isvalid_v(ϕ, I...) # interpolate to ith vertex - Pv_ij = av_clamped(Pr, Ic...) - εxxv_ij = av_clamped(ε[1], Ic...) - εyyv_ij = av_clamped(ε[2], Ic...) - τxxv_ij = av_clamped(τ[1], Ic...) - τyyv_ij = av_clamped(τ[2], Ic...) + Pv_ij = av_clamped(Pr, Ic...) + εxxv_ij = av_clamped(ε[1], Ic...) + εyyv_ij = av_clamped(ε[2], Ic...) + τxxv_ij = av_clamped(τ[1], Ic...) + τyyv_ij = av_clamped(τ[2], Ic...) τxxv_old_ij = av_clamped(τ_o[1], Ic...) τyyv_old_ij = av_clamped(τ_o[2], Ic...) - EIIv_ij = av_clamped(EII, Ic...) + EIIv_ij = av_clamped(EII, Ic...) ## vertex phase = @inbounds phase_vertex[I...] - is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(rheology, EIIv_ij, phase) + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( + rheology, EIIv_ij, phase + ) _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) Kv = fn_ratio(get_bulk_modulus, rheology, phase) volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ @@ -49,15 +51,18 @@ # stress increments @ vertex dτxxv = - (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * dτ_rv + (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * + dτ_rv dτyyv = - (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * dτ_rv + (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * + dτ_rv dτxyv = ( -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + 2.0 * ηv_ij * ε[3][I...] ) * dτ_rv - τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) + τIIv_ij = + √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) # yield function @ center Fv = τIIv_ij - Cv - Pv_ij * sinϕv @@ -82,7 +87,9 @@ # Material properties phase = @inbounds phase_center[I...] _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) - is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase(rheology, EII[I...], phase) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase( + rheology, EII[I...], phase + ) K = fn_ratio(get_bulk_modulus, rheology, phase) volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ ηij = η[I...] @@ -132,4 +139,4 @@ end return nothing -end \ No newline at end of file +end diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index eda4c3474..6ff77e206 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,6 +1,6 @@ @parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T, N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} -) where {T, N} + ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} +) where {T,N} @inline d_xi(A) = _d_xi(A, _di[1], I...) @inline d_yi(A) = _d_yi(A, _di[2], I...) @inline d_zi(A) = _d_zi(A, _di[3], I...) @@ -9,14 +9,28 @@ if isvalid_c(ϕ, I...) @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) - else + else @inbounds ∇V[I...] = zero(T) end return nothing end @parallel_indices (i, j) function compute_V!( - Vx::AbstractArray{T,2}, Vy, Rx, Ry, P, τxx, τyy, τxy, ηdτ, ρgx, ρgy, ητ, ϕ::JustRelax.RockRatio, _dx, _dy + Vx::AbstractArray{T,2}, + Vy, + Rx, + Ry, + P, + τxx, + τyy, + τxy, + ηdτ, + ρgx, + ρgy, + ητ, + ϕ::JustRelax.RockRatio, + _dx, + _dy, ) where {T} d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) @@ -29,23 +43,30 @@ end if all((i, j) .< size(Vx) .- 1) if isvalid_vx(ϕ, i + 1, j) - Rx[i, j]= R_Vx = (-d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx)) + Rx[i, j] = + R_Vx = ( + -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - + av_xa(ρgx) + ) Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) else - Rx[i, j] = zero(T) + Rx[i, j] = zero(T) Vx[i + 1, j + 1] = zero(T) end end - if all((i, j) .< size(Vy) .- 1) + if all((i, j) .< size(Vy) .- 1) if isvalid_vy(ϕ, i, j + 1) - Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy) + Ry[i, j] = + R_Vy = + -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - + av_ya(ρgy) Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) else - Ry[i, j] = zero(T) + Ry[i, j] = zero(T) Vy[i + 1, j + 1] = zero(T) end end return nothing -end \ No newline at end of file +end diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 0e5954f65..27ba90c86 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -1,30 +1,30 @@ function RockRatio(nx, ny) - ni = nx, ny - center = @zeros(ni...) - vertex = @zeros(ni.+1...) - Vx = @zeros(nx+1, ny) # no ghost nodes! - Vy = @zeros(nx, ny+1) # no ghost nodes! + ni = nx, ny + center = @zeros(ni...) + vertex = @zeros(ni .+ 1...) + Vx = @zeros(nx+1, ny) # no ghost nodes! + Vy = @zeros(nx, ny+1) # no ghost nodes! return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) end function RockRatio(nx, ny, nz) - ni = nx, ny, nz - center = @zeros(ni...) - vertex = @zeros(ni.+1...) - Vx = @zeros(nx+1, ny, nz) # no ghost nodes! - Vy = @zeros(nx, ny+1, nz) # no ghost nodes! - Vz = @zeros(nx, ny, nz+1) # no ghost nodes! - yz = @zeros(nx, ny + 1, nz + 1) - xz = @zeros(nx + 1, ny, nz + 1) - xy = @zeros(nx + 1, ny + 1, nz) + ni = nx, ny, nz + center = @zeros(ni...) + vertex = @zeros(ni .+ 1...) + Vx = @zeros(nx+1, ny, nz) # no ghost nodes! + Vy = @zeros(nx, ny+1, nz) # no ghost nodes! + Vz = @zeros(nx, ny, nz+1) # no ghost nodes! + yz = @zeros(nx, ny + 1, nz + 1) + xz = @zeros(nx + 1, ny, nz + 1) + xy = @zeros(nx + 1, ny + 1, nz) return JustRelax.RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) end -@inline size_c(x::JustRelax.AbstractMask) = size(x.center) -@inline size_v(x::JustRelax.AbstractMask) = size(x.vertex) +@inline size_c(x::JustRelax.AbstractMask) = size(x.center) +@inline size_v(x::JustRelax.AbstractMask) = size(x.vertex) @inline size_vx(x::JustRelax.AbstractMask) = size(x.Vx) @inline size_vy(x::JustRelax.AbstractMask) = size(x.Vy) @inline size_vz(x::JustRelax.AbstractMask) = size(x.Vz) @@ -44,15 +44,23 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. """ function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) nvi = size_v(ϕ) - @parallel (@idx nvi) update_rock_ratio_cv!(ϕ, phase_ratios.center, phase_ratios.vertex, air_phase) + @parallel (@idx nvi) update_rock_ratio_cv!( + ϕ, phase_ratios.center, phase_ratios.vertex, air_phase + ) @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) return nothing end -@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = 1 - @index phase_ratio[air_phase, I...] -@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer, N}) where N = @index phase_ratio[air_phase, I...] +@inline compute_rock_ratio( + phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} +) where {N} = 1 - @index phase_ratio[air_phase, I...] +@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[ + air_phase, I... +] -@parallel_indices (I...) function update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase) +@parallel_indices (I...) function update_rock_ratio_cv!( + ϕ, ratio_center, ratio_vertex, air_phase +) if all(I .≤ size(ratio_center)) ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) end @@ -60,14 +68,16 @@ end return nothing end -@parallel_indices (I...) function update_rock_ratio_vel!(ϕ::JustRelax.RockRatio{T, N}) where {T, N} +@parallel_indices (I...) function update_rock_ratio_vel!( + ϕ::JustRelax.RockRatio{T,N} +) where {T,N} # 2D - @inline av_x(A::AbstractArray{T, 2}) where T = _av_xa(A, I...) - @inline av_y(A::AbstractArray{T, 2}) where T = _av_ya(A, I...) + @inline av_x(A::AbstractArray{T,2}) where {T} = _av_xa(A, I...) + @inline av_y(A::AbstractArray{T,2}) where {T} = _av_ya(A, I...) # 3D - @inline av_x(A::AbstractArray{T, 3}) where T = _av_yz(A, I...) - @inline av_y(A::AbstractArray{T, 3}) where T = _av_xz(A, I...) - @inline av_z(A::AbstractArray{T, 3}) where T = _av_xy(A, I...) + @inline av_x(A::AbstractArray{T,3}) where {T} = _av_yz(A, I...) + @inline av_y(A::AbstractArray{T,3}) where {T} = _av_xz(A, I...) + @inline av_z(A::AbstractArray{T,3}) where {T} = _av_xy(A, I...) all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) @@ -87,8 +97,8 @@ Check if `ϕ.center[inds...]` is a not a nullspace. - `inds`: Cartesian indices to check. """ function isvalid_c(ϕ::JustRelax.RockRatio, i, j) - vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) - vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) + vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) v = vx * vy return v * (ϕ.center[i, j] > 0) end @@ -104,15 +114,15 @@ Check if `ϕ.vertex[inds...]` is a not a nullspace. """ function isvalid_v(ϕ::JustRelax.RockRatio, i, j) nx, ny = size(ϕ.Vx) - j_bot = max(j - 1, 1) - j0 = min(j, ny) - vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - - nx, ny = size(ϕ.Vy) - i_left = max(i - 1, 1) - i0 = min(i, nx) - vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) - v = vx * vy + j_bot = max(j - 1, 1) + j0 = min(j, ny) + vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + + nx, ny = size(ϕ.Vy) + i_left = max(i - 1, 1) + i0 = min(i, nx) + vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + v = vx * vy return v * (ϕ.vertex[i, j] > 0) end @@ -126,8 +136,8 @@ Check if `ϕ.Vx[inds...]` is a not a nullspace. - `inds`: Cartesian indices to check. """ function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i-1, j] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j+1] > 0) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) cv = c * v return cv * (ϕ.Vx[i, j] > 0) end @@ -142,8 +152,8 @@ Check if `ϕ.Vy[inds...]` is a not a nullspace. - `inds`: Cartesian indices to check. """ function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) + c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) + v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) cv = c * v return cv * (ϕ.Vy[i, j] > 0) -end \ No newline at end of file +end diff --git a/src/variational_stokes/types.jl b/src/variational_stokes/types.jl index 1b3c85372..aa8ef3642 100644 --- a/src/variational_stokes/types.jl +++ b/src/variational_stokes/types.jl @@ -1,20 +1,29 @@ abstract type AbstractMask end -struct RockRatio{T, N} <: AbstractMask +struct RockRatio{T,N} <: AbstractMask center::T vertex::T Vx::T Vy::T - Vz::Union{Nothing, T} - yz::Union{Nothing, T} - xz::Union{Nothing, T} - xy::Union{Nothing, T} - - function RockRatio(center::AbstractArray{F, N}, vertex::T, Vx::T, Vy::T, Vz::Union{Nothing, T}, yz::Union{Nothing, T}, xz::Union{Nothing, T}, xy::Union{Nothing, T}) where {F, N, T} - new{T, N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) + Vz::Union{Nothing,T} + yz::Union{Nothing,T} + xz::Union{Nothing,T} + xy::Union{Nothing,T} + + function RockRatio( + center::AbstractArray{F,N}, + vertex::T, + Vx::T, + Vy::T, + Vz::Union{Nothing,T}, + yz::Union{Nothing,T}, + xz::Union{Nothing,T}, + xy::Union{Nothing,T}, + ) where {F,N,T} + return new{T,N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) end end -RockRatio(ni::NTuple{N, Integer}) where N = RockRatio(ni...) +RockRatio(ni::NTuple{N,Integer}) where {N} = RockRatio(ni...) -Adapt.@adapt_structure RockRatio \ No newline at end of file +Adapt.@adapt_structure RockRatio From 264d1597eeb79c6162fc855f13950354c8ca3733 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 14:30:21 +0200 Subject: [PATCH 09/77] remove files --- Cavity2D.jl | 396 --------------------------------------------- Duretz2016.jl | 438 -------------------------------------------------- 2 files changed, 834 deletions(-) delete mode 100644 Cavity2D.jl delete mode 100644 Duretz2016.jl diff --git a/Cavity2D.jl b/Cavity2D.jl deleted file mode 100644 index d7e1d57a4..000000000 --- a/Cavity2D.jl +++ /dev/null @@ -1,396 +0,0 @@ -using Printf -using JustRelax, JustRelax.JustRelax2D -import JustRelax.JustRelax2D as JR -const backend_JR = CPUBackend - -using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) - -# Load script dependencies -using LinearAlgebra, GeoParams, GLMakie - -include("mask.jl") -include("MiniKernels.jl") - -# Velocity helper grids for the particle advection -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - @parallel f_x(A, B) -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 - -function init_phases!(phases, particles) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index) - r = 5 - - @inbounds for ip in cellaxes(phases) - # quick escape - @index(index[ip, i, j]) == 0 && continue - - x = @index px[ip, i, j] - depth = abs(@index py[ip, i, j]) - @index phases[ip, i, j] = 2.0 - - if 0e0 ≤ depth ≤ 10e0 || ((x - 25)^2 + (depth - 55)^2 ≤ r^2) - @index phases[ip, i, j] = 1.0 - end - - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) -end -## END OF HELPER FUNCTION ------------------------------------------------------------ - -# (Path)/folder where output data and figures are stored -n = 101 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 5e0 # thickness of sticky air layer - ly = 50e0 + thick_air # domain length in y - lx = 50e0 # 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) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - # Name = "Air", - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=2.7e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Rock", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=2.7e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - ) - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 120, 120, 100 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi, di, ni - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - init_phases!(pPhases, particles) - phase_ratios = PhaseRatios(backend, length(rheology), ni) - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # ---------------------------------------------------- - - # RockRatios - air_phase = 1 - ϕ = RockRatio(ni...) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0*π, r=0.7, CFL = 0.95 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) - - # Boundary conditions - flow_bcs = VelocityBoundaryConditions(; - free_slip = (left = true, right = true, top = true, bot = false), - no_slip = (left = false, right = false, top = false, bot = true), - ) - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "FreeSurfacePlume" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - iterMax = 5e3 - nout = 1e3 - viscosity_cutoff = (-Inf, Inf) - free_surface = false - ητ = @zeros(ni...) - while it < 1 - - ## variational solver - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - relλ = 0.2 - θ = deepcopy(stokes.P) - λ = @zeros(ni...) - λv = @zeros(ni .+ 1...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - # Vx_on_Vy = @zeros(size(stokes.V.Vy)) - - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - - # for a in 1:100 - while iter ≤ iterMax - err < ϵ && break - # for _ in 1:100 - JR.compute_maxloc!(ητ, η; window=(1, 1)) - # update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - @parallel (@idx ni) compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ - ) - - JR.update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) JR.compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - # if rem(iter, nout) == 0 - # @copy η0 η - # end - # if do_visc - # update_viscosity!( - # stokes, - # phase_ratios, - # args, - # rheology, - # viscosity_cutoff; - # relaxation=viscosity_relaxation, - # ) - # end - - @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - # update_halo!(stokes.τ.xy) - - # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( - # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] - # ) - - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni.+1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - # ϕ.Vx, - # ϕ.Vy, - _di..., - ) - # apply boundary conditions - # velocity2displacement!(stokes, dt) - # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - # end - # f,ax,h=heatmap(stokes.V.Vy) - # # f,ax,h=heatmap(stokes.V.Vx) - # Colorbar(f[1,2], h, label="Vy"); f - # update_halo!(@velocity(stokes)...) - # end - - iter += 1 - - if iter % nout == 0 && iter > 1 - # er_η = norm_mpi(@.(log10(η) - log10(η0))) - # er_η < 1e-3 && (do_visc = false) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), - norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - # end - isnan(err) && error("NaN(s)") - isinf(err) && error("Inf(s)") - end - end - heatmap(stokes.V.Vy) - - dt = compute_dt(stokes, di) / 2 - # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - @show it += 1 - t += dt - - if it == 1 || rem(it, 1) == 0 - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 5 - fig = Figure(size = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - # arrows!( - # ax, - # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - # color = :red, - # ) - # heatmap!(ax, stokes.V.Vy) - # heatmap!(ax, stokes.τ.xx) - - # ind = iszero.(stokes.V.Vy) - # stokes.V.Vy[ind] .= NaN - heatmap!(ax, stokes.V.Vy) - display(fig) - # save(joinpath(figdir, "$(it).png"), fig) - - end - end - return stokes, ϕ -end -# ## END OF MAIN SCRIPT ---------------------------------------------------------------- -stoke, ϕ = main(igg, nx, ny); - -# # # (Path)/folder where output data and figures are stored -# # n = 100 -# # nx = n -# # ny = n -# # igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid -# # IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -# # else -# # igg -# # end - diff --git a/Duretz2016.jl b/Duretz2016.jl deleted file mode 100644 index 674a7dcb9..000000000 --- a/Duretz2016.jl +++ /dev/null @@ -1,438 +0,0 @@ -using Printf -using JustRelax, JustRelax.JustRelax2D -import JustRelax.JustRelax2D as JR -const backend_JR = CPUBackend - -using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) - -# Load script dependencies -using LinearAlgebra, GeoParams, GLMakie - -include("mask.jl") -include("MiniKernels.jl") - -# Velocity helper grids for the particle advection -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - @parallel f_x(A, B) -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 - -topo_fun(x, r) = -√(r^2 - x^2) - -function init_phases!(phases, particles) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index) - r=0.5 - rinc=0.1 - - @inbounds for ip in cellaxes(phases) - # quick escape - @index(index[ip, i, j]) == 0 && continue - - x = @index px[ip, i, j] - depth = (@index py[ip, i, j]) - - # h = √(r^2 - x^2) - 0.4330127018922193*2 - h = -√(r^2 - x^2) - if depth ≤ h - @index phases[ip, i, j] = 2.0 - else - @index phases[ip, i, j] = 1.0 - end - - # h = √(r^2 - x^2) - 0.55*2 - # h = -√(r^2 - x^2) - 0.2 - # if depth ≤ h - # @index phases[ip, i, j] = 1.0 - # end - - # if (x^2 + (depth + 0.75)^2 ≤ rinc^2) - # @index phases[ip, i, j] = 1.0 - # end - - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) -end - -# init_phases!(pPhases, particles) -# phase_ratios = PhaseRatios(backend, length(rheology), ni) -# update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) -# p=[argmax(p) for p in phase_ratios.center] -# heatmap(p) - - -## END OF HELPER FUNCTION ------------------------------------------------------------ - -# (Path)/folder where output data and figures are stored -n = 101 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 0 # thickness of sticky air layer - ly = 0.5 # domain length in y - lx = 0.5 # 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.25, -0.75 # origin coordinates (15km f sticky air layer) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - L = 0.5 - H = 0.5 - r = 0.5 - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=1), - CompositeRheology = CompositeRheology((LinearViscous(; η=1),)), - Gravity = ConstantGravity(; g=1), - ), - # Name = "Mantle", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=1), - CompositeRheology = CompositeRheology((LinearViscous(; η=1),)), - Gravity = ConstantGravity(; g=1), - ), - ) - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 40, 70, 15 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi, di, ni - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - init_phases!(pPhases, particles) - phase_ratios = PhaseRatios(backend, length(rheology), ni) - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # ---------------------------------------------------- - - # RockRatios - air_phase = 1 - ϕ = RockRatio(ni...) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) - - # Boundary conditions - flow_bcs = VelocityBoundaryConditions(; - free_slip = (left = true, right = true, top = true, bot = true), - free_surface = true - ) - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "Duretz2016" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - iterMax = 5e3 - nout = 1e3 - viscosity_cutoff = (-Inf, Inf) - free_surface = false - ητ = @zeros(ni...) - # while it < 1 - - ## variational solver - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - relλ = 0.2 - θ = deepcopy(stokes.P) - λ = @zeros(ni...) - λv = @zeros(ni .+ 1...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - # Vx_on_Vy = @zeros(size(stokes.V.Vy)) - - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - - while iter ≤ iterMax - err < ϵ && break - # for _ in 1:100 - JR.compute_maxloc!(ητ, η; window=(1, 1)) - # update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - @parallel (@idx ni) compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - ϕ, - dt, - pt_stokes.r, - pt_stokes.θ_dτ - ) - - JR.update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) JR.compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - # if rem(iter, nout) == 0 - # @copy η0 η - # end - # if do_visc - # update_viscosity!( - # stokes, - # phase_ratios, - # args, - # rheology, - # viscosity_cutoff; - # relaxation=viscosity_relaxation, - # ) - # end - - @parallel (@idx ni .+ 1) update_stresses_center_vertex_ps!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - # update_halo!(stokes.τ.xy) - - # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) JR.interp_Vx∂ρ∂x_on_Vy!( - # Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] - # ) - - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni.+1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - # ϕ.Vx, - # ϕ.Vy, - _di..., - ) - # apply boundary conditions - # velocity2displacement!(stokes, dt) - # JR.free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - # end - # f,ax,h=heatmap(stokes.V.Vy) - # # f,ax,h=heatmap(stokes.V.Vx) - # Colorbar(f[1,2], h, label="Vy"); f - # update_halo!(@velocity(stokes)...) - # end - - iter += 1 - - if iter % nout == 0 && iter > 1 - # er_η = norm_mpi(@.(log10(η) - log10(η0))) - # er_η < 1e-3 && (do_visc = false) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Rx), - norm(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - # if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - # end - isnan(err) && error("NaN(s)") - isinf(err) && error("Inf(s)") - end - end - - dt = compute_dt(stokes, di) / 2 - # ------------------------------ - heatmap(stokes.τ.xy) - - # # Advection -------------------- - # # advect particles in space - # advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # # advect particles in memory - # move_particles!(particles, xvi, particle_args) - # # check if we need to inject particles - # inject_particles_phase!(particles, pPhases, (), (), xvi) - # # update phase ratios - # update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # update_rock_ratio!(ϕ, phase_ratios, air_phase) - - # @show it += 1 - # t += dt - - # if it == 1 || rem(it, 1) == 0 - # velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - # nt = 5 - # fig = Figure(size = (900, 900), title = "t = $t") - # ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - # # arrows!( - # # ax, - # # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - # # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - # # color = :red, - # # ) - # fig - # save(joinpath(figdir, "$(it).png"), fig) - - # end - # # end - # return nothing - # heatmap(stokes.τ.xy) - ind = iszero.(stokes.P) - stokes.P[ind] .= NaN - f,ax,=heatmap(xci...,stokes.P) - ind = iszero.(stokes.V.Vy) - stokes.V.Vy[ind] .= NaN - f,ax,=heatmap(xci[1],xvi[2], stokes.V.Vy[2:end-1,:]) - - px, py = particles.coords - scatter!( - px.data[:], py.data[:],color=pPhases.data[:], - markersize=3) - h = @. -√(0.5^2 - xci[1]^2) - lines!(ax,xci[1], h, color=:red) - xlims!(ax, -0.25, -0.15) - ylims!(ax, -0.5, -0.4) - f - # ind = iszero.(stokes.V.Vx) - # stokes.V.Vx[ind] .= NaN - # heatmap(stokes.V.Vx) - # stokes -end -# ## END OF MAIN SCRIPT ---------------------------------------------------------------- -main(igg, nx, ny) - - -# heatmap(stokes.τ.xy) -# # ind = iszero.(stokes.P) -# # stokes.P[ind] .= NaN -# heatmap(stokes.P) -# ind = iszero.(stokes.V.Vy) -# stokes.V.Vy[ind] .= NaN -# heatmap(stokes.V.Vy) -# # ind = iszero.(stokes.V.Vx) -# # stokes.V.Vx[ind] .= NaN -# # heatmap(stokes.V.Vx) \ No newline at end of file From bda0a76f56a6234d8004b049f9445a609c409a8a Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 16:20:41 +0200 Subject: [PATCH 10/77] delete wrong files --- src/VariationalStokes/MiniKernels.jl | 20 -- src/VariationalStokes/PressureKernels.jl | 74 -------- src/VariationalStokes/Stokes2D.jl | 221 ----------------------- src/VariationalStokes/StressKernels.jl | 142 --------------- src/VariationalStokes/VelocityKernels.jl | 72 -------- src/VariationalStokes/mask.jl | 159 ---------------- src/VariationalStokes/types.jl | 29 --- 7 files changed, 717 deletions(-) delete mode 100644 src/VariationalStokes/MiniKernels.jl delete mode 100644 src/VariationalStokes/PressureKernels.jl delete mode 100644 src/VariationalStokes/Stokes2D.jl delete mode 100644 src/VariationalStokes/StressKernels.jl delete mode 100644 src/VariationalStokes/VelocityKernels.jl delete mode 100644 src/VariationalStokes/mask.jl delete mode 100644 src/VariationalStokes/types.jl diff --git a/src/VariationalStokes/MiniKernels.jl b/src/VariationalStokes/MiniKernels.jl deleted file mode 100644 index 5584994cc..000000000 --- a/src/VariationalStokes/MiniKernels.jl +++ /dev/null @@ -1,20 +0,0 @@ -# masked versions -for fn in (:center, :next, :left, :right, :back, :front) - @eval begin - Base.@propagate_inbounds @inline ($fn)( - A::T, ϕ::T, inds::Vararg{Integer,N} - ) where {T<:AbstractArray,N} = ($fn)(A, inds...) * ($fn)(ϕ, inds...) - end -end - -# finite differences -Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = - (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = - (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy -Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer,N}) where {N,T} = - (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz -Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = - (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx -Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = - (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy diff --git a/src/VariationalStokes/PressureKernels.jl b/src/VariationalStokes/PressureKernels.jl deleted file mode 100644 index 86c49f6cf..000000000 --- a/src/VariationalStokes/PressureKernels.jl +++ /dev/null @@ -1,74 +0,0 @@ -function compute_P!( - P, - P0, - RP, - ∇V, - η, - rheology::NTuple{N,MaterialParams}, - phase_ratio::JustPIC.PhaseRatios, - ϕ::JustRelax.RockRatio, - dt, - r, - θ_dτ; - ΔTc=nothing, - kwargs..., -) where {N} - ni = size(P) - @parallel (@idx ni) compute_P_kernel!( - P, P0, RP, ∇V, η, rheology, phase_ratio.center, ϕ, dt, r, θ_dτ, ΔTc - ) - return nothing -end - -@parallel_indices (I...) function compute_P!( - P, - P0, - RP, - ∇V, - η, - rheology::NTuple{N,MaterialParams}, - phase_ratio, - ϕ::JustRelax.RockRatio, - dt, - r, - θ_dτ, - ::Nothing, -) where {N} - if isvalid_c(ϕ, I...) - K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...])) - RP[I...], P[I...] = _compute_P!( - P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ - ) - else - RP[I...] = P[I...] = zero(eltype(P)) - end - return nothing -end - -@parallel_indices (I...) function compute_P_kernel!( - P, - P0, - RP, - ∇V, - η, - rheology::NTuple{N,MaterialParams}, - phase_ratio::C, - ϕ::JustRelax.RockRatio, - dt, - r, - θ_dτ, - ΔTc, - ::Nothing, -) where {N,C<:JustRelax.CellArray} - if isvalid_c(ϕ, I...) - phase_ratio_I = phase_ratio[I...] - K = fn_ratio(get_bulk_modulus, rheology, phase_ratio_I) - α = fn_ratio(get_thermal_expansion, rheology, phase_ratio_I) - RP[I...], P[I...] = _compute_P!( - P[I...], P0[I...], ∇V[I...], ΔTc[I...], α, η[I...], K, dt, r, θ_dτ - ) - else - RP[I...] = P[I...] = zero(eltype(P)) - end - return nothing -end diff --git a/src/VariationalStokes/Stokes2D.jl b/src/VariationalStokes/Stokes2D.jl deleted file mode 100644 index 6121b698b..000000000 --- a/src/VariationalStokes/Stokes2D.jl +++ /dev/null @@ -1,221 +0,0 @@ -## 2D VISCO-ELASTIC STOKES SOLVER - -# backend trait -function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs) - solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) - return nothing -end - -# entry point for extensions -function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) - return _solve_VS!(stokes, args...; kwargs...) -end - -function _solve_VS!( - stokes::JustRelax.StokesArrays, - pt_stokes, - di::NTuple{2,T}, - flow_bcs::AbstractFlowBoundaryConditions, - ρg, - phase_ratios::JustPIC.PhaseRatios, - ϕ::JustRelax.RockRatio, - rheology, - args, - dt, - igg::IGG; - viscosity_cutoff=(-Inf, Inf), - viscosity_relaxation=1e-2, - iterMax=50e3, - iterMin=1e2, - nout=500, - b_width=(4, 4, 0), - verbose=true, - kwargs..., -) where {T} - - # unpack - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - relλ = 0.2 - θ = deepcopy(stokes.P) - λ = @zeros(ni...) - λv = @zeros(ni .+ 1...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) - displacement2velocity!(stokes, dt, flow_bcs) - - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) - - while iter ≤ iterMax - iterMin < iter && err < ϵ && break - - wtime0 += @elapsed begin - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) - - compute_P!( - θ, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - dt, - r, - θ_dτ, - args, - ) - - update_ρg!(ρg[2], phase_ratios, rheology, args) - - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - update_viscosity!( - stokes, - phase_ratios, - args, - rheology, - viscosity_cutoff; - relaxation=viscosity_relaxation, - ) - - @parallel (@idx ni .+ 1) update_stresses_center_vertex!( - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - @tensor_center(stokes.τ), - (stokes.τ.xy,), - @tensor_center(stokes.τ_o), - (stokes.τ_o.xy,), - θ, - stokes.P, - stokes.viscosity.η, - λ, - λv, - stokes.τ.II, - stokes.viscosity.η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_ratios.center, - phase_ratios.vertex, - ϕ, - ) - update_halo!(stokes.τ.xy) - - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni .+ 1) compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - _di..., - ) - # apply boundary conditions - velocity2displacement!(stokes, dt) - free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) - update_halo!(@velocity(stokes)...) - # end - end - - iter += 1 - - if iter % nout == 0 && iter > 1 - errs = ( - norm_mpi(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) / - length(stokes.R.Rx), - norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / - length(stokes.R.Ry), - norm_mpi(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum_mpi(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - # compute vorticity - @parallel (@idx ni .+ 1) compute_vorticity!( - stokes.ω.xy, @velocity(stokes)..., inv.(di)... - ) - - # accumulate plastic strain tensor - @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end diff --git a/src/VariationalStokes/StressKernels.jl b/src/VariationalStokes/StressKernels.jl deleted file mode 100644 index 6b1d848a3..000000000 --- a/src/VariationalStokes/StressKernels.jl +++ /dev/null @@ -1,142 +0,0 @@ -# 2D kernel -@parallel_indices (I...) function update_stresses_center_vertex!( - ε::NTuple{3,T}, # normal components @ centers; shear components @ vertices - ε_pl::NTuple{3}, # whole Voigt tensor @ centers - EII, # accumulated plastic strain rate @ centers - τ::NTuple{3}, # whole Voigt tensor @ centers - τshear_v::NTuple{1}, # shear tensor components @ vertices - τ_o::NTuple{3}, - τshear_ov::NTuple{1}, # shear tensor components @ vertices - Pr, - Pr_c, - η, - λ, - λv, - τII, - η_vep, - relλ, - dt, - θ_dτ, - rheology, - phase_center, - phase_vertex, - ϕ::JustRelax.RockRatio, -) where {T} - τxyv = τshear_v[1] - τxyv_old = τshear_ov[1] - ni = size(Pr) - Ic = clamped_indices(ni, I...) - - if isvalid_v(ϕ, I...) - # interpolate to ith vertex - Pv_ij = av_clamped(Pr, Ic...) - εxxv_ij = av_clamped(ε[1], Ic...) - εyyv_ij = av_clamped(ε[2], Ic...) - τxxv_ij = av_clamped(τ[1], Ic...) - τyyv_ij = av_clamped(τ[2], Ic...) - τxxv_old_ij = av_clamped(τ_o[1], Ic...) - τyyv_old_ij = av_clamped(τ_o[2], Ic...) - EIIv_ij = av_clamped(EII, Ic...) - - ## vertex - phase = @inbounds phase_vertex[I...] - is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( - rheology, EIIv_ij, phase - ) - _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) - Kv = fn_ratio(get_bulk_modulus, rheology, phase) - volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ - ηv_ij = av_clamped(η, Ic...) - dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) - - # stress increments @ vertex - dτxxv = - (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * - dτ_rv - dτyyv = - (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * - dτ_rv - dτxyv = - ( - -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + - 2.0 * ηv_ij * ε[3][I...] - ) * dτ_rv - τIIv_ij = - √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) - - # yield function @ center - Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) - # stress correction @ vertex - λv[I...] = - (1.0 - relλ) * λv[I...] + - relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) - dQdτxy = 0.5 * (τxyv[I...] + dτxyv) / τIIv_ij - τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[I...] * dQdτxy * dτ_rv - else - # stress correction @ vertex - τxyv[I...] += dτxyv - end - else - τxyv[I...] = zero(eltype(T)) - end - - ## center - if all(I .≤ ni) - if isvalid_c(ϕ, I...) - # Material properties - phase = @inbounds phase_center[I...] - _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) - is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase( - rheology, EII[I...], phase - ) - K = fn_ratio(get_bulk_modulus, rheology, phase) - volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ - ηij = η[I...] - dτ_r = 1.0 / (θ_dτ + ηij * _Gdt + 1.0) - - # cache strain rates for center calculations - τij, τij_o, εij = cache_tensors(τ, τ_o, ε, I...) - - # visco-elastic strain rates @ center - εij_ve = @. εij + 0.5 * τij_o * _Gdt - εII_ve = GeoParams.second_invariant(εij_ve) - # stress increments @ center - dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r - τII_ij = GeoParams.second_invariant(dτij .+ τij) - # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ - - if is_pl && !iszero(τII_ij) - # stress correction @ center - λ[I...] = - (1.0 - relλ) * λ[I...] + - relλ .* (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) - dQdτij = @. 0.5 * (τij + dτij) / τII_ij - εij_pl = λ[I...] .* dQdτij - dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r - τij = dτij .+ τij - setindex!.(τ, τij, I...) - setindex!.(ε_pl, εij_pl, I...) - τII[I...] = GeoParams.second_invariant(τij) - Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ - η_vep[I...] = 0.5 * τII_ij / εII_ve - else - # stress correction @ center - setindex!.(τ, dτij .+ τij, I...) - η_vep[I...] = ηij - τII[I...] = τII_ij - end - - Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) - else - Pr_c[I...] = zero(eltype(T)) - # τij, = cache_tensors(τ, τ_o, ε, I...) - dτij = zero(eltype(T)), zero(eltype(T)), zero(eltype(T)) - # setindex!.(τ, dτij .+ τij, I...) - setindex!.(τ, dτij, I...) - end - end - - return nothing -end diff --git a/src/VariationalStokes/VelocityKernels.jl b/src/VariationalStokes/VelocityKernels.jl deleted file mode 100644 index 6ff77e206..000000000 --- a/src/VariationalStokes/VelocityKernels.jl +++ /dev/null @@ -1,72 +0,0 @@ -@parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} -) where {T,N} - @inline d_xi(A) = _d_xi(A, _di[1], I...) - @inline d_yi(A) = _d_yi(A, _di[2], I...) - @inline d_zi(A) = _d_zi(A, _di[3], I...) - - f = d_xi, d_yi, d_zi - - if isvalid_c(ϕ, I...) - @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) - else - @inbounds ∇V[I...] = zero(T) - end - return nothing -end - -@parallel_indices (i, j) function compute_V!( - Vx::AbstractArray{T,2}, - Vy, - Rx, - Ry, - P, - τxx, - τyy, - τxy, - ηdτ, - ρgx, - ρgy, - ητ, - ϕ::JustRelax.RockRatio, - _dx, - _dy, -) where {T} - d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) - d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) - d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) - av_xa(A) = _av_xa(A, i, j) - av_ya(A) = _av_ya(A, i, j) - harm_xa(A) = _av_xa(A, i, j) - harm_ya(A) = _av_ya(A, i, j) - - if all((i, j) .< size(Vx) .- 1) - if isvalid_vx(ϕ, i + 1, j) - Rx[i, j] = - R_Vx = ( - -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - - av_xa(ρgx) - ) - Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) - else - Rx[i, j] = zero(T) - Vx[i + 1, j + 1] = zero(T) - end - end - - if all((i, j) .< size(Vy) .- 1) - if isvalid_vy(ϕ, i, j + 1) - Ry[i, j] = - R_Vy = - -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - - av_ya(ρgy) - Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) - else - Ry[i, j] = zero(T) - Vy[i + 1, j + 1] = zero(T) - end - end - - return nothing -end diff --git a/src/VariationalStokes/mask.jl b/src/VariationalStokes/mask.jl deleted file mode 100644 index 27ba90c86..000000000 --- a/src/VariationalStokes/mask.jl +++ /dev/null @@ -1,159 +0,0 @@ - -function RockRatio(nx, ny) - ni = nx, ny - center = @zeros(ni...) - vertex = @zeros(ni .+ 1...) - Vx = @zeros(nx+1, ny) # no ghost nodes! - Vy = @zeros(nx, ny+1) # no ghost nodes! - - return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) -end - -function RockRatio(nx, ny, nz) - ni = nx, ny, nz - center = @zeros(ni...) - vertex = @zeros(ni .+ 1...) - Vx = @zeros(nx+1, ny, nz) # no ghost nodes! - Vy = @zeros(nx, ny+1, nz) # no ghost nodes! - Vz = @zeros(nx, ny, nz+1) # no ghost nodes! - yz = @zeros(nx, ny + 1, nz + 1) - xz = @zeros(nx + 1, ny, nz + 1) - xy = @zeros(nx + 1, ny + 1, nz) - - return JustRelax.RockRatio(center, vertex, Vx, Vy, Vz, yz, xz, xy) -end - -@inline size_c(x::JustRelax.AbstractMask) = size(x.center) -@inline size_v(x::JustRelax.AbstractMask) = size(x.vertex) -@inline size_vx(x::JustRelax.AbstractMask) = size(x.Vx) -@inline size_vy(x::JustRelax.AbstractMask) = size(x.Vy) -@inline size_vz(x::JustRelax.AbstractMask) = size(x.Vz) -@inline size_yz(x::JustRelax.AbstractMask) = size(x.yz) -@inline size_xz(x::JustRelax.AbstractMask) = size(x.xz) -@inline size_xy(x::JustRelax.AbstractMask) = size(x.xy) - -""" - update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) - -Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The rock ratio object to be updated. -- `phase_ratios`: The ratios of different phases present. -- `air_phase`: The phase representing air. -""" -function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) - nvi = size_v(ϕ) - @parallel (@idx nvi) update_rock_ratio_cv!( - ϕ, phase_ratios.center, phase_ratios.vertex, air_phase - ) - @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) - return nothing -end - -@inline compute_rock_ratio( - phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} -) where {N} = 1 - @index phase_ratio[air_phase, I...] -@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[ - air_phase, I... -] - -@parallel_indices (I...) function update_rock_ratio_cv!( - ϕ, ratio_center, ratio_vertex, air_phase -) - if all(I .≤ size(ratio_center)) - ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) - end - ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) - return nothing -end - -@parallel_indices (I...) function update_rock_ratio_vel!( - ϕ::JustRelax.RockRatio{T,N} -) where {T,N} - # 2D - @inline av_x(A::AbstractArray{T,2}) where {T} = _av_xa(A, I...) - @inline av_y(A::AbstractArray{T,2}) where {T} = _av_ya(A, I...) - # 3D - @inline av_x(A::AbstractArray{T,3}) where {T} = _av_yz(A, I...) - @inline av_y(A::AbstractArray{T,3}) where {T} = _av_xz(A, I...) - @inline av_z(A::AbstractArray{T,3}) where {T} = _av_xy(A, I...) - - all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - end - return nothing -end - -""" - isvalid_c(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.center[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_c(ϕ::JustRelax.RockRatio, i, j) - vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) - vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) - v = vx * vy - return v * (ϕ.center[i, j] > 0) -end - -""" - isvalid_v(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.vertex[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_v(ϕ::JustRelax.RockRatio, i, j) - nx, ny = size(ϕ.Vx) - j_bot = max(j - 1, 1) - j0 = min(j, ny) - vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - - nx, ny = size(ϕ.Vy) - i_left = max(i - 1, 1) - i0 = min(i, nx) - vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) - v = vx * vy - return v * (ϕ.vertex[i, j] > 0) -end - -""" - isvalid_vx(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vx[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) - cv = c * v - return cv * (ϕ.Vx[i, j] > 0) -end - -""" - isvalid_vy(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vy[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) - cv = c * v - return cv * (ϕ.Vy[i, j] > 0) -end diff --git a/src/VariationalStokes/types.jl b/src/VariationalStokes/types.jl deleted file mode 100644 index aa8ef3642..000000000 --- a/src/VariationalStokes/types.jl +++ /dev/null @@ -1,29 +0,0 @@ -abstract type AbstractMask end - -struct RockRatio{T,N} <: AbstractMask - center::T - vertex::T - Vx::T - Vy::T - Vz::Union{Nothing,T} - yz::Union{Nothing,T} - xz::Union{Nothing,T} - xy::Union{Nothing,T} - - function RockRatio( - center::AbstractArray{F,N}, - vertex::T, - Vx::T, - Vy::T, - Vz::Union{Nothing,T}, - yz::Union{Nothing,T}, - xz::Union{Nothing,T}, - xy::Union{Nothing,T}, - ) where {F,N,T} - return new{T,N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) - end -end - -RockRatio(ni::NTuple{N,Integer}) where {N} = RockRatio(ni...) - -Adapt.@adapt_structure RockRatio From f1d2530e37a699ca4717f51abed71db842241d15 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 16:21:09 +0200 Subject: [PATCH 11/77] up miniapp --- .../RayleighTaylor2D.jl | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 203d05eb2..db3c5db1c 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -1,11 +1,14 @@ +using CUDA using JustRelax, JustRelax.JustRelax2D -const backend_JR = CPUBackend +const backend_JR = CUDABackend +# const backend_JR = CPUBackend using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend +const backend = CUDABackend +# const backend = JustPIC.CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) +@init_parallel_stencil(CUDA, Float64, 2) # Load script dependencies using LinearAlgebra, GeoParams, GLMakie @@ -133,7 +136,7 @@ function main(igg, nx, ny) # RockRatios air_phase = 1 - ϕ = RockRatio(ni...) + ϕ = RockRatio(backend, ni) update_rock_ratio!(ϕ, phase_ratios, air_phase) # STOKES --------------------------------------------- @@ -180,7 +183,7 @@ function main(igg, nx, ny) viscosity_cutoff = (-Inf, Inf) free_surface = false ητ = @zeros(ni...) - while it < 1 + while it < 1000 ## variational solver # Stokes solver ---------------- @@ -229,7 +232,7 @@ function main(igg, nx, ny) fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) - scatter!(ax, px.data[:].*1e-3, py.data[:].*1e-3, color = pPhases.data[:], colormap = :grayC) + scatter!(ax, Array(px.data[:]).*1e-3, Array(py.data[:]).*1e-3, color =Array(pPhases.data[:]), colormap = :grayC) arrows!( ax, xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., @@ -243,6 +246,6 @@ function main(igg, nx, ny) end return nothing end -# ## END OF MAIN SCRIPT ---------------------------------------------------------------- -main(igg, nx, ny) +## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) \ No newline at end of file From 1afc888e260cd539a946ccb59003655977f8308d Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 22 Oct 2024 16:21:23 +0200 Subject: [PATCH 12/77] add GPU compatibility --- src/ext/AMDGPU/2D.jl | 17 ++++++++ src/ext/AMDGPU/3D.jl | 18 +++++++++ src/ext/CUDA/2D.jl | 22 +++++++++++ src/ext/CUDA/3D.jl | 17 ++++++++ src/variational_stokes/Stokes2D.jl | 2 +- src/variational_stokes/VelocityKernels.jl | 48 +++++++++++++++++++---- src/variational_stokes/mask.jl | 8 +++- src/variational_stokes/types.jl | 27 +++++++------ 8 files changed, 136 insertions(+), 23 deletions(-) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index b48fcb197..712e3b17a 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -55,6 +55,10 @@ function JR2D.WENO5( return WENO5(method, tuple(ni...)) end +function JR2D.RockRatio(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N} + return RockRatio(ni...) +end + function JR2D.PTThermalCoeffs( ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -286,6 +290,10 @@ function JR2D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) end +function JR2D.solve_VariationalStokes!(::AMDGPUBackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + function JR2D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs) return _heatdiffusion_PT!(thermal, args...; kwargs...) end @@ -391,4 +399,13 @@ function JR2D.rotate_stress_particles!( return nothing end +# rock ratios + +function JR2D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase +) where {T,nD,N} + update_rock_ratio!(ϕ, phase_ratios, air_phase) + return nothing +end + end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 76fce83d3..6a566f153 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -55,6 +55,10 @@ function JR3D.WENO5( return WENO5(method, tuple(ni...)) end +function JR3D.RockRatio(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N} + return RockRatio(ni...) +end + function JR3D.PTThermalCoeffs( ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -285,11 +289,16 @@ end function displacement2velocity!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, dt) return _displacement2velocity!(stokes, dt) end + # Solvers function JR3D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) end +function JR3D.solve_VariationalStokes!(::AMDGPUBackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + function JR3D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs) return _heatdiffusion_PT!(thermal, args...; kwargs...) end @@ -407,4 +416,13 @@ function JR3D.rotate_stress_particles!( return nothing end +# rock ratios + +function JR3D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase +) where {T,nD,N} + update_rock_ratio!(ϕ, phase_ratios, air_phase) + return nothing +end + end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index c4315698e..0b83f7e65 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -30,6 +30,7 @@ include("../../common.jl") include("../../stokes/Stokes2D.jl") # Types + function JR2D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} return StokesArrays(ni) end @@ -46,6 +47,14 @@ function JR2D.WENO5(::Type{CUDABackend}, method::Val{T}, ni::NTuple{N,Integer}) return WENO5(method, tuple(ni...)) end +function JR2D.RockRatio(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} + return RockRatio(ni...) +end + +function JR2D.RockRatio(::Type{CUDABackend}, ni::Vararg{Integer,N}) where {N} + return RockRatio(ni...) +end + function JR2D.PTThermalCoeffs( ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -270,6 +279,10 @@ function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) end +function JR2D.solve_VariationalStokes!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) return _heatdiffusion_PT!(thermal, args...; kwargs...) end @@ -372,4 +385,13 @@ function JR2D.rotate_stress_particles!( return nothing end +# rock ratios + +function JR2D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, air_phase) + return nothing +end + end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 4de3b3529..c47254502 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -45,6 +45,10 @@ function JR3D.WENO5(::Type{CUDABackend}, method::Val{T}, ni::NTuple{N,Integer}) return WENO5(method, tuple(ni...)) end +function JR3D.RockRatio(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} + return RockRatio(ni...) +end + function JR3D.PTThermalCoeffs( ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -289,6 +293,10 @@ function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) end +function JR3D.solve_VariationalStokes!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) return _heatdiffusion_PT!(thermal, args...; kwargs...) end @@ -404,4 +412,13 @@ function JR3D.rotate_stress_particles!( return nothing end +# rock ratios + +function JR3D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, air_phase) + return nothing +end + end diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 6121b698b..fcff656bd 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -90,7 +90,7 @@ function _solve_VS!( compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., ϕ, _di...) compute_P!( θ, diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 6ff77e206..eb1ae8e1c 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,20 +1,54 @@ @parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} -) where {T,N} - @inline d_xi(A) = _d_xi(A, _di[1], I...) - @inline d_yi(A) = _d_yi(A, _di[2], I...) - @inline d_zi(A) = _d_zi(A, _di[3], I...) + ∇V::AbstractArray{T,2}, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy +) where {T} + @inline d_xi(A) = _d_xi(A, _dx, I...) + @inline d_yi(A) = _d_yi(A, _dy, I...) - f = d_xi, d_yi, d_zi + if isvalid_c(ϕ, I...) + @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + else + @inbounds ∇V[I...] = zero(T) + end + return nothing +end + +@parallel_indices (I...) function compute_∇V!( + ∇V::AbstractArray{T,2}, Vx, Vy, Vz, ϕ::JustRelax.RockRatio, _dx, _dy, _dz +) where {T} + @inline d_xi(A) = _d_xi(A, _dx, I...) + @inline d_yi(A) = _d_yi(A, _dy, I...) + @inline d_zi(A) = _d_zi(A, _dz, I...) if isvalid_c(ϕ, I...) - @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N) + @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) else @inbounds ∇V[I...] = zero(T) end return nothing end +# @parallel_indices (I...) function compute_∇V!( +# ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} +# ) where {T,N} +# # @inline d_xi(A) = _d_xi(A, _di, I...) +# # @inline d_yi(A) = _d_yi(A, _di, I...) +# # @inline d_zi(A) = _d_zi(A, _di, I...) + +# # f = d_xi, d_yi, d_zi +# f = _d_xi, _d_yi, _d_zi + +# if isvalid_c(ϕ, I...) +# v = zero(T) +# for i in 1:N +# v += f[i](V[i], _di[I], I...) +# end +# @inbounds ∇V[I...] = v +# else +# @inbounds ∇V[I...] = zero(T) +# end +# return nothing +# end + @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 27ba90c86..6c3abc4de 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -1,3 +1,6 @@ +function RockRatio(::Type{CPUBackend}, ni::NTuple{N,Integer}) where {N} + return RockRatio(ni...) +end function RockRatio(nx, ny) ni = nx, ny @@ -5,8 +8,8 @@ function RockRatio(nx, ny) vertex = @zeros(ni .+ 1...) Vx = @zeros(nx+1, ny) # no ghost nodes! Vy = @zeros(nx, ny+1) # no ghost nodes! - - return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing) + dummy = @zeros(1, 1) # because it cant be a Union{T, Nothing} type on the GPU.... + return JustRelax.RockRatio(center, vertex, Vx, Vy, dummy, dummy, dummy, dummy) end function RockRatio(nx, ny, nz) @@ -54,6 +57,7 @@ end @inline compute_rock_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} ) where {N} = 1 - @index phase_ratio[air_phase, I...] + @inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[ air_phase, I... ] diff --git a/src/variational_stokes/types.jl b/src/variational_stokes/types.jl index aa8ef3642..f9dd2dd82 100644 --- a/src/variational_stokes/types.jl +++ b/src/variational_stokes/types.jl @@ -5,25 +5,26 @@ struct RockRatio{T,N} <: AbstractMask vertex::T Vx::T Vy::T - Vz::Union{Nothing,T} - yz::Union{Nothing,T} - xz::Union{Nothing,T} - xy::Union{Nothing,T} + Vz::T + yz::T + xz::T + xy::T function RockRatio( - center::AbstractArray{F,N}, - vertex::T, - Vx::T, - Vy::T, - Vz::Union{Nothing,T}, - yz::Union{Nothing,T}, - xz::Union{Nothing,T}, - xy::Union{Nothing,T}, + center::AbstractArray{F,N}, vertex::T, Vx::T, Vy::T, Vz::T, yz::T, xz::T, xy::T ) where {F,N,T} return new{T,N}(center, vertex, Vx, Vy, Vz, yz, xz, xy) end end -RockRatio(ni::NTuple{N,Integer}) where {N} = RockRatio(ni...) +RockRatio(::Type{CPUBackend}, ni::NTuple{N,Number}) where {N} = RockRatio(ni...) + +function RockRatio(::Number, ::Number) + throw(ArgumentError("RockRatio dimensions must be given as integers")) +end + +function RockRatio(::Number, ::Number, ::Number) + throw(ArgumentError("RockRatio dimensions must be given as integers")) +end Adapt.@adapt_structure RockRatio From 05b85ef19e2f08f9d44cab5cf1222aa956ef1510 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 23 Oct 2024 11:11:01 +0200 Subject: [PATCH 13/77] try out something --- src/variational_stokes/mask.jl | 66 ++++++++++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 6c3abc4de..21402d712 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -143,7 +143,8 @@ function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) cv = c * v - return cv * (ϕ.Vx[i, j] > 0) + # return cv * (ϕ.Vx[i, j] > 0) + return ϕ.Vx[i, j] > 0 end """ @@ -159,5 +160,66 @@ function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) cv = c * v - return cv * (ϕ.Vy[i, j] > 0) + # return cv * (ϕ.Vy[i, j] > 0) + return ϕ.Vy[i, j] > 0 end + + + +# """ +# isvalid_c(ϕ::JustRelax.RockRatio, inds...) + +# Check if `ϕ.center[inds...]` is a not a nullspace. + +# # Arguments +# - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +# - `inds`: Cartesian indices to check. +# """ +# function isvalid_c(ϕ::JustRelax.RockRatio, i, j) +# return (ϕ.center[i, j] > 0) +# end + +# """ +# isvalid_v(ϕ::JustRelax.RockRatio, inds...) + +# Check if `ϕ.vertex[inds...]` is a not a nullspace. + +# # Arguments +# - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +# - `inds`: Cartesian indices to check. +# """ +# function isvalid_v(ϕ::JustRelax.RockRatio, i, j) +# return (ϕ.vertex[i, j] > 0) +# end + +# """ +# isvalid_vx(ϕ::JustRelax.RockRatio, inds...) + +# Check if `ϕ.Vx[inds...]` is a not a nullspace. + +# # Arguments +# - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +# - `inds`: Cartesian indices to check. +# """ +# function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) +# c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0) +# v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0) +# cv = c || v +# return cv || (ϕ.Vx[i, j] > 0) +# end + +# """ +# isvalid_vy(ϕ::JustRelax.RockRatio, inds...) + +# Check if `ϕ.Vy[inds...]` is a not a nullspace. + +# # Arguments +# - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +# - `inds`: Cartesian indices to check. +# """ +# function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) +# c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) +# v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) +# cv = c || v +# return cv || (ϕ.Vy[i, j] > 0) +# end From f2d612b992e3aa42f790cb100792445902ec906e Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Fri, 25 Oct 2024 10:49:02 +0200 Subject: [PATCH 14/77] need to mask bouyancy forces as well --- src/MiniKernels.jl | 6 +- src/variational_stokes/MiniKernels.jl | 41 +++++ src/variational_stokes/VelocityKernels.jl | 6 +- src/variational_stokes/mask.jl | 207 +++++++++++----------- 4 files changed, 150 insertions(+), 110 deletions(-) diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl index 22f3e05aa..98120fcdc 100644 --- a/src/MiniKernels.jl +++ b/src/MiniKernels.jl @@ -155,7 +155,7 @@ end ## Because mysum(::generator) does not work inside CUDA kernels... @inline mysum(A, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ranges...) -@inline function mysum(f::F, A, ranges_i) where {F<:Function} +@inline function mysum(f::F, A::AbstractArray, ranges_i) where {F<:Function} s = 0.0 for i in ranges_i s += f(A[i]) @@ -163,7 +163,7 @@ end return s end -@inline function mysum(f::F, A, ranges_i, ranges_j) where {F<:Function} +@inline function mysum(f::F, A::AbstractArray, ranges_i, ranges_j) where {F<:Function} s = 0.0 for i in ranges_i, j in ranges_j s += f(A[i, j]) @@ -171,7 +171,7 @@ end return s end -@inline function mysum(f::F, A, ranges_i, ranges_j, ranges_k) where {F<:Function} +@inline function mysum(f::F, A::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function} s = 0.0 for i in ranges_i, j in ranges_j, k in ranges_k s += f(A[i, j, k]) diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 5584994cc..4fdf7c43b 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -18,3 +18,44 @@ Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) w (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy + +# averages +Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j) where {T<:T2} = + 0.25 * mysum(A, ϕ, (i+1):(i+2), (j+1):(j+2)) +Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j) where {T<:T2} = + 0.25 * mysum(A, ϕ, (i):(i+1), (j):(j+1)) +Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = + (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_ya(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = + (center(A, ϕ, I...) + front(A, ϕ, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_xi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = + (front(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 +Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = + (right(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 + +## Because mysum(::generator) does not work inside CUDA kernels... +@inline mysum(A, ϕ, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ϕ, ranges...) + +@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i) where {F<:Function} + s = 0.0 + for i in ranges_i + s += f(A[i]) * ϕ[i] + end + return s +end + +@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j) where {F<:Function} + s = 0.0 + for i in ranges_i, j in ranges_j + s += f(A[i, j]) * ϕ[i, j] + end + return s +end + +@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function} + s = 0.0 + for i in ranges_i, j in ranges_j, k in ranges_k + s += f(A[i, j, k]) * ϕ[i, j, k] + end + return s +end \ No newline at end of file diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index eb1ae8e1c..9f7746210 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -70,6 +70,8 @@ end d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A, ϕ) = _av_xa(A, ϕ, i, j) + av_ya(A, ϕ) = _av_ya(A, ϕ, i, j) av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -80,7 +82,7 @@ end Rx[i, j] = R_Vx = ( -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - - av_xa(ρgx) + av_xa(ρgx, ϕ.center) ) Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) else @@ -94,7 +96,7 @@ end Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - - av_ya(ρgy) + av_ya(ρgy, ϕ.center) Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) else Ry[i, j] = zero(T) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 21402d712..cf59e1a5d 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -45,12 +45,19 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. - `phase_ratios`: The ratios of different phases present. - `air_phase`: The phase representing air. """ -function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase) +function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase) where N nvi = size_v(ϕ) @parallel (@idx nvi) update_rock_ratio_cv!( ϕ, phase_ratios.center, phase_ratios.vertex, air_phase ) - @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) + # @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) + + @parallel (@idx size(ϕ.Vx)) update_rock_ratio_vel!(ϕ.Vx, ratio_vel[1], air_phase) + @parallel (@idx size(ϕ.Vy)) update_rock_ratio_vel!(ϕ.Vy, ratio_vel[2], air_phase) + if N === 3 + @parallel (@idx size(ϕ.Vz)) update_rock_ratio_vel!(ϕ.Vz, ratio_vel[3], air_phase) + end + return nothing end @@ -72,98 +79,11 @@ end return nothing end -@parallel_indices (I...) function update_rock_ratio_vel!( - ϕ::JustRelax.RockRatio{T,N} -) where {T,N} - # 2D - @inline av_x(A::AbstractArray{T,2}) where {T} = _av_xa(A, I...) - @inline av_y(A::AbstractArray{T,2}) where {T} = _av_ya(A, I...) - # 3D - @inline av_x(A::AbstractArray{T,3}) where {T} = _av_yz(A, I...) - @inline av_y(A::AbstractArray{T,3}) where {T} = _av_xz(A, I...) - @inline av_z(A::AbstractArray{T,3}) where {T} = _av_xy(A, I...) - - all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex)) - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case - all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex)) - end +@parallel_indices (I...) function update_rock_ratio_vel!(ϕ, ratio, air_phase) + ϕ[I...] = Float64(Float16(compute_rock_ratio(ratio, air_phase, I...))) return nothing end -""" - isvalid_c(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.center[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_c(ϕ::JustRelax.RockRatio, i, j) - vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) - vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) - v = vx * vy - return v * (ϕ.center[i, j] > 0) -end - -""" - isvalid_v(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.vertex[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_v(ϕ::JustRelax.RockRatio, i, j) - nx, ny = size(ϕ.Vx) - j_bot = max(j - 1, 1) - j0 = min(j, ny) - vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - - nx, ny = size(ϕ.Vy) - i_left = max(i - 1, 1) - i0 = min(i, nx) - vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) - v = vx * vy - return v * (ϕ.vertex[i, j] > 0) -end - -""" - isvalid_vx(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vx[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) - cv = c * v - # return cv * (ϕ.Vx[i, j] > 0) - return ϕ.Vx[i, j] > 0 -end - -""" - isvalid_vy(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vy[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) - v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) - cv = c * v - # return cv * (ϕ.Vy[i, j] > 0) - return ϕ.Vy[i, j] > 0 -end - # """ @@ -175,8 +95,11 @@ end # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. # - `inds`: Cartesian indices to check. # """ -# function isvalid_c(ϕ::JustRelax.RockRatio, i, j) -# return (ϕ.center[i, j] > 0) +# Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) +# vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) +# vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) +# v = vx * vy +# return v * (ϕ.center[i, j] > 0) # end # """ @@ -188,8 +111,18 @@ end # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. # - `inds`: Cartesian indices to check. # """ -# function isvalid_v(ϕ::JustRelax.RockRatio, i, j) -# return (ϕ.vertex[i, j] > 0) +# Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) +# nx, ny = size(ϕ.Vx) +# j_bot = max(j - 1, 1) +# j0 = min(j, ny) +# vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + +# nx, ny = size(ϕ.Vy) +# i_left = max(i - 1, 1) +# i0 = min(i, nx) +# vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) +# v = vx * vy +# return v * (ϕ.vertex[i, j] > 0) # end # """ @@ -201,11 +134,12 @@ end # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. # - `inds`: Cartesian indices to check. # """ -# function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) -# c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0) -# v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0) -# cv = c || v -# return cv || (ϕ.Vx[i, j] > 0) +# Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) +# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) +# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) +# # cv = c * v +# # return cv * (ϕ.Vx[i, j] > 0) +# return (ϕ.Vx[i, j] > 0) # end # """ @@ -217,9 +151,72 @@ end # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. # - `inds`: Cartesian indices to check. # """ -# function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) -# c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) -# v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) -# cv = c || v -# return cv || (ϕ.Vy[i, j] > 0) +# Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) +# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) +# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) +# # cv = c * v +# # return cv * (ϕ.Vy[i, j] > 0) +# return (ϕ.Vy[i, j] > 0) # end + +####### + +""" + isvalid_c(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.center[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) + (ϕ.center[i, j] > 0) +end + +""" + isvalid_v(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.vertex[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) + (ϕ.vertex[i, j] > 0) +end + +""" + isvalid_vx(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vx[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0) + v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0) + cv = c || v + return cv || (ϕ.Vx[i, j] > 0) +end + +""" + isvalid_vy(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vy[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) + c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) + v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) + cv = c || v + return cv || (ϕ.Vy[i, j] > 0) +end + +Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer, N}) where N = ϕ[I...] > 0 From f4a78f167f1c37d113f7f5bd58ad41c48a6e1fb0 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Fri, 25 Oct 2024 15:35:20 +0200 Subject: [PATCH 15/77] save state --- src/rheology/Viscosity.jl | 46 +++++++++++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 66b520db9..1e513bfc1 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -35,7 +35,7 @@ end # with phase ratios @inline function update_viscosity!( - stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; relaxation=1e0 + stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; air_phase = 0, relaxation=1e0 ) update_viscosity!( islinear(rheology), @@ -43,6 +43,7 @@ end phase_ratios, args, rheology, + air_phase, cutoff; relaxation=relaxation, ) @@ -55,6 +56,7 @@ end phase_ratios, args, rheology, + air_phase, cutoff; relaxation=1e0, ) @@ -67,10 +69,12 @@ end phase_ratios, args, rheology, + air_phase, cutoff; relaxation=1e0, ) - compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff; relaxation=relaxation) + air_phase, + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff; relaxation=relaxation) return nothing end @@ -149,10 +153,10 @@ end end function compute_viscosity!( - stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; relaxation=1e0 + stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, air_phase, cutoff; relaxation=1e0 ) return compute_viscosity!( - backend(stokes), stokes, relaxation, phase_ratios, args, rheology, cutoff + backend(stokes), stokes, relaxation, phase_ratios, args, rheology, air_phase, cutoff ) end @@ -163,9 +167,10 @@ function compute_viscosity!( phase_ratios, args, rheology, + air_phase, cutoff, ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, air_phase, cutoff) end function _compute_viscosity!( @@ -174,8 +179,10 @@ function _compute_viscosity!( phase_ratios::JustPIC.PhaseRatios, args, rheology, - cutoff, + air_phase, + cutoff; ) + air_phase ni = size(stokes.viscosity.η) @parallel (@idx ni) compute_viscosity_kernel!( stokes.viscosity.η, @@ -190,8 +197,8 @@ function _compute_viscosity!( end @parallel_indices (I...) function compute_viscosity_kernel!( - η, ν, ratios_center, εxx, εyy, εxyv, args, rheology, cutoff -) + η, ν, ratios_center, εxx, εyy, εxyv, args, rheology, air_phase::Integer, cutoff +) where N # convenience closure @inline gather(A) = _gather(A, I...) @@ -207,7 +214,9 @@ end args_ij = local_viscosity_args(args, I...) # local phase ratio - ratio_ij = @cell(ratios_center[I...]) + ratio_ij = @cell ratios_center[I...] + # remove phase ratio of the air if necessary & normalize ratios + ratio_ij = correct_phase_ratio(air_phase, ratio_ij) # compute second invariant of strain rate tensor εij = εII_0 + ε[1], -εII_0 + ε[1], gather(εxyv) @@ -258,7 +267,7 @@ end end @parallel_indices (I...) function compute_viscosity_kernel!( - η, ν, ratios_center, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, cutoff + η, ν, ratios_center, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, air_phase, cutoff ) # convenience closures @@ -276,7 +285,9 @@ end args_ijk = local_viscosity_args(args, I...) # local phase ratio - ratio_ijk = @cell(ratios_center[I...]) + ratio_ijk = @cell ratios_center[I...] + # remove phase ratio of the air if necessary & normalize ratios + ratio_ijk = correct_phase_ratio(air_phase, ratio_ijk) # compute second invariant of strain rate tensor εij_normal = εij_normal .+ (εII_0, -εII_0 * 0.5, -εII_0 * 0.5) @@ -340,3 +351,16 @@ end # η # end # end + + +function correct_phase_ratio(air_phase, ratio::SVector{N}) where N + corrected_ratio = if iszero(air_phase) + ratio + else + mask = ntuple(i-> (i !== air_phase), Val(N)) + # set air phase ratio to zero + corrected_ratio = ratio .* mask + # normalize phase ratios without air + corrected_ratio ./ sum(corrected_ratio) + end +end \ No newline at end of file From 1626d68fa0ba73e0d40428b421a507746fc65197 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 27 Oct 2024 18:26:20 +0100 Subject: [PATCH 16/77] many fixes & updates --- src/MiniKernels.jl | 105 ++++++------ src/Utils.jl | 58 +++---- src/rheology/BuoyancyForces.jl | 21 +-- src/rheology/Viscosity.jl | 54 ++++-- src/stokes/Stokes2D.jl | 4 +- src/stokes/StressKernels.jl | 12 +- src/thermal_diffusion/DiffusionExplicit.jl | 8 +- src/variational_stokes/MiniKernels.jl | 18 +- src/variational_stokes/Stokes2D.jl | 8 +- src/variational_stokes/StressKernels.jl | 9 +- src/variational_stokes/VelocityKernels.jl | 73 ++++---- src/variational_stokes/mask.jl | 190 ++++++++++----------- 12 files changed, 289 insertions(+), 271 deletions(-) diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl index 98120fcdc..456cf4a30 100644 --- a/src/MiniKernels.jl +++ b/src/MiniKernels.jl @@ -1,36 +1,35 @@ # unmasked versions -Base.@propagate_inbounds @inline center(A::AbstractArray, inds::Vararg{Integer,N}) where {N} = A[inds...] -Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer,N}) where {N} = A[inds .+ 1...] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ - i - 1, j -] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ - i + 1, j -] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ - i, j - 1 -] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where {I<:Integer} = A[ - i, j + 1 -] -Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i - 1, j, k -] -Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i + 1, j, k -] -Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i, j - 1, k -] -Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i, j + 1, k -] -Base.@propagate_inbounds @inline bot(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i, j, k - 1 -] -Base.@propagate_inbounds @inline top(A::AbstractArray, i::I, j::I, k::I) where {I<:Integer} = A[ - i, j, k + 1 -] +Base.@propagate_inbounds @inline center( + A::AbstractArray, inds::Vararg{Integer,N} +) where {N} = A[inds...] +Base.@propagate_inbounds @inline next(A::AbstractArray, inds::Vararg{Integer,N}) where {N} = + A[inds .+ 1...] +Base.@propagate_inbounds @inline left(A::AbstractArray, i::I, j::I) where {I<:Integer} = + A[i - 1, j] +Base.@propagate_inbounds @inline right(A::AbstractArray, i::I, j::I) where {I<:Integer} = + A[i + 1, j] +Base.@propagate_inbounds @inline back(A::AbstractArray, i::I, j::I) where {I<:Integer} = + A[i, j - 1] +Base.@propagate_inbounds @inline front(A::AbstractArray, i::I, j::I) where {I<:Integer} = + A[i, j + 1] +Base.@propagate_inbounds @inline left( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i - 1, j, k] +Base.@propagate_inbounds @inline right( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i + 1, j, k] +Base.@propagate_inbounds @inline back( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i, j - 1, k] +Base.@propagate_inbounds @inline front( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i, j + 1, k] +Base.@propagate_inbounds @inline bot( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i, j, k - 1] +Base.@propagate_inbounds @inline top( + A::AbstractArray, i::I, j::I, k::I +) where {I<:Integer} = A[i, j, k + 1] ## 2D mini kernels const T2 = AbstractArray{T,2} where {T} @@ -52,11 +51,13 @@ Base.@propagate_inbounds @inline _d_yi( A::AbstractArray, _dy, I::Vararg{Integer,N} ) where {N} = (-right(A, I...) + next(A, I...)) * _dy +Base.@propagate_inbounds @inline div(Ax, Ay, _dx, _dy, I::Vararg{Integer, 2}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + # averages Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = - 0.25 * mysum(A, (i+1):(i+2), (j+1):(j+2)) + 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2)) Base.@propagate_inbounds @inline _av_a(A::T, i, j) where {T<:T2} = - 0.25 * mysum(A, (i):(i+1), (j):(j+1)) + 0.25 * mysum(A, (i):(i + 1), (j):(j + 1)) Base.@propagate_inbounds @inline _av_xa(A::T, I::Vararg{Integer,2}) where {T<:T2} = (center(A, I...) + right(A, I...)) * 0.5 Base.@propagate_inbounds @inline _av_ya(A::T, I::Vararg{Integer,2}) where {T<:T2} = @@ -67,10 +68,10 @@ Base.@propagate_inbounds @inline _av_yi(A::T, I::Vararg{Integer,2}) where {T<:T2 (right(A, I...), next(A, I...)) * 0.5 # harmonic averages Base.@propagate_inbounds @inline function _harm(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i+1):(i+2), (j+1):(j+2)) + return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2)) end Base.@propagate_inbounds @inline function _harm_a(A::T, i, j) where {T<:T2} - return eltype(A)(4) * mysum(inv, A, (i):(i+1), (j):(j+1)) + return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1)) end Base.@propagate_inbounds @inline function _harm_xa(A::T, I::Vararg{Integer,2}) where {T<:T2} return eltype(A)(2) * (inv(right(A, I...)) + inv(center(A, I...))) @@ -90,9 +91,11 @@ const T3 = AbstractArray{T,3} where {T} @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3} return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz end +Base.@propagate_inbounds @inline div(Ax, Ay, Az, _dx, _dy, _dz, I::Vararg{Integer, 3}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + _d_zi(Az, _dz, I...) + # averages Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = - 0.125 * mysum(A, i:(i+1), j:(j+1), k:(k+1)) + 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1)) Base.@propagate_inbounds @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + right(A, i, j, k)) Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = @@ -100,17 +103,17 @@ Base.@propagate_inbounds @inline _av_y(A::T, i, j, k) where {T<:T3} = Base.@propagate_inbounds @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (center(A, i, j, k) + top(A, i, j, k)) Base.@propagate_inbounds @inline _av_xy(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, i:(i+1), j:(j+1), k:k) + 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k) Base.@propagate_inbounds @inline _av_xz(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, i:(i+1), j:j, k:(k+1)) + 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1)) Base.@propagate_inbounds @inline _av_yz(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, i:i, j:(j+1), k:(k+1)) + 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1)) Base.@propagate_inbounds @inline _av_xyi(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, (i-1):i, (j-1):j, k:k) + 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k) Base.@propagate_inbounds @inline _av_xzi(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, (i-1):i, j:j, (k-1):k) + 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k) Base.@propagate_inbounds @inline _av_yzi(A::T, i, j, k) where {T<:T3} = - 0.25 * mysum(A, i:i, (j-1):j, (k-1):k) + 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k) # harmonic averages @inline function _harm_x(A::T, i, j, k) where {T<:T3} return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(right(A, i, j, k))) @@ -122,22 +125,22 @@ end return eltype(A)(2) * inv(inv(center(A, i, j, k)) + inv(top(A, i, j, k))) end @inline function _harm_xy(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i+1), j:(j+1), k:k)) + return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k)) end @inline function _harm_xz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:(i+1), j:j, k:(k+1))) + return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:j, k:(k + 1))) end @inline function _harm_yz(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, j:(j+1), k:(k+1))) + return eltype(A)(4) * inv(mysum(A, i:i, j:(j + 1), k:(k + 1))) end @inline function _harm_xyi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i-1):i, (j-1):j, k:k)) + return eltype(A)(4) * inv(mysum(A, (i - 1):i, (j - 1):j, k:k)) end @inline function _harm_xzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, (i-1):i, j:j, (k-1):k)) + return eltype(A)(4) * inv(mysum(A, (i - 1):i, j:j, (k - 1):k)) end @inline function _harm_yzi(A::T, i, j, k) where {T<:T3} - return eltype(A)(4) * inv(mysum(A, i:i, (j-1):j, (k-1):k)) + return eltype(A)(4) * inv(mysum(A, i:i, (j - 1):j, (k - 1):k)) end # others @@ -171,7 +174,9 @@ end return s end -@inline function mysum(f::F, A::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function} +@inline function mysum( + f::F, A::AbstractArray, ranges_i, ranges_j, ranges_k +) where {F<:Function} s = 0.0 for i in ranges_i, j in ranges_j, k in ranges_k s += f(A[i, j, k]) diff --git a/src/Utils.jl b/src/Utils.jl index 3ec3c1380..f9cce928e 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -83,9 +83,8 @@ end @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz -@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = A.xx, -A.yy, -A.xy_c +@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = + A.xx, A.yy, A.xy_c @inline function _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}}) where {T} return A.xx, A.yy, A.zz, A.yz_c, A.xz_c, A.xy_c end @@ -102,9 +101,8 @@ macro velocity(A) end @inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy -@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, -V.Vy, -V.Vz +@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = + V.Vx, V.Vy, V.Vz """ @displacement(U) @@ -117,11 +115,10 @@ macro displacement(A) end end -@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} = U.Ux, -U.Uy -@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,3}}) where {T} = U.Ux, -U.Uy, -U.Uz +@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} = + U.Ux, U.Uy +@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,3}}) where {T} = + U.Ux, U.Uy, U.Uz """ @qT(V) @@ -135,9 +132,8 @@ macro qT(A) end @inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx, A.qTy -@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = A.qTx, -A.qTy, -A.qTz +@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = + A.qTx, A.qTy, A.qTz """ @qT2(V) @@ -150,8 +146,8 @@ macro qT2(A) end end -@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx2, -A.qTy2 +@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = + A.qTx2, A.qTy2 @inline function unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} return A.qTx2, A.qTy2, A.qTz2 end @@ -364,8 +360,8 @@ end @inline function _maxloc_window_clamped(A, I, J, width_x, width_y) nx, ny = size(A) - I_range = (I-width_x):(I+width_x) - J_range = (J-width_y):(J+width_y) + I_range = (I - width_x):(I + width_x) + J_range = (J - width_y):(J + width_y) x = -Inf for j in J_range @@ -383,9 +379,9 @@ end @inline function _maxloc_window_clamped(A, I, J, K, width_x, width_y, width_z) nx, ny, nz = size(A) - I_range = (I-width_x):(I+width_x) - J_range = (J-width_y):(J+width_y) - K_range = (K-width_z):(K+width_z) + I_range = (I - width_x):(I + width_x) + J_range = (J - width_y):(J + width_y) + K_range = (K - width_z):(K + width_z) x = -Inf for k in K_range @@ -432,21 +428,17 @@ function compute_dt(::CPUBackendTrait, S::JustRelax.StokesArrays, args...) return _compute_dt(S, args...) end -@inline _compute_dt(S::JustRelax.StokesArrays, di) = _compute_dt( - @velocity(S), di, Inf, maximum -) +@inline _compute_dt(S::JustRelax.StokesArrays, di) = + _compute_dt(@velocity(S), di, Inf, maximum) -@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = _compute_dt( - @velocity(S), di, dt_diff, maximum -) +@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = + _compute_dt(@velocity(S), di, dt_diff, maximum) -@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) = _compute_dt( - @velocity(S), di, dt_diff, maximum_mpi -) +@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) = + _compute_dt(@velocity(S), di, dt_diff, maximum_mpi) -@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) = _compute_dt( - @velocity(S), di, Inf, maximum_mpi -) +@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) = + _compute_dt(@velocity(S), di, Inf, maximum_mpi) @inline function _compute_dt(V::NTuple, di, dt_diff, max_fun::F) where {F<:Function} n = inv(length(V) + 0.1) diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 8dcb4282e..946cab3b8 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -91,20 +91,17 @@ Compute the buoyancy forces based on the given rheology, arguments, and phase ra end # without phase ratios -@inline update_ρg!(ρg::AbstractArray, rheology, args) = update_ρg!( - isconstant(rheology), ρg, rheology, args -) +@inline update_ρg!(ρg::AbstractArray, rheology, args) = + update_ρg!(isconstant(rheology), ρg, rheology, args) @inline update_ρg!(::ConstantDensityTrait, ρg, rheology, args) = nothing -@inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = compute_ρg!( - ρg, rheology, args -) +@inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = + compute_ρg!(ρg, rheology, args) # with phase ratios -@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = update_ρg!( - isconstant(rheology), ρg, phase_ratios, rheology, args -) +@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = + update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args) @inline update_ρg!( ::ConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args ) = nothing -@inline update_ρg!(::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args) = compute_ρg!( - ρg, phase_ratios, rheology, args -) +@inline update_ρg!( + ::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args +) = compute_ρg!(ρg, phase_ratios, rheology, args) diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 1e513bfc1..d39a34092 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -35,7 +35,13 @@ end # with phase ratios @inline function update_viscosity!( - stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; air_phase = 0, relaxation=1e0 + stokes::JustRelax.StokesArrays, + phase_ratios, + args, + rheology, + cutoff; + air_phase::Integer=0, + relaxation=1e0, ) update_viscosity!( islinear(rheology), @@ -73,8 +79,9 @@ end cutoff; relaxation=1e0, ) - air_phase, - compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff; relaxation=relaxation) + compute_viscosity!( + stokes, phase_ratios, args, rheology, air_phase, cutoff; relaxation=relaxation + ) return nothing end @@ -153,7 +160,13 @@ end end function compute_viscosity!( - stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, air_phase, cutoff; relaxation=1e0 + stokes::JustRelax.StokesArrays, + phase_ratios, + args, + rheology, + air_phase::Integer, + cutoff; + relaxation=1e0, ) return compute_viscosity!( backend(stokes), stokes, relaxation, phase_ratios, args, rheology, air_phase, cutoff @@ -182,7 +195,6 @@ function _compute_viscosity!( air_phase, cutoff; ) - air_phase ni = size(stokes.viscosity.η) @parallel (@idx ni) compute_viscosity_kernel!( stokes.viscosity.η, @@ -191,6 +203,7 @@ function _compute_viscosity!( @strain(stokes)..., args, rheology, + air_phase, cutoff, ) return nothing @@ -198,7 +211,7 @@ end @parallel_indices (I...) function compute_viscosity_kernel!( η, ν, ratios_center, εxx, εyy, εxyv, args, rheology, air_phase::Integer, cutoff -) where N +) # convenience closure @inline gather(A) = _gather(A, I...) @@ -267,7 +280,19 @@ end end @parallel_indices (I...) function compute_viscosity_kernel!( - η, ν, ratios_center, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, air_phase, cutoff + η, + ν, + ratios_center, + εxx, + εyy, + εzz, + εyzv, + εxzv, + εxyv, + args, + rheology, + air_phase::Integer, + cutoff, ) # convenience closures @@ -352,15 +377,16 @@ end # end # end - -function correct_phase_ratio(air_phase, ratio::SVector{N}) where N - corrected_ratio = if iszero(air_phase) - ratio +function correct_phase_ratio(air_phase, ratio::SVector{N, T}) where {N, T} + if iszero(air_phase) + return ratio + elseif ratio[air_phase] ≈ 1 + return SVector{N,T}(zero(T) for _ in 1:N) else - mask = ntuple(i-> (i !== air_phase), Val(N)) + mask = ntuple(i -> (i !== air_phase), Val(N)) # set air phase ratio to zero corrected_ratio = ratio .* mask # normalize phase ratios without air - corrected_ratio ./ sum(corrected_ratio) + return corrected_ratio ./ sum(corrected_ratio) end -end \ No newline at end of file +end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index ce65dd0b7..ad2eb38ae 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -377,7 +377,7 @@ function _solve!( center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) - @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( Vx_on_Vy, stokes.V.Vx ) @@ -601,7 +601,7 @@ function _solve!( ) update_halo!(stokes.τ.xy) - @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!( + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!( Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1] ) diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index c4b53685c..59fb01217 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -653,7 +653,7 @@ end # yield function @ vertex Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) + if is_pl && !iszero(τIIv_ij) && Fv >0 # stress correction @ vertex λv[1][I...] = (1.0 - relλ) * λv[1][I...] + @@ -728,7 +728,7 @@ end # yield function @ vertex Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) + if is_pl && !iszero(τIIv_ij) && Fv >0 # stress correction @ vertex λv[2][I...] = (1.0 - relλ) * λv[2][I...] + @@ -805,7 +805,7 @@ end # yield function @ vertex Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) + if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[3][I...] = (1.0 - relλ) * λv[3][I...] + @@ -842,7 +842,7 @@ end # yield function @ center F = τII_ij - C - Pr[I...] * sinϕ - if is_pl && !iszero(τII_ij) + if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center λ[I...] = (1.0 - relλ) * λ[I...] + @@ -930,7 +930,7 @@ end # yield function @ center Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) + if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[I...] = (1.0 - relλ) * λv[I...] + @@ -965,7 +965,7 @@ end # yield function @ center F = τII_ij - C - Pr[I...] * sinϕ - if is_pl && !iszero(τII_ij) + if is_pl && !iszero(τII_ij) && F >0 # stress correction @ center λ[I...] = (1.0 - relλ) * λ[I...] + diff --git a/src/thermal_diffusion/DiffusionExplicit.jl b/src/thermal_diffusion/DiffusionExplicit.jl index 97a0fee4c..2fa055fb8 100644 --- a/src/thermal_diffusion/DiffusionExplicit.jl +++ b/src/thermal_diffusion/DiffusionExplicit.jl @@ -381,7 +381,7 @@ function JustRelax.solve!( # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( + @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy ) @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _dx, _dy) @@ -413,7 +413,7 @@ function JustRelax.solve!( # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( + @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, phase_ratios.center, args, _di... ) @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _di...) @@ -442,7 +442,7 @@ function JustRelax.solve!( nx, ny = size(thermal.T) # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( + @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy ) @parallel advect_T!( @@ -481,7 +481,7 @@ function JustRelax.solve!( nx, ny = size(thermal.T) # solve heat diffusion @parallel assign!(thermal.Told, thermal.T) - @parallel (1:(nx-1), 1:(ny-1)) compute_flux!( + @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!( thermal.qTx, thermal.qTy, thermal.T, phases, rheology, args, _dx, _dy ) @parallel advect_T!( diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 4fdf7c43b..3e2055b86 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -21,9 +21,9 @@ Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) w # averages Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j) where {T<:T2} = - 0.25 * mysum(A, ϕ, (i+1):(i+2), (j+1):(j+2)) + 0.25 * mysum(A, ϕ, (i + 1):(i + 2), (j + 1):(j + 2)) Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j) where {T<:T2} = - 0.25 * mysum(A, ϕ, (i):(i+1), (j):(j+1)) + 0.25 * mysum(A, ϕ, (i):(i + 1), (j):(j + 1)) Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5 Base.@propagate_inbounds @inline _av_ya(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = @@ -36,7 +36,9 @@ Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where ## Because mysum(::generator) does not work inside CUDA kernels... @inline mysum(A, ϕ, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ϕ, ranges...) -@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i) where {F<:Function} +@inline function mysum( + f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i +) where {F<:Function} s = 0.0 for i in ranges_i s += f(A[i]) * ϕ[i] @@ -44,7 +46,9 @@ Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where return s end -@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j) where {F<:Function} +@inline function mysum( + f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j +) where {F<:Function} s = 0.0 for i in ranges_i, j in ranges_j s += f(A[i, j]) * ϕ[i, j] @@ -52,10 +56,12 @@ end return s end -@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function} +@inline function mysum( + f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j, ranges_k +) where {F<:Function} s = 0.0 for i in ranges_i, j in ranges_j, k in ranges_k s += f(A[i, j, k]) * ϕ[i, j, k] end return s -end \ No newline at end of file +end diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index fcff656bd..c9a189626 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -23,6 +23,7 @@ function _solve_VS!( args, dt, igg::IGG; + air_phase::Integer=0, viscosity_cutoff=(-Inf, Inf), viscosity_relaxation=1e-2, iterMax=50e3, @@ -77,7 +78,7 @@ function _solve_VS!( # compute buoyancy forces and viscosity compute_ρg!(ρg[end], phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, viscosity_cutoff) displacement2velocity!(stokes, dt, flow_bcs) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @@ -90,7 +91,7 @@ function _solve_VS!( compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., ϕ, _di...) + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di) compute_P!( θ, @@ -118,6 +119,7 @@ function _solve_VS!( args, rheology, viscosity_cutoff; + air_phase=air_phase, relaxation=viscosity_relaxation, ) @@ -175,7 +177,7 @@ function _solve_VS!( length(stokes.R.Rx), norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm_mpi(stokes.R.RP) / length(stokes.R.RP), + norm_mpi(@views stokes.R.RP[ϕ.center .> 0]) / length(@views stokes.R.RP[ϕ.center .> 0]), ) push!(norm_Rx, errs[1]) push!(norm_Ry, errs[2]) diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl index 6b1d848a3..9aa70acdf 100644 --- a/src/variational_stokes/StressKernels.jl +++ b/src/variational_stokes/StressKernels.jl @@ -61,12 +61,13 @@ -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] + 2.0 * ηv_ij * ε[3][I...] ) * dτ_rv - τIIv_ij = - √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) + τIIv_ij = √( + 0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2 + ) # yield function @ center Fv = τIIv_ij - Cv - Pv_ij * sinϕv - if is_pl && !iszero(τIIv_ij) + if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[I...] = (1.0 - relλ) * λv[I...] + @@ -107,7 +108,7 @@ # yield function @ center F = τII_ij - C - Pr[I...] * sinϕ - if is_pl && !iszero(τII_ij) + if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center λ[I...] = (1.0 - relλ) * λ[I...] + diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 9f7746210..7624feaaf 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,54 +1,43 @@ -@parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T,2}, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy -) where {T} - @inline d_xi(A) = _d_xi(A, _dx, I...) - @inline d_yi(A) = _d_yi(A, _dy, I...) - - if isvalid_c(ϕ, I...) - @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) - else - @inbounds ∇V[I...] = zero(T) - end - return nothing -end - -@parallel_indices (I...) function compute_∇V!( - ∇V::AbstractArray{T,2}, Vx, Vy, Vz, ϕ::JustRelax.RockRatio, _dx, _dy, _dz -) where {T} - @inline d_xi(A) = _d_xi(A, _dx, I...) - @inline d_yi(A) = _d_yi(A, _dy, I...) - @inline d_zi(A) = _d_zi(A, _dz, I...) +# @parallel_indices (I...) function compute_∇V!( +# ∇V::AbstractArray{T,2}, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy +# ) where {T} +# @inline d_xi(A) = _d_xi(A, _dx, I...) +# @inline d_yi(A) = _d_yi(A, _dy, I...) - if isvalid_c(ϕ, I...) - @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) - else - @inbounds ∇V[I...] = zero(T) - end - return nothing -end +# if isvalid_c(ϕ, I...) +# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) +# else +# @inbounds ∇V[I...] = zero(T) +# end +# return nothing +# end # @parallel_indices (I...) function compute_∇V!( -# ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} -# ) where {T,N} -# # @inline d_xi(A) = _d_xi(A, _di, I...) -# # @inline d_yi(A) = _d_yi(A, _di, I...) -# # @inline d_zi(A) = _d_zi(A, _di, I...) - -# # f = d_xi, d_yi, d_zi -# f = _d_xi, _d_yi, _d_zi +# ∇V::AbstractArray{T,2}, Vx, Vy, Vz, ϕ::JustRelax.RockRatio, _dx, _dy, _dz +# ) where {T} +# @inline d_xi(A) = _d_xi(A, _dx, I...) +# @inline d_yi(A) = _d_yi(A, _dy, I...) +# @inline d_zi(A) = _d_zi(A, _dz, I...) # if isvalid_c(ϕ, I...) -# v = zero(T) -# for i in 1:N -# v += f[i](V[i], _di[I], I...) -# end -# @inbounds ∇V[I...] = v +# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) # else # @inbounds ∇V[I...] = zero(T) # end # return nothing # end +@parallel_indices (I...) function compute_∇V!( + ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} +) where {T,N} + if isvalid_c(ϕ, I...) + @inbounds ∇V[I...] = div(V..., _di..., I...) + else + @inbounds ∇V[I...] = zero(T) + end + return nothing +end + @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, @@ -84,7 +73,7 @@ end -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx, ϕ.center) ) - Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ, ϕ.center) else Rx[i, j] = zero(T) Vx[i + 1, j + 1] = zero(T) @@ -97,7 +86,7 @@ end R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy, ϕ.center) - Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ, ϕ.center) else Ry[i, j] = zero(T) Vy[i + 1, j + 1] = zero(T) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index cf59e1a5d..3ad186c50 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -6,8 +6,8 @@ function RockRatio(nx, ny) ni = nx, ny center = @zeros(ni...) vertex = @zeros(ni .+ 1...) - Vx = @zeros(nx+1, ny) # no ghost nodes! - Vy = @zeros(nx, ny+1) # no ghost nodes! + Vx = @zeros(nx + 1, ny) # no ghost nodes! + Vy = @zeros(nx, ny + 1) # no ghost nodes! dummy = @zeros(1, 1) # because it cant be a Union{T, Nothing} type on the GPU.... return JustRelax.RockRatio(center, vertex, Vx, Vy, dummy, dummy, dummy, dummy) end @@ -16,9 +16,9 @@ function RockRatio(nx, ny, nz) ni = nx, ny, nz center = @zeros(ni...) vertex = @zeros(ni .+ 1...) - Vx = @zeros(nx+1, ny, nz) # no ghost nodes! - Vy = @zeros(nx, ny+1, nz) # no ghost nodes! - Vz = @zeros(nx, ny, nz+1) # no ghost nodes! + Vx = @zeros(nx + 1, ny, nz) # no ghost nodes! + Vy = @zeros(nx, ny + 1, nz) # no ghost nodes! + Vz = @zeros(nx, ny, nz + 1) # no ghost nodes! yz = @zeros(nx, ny + 1, nz + 1) xz = @zeros(nx + 1, ny, nz + 1) xy = @zeros(nx + 1, ny + 1, nz) @@ -45,7 +45,9 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. - `phase_ratios`: The ratios of different phases present. - `air_phase`: The phase representing air. """ -function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase) where N +function update_rock_ratio!( + ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase +) where {N} nvi = size_v(ϕ) @parallel (@idx nvi) update_rock_ratio_cv!( ϕ, phase_ratios.center, phase_ratios.vertex, air_phase @@ -65,9 +67,9 @@ end phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} ) where {N} = 1 - @index phase_ratio[air_phase, I...] -@inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[ - air_phase, I... -] +@inline compute_air_ratio( + phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} +) where {N} = @index phase_ratio[air_phase, I...] @parallel_indices (I...) function update_rock_ratio_cv!( ϕ, ratio_center, ratio_vertex, air_phase @@ -84,7 +86,80 @@ end return nothing end +""" + isvalid_c(ϕ::JustRelax.RockRatio, inds...) +Check if `ϕ.center[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) + vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) + vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + v = vx * vy + return v * (ϕ.center[i, j] > 0) +end + +""" + isvalid_v(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.vertex[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) + nx, ny = size(ϕ.Vx) + j_bot = max(j - 1, 1) + j0 = min(j, ny) + vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + + nx, ny = size(ϕ.Vy) + i_left = max(i - 1, 1) + i0 = min(i, nx) + vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + v = vx * vy + return v * (ϕ.vertex[i, j] > 0) +end + +""" + isvalid_vx(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vx[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) + # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) + # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) + # cv = c * v + # return cv * (ϕ.Vx[i, j] > 0) + return (ϕ.Vx[i, j] > 0) +end + +""" + isvalid_vy(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vy[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) + # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) + # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) + # cv = c * v + # return cv * (ϕ.Vy[i, j] > 0) + return (ϕ.Vy[i, j] > 0) +end + +###### # """ # isvalid_c(ϕ::JustRelax.RockRatio, inds...) @@ -96,10 +171,7 @@ end # - `inds`: Cartesian indices to check. # """ # Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) -# vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) -# vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) -# v = vx * vy -# return v * (ϕ.center[i, j] > 0) +# return isvalid(ϕ.center, i, j) # end # """ @@ -112,17 +184,7 @@ end # - `inds`: Cartesian indices to check. # """ # Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) -# nx, ny = size(ϕ.Vx) -# j_bot = max(j - 1, 1) -# j0 = min(j, ny) -# vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) - -# nx, ny = size(ϕ.Vy) -# i_left = max(i - 1, 1) -# i0 = min(i, nx) -# vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) -# v = vx * vy -# return v * (ϕ.vertex[i, j] > 0) +# return isvalid(ϕ.vertex, i, j) # end # """ @@ -135,11 +197,10 @@ end # - `inds`: Cartesian indices to check. # """ # Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) -# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) -# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) -# # cv = c * v -# # return cv * (ϕ.Vx[i, j] > 0) -# return (ϕ.Vx[i, j] > 0) +# c = isvalid(ϕ.center, i, j) || isvalid(ϕ.center, i - 1, j) +# v = isvalid(ϕ.vertex, i, j) || isvalid(ϕ.vertex, i, j + 1) +# cv = c || v +# return cv || isvalid(ϕ.Vx, i, j) # end # """ @@ -152,71 +213,10 @@ end # - `inds`: Cartesian indices to check. # """ # Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) -# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) -# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) -# # cv = c * v -# # return cv * (ϕ.Vy[i, j] > 0) -# return (ϕ.Vy[i, j] > 0) +# c = isvalid(ϕ.center, i, j) || isvalid(ϕ.center, i, j - 1) +# v = isvalid(ϕ.vertex, i, j) || isvalid(ϕ.vertex, i + 1, j) +# cv = c || v +# return cv || isvalid(ϕ.Vy, i, j) # end -####### - -""" - isvalid_c(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.center[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) - (ϕ.center[i, j] > 0) -end - -""" - isvalid_v(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.vertex[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) - (ϕ.vertex[i, j] > 0) -end - -""" - isvalid_vx(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vx[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0) - v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0) - cv = c || v - return cv || (ϕ.Vx[i, j] > 0) -end - -""" - isvalid_vy(ϕ::JustRelax.RockRatio, inds...) - -Check if `ϕ.Vy[inds...]` is a not a nullspace. - -# Arguments -- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. -- `inds`: Cartesian indices to check. -""" -Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) - c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0) - v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0) - cv = c || v - return cv || (ϕ.Vy[i, j] > 0) -end - -Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer, N}) where N = ϕ[I...] > 0 +Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer,N}) where {N} = ϕ[I...] > 0 From 67440305a10d502e73697f3a85b53dd49d2f70db Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 1 Nov 2024 16:36:33 +0100 Subject: [PATCH 17/77] updates and fixes --- src/MiniKernels.jl | 6 ++++-- src/ext/AMDGPU/2D.jl | 10 ++++++++++ src/ext/AMDGPU/3D.jl | 10 ++++++++++ src/ext/CUDA/2D.jl | 17 +++++++++++++++-- src/ext/CUDA/3D.jl | 7 +++++++ src/rheology/Viscosity.jl | 2 +- src/stokes/StressKernels.jl | 11 +++++++---- src/variational_stokes/Stokes2D.jl | 3 ++- 8 files changed, 56 insertions(+), 10 deletions(-) diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl index 456cf4a30..e8a4c2fb1 100644 --- a/src/MiniKernels.jl +++ b/src/MiniKernels.jl @@ -51,7 +51,8 @@ Base.@propagate_inbounds @inline _d_yi( A::AbstractArray, _dy, I::Vararg{Integer,N} ) where {N} = (-right(A, I...) + next(A, I...)) * _dy -Base.@propagate_inbounds @inline div(Ax, Ay, _dx, _dy, I::Vararg{Integer, 2}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) +Base.@propagate_inbounds @inline div(Ax, Ay, _dx, _dy, I::Vararg{Integer,2}) = + _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) # averages Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} = @@ -91,7 +92,8 @@ const T3 = AbstractArray{T,3} where {T} @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3} return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz end -Base.@propagate_inbounds @inline div(Ax, Ay, Az, _dx, _dy, _dz, I::Vararg{Integer, 3}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + _d_zi(Az, _dz, I...) +Base.@propagate_inbounds @inline div(Ax, Ay, Az, _dx, _dy, _dz, I::Vararg{Integer,3}) = + _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + _d_zi(Az, _dz, I...) # averages Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} = diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index d93070f49..b78c4143f 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -408,4 +408,14 @@ function JR2D.update_rock_ratio!( return nothing end +function JR2D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N}, + phase_ratios, + ratio_vel::NTuple{N}, + air_phase, +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) + return nothing +end + end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 7b00514e3..a55e0cde2 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -425,4 +425,14 @@ function JR3D.update_rock_ratio!( return nothing end +function JR3D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N}, + phase_ratios, + ratio_vel::NTuple{N}, + air_phase, +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) + return nothing +end + end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 6e8711695..888a2f327 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -185,6 +185,12 @@ function JR2D.compute_viscosity!( return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end +function JR2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, air_phase, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, air_phase, cutoff) +end + function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end @@ -194,9 +200,9 @@ function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cuto end function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, air_phase, cutoff ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, air_phase, cutoff) end function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) @@ -394,4 +400,11 @@ function JR2D.update_rock_ratio!( return nothing end +function JR2D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) + return nothing +end + end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index ee4a0007e..b44e3527e 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -421,4 +421,11 @@ function JR3D.update_rock_ratio!( return nothing end +function JR3D.update_rock_ratio!( + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase +) where {T,nD,N,D} + update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) + return nothing +end + end diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index d39a34092..57116c626 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -377,7 +377,7 @@ end # end # end -function correct_phase_ratio(air_phase, ratio::SVector{N, T}) where {N, T} +function correct_phase_ratio(air_phase, ratio::SVector{N,T}) where {N,T} if iszero(air_phase) return ratio elseif ratio[air_phase] ≈ 1 diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index ca51605a5..dd5149d51 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -622,7 +622,8 @@ end τxyv_old_ij = av_clamped_yz_z(τxyv_old, Ic...) # vertex parameters - phase = @inbounds phase_vertex[I...] + phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1] + 1, I[2], I[3]]) * 0.5 + # phase = @inbounds phase_vertex[I...] is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( rheology, EIIv_ij, phase ) @@ -697,7 +698,8 @@ end τxyv_old_ij = av_clamped_xz_z(τxyv_old, Ic...) # vertex parameters - phase = @inbounds phase_vertex[I...] + phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1], I[2] + 1, I[3]]) * 0.5 + # phase = @inbounds phase_vertex[I...] is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( rheology, EIIv_ij, phase ) @@ -774,7 +776,8 @@ end τxyv_old_ij = τxyv_old[I...] # vertex parameters - phase = @inbounds phase_vertex[I...] + phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1], I[2], I[3] + 1]) * 0.5 + # phase = @inbounds phase_vertex[I...] is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( rheology, EIIv_ij, phase ) @@ -969,7 +972,7 @@ end # yield function @ center F = τII_ij - C - Pr[I...] * sinϕ - if is_pl && !iszero(τII_ij) && F >0 + if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center λ[I...] = (1.0 - relλ) * λ[I...] + diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index c9a189626..bf94910c0 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -177,7 +177,8 @@ function _solve_VS!( length(stokes.R.Rx), norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) / length(stokes.R.Ry), - norm_mpi(@views stokes.R.RP[ϕ.center .> 0]) / length(@views stokes.R.RP[ϕ.center .> 0]), + norm_mpi(@views stokes.R.RP[ϕ.center .> 0]) / + length(@views stokes.R.RP[ϕ.center .> 0]), ) push!(norm_Rx, errs[1]) push!(norm_Ry, errs[2]) From 70de2f147a70fc22ae37bb249967dfc3c1bfdbf3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 1 Nov 2024 16:36:49 +0100 Subject: [PATCH 18/77] miniapp --- .../VariationalSubduction2D.jl | 355 ++++++++++++++++++ .../VariationalSubduction2D_rheology.jl | 69 ++++ 2 files changed, 424 insertions(+) create mode 100644 miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D.jl create mode 100644 miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D_rheology.jl diff --git a/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D.jl b/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D.jl new file mode 100644 index 000000000..83a48b4e1 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D.jl @@ -0,0 +1,355 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +# Load script dependencies +using GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Subduction2D_setup.jl") +include("VariationalSubduction2D_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +using StaticArrays +function correct_phase_ratio(air_phase, ratio::SVector{N, T}) where {N, T} + if iszero(air_phase) + return ratio + elseif ratio[air_phase] ≈ 1 + return SVector{N,T}(zero(T) for _ in 1:N) + else + mask = ntuple(i -> (i !== air_phase), Val(N)) + # set air phase ratio to zero + corrected_ratio = ratio .* mask + # normalize phase ratios without air + # return corrected_ratio ./ sum(corrected_ratio) + end +end + +@parallel_indices (I...) function renormalize_phase_ratios(ratios_center, ratios_vertex, air_phase) + # renormalize centers + if all(I .≤ size(ratios_center)) + # local phase ratio + ratio_ij = @cell ratios_center[I...] + # remove phase ratio of the air if necessary & normalize ratios + @cell ratios_center[I...] = correct_phase_ratio(air_phase, ratio_ij) + end + + # renormalize centers + # local phase ratio + ratio_ij = @cell ratios_vertex[I...] + # remove phase ratio of the air if necessary & normalize ratios + @cell ratios_vertex[I...] = correct_phase_ratio(air_phase, ratio_ij) + return nothing +end + +import ParallelStencil.INDICES +const idx_k = INDICES[2] +macro all_k(A) + esc(:($A[$idx_k])) +end + +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + + @parallel f_x(A, B) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) + + # Physical domain ------------------------------------ + ni = nx, ny # number of cells + di = @. li / ni # grid steps + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies() + dt = 10e3 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell = 50 + max_xcell = 75 + min_xcell = 30 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # material phase & temperature + pPhases, = init_cell_arrays(particles, Val(1)) + particle_args = (pPhases,) + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); + init_phases!(pPhases, phases_device, particles, xvi) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + phase_ratio_vx = @fill(0.0, nx+1, ny, celldims = length(rheology)) + phase_ratio_vy = @fill(0.0, nx, ny+1, celldims = length(rheology)) + + phase_ratios_midpoint!(phase_ratio_vx, particles, xci, pPhases, :x) + phase_ratios_midpoint!(phase_ratio_vy, particles, xci, pPhases, :y) + # ---------------------------------------------------- + + # RockRatios + air_phase = 3 + ϕ_R = RockRatio(backend, ni) + update_rock_ratio!(ϕ_R, phase_ratios, (phase_ratio_vx, phase_ratio_vy), air_phase) + # @parallel (@idx ni.+1) renormalize_phase_ratios(phase_ratios.center, phase_ratios.vertex, air_phase) + + # marker chain + nxcell, min_xcell, max_xcell = 12, 6, 24 + initial_elevation = 0e0 + chain = init_markerchain(JustPIC.CPUBackend, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend, ni) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + + # Rheology + args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) + viscosity_cutoff = (1e18, 1e23) + compute_viscosity!(stokes, phase_ratios, args0, rheology, air_phase, viscosity_cutoff) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true), + free_surface = false, + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + end + + τxx_v = @zeros(ni.+1...) + τyy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + + while it < 100 # run only for 5 Myrs + + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + # out = solve!( + # stokes, + # pt_stokes, + # di, + # flow_bcs, + # ρg, + # phase_ratios, + # rheology, + # args, + # dt, + # igg; + # kwargs = ( + # iterMax = 100e3, + # nout = 2e3, + # viscosity_cutoff = viscosity_cutoff, + # free_surface = false, + # viscosity_relaxation = 1e-2 + # ) + # ); + solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ_R, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 50e3,#250e3, + # free_surface = false, + nout = 2e3,#5e3, + viscosity_cutoff = viscosity_cutoff, + ) + ) + end + + println("Stokes solver time ") + println(" Total time: $t_stokes s") + # println(" Time/iteration: $(t_stokes / out.iter) s") + tensor_invariant!(stokes.ε) + dt = compute_dt(stokes, di) * 0.8 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_MQS!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + phase_ratios_midpoint!(phase_ratio_vx, particles, xci, pPhases, :x) + phase_ratios_midpoint!(phase_ratio_vy, particles, xci, pPhases, :y) + update_rock_ratio!(ϕ_R, phase_ratios, (phase_ratio_vx, phase_ratio_vy), air_phase) + @parallel (@idx ni.+1) renormalize_phase_ratios(phase_ratios.center, phase_ratios.vertex, air_phase) + + advect_markerchain!(chain, method, V, grid_vxi, dt) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 1) == 0 + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + ) + data_c = (; + P = Array(stokes.P), + η = Array(η_vep), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + # clr = pT.data[:] + idxv = particles.index.data[:]; + + chain_x = chain.coords[1].data[:]./1e3 + chain_y = chain.coords[2].data[:]./1e3 + + # Make Makie figure + ar = 3 + fig = Figure(size = (1200, 900), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "log10(εII) (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") + ax3 = Axis(fig[1,3], aspect = ar, title = "τII") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + scatter!(ax1, chain_x, chain_y, markersize = 3) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array((stokes.τ.II)) , colormap=:batlow) + scatter!(ax3, chain_x, chain_y, markersize = 3) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + display(fig) + save(joinpath(figdir, "$(it).png"), fig) + end + # ------------------------------ + + end + + return stokes +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction2D_MQS_variational" +nx, ny = 250, 100 +li, origin, phases_GMG, T_GMG = GMG_subduction_2D(nx+1, ny+1) +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +stokes = main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); \ No newline at end of file diff --git a/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D_rheology.jl b/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D_rheology.jl new file mode 100644 index 000000000..4804f8ac6 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/StickyAirSubduction/VariationalSubduction2D_rheology.jl @@ -0,0 +1,69 @@ +function init_rheologies() + + # Define rheolgy struct + rheology = ( + # Name = "Asthenoshpere", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Oceanic lithosphere", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23),)), + Gravity = ConstantGravity(; g=9.81), + ), + ) +end + +function init_phases!(phases, phase_grid, particles, xvi) + ni = size(phases) + @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) +end + +@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} + + ni = size(phases) + + for ip in cellaxes(phases) + # quick escape + @index(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1 + ii = I[1] + offi + jj = I[2] + offj + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + ) + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj] + end + end + @index phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end From 33068502c65ca7054268b23539b59687fe2ab252 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 13 Nov 2024 09:56:38 +0100 Subject: [PATCH 19/77] up --- src/variational_stokes/Stokes2D.jl | 2 +- src/variational_stokes/VelocityKernels.jl | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index bf94910c0..b4f7bac1c 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -163,7 +163,7 @@ function _solve_VS!( ) # apply boundary conditions velocity2displacement!(stokes, dt) - free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + # free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) flow_bcs!(stokes, flow_bcs) update_halo!(@velocity(stokes)...) # end diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 7624feaaf..788820d0d 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -73,7 +73,7 @@ end -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - av_xa(ρgx, ϕ.center) ) - Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ, ϕ.center) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) else Rx[i, j] = zero(T) Vx[i + 1, j + 1] = zero(T) @@ -86,7 +86,7 @@ end R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy, ϕ.center) - Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ, ϕ.center) + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) else Ry[i, j] = zero(T) Vy[i + 1, j + 1] = zero(T) From 11849d93acba49c2bc2d99877cdc245d974ad2e3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 14 Nov 2024 10:45:55 +0100 Subject: [PATCH 20/77] test scripts --- volcano2D/Caldera2D.jl | 460 ++++++++++++++++++++++++++++++++ volcano2D/Caldera_rheology.jl | 149 +++++++++++ volcano2D/Caldera_setup.jl | 82 ++++++ volcano2D/DK2D.jl | 489 ++++++++++++++++++++++++++++++++++ volcano2D/DK_rheology.jl | 152 +++++++++++ volcano2D/DK_setup.jl | 82 ++++++ 6 files changed, 1414 insertions(+) create mode 100644 volcano2D/Caldera2D.jl create mode 100644 volcano2D/Caldera_rheology.jl create mode 100644 volcano2D/Caldera_setup.jl create mode 100644 volcano2D/DK2D.jl create mode 100644 volcano2D/DK_rheology.jl create mode 100644 volcano2D/DK_setup.jl diff --git a/volcano2D/Caldera2D.jl b/volcano2D/Caldera2D.jl new file mode 100644 index 000000000..4edcd8df2 --- /dev/null +++ b/volcano2D/Caldera2D.jl @@ -0,0 +1,460 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +# Load script dependencies +using GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Caldera_setup.jl") +include("Caldera_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[2] +macro all_k(A) + esc(:($A[$idx_k])) +end + +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + + @parallel f_x(A, B) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end + +function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) + xv, yv = xvi + + @parallel_indices (i, j) function pure_shear_x!(Vx, εbg, lx) + xi = xv[i] + Vx[i, j + 1] = εbg * (xi - lx * 0.5) + return nothing + end + + @parallel_indices (i, j) function pure_shear_y!(Vy, εbg, ly) + yi = yv[j] + Vy[i + 1, j] = abs(yi) * εbg + return nothing + end + + nx, ny = size(Vx) + @parallel (1:nx, 1:(ny - 2)) pure_shear_x!(Vx, εbg, lx) + nx, ny = size(Vy) + @parallel (1:(nx - 2), 1:ny) pure_shear_y!(Vy, εbg, ly) + + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) + + # Physical domain ------------------------------------ + ni = nx, ny # number of cells + di = @. li / ni # grid steps + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies() + dt = 1e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # dt = Inf # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell = 40 + max_xcell = 60 + min_xcell = 20 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # material phase & temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + + # particle fields for the stress rotation + pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress + # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress + pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity + particle_args = (pT, pPhases, pτ..., pω...) + particle_args_reduced = (pT, pτ..., pω...) + + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); + init_phases!(pPhases, phases_device, particles, xvi) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + # rock ratios for variational stokes + # RockRatios + air_phase = 5 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.98 / √2.1) # Re=3π, r=0.7 + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + Ttop = T_GMG[1, end] + Tbot = T_GMG[1, 1] + thermal = ThermalArrays(backend, ni) + @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + ) + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + + # Rheology + args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) + viscosity_cutoff = (1e18, 1e23) + compute_viscosity!(stokes, phase_ratios, args0, rheology, air_phase, viscosity_cutoff) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + backend, rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-8, CFL=0.95 / √2 + ) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true), + free_surface = false, + ) + εbg = 1e-15 + apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) + # ---------------------------------------------------- + + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + end + + T_buffer = @zeros(ni.+1) + Told_buffer = similar(T_buffer) + dt₀ = similar(stokes.P) + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle!(pT, xvi, T_buffer, particles) + + ## Plot initial T and P profile + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:]), + Yv, + ) + lines!( + ax2, + Array(stokes.P[:]), + Y, + ) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + τxx_v = @zeros(ni.+1...) + τyy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + thermal.Told .= thermal.T + + while it < 100 #000 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + + particle2centroid!(stokes.τ.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ.yy, pτyy, xci, particles) + particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + + t_stokes = @elapsed solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = viscosity_cutoff, + ) + ) + + center2vertex!(τxx_v, stokes.τ.xx) + center2vertex!(τyy_v, stokes.τ.yy) + centroid2particle!(pτxx , xci, stokes.τ.xx, particles) + centroid2particle!(pτyy , xci, stokes.τ.yy, particles) + grid2particle!(pτxy, xvi, stokes.τ.xy, particles) + rotate_stress_particles!(pτ, pω, particles, dt) + + println("Stokes solver time ") + println(" Total time: $t_stokes s") + tensor_invariant!(stokes.ε) + tensor_invariant!(stokes.ε_pl) + dtmax = 1e3 * 3600 * 24 * 365.25 + dt = compute_dt(stokes, di) * 0.8 + + println("dt = $(dt/(3600 * 24 *365.25)) years") + # ------------------------------ + + # # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 50e3, + nout = 1e2, + verbose = true, + ) + ) + thermal.ΔT .= thermal.T .- thermal.Told + vertex2center!(thermal.ΔTc, thermal.ΔT[2:end-1, :]) + + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + inject_particles_phase!( + particles, + pPhases, + particle_args_reduced, + (T_buffer, τxx_v, τyy_v, stokes.τ.xy, stokes.ω.xy), + xvi + ) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 2) == 0 + # checkpointing(figdir, stokes, thermal.T, η, t) + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(T_buffer), + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + ) + data_c = (; + P = Array(stokes.P), + η = Array(η_vep), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + # clr = pT.data[:] + idxv = particles.index.data[:]; + # Make Makie figure + ar = 2 + fig = Figure(size = (1200, 900), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") + # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") + ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII_plastic)") + # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) + # Plot particles phase + # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) + # Plot 2nd invariant of strain rate + # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:lipari) + # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + fig + save(joinpath(figdir, "$(it).png"), fig) + + ## Plot initial T and P profile + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:]), + Yv, + ) + lines!( + ax2, + Array(stokes.P[:]), + Y, + ) + hideydecorations!(ax2) + save(joinpath(figdir, "thermal_profile_$it.png"), fig) + fig + end + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Caldera2D" +n = 256 +nx, ny = n, n >>> 1 +li, origin, phases_GMG, T_GMG = setup2D( + nx+1, ny+1; + sticky_air = 2.5, + flat = true, + chimney = false, + chamber_T = 1e3, + chamber_depth = 7e0, + chamber_radius = 1e0/2, + aspect_x = 10, +) + +heatmap(phases_GMG) +# heatmap(T_GMG) + +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); \ No newline at end of file diff --git a/volcano2D/Caldera_rheology.jl b/volcano2D/Caldera_rheology.jl new file mode 100644 index 000000000..16734f928 --- /dev/null +++ b/volcano2D/Caldera_rheology.jl @@ -0,0 +1,149 @@ +using GeoParams.Dislocation +using GeoParams.Diffusion + +function init_rheology_nonNewtonian() + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + + el = ConstantElasticity(; G = 40e9) + + lithosphere_rheology = CompositeRheology( (el, disl_wet_olivine, diff_wet_olivine)) + init_rheologies(lithosphere_rheology) +end + +function init_rheology_nonNewtonian_plastic() + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + # plasticity + ϕ_wet_olivine = asind(0.1) + C_wet_olivine = 1e6 + η_reg = 1e16 + el = ConstantElasticity(; G = 40e9, ν = 0.45) + lithosphere_rheology = CompositeRheology( + ( + el, + disl_wet_olivine, + diff_wet_olivine, + DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ) + init_rheologies(lithosphere_rheology) +end + +function init_rheology_linear() + el = ConstantElasticity(; G = 40e9, ν = 0.49) + # lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), )) + lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), el)) + init_rheologies(lithosphere_rheology) +end + +function init_rheologies(; linear=false) + + η_reg = 1e18 + C = linear ? Inf : 10e6 + ϕ = 15 + soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 2) # nonlinear softening law + # soft_C = NonLinearSoftening() # nonlinear softening law + pl = DruckerPrager_regularised(; C=C, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C=soft_C) + el = ConstantElasticity(; G = 25e9, ν = 0.45) + β = 1 / el.Kb.val + Cp = 1200.0 + + # Define rheolgy struct + rheology = ( + # Name = "Upper crust", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=2.7e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23), el, pl)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Lower crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=2.7e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), + Gravity = ConstantGravity(; g=9.81), + ), + + # Name = "magma chamber", + SetMaterialParams(; + Phase = 3, + Density = T_Density(; ρ0=2.5e3, T0=273.15), + Conductivity = ConstantConductivity(; k = 1.5), + # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), + HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), + LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), )), + ), + # Name = "magma chamber - hot anomaly", + SetMaterialParams(; + Phase = 4, + Density = T_Density(; ρ0=2.5e3, T0=273.15), + Conductivity = ConstantConductivity(; k = 1.5), + # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), + HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), + LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), )), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 5, + Density = ConstantDensity(; ρ=0e0), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), + Gravity = ConstantGravity(; g=9.81), + ), + ) +end + +function init_phases!(phases, phase_grid, particles, xvi) + ni = size(phases) + @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) +end + +@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} + + ni = size(phases) + + for ip in cellaxes(phases) + # quick escape + @index(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1 + ii = I[1] + offi + jj = I[2] + offj + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + ) + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj] + end + end + @index phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/volcano2D/Caldera_setup.jl b/volcano2D/Caldera_setup.jl new file mode 100644 index 000000000..885b0c84d --- /dev/null +++ b/volcano2D/Caldera_setup.jl @@ -0,0 +1,82 @@ +## Model Setup +using GeophysicalModelGenerator + +function setup2D( + nx, nz; + sticky_air = 2.5, + flat = true, + chimney = false, + chamber_T = 1e3, + chamber_depth = 5e0, + chamber_radius = 2e0, + aspect_x = 1.5, +) + + Lx = Ly = 50 + x = range(0.0, Lx, nx); + y = range(0.0, Ly, 2); + z = range(-25, sticky_air, nz); + Grid = CartData(xyz_grid(x,y,z)); + + # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) + Phases = fill(5, nx, 2, nz); + + # In many (geodynamic) models, one also has to define the temperature, so lets define it as well + Temp = fill(0.0, nx, 2, nz); + + add_box!(Phases, Temp, Grid; + xlim=(minimum(Grid.x.val), maximum(Grid.x.val)), + ylim=(minimum(Grid.y.val), maximum(Grid.y.val)), + zlim=(minimum(Grid.z.val), 0.0), + phase = LithosphericPhases(Layers=[chamber_depth], Phases=[1, 2]), + T = HalfspaceCoolingTemp(Age=20) + ) + + if !flat + add_volcano!(Phases, Temp, Grid; + volcanic_phase = 1, + center = (mean(Grid.x.val), 0.0), + height = 3, + radius = 5, + crater = 0.5, + base = 0.0, + background = nothing, + T = HalfspaceCoolingTemp(Age=20) + ) + end + + add_ellipsoid!(Phases, Temp, Grid; + cen = (mean(Grid.x.val), 0, -chamber_depth), + axes = (chamber_radius * aspect_x, 2.5, chamber_radius), + phase = ConstantPhase(3), + T = ConstantTemp(T=chamber_T) + ) + + # add_sphere!(Phases, Temp, Grid; + # cen = (mean(Grid.x.val), 0, -chamber_depth), + # radius = 0.5, + # phase = ConstantPhase(4), + # T = ConstantTemp(T=chamber_T) + # ) + + if chimney + add_cylinder!(Phases, Temp, Grid; + base = (mean(Grid.x.val), 0, -chamber_depth), + cap = (mean(Grid.x.val), 0, 0e0), + radius = 0.05, + phase = ConstantPhase(2), + # T = LinearTemp(Ttop=20, Tbot=1000), + # T = ConstantTemp(T=800), + T = ConstantTemp(T=chamber_T), + ) + end + + Grid = addfield(Grid,(; Phases, Temp)) + li = (abs(last(x)-first(x)), abs(last(z)-first(z))) .* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] + T = Temp[:,1,:] .+ 273 + # write_paraview(Grid, "Volcano2D") + return li, origin, ph, T, Grid +end diff --git a/volcano2D/DK2D.jl b/volcano2D/DK2D.jl new file mode 100644 index 000000000..c8f1da6cc --- /dev/null +++ b/volcano2D/DK2D.jl @@ -0,0 +1,489 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +# Load script dependencies +using GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("DK_setup.jl") +include("DK_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[2] +macro all_k(A) + esc(:($A[$idx_k])) +end + +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + + @parallel f_x(A, B) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end + +function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) + xv, yv = xvi + + @parallel_indices (i, j) function pure_shear_x!(Vx, εbg, lx) + xi = xv[i] + Vx[i, j + 1] = εbg * (xi - lx * 0.5) + return nothing + end + + @parallel_indices (i, j) function pure_shear_y!(Vy, εbg, ly) + yi = yv[j] + Vy[i + 1, j] = abs(yi) * εbg + return nothing + end + + nx, ny = size(Vx) + @parallel (1:nx, 1:(ny - 2)) pure_shear_x!(Vx, εbg, lx) + nx, ny = size(Vy) + @parallel (1:(nx - 2), 1:ny) pure_shear_y!(Vy, εbg, ly) + + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) + + # Physical domain ------------------------------------ + ni = nx, ny # number of cells + di = @. li / ni # grid steps + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies() + rheology_linear = init_rheologies(; linear=true) + dt = 1e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # dt = Inf # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell = 40 + max_xcell = 60 + min_xcell = 20 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # material phase & temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + + # particle fields for the stress rotation + pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress + # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress + pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity + particle_args = (pT, pPhases, pτ..., pω...) + particle_args_reduced = (pT, pτ..., pω...) + + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); + init_phases!(pPhases, phases_device, particles, xvi) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + # rock ratios for variational stokes + # RockRatios + air_phase = 5 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 1e0 , r=0.7, CFL = 0.98 / √2.1) # Re=3π, r=0.7 + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + Ttop = T_GMG[1, end] + Tbot = T_GMG[1, 1] + thermal = ThermalArrays(backend, ni) + @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + ) + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + + # Rheology + args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) + viscosity_cutoff = (1e18, 1e23) + compute_viscosity!(stokes, phase_ratios, args0, rheology, air_phase, viscosity_cutoff) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + backend, rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-8, CFL=0.95 / √2 + ) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true), + free_surface = false, + ) + εbg = 1e-15 * 0 + apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) + # ---------------------------------------------------- + + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + end + + T_buffer = @zeros(ni.+1) + Told_buffer = similar(T_buffer) + dt₀ = similar(stokes.P) + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle!(pT, xvi, T_buffer, particles) + id = @. pPhases.data == 3 + @views pT.data[id] .= 750+273 + + ## Plot initial T and P profile + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:]), + Yv, + ) + lines!( + ax2, + # Array(stokes.P[:]), + Array(log10.(stokes.viscosity.η[:])), + Y, + ) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + τxx_v = @zeros(ni.+1...) + τyy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + thermal.Told .= thermal.T + + println("Starting warm up linear solver...") + args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + t_stokes = @elapsed solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology_linear, + args, + dt, + igg; + kwargs = (; + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = viscosity_cutoff, + ) + ) + println("... finished warm up linear solver") + while it < 300 #000 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + + particle2centroid!(stokes.τ.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ.yy, pτyy, xci, particles) + particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + + t_stokes = @elapsed solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = viscosity_cutoff, + ) + ) + + center2vertex!(τxx_v, stokes.τ.xx) + center2vertex!(τyy_v, stokes.τ.yy) + centroid2particle!(pτxx , xci, stokes.τ.xx, particles) + centroid2particle!(pτyy , xci, stokes.τ.yy, particles) + grid2particle!(pτxy, xvi, stokes.τ.xy, particles) + rotate_stress_particles!(pτ, pω, particles, dt) + + println("Stokes solver time ") + println(" Total time: $t_stokes s") + tensor_invariant!(stokes.ε) + tensor_invariant!(stokes.ε_pl) + dt_max = 1e3 * 3600 * 24 * 365.25 + dt = compute_dt(stokes, di, dt_max) * 0.8 + + println("dt = $(dt/(3600 * 24 *365.25)) years") + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 50e3, + nout = 1e2, + verbose = true, + ) + ) + thermal.ΔT .= thermal.T .- thermal.Told + vertex2center!(thermal.ΔTc, thermal.ΔT[2:end-1, :]) + + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + inject_particles_phase!( + particles, + pPhases, + particle_args_reduced, + (T_buffer, τxx_v, τyy_v, stokes.τ.xy, stokes.ω.xy), + xvi + ) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + + @show it += 1 + t += dt + + if it > 1 && rem(it, 50) == 0 + id = @. pPhases.data == 3 + @views pT.data[id] .= 750+273 + end + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 2) == 0 + # checkpointing(figdir, stokes, thermal.T, η, t) + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(T_buffer), + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + ) + data_c = (; + P = Array(stokes.P), + η = Array(η_vep), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + # clr = pT.data[:] + idxv = particles.index.data[:]; + # Make Makie figure + ar = 1 + fig = Figure(size = (1200, 900), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "EII plastic") + # ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") + # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") + ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII_plastic)") + # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) + # Plot particles phase + # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) + h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.EII_pl) , colormap=:batlow) + # Plot 2nd invariant of strain rate + # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:lipari) + # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + fig + save(joinpath(figdir, "$(it).png"), fig) + + # ## Plot initial T and P profile + # fig = 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(; size=(1200, 900)) + # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + # scatter!( + # ax1, + # Array(thermal.T[2:(end - 1), :][:]), + # Yv, + # ) + # lines!( + # ax2, + # Array(stokes.P[:]), + # Y, + # ) + # hideydecorations!(ax2) + # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + # fig + # end + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "DK_2D_pulses" +n = 128 +nx, ny = n, n +li, origin, phases_GMG, T_GMG = setup2D( + nx+1, ny+1; + sticky_air = 1, + flat = true, + chimney = false, + chamber_T = 750, + chamber_depth = 5e0, + chamber_radius = 1.5, + aspect_x = 1, +) + +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/volcano2D/DK_rheology.jl b/volcano2D/DK_rheology.jl new file mode 100644 index 000000000..138c52d8b --- /dev/null +++ b/volcano2D/DK_rheology.jl @@ -0,0 +1,152 @@ +using GeoParams.Dislocation +using GeoParams.Diffusion + +function init_rheology_nonNewtonian() + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + + el = ConstantElasticity(; G = 40e9) + + lithosphere_rheology = CompositeRheology( (el, disl_wet_olivine, diff_wet_olivine)) + init_rheologies(lithosphere_rheology) +end + +function init_rheology_nonNewtonian_plastic() + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + # plasticity + ϕ_wet_olivine = asind(0.1) + C_wet_olivine = 1e6 + η_reg = 1e16 + el = ConstantElasticity(; G = 40e9, ν = 0.45) + lithosphere_rheology = CompositeRheology( + ( + el, + disl_wet_olivine, + diff_wet_olivine, + DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ) + init_rheologies(lithosphere_rheology) +end + +function init_rheologies(; linear=false) + + η_reg = 1e18 + C = linear ? Inf : 15e6 + ϕ = 30 + soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 2) # nonlinear softening law + # soft_C = NonLinearSoftening() # nonlinear softening law + pl = DruckerPrager_regularised(; C=C, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C=soft_C) + β = 1e-11 + el = ConstantElasticity(; G = 60e9, Kb = 1 / β) + Cp = 1050.0 + + disl_creep = DislocationCreep(; + A = 1.67e-24, + n = 3.5, + V = 0, + E = 1.87e5, + R = 8.314, + ) + g = 0 + # Define rheolgy struct + rheology = ( + # Name = "Upper crust", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (disl_creep, el, pl)), + RadioactiveHeat = ConstantRadioactiveHeat(1e-6), + Gravity = ConstantGravity(; g=g), + ), + # Name = "Lower crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (disl_creep, el, pl)), + RadioactiveHeat = ConstantRadioactiveHeat(1e-6), + Gravity = ConstantGravity(; g=g), + ), + + # Name = "magma chamber", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (disl_creep, el, pl)), + RadioactiveHeat = ConstantRadioactiveHeat(1e-6), + Gravity = ConstantGravity(; g=g), + ), + # Name = "magma chamber - hot anomaly", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (disl_creep, el, pl)), + RadioactiveHeat = ConstantRadioactiveHeat(1e-6), + Gravity = ConstantGravity(; g=g), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 5, + Density = ConstantDensity(; ρ=0e0), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), + Gravity = ConstantGravity(; g=g), + ), + ) +end + +function init_phases!(phases, phase_grid, particles, xvi) + ni = size(phases) + @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) +end + +@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} + + ni = size(phases) + + for ip in cellaxes(phases) + # quick escape + @index(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1 + ii = I[1] + offi + jj = I[2] + offj + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + ) + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj] + end + end + @index phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/volcano2D/DK_setup.jl b/volcano2D/DK_setup.jl new file mode 100644 index 000000000..03e5c1c84 --- /dev/null +++ b/volcano2D/DK_setup.jl @@ -0,0 +1,82 @@ +## Model Setup +using GeophysicalModelGenerator + +function setup2D( + nx, nz; + sticky_air = 2.5, + flat = true, + chimney = false, + chamber_T = 1e3, + chamber_depth = 5e0, + chamber_radius = 2e0, + aspect_x = 1.5, +) + + Lx = Ly = 10.5 + x = range(0.0, Lx, nx); + y = range(0.0, Ly, 2); + z = range(-Lx, sticky_air, nz); + Grid = CartData(xyz_grid(x,y,z)); + + # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) + Phases = fill(5, nx, 2, nz); + + # In many (geodynamic) models, one also has to define the temperature, so lets define it as well + Temp = fill(0.0, nx, 2, nz); + + add_box!(Phases, Temp, Grid; + xlim =(minimum(Grid.x.val), maximum(Grid.x.val)), + ylim =(minimum(Grid.y.val), maximum(Grid.y.val)), + zlim =(minimum(Grid.z.val), 0.0), + phase = LithosphericPhases(Layers=[chamber_depth], Phases=[1, 2]), + T = LinearTemp(Ttop=0, Tbot=450) + ) + + if !flat + add_volcano!(Phases, Temp, Grid; + volcanic_phase = 1, + center = (mean(Grid.x.val), 0.0), + height = 3, + radius = 5, + crater = 0.5, + base = 0.0, + background = nothing, + T = LinearTemp(Ttop=0, Tbot=450) + ) + end + + add_ellipsoid!(Phases, Temp, Grid; + cen = (mean(Grid.x.val), 0, -chamber_depth), + axes = (chamber_radius * aspect_x, 2.5, chamber_radius), + phase = ConstantPhase(3), + T = ConstantTemp(T=chamber_T) + ) + + # add_sphere!(Phases, Temp, Grid; + # cen = (mean(Grid.x.val), 0, -chamber_depth), + # radius = 0.5, + # phase = ConstantPhase(4), + # T = ConstantTemp(T=chamber_T) + # ) + + if chimney + add_cylinder!(Phases, Temp, Grid; + base = (mean(Grid.x.val), 0, -chamber_depth), + cap = (mean(Grid.x.val), 0, 0e0), + radius = 0.05, + phase = ConstantPhase(2), + # T = LinearTemp(Ttop=20, Tbot=1000), + # T = ConstantTemp(T=800), + T = ConstantTemp(T=chamber_T), + ) + end + + Grid = addfield(Grid,(; Phases, Temp)) + li = (abs(last(x)-first(x)), abs(last(z)-first(z))) .* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] + T = Temp[:,1,:] .+ 273 + # write_paraview(Grid, "Volcano2D") + return li, origin, ph, T, Grid +end From 9cfe89245d60163be0e230a95e144619959fab29 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 22 Nov 2024 16:48:22 +0100 Subject: [PATCH 21/77] some fixes --- src/thermal_diffusion/DiffusionPT_kernels.jl | 2 +- src/variational_stokes/Stokes2D.jl | 2 +- src/variational_stokes/VelocityKernels.jl | 81 +++++++++++++------- 3 files changed, 54 insertions(+), 31 deletions(-) diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl index 05b3b3d04..b1b5d47a0 100644 --- a/src/thermal_diffusion/DiffusionPT_kernels.jl +++ b/src/thermal_diffusion/DiffusionPT_kernels.jl @@ -1,5 +1,5 @@ isNotDirichlet(m, inds::Vararg{Int,N}) where {N} = iszero(m[inds...]) -isNotDirichlet(::Nothing, ::Vararg{Int,N}) where {N} = false +isNotDirichlet(::Nothing, ::Vararg{Int,N}) where {N} = true ## 3D KERNELS diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index b4f7bac1c..9be98e045 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -110,7 +110,7 @@ function _solve_VS!( update_ρg!(ρg[2], phase_ratios, rheology, args) @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., ϕ, _di... ) update_viscosity!( diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 788820d0d..4a4bc7677 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,32 +1,3 @@ -# @parallel_indices (I...) function compute_∇V!( -# ∇V::AbstractArray{T,2}, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy -# ) where {T} -# @inline d_xi(A) = _d_xi(A, _dx, I...) -# @inline d_yi(A) = _d_yi(A, _dy, I...) - -# if isvalid_c(ϕ, I...) -# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) -# else -# @inbounds ∇V[I...] = zero(T) -# end -# return nothing -# end - -# @parallel_indices (I...) function compute_∇V!( -# ∇V::AbstractArray{T,2}, Vx, Vy, Vz, ϕ::JustRelax.RockRatio, _dx, _dy, _dz -# ) where {T} -# @inline d_xi(A) = _d_xi(A, _dx, I...) -# @inline d_yi(A) = _d_yi(A, _dy, I...) -# @inline d_zi(A) = _d_zi(A, _dz, I...) - -# if isvalid_c(ϕ, I...) -# @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz) -# else -# @inbounds ∇V[I...] = zero(T) -# end -# return nothing -# end - @parallel_indices (I...) function compute_∇V!( ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} ) where {T,N} @@ -38,6 +9,58 @@ return nothing end +@parallel_indices (i, j) function compute_strain_rate!( + εxx::AbstractArray{T,2}, εyy, εxy, ∇V, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy +) where {T} + @inline d_xi(A) = _d_xi(A, _dx, i, j) + @inline d_yi(A) = _d_yi(A, _dy, i, j) + @inline d_xa(A) = _d_xa(A, _dx, i, j) + @inline d_ya(A) = _d_ya(A, _dy, i, j) + @inline d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + @inline d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + @inline d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + @inline d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + + if all((i, j) .≤ size(εxx)) + if isvalid_c(ϕ, i, j) + ∇V_ij = ∇V[i, j] / 3 + εxx[i, j] = d_xi(Vx) - ∇V_ij + εyy[i, j] = d_yi(Vy) - ∇V_ij + else + εxx[i, j] = zero(T) + εyy[i, j] = zero(T) + end + # εxx[i, j] = (Vx[i + 1, j + 1] * ϕ.Vx[i + 1, j] - Vx[i, j + 1] * ϕ.Vx[i, j]) * _dx - ∇V_ij + # εyy[i, j] = (Vy[i + 1, j + 1] * ϕ.Vy[i, j + 1] - Vy[i + 1, j] * ϕ.Vy[i, j]) * _dy - ∇V_ij + end + + εxy[i, j] = if isvalid_v(ϕ, i, j) + 0.5 * ( + d_ya(Vx) + + d_xa(Vy) + ) + else + zero(T) + end + + # εxy[i, j] = 0.5 * ( + # d_ya(Vx) + + # d_xa(Vy) + # ) + # vy_mask_left = ϕ.Vy[max(i - 1, 1), j] + # vy_mask_right = ϕ.Vy[min(i + 1, size(ϕ.Vy, 1)), j] + + # vx_mask_bot = ϕ.Vx[i, max(j - 1, 1)] + # vx_mask_top = ϕ.Vx[i, min(j + 1, size(ϕ.Vx, 2))] + + # εxy[i, j] = 0.5 * ( + # (Vx[i, j+1] * vx_mask_top - Vx[i, j] * vx_mask_bot) * _dy + + # (Vy[i+1, j] * vy_mask_right - Vy[i, j] * vy_mask_left) * _dx + # ) + + return nothing +end + @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, From c4abfa7669d42ea7c1d9ed043e8b41e8794dd3c2 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 26 Nov 2024 13:26:55 +0100 Subject: [PATCH 22/77] up --- volcano2D/Caldera2D.jl | 155 ++++++++++++++++++---------------- volcano2D/Caldera_rheology.jl | 17 ++-- volcano2D/Caldera_setup.jl | 4 +- 3 files changed, 97 insertions(+), 79 deletions(-) diff --git a/volcano2D/Caldera2D.jl b/volcano2D/Caldera2D.jl index 4edcd8df2..e05ded292 100644 --- a/volcano2D/Caldera2D.jl +++ b/volcano2D/Caldera2D.jl @@ -97,7 +97,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Physical properties using GeoParams ---------------- rheology = init_rheologies() - dt = 1e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + dt = 5e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter # dt = Inf # diffusive CFL timestep limiter # ---------------------------------------------------- @@ -138,12 +138,9 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.98 / √2.1) # Re=3π, r=0.7 - # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - Ttop = T_GMG[1, end] - Tbot = T_GMG[1, 1] thermal = ThermalArrays(backend, ni) @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) thermal_bc = TemperatureBoundaryConditions(; @@ -151,12 +148,14 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) + Ttop = thermal.T[2:end-1, end] + Tbot = thermal.T[2:end-1, 1] # ---------------------------------------------------- # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(2)) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) # Rheology args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) @@ -169,12 +168,21 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ) # Boundary conditions + # flow_bcs = DisplacementBoundaryConditions(; flow_bcs = VelocityBoundaryConditions(; free_slip = (left = true , right = true , top = true , bot = true), free_surface = false, ) - εbg = 1e-15 + + # U = 0.02 + # stokes.U.Ux .= PTArray(backend)([(x - li[1] * 0.5) * U / dt for x in xvi[1], _ in 1:ny+2]) + # stokes.U.Uy .= PTArray(backend)([-y * U / dt for _ in 1:nx+2, y in xvi[2]]) + # flow_bcs!(stokes, flow_bcs) # apply boundary conditions + # displacement2velocity!(stokes, dt) + + εbg = 1e-15 * 1 apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(@velocity(stokes)...) @@ -203,26 +211,26 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk grid2particle!(pT, xvi, T_buffer, particles) ## Plot initial T and P profile - fig = 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:]), - Yv, - ) - lines!( - ax2, - Array(stokes.P[:]), - Y, - ) - hideydecorations!(ax2) - # save(joinpath(figdir, "initial_profile.png"), fig) - fig - end + # fig = 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(; size=(1200, 900)) + # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + # scatter!( + # ax1, + # Array(thermal.T[2:(end - 1), :][:]), + # Yv, + # ) + # lines!( + # ax2, + # Array(stokes.P[:]), + # Y, + # ) + # hideydecorations!(ax2) + # # save(joinpath(figdir, "initial_profile.png"), fig) + # fig + # end τxx_v = @zeros(ni.+1...) τyy_v = @zeros(ni.+1...) @@ -231,7 +239,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk t, it = 0.0, 0 thermal.Told .= thermal.T - while it < 100 #000 # run only for 5 Myrs + while it < 500 #000 # run only for 5 Myrs # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) @@ -241,12 +249,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) - args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) - # args = (; T=thermal.Tc, P=stokes.P, dt=Inf) - - particle2centroid!(stokes.τ.xx, pτxx, xci, particles) - particle2centroid!(stokes.τ.yy, pτyy, xci, particles) - particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + args = (; T=thermal.Tc, P=stokes.P, dt=Inf) t_stokes = @elapsed solve_VariationalStokes!( stokes, @@ -278,13 +282,13 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk println(" Total time: $t_stokes s") tensor_invariant!(stokes.ε) tensor_invariant!(stokes.ε_pl) - dtmax = 1e3 * 3600 * 24 * 365.25 - dt = compute_dt(stokes, di) * 0.8 + dtmax = 2e3 * 3600 * 24 * 365.25 + dt = compute_dt(stokes, di, dtmax) println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ - # # Thermal solver --------------- + # Thermal solver --------------- heatdiffusion_PT!( thermal, pt_thermal, @@ -332,11 +336,15 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + particle2centroid!(stokes.τ_o.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ_o.yy, pτyy, xci, particles) + particle2grid!(stokes.τ_o.xy, pτxy, xvi, particles) + @show it += 1 t += dt # Data I/O and plotting --------------------- - if it == 1 || rem(it, 2) == 0 + if it == 1 || rem(it, 1) == 0 # checkpointing(figdir, stokes, thermal.T, η, t) (; η_vep, η) = stokes.viscosity if do_vtk @@ -345,6 +353,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk T = Array(T_buffer), τII = Array(stokes.τ.II), εII = Array(stokes.ε.II), + εII_pl = Array(stokes.ε_pl.II), ) data_c = (; P = Array(stokes.P), @@ -360,7 +369,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk xci, data_v, data_c, - velocity_v + velocity_v; + t = 0 ) end @@ -379,19 +389,19 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") - ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII_plastic)") - # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII)") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") # Plot temperature h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) # Plot particles phase - # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) - h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) # Plot 2nd invariant of strain rate # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) # Plot effective viscosity - h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:lipari) - # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:lipari) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) @@ -403,27 +413,27 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk fig save(joinpath(figdir, "$(it).png"), fig) - ## Plot initial T and P profile - fig = 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:]), - Yv, - ) - lines!( - ax2, - Array(stokes.P[:]), - Y, - ) - hideydecorations!(ax2) - save(joinpath(figdir, "thermal_profile_$it.png"), fig) - fig - end + # ## Plot initial T and P profile + # fig = 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(; size=(1200, 900)) + # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + # scatter!( + # ax1, + # Array(thermal.T[2:(end - 1), :][:]), + # Yv, + # ) + # lines!( + # ax2, + # Array(stokes.P[:]), + # Y, + # ) + # hideydecorations!(ax2) + # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + # fig + # end end # ------------------------------ @@ -434,21 +444,22 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- do_vtk = true # set to true to generate VTK files for ParaView +# figdir = "Caldera2D_noPguess" figdir = "Caldera2D" -n = 256 +n = 128 * 2 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; - sticky_air = 2.5, - flat = true, + sticky_air = 4, + flat = false, chimney = false, chamber_T = 1e3, chamber_depth = 7e0, - chamber_radius = 1e0/2, - aspect_x = 10, + chamber_radius = 0.5, + aspect_x = 6, ) -heatmap(phases_GMG) +# heatmap(phases_GMG) # heatmap(T_GMG) igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -457,4 +468,4 @@ else igg end -main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); \ No newline at end of file +main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/volcano2D/Caldera_rheology.jl b/volcano2D/Caldera_rheology.jl index 16734f928..9f7d6bd3f 100644 --- a/volcano2D/Caldera_rheology.jl +++ b/volcano2D/Caldera_rheology.jl @@ -52,6 +52,11 @@ function init_rheologies(; linear=false) el = ConstantElasticity(; G = 25e9, ν = 0.45) β = 1 / el.Kb.val Cp = 1200.0 + + #dislocation laws + disl_top = SetDislocationCreep(Dislocation.dry_olivine_Karato_2003) + # diffusion laws + disl_bot = SetDislocationCreep(Dislocation.wet_quartzite_Hirth_2001) # Define rheolgy struct rheology = ( @@ -61,7 +66,8 @@ function init_rheologies(; linear=false) Density = PT_Density(; ρ0=2.7e3, T0=273.15, β=β), HeatCapacity = ConstantHeatCapacity(; Cp = Cp), Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23), el, pl)), + CompositeRheology = CompositeRheology( (disl_top, el, pl)), + # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23), el, pl)), Gravity = ConstantGravity(; g=9.81), ), # Name = "Lower crust", @@ -70,7 +76,8 @@ function init_rheologies(; linear=false) Density = PT_Density(; ρ0=2.7e3, T0=273.15, β=β), HeatCapacity = ConstantHeatCapacity(; Cp = Cp), Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), + CompositeRheology = CompositeRheology( (disl_bot, el, pl)), + # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), Gravity = ConstantGravity(; g=9.81), ), @@ -82,7 +89,7 @@ function init_rheologies(; linear=false) # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), )), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el)), ), # Name = "magma chamber - hot anomaly", SetMaterialParams(; @@ -92,12 +99,12 @@ function init_rheologies(; linear=false) # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), )), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el, )), ), # Name = "StickyAir", SetMaterialParams(; Phase = 5, - Density = ConstantDensity(; ρ=0e0), + Density = ConstantDensity(; ρ=1e0), HeatCapacity = ConstantHeatCapacity(; Cp = Cp), Conductivity = ConstantConductivity(; k = 2.5), CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), diff --git a/volcano2D/Caldera_setup.jl b/volcano2D/Caldera_setup.jl index 885b0c84d..32f60911f 100644 --- a/volcano2D/Caldera_setup.jl +++ b/volcano2D/Caldera_setup.jl @@ -62,9 +62,9 @@ function setup2D( if chimney add_cylinder!(Phases, Temp, Grid; base = (mean(Grid.x.val), 0, -chamber_depth), - cap = (mean(Grid.x.val), 0, 0e0), + cap = (mean(Grid.x.val), 0, 3e0), radius = 0.05, - phase = ConstantPhase(2), + phase = ConstantPhase(3), # T = LinearTemp(Ttop=20, Tbot=1000), # T = ConstantTemp(T=800), T = ConstantTemp(T=chamber_T), From 1a9fc5035046e1a8597f2bb1e15956a546128342 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 26 Nov 2024 16:24:29 +0100 Subject: [PATCH 23/77] up --- volcano2D/Caldera2D.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/volcano2D/Caldera2D.jl b/volcano2D/Caldera2D.jl index e05ded292..70c7da795 100644 --- a/volcano2D/Caldera2D.jl +++ b/volcano2D/Caldera2D.jl @@ -276,6 +276,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk centroid2particle!(pτxx , xci, stokes.τ.xx, particles) centroid2particle!(pτyy , xci, stokes.τ.yy, particles) grid2particle!(pτxy, xvi, stokes.τ.xy, particles) + grid2particle!(pωxy, xvi, stokes.ω.xy, particles) rotate_stress_particles!(pτ, pω, particles, dt) println("Stokes solver time ") From 3a55952660ce8a04c9ef07afd0a835dd59900636 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 26 Nov 2024 16:41:01 +0100 Subject: [PATCH 24/77] up script again --- volcano2D/Caldera2D.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/volcano2D/Caldera2D.jl b/volcano2D/Caldera2D.jl index 70c7da795..4d0ee07ca 100644 --- a/volcano2D/Caldera2D.jl +++ b/volcano2D/Caldera2D.jl @@ -319,6 +319,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # ------------------------------ # Advection -------------------- + copyinn_x!(T_buffer, thermal.T) # advect particles in space advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) # advect particles in memory From dc88598c87c726021f22348ae26bd62d05295ab7 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 29 Nov 2024 17:34:47 +0100 Subject: [PATCH 25/77] path way to 3D variational stokes --- .../stress_rotation_particles.jl | 10 +- src/variational_stokes/MiniKernels.jl | 39 ++- src/variational_stokes/Stokes2D.jl | 8 +- src/variational_stokes/Stokes3D.jl | 221 ++++++++++++++ src/variational_stokes/StressKernels.jl | 275 ++++++++++++++++++ src/variational_stokes/VelocityKernels.jl | 168 +++++++++-- src/variational_stokes/mask.jl | 168 +++++++++-- 7 files changed, 835 insertions(+), 54 deletions(-) create mode 100644 src/variational_stokes/Stokes3D.jl diff --git a/src/stress_rotation/stress_rotation_particles.jl b/src/stress_rotation/stress_rotation_particles.jl index 34441de96..ab874237a 100644 --- a/src/stress_rotation/stress_rotation_particles.jl +++ b/src/stress_rotation/stress_rotation_particles.jl @@ -3,8 +3,8 @@ using StaticArrays # Vorticity tensor @parallel_indices (I...) function compute_vorticity!(ωxy, Vx, Vy, _dx, _dy) - @inline dx(A) = _d_xa(A, I..., _dx) - @inline dy(A) = _d_ya(A, I..., _dy) + @inline dx(A) = _d_xa(A, _dx, I...) + @inline dy(A) = _d_ya(A, _dy, I...) ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) @@ -14,9 +14,9 @@ end @parallel_indices (I...) function compute_vorticity!( ωyz, ωxz, ωxy, Vx, Vy, Vz, _dx, _dy, _dz ) - dx(A) = _d_xa(A, I..., _dx) - dy(A) = _d_ya(A, I..., _dy) - dz(A) = _d_za(A, I..., _dz) + dx(A) = _d_xa(A, _dx, I...) + dy(A) = _d_ya(A, _dy, I...) + dz(A) = _d_za(A, _dz, I...) if all(I .≤ size(ωyz)) ωyz[I...] = 0.5 * (dy(Vz) - dz(Vy)) diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 3e2055b86..6f9d217b4 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -10,29 +10,66 @@ end # finite differences Base.@propagate_inbounds @inline _d_xa(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = (-center(A, ϕ, I...) + right(A, ϕ, I...)) * _dx + Base.@propagate_inbounds @inline _d_ya(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dy + Base.@propagate_inbounds @inline _d_za(A::T, ϕ::T, _dz, I::Vararg{Integer,N}) where {N,T} = (-center(A, ϕ, I...) + front(A, ϕ, I...)) * _dz + Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) where {N,T} = (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx + Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} = (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy -# averages +Base.@propagate_inbounds @inline _d_zi(A::T, ϕ::T, _dz, I::Vararg{Integer,N}) where {N,T} = + (-top(A, ϕ, I...) + next(A, ϕ, I...)) * _dz + +# averages 2D Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j) where {T<:T2} = 0.25 * mysum(A, ϕ, (i + 1):(i + 2), (j + 1):(j + 2)) + Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j) where {T<:T2} = 0.25 * mysum(A, ϕ, (i):(i + 1), (j):(j + 1)) + Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5 + Base.@propagate_inbounds @inline _av_ya(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = (center(A, ϕ, I...) + front(A, ϕ, I...)) * 0.5 + Base.@propagate_inbounds @inline _av_xi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = (front(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 + Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} = (right(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 +# averages 3D +Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j, k) where {T<:T3} = + 0.125 * mysum(A, ϕ, (i + 1):(i + 2), (j + 1):(j + 2), (k + 1):(k + 2)) + +Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j, k) where {T<:T3} = + 0.125 * mysum(A, ϕ, (i):(i + 1), (j):(j + 1), (k):(k + 1)) + +Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5 + +Base.@propagate_inbounds @inline _av_ya(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (center(A, ϕ, I...) + front(A, ϕ, I...)) * 0.5 + +Base.@propagate_inbounds @inline _av_za(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (center(A, ϕ, I...) + top(A, ϕ, I...)) * 0.5 + +Base.@propagate_inbounds @inline _av_xi(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (front(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 + +Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (right(A, ϕ, I...), next(A, ϕ, I...)) * 0.5 + +Base.@propagate_inbounds @inline _av_zi(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = + (top(A, ϕ, I...) + next(A, ϕ, I...)) * 0.5 + ## Because mysum(::generator) does not work inside CUDA kernels... @inline mysum(A, ϕ, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ϕ, ranges...) diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 9be98e045..52feeb09d 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -81,9 +81,6 @@ function _solve_VS!( compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, viscosity_cutoff) displacement2velocity!(stokes, dt, flow_bcs) - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) - while iter ≤ iterMax iterMin < iter && err < ϵ && break @@ -100,7 +97,7 @@ function _solve_VS!( stokes.∇V, ητ, rheology, - phase_ratios.center, + phase_ratios, dt, r, θ_dτ, @@ -213,6 +210,9 @@ function _solve_VS!( # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + return ( iter=iter, err_evo1=err_evo1, diff --git a/src/variational_stokes/Stokes3D.jl b/src/variational_stokes/Stokes3D.jl new file mode 100644 index 000000000..20d28592a --- /dev/null +++ b/src/variational_stokes/Stokes3D.jl @@ -0,0 +1,221 @@ +## 3D VISCO-ELASTIC STOKES SOLVER + +# backend trait +function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs) + solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) + return nothing +end + +# entry point for extensions +function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) + return _solve_VS!(stokes, args...; kwargs...) +end + + +# GeoParams and multiple phases +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, + di::NTuple{3,T}, + flow_bcs::AbstractFlowBoundaryConditions, + ρg, + phase_ratios::JustPIC.PhaseRatios, + ϕ::JustRelax.RockRatio, + rheology::NTuple{N,AbstractMaterialParamsStruct}, + args, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 4), + verbose=true, + viscosity_relaxation=1e-2, + viscosity_cutoff=(-Inf, Inf), + kwargs..., +) where {T,N} + + ## UNPACK + + # solver related + ϵ = pt_stokes.ϵ + # geometry + _di = @. 1 / di + ni = size(stokes.P) + (; η, η_vep) = stokes.viscosity + + # errors + err = Inf + iter = 0 + cont = 0 + err_evo1 = Float64[] + err_evo2 = Int64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_Rz = Float64[] + norm_∇V = Float64[] + + @copy stokes.P0 stokes.P + θ = deepcopy(stokes.P) + λ = @zeros(ni...) + λv_yz = @zeros(size(stokes.τ.yz)...) + λv_xz = @zeros(size(stokes.τ.xz)...) + λv_xy = @zeros(size(stokes.τ.xy)...) + + # solver loop + wtime0 = 0.0 + ητ = deepcopy(η) + + # compute buoyancy forces and viscosity + compute_ρg!(ρg, phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + # convert displacement to velocity + displacement2velocity!(stokes, dt, flow_bcs) + + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + # ~preconditioner + compute_maxloc!(ητ, η) + update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + compute_P!( + θ, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + dt, + pt_stokes.r, + pt_stokes.θ_dτ, + args, + ) + + @parallel (@idx ni) compute_strain_rate!( + stokes.∇V, @strain(stokes)..., @velocity(stokes)..., ϕ, _di... + ) + + # Update buoyancy + update_ρg!(ρg, phase_ratios, rheology, args) + + # Update viscosity + update_viscosity!( + stokes, + phase_ratios, + args, + rheology, + viscosity_cutoff; + air_phase=air_phase, + relaxation=viscosity_relaxation, + ) + # update_stress!(stokes, θ, λ, phase_ratios, rheology, dt, pt_stokes.θ_dτ) + + @parallel (@idx ni .+ 1) update_stresses_center_vertex!( + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + @tensor_center(stokes.τ), + (stokes.τ.yz, stokes.τ.xz, stokes.τ.xy), + @tensor_center(stokes.τ_o), + (stokes.τ_o.yz, stokes.τ_o.xz, stokes.τ_o.xy), + θ, + stokes.P, + stokes.viscosity.η, + λ, + (λv_yz, λv_xz, λv_xy), + stokes.τ.II, + stokes.viscosity.η_vep, + 0.2, + dt, + pt_stokes.θ_dτ, + rheology, + phase_ratios.center, + phase_ratios.vertex, + phase_ratios.xy, + phase_ratios.yz, + phase_ratios.xz, + ϕ, + ) + update_halo!(stokes.τ.yz) + update_halo!(stokes.τ.xz) + update_halo!(stokes.τ.xy) + + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + @residuals(stokes.R)..., + stokes.P, + ρg..., + @stress(stokes)..., + ητ, + pt_stokes.ηdτ, + ϕ, + _di..., + ) + # apply boundary conditions + velocity2displacement!(stokes, dt) + free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + cont += 1 + for (norm_Ri, Ri) in zip((norm_Rx, norm_Ry, norm_Rz), @residuals(stokes.R)) + push!( + norm_Ri, + norm_mpi(Ri[2:(end - 1), 2:(end - 1), 2:(end - 1)]) / length(Ri), + ) + end + push!(norm_∇V, norm_mpi(stokes.R.RP) / length(stokes.R.RP)) + err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && (verbose || iter == iterMax) + @printf( + "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[cont], + norm_Ry[cont], + norm_Rz[cont], + norm_∇V[cont] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + av_time = wtime0 / (iter - 1) # average time per iteration + + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.yz, stokes.ω.xz, stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) + + # accumulate plastic strain tensor + @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_Rz=norm_Rz, + norm_∇V=norm_∇V, + time=wtime0, + av_time=av_time, + ) +end diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl index 9aa70acdf..661a0cf05 100644 --- a/src/variational_stokes/StressKernels.jl +++ b/src/variational_stokes/StressKernels.jl @@ -141,3 +141,278 @@ return nothing end + +# 3D kernel +@parallel_indices (I...) function update_stresses_center_vertex!( + ε::NTuple{6}, # normal components @ centers; shear components @ vertices + ε_pl::NTuple{6}, # whole Voigt tensor @ centers + EII, # accumulated plastic strain rate @ centers + τ::NTuple{6}, # whole Voigt tensor @ centers + τshear_v::NTuple{3}, # shear tensor components @ vertices + τ_o::NTuple{6}, + τshear_ov::NTuple{3}, # shear tensor components @ vertices + Pr, + Pr_c, + η, + λ, + λv::NTuple{3}, + τII, + η_vep, + relλ, + dt, + θ_dτ, + rheology, + phase_center, + phase_vertex, + phase_xy, + phase_yz, + phase_xz, + ϕ::JustRelax.RockRatio, +) + τyzv, τxzv, τxyv = τshear_v + τyzv_old, τxzv_old, τxyv_old = τshear_ov + + ni = size(Pr) + Ic = clamped_indices(ni, I...) + + ## yz + if all(I .≤ size(ε[4])) && isvalid_yz(ϕ, I...) + # interpolate to ith vertex + ηv_ij = av_clamped_yz(η, Ic...) + Pv_ij = av_clamped_yz(Pr, Ic...) + EIIv_ij = av_clamped_yz(EII, Ic...) + εxxv_ij = av_clamped_yz(ε[1], Ic...) + εyyv_ij = av_clamped_yz(ε[2], Ic...) + εzzv_ij = av_clamped_yz(ε[3], Ic...) + εyzv_ij = ε[4][I...] + εxzv_ij = av_clamped_yz_y(ε[5], Ic...) + εxyv_ij = av_clamped_yz_z(ε[6], Ic...) + + τxxv_ij = av_clamped_yz(τ[1], Ic...) + τyyv_ij = av_clamped_yz(τ[2], Ic...) + τzzv_ij = av_clamped_yz(τ[3], Ic...) + τyzv_ij = τyzv[I...] + τxzv_ij = av_clamped_yz_y(τxzv, Ic...) + τxyv_ij = av_clamped_yz_z(τxyv, Ic...) + + τxxv_old_ij = av_clamped_yz(τ_o[1], Ic...) + τyyv_old_ij = av_clamped_yz(τ_o[2], Ic...) + τzzv_old_ij = av_clamped_yz(τ_o[3], Ic...) + τyzv_old_ij = τyzv_old[I...] + τxzv_old_ij = av_clamped_yz_y(τxzv_old, Ic...) + τxyv_old_ij = av_clamped_yz_z(τxyv_old, Ic...) + + # vertex parameters + phase = @inbounds phase_yz[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( + rheology, EIIv_ij, phase + ) + _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + Kv = fn_ratio(get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv) + dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv) + dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv) + dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv) + dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv) + dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv) + + dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv + τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij + τIIv_ij = second_invariant(τijv .+ dτijv) + + # yield function @ vertex + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) && Fv > 0 + # stress correction @ vertex + λv[1][I...] = + (1.0 - relλ) * λv[1][I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + + dQdτyz = 0.5 * (τyzv_ij + dτyzv) / τIIv_ij + τyzv[I...] += dτyzv - 2.0 * ηv_ij * 0.5 * λv[1][I...] * dQdτyz * dτ_rv + else + # stress correction @ vertex + τyzv[I...] += dτyzv + end + end + + ## xz + if all(I .≤ size(ε[5])) && isvalid_xz(ϕ, I...) + # interpolate to ith vertex + ηv_ij = av_clamped_xz(η, Ic...) + EIIv_ij = av_clamped_xz(EII, Ic...) + Pv_ij = av_clamped_xz(Pr, Ic...) + εxxv_ij = av_clamped_xz(ε[1], Ic...) + εyyv_ij = av_clamped_xz(ε[2], Ic...) + εzzv_ij = av_clamped_xz(ε[3], Ic...) + εyzv_ij = av_clamped_xz_x(ε[4], Ic...) + εxzv_ij = ε[5][I...] + εxyv_ij = av_clamped_xz_z(ε[6], Ic...) + τxxv_ij = av_clamped_xz(τ[1], Ic...) + τyyv_ij = av_clamped_xz(τ[2], Ic...) + τzzv_ij = av_clamped_xz(τ[3], Ic...) + τyzv_ij = av_clamped_xz_x(τyzv, Ic...) + τxzv_ij = τxzv[I...] + τxyv_ij = av_clamped_xz_z(τxyv, Ic...) + τxxv_old_ij = av_clamped_xz(τ_o[1], Ic...) + τyyv_old_ij = av_clamped_xz(τ_o[2], Ic...) + τzzv_old_ij = av_clamped_xz(τ_o[3], Ic...) + τyzv_old_ij = av_clamped_xz_x(τyzv_old, Ic...) + τxzv_old_ij = τxzv_old[I...] + τxyv_old_ij = av_clamped_xz_z(τxyv_old, Ic...) + + # vertex parameters + phase = @inbounds phase_xz[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( + rheology, EIIv_ij, phase + ) + _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + Kv = fn_ratio(get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv) + dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv) + dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv) + dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv) + dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv) + dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv) + + dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv + τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij + τIIv_ij = second_invariant(τijv .+ dτijv) + + # yield function @ vertex + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) && Fv > 0 + # stress correction @ vertex + λv[2][I...] = + (1.0 - relλ) * λv[2][I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + + dQdτxz = 0.5 * (τxzv_ij + dτxzv) / τIIv_ij + τxzv[I...] += dτxzv - 2.0 * ηv_ij * 0.5 * λv[2][I...] * dQdτxz * dτ_rv + else + # stress correction @ vertex + τxzv[I...] += dτxzv + end + end + + ## xy + if all(I .≤ size(ε[6])) && isvalid_xy(ϕ, I...) + # interpolate to ith vertex + ηv_ij = av_clamped_xy(η, Ic...) + EIIv_ij = av_clamped_xy(EII, Ic...) + Pv_ij = av_clamped_xy(Pr, Ic...) + εxxv_ij = av_clamped_xy(ε[1], Ic...) + εyyv_ij = av_clamped_xy(ε[2], Ic...) + εzzv_ij = av_clamped_xy(ε[3], Ic...) + εyzv_ij = av_clamped_xy_x(ε[4], Ic...) + εxzv_ij = av_clamped_xy_y(ε[5], Ic...) + εxyv_ij = ε[6][I...] + + τxxv_ij = av_clamped_xy(τ[1], Ic...) + τyyv_ij = av_clamped_xy(τ[2], Ic...) + τzzv_ij = av_clamped_xy(τ[3], Ic...) + τyzv_ij = av_clamped_xy_x(τyzv, Ic...) + τxzv_ij = av_clamped_xy_y(τxzv, Ic...) + τxyv_ij = τxyv[I...] + + τxxv_old_ij = av_clamped_xy(τ_o[1], Ic...) + τyyv_old_ij = av_clamped_xy(τ_o[2], Ic...) + τzzv_old_ij = av_clamped_xy(τ_o[3], Ic...) + τyzv_old_ij = av_clamped_xy_x(τyzv_old, Ic...) + τxzv_old_ij = av_clamped_xy_y(τxzv_old, Ic...) + τxyv_old_ij = τxyv_old[I...] + + # vertex parameters + phase = @inbounds phase_xy[I...] + is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase( + rheology, EIIv_ij, phase + ) + _Gvdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + Kv = fn_ratio(get_bulk_modulus, rheology, phase) + volumev = isinf(Kv) ? 0.0 : Kv * dt * sinϕv * sinψv # plastic volumetric change K * dt * sinϕ * sinψ + dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0) + + # stress increments @ vertex + dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv) + dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv) + dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv) + dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv) + dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv) + dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv) + dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv + τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij + τIIv_ij = second_invariant(τijv .+ dτijv) + + # yield function @ vertex + Fv = τIIv_ij - Cv - Pv_ij * sinϕv + if is_pl && !iszero(τIIv_ij) && Fv > 0 + # stress correction @ vertex + λv[3][I...] = + (1.0 - relλ) * λv[3][I...] + + relλ * (max(Fv, 0.0) / (ηv_ij * dτ_rv + η_regv + volumev)) + + dQdτxy = 0.5 * (τxyv_ij + dτxyv) / τIIv_ij + τxyv[I...] += dτxyv - 2.0 * ηv_ij * 0.5 * λv[3][I...] * dQdτxy * dτ_rv + else + # stress correction @ vertex + τxyv[I...] += dτxyv + end + end + + ## center + if all(I .≤ ni) && isvalid_c(ϕ, I...) + # Material properties + phase = @inbounds phase_center[I...] + _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase(rheology, EII[I...], phase) + K = fn_ratio(get_bulk_modulus, rheology, phase) + volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ # plastic volumetric change K * dt * sinϕ * sinψ + ηij = η[I...] + dτ_r = inv(θ_dτ + ηij * _Gdt + 1.0) + + # cache strain rates for center calculations + τij, τij_o, εij = cache_tensors(τ, τ_o, ε, I...) + + # visco-elastic strain rates @ center + εij_ve = @. εij + 0.5 * τij_o * _Gdt + εII_ve = second_invariant(εij_ve) + # stress increments @ center + dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij + 2.0 * ηij * εij) * dτ_r + τII_ij = second_invariant(dτij .+ τij) + # yield function @ center + F = τII_ij - C - Pr[I...] * sinϕ + + if is_pl && !iszero(τII_ij) && F > 0 + # stress correction @ center + λ[I...] = + (1.0 - relλ) * λ[I...] + + relλ * (max(F, 0.0) / (η[I...] * dτ_r + η_reg + volume)) + dQdτij = @. 0.5 * (τij + dτij) / τII_ij + εij_pl = λ[I...] .* dQdτij + dτij = @. dτij - 2.0 * ηij * εij_pl * dτ_r + τij = dτij .+ τij + setindex!.(τ, τij, I...) + setindex!.(ε_pl, εij_pl, I...) + τII[I...] = second_invariant(τij) + Pr_c[I...] = Pr[I...] + K * dt * λ[I...] * sinψ + η_vep[I...] = 0.5 * τII_ij / εII_ve + else + # stress correction @ center + setindex!.(τ, dτij .+ τij, I...) + η_vep[I...] = ηij + τII[I...] = τII_ij + end + + Pr_c[I...] = Pr[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) + end + + return nothing +end diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 4a4bc7677..4acfd6b2f 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -16,10 +16,6 @@ end @inline d_yi(A) = _d_yi(A, _dy, i, j) @inline d_xa(A) = _d_xa(A, _dx, i, j) @inline d_ya(A) = _d_ya(A, _dy, i, j) - @inline d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) - @inline d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) - @inline d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) - @inline d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) if all((i, j) .≤ size(εxx)) if isvalid_c(ϕ, i, j) @@ -30,34 +26,75 @@ end εxx[i, j] = zero(T) εyy[i, j] = zero(T) end - # εxx[i, j] = (Vx[i + 1, j + 1] * ϕ.Vx[i + 1, j] - Vx[i, j + 1] * ϕ.Vx[i, j]) * _dx - ∇V_ij - # εyy[i, j] = (Vy[i + 1, j + 1] * ϕ.Vy[i, j + 1] - Vy[i + 1, j] * ϕ.Vy[i, j]) * _dy - ∇V_ij end εxy[i, j] = if isvalid_v(ϕ, i, j) - 0.5 * ( - d_ya(Vx) + - d_xa(Vy) - ) + 0.5 * (d_ya(Vx) + d_xa(Vy)) else zero(T) end - # εxy[i, j] = 0.5 * ( - # d_ya(Vx) + - # d_xa(Vy) - # ) - # vy_mask_left = ϕ.Vy[max(i - 1, 1), j] - # vy_mask_right = ϕ.Vy[min(i + 1, size(ϕ.Vy, 1)), j] - - # vx_mask_bot = ϕ.Vx[i, max(j - 1, 1)] - # vx_mask_top = ϕ.Vx[i, min(j + 1, size(ϕ.Vx, 2))] - - # εxy[i, j] = 0.5 * ( - # (Vx[i, j+1] * vx_mask_top - Vx[i, j] * vx_mask_bot) * _dy + - # (Vy[i+1, j] * vy_mask_right - Vy[i, j] * vy_mask_left) * _dx - # ) + return nothing +end + +@parallel_indices (i, j, k) function compute_strain_rate!( + ∇V::AbstractArray{T,3}, + εxx, + εyy, + εzz, + εyz, + εxz, + εxy, + Vx, + Vy, + Vz, + ϕ::JustRelax.RockRatio, + _dx, + _dy, + _dz, +) where {T} + d_xi(A) = _d_xi(A, _dx, i, j, k) + d_yi(A) = _d_yi(A, _dy, i, j, k) + d_zi(A) = _d_zi(A, _dz, i, j, k) + @inbounds begin + # normal components are all located @ cell centers + if all((i, j, k) .≤ size(εxx)) + if isvalid_c(ϕ, i, j, k) + ∇Vijk = ∇V[i, j, k] * inv(3) + # Compute ε_xx + εxx[i, j, k] = d_xi(Vx) - ∇Vijk + # Compute ε_yy + εyy[i, j, k] = d_yi(Vy) - ∇Vijk + # Compute ε_zz + εzz[i, j, k] = d_zi(Vz) - ∇Vijk + end + end + # Compute ε_yz + if all((i, j, k) .≤ size(εyz)) && isvalid_yz(ϕ, i, j, k) + εyz[i, j, k] = + 0.5 * ( + _dz * (Vy[i + 1, j, k + 1] - Vy[i + 1, j, k]) + + _dy * (Vz[i + 1, j + 1, k] - Vz[i + 1, j, k]) + ) + end + # Compute ε_xz + if all((i, j, k) .≤ size(εxz)) && isvalid_xz(ϕ, i, j, k) + εxz[i, j, k] = + 0.5 * ( + _dz * (Vx[i, j + 1, k + 1] - Vx[i, j + 1, k]) + + _dx * (Vz[i + 1, j + 1, k] - Vz[i, j + 1, k]) + ) + end + # Compute ε_xy + if all((i, j, k) .≤ size(εxy)) && isvalid_xy(ϕ, i, j, k) + εxy[i, j, k] = + 0.5 * ( + _dy * (Vx[i, j + 1, k + 1] - Vx[i, j, k + 1]) + + _dx * (Vy[i + 1, j, k + 1] - Vy[i, j, k + 1]) + ) + end + end return nothing end @@ -84,6 +121,7 @@ end d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) av_xa(A, ϕ) = _av_xa(A, ϕ, i, j) av_ya(A, ϕ) = _av_ya(A, ϕ, i, j) + av_xa(A) = _av_xa(A, i, j) av_ya(A) = _av_ya(A, i, j) harm_xa(A) = _av_xa(A, i, j) @@ -118,3 +156,85 @@ end return nothing end + +@parallel_indices (i, j, k) function compute_V!( + Vx::AbstractArray{T,3}, + Vy, + Vz, + Rx, + Ry, + Rz, + P, + fx, + fy, + fz, + τxx, + τyy, + τzz, + τyz, + τxz, + τxy, + ητ, + ηdτ, + ϕ::JustRelax.RockRatio, + _dx, + _dy, + _dz, +) where {T} + @inline harm_x(A) = _harm_x(A, i, j, k) + @inline harm_y(A) = _harm_y(A, i, j, k) + @inline harm_z(A) = _harm_z(A, i, j, k) + @inline d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j, k) + @inline d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j, k) + @inline d_za(A, ϕ) = _d_za(A, ϕ, _dz, i, j, k) + @inline d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j, k) + @inline d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j, k) + @inline d_zi(A, ϕ) = _d_zi(A, ϕ, _dz, i, j, k) + @inline av_x(A) = _av_x(A, i, j, k) + @inline av_y(A) = _av_y(A, i, j, k) + @inline av_z(A) = _av_z(A, i, j, k) + @inline av_x(A, ϕ) = _av_x(A, ϕ, i, j, k) + @inline av_y(A, ϕ) = _av_y(A, ϕ, i, j, k) + @inline av_z(A, ϕ) = _av_z(A, ϕ, i, j, k) + + @inbounds begin + if all((i, j, k) .< size(Vx) .- 1) + if isvalid_vx(ϕ, i + 1, j, k) + Rx_ijk = + Rx[i, j, k] = + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.xy) + d_zi(τxz, ϕ.xz) - + d_xa(P, ϕ.center) - av_x(fx, ϕ.center) + Vx[i + 1, j + 1, k + 1] += Rx_ijk * ηdτ / av_x(ητ) + else + Rx[i, j, k] = zero(T) + Vx[i + 1, j + 1, k + 1] = zero(T) + end + end + if all((i, j, k) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j + 1, k) + Ry_ijk = + Ry[i, j, k] = + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.xy) + d_zi(τyz, ϕ.yz) - + d_ya(P, ϕ.center) - av_y(fy, ϕ.center) + Vy[i + 1, j + 1, k + 1] += Ry_ijk * ηdτ / av_y(ητ) + else + Ry[i, j, k] = zero(T) + Vy[i + 1, j + 1, k + 1] = zero(T) + end + end + if all((i, j, k) .< size(Vz) .- 1) + if isvalid_vz(ϕ, i, j, k + 1) + Rz_ijk = + Rz[i, j, k] = + d_za(τzz, ϕ.center) + d_xi(τxz, ϕ.xz) + d_yi(τyz, ϕ.yz) - + d_za(P, ϕ.center) - av_z(fz, ϕ.center) + Vz[i + 1, j + 1, k + 1] += Rz_ijk * ηdτ / av_z(ητ) + else + Rz[i, j, k] = zero(T) + Vz[i + 1, j + 1, k + 1] = zero(T) + end + end + end + + return nothing +end diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 3ad186c50..01be3cc98 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -96,10 +96,18 @@ Check if `ϕ.center[inds...]` is a not a nullspace. - `inds`: Cartesian indices to check. """ Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) - vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0) - vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0) + vx = isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vx[i + 1, j]) + vy = isvalid(ϕ.Vy, i, j) * isvalid(ϕ.Vy[i, j + 1]) v = vx * vy - return v * (ϕ.center[i, j] > 0) + return v * isvalid(ϕ.center, i, j) +end + +Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j, k) + vx = isvalid(ϕ.Vx, i, j, k) * isvalid(ϕ.Vx, i + 1, j, k) + vy = isvalid(ϕ.Vy, i, j, k) * isvalid(ϕ.Vy, i, j + 1, k) + vz = isvalid(ϕ.Vz, i, j, k) * isvalid(ϕ.Vz, i, j, k + 1) + v = vx * vy * vz + return v * isvalid(ϕ.center, i, j, k) end """ @@ -115,16 +123,17 @@ Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, nx, ny = size(ϕ.Vx) j_bot = max(j - 1, 1) j0 = min(j, ny) - vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0) + vx = isvalid(ϕ.Vx, i, j0) * isvalid(ϕ.Vx, i, j_bot) nx, ny = size(ϕ.Vy) i_left = max(i - 1, 1) i0 = min(i, nx) - vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0) + vy = isvalid(ϕ.Vy, i0, j) * isvalid(ϕ.Vy, i_left, j) v = vx * vy - return v * (ϕ.vertex[i, j] > 0) + return v * isvalid(ϕ.vertex, i, j) end + """ isvalid_vx(ϕ::JustRelax.RockRatio, inds...) @@ -134,14 +143,20 @@ Check if `ϕ.Vx[inds...]` is a not a nullspace. - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. - `inds`: Cartesian indices to check. """ -Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j) - # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) - # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) - # cv = c * v - # return cv * (ϕ.Vx[i, j] > 0) - return (ϕ.Vx[i, j] > 0) +Base.@propagate_inbounds @inline function isvalid_vx( + ϕ::JustRelax.RockRatio, I::Vararg{Integer,N} +) where {N} + return isvalid(ϕ.Vx, I...) end +# Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, I::Vararg{Integer,N}) where {N} +# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) +# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) +# # cv = c * v +# # return cv * (ϕ.Vx[i, j] > 0) +# return (ϕ.Vx[I...] > 0) +# end + """ isvalid_vy(ϕ::JustRelax.RockRatio, inds...) @@ -151,14 +166,129 @@ Check if `ϕ.Vy[inds...]` is a not a nullspace. - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. - `inds`: Cartesian indices to check. """ -Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) - # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) - # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) - # cv = c * v - # return cv * (ϕ.Vy[i, j] > 0) - return (ϕ.Vy[i, j] > 0) +# Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j) +# # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0) +# # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0) +# # cv = c * v +# # return cv * (ϕ.Vy[i, j] > 0) +# return (ϕ.Vy[i, j] > 0) +# end +Base.@propagate_inbounds @inline function isvalid_vy( + ϕ::JustRelax.RockRatio, I::Vararg{Integer,N} +) where {N} + return isvalid(ϕ.Vy, I...) +end + +""" + isvalid_vz(ϕ::JustRelax.RockRatio, inds...) + +Check if `ϕ.Vz[inds...]` is a not a nullspace. + +# Arguments +- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. +- `inds`: Cartesian indices to check. +""" +Base.@propagate_inbounds @inline function isvalid_vz( + ϕ::JustRelax.RockRatio, I::Vararg{Integer,N} +) where {N} + return isvalid(ϕ.Vz, I...) +end + +Base.@propagate_inbounds @inline function isvalid_velocity(ϕ::JustRelax.RockRatio, i, j) + return isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vy, i, j) +end + +Base.@propagate_inbounds @inline function isvalid_velocity(ϕ::JustRelax.RockRatio, i, j, k) + return isvalid(ϕ.Vx, i, j, k) * isvalid(ϕ.Vy, i, j, k) * isvalid(ϕ.Vz, i, j, k) +end + +Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j, k) + # yz + nx, ny, nz = size(ϕ.yz) + i_left = max(i - 1, 1) + i_right = min(i, nx) + yz = isvalid(ϕ.yz, i_left, j, k) * isvalid(ϕ.yz, i_right, j, k) + + # xz + nx, ny, nz = size(ϕ.xz) + j_front = max(j - 1, 1) + j_back = min(j, ny) + xz = isvalid(ϕ.xz, i, j_front, k) * isvalid(ϕ.xz, i, j_back, k) + + # xy + nx, ny, nz = size(ϕ.xy) + k_top = max(k - 1, 1) + k_bot = min(k, nz) + xy = isvalid(ϕ.xy, i, j, k_top) * isvalid(ϕ.xy, i, j, k_back) + + # V + v = yz * xz * xy + + return v * isvalid(ϕ.vertex, i, j, k) +end + +Base.@propagate_inbounds @inline function isvalid_xz(ϕ::JustRelax.RockRatio, i, j, k) + + # check vertices + v = isvalid(ϕ.vertex, i, j, k) * isvalid(ϕ.vertex, i, j + 1, k) + + # check vz + nx, ny, nz = size(ϕ.vz) + i_left = max(i - 1, 1) + i_right = min(i, nx) + vz = isvalid(ϕ.vz, i_left, j, k) * isvalid(ϕ.vz, i_right, j, k) + + # check vx + nx, ny, nz = size(ϕ.vx) + k_top = max(k - 1, 1) + k_bot = min(k, nz) + vx = isvalid(ϕ.vx, i, j, k_top) * isvalid(ϕ.vx, i, j, k_back) + + return v * vx * vz * isvalid(ϕ.vertex, i, j, k) end +Base.@propagate_inbounds @inline function isvalid_xy(ϕ::JustRelax.RockRatio, i, j, k) + + # check vertices + v = isvalid(ϕ.vertex, i, j, k) * isvalid(ϕ.vertex, i, j, k + 1) + + # check vx + nx, ny, nz = size(ϕ.vx) + j_front = max(j - 1, 1) + j_back = min(j, ny) + vx = isvalid(ϕ.vx, i, j_front, k) * isvalid(ϕ.vx, i, j_back, k) + + # check vy + nx, ny, nz = size(ϕ.vy) + i_left = max(i - 1, 1) + i_right = min(i, nx) + vy = isvalid(ϕ.vy, i_left, j, k) * isvalid(ϕ.vy, i_right, j, k) + + return v * vy * vz * isvalid(ϕ.vertex, i, j, k) +end + +Base.@propagate_inbounds @inline function isvalid_yz(ϕ::JustRelax.RockRatio, i, j, k) + + # check vertices + v = isvalid(ϕ.vertex, i, j, k) * isvalid(ϕ.vertex, i + 1, j, k) + + # check vz + nx, ny, nz = size(ϕ.vz) + j_front = max(j - 1, 1) + j_back = min(j, ny) + vz = isvalid(ϕ.vz, i, j_front, k) * isvalid(ϕ.vz, i, j_back, k) + + # check vy + nx, ny, nz = size(ϕ.vy) + k_top = max(k - 1, 1) + k_bot = min(k, nz) + vy = isvalid(ϕ.vy, i, j, k_top) * isvalid(ϕ.vy, i, j, k_back) + + return v * vy * vz * isvalid(ϕ.vertex, i, j, k) +end + +Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer,N}) where {N} = ϕ[I...] > 0 + ###### # """ @@ -218,5 +348,3 @@ end # cv = c || v # return cv || isvalid(ϕ.Vy, i, j) # end - -Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer,N}) where {N} = ϕ[I...] > 0 From f0843160f43e77a9b5fdacc4a52c0142cc4e189c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 29 Nov 2024 17:35:00 +0100 Subject: [PATCH 26/77] update miniapps --- .../free_surface_stabilization/Crameri2D.jl | 166 ++++++---- .../PlumeFreeSurface_VariationalStokes.jl | 284 ++++++++++++++++++ 2 files changed, 388 insertions(+), 62 deletions(-) create mode 100644 miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/Crameri2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/Crameri2D.jl index 2447b3fb2..a9310b234 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/Crameri2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/Crameri2D.jl @@ -1,39 +1,39 @@ -using CUDA -using JustRelax, JustRelax.JustRelax2D -const backend_JR = CUDABackend -# const backend_JR = CPUBackend +const isCUDA = false +# const isCUDA = true -using JustPIC, JustPIC._2D -const backend = CUDABackend -# const backend = JustPIC.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(CUDA, Float64, 2) +@static if isCUDA + using CUDA +end -# Load script dependencies -using LinearAlgebra, GeoParams, GLMakie +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -# Velocity helper grids for the particle advection -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - @parallel f_x(A, B) +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend end -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) +using ParallelStencil + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) 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 +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend end +# Load script dependencies +using GeoParams, GLMakie + function init_phases!(phases, particles, A) ni = size(phases) @@ -57,8 +57,8 @@ function init_phases!(phases, particles, A) @index phases[ip, i, j] = 2.0 end - if depth < (-cos(x * 2π/700e3) * 7e3 + 100e3) - # if depth < (-cos(x * 2π/2800e3) * 7e3 + 100e3) + # if depth < (-cos(x * 2π/700e3) * 7e3 + 100e3) + if depth < (-cos(x * 2π/2800e3) * 7e3 + 100e3) @index phases[ip, i, j] = 1.0 end @@ -72,7 +72,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ # (Path)/folder where output data and figures are stored -n = 101 +n = 64 nx = n ny = n igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -101,8 +101,8 @@ function main(igg, nx, ny) # Name = "Air", SetMaterialParams(; Phase = 1, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e23),)), + Density = ConstantDensity(; ρ=0e0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e22),)), Gravity = ConstantGravity(; g=10), ), # Name = "Crust", @@ -125,7 +125,7 @@ function main(igg, nx, ny) # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 60, 80, 40 particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi, di, ni + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni ) # velocity grids grid_vx, grid_vy = velocity_grids(xci, xvi, di) @@ -135,32 +135,33 @@ function main(igg, nx, ny) # Elliptical temperature anomaly A = 5e3 # Amplitude of the anomaly - phase_ratios = PhaseRatios(backend, length(rheology), ni) + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) init_phases!(pPhases, particles, A) update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - # ---------------------------------------------------- - + + # rock ratios for variational stokes # RockRatios - air_phase = 1 - ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + air_phase = 1 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.98 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) + thermal = ThermalArrays(backend, ni) # ---------------------------------------------------- # Buoyancy forces & rheology ρg = @zeros(ni...), @zeros(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) compute_ρg!(ρg[2], phase_ratios, rheology, args) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, (1e18, 1e24)) # Boundary conditions flow_bcs = VelocityBoundaryConditions(; @@ -178,21 +179,10 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 dt = 10e3 * (3600 * 24 * 365.25) - dt_max = 50e3 * (3600 * 24 * 365.25) - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - (; η, η_vep) = stokes.viscosity - ni = size(stokes.P) - iterMax = 15e3 - nout = 1e3 - viscosity_cutoff = (-Inf, Inf) - free_surface = false - ητ = @zeros(ni...) while it < 20 - ## variational solver - # Stokes solver ---------------- + # Variational Stokes solver ---------------- solve_VariationalStokes!( stokes, pt_stokes, @@ -203,14 +193,14 @@ function main(igg, nx, ny) ϕ, rheology, args, - dt, + Inf, igg; kwargs = ( iterMax = 50e3, iterMin = 1e3, - viscosity_relaxation = 1e-2, + viscosity_relaxation = 1e-2, nout = 2e3, - viscosity_cutoff = (-Inf, Inf) + viscosity_cutoff = (1e18, 1e24) ) ) dt = compute_dt(stokes, di, dt_max) @@ -225,7 +215,7 @@ function main(igg, nx, ny) inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) @show it += 1 t += dt @@ -256,14 +246,15 @@ function main(igg, nx, ny) xci, data_v, data_c, - velocity_v + velocity_v; + t = t ) # end # if it == 1 || rem(it, 1) == 0 # px, py = particles.coords - # velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) # nt = 5 # fig = Figure(size = (900, 900), title = "t = $t") # ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") @@ -279,6 +270,27 @@ function main(igg, nx, ny) # fig # save(joinpath(figdir, "$(it).png"), fig) + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + + # Make particles plottable + nt = 5 + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + idxv = particles.index.data[:]; + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.V.Vy)) + scatter!(ax, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 2, colormap=:grayC) + arrows!( + ax, + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + color = :red, + ) + save(joinpath(figdir, "$(it).png"), fig) + fig # end end return nothing @@ -287,4 +299,34 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- main(igg, nx, ny) -# heatmap(Array(ρg[2]), colormap = :vikO) \ No newline at end of file + + +# @parallel_indices (I...) function compute_P!( +# P, +# P0, +# RP, +# ∇V, +# η, +# rheology::NTuple{N,MaterialParams}, +# phase_ratio, +# ϕ::JustRelax.RockRatio, +# dt, +# r, +# θ_dτ, +# ::Nothing, +# ) where {N} +# # if isvalid_c(ϕ, I...) +# # K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...])) +# # RP[I...], P[I...] = _compute_P!( +# # P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ +# # ) +# # else +# # RP[I...] = P[I...] = zero(eltype(P)) +# # end +# return nothing +# end + +# @parallel (@idx ni) compute_P_kernel!( +# stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, stokes.viscosityη, +# rheology, phase_ratio.center, ϕ, dt, pt_stokes.r, pt_stokes.θ_dτ, nothing +# ) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl new file mode 100644 index 000000000..0a2fafff8 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl @@ -0,0 +1,284 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +# Load script dependencies +using GeoParams, GLMakie + +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 + +function init_phases!(phases, particles) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index) + r=100e3 + f(x, A, λ) = A * sin(π*x/λ) + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) + @index phases[ip, i, j] = 2.0 + + if 0e0 ≤ depth ≤ 100e3 + @index phases[ip, i, j] = 1.0 + + else + @index phases[ip, i, j] = 2.0 + + if ((x - 250e3)^2 + (depth - 250e3)^2 ≤ r^2) + @index phases[ip, i, j] = 3.0 + end + end + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 100e3 # thickness of sticky air layer + ly = 400e3 + thick_air # domain length in y + lx = 500e3 # 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) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + # Name = "Air", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=1e1), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Mantle", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Plume", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ) + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 30, 40, 15 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + init_phases!(pPhases, particles) + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + # rock ratios for variational stokes + # RockRatios + air_phase = 1 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + free_surface = true + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "FreeSurfacePlume" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + while it < 150 + + # solve!( + # stokes, + # pt_stokes, + # di, + # flow_bcs, + # ρg, + # phase_ratios, + # rheology, + # args, + # dt, + # igg; + # kwargs = (; + # iterMax = 50e3, + # nout = 1e3, + # viscosity_cutoff = (-Inf, Inf), + # free_surface = true + # ) + # ) + + solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = (-Inf, Inf), + ) + ) + dt = compute_dt(stokes, di) / 2 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 10) == 0 + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + + # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) + # arrows!( + # ax, + # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + # color = :red, + # ) + + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + idxv = particles.index.data[:]; + scatter!(ax, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 2) + arrows!( + ax, + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + color = :red, + ) + + fig + save(joinpath(figdir, "$(it).png"), fig) + + end + end + return nothing +end +## END OF MAIN SCRIPT ---------------------------------------------------------------- + +# (Path)/folder where output data and figures are stored +n = 256 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +main(igg, nx, ny) From f5efe36ef2d95e1c08fc3407035d587f228fbcb5 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 2 Dec 2024 11:44:00 +0100 Subject: [PATCH 27/77] fix P bug --- src/stokes/PressureKernels.jl | 2 +- src/variational_stokes/PressureKernels.jl | 54 +++++++++++------------ 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/stokes/PressureKernels.jl b/src/stokes/PressureKernels.jl index 468fa9c7c..14961e8ba 100644 --- a/src/stokes/PressureKernels.jl +++ b/src/stokes/PressureKernels.jl @@ -64,7 +64,7 @@ function compute_P!( ) where {N} ni = size(P) @parallel (@idx ni) compute_P_kernel!( - P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ, ΔTc, melt_fraction + P, P0, RP, ∇V, η, rheology, phase_ratio.center, dt, r, θ_dτ, ΔTc, melt_fraction ) return nothing end diff --git a/src/variational_stokes/PressureKernels.jl b/src/variational_stokes/PressureKernels.jl index 86c49f6cf..c7c90c1dd 100644 --- a/src/variational_stokes/PressureKernels.jl +++ b/src/variational_stokes/PressureKernels.jl @@ -45,30 +45,30 @@ end return nothing end -@parallel_indices (I...) function compute_P_kernel!( - P, - P0, - RP, - ∇V, - η, - rheology::NTuple{N,MaterialParams}, - phase_ratio::C, - ϕ::JustRelax.RockRatio, - dt, - r, - θ_dτ, - ΔTc, - ::Nothing, -) where {N,C<:JustRelax.CellArray} - if isvalid_c(ϕ, I...) - phase_ratio_I = phase_ratio[I...] - K = fn_ratio(get_bulk_modulus, rheology, phase_ratio_I) - α = fn_ratio(get_thermal_expansion, rheology, phase_ratio_I) - RP[I...], P[I...] = _compute_P!( - P[I...], P0[I...], ∇V[I...], ΔTc[I...], α, η[I...], K, dt, r, θ_dτ - ) - else - RP[I...] = P[I...] = zero(eltype(P)) - end - return nothing -end +# @parallel_indices (I...) function compute_P_kernel!( +# P, +# P0, +# RP, +# ∇V, +# η, +# rheology::NTuple{N,MaterialParams}, +# phase_ratio::C, +# ϕ::JustRelax.RockRatio, +# dt, +# r, +# θ_dτ, +# ΔTc, +# ::Nothing, +# ) where {N,C<:JustRelax.CellArray} +# if isvalid_c(ϕ, I...) +# phase_ratio_I = phase_ratio[I...] +# K = fn_ratio(get_bulk_modulus, rheology, phase_ratio_I) +# α = fn_ratio(get_thermal_expansion, rheology, phase_ratio_I) +# RP[I...], P[I...] = _compute_P!( +# P[I...], P0[I...], ∇V[I...], ΔTc[I...], α, η[I...], K, dt, r, θ_dτ +# ) +# else +# RP[I...] = P[I...] = zero(eltype(P)) +# end +# return nothing +# end From 7fecf7d657d63f49872f0b5d9100b06fd2983a9e Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 2 Dec 2024 11:46:32 +0100 Subject: [PATCH 28/77] format --- src/variational_stokes/Stokes3D.jl | 1 - src/variational_stokes/StressKernels.jl | 2 +- src/variational_stokes/mask.jl | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/variational_stokes/Stokes3D.jl b/src/variational_stokes/Stokes3D.jl index 20d28592a..2b55b57a9 100644 --- a/src/variational_stokes/Stokes3D.jl +++ b/src/variational_stokes/Stokes3D.jl @@ -11,7 +11,6 @@ function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) return _solve_VS!(stokes, args...; kwargs...) end - # GeoParams and multiple phases function _solve!( stokes::JustRelax.StokesArrays, diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl index 661a0cf05..9ac0c3a3a 100644 --- a/src/variational_stokes/StressKernels.jl +++ b/src/variational_stokes/StressKernels.jl @@ -304,7 +304,7 @@ end end ## xy - if all(I .≤ size(ε[6])) && isvalid_xy(ϕ, I...) + if all(I .≤ size(ε[6])) && isvalid_xy(ϕ, I...) # interpolate to ith vertex ηv_ij = av_clamped_xy(η, Ic...) EIIv_ij = av_clamped_xy(EII, Ic...) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 01be3cc98..07b56d794 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -133,7 +133,6 @@ Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, return v * isvalid(ϕ.vertex, i, j) end - """ isvalid_vx(ϕ::JustRelax.RockRatio, inds...) From 6510117254204183e1c949ad2f01d12128f0dc09 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 2 Dec 2024 16:02:43 +0100 Subject: [PATCH 29/77] add return to variational stokes solver --- src/variational_stokes/Stokes2D.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 52feeb09d..135be45e0 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -2,8 +2,8 @@ # backend trait function solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs) - solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) - return nothing + out = solve_VariationalStokes!(backend(stokes), stokes, args...; kwargs) + return out end # entry point for extensions From 7a5bd8eb4a172747925379017ddbf4053168be7d Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 2 Dec 2024 16:15:03 +0100 Subject: [PATCH 30/77] add Volcano 2D setup to miniapps; add test; update compats in project --- Project.toml | 4 +- .../stokes2D/Volcano2D}/Caldera2D.jl | 243 +++++++------- .../stokes2D/Volcano2D}/Caldera_rheology.jl | 0 .../stokes2D/Volcano2D}/Caldera_setup.jl | 0 test/Project.toml | 6 +- volcano2D/DK2D.jl => test/test_Volcano2D.jl | 317 +++++------------- volcano2D/DK_rheology.jl | 152 --------- volcano2D/DK_setup.jl | 82 ----- 8 files changed, 217 insertions(+), 587 deletions(-) rename {volcano2D => miniapps/benchmarks/stokes2D/Volcano2D}/Caldera2D.jl (67%) rename {volcano2D => miniapps/benchmarks/stokes2D/Volcano2D}/Caldera_rheology.jl (100%) rename {volcano2D => miniapps/benchmarks/stokes2D/Volcano2D}/Caldera_setup.jl (100%) rename volcano2D/DK2D.jl => test/test_Volcano2D.jl (50%) delete mode 100644 volcano2D/DK_rheology.jl delete mode 100644 volcano2D/DK_setup.jl diff --git a/Project.toml b/Project.toml index cc5b481ed..acc9ae688 100644 --- a/Project.toml +++ b/Project.toml @@ -34,11 +34,11 @@ AMDGPU = "0.8, 0.9, 1" Adapt = "4" CUDA = "5" CellArrays = "0.2" -GeoParams = "0.6.7" +GeoParams = "0.6" HDF5 = "0.17.1" ImplicitGlobalGrid = "0.15" JLD2 = "0.4, 0.5" -JustPIC = "0.5.3" +JustPIC = "0.5" MPI = "0.20" MuladdMacro = "0.2" ParallelStencil = "0.13.6" diff --git a/volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl similarity index 67% rename from volcano2D/Caldera2D.jl rename to miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 4d0ee07ca..862dca68d 100644 --- a/volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, GLMakie, CellArrays +using GeoParams, CairoMakie, CellArrays, Statistics # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -72,7 +72,7 @@ function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) @parallel_indices (i, j) function pure_shear_y!(Vy, εbg, ly) yi = yv[j] - Vy[i + 1, j] = abs(yi) * εbg + Vy[i + 1, j] = abs(yi) * εbg return nothing end @@ -86,7 +86,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) +function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) # Physical domain ------------------------------------ ni = nx, ny # number of cells @@ -148,8 +148,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) - Ttop = thermal.T[2:end-1, end] - Tbot = thermal.T[2:end-1, 1] + Ttop = thermal.T[2:end-1, end] + Tbot = thermal.T[2:end-1, 1] # ---------------------------------------------------- # Buoyancy forces @@ -180,7 +180,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # flow_bcs!(stokes, flow_bcs) # apply boundary conditions # displacement2velocity!(stokes, dt) - εbg = 1e-15 * 1 + εbg = 1e-15 * 1 apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -188,12 +188,15 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # IO ------------------------------------------------- # if it does not exist, make folder where figures are stored - if do_vtk - vtk_dir = joinpath(figdir, "vtk") - take(vtk_dir) + if plotting + checkpoint = joinpath(figdir, "checkpoint") + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) end - take(figdir) - cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) # ---------------------------------------------------- local Vx_v, Vy_v @@ -248,7 +251,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk @views thermal.T[2:end-1, :] .= T_buffer thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) - + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) args = (; T=thermal.Tc, P=stokes.P, dt=Inf) @@ -266,7 +269,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk igg; kwargs = (; iterMax = 100e3, - nout = 2e3, + nout = 2e3, viscosity_cutoff = viscosity_cutoff, ) ) @@ -285,7 +288,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk tensor_invariant!(stokes.ε_pl) dtmax = 2e3 * 3600 * 24 * 365.25 dt = compute_dt(stokes, di, dtmax) - + println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ @@ -327,9 +330,9 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # check if we need to inject particles # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) inject_particles_phase!( - particles, - pPhases, - particle_args_reduced, + particles, + pPhases, + particle_args_reduced, (T_buffer, τxx_v, τyy_v, stokes.τ.xy, stokes.ω.xy), xvi ) @@ -338,122 +341,128 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) - particle2centroid!(stokes.τ_o.xx, pτxx, xci, particles) - particle2centroid!(stokes.τ_o.yy, pτyy, xci, particles) - particle2grid!(stokes.τ_o.xy, pτxy, xvi, particles) + particle2centroid!(stokes.τ.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ.yy, pτyy, xci, particles) + particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + tensor_invariant!(stokes.τ) @show it += 1 t += dt - - # Data I/O and plotting --------------------- - if it == 1 || rem(it, 1) == 0 - # checkpointing(figdir, stokes, thermal.T, η, t) - (; η_vep, η) = stokes.viscosity - if do_vtk - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - data_v = (; - T = Array(T_buffer), - τII = Array(stokes.τ.II), - εII = Array(stokes.ε.II), - εII_pl = Array(stokes.ε_pl.II), - ) - data_c = (; - P = Array(stokes.P), - η = Array(η_vep), - ) - velocity_v = ( - Array(Vx_v), - Array(Vy_v), - ) - save_vtk( - joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), - xvi, - xci, - data_v, - data_c, - velocity_v; - t = 0 - ) + if plotting + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 1) == 0 + if igg.me == 0 && it == 1 + metadata(pwd(), checkpoint, basename(@__FILE__), "CalderaModelSetup.jl", "CalderaRheology.jl") + end + checkpointing_jld2(checkpoint, stokes, thermal, t, dt, igg) + + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(T_buffer), + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + εII_pl = Array(stokes.ε_pl.II), + ) + data_c = (; + P = Array(stokes.P), + η = Array(η_vep), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v; + t = 0 + ) + end + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + # clr = pT.data[:] + idxv = particles.index.data[:]; + # Make Makie figure + ar = 2 + fig = Figure(size = (1200, 900), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") + # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") + ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") + # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII)") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) + # Plot 2nd invariant of strain rate + # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) + # Plot effective viscosity + # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:lipari) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + fig + save(joinpath(figdir, "$(it).png"), fig) + + # ## Plot initial T and P profile + # fig = 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(; size=(1200, 900)) + # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + # scatter!( + # ax1, + # Array(thermal.T[2:(end - 1), :][:]), + # Yv, + # ) + # lines!( + # ax2, + # Array(stokes.P[:]), + # Y, + # ) + # hideydecorations!(ax2) + # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + # fig + # end end - - # Make particles plottable - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:]./1e3 - pyv = ppy.data[:]./1e3 - clr = pPhases.data[:] - # clr = pT.data[:] - idxv = particles.index.data[:]; - # Make Makie figure - ar = 2 - fig = Figure(size = (1200, 900), title = "t = $t") - ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") - ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") - # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") - ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") - # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII)") - ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") - # Plot temperature - h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) - # Plot particles phase - h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) - # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) - # Plot 2nd invariant of strain rate - # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) - h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) - # Plot effective viscosity - # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:lipari) - h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) - hidexdecorations!(ax1) - hidexdecorations!(ax2) - hidexdecorations!(ax3) - Colorbar(fig[1,2], h1) - Colorbar(fig[2,2], h2) - Colorbar(fig[1,4], h3) - Colorbar(fig[2,4], h4) - linkaxes!(ax1, ax2, ax3, ax4) - fig - save(joinpath(figdir, "$(it).png"), fig) - - # ## Plot initial T and P profile - # fig = 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(; size=(1200, 900)) - # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - # scatter!( - # ax1, - # Array(thermal.T[2:(end - 1), :][:]), - # Yv, - # ) - # lines!( - # ax2, - # Array(stokes.P[:]), - # Y, - # ) - # hideydecorations!(ax2) - # save(joinpath(figdir, "thermal_profile_$it.png"), fig) - # fig - # end + # ------------------------------ end - # ------------------------------ - end return nothing end ## END OF MAIN SCRIPT ---------------------------------------------------------------- +const plotting = true do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "Caldera2D" n = 128 * 2 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( - nx+1, ny+1; - sticky_air = 4, - flat = false, + nx+1, ny+1; + sticky_air = 4, + flat = false, chimney = false, chamber_T = 1e3, chamber_depth = 7e0, @@ -470,4 +479,4 @@ else igg end -main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); +main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/volcano2D/Caldera_rheology.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl similarity index 100% rename from volcano2D/Caldera_rheology.jl rename to miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl diff --git a/volcano2D/Caldera_setup.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl similarity index 100% rename from volcano2D/Caldera_setup.jl rename to miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl diff --git a/test/Project.toml b/test/Project.toml index 019ed2ec4..e517a7b3f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" +GeophysicalModelGenerator = "3700c31b-fa53-48a6-808a-ef22d5a84742" JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" ParallelStencil = "94395366-693c-11ea-3b26-d9b7aac5d958" @@ -21,8 +22,9 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" AMDGPU = "0.8, 0.9, 1" CUDA = "5" CellArrays = "0.2" -GeoParams = "0.6.7" -JustPIC = "0.5.3" +GeoParams = "0.6" +GeophysicalModelGenerator = "0.7" +JustPIC = "0.5" MPI = "0.20" ParallelStencil = "0.11, 0.12, 0.13" SpecialFunctions = "2" diff --git a/volcano2D/DK2D.jl b/test/test_Volcano2D.jl similarity index 50% rename from volcano2D/DK2D.jl rename to test/test_Volcano2D.jl index c8f1da6cc..ca0f64724 100644 --- a/volcano2D/DK2D.jl +++ b/test/test_Volcano2D.jl @@ -1,42 +1,41 @@ -# const isCUDA = false -const isCUDA = true - -@static if isCUDA +push!(LOAD_PATH, "..") +@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + using AMDGPU +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" using CUDA end -using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO - -const backend = @static if isCUDA - CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -else - JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -end - +using Test, Suppressor +using GeoParams +using JustRelax, JustRelax.JustRelax2D using ParallelStencil, ParallelStencil.FiniteDifferences2D -@static if isCUDA +const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + @init_parallel_stencil(AMDGPU, Float64, 2) + AMDGPUBackend +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" @init_parallel_stencil(CUDA, Float64, 2) + CUDABackend else @init_parallel_stencil(Threads, Float64, 2) + CPUBackend end using JustPIC, JustPIC._2D -# Threads is the default backend, -# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, -# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend_JP = @static if isCUDA - CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +const backend_JP = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + JustPIC.AMDGPUBackend +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" + CUDABackend else - JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + JustPIC.CPUBackend end - # Load script dependencies -using GeoParams, GLMakie, CellArrays +using GeoParams, CellArrays, Statistics # Load file with all the rheology configurations -include("DK_setup.jl") -include("DK_rheology.jl") +include("../miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl") +include("../miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- @@ -72,7 +71,7 @@ function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) @parallel_indices (i, j) function pure_shear_y!(Vy, εbg, ly) yi = yv[j] - Vy[i + 1, j] = abs(yi) * εbg + Vy[i + 1, j] = abs(yi) * εbg return nothing end @@ -86,7 +85,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) +function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) # Physical domain ------------------------------------ ni = nx, ny # number of cells @@ -97,8 +96,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Physical properties using GeoParams ---------------- rheology = init_rheologies() - rheology_linear = init_rheologies(; linear=true) - dt = 1e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + dt = 5e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter # dt = Inf # diffusive CFL timestep limiter # ---------------------------------------------------- @@ -138,13 +136,10 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 1e0 , r=0.7, CFL = 0.98 / √2.1) # Re=3π, r=0.7 - # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.98 / √2.1) # Re=3π, r=0.7 # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - Ttop = T_GMG[1, end] - Tbot = T_GMG[1, 1] thermal = ThermalArrays(backend, ni) @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) thermal_bc = TemperatureBoundaryConditions(; @@ -152,12 +147,14 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) + Ttop = thermal.T[2:end-1, end] + Tbot = thermal.T[2:end-1, 1] # ---------------------------------------------------- # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(2)) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) # Rheology args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) @@ -170,30 +167,27 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk ) # Boundary conditions + # flow_bcs = DisplacementBoundaryConditions(; flow_bcs = VelocityBoundaryConditions(; free_slip = (left = true , right = true , top = true , bot = true), free_surface = false, ) - εbg = 1e-15 * 0 + + # U = 0.02 + # stokes.U.Ux .= PTArray(backend)([(x - li[1] * 0.5) * U / dt for x in xvi[1], _ in 1:ny+2]) + # stokes.U.Uy .= PTArray(backend)([-y * U / dt for _ in 1:nx+2, y in xvi[2]]) + # flow_bcs!(stokes, flow_bcs) # apply boundary conditions + # displacement2velocity!(stokes, dt) + + εbg = 1e-15 * 1 apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(@velocity(stokes)...) - - # IO ------------------------------------------------- - # if it does not exist, make folder where figures are stored - if do_vtk - vtk_dir = joinpath(figdir, "vtk") - take(vtk_dir) - end - take(figdir) - cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) # ---------------------------------------------------- - local Vx_v, Vy_v - if do_vtk - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - end + local iters + T_buffer = @zeros(ni.+1) Told_buffer = similar(T_buffer) @@ -202,31 +196,6 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk copyinn_x!(dst, src) end grid2particle!(pT, xvi, T_buffer, particles) - id = @. pPhases.data == 3 - @views pT.data[id] .= 750+273 - - ## Plot initial T and P profile - fig = 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:]), - Yv, - ) - lines!( - ax2, - # Array(stokes.P[:]), - Array(log10.(stokes.viscosity.η[:])), - Y, - ) - hideydecorations!(ax2) - # save(joinpath(figdir, "initial_profile.png"), fig) - fig - end τxx_v = @zeros(ni.+1...) τyy_v = @zeros(ni.+1...) @@ -234,29 +203,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Time loop t, it = 0.0, 0 thermal.Told .= thermal.T - - println("Starting warm up linear solver...") - args = (; T=thermal.Tc, P=stokes.P, dt=Inf) - t_stokes = @elapsed solve_VariationalStokes!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - phase_ratios, - ϕ, - rheology_linear, - args, - dt, - igg; - kwargs = (; - iterMax = 100e3, - nout = 2e3, - viscosity_cutoff = viscosity_cutoff, - ) - ) - println("... finished warm up linear solver") - while it < 300 #000 # run only for 5 Myrs + + while it < 5 #000 # run only for 5 Myrs # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) @@ -265,15 +213,11 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk @views thermal.T[2:end-1, :] .= T_buffer thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) - - args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) - # args = (; T=thermal.Tc, P=stokes.P, dt=Inf) - particle2centroid!(stokes.τ.xx, pτxx, xci, particles) - particle2centroid!(stokes.τ.yy, pτyy, xci, particles) - particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + args = (; T=thermal.Tc, P=stokes.P, dt=Inf) - t_stokes = @elapsed solve_VariationalStokes!( + iters = solve_VariationalStokes!( stokes, pt_stokes, di, @@ -287,7 +231,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk igg; kwargs = (; iterMax = 100e3, - nout = 2e3, + nout = 2e3, viscosity_cutoff = viscosity_cutoff, ) ) @@ -297,15 +241,14 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk centroid2particle!(pτxx , xci, stokes.τ.xx, particles) centroid2particle!(pτyy , xci, stokes.τ.yy, particles) grid2particle!(pτxy, xvi, stokes.τ.xy, particles) + grid2particle!(pωxy, xvi, stokes.ω.xy, particles) rotate_stress_particles!(pτ, pω, particles, dt) - println("Stokes solver time ") - println(" Total time: $t_stokes s") tensor_invariant!(stokes.ε) tensor_invariant!(stokes.ε_pl) - dt_max = 1e3 * 3600 * 24 * 365.25 - dt = compute_dt(stokes, di, dt_max) * 0.8 - + dtmax = 2e3 * 3600 * 24 * 365.25 + dt = compute_dt(stokes, di, dtmax) + println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ @@ -339,6 +282,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # ------------------------------ # Advection -------------------- + copyinn_x!(T_buffer, thermal.T) # advect particles in space advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) # advect particles in memory @@ -346,9 +290,9 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # check if we need to inject particles # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) inject_particles_phase!( - particles, - pPhases, - particle_args_reduced, + particles, + pPhases, + particle_args_reduced, (T_buffer, τxx_v, τyy_v, stokes.τ.xy, stokes.ω.xy), xvi ) @@ -357,133 +301,42 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + particle2centroid!(stokes.τ.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ.yy, pτyy, xci, particles) + particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + tensor_invariant!(stokes.τ) + @show it += 1 t += dt - - if it > 1 && rem(it, 50) == 0 - id = @. pPhases.data == 3 - @views pT.data[id] .= 750+273 - end - - # Data I/O and plotting --------------------- - if it == 1 || rem(it, 2) == 0 - # checkpointing(figdir, stokes, thermal.T, η, t) - (; η_vep, η) = stokes.viscosity - if do_vtk - velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - data_v = (; - T = Array(T_buffer), - τII = Array(stokes.τ.II), - εII = Array(stokes.ε.II), - ) - data_c = (; - P = Array(stokes.P), - η = Array(η_vep), - ) - velocity_v = ( - Array(Vx_v), - Array(Vy_v), - ) - save_vtk( - joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), - xvi, - xci, - data_v, - data_c, - velocity_v - ) - end - - # Make particles plottable - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:]./1e3 - pyv = ppy.data[:]./1e3 - clr = pPhases.data[:] - # clr = pT.data[:] - idxv = particles.index.data[:]; - # Make Makie figure - ar = 1 - fig = Figure(size = (1200, 900), title = "t = $t") - ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") - ax2 = Axis(fig[2,1], aspect = ar, title = "EII plastic") - # ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") - # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") - ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") - ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII_plastic)") - # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") - # Plot temperature - h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) - # Plot particles phase - # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) - # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) - h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.EII_pl) , colormap=:batlow) - # Plot 2nd invariant of strain rate - # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) - h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) - # Plot effective viscosity - h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:lipari) - # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) - hidexdecorations!(ax1) - hidexdecorations!(ax2) - hidexdecorations!(ax3) - Colorbar(fig[1,2], h1) - Colorbar(fig[2,2], h2) - Colorbar(fig[1,4], h3) - Colorbar(fig[2,4], h4) - linkaxes!(ax1, ax2, ax3, ax4) - fig - save(joinpath(figdir, "$(it).png"), fig) - - # ## Plot initial T and P profile - # fig = 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(; size=(1200, 900)) - # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - # scatter!( - # ax1, - # Array(thermal.T[2:(end - 1), :][:]), - # Yv, - # ) - # lines!( - # ax2, - # Array(stokes.P[:]), - # Y, - # ) - # hideydecorations!(ax2) - # save(joinpath(figdir, "thermal_profile_$it.png"), fig) - # fig - # end - end - # ------------------------------ end - return nothing + return iters end ## END OF MAIN SCRIPT ---------------------------------------------------------------- -do_vtk = true # set to true to generate VTK files for ParaView -figdir = "DK_2D_pulses" -n = 128 -nx, ny = n, n -li, origin, phases_GMG, T_GMG = setup2D( - nx+1, ny+1; - sticky_air = 1, - flat = true, - chimney = false, - chamber_T = 750, - chamber_depth = 5e0, - chamber_radius = 1.5, - aspect_x = 1, -) - -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end +@testset "Volcano 2D" begin + @suppress begin + n = 64 + nx, ny = n, n >>> 1 + li, origin, phases_GMG, T_GMG = setup2D( + nx+1, ny+1; + sticky_air = 4, + flat = false, + chimney = false, + chamber_T = 1e3, + chamber_depth = 7e0, + chamber_radius = 0.5, + aspect_x = 6, + ) + + igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) + else + igg + end -main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); + iters = main(li, origin, phases_GMG, T_GMG, igg; nx = nx, ny = ny); + @test passed = iters.err_evo1[end] < 1e-4 + end +end diff --git a/volcano2D/DK_rheology.jl b/volcano2D/DK_rheology.jl deleted file mode 100644 index 138c52d8b..000000000 --- a/volcano2D/DK_rheology.jl +++ /dev/null @@ -1,152 +0,0 @@ -using GeoParams.Dislocation -using GeoParams.Diffusion - -function init_rheology_nonNewtonian() - #dislocation laws - disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) - # diffusion laws - diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) - - el = ConstantElasticity(; G = 40e9) - - lithosphere_rheology = CompositeRheology( (el, disl_wet_olivine, diff_wet_olivine)) - init_rheologies(lithosphere_rheology) -end - -function init_rheology_nonNewtonian_plastic() - #dislocation laws - disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) - # diffusion laws - diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) - # plasticity - ϕ_wet_olivine = asind(0.1) - C_wet_olivine = 1e6 - η_reg = 1e16 - el = ConstantElasticity(; G = 40e9, ν = 0.45) - lithosphere_rheology = CompositeRheology( - ( - el, - disl_wet_olivine, - diff_wet_olivine, - DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity - ) - ) - init_rheologies(lithosphere_rheology) -end - -function init_rheologies(; linear=false) - - η_reg = 1e18 - C = linear ? Inf : 15e6 - ϕ = 30 - soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 2) # nonlinear softening law - # soft_C = NonLinearSoftening() # nonlinear softening law - pl = DruckerPrager_regularised(; C=C, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C=soft_C) - β = 1e-11 - el = ConstantElasticity(; G = 60e9, Kb = 1 / β) - Cp = 1050.0 - - disl_creep = DislocationCreep(; - A = 1.67e-24, - n = 3.5, - V = 0, - E = 1.87e5, - R = 8.314, - ) - g = 0 - # Define rheolgy struct - rheology = ( - # Name = "Upper crust", - SetMaterialParams(; - Phase = 1, - Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), - HeatCapacity = ConstantHeatCapacity(; Cp = Cp), - Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (disl_creep, el, pl)), - RadioactiveHeat = ConstantRadioactiveHeat(1e-6), - Gravity = ConstantGravity(; g=g), - ), - # Name = "Lower crust", - SetMaterialParams(; - Phase = 2, - Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), - HeatCapacity = ConstantHeatCapacity(; Cp = Cp), - Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (disl_creep, el, pl)), - RadioactiveHeat = ConstantRadioactiveHeat(1e-6), - Gravity = ConstantGravity(; g=g), - ), - - # Name = "magma chamber", - SetMaterialParams(; - Phase = 3, - Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), - HeatCapacity = ConstantHeatCapacity(; Cp = Cp), - Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (disl_creep, el, pl)), - RadioactiveHeat = ConstantRadioactiveHeat(1e-6), - Gravity = ConstantGravity(; g=g), - ), - # Name = "magma chamber - hot anomaly", - SetMaterialParams(; - Phase = 4, - Density = PT_Density(; ρ0=2.650e3, T0=273.15, β=β), - HeatCapacity = ConstantHeatCapacity(; Cp = Cp), - Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (disl_creep, el, pl)), - RadioactiveHeat = ConstantRadioactiveHeat(1e-6), - Gravity = ConstantGravity(; g=g), - ), - # Name = "StickyAir", - SetMaterialParams(; - Phase = 5, - Density = ConstantDensity(; ρ=0e0), - HeatCapacity = ConstantHeatCapacity(; Cp = Cp), - Conductivity = ConstantConductivity(; k = 2.5), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), - Gravity = ConstantGravity(; g=g), - ), - ) -end - -function init_phases!(phases, phase_grid, particles, xvi) - ni = size(phases) - @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) -end - -@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} - - ni = size(phases) - - for ip in cellaxes(phases) - # quick escape - @index(index[ip, I...]) == 0 && continue - - pᵢ = ntuple(Val(N)) do i - @index pcoords[i][ip, I...] - end - - d = Inf # distance to the nearest particle - particle_phase = -1 - for offi in 0:1, offj in 0:1 - ii = I[1] + offi - jj = I[2] + offj - - !(ii ≤ ni[1]) && continue - !(jj ≤ ni[2]) && continue - - xvᵢ = ( - xvi[1][ii], - xvi[2][jj], - ) - d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) - if d_ijk < d - d = d_ijk - particle_phase = phase_grid[ii, jj] - end - end - @index phases[ip, I...] = Float64(particle_phase) - end - - return nothing -end diff --git a/volcano2D/DK_setup.jl b/volcano2D/DK_setup.jl deleted file mode 100644 index 03e5c1c84..000000000 --- a/volcano2D/DK_setup.jl +++ /dev/null @@ -1,82 +0,0 @@ -## Model Setup -using GeophysicalModelGenerator - -function setup2D( - nx, nz; - sticky_air = 2.5, - flat = true, - chimney = false, - chamber_T = 1e3, - chamber_depth = 5e0, - chamber_radius = 2e0, - aspect_x = 1.5, -) - - Lx = Ly = 10.5 - x = range(0.0, Lx, nx); - y = range(0.0, Ly, 2); - z = range(-Lx, sticky_air, nz); - Grid = CartData(xyz_grid(x,y,z)); - - # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) - Phases = fill(5, nx, 2, nz); - - # In many (geodynamic) models, one also has to define the temperature, so lets define it as well - Temp = fill(0.0, nx, 2, nz); - - add_box!(Phases, Temp, Grid; - xlim =(minimum(Grid.x.val), maximum(Grid.x.val)), - ylim =(minimum(Grid.y.val), maximum(Grid.y.val)), - zlim =(minimum(Grid.z.val), 0.0), - phase = LithosphericPhases(Layers=[chamber_depth], Phases=[1, 2]), - T = LinearTemp(Ttop=0, Tbot=450) - ) - - if !flat - add_volcano!(Phases, Temp, Grid; - volcanic_phase = 1, - center = (mean(Grid.x.val), 0.0), - height = 3, - radius = 5, - crater = 0.5, - base = 0.0, - background = nothing, - T = LinearTemp(Ttop=0, Tbot=450) - ) - end - - add_ellipsoid!(Phases, Temp, Grid; - cen = (mean(Grid.x.val), 0, -chamber_depth), - axes = (chamber_radius * aspect_x, 2.5, chamber_radius), - phase = ConstantPhase(3), - T = ConstantTemp(T=chamber_T) - ) - - # add_sphere!(Phases, Temp, Grid; - # cen = (mean(Grid.x.val), 0, -chamber_depth), - # radius = 0.5, - # phase = ConstantPhase(4), - # T = ConstantTemp(T=chamber_T) - # ) - - if chimney - add_cylinder!(Phases, Temp, Grid; - base = (mean(Grid.x.val), 0, -chamber_depth), - cap = (mean(Grid.x.val), 0, 0e0), - radius = 0.05, - phase = ConstantPhase(2), - # T = LinearTemp(Ttop=20, Tbot=1000), - # T = ConstantTemp(T=800), - T = ConstantTemp(T=chamber_T), - ) - end - - Grid = addfield(Grid,(; Phases, Temp)) - li = (abs(last(x)-first(x)), abs(last(z)-first(z))) .* 1e3 - origin = (x[1], z[1]) .* 1e3 - - ph = Phases[:,1,:] - T = Temp[:,1,:] .+ 273 - # write_paraview(Grid, "Volcano2D") - return li, origin, ph, T, Grid -end From effaf5bd8bcba950c1445c05c9bc00a5f614ad06 Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 2 Dec 2024 16:54:36 +0100 Subject: [PATCH 31/77] typos --- miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 862dca68d..32547c548 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -195,7 +195,6 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", take(vtk_dir) end take(figdir) - cp(@__FILE__, joinpath(figdir, "script.jl"); force = true) end # ---------------------------------------------------- @@ -352,7 +351,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Data I/O and plotting --------------------- if it == 1 || rem(it, 1) == 0 if igg.me == 0 && it == 1 - metadata(pwd(), checkpoint, basename(@__FILE__), "CalderaModelSetup.jl", "CalderaRheology.jl") + metadata(pwd(), checkpoint, basename(@__FILE__), "Caldera_setup.jl", "Caldera_rheology.jl") end checkpointing_jld2(checkpoint, stokes, thermal, t, dt, igg) @@ -457,7 +456,7 @@ const plotting = true do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "Caldera2D" -n = 128 * 2 +n = 128 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; From 2eecb570c36f0c7e5f09b1aa2a43502597ea8d47 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 2 Dec 2024 19:28:36 +0100 Subject: [PATCH 32/77] up 3D variational stuff --- src/JustRelax_CPU.jl | 6 ++++ src/common.jl | 3 -- src/ext/AMDGPU/2D.jl | 10 ------ src/ext/AMDGPU/3D.jl | 10 ------ src/ext/CUDA/2D.jl | 7 ---- src/ext/CUDA/3D.jl | 7 ---- src/variational_stokes/Stokes3D.jl | 7 ++-- src/variational_stokes/mask.jl | 52 +++++++++++++++++++++++++----- test/test_Volcano2D.jl | 4 +-- 9 files changed, 56 insertions(+), 50 deletions(-) diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index b802f7a5b..67ada07a9 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -31,6 +31,9 @@ include("common.jl") include("stokes/Stokes2D.jl") export solve! +include("variational_stokes/Stokes2D.jl") +export solve_VariationalStokes! + end module JustRelax3D @@ -65,4 +68,7 @@ include("common.jl") include("stokes/Stokes3D.jl") export solve! +include("variational_stokes/Stokes3D.jl") +export solve_VariationalStokes! + end diff --git a/src/common.jl b/src/common.jl index 164e09682..d1fea9fae 100644 --- a/src/common.jl +++ b/src/common.jl @@ -97,9 +97,6 @@ include("variational_stokes/StressKernels.jl") include("variational_stokes/VelocityKernels.jl") -include("variational_stokes/Stokes2D.jl") -export solve_VariationalStokes! - include("stress_rotation/constructors.jl") export StressParticles diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index e03d3ef59..6df657537 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -420,16 +420,6 @@ function JR2D.update_rock_ratio!( return nothing end -function JR2D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N}, - phase_ratios, - ratio_vel::NTuple{N}, - air_phase, -) where {T,nD,N,D} - update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) - return nothing -end - function JR2D.stress2grid!( stokes, τ_particles::JustRelax.StressParticles{JustPIC.AMDGPUBackend}, diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index aa5b374f1..a832d4e2c 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -442,16 +442,6 @@ function JR3D.update_rock_ratio!( return nothing end -function JR3D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N}, - phase_ratios, - ratio_vel::NTuple{N}, - air_phase, -) where {T,nD,N,D} - update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) - return nothing -end - function JR3D.stress2grid!( stokes, τ_particles::JustRelax.StressParticles{JustPIC.AMDGPUBackend}, diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index eb258f372..b5845dbf3 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -410,13 +410,6 @@ function JR2D.update_rock_ratio!( return nothing end -function JR2D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase -) where {T,nD,N,D} - update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) - return nothing -end - function JR2D.stress2grid!( stokes, τ_particles::JustRelax.StressParticles{CUDABackend}, xvi, xci, particles ) diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 65a38a342..e055082b3 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -435,13 +435,6 @@ function JR3D.update_rock_ratio!( return nothing end -function JR3D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase -) where {T,nD,N,D} - update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase) - return nothing -end - function JR3D.stress2grid!( stokes, τ_particles::JustRelax.StressParticles{CUDABackend}, xvi, xci, particles ) diff --git a/src/variational_stokes/Stokes3D.jl b/src/variational_stokes/Stokes3D.jl index 2b55b57a9..4d6f98dd1 100644 --- a/src/variational_stokes/Stokes3D.jl +++ b/src/variational_stokes/Stokes3D.jl @@ -12,7 +12,7 @@ function solve_VariationalStokes!(::CPUBackendTrait, stokes, args...; kwargs) end # GeoParams and multiple phases -function _solve!( +function _solve_VS!( stokes::JustRelax.StokesArrays, pt_stokes, di::NTuple{3,T}, @@ -24,6 +24,7 @@ function _solve!( args, dt, igg::IGG; + air_phase::Integer=0, iterMax=10e3, nout=500, b_width=(4, 4, 4), @@ -66,7 +67,7 @@ function _solve!( # compute buoyancy forces and viscosity compute_ρg!(ρg, phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, viscosity_cutoff) # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) @@ -77,7 +78,7 @@ function _solve!( compute_maxloc!(ητ, η) update_halo!(ητ) - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di...) compute_P!( θ, stokes.P0, diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 07b56d794..f5e5c523e 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -46,23 +46,59 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. - `air_phase`: The phase representing air. """ function update_rock_ratio!( - ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase -) where {N} + ϕ::JustRelax.RockRatio{T, 2}, phase_ratios, air_phase +) where {T} nvi = size_v(ϕ) @parallel (@idx nvi) update_rock_ratio_cv!( ϕ, phase_ratios.center, phase_ratios.vertex, air_phase ) - # @parallel (@idx nvi) update_rock_ratio_vel!(ϕ) - @parallel (@idx size(ϕ.Vx)) update_rock_ratio_vel!(ϕ.Vx, ratio_vel[1], air_phase) - @parallel (@idx size(ϕ.Vy)) update_rock_ratio_vel!(ϕ.Vy, ratio_vel[2], air_phase) - if N === 3 - @parallel (@idx size(ϕ.Vz)) update_rock_ratio_vel!(ϕ.Vz, ratio_vel[3], air_phase) + dst = ϕ.Vx, ϕ.Vy + src = phase_ratios.Vx, phase_ratios.Vy + + for (dstᵢ, srcᵢ) in zip(dst, src) + @parallel (@idx size(dstᵢ)) _update_rock_ratio!(dstᵢ, srcᵢ, air_phase) end return nothing end +function update_rock_ratio!( + ϕ::JustRelax.RockRatio{T, 3}, phase_ratios, air_phase +) where {T} + nvi = size_v(ϕ) + @parallel (@idx nvi) update_rock_ratio_cv!( + ϕ, phase_ratios.center, phase_ratios.vertex, air_phase + ) + + dst = ϕ.Vx, ϕ.Vy, ϕ.Vz, ϕ.xy, ϕ.yz, ϕ.xz + src = phase_ratios.Vx, phase_ratios.Vy, phase_ratios.Vz, phase_ratios.xy, phase_ratios.yz, phase_ratios.xz + + for (dstᵢ, srcᵢ) in zip(dst, src) + @parallel (@idx size(dstᵢ)) _update_rock_ratio!(dstᵢ, srcᵢ, air_phase) + end + + return nothing +end + +# function update_rock_ratio!( +# ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase +# ) where {N} +# nvi = size_v(ϕ) +# @parallel (@idx nvi) update_rock_ratio_cv!( +# ϕ, phase_ratios.center, phase_ratios.vertex, air_phase +# ) +# # @parallel (@idx nvi) _update_rock_ratio!(ϕ) + +# @parallel (@idx size(ϕ.Vx)) _update_rock_ratio!(ϕ.Vx, ratio_vel[1], air_phase) +# @parallel (@idx size(ϕ.Vy)) _update_rock_ratio!(ϕ.Vy, ratio_vel[2], air_phase) +# if N === 3 +# @parallel (@idx size(ϕ.Vz)) _update_rock_ratio!(ϕ.Vz, ratio_vel[3], air_phase) +# end + +# return nothing +# end + @inline compute_rock_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} ) where {N} = 1 - @index phase_ratio[air_phase, I...] @@ -81,7 +117,7 @@ end return nothing end -@parallel_indices (I...) function update_rock_ratio_vel!(ϕ, ratio, air_phase) +@parallel_indices (I...) function _update_rock_ratio!(ϕ, ratio, air_phase) ϕ[I...] = Float64(Float16(compute_rock_ratio(ratio, air_phase, I...))) return nothing end diff --git a/test/test_Volcano2D.jl b/test/test_Volcano2D.jl index ca0f64724..f5654b1ce 100644 --- a/test/test_Volcano2D.jl +++ b/test/test_Volcano2D.jl @@ -130,7 +130,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # RockRatios air_phase = 5 ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -299,7 +299,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) particle2centroid!(stokes.τ.xx, pτxx, xci, particles) particle2centroid!(stokes.τ.yy, pτyy, xci, particles) From a763d8494e039a018d4fba6d5ce8d38737f0b96c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 3 Dec 2024 09:39:35 +0100 Subject: [PATCH 33/77] fix ext --- src/ext/AMDGPU/2D.jl | 4 ++-- src/ext/AMDGPU/3D.jl | 4 ++-- src/ext/CUDA/2D.jl | 4 ++-- src/ext/CUDA/3D.jl | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 6df657537..319a8120c 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -414,8 +414,8 @@ end # rock ratios function JR2D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase -) where {T,nD,N} + ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},2}, phase_ratios, air_phase +) where {T,nD,D} update_rock_ratio!(ϕ, phase_ratios, air_phase) return nothing end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index a832d4e2c..8a1b072ea 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -436,8 +436,8 @@ end # rock ratios function JR3D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase -) where {T,nD,N} + ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},3}, phase_ratios, air_phase +) where {T,nD,D} update_rock_ratio!(ϕ, phase_ratios, air_phase) return nothing end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index b5845dbf3..0c16af1a2 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -404,8 +404,8 @@ end # rock ratios function JR2D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase -) where {T,nD,N,D} + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},2}, phase_ratios, air_phase +) where {T,nD,D} update_rock_ratio!(ϕ, phase_ratios, air_phase) return nothing end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index e055082b3..0a4bfbde4 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -429,8 +429,8 @@ end # rock ratios function JR3D.update_rock_ratio!( - ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase -) where {T,nD,N,D} + ϕ::JustRelax.RockRatio{CuArray{T,nD,D},3}, phase_ratios, air_phase +) where {T,nD,D} update_rock_ratio!(ϕ, phase_ratios, air_phase) return nothing end From 29767b0207f929063e6392c6267c65daaeb3847c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 3 Dec 2024 10:05:30 +0100 Subject: [PATCH 34/77] fix extensions --- src/ext/AMDGPU/2D.jl | 1 + src/ext/CUDA/2D.jl | 1 + 2 files changed, 2 insertions(+) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 319a8120c..fe18347b8 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -39,6 +39,7 @@ __init__() = @init_parallel_stencil(AMDGPU, Float64, 2) include("../../common.jl") include("../../stokes/Stokes2D.jl") +include("../../variational_stokes/Stokes2D.jl") # Types function JR2D.StokesArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N} diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 0c16af1a2..fa724a72c 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -32,6 +32,7 @@ __init__() = @init_parallel_stencil(CUDA, Float64, 2) include("../../common.jl") include("../../stokes/Stokes2D.jl") +include("../../variational_stokes/Stokes2D.jl") # Types From 75153e4f9a52b03e5ed2e78d7c31a5ff2357cf23 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 3 Dec 2024 12:10:03 +0100 Subject: [PATCH 35/77] phase correction given marker chain --- src/JustRelax_CPU.jl | 3 + src/phases/topography_correction.jl | 106 ++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 src/phases/topography_correction.jl diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 67ada07a9..89f5056e3 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -34,6 +34,9 @@ export solve! include("variational_stokes/Stokes2D.jl") export solve_VariationalStokes! +include("phases/topography_correction.jl") +export update_phases_given_markerchain! + end module JustRelax3D diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl new file mode 100644 index 000000000..fdfc10b26 --- /dev/null +++ b/src/phases/topography_correction.jl @@ -0,0 +1,106 @@ +import JustPIC._2D: is_above_surface, cell_index, interp1D_inner, interp1D_extremas, distance +using StaticArrays + +function update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) + (; coords, index) = particles; + dy = di[2] + @parallel (1:size(index,1)) _update_phases_given_markerchain!(phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase) +end + +@parallel_indices (icell) function _update_phases_given_markerchain!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase) + _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords,cell_vertices, origin, dy, air_phase, icell) + return nothing +end + +function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell) + chain_yi = @cell chain_coords[2][icell] + cell_range = find_minmax_cell_indices(chain_yi, origin[2], dy) + T = eltype(eltype(phase)) + + # iterate over cells with marker chain on them + for j in cell_range + # iterate over particles j-th cell + for ip in cellaxes(index) + (@index index[ip, icell, j]) || continue + xq = @index coords[1][ip, icell, j] + yq = @index coords[2][ip, icell, j] + phaseq = @index phase[ip, icell, j] + + # check if particle is above the marker chain + above = is_above_surface2(xq, yq, chain_coords, cell_vertices) + # if the particle is above the surface and the phase is not air, set the phase to air + if above && phaseq != air_phase + @index phase[ip, icell, j] = T(air_phase) + end + # if the particle is above the surface and the phase is air, set the phase to the closes rock phase + if !above && phaseq == air_phase + @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) + end + end + end + + return nothing +end + +## Utils + +function extrema_CA(x::AbstractArray) + max_val = x[1] + min_val = x[1] + for i in 2:length(x) + xᵢ = x[i] + isnan(xᵢ) && continue + if xᵢ > max_val + max_val = xᵢ + end + if xᵢ < min_val + min_val = xᵢ + end + end + return min_val, max_val +end + +function find_minmax_cell_indices(chain_yi, origin_y, dy) + ymin, ymax = extrema_CA(chain_yi) + min_cell_j = Int((ymin - origin_y) ÷ dy) + 1 + max_cell_j = Int((ymax - origin_y) ÷ dy) + 1 + return min_cell_j:max_cell_j +end + +@inline function is_above_surface2(xq, yq, coords, cell_vertices) + I = cell_index(xq, cell_vertices) + x_cell, y_cell = coords[1][I], coords[2][I] + ychain = if 1 < I[1] < length(cell_vertices) - 1 + interp1D_inner(xq, x_cell, y_cell, coords, I) + else + interp1D_extremas(xq, x_cell, y_cell) + end + return yq > ychain +end + +# find closest phase different than the given `skip_phase` +function closest_phase(coords, pn, index, current_particle, phases, skip_phase, I::Vararg{Int, N}) where N + new_phase = @index phases[current_particle, I...] + dist_min = Inf + px, py = coords + + for ip in cellaxes(index) + # early escape conditions + (ip == current_particle) && continue # current particle + (@index index[ip, I...]) || continue + # get the phase of the particle and skip if it is the same as the `skip_phase` + phaseᵢ = @index phases[ip, I...] + phaseᵢ == skip_phase && continue + + # distance from new point to the existing particle + pxi = @index(px[ip, I...]), @index(py[ip, I...]) + d = distance(pxi, pn) + # update the closest phase + if d < dist_min + new_phase = phaseᵢ + dist_min = d + end + end + + return new_phase +end From e2ef92e12c7d4439900d1efb2973d274f8c35eef Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 3 Dec 2024 12:10:12 +0100 Subject: [PATCH 36/77] up miniapp --- .../PlumeFreeSurface_VariationalStokes.jl | 54 +++++++++++-------- 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl index 0a2fafff8..e65c80c70 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl @@ -1,5 +1,5 @@ -# const isCUDA = false -const isCUDA = true +const isCUDA = false +# const isCUDA = true @static if isCUDA using CUDA @@ -123,14 +123,14 @@ function main(igg, nx, ny) # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 30, 40, 15 - particles = init_particles( + particles = init_particles( backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni ) # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) + grid_vxi = velocity_grids(xci, xvi, di) # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) # Elliptical temperature anomaly init_phases!(pPhases, particles) @@ -141,7 +141,15 @@ function main(igg, nx, ny) # RockRatios air_phase = 1 ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + # ---------------------------------------------------- + + # Initialize marker chain------------------------------- + nxcell, max_xcell, min_xcell = 24, 40, 12 + nxcell, min_xcell, max_xcell = 12, 6, 24 + + initial_elevation = -100e3 + chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -220,32 +228,27 @@ function main(igg, nx, ny) # Advection -------------------- # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) # advect particles in memory move_particles!(particles, xvi, particle_args) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + # advect marker chain + advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + # ------------------------------ @show it += 1 t += dt - if it == 1 || rem(it, 10) == 0 + if it == 1 || rem(it, 5) == 0 velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) nt = 5 fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") - - # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) - # arrows!( - # ax, - # xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - # lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - # color = :red, - # ) - # Make particles plottable p = particles.coords @@ -254,13 +257,18 @@ function main(igg, nx, ny) pyv = ppy.data[:]./1e3 clr = pPhases.data[:] idxv = particles.index.data[:]; - scatter!(ax, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 2) + + chain_x = chain.coords[1].data[:]./1e3 + chain_y = chain.coords[2].data[:]./1e3 + + scatter!(ax, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 5) arrows!( ax, xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :red, + color = :gray, ) + scatter!(ax, Array(chain_x), Array(chain_y), color=:red, markersize = 5) fig save(joinpath(figdir, "$(it).png"), fig) @@ -272,7 +280,7 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- # (Path)/folder where output data and figures are stored -n = 256 +n = 128 nx = n ny = n igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -281,4 +289,4 @@ else igg end -main(igg, nx, ny) +main(igg, nx, ny) \ No newline at end of file From 7d53e2aa6a7ca72e52de4044314fba7e89896dc5 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 4 Dec 2024 15:46:18 +0100 Subject: [PATCH 37/77] improve `update_phases_given_markerchain!` --- src/phases/topography_correction.jl | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index fdfc10b26..d923d8af9 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -1,7 +1,7 @@ -import JustPIC._2D: is_above_surface, cell_index, interp1D_inner, interp1D_extremas, distance +import JustPIC._2D: cell_index, interp1D_inner, interp1D_extremas, distance using StaticArrays -function update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) +function update_phases_given_markerchain!(phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase) where {backend} (; coords, index) = particles; dy = di[2] @parallel (1:size(index,1)) _update_phases_given_markerchain!(phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase) @@ -13,10 +13,13 @@ end end function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell) - chain_yi = @cell chain_coords[2][icell] - cell_range = find_minmax_cell_indices(chain_yi, origin[2], dy) - T = eltype(eltype(phase)) - + T = eltype(eltype(phase)) + chain_yi = @cell chain_coords[2][icell] + min_cell_j, max_cell_j = find_minmax_cell_indices(chain_yi, origin[2], dy) + min_cell_j = max(1, min_cell_j - 10) + max_cell_j = max(1, max_cell_j + 10) + cell_range = min_cell_j:max_cell_j + # iterate over cells with marker chain on them for j in cell_range # iterate over particles j-th cell @@ -27,7 +30,7 @@ function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_co phaseq = @index phase[ip, icell, j] # check if particle is above the marker chain - above = is_above_surface2(xq, yq, chain_coords, cell_vertices) + above = is_above_chain(xq, yq, chain_coords, cell_vertices) # if the particle is above the surface and the phase is not air, set the phase to air if above && phaseq != air_phase @index phase[ip, icell, j] = T(air_phase) @@ -64,10 +67,10 @@ function find_minmax_cell_indices(chain_yi, origin_y, dy) ymin, ymax = extrema_CA(chain_yi) min_cell_j = Int((ymin - origin_y) ÷ dy) + 1 max_cell_j = Int((ymax - origin_y) ÷ dy) + 1 - return min_cell_j:max_cell_j + return min_cell_j, max_cell_j end -@inline function is_above_surface2(xq, yq, coords, cell_vertices) +@inline function is_above_chain(xq, yq, coords, cell_vertices) I = cell_index(xq, cell_vertices) x_cell, y_cell = coords[1][I], coords[2][I] ychain = if 1 < I[1] < length(cell_vertices) - 1 From 5d55199f36ae59245d9017d7ce157efc48fe7871 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 4 Dec 2024 15:46:41 +0100 Subject: [PATCH 38/77] add & export `update_phases_given_markerchain!` --- src/JustRelax_CPU.jl | 3 --- src/common.jl | 3 +++ src/ext/AMDGPU/2D.jl | 6 ++++++ src/ext/CUDA/2D.jl | 6 ++++++ 4 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 89f5056e3..67ada07a9 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -34,9 +34,6 @@ export solve! include("variational_stokes/Stokes2D.jl") export solve_VariationalStokes! -include("phases/topography_correction.jl") -export update_phases_given_markerchain! - end module JustRelax3D diff --git a/src/common.jl b/src/common.jl index d1fea9fae..23feaf4fb 100644 --- a/src/common.jl +++ b/src/common.jl @@ -71,6 +71,9 @@ export vertex2center!, include("advection/weno5.jl") export WENO_advection! +include("phases/topography_correction.jl") +export update_phases_given_markerchain! + # Stokes include("rheology/GeoParams.jl") diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index fe18347b8..31f0138d0 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -444,4 +444,10 @@ function JR2D.rotate_stress!( return nothing end +# marker chain + +function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{JustPIC.AMDGPUBackend}, particles::Particles{JustPIC.AMDGPUBackend}, origin, di, air_phase) + update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) +end + end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index fa724a72c..e3e15382e 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -425,4 +425,10 @@ function JR2D.rotate_stress!( return nothing end +# marker chain + +function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{CUDABackend}, particles::Particles{CUDABackend}, origin, di, air_phase) + update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) +end + end From 2f8358cdbe25188aded102161b7391ca6eefa1e6 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 4 Dec 2024 16:45:35 +0100 Subject: [PATCH 39/77] update miniapp --- .../PlumeFreeSurface_VariationalStokes.jl | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl index e65c80c70..b320c873b 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl @@ -145,9 +145,7 @@ function main(igg, nx, ny) # ---------------------------------------------------- # Initialize marker chain------------------------------- - nxcell, max_xcell, min_xcell = 24, 40, 12 - nxcell, min_xcell, max_xcell = 12, 6, 24 - + nxcell, max_xcell, min_xcell = 100, 150, 75 initial_elevation = -100e3 chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); # ---------------------------------------------------- @@ -184,7 +182,7 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 dt = 1e3 * (3600 * 24 * 365.25) - while it < 150 + while it < 40 # solve!( # stokes, @@ -233,14 +231,16 @@ function main(igg, nx, ny) move_particles!(particles, xvi, particle_args) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) - + # advect marker chain advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), grid_vxi, dt) update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) # ------------------------------ + @show it += 1 t += dt @@ -275,6 +275,7 @@ function main(igg, nx, ny) end end + return nothing end ## END OF MAIN SCRIPT ---------------------------------------------------------------- From f39ebc0cdd9bcb59962bcb37be82f70b45e0e32f Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 4 Dec 2024 17:28:19 +0100 Subject: [PATCH 40/77] tweak masking --- src/variational_stokes/mask.jl | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index f5e5c523e..ae1bf8941 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -131,10 +131,17 @@ Check if `ϕ.center[inds...]` is a not a nullspace. - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. - `inds`: Cartesian indices to check. """ +# Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) +# vx = isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vx[i + 1, j]) +# vy = isvalid(ϕ.Vy, i, j) * isvalid(ϕ.Vy[i, j + 1]) +# v = vx * vy +# return v * isvalid(ϕ.center, i, j) +# end + Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) vx = isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vx[i + 1, j]) vy = isvalid(ϕ.Vy, i, j) * isvalid(ϕ.Vy[i, j + 1]) - v = vx * vy + v = vx || vy return v * isvalid(ϕ.center, i, j) end @@ -165,10 +172,24 @@ Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, i_left = max(i - 1, 1) i0 = min(i, nx) vy = isvalid(ϕ.Vy, i0, j) * isvalid(ϕ.Vy, i_left, j) - v = vx * vy + v = vx || vy return v * isvalid(ϕ.vertex, i, j) end +# Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) +# nx, ny = size(ϕ.Vx) +# j_bot = max(j - 1, 1) +# j0 = min(j, ny) +# vx = isvalid(ϕ.Vx, i, j0) * isvalid(ϕ.Vx, i, j_bot) + +# nx, ny = size(ϕ.Vy) +# i_left = max(i - 1, 1) +# i0 = min(i, nx) +# vy = isvalid(ϕ.Vy, i0, j) * isvalid(ϕ.Vy, i_left, j) +# v = vx * vy +# return v * isvalid(ϕ.vertex, i, j) +# end + """ isvalid_vx(ϕ::JustRelax.RockRatio, inds...) @@ -184,6 +205,15 @@ Base.@propagate_inbounds @inline function isvalid_vx( return isvalid(ϕ.Vx, I...) end + +Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, I::Vararg{Integer,N}) where {N} + # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) + # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) + # cv = c * v + # return cv * (ϕ.Vx[i, j] > 0) + return (ϕ.Vx[I...] > 0) +end + # Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, I::Vararg{Integer,N}) where {N} # # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) # # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) From f035e196d2020f8baef609b1612132ad81dfeb6c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 4 Dec 2024 17:30:52 +0100 Subject: [PATCH 41/77] remove duplicated code --- src/variational_stokes/mask.jl | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index ae1bf8941..bd3f76772 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -205,15 +205,6 @@ Base.@propagate_inbounds @inline function isvalid_vx( return isvalid(ϕ.Vx, I...) end - -Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, I::Vararg{Integer,N}) where {N} - # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) - # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) - # cv = c * v - # return cv * (ϕ.Vx[i, j] > 0) - return (ϕ.Vx[I...] > 0) -end - # Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, I::Vararg{Integer,N}) where {N} # # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0) # # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0) From e2d040b59992cc65985b7bae53c9ca07974aec8f Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 5 Dec 2024 13:51:15 +0100 Subject: [PATCH 42/77] update caldera miniapp --- .../stokes2D/Volcano2D/Caldera2D.jl | 109 ++++++++++-------- 1 file changed, 63 insertions(+), 46 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 32547c548..edcea751e 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -13,7 +13,7 @@ else JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend end -using ParallelStencil, ParallelStencil.FiniteDifferences2D +using ParallelStencil @static if isCUDA @init_parallel_stencil(CUDA, Float64, 2) @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, CairoMakie, CellArrays, Statistics +using GeoParams, GLMakie, Statistics # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -40,12 +40,6 @@ include("Caldera_rheology.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- -import ParallelStencil.INDICES -const idx_k = INDICES[2] -macro all_k(A) - esc(:($A[$idx_k])) -end - function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -55,12 +49,6 @@ function copyinn_x!(A, B) @parallel f_x(A, B) end -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) - return nothing -end - function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) xv, yv = xvi @@ -85,6 +73,13 @@ function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) end ## END OF HELPER FUNCTION ------------------------------------------------------------ +function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) + topo_idx = [findfirst(x->x==air_phase, row) - 1 for row in eachrow(phases_GMG)] + yv = xvi[2] + topo_y = yv[topo_idx] + return topo_y +end + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) @@ -92,19 +87,19 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ni = nx, ny # number of cells di = @. li / ni # grid steps grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # Physical properties using GeoParams ---------------- rheology = init_rheologies() - dt = 5e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + dt = 5e2 * 3600 * 24 * 365 # dt = Inf # diffusive CFL timestep limiter # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell = 40 - max_xcell = 60 - min_xcell = 20 + nxcell = 100 + max_xcell = 150 + min_xcell = 75 particles = init_particles( backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni ) @@ -114,24 +109,39 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # material phase & temperature pPhases, pT = init_cell_arrays(particles, Val(2)) - # particle fields for the stress rotation - pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress - # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress - pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity - particle_args = (pT, pPhases, pτ..., pω...) - particle_args_reduced = (pT, pτ..., pω...) - # Assign particles phases anomaly phases_device = PTArray(backend)(phases_GMG) phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); init_phases!(pPhases, phases_device, particles, xvi) + + # Initialize marker chain + nxcell, max_xcell, min_xcell = 100, 150, 75 + initial_elevation = 0e0 + chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); + air_phase = 5 + topo_y = extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) + for _ in 1:3 + @views hn = 0.5 .* (topo_y[1:end-1] .+ topo_y[2:end]) + @views topo_y[2:end-1] .= 0.5 .* (hn[1:end-1] .+ hn[2:end]) + fill_chain_from_vertices!(chain, PTArray(backend)(topo_y)) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + end update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # particle fields for the stress rotation + # pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress + # # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress + # pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity + # particle_args = (pT, pPhases, pτ..., pω...) + # particle_args_reduced = (pT, pτ..., pω...) + pτ = StressParticles(particles) + particle_args = (pT, pPhases, unwrap(pτ)...) + particle_args_reduced = (pT, unwrap(pτ)...) + # rock ratios for variational stokes # RockRatios - air_phase = 5 ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -241,7 +251,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", t, it = 0.0, 0 thermal.Told .= thermal.T - while it < 500 #000 # run only for 5 Myrs + while it < 1000 #000 # run only for 5 Myrs # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) @@ -254,6 +264,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + stress2grid!(stokes, pτ, xvi, xci, particles) + t_stokes = @elapsed solve_VariationalStokes!( stokes, pt_stokes, @@ -273,20 +285,15 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ) ) - center2vertex!(τxx_v, stokes.τ.xx) - center2vertex!(τyy_v, stokes.τ.yy) - centroid2particle!(pτxx , xci, stokes.τ.xx, particles) - centroid2particle!(pτyy , xci, stokes.τ.yy, particles) - grid2particle!(pτxy, xvi, stokes.τ.xy, particles) - grid2particle!(pωxy, xvi, stokes.ω.xy, particles) - rotate_stress_particles!(pτ, pω, particles, dt) + # rotate stresses + rotate_stress!(pτ, stokes, particles, xci, xvi, dt) println("Stokes solver time ") println(" Total time: $t_stokes s") tensor_invariant!(stokes.ε) tensor_invariant!(stokes.ε_pl) dtmax = 2e3 * 3600 * 24 * 365.25 - dt = compute_dt(stokes, di, dtmax) + dt = compute_dt(stokes, di, dtmax) * 0.5 println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ @@ -303,7 +310,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", kwargs = ( igg = igg, phase = phase_ratios, - iterMax = 50e3, + iterMax = 100e3, nout = 1e2, verbose = true, ) @@ -328,6 +335,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", move_particles!(particles, xvi, particle_args) # check if we need to inject particles # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + center2vertex!(τxx_v, stokes.τ.xx) + center2vertex!(τyy_v, stokes.τ.yy) inject_particles_phase!( particles, pPhases, @@ -336,20 +345,21 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", xvi ) + # advect marker chain + advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase) + update_rock_ratio!(ϕ, phase_ratios, air_phase) - particle2centroid!(stokes.τ.xx, pτxx, xci, particles) - particle2centroid!(stokes.τ.yy, pτyy, xci, particles) - particle2grid!(stokes.τ.xy, pτxy, xvi, particles) tensor_invariant!(stokes.τ) @show it += 1 t += dt if plotting # Data I/O and plotting --------------------- - if it == 1 || rem(it, 1) == 0 + if it == 1 || rem(it, 5) == 0 if igg.me == 0 && it == 1 metadata(pwd(), checkpoint, basename(@__FILE__), "Caldera_setup.jl", "Caldera_rheology.jl") end @@ -391,6 +401,10 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", clr = pPhases.data[:] # clr = pT.data[:] idxv = particles.index.data[:]; + + chain_x = chain.coords[1].data[:]./1e3 + chain_y = chain.coords[2].data[:]./1e3 + # Make Makie figure ar = 2 fig = Figure(size = (1200, 900), title = "t = $t") @@ -403,7 +417,9 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Plot temperature h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) # Plot particles phase - h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 3) + scatter!(ax2, Array(chain_x), Array(chain_y), color=:red, markersize = 3) + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) # Plot 2nd invariant of strain rate # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) @@ -453,16 +469,17 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- const plotting = true + do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "Caldera2D" -n = 128 +n = 256 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; sticky_air = 4, flat = false, - chimney = false, + chimney = true, chamber_T = 1e3, chamber_depth = 7e0, chamber_radius = 0.5, From a8f2a66fb90f422265c1196fc572b4f40fc8cd16 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 5 Dec 2024 14:33:36 +0100 Subject: [PATCH 43/77] fix spurious assymetries in the masking --- src/variational_stokes/mask.jl | 30 +++++------------------------- 1 file changed, 5 insertions(+), 25 deletions(-) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index bd3f76772..4247d7717 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -101,7 +101,7 @@ end @inline compute_rock_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} -) where {N} = 1 - @index phase_ratio[air_phase, I...] +) where {N} = (x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5) @inline compute_air_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} @@ -118,7 +118,8 @@ end end @parallel_indices (I...) function _update_rock_ratio!(ϕ, ratio, air_phase) - ϕ[I...] = Float64(Float16(compute_rock_ratio(ratio, air_phase, I...))) + # ϕ[I...] = Float64(Float16(compute_rock_ratio(ratio, air_phase, I...))) + ϕ[I...] = clamp(compute_rock_ratio(ratio, air_phase, I...), 0, 1) return nothing end @@ -131,17 +132,10 @@ Check if `ϕ.center[inds...]` is a not a nullspace. - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against. - `inds`: Cartesian indices to check. """ -# Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) -# vx = isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vx[i + 1, j]) -# vy = isvalid(ϕ.Vy, i, j) * isvalid(ϕ.Vy[i, j + 1]) -# v = vx * vy -# return v * isvalid(ϕ.center, i, j) -# end - Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j) vx = isvalid(ϕ.Vx, i, j) * isvalid(ϕ.Vx[i + 1, j]) vy = isvalid(ϕ.Vy, i, j) * isvalid(ϕ.Vy[i, j + 1]) - v = vx || vy + v = vx * vy return v * isvalid(ϕ.center, i, j) end @@ -172,24 +166,10 @@ Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, i_left = max(i - 1, 1) i0 = min(i, nx) vy = isvalid(ϕ.Vy, i0, j) * isvalid(ϕ.Vy, i_left, j) - v = vx || vy + v = vx * vy return v * isvalid(ϕ.vertex, i, j) end -# Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j) -# nx, ny = size(ϕ.Vx) -# j_bot = max(j - 1, 1) -# j0 = min(j, ny) -# vx = isvalid(ϕ.Vx, i, j0) * isvalid(ϕ.Vx, i, j_bot) - -# nx, ny = size(ϕ.Vy) -# i_left = max(i - 1, 1) -# i0 = min(i, nx) -# vy = isvalid(ϕ.Vy, i0, j) * isvalid(ϕ.Vy, i_left, j) -# v = vx * vy -# return v * isvalid(ϕ.vertex, i, j) -# end - """ isvalid_vx(ϕ::JustRelax.RockRatio, inds...) From 0d22ef9cbfa60b783276e091eedc390c5e4afbcd Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 5 Dec 2024 16:52:29 +0100 Subject: [PATCH 44/77] up caldera2D --- .../stokes2D/Volcano2D/Caldera2D.jl | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index edcea751e..febcb53f9 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -1,5 +1,5 @@ -# const isCUDA = false -const isCUDA = true +const isCUDA = false +# const isCUDA = true @static if isCUDA using CUDA @@ -13,7 +13,7 @@ else JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend end -using ParallelStencil +using ParallelStencil, ParallelStencil.FiniteDifferences2D @static if isCUDA @init_parallel_stencil(CUDA, Float64, 2) @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, GLMakie, Statistics +using GeoParams, GLMakie, CellArrays, Statistics # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -40,6 +40,12 @@ include("Caldera_rheology.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +import ParallelStencil.INDICES +const idx_k = INDICES[2] +macro all_k(A) + esc(:($A[$idx_k])) +end + function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -49,6 +55,12 @@ function copyinn_x!(A, B) @parallel f_x(A, B) end +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end + function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) xv, yv = xvi @@ -473,7 +485,7 @@ const plotting = true do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "Caldera2D" -n = 256 +n = 128 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; @@ -486,9 +498,6 @@ li, origin, phases_GMG, T_GMG = setup2D( aspect_x = 6, ) -# heatmap(phases_GMG) -# heatmap(T_GMG) - igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) else From c199a0fe8a211cb06ab2c84e475423dd4d2e7279 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 5 Dec 2024 16:58:31 +0100 Subject: [PATCH 45/77] fix --- src/phases/topography_correction.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index d923d8af9..84bbe7107 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -1,6 +1,7 @@ import JustPIC._2D: cell_index, interp1D_inner, interp1D_extremas, distance using StaticArrays + function update_phases_given_markerchain!(phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase) where {backend} (; coords, index) = particles; dy = di[2] @@ -17,9 +18,9 @@ function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_co chain_yi = @cell chain_coords[2][icell] min_cell_j, max_cell_j = find_minmax_cell_indices(chain_yi, origin[2], dy) min_cell_j = max(1, min_cell_j - 10) - max_cell_j = max(1, max_cell_j + 10) + max_cell_j = min(size(index, 2), max_cell_j + 10) cell_range = min_cell_j:max_cell_j - + # iterate over cells with marker chain on them for j in cell_range # iterate over particles j-th cell From 0f07e20d2b346bd111a0afdf3020f5f68f37ee45 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 5 Dec 2024 16:59:09 +0100 Subject: [PATCH 46/77] small corrections --- src/variational_stokes/MiniKernels.jl | 14 +++++++------- src/variational_stokes/mask.jl | 18 ------------------ 2 files changed, 7 insertions(+), 25 deletions(-) diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 6f9d217b4..7d0c13ebe 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -47,10 +47,10 @@ Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where # averages 3D Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j, k) where {T<:T3} = - 0.125 * mysum(A, ϕ, (i + 1):(i + 2), (j + 1):(j + 2), (k + 1):(k + 2)) + 0.125 * mymaskedsum(A, ϕ, (i + 1):(i + 2), (j + 1):(j + 2), (k + 1):(k + 2)) Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j, k) where {T<:T3} = - 0.125 * mysum(A, ϕ, (i):(i + 1), (j):(j + 1), (k):(k + 1)) + 0.125 * mymaskedsum(A, ϕ, (i):(i + 1), (j):(j + 1), (k):(k + 1)) Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5 @@ -70,10 +70,10 @@ Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,3}) where Base.@propagate_inbounds @inline _av_zi(A::T, ϕ::T, I::Vararg{Integer,3}) where {T<:T3} = (top(A, ϕ, I...) + next(A, ϕ, I...)) * 0.5 -## Because mysum(::generator) does not work inside CUDA kernels... -@inline mysum(A, ϕ, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ϕ, ranges...) +## Because mymaskedsum(::generator) does not work inside CUDA kernels... +@inline mymaskedsum(A::AbstractArray, ϕ::AbstractArray, ranges::Vararg{T,N}) where {T,N} = mymaskedsum(identity, A, ϕ, ranges...) -@inline function mysum( +@inline function mymaskedsum( f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i ) where {F<:Function} s = 0.0 @@ -83,7 +83,7 @@ Base.@propagate_inbounds @inline _av_zi(A::T, ϕ::T, I::Vararg{Integer,3}) where return s end -@inline function mysum( +@inline function mymaskedsum( f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j ) where {F<:Function} s = 0.0 @@ -93,7 +93,7 @@ end return s end -@inline function mysum( +@inline function mymaskedsum( f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j, ranges_k ) where {F<:Function} s = 0.0 diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 4247d7717..69bae8b65 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -81,24 +81,6 @@ function update_rock_ratio!( return nothing end -# function update_rock_ratio!( -# ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase -# ) where {N} -# nvi = size_v(ϕ) -# @parallel (@idx nvi) update_rock_ratio_cv!( -# ϕ, phase_ratios.center, phase_ratios.vertex, air_phase -# ) -# # @parallel (@idx nvi) _update_rock_ratio!(ϕ) - -# @parallel (@idx size(ϕ.Vx)) _update_rock_ratio!(ϕ.Vx, ratio_vel[1], air_phase) -# @parallel (@idx size(ϕ.Vy)) _update_rock_ratio!(ϕ.Vy, ratio_vel[2], air_phase) -# if N === 3 -# @parallel (@idx size(ϕ.Vz)) _update_rock_ratio!(ϕ.Vz, ratio_vel[3], air_phase) -# end - -# return nothing -# end - @inline compute_rock_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} ) where {N} = (x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5) From cbde08f7fa4239db9177e4c2a6931dda4b50ef41 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Wed, 11 Dec 2024 10:01:49 +0100 Subject: [PATCH 47/77] update caldera miniapp --- .../stokes2D/Volcano2D/Caldera2D.jl | 205 ++++++++++++------ .../stokes2D/Volcano2D/Caldera_rheology.jl | 24 +- .../stokes2D/Volcano2D/Caldera_setup.jl | 35 +-- 3 files changed, 185 insertions(+), 79 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index febcb53f9..e8eec2016 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, GLMakie, CellArrays, Statistics +using GeoParams, CairoMakie, CellArrays, Statistics, Dates # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -92,6 +92,25 @@ function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) return topo_y end +function thermal_anomaly!(Temp, mask, Ω_T, phases, particles, conduit_phase, magma_phase) + + @parallel_indices (i, j) function _thermal_anomaly!(Temp, mask, Ω_T, phases, index, conduit_phase, magma_phase) + @inbounds for ip in cellaxes(phases) + #quick escape + @index(index[ip, i, j]) == 0 && continue + phase_ij = @index phases[ip, i, j] + + if phase_ij == conduit_phase || phase_ij == magma_phase + Temp[i+1, j+1] = Ω_T + mask[i+1,j+1] = 1 + end + end + return nothing + end + + @parallel _thermal_anomaly!(Temp, mask, Ω_T, phases, particles.index, conduit_phase, magma_phase) +end + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) @@ -125,12 +144,12 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", phases_device = PTArray(backend)(phases_GMG) phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); init_phases!(pPhases, phases_device, particles, xvi) - + # Initialize marker chain nxcell, max_xcell, min_xcell = 100, 150, 75 initial_elevation = 0e0 chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); - air_phase = 5 + air_phase = 6 topo_y = extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) for _ in 1:3 @views hn = 0.5 .* (topo_y[1:end-1] .+ topo_y[2:end]) @@ -165,8 +184,14 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # TEMPERATURE PROFILE -------------------------------- thermal = ThermalArrays(backend, ni) @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) + + # Add thermal anomaly BC's + Ω_T = 1223e0 # inner BCs temperature + mask = @zeros(size(thermal.T)...) + thermal_anomaly!(thermal.T, mask, Ω_T, pPhases, particles, 5, 3) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), + dirichlet = (constant = Ω_T, mask=mask) ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -176,12 +201,17 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(2)) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + # Melt fraction + ϕ_m = @zeros(ni...) + compute_melt_fraction!( + ϕ_m, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P) + ) # Rheology - args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) - viscosity_cutoff = (1e18, 1e23) + args0 = (; ϕ=ϕ_m,T=thermal.Tc, P=stokes.P, dt = Inf) + viscosity_cutoff = (1e16, 1e23) compute_viscosity!(stokes, phase_ratios, args0, rheology, air_phase, viscosity_cutoff) # PT coefficients for thermal diffusion @@ -235,26 +265,28 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", grid2particle!(pT, xvi, T_buffer, particles) ## Plot initial T and P profile - # fig = 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(; size=(1200, 900)) - # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - # scatter!( - # ax1, - # Array(thermal.T[2:(end - 1), :][:]), - # Yv, - # ) - # lines!( - # ax2, - # Array(stokes.P[:]), - # Y, - # ) - # hideydecorations!(ax2) - # # save(joinpath(figdir, "initial_profile.png"), fig) - # fig - # end + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Density") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:]), + Yv./1e3, + ) + # lines!( + scatter!( + ax2, + # Array(stokes.P[:]./1e6), + Array(ρg[2][:]./9.81), + Y./1e3, + ) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end τxx_v = @zeros(ni.+1...) τyy_v = @zeros(ni.+1...) @@ -263,18 +295,21 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", t, it = 0.0, 0 thermal.Told .= thermal.T - while it < 1000 #000 # run only for 5 Myrs + while it < 10 #000 # run only for 5 Myrs # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end-1, :] .= T_buffer + if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) + thermal_anomaly!(thermal.T, mask, Ω_T, pPhases, particles, 5, 3) + end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) - args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + args = (; ϕ=ϕ_m, T=thermal.Tc, P=stokes.P, dt=Inf) stress2grid!(stokes, pτ, xvi, xci, particles) @@ -360,18 +395,25 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # advect marker chain advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), grid_vxi, dt) update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) - + + compute_melt_fraction!( + ϕ_m, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P) + ) + # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) update_rock_ratio!(ϕ, phase_ratios, air_phase) tensor_invariant!(stokes.τ) + # track deformation of free_surface + push!(deformation_x, chain.coords[1].data[:]./1e3) + push!(deformation_y, chain.coords[2].data[:]./1e3) @show it += 1 t += dt if plotting # Data I/O and plotting --------------------- - if it == 1 || rem(it, 5) == 0 + if it == 1 || rem(it, 1) == 0 if igg.me == 0 && it == 1 metadata(pwd(), checkpoint, basename(@__FILE__), "Caldera_setup.jl", "Caldera_rheology.jl") end @@ -401,7 +443,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", data_v, data_c, velocity_v; - t = 0 + t = round(t/(1e3 * 3600 * 24 *365.25); digits=3) ) end @@ -420,8 +462,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Make Makie figure ar = 2 fig = Figure(size = (1200, 900), title = "t = $t") - ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") - ax2 = Axis(fig[2,1], aspect = ar, title = "Vy") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(round(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [cm/yr]") # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII)") @@ -429,10 +471,10 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Plot temperature h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) # Plot particles phase - h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 3) - scatter!(ax2, Array(chain_x), Array(chain_y), color=:red, markersize = 3) + # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 3) + # scatter!(ax2, Array(chain_x), Array(chain_y), color=:red, markersize = 3) - # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) + h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, uconvert.(u"cm/yr",Array(stokes.V.Vy)u"m/s") , colormap=:batlow) # Plot 2nd invariant of strain rate # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) @@ -451,26 +493,47 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", save(joinpath(figdir, "$(it).png"), fig) # ## Plot initial T and P profile - # fig = 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(; size=(1200, 900)) - # ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - # ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - # scatter!( - # ax1, - # Array(thermal.T[2:(end - 1), :][:]), - # Yv, - # ) - # lines!( - # ax2, - # Array(stokes.P[:]), - # Y, - # ) - # hideydecorations!(ax2) - # save(joinpath(figdir, "thermal_profile_$it.png"), fig) - # fig - # end + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:].-273.15), + Yv./1e3, + ) + lines!( + ax2, + Array(stokes.P[:]./1e6), + Y./1e3, + ) + hideydecorations!(ax2) + # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + fig + end + + fig = 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(; size=(1200, 900)) + ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") + ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") + scatter!( + ax1, + Array(thermal.T[2:(end - 1), :][:].-273.15), + Yv./1e3, + ) + lines!( + ax2, + Array(stokes.P[:]./1e6), + Y./1e3, + ) + hideydecorations!(ax2) + # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + fig + end end # ------------------------------ end @@ -484,17 +547,20 @@ const plotting = true do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" -figdir = "Caldera2D" +figdir = "$(today())_Conduit" n = 128 -nx, ny = n, n >>> 1 +nx, ny = n, n #>>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; - sticky_air = 4, + sticky_air = 4e0, + dimensions = (30e0, 20e0), # extent in x and y in km flat = false, chimney = true, - chamber_T = 1e3, + volcano_size = (3e0, 5e0), + conduit_radius = 0.5e0, + chamber_T = 900e0, chamber_depth = 7e0, - chamber_radius = 0.5, + chamber_radius = 1e0, aspect_x = 6, ) @@ -505,3 +571,20 @@ else end main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); + +function plot_particles(particles, pPhases, chain) + p = particles.coords + # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter + ppx, ppy = p + # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) + # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) + pxv = ppx.data[:] + pyv = ppy.data[:] + clr = pPhases.data[:] + # clr = pϕ.data[:] + idxv = particles.index.data[:] + f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma, markersize=1) + + Colorbar(f[1,2], h) + f +end diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl index 9f7d6bd3f..7bb80e260 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl @@ -42,7 +42,7 @@ function init_rheology_linear() end function init_rheologies(; linear=false) - + η_reg = 1e18 C = linear ? Inf : 10e6 ϕ = 15 @@ -53,11 +53,12 @@ function init_rheologies(; linear=false) β = 1 / el.Kb.val Cp = 1200.0 + magma_visc = ViscosityPartialMelt_Costa_etal_2009(η=LinearMeltViscosity(A = -8.1590, B = 2.4050e+04K, T0 = -430.9606K,η0=1e3Pa*s)) #dislocation laws disl_top = SetDislocationCreep(Dislocation.dry_olivine_Karato_2003) # diffusion laws disl_bot = SetDislocationCreep(Dislocation.wet_quartzite_Hirth_2001) - + # Define rheolgy struct rheology = ( # Name = "Upper crust", @@ -68,6 +69,7 @@ function init_rheologies(; linear=false) Conductivity = ConstantConductivity(; k = 2.5), CompositeRheology = CompositeRheology( (disl_top, el, pl)), # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23), el, pl)), + Melting = MeltingParam_Smooth3rdOrder(a=517.9, b=-1619.0, c=1699.0, d = -597.4), #mafic melting curve Gravity = ConstantGravity(; g=9.81), ), # Name = "Lower crust", @@ -78,6 +80,7 @@ function init_rheologies(; linear=false) Conductivity = ConstantConductivity(; k = 2.5), CompositeRheology = CompositeRheology( (disl_bot, el, pl)), # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), + Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve Gravity = ConstantGravity(; g=9.81), ), @@ -90,6 +93,7 @@ function init_rheologies(; linear=false) HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el)), + Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), # Name = "magma chamber - hot anomaly", SetMaterialParams(; @@ -100,15 +104,29 @@ function init_rheologies(; linear=false) HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el, )), + Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), - # Name = "StickyAir", + + # Name = "Conduit", SetMaterialParams(; Phase = 5, + Density = T_Density(; ρ0=1.5e3, T0=273.15), + Conductivity = ConstantConductivity(; k = 1.5), + # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), + HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), + LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el, )), + Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 6, Density = ConstantDensity(; ρ=1e0), HeatCapacity = ConstantHeatCapacity(; Cp = Cp), Conductivity = ConstantConductivity(; k = 2.5), CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), Gravity = ConstantGravity(; g=9.81), + # Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), ) end diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl index 32f60911f..d08355f2a 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl @@ -2,24 +2,27 @@ using GeophysicalModelGenerator function setup2D( - nx, nz; - sticky_air = 2.5, - flat = true, + nx, nz; + sticky_air = 5e0, + dimensions = (30e0, 20e0), # extent in x and y in km + flat = true, chimney = false, + volcano_size = (3e0, 5e0), + conduit_radius = 0.25, chamber_T = 1e3, chamber_depth = 5e0, chamber_radius = 2e0, aspect_x = 1.5, ) - Lx = Ly = 50 + Lx = Ly = dimensions[1] x = range(0.0, Lx, nx); y = range(0.0, Ly, 2); - z = range(-25, sticky_air, nz); + z = range(-dimensions[2], sticky_air, nz); Grid = CartData(xyz_grid(x,y,z)); # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) - Phases = fill(5, nx, 2, nz); + Phases = fill(6, nx, 2, nz); # In many (geodynamic) models, one also has to define the temperature, so lets define it as well Temp = fill(0.0, nx, 2, nz); @@ -36,8 +39,8 @@ function setup2D( add_volcano!(Phases, Temp, Grid; volcanic_phase = 1, center = (mean(Grid.x.val), 0.0), - height = 3, - radius = 5, + height = volcano_size[1], + radius = volcano_size[2], crater = 0.5, base = 0.0, background = nothing, @@ -53,18 +56,18 @@ function setup2D( ) # add_sphere!(Phases, Temp, Grid; - # cen = (mean(Grid.x.val), 0, -chamber_depth), - # radius = 0.5, + # cen = (mean(Grid.x.val), 0, -(chamber_depth-(chamber_radius/2))), + # radius = (chamber_radius/2), # phase = ConstantPhase(4), - # T = ConstantTemp(T=chamber_T) + # T = ConstantTemp(T=chamber_T+100) # ) if chimney add_cylinder!(Phases, Temp, Grid; - base = (mean(Grid.x.val), 0, -chamber_depth), - cap = (mean(Grid.x.val), 0, 3e0), - radius = 0.05, - phase = ConstantPhase(3), + base = (mean(Grid.x.val), 0, -(chamber_depth-chamber_radius)), + cap = (mean(Grid.x.val), 0, flat ? 0e0 : volcano_size[1]), + radius = conduit_radius, + phase = ConstantPhase(5), # T = LinearTemp(Ttop=20, Tbot=1000), # T = ConstantTemp(T=800), T = ConstantTemp(T=chamber_T), @@ -77,6 +80,8 @@ function setup2D( ph = Phases[:,1,:] T = Temp[:,1,:] .+ 273 + V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * 1.0 + printstyled("Magma volume of the initial chamber: $(round(V; digits=3)) km³ \n"; bold=true, color=:red) # write_paraview(Grid, "Volcano2D") return li, origin, ph, T, Grid end From a56c5e81df78a50ea63d0ab44f2b31e1fc577353 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 11 Dec 2024 12:50:28 +0100 Subject: [PATCH 48/77] fix Dirichlet T mask --- .../stokes2D/Volcano2D/Caldera2D.jl | 69 ++++++++++--------- .../stokes2D/Volcano2D/Caldera_setup.jl | 2 +- 2 files changed, 38 insertions(+), 33 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index e8eec2016..1bfd267a5 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, CairoMakie, CellArrays, Statistics, Dates +using GeoParams, GLMakie, CellArrays, Statistics, Dates # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -92,25 +92,30 @@ function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) return topo_y end -function thermal_anomaly!(Temp, mask, Ω_T, phases, particles, conduit_phase, magma_phase) +function thermal_anomaly!(Temp, mask, Ω_T, phase_ratios, conduit_phase, magma_phase) - @parallel_indices (i, j) function _thermal_anomaly!(Temp, mask, Ω_T, phases, index, conduit_phase, magma_phase) - @inbounds for ip in cellaxes(phases) - #quick escape - @index(index[ip, i, j]) == 0 && continue - phase_ij = @index phases[ip, i, j] + @parallel_indices (i, j) function _thermal_anomaly!(Temp, mask, Ω_T, vertex_ratio, conduit_phase, magma_phase) + # quick escape + conduit_ratio_ij = @index vertex_ratio[conduit_phase, i, j] + magma_ratio_ij = @index vertex_ratio[magma_phase, i, j] - if phase_ij == conduit_phase || phase_ij == magma_phase - Temp[i+1, j+1] = Ω_T - mask[i+1,j+1] = 1 - end + if conduit_ratio_ij > 0.5 || magma_ratio_ij > 0.5 + # if isone(conduit_ratio_ij) || isone(magma_ratio_ij) + Temp[i+1, j] = Ω_T + mask[i+1, j] = 1 end + return nothing end + + ni = size(phase_ratios.vertex) - @parallel _thermal_anomaly!(Temp, mask, Ω_T, phases, particles.index, conduit_phase, magma_phase) + @parallel (@idx ni) _thermal_anomaly!(Temp, mask, Ω_T, phase_ratios.vertex, conduit_phase, magma_phase) + + return nothing end + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) @@ -188,7 +193,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Add thermal anomaly BC's Ω_T = 1223e0 # inner BCs temperature mask = @zeros(size(thermal.T)...) - thermal_anomaly!(thermal.T, mask, Ω_T, pPhases, particles, 5, 3) + thermal_anomaly!(thermal.T, mask, Ω_T, phase_ratios, 5, 3) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), dirichlet = (constant = Ω_T, mask=mask) @@ -549,7 +554,7 @@ do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "$(today())_Conduit" n = 128 -nx, ny = n, n #>>> 1 +nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; sticky_air = 4e0, @@ -570,21 +575,21 @@ else igg end -main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); - -function plot_particles(particles, pPhases, chain) - p = particles.coords - # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter - ppx, ppy = p - # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) - # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) - pxv = ppx.data[:] - pyv = ppy.data[:] - clr = pPhases.data[:] - # clr = pϕ.data[:] - idxv = particles.index.data[:] - f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma, markersize=1) - - Colorbar(f[1,2], h) - f -end +# main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); + +# function plot_particles(particles, pPhases, chain) +# p = particles.coords +# # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter +# ppx, ppy = p +# # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) +# # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) +# pxv = ppx.data[:] +# pyv = ppy.data[:] +# clr = pPhases.data[:] +# # clr = pϕ.data[:] +# idxv = particles.index.data[:] +# f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma, markersize=1) + +# Colorbar(f[1,2], h) +# f +# end diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl index d08355f2a..d88f10878 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl @@ -80,7 +80,7 @@ function setup2D( ph = Phases[:,1,:] T = Temp[:,1,:] .+ 273 - V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * 1.0 + V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * 1.0 printstyled("Magma volume of the initial chamber: $(round(V; digits=3)) km³ \n"; bold=true, color=:red) # write_paraview(Grid, "Volcano2D") return li, origin, ph, T, Grid From 7ad8d6075255f45e1ffbef0d3f1c78f2171060af Mon Sep 17 00:00:00 2001 From: aelligp Date: Wed, 11 Dec 2024 17:16:12 +0100 Subject: [PATCH 49/77] up script --- .../benchmarks/stokes2D/Volcano2D/Caldera2D.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 1bfd267a5..ee4e41d3e 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -1,5 +1,5 @@ -const isCUDA = false -# const isCUDA = true +# const isCUDA = false +const isCUDA = true @static if isCUDA using CUDA @@ -107,11 +107,11 @@ function thermal_anomaly!(Temp, mask, Ω_T, phase_ratios, conduit_phase, magma_p return nothing end - + ni = size(phase_ratios.vertex) @parallel (@idx ni) _thermal_anomaly!(Temp, mask, Ω_T, phase_ratios.vertex, conduit_phase, magma_phase) - + return nothing end @@ -307,8 +307,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end-1, :] .= T_buffer - if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) - thermal_anomaly!(thermal.T, mask, Ω_T, pPhases, particles, 5, 3) + if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 + thermal_anomaly!(thermal.T, mask, Ω_T, phase_ratios, 5, 3) end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -544,7 +544,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", end end - return nothing + return nothing end ## END OF MAIN SCRIPT ---------------------------------------------------------------- From 9fb3249f9fcea3a0c9281c5d2be4dc40defd10da Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 11 Dec 2024 17:45:16 +0100 Subject: [PATCH 50/77] fix Dirichlet on GPUs --- src/boundaryconditions/Dirichlet.jl | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl index 0c8d259fe..52978161f 100644 --- a/src/boundaryconditions/Dirichlet.jl +++ b/src/boundaryconditions/Dirichlet.jl @@ -25,22 +25,29 @@ function Base.getindex(x::DirichletBoundaryCondition, inds::Vararg{Int,N}) where return x.value[inds...] * x.mask[inds...] end -struct ConstantArray{T,N} <: AbstractArray{T,N} +# struct ConstantArray{T,N} <: AbstractArray{T,N} +# val::T + +# ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val) +# end +# Adapt.@adapt_structure ConstantArray + +struct ConstantArray{T} val::T - - ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val) + + ConstantArray(val::T) where {T<:Number} = new{T}(val) end Adapt.@adapt_structure ConstantArray Base.getindex(A::ConstantArray, ::Vararg{Int,N}) where {N} = A.val Base.setindex!(::ConstantArray, ::Any, ::Vararg{Int,N}) where {N} = nothing -function Base.show(io::IO, ::MIME"text/plain", A::ConstantArray{T,N}) where {T,N} - println(io, "ConstantArray{$T,$N}:") +function Base.show(io::IO, ::MIME"text/plain", A::ConstantArray{T}) where {T} + println(io, "ConstantArray{$T}:") return println(io, " ", A.val) end -function Base.show(io::IO, A::ConstantArray{T,N}) where {T,N} - println(io, "ConstantArray{$T,$N}:") +function Base.show(io::IO, A::ConstantArray{T}) where {T} + println(io, "ConstantArray{$T}:") return println(io, " ", A.val) end @@ -49,7 +56,7 @@ struct ConstantDirichletBoundaryCondition{T,M} <: AbstractDirichletBoundaryCondi mask::M function ConstantDirichletBoundaryCondition(value::N, mask::M) where {N<:Number,M} - v = ConstantArray(value, Val(dims(mask))) + v = ConstantArray(value) T = typeof(v) return new{T,M}(v, mask) end From 46ac3c0aca0182a81dab39b752e2bd6aef38127d Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 11 Dec 2024 18:17:16 +0100 Subject: [PATCH 51/77] fix Dirichlet on GPUs --- src/boundaryconditions/Dirichlet.jl | 32 ++++++++++++++--------------- src/mask/mask.jl | 16 +++++++-------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl index 52978161f..14dbd2513 100644 --- a/src/boundaryconditions/Dirichlet.jl +++ b/src/boundaryconditions/Dirichlet.jl @@ -2,14 +2,14 @@ abstract type AbstractDirichletBoundaryCondition{T,M} end struct DirichletBoundaryCondition{T,M} <: AbstractDirichletBoundaryCondition{T,M} value::T mask::M +end - function DirichletBoundaryCondition(value::T, mask::M) where {T,M} - return new{T,M}(value, mask) - end +# function DirichletBoundaryCondition(value::T, mask::M) where {T,M} +# return DirichletBoundaryCondition{T,M}(value, mask) +# end - function DirichletBoundaryCondition() - return new{Nothing,Nothing}(nothing, nothing) - end +function DirichletBoundaryCondition() + return DirichletBoundaryCondition{Nothing,Nothing}(nothing, nothing) end Adapt.@adapt_structure DirichletBoundaryCondition @@ -54,16 +54,15 @@ end struct ConstantDirichletBoundaryCondition{T,M} <: AbstractDirichletBoundaryCondition{T,M} value::T mask::M +end - function ConstantDirichletBoundaryCondition(value::N, mask::M) where {N<:Number,M} - v = ConstantArray(value) - T = typeof(v) - return new{T,M}(v, mask) - end +function ConstantDirichletBoundaryCondition(value::T, mask::M) where {T<:Number,M} + v = ConstantArray(value) + return ConstantDirichletBoundaryCondition{typeof(v),M}(v, mask) +end - function ConstantDirichletBoundaryCondition() - return new{Nothing,Nothing}(nothing, nothing) - end +function ConstantDirichletBoundaryCondition() + return ConstantDirichletBoundaryCondition{Nothing,Nothing}(nothing, nothing) end Adapt.@adapt_structure ConstantDirichletBoundaryCondition @@ -72,13 +71,14 @@ function Base.getindex(x::ConstantDirichletBoundaryCondition, inds::Vararg{Int,N return x.value * x.mask[inds...] end function Base.getindex( - x::ConstantDirichletBoundaryCondition{Nothing,Nothing}, ::Vararg{Int,N} + ::ConstantDirichletBoundaryCondition{Nothing,Nothing}, ::Vararg{Int,N} ) where {N} return 0 end @inline function apply_dirichlet!(A::AbstractArray, bc::AbstractDirichletBoundaryCondition) - return apply_mask!(A, bc.value, bc.mask) + apply_mask!(A, bc.value, bc.mask) + return nothing end @inline function apply_dirichlet!( diff --git a/src/mask/mask.jl b/src/mask/mask.jl index 6147cefe9..8370a6f41 100644 --- a/src/mask/mask.jl +++ b/src/mask/mask.jl @@ -35,21 +35,21 @@ Base.similar(m::Mask) = Mask(size(m)...) @inline dims(::Mask{A}) where {A<:AbstractArray{T,N}} where {T,N} = N -@inline apply_mask!(A::AbstractArray, B::AbstractArray, m::Mask) = +@inline apply_mask!(A::AbstractArray, B::Any, m::Mask) = (A .= inv(m) .* A .+ m.mask .* B) -@inline apply_mask!(::AbstractArray, ::AbstractArray, ::Nothing) = nothing +@inline apply_mask!(::AbstractArray, ::Any, ::Nothing) = nothing @inline apply_mask!( - A::AbstractArray, B::AbstractArray, m::Mask, inds::Vararg{Int,N} + A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N} ) where {N} = (A[inds...] = inv(m, inds...) * A[inds...] + m[inds...] * B[inds...]) @inline apply_mask!( - ::AbstractArray, ::AbstractArray, ::Nothing, inds::Vararg{Int,N} + ::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N} ) where {N} = nothing -@inline apply_mask(A::AbstractArray, B::AbstractArray, m::Mask) = inv(m) .* A .+ m.mask .* B +@inline apply_mask(A::AbstractArray, B::Any, m::Mask) = inv(m) .* A .+ m.mask .* B @inline apply_mask( - A::AbstractArray, B::AbstractArray, m::Mask, inds::Vararg{Int,N} + A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N} ) where {N} = @muladd inv(m, inds...) * A[inds...] + m[inds...] * B[inds...] -@inline apply_mask(A::AbstractArray, ::AbstractArray, ::Nothing) = A +@inline apply_mask(A::AbstractArray, ::Any, ::Nothing) = A @inline apply_mask( - A::AbstractArray, ::AbstractArray, ::Nothing, inds::Vararg{Int,N} + A::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N} ) where {N} = A[inds...] From 47fa0bfc66501a1c3e0c9a38da1c175f3d152cf3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 12 Dec 2024 10:46:51 +0100 Subject: [PATCH 52/77] remove particles that cross the marker chain --- src/phases/topography_correction.jl | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index 84bbe7107..6c96e01a0 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -34,11 +34,17 @@ function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_co above = is_above_chain(xq, yq, chain_coords, cell_vertices) # if the particle is above the surface and the phase is not air, set the phase to air if above && phaseq != air_phase - @index phase[ip, icell, j] = T(air_phase) + # @index phase[ip, icell, j] = T(air_phase) + @index coords[1][ip, icell, j] = NaN + @index coords[2][ip, icell, j] = NaN + @index index[ip, icell, j] = 0 end # if the particle is above the surface and the phase is air, set the phase to the closes rock phase if !above && phaseq == air_phase - @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) + # @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) + @index coords[1][ip, icell, j] = NaN + @index coords[2][ip, icell, j] = NaN + @index index[ip, icell, j] = 0 end end end From 45a05f44a6c58ac1b52d0e4b3a9e45837a624331 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 13 Dec 2024 14:26:36 +0100 Subject: [PATCH 53/77] up --- .../PlumeFreeSurface_VariationalStokes.jl | 34 ++++--------------- 1 file changed, 7 insertions(+), 27 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl index b320c873b..c414762f8 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl @@ -1,5 +1,5 @@ -const isCUDA = false -# const isCUDA = true +# const isCUDA = false +const isCUDA = true @static if isCUDA using CUDA @@ -181,28 +181,9 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) - while it < 40 - - # solve!( - # stokes, - # pt_stokes, - # di, - # flow_bcs, - # ρg, - # phase_ratios, - # rheology, - # args, - # dt, - # igg; - # kwargs = (; - # iterMax = 50e3, - # nout = 1e3, - # viscosity_cutoff = (-Inf, Inf), - # free_surface = true - # ) - # ) - + dt = 10e3 * (3600 * 24 * 365.25) + while it < 150 + # Stokes ----------------------- solve_VariationalStokes!( stokes, pt_stokes, @@ -221,12 +202,12 @@ function main(igg, nx, ny) viscosity_cutoff = (-Inf, Inf), ) ) - dt = compute_dt(stokes, di) / 2 + # dt = compute_dt(stokes, di) / 2 # ------------------------------ # Advection -------------------- # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + advection_MQS!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) # advect particles in memory move_particles!(particles, xvi, particle_args) # check if we need to inject particles @@ -272,7 +253,6 @@ function main(igg, nx, ny) fig save(joinpath(figdir, "$(it).png"), fig) - end end From 191c03488deb1148af139679fe0c0d87b91350ff Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 13 Dec 2024 14:36:40 +0100 Subject: [PATCH 54/77] fixup --- src/phases/topography_correction.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index 6c96e01a0..91a2e91a0 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -37,14 +37,14 @@ function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_co # @index phase[ip, icell, j] = T(air_phase) @index coords[1][ip, icell, j] = NaN @index coords[2][ip, icell, j] = NaN - @index index[ip, icell, j] = 0 + @index index[ip, icell, j] = false end # if the particle is above the surface and the phase is air, set the phase to the closes rock phase if !above && phaseq == air_phase # @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) @index coords[1][ip, icell, j] = NaN @index coords[2][ip, icell, j] = NaN - @index index[ip, icell, j] = 0 + @index index[ip, icell, j] = false end end end From 55a3c969d05802d4220ad50f7c6a0713658a5e55 Mon Sep 17 00:00:00 2001 From: aelligp Date: Fri, 13 Dec 2024 16:43:54 +0100 Subject: [PATCH 55/77] Variational stokes shearband miniapp --- .../shear_band/ShearBand2D_variational.jl | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl new file mode 100644 index 000000000..e6cd5e5d7 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl @@ -0,0 +1,206 @@ +using GeoParams, CairoMakie, CellArrays +using JustRelax, JustRelax.JustRelax2D +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) + +const backend = CPUBackend + +using JustPIC, JustPIC._2D +const backend_JP = JustPIC.CPUBackend + +# HELPER FUNCTIONS ----------------------------------- ---------------------------- +solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) + +# Initialize phases on the particles +function init_phases!(phase_ratios, xci, xvi, radius) + ni = size(phase_ratios.center) + origin = 0.5, 0.5 + + @parallel_indices (i, j) function init_phases!(phases, xc, yc, o_x, o_y, radius) + x, y = xc[i], yc[j] + if ((x-o_x)^2 + (y-o_y)^2) > radius^2 + @index phases[1, i, j] = 1.0 + @index phases[2, i, j] = 0.0 + + else + @index phases[1, i, j] = 0.0 + @index phases[2, i, j] = 1.0 + end + return nothing + end + + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni.+1) init_phases!(phase_ratios.vertex, xvi..., origin..., radius) + return nothing +end + +# MAIN SCRIPT -------------------------------------------------------------------- +function main(igg; nx=64, ny=64, figdir="model_figs") + + # Physical domain ------------------------------------ + ly = 1e0 # domain length in y + lx = ly # 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, 0.0 # origin coordinates + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + dt = Inf + + # Physical properties using GeoParams ---------------- + τ_y = 1.6 # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) + ϕ = 30 # friction angle + C = τ_y # Cohesion + η0 = 1.0 # viscosity + G0 = 1.0 # elastic shear modulus + Gi = G0/(6.0-4.0) # elastic shear modulus perturbation + εbg = 1.0 # background strain-rate + η_reg = 8e-3 # regularisation "viscosity" + dt = η0/G0/4.0 # assumes Maxwell time of 4 + el_bg = ConstantElasticity(; G=G0, Kb=4) + el_inc = ConstantElasticity(; G=Gi, Kb=4) + visc = LinearViscous(; η=η0) + pl = DruckerPrager_regularised(; # non-regularized plasticity + C = C, + ϕ = ϕ, + η_vp = η_reg, + Ψ = 15 + ) + + rheology = ( + # Low density phase + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + # CompositeRheology = CompositeRheology((visc, el_bg, )), + CompositeRheology = CompositeRheology((visc, el_bg, pl)), + Elasticity = el_bg, + + ), + # High density phase + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + # CompositeRheology = CompositeRheology((visc, el_inc, )), + CompositeRheology = CompositeRheology((visc, el_inc, pl)), + Elasticity = el_inc, + ), + ) + + # perturbation array for the cohesion + perturbation_C = @rand(ni...) + + # Initialize phase ratios ------------------------------- + radius = 0.1 + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) + init_phases!(phase_ratios, xci, xvi, radius) + air_phase = 2 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.95 / √2.1) + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, Re=3e0, r=0.7, CFL = 0.95 / √2.1) + + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt, perturbation_C = perturbation_C) + + # Rheology + compute_viscosity!( + stokes, phase_ratios, args, rheology, air_phase, (-Inf, Inf) + ) + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + no_slip = (left = false, right = false, top = false, bot=false), + ) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + take(figdir) + + # Time loop + t, it = 0.0, 0 + tmax = 5 + τII = Float64[] + sol = Float64[] + ttot = Float64[] + + # while t < tmax + for _ in 1:15 + + # Stokes solver ---------------- + iters = solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 50e3, + nout = 2e3, + viscosity_cutoff = (-Inf, Inf), + ) + ) + tensor_invariant!(stokes.ε) + push!(τII, maximum(stokes.τ.xx)) + + it += 1 + t += dt + + push!(sol, solution(εbg, t, G0, η0)) + push!(ttot, t) + + println("it = $it; t = $t \n") + + # visualisation + th = 0:pi/50:3*pi; + xunit = @. radius * cos(th) + 0.5; + yunit = @. radius * sin(th) + 0.5; + + fig = Figure(size = (1600, 1600), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = 1, title = L"\tau_{II}", titlesize=35) + # ax2 = Axis(fig[2,1], aspect = 1, title = "η_vep") + ax2 = Axis(fig[2,1], aspect = 1, title = L"E_{II}", titlesize=35) + ax3 = Axis(fig[1,2], aspect = 1, title = L"\log_{10}(\varepsilon_{II})", titlesize=35) + ax4 = Axis(fig[2,2], aspect = 1) + heatmap!(ax1, xci..., Array(stokes.τ.II) , colormap=:batlow) + # heatmap!(ax2, xci..., Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) + heatmap!(ax2, xci..., Array(log10.(stokes.EII_pl)) , colormap=:batlow) + heatmap!(ax3, xci..., Array(log10.(stokes.ε.II)) , colormap=:batlow) + lines!(ax2, xunit, yunit, color = :black, linewidth = 5) + lines!(ax4, ttot, τII, color = :black) + lines!(ax4, ttot, sol, color = :red) + hidexdecorations!(ax1) + hidexdecorations!(ax3) + save(joinpath(figdir, "$(it).png"), fig) + + end + + return nothing +end + +n = 128 +nx = n +ny = n +figdir = "Variational_ShearBands2D" +igg = if !(JustRelax.MPI.Initialized()) + IGG(init_global_grid(nx, ny, 1; init_MPI = true)...) +else + igg +end +main(igg; figdir = figdir, nx = nx, ny = ny); From 1d606f1b4b098a02acc6433cc55959c33f3a7b56 Mon Sep 17 00:00:00 2001 From: aelligp Date: Fri, 13 Dec 2024 17:30:31 +0100 Subject: [PATCH 56/77] fix stokes solve for non variational; typos; format --- _typos.toml | 1 + .../stokes2D/shear_band/ShearBand2D.jl | 2 +- src/boundaryconditions/Dirichlet.jl | 4 +- src/ext/AMDGPU/2D.jl | 11 ++++- src/ext/CUDA/2D.jl | 11 ++++- src/mask/mask.jl | 23 ++++----- src/phases/topography_correction.jl | 47 ++++++++++++------- src/stokes/Stokes2D.jl | 4 +- src/stokes/Stokes3D.jl | 6 +-- src/variational_stokes/MiniKernels.jl | 3 +- src/variational_stokes/mask.jl | 20 ++++---- test/test_Blankenbach.jl | 4 +- test/test_VanKeken.jl | 2 +- test/test_WENO5.jl | 4 +- test/test_shearband2D.jl | 2 +- test/test_shearband2D_MPI.jl | 2 +- test/test_shearband2D_softening.jl | 2 +- test/test_shearband3D_MPI.jl | 2 +- test/test_shearheating2D.jl | 4 +- test/test_shearheating3D.jl | 2 +- test/test_sinking_block.jl | 2 +- test/test_thermalstresses.jl | 6 +-- 22 files changed, 92 insertions(+), 72 deletions(-) diff --git a/_typos.toml b/_typos.toml index ac343ee72..434c58525 100644 --- a/_typos.toml +++ b/_typos.toml @@ -2,3 +2,4 @@ Ths = "Ths" oly = "oly" iy = "iy" +pn = "pn" diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 24d575f63..72af21f3a 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -110,7 +110,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Rheology compute_viscosity!( - stokes, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, 0, (-Inf, Inf) ) # Boundary conditions flow_bcs = VelocityBoundaryConditions(; diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl index 14dbd2513..ba981c10d 100644 --- a/src/boundaryconditions/Dirichlet.jl +++ b/src/boundaryconditions/Dirichlet.jl @@ -27,14 +27,14 @@ end # struct ConstantArray{T,N} <: AbstractArray{T,N} # val::T - + # ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val) # end # Adapt.@adapt_structure ConstantArray struct ConstantArray{T} val::T - + ConstantArray(val::T) where {T<:Number} = new{T}(val) end Adapt.@adapt_structure ConstantArray diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 31f0138d0..5fae85403 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -446,8 +446,15 @@ end # marker chain -function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{JustPIC.AMDGPUBackend}, particles::Particles{JustPIC.AMDGPUBackend}, origin, di, air_phase) - update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) +function JR2D.update_phases_given_markerchain!( + phase, + chain::MarkerChain{JustPIC.AMDGPUBackend}, + particles::Particles{JustPIC.AMDGPUBackend}, + origin, + di, + air_phase, +) + return update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) end end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index e3e15382e..f174c2718 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -427,8 +427,15 @@ end # marker chain -function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{CUDABackend}, particles::Particles{CUDABackend}, origin, di, air_phase) - update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) +function JR2D.update_phases_given_markerchain!( + phase, + chain::MarkerChain{CUDABackend}, + particles::Particles{CUDABackend}, + origin, + di, + air_phase, +) + return update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase) end end diff --git a/src/mask/mask.jl b/src/mask/mask.jl index 8370a6f41..3ffca7a86 100644 --- a/src/mask/mask.jl +++ b/src/mask/mask.jl @@ -35,21 +35,16 @@ Base.similar(m::Mask) = Mask(size(m)...) @inline dims(::Mask{A}) where {A<:AbstractArray{T,N}} where {T,N} = N -@inline apply_mask!(A::AbstractArray, B::Any, m::Mask) = - (A .= inv(m) .* A .+ m.mask .* B) +@inline apply_mask!(A::AbstractArray, B::Any, m::Mask) = (A .= inv(m) .* A .+ m.mask .* B) @inline apply_mask!(::AbstractArray, ::Any, ::Nothing) = nothing -@inline apply_mask!( - A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N} -) where {N} = (A[inds...] = inv(m, inds...) * A[inds...] + m[inds...] * B[inds...]) -@inline apply_mask!( - ::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N} -) where {N} = nothing +@inline apply_mask!(A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}) where {N} = + (A[inds...] = inv(m, inds...) * A[inds...] + m[inds...] * B[inds...]) +@inline apply_mask!(::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}) where {N} = + nothing @inline apply_mask(A::AbstractArray, B::Any, m::Mask) = inv(m) .* A .+ m.mask .* B -@inline apply_mask( - A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N} -) where {N} = @muladd inv(m, inds...) * A[inds...] + m[inds...] * B[inds...] +@inline apply_mask(A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}) where {N} = + @muladd inv(m, inds...) * A[inds...] + m[inds...] * B[inds...] @inline apply_mask(A::AbstractArray, ::Any, ::Nothing) = A -@inline apply_mask( - A::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N} -) where {N} = A[inds...] +@inline apply_mask(A::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}) where {N} = + A[inds...] diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index 91a2e91a0..3155401e7 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -1,25 +1,34 @@ import JustPIC._2D: cell_index, interp1D_inner, interp1D_extremas, distance using StaticArrays - -function update_phases_given_markerchain!(phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase) where {backend} - (; coords, index) = particles; +function update_phases_given_markerchain!( + phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase +) where {backend} + (; coords, index) = particles dy = di[2] - @parallel (1:size(index,1)) _update_phases_given_markerchain!(phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase) + @parallel (1:size(index, 1)) _update_phases_given_markerchain!( + phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase + ) end -@parallel_indices (icell) function _update_phases_given_markerchain!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase) - _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords,cell_vertices, origin, dy, air_phase, icell) +@parallel_indices (icell) function _update_phases_given_markerchain!( + phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase +) + _update_phases_given_markerchain_kernel!( + phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell + ) return nothing end -function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell) - T = eltype(eltype(phase)) - chain_yi = @cell chain_coords[2][icell] - min_cell_j, max_cell_j = find_minmax_cell_indices(chain_yi, origin[2], dy) - min_cell_j = max(1, min_cell_j - 10) - max_cell_j = min(size(index, 2), max_cell_j + 10) - cell_range = min_cell_j:max_cell_j +function _update_phases_given_markerchain_kernel!( + phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell +) + T = eltype(eltype(phase)) + chain_yi = @cell chain_coords[2][icell] + min_cell_j, max_cell_j = find_minmax_cell_indices(chain_yi, origin[2], dy) + min_cell_j = max(1, min_cell_j - 10) + max_cell_j = min(size(index, 2), max_cell_j + 10) + cell_range = min_cell_j:max_cell_j # iterate over cells with marker chain on them for j in cell_range @@ -58,7 +67,7 @@ function extrema_CA(x::AbstractArray) max_val = x[1] min_val = x[1] for i in 2:length(x) - xᵢ = x[i] + xᵢ = x[i] isnan(xᵢ) && continue if xᵢ > max_val max_val = xᵢ @@ -89,10 +98,12 @@ end end # find closest phase different than the given `skip_phase` -function closest_phase(coords, pn, index, current_particle, phases, skip_phase, I::Vararg{Int, N}) where N +function closest_phase( + coords, pn, index, current_particle, phases, skip_phase, I::Vararg{Int,N} +) where {N} new_phase = @index phases[current_particle, I...] - dist_min = Inf - px, py = coords + dist_min = Inf + px, py = coords for ip in cellaxes(index) # early escape conditions @@ -108,7 +119,7 @@ function closest_phase(coords, pn, index, current_particle, phases, skip_phase, # update the closest phase if d < dist_min new_phase = phaseᵢ - dist_min = d + dist_min = d end end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index ebe118cbd..1071e81e5 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -528,7 +528,7 @@ function _solve!( # compute buoyancy forces and viscosity compute_ρg!(ρg, phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff) displacement2velocity!(stokes, dt, flow_bcs) while iter ≤ iterMax @@ -547,7 +547,7 @@ function _solve!( stokes.∇V, ητ, rheology, - phase_ratios.center, + phase_ratios, dt, r, θ_dτ, diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 1053d0341..c75af2943 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -211,7 +211,7 @@ function _solve!( # compute buoyancy forces and viscosity compute_ρg!(ρg, phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff) # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) @@ -425,7 +425,7 @@ function _solve!( # compute buoyancy forces and viscosity compute_ρg!(ρg, phase_ratios, rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff) # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) @@ -444,7 +444,7 @@ function _solve!( stokes.∇V, ητ, rheology, - phase_ratios.center, + phase_ratios, dt, pt_stokes.r, pt_stokes.θ_dτ, diff --git a/src/variational_stokes/MiniKernels.jl b/src/variational_stokes/MiniKernels.jl index 7d0c13ebe..7bfbe4a7e 100644 --- a/src/variational_stokes/MiniKernels.jl +++ b/src/variational_stokes/MiniKernels.jl @@ -71,7 +71,8 @@ Base.@propagate_inbounds @inline _av_zi(A::T, ϕ::T, I::Vararg{Integer,3}) where (top(A, ϕ, I...) + next(A, ϕ, I...)) * 0.5 ## Because mymaskedsum(::generator) does not work inside CUDA kernels... -@inline mymaskedsum(A::AbstractArray, ϕ::AbstractArray, ranges::Vararg{T,N}) where {T,N} = mymaskedsum(identity, A, ϕ, ranges...) +@inline mymaskedsum(A::AbstractArray, ϕ::AbstractArray, ranges::Vararg{T,N}) where {T,N} = + mymaskedsum(identity, A, ϕ, ranges...) @inline function mymaskedsum( f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index 69bae8b65..e74179f00 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -45,9 +45,7 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`. - `phase_ratios`: The ratios of different phases present. - `air_phase`: The phase representing air. """ -function update_rock_ratio!( - ϕ::JustRelax.RockRatio{T, 2}, phase_ratios, air_phase -) where {T} +function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,2}, phase_ratios, air_phase) where {T} nvi = size_v(ϕ) @parallel (@idx nvi) update_rock_ratio_cv!( ϕ, phase_ratios.center, phase_ratios.vertex, air_phase @@ -63,17 +61,17 @@ function update_rock_ratio!( return nothing end -function update_rock_ratio!( - ϕ::JustRelax.RockRatio{T, 3}, phase_ratios, air_phase -) where {T} +function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,3}, phase_ratios, air_phase) where {T} nvi = size_v(ϕ) @parallel (@idx nvi) update_rock_ratio_cv!( ϕ, phase_ratios.center, phase_ratios.vertex, air_phase ) dst = ϕ.Vx, ϕ.Vy, ϕ.Vz, ϕ.xy, ϕ.yz, ϕ.xz - src = phase_ratios.Vx, phase_ratios.Vy, phase_ratios.Vz, phase_ratios.xy, phase_ratios.yz, phase_ratios.xz - + src = phase_ratios.Vx, + phase_ratios.Vy, phase_ratios.Vz, phase_ratios.xy, phase_ratios.yz, + phase_ratios.xz + for (dstᵢ, srcᵢ) in zip(dst, src) @parallel (@idx size(dstᵢ)) _update_rock_ratio!(dstᵢ, srcᵢ, air_phase) end @@ -81,9 +79,9 @@ function update_rock_ratio!( return nothing end -@inline compute_rock_ratio( - phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} -) where {N} = (x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5) +@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = ( + x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5 +) @inline compute_air_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} diff --git a/test/test_Blankenbach.jl b/test/test_Blankenbach.jl index b7c53390b..5bb09b4e5 100644 --- a/test/test_Blankenbach.jl +++ b/test/test_Blankenbach.jl @@ -142,7 +142,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10) η = @ones(ni...) compute_ρg!(ρg[2], phase_ratios, rheology, args) compute_viscosity!( - stokes, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, 0, (-Inf, Inf) ) # PT coefficients for thermal diffusion ------------- @@ -182,7 +182,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10) # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)) compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index 08326760a..0e8ea3212 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -127,7 +127,7 @@ function VanKeken2D(ny=32, nx=32) compute_ρg!(ρg[2], phase_ratios, rheology, args) # Rheology - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)) # Boundary conditions flow_bcs = VelocityBoundaryConditions(; diff --git a/test/test_WENO5.jl b/test/test_WENO5.jl index ccd7c2d42..9457c06c3 100644 --- a/test/test_WENO5.jl +++ b/test/test_WENO5.jl @@ -195,7 +195,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, thermal_perturbation = # Rheology viscosity_cutoff = (1e16, 1e24) - compute_viscosity!(stokes, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, args, rheology, 0, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -225,7 +225,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, thermal_perturbation = args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, args, rheology, viscosity_cutoff + stokes, args, rheology, 0, viscosity_cutoff ) # ------------------------------ iters = solve!( diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index f06b5618c..250757706 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -130,7 +130,7 @@ function ShearBand2D() # Rheology compute_viscosity!( - stokes, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, 0, (-Inf, Inf) ) # Boundary conditions flow_bcs = VelocityBoundaryConditions(; diff --git a/test/test_shearband2D_MPI.jl b/test/test_shearband2D_MPI.jl index aa47fc5d7..29f09e13f 100644 --- a/test/test_shearband2D_MPI.jl +++ b/test/test_shearband2D_MPI.jl @@ -126,7 +126,7 @@ function main(igg; nx=64, ny=64) # Rheology compute_viscosity!( - stokes, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, 0, (-Inf, Inf) ) # Boundary conditions diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl index 441091350..95fc14f69 100644 --- a/test/test_shearband2D_softening.jl +++ b/test/test_shearband2D_softening.jl @@ -135,7 +135,7 @@ function ShearBand2D() # Rheology compute_viscosity!( - stokes, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, 0, (-Inf, Inf) ) # Boundary conditions diff --git a/test/test_shearband3D_MPI.jl b/test/test_shearband3D_MPI.jl index 9f818a27c..59c81588a 100644 --- a/test/test_shearband3D_MPI.jl +++ b/test/test_shearband3D_MPI.jl @@ -135,7 +135,7 @@ function main(igg; nx=64, ny=64, nz=64) args = (; T = @zeros(ni...), P = stokes.P, dt = Inf) # Rheology cutoff_visc = -Inf, Inf - compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, cutoff_visc) # Boundary conditions flow_bcs = VelocityBoundaryConditions(; diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index e71f5578c..d2eabcbdb 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -127,7 +127,7 @@ function Shearheating2D(igg; nx=32, ny=32) # Rheology args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -226,7 +226,7 @@ function Shearheating2D(igg; nx=32, ny=32) @views T_buffer[:, end] .= 273.0 + 400 @views thermal.T[2:end-1, :] .= T_buffer temperature2center!(thermal) - + @show it += 1 t += dt diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index 939322e63..c612f1ba8 100644 --- a/test/test_shearheating3D.jl +++ b/test/test_shearheating3D.jl @@ -125,7 +125,7 @@ function Shearheating3D(igg; nx=16, ny=16, nz=16) # Rheology args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0,(-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index 50a823fad..cda576b37 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -160,7 +160,7 @@ function Sinking_Block2D() # Viscosity args = (; T = @ones(ni...), P = stokes.P, dt=Inf) η_cutoff = -Inf, Inf - compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions diff --git a/test/test_thermalstresses.jl b/test/test_thermalstresses.jl index 56a9099b1..a5e69d577 100644 --- a/test/test_thermalstresses.jl +++ b/test/test_thermalstresses.jl @@ -308,7 +308,7 @@ function main2D(igg; nx=32, ny=32) ) flow_bcs!(stokes, flow_bcs) - compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, cutoff_visc) ϕ = @zeros(ni...) compute_melt_fraction!( @@ -347,7 +347,7 @@ function main2D(igg; nx=32, ny=32) # Update buoyancy and viscosity - args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, 0, cutoff_visc) # Stokes solver ----------------- solve!( @@ -456,4 +456,4 @@ end @test Array(thermal.T)[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 0.5369 rtol = 1e-2 @test Array(ϕ)[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 9.351e-9 rtol = 1e-1 end -end \ No newline at end of file +end From 6ac498b9bb07a236ddc3d4e2b36037c69d13fd47 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 14 Dec 2024 14:38:03 +0100 Subject: [PATCH 57/77] add masking checks --- src/variational_stokes/mask.jl | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index e74179f00..cda4fac9f 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -79,21 +79,27 @@ function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,3}, phase_ratios, air_phas return nothing end -@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = ( - x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5 -) +@inline function compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} + 1 ≤ air_phase ≤ numphases(phase_ratio) || return 1e0 + x = 1 - @index phase_ratio[air_phase, I...] + x *= x > 1e-5 + return x +end -@inline compute_air_ratio( +@inline function compute_air_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} -) where {N} = @index phase_ratio[air_phase, I...] +) where {N} + 1 ≤ air_phase ≤ numphases(phase_ratio) || return 1e0 + return @index phase_ratio[air_phase, I...] +end @parallel_indices (I...) function update_rock_ratio_cv!( ϕ, ratio_center, ratio_vertex, air_phase ) if all(I .≤ size(ratio_center)) - ϕ.center[I...] = Float64(Float16(compute_rock_ratio(ratio_center, air_phase, I...))) + ϕ.center[I...] = compute_rock_ratio(ratio_center, air_phase, I...) end - ϕ.vertex[I...] = Float64(Float16(compute_rock_ratio(ratio_vertex, air_phase, I...))) + ϕ.vertex[I...] = compute_rock_ratio(ratio_vertex, air_phase, I...) return nothing end @@ -364,3 +370,4 @@ Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer,N}) where {N} = # cv = c || v # return cv || isvalid(ϕ.Vy, i, j) # end + From c30ab31a075a05a46d6d3e01ad29e40c6b3dac15 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 16 Dec 2024 10:31:33 +0100 Subject: [PATCH 58/77] up miniapp --- .../stokes2D/shear_band/ShearBand2D_variational.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl index e6cd5e5d7..af84e740e 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational.jl @@ -1,4 +1,4 @@ -using GeoParams, CairoMakie, CellArrays +using GeoParams, GLMakie, CellArrays using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) @@ -81,7 +81,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") ), # High density phase SetMaterialParams(; - Phase = 2, + Phase = 1, Density = ConstantDensity(; ρ = 0.0), Gravity = ConstantGravity(; g = 0.0), # CompositeRheology = CompositeRheology((visc, el_inc, )), @@ -97,7 +97,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") radius = 0.1 phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) init_phases!(phase_ratios, xci, xvi, radius) - air_phase = 2 + air_phase = 0 ϕ = RockRatio(backend, ni) update_rock_ratio!(ϕ, phase_ratios, air_phase) # STOKES --------------------------------------------- From f09df9298262284b32b71aed9ecc956623450f66 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 16 Dec 2024 13:12:32 +0100 Subject: [PATCH 59/77] couple of Dirichlet fixes --- src/boundaryconditions/Dirichlet.jl | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl index ba981c10d..357567eb7 100644 --- a/src/boundaryconditions/Dirichlet.jl +++ b/src/boundaryconditions/Dirichlet.jl @@ -15,7 +15,7 @@ end Adapt.@adapt_structure DirichletBoundaryCondition function DirichletBoundaryCondition(A::AbstractArray{T}) where {T} - m = Mask(size(A)...) + m = Mask(A) copyto!(m.mask, T.(A .!= 0)) return DirichletBoundaryCondition(A, m) end @@ -25,13 +25,6 @@ function Base.getindex(x::DirichletBoundaryCondition, inds::Vararg{Int,N}) where return x.value[inds...] * x.mask[inds...] end -# struct ConstantArray{T,N} <: AbstractArray{T,N} -# val::T - -# ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val) -# end -# Adapt.@adapt_structure ConstantArray - struct ConstantArray{T} val::T From 90308f406914d5ae9e3048bbc6697576df7ad467 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 16 Dec 2024 13:31:19 +0100 Subject: [PATCH 60/77] another Dirichlet fix --- src/boundaryconditions/Dirichlet.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl index 357567eb7..28c9b1e73 100644 --- a/src/boundaryconditions/Dirichlet.jl +++ b/src/boundaryconditions/Dirichlet.jl @@ -15,8 +15,9 @@ end Adapt.@adapt_structure DirichletBoundaryCondition function DirichletBoundaryCondition(A::AbstractArray{T}) where {T} - m = Mask(A) - copyto!(m.mask, T.(A .!= 0)) + m = Mask(copy(A)) + copyto!(m.mask, T.(.!iszero.(A))) + return DirichletBoundaryCondition(A, m) end From 6ccbc1e24e5dd69039690106a21128e08d745889 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 16 Dec 2024 13:32:28 +0100 Subject: [PATCH 61/77] add inner T BCs also on the air --- .../stokes2D/Volcano2D/Caldera2D.jl | 57 ++++++++++++------- 1 file changed, 35 insertions(+), 22 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index ee4e41d3e..e76fb3e3b 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -1,5 +1,5 @@ -# const isCUDA = false -const isCUDA = true +const isCUDA = false +# const isCUDA = true @static if isCUDA using CUDA @@ -92,17 +92,20 @@ function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) return topo_y end -function thermal_anomaly!(Temp, mask, Ω_T, phase_ratios, conduit_phase, magma_phase) +function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_phase, magma_phase, air_phase) - @parallel_indices (i, j) function _thermal_anomaly!(Temp, mask, Ω_T, vertex_ratio, conduit_phase, magma_phase) + @parallel_indices (i, j) function _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, vertex_ratio, conduit_phase, magma_phase, air_phase) # quick escape conduit_ratio_ij = @index vertex_ratio[conduit_phase, i, j] magma_ratio_ij = @index vertex_ratio[magma_phase, i, j] + air_ratio_ij = @index vertex_ratio[air_phase, i, j] if conduit_ratio_ij > 0.5 || magma_ratio_ij > 0.5 # if isone(conduit_ratio_ij) || isone(magma_ratio_ij) - Temp[i+1, j] = Ω_T - mask[i+1, j] = 1 + Ω_T[i+1, j] = Temp[i+1, j] = T_chamber + + elseif air_ratio_ij > 0.5 + Ω_T[i+1, j] = Temp[i+1, j] = T_air end return nothing @@ -110,12 +113,16 @@ function thermal_anomaly!(Temp, mask, Ω_T, phase_ratios, conduit_phase, magma_p ni = size(phase_ratios.vertex) - @parallel (@idx ni) _thermal_anomaly!(Temp, mask, Ω_T, phase_ratios.vertex, conduit_phase, magma_phase) + @parallel (@idx ni) _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, phase_ratios.vertex, conduit_phase, magma_phase, air_phase) + + @views Ω_T[1, :] .= Ω_T[2, :] + @views Ω_T[end, :] .= Ω_T[end-1, :] + @views Temp[1, :] .= Temp[2, :] + @views Temp[end, :] .= Temp[end-1, :] return nothing end - ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) @@ -177,7 +184,9 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # rock ratios for variational stokes # RockRatios ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + # update_rock_ratio!(ϕ, phase_ratios, air_phase) + compute_rock_fraction!(ϕ, chain, xvi, di) + # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -191,12 +200,15 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) # Add thermal anomaly BC's - Ω_T = 1223e0 # inner BCs temperature - mask = @zeros(size(thermal.T)...) - thermal_anomaly!(thermal.T, mask, Ω_T, phase_ratios, 5, 3) + T_chamber = 1223e0 + T_air = 273e0 + Ω_T = @zeros(size(thermal.T)...) + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + JustRelax.DirichletBoundaryCondition(Ω_T) + thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true, right = true, top = false, bot = false), - dirichlet = (constant = Ω_T, mask=mask) + no_flux = (; left = true, right = true, top = false, bot = false), + dirichlet = (; mask = Ω_T) ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -307,9 +319,9 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end-1, :] .= T_buffer - if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 - thermal_anomaly!(thermal.T, mask, Ω_T, phase_ratios, 5, 3) - end + # if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + # end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -407,13 +419,14 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + # update_rock_ratio!(ϕ, phase_ratios, air_phase) + compute_rock_fraction!(ϕ, chain, xvi, di) tensor_invariant!(stokes.τ) # track deformation of free_surface - push!(deformation_x, chain.coords[1].data[:]./1e3) - push!(deformation_y, chain.coords[2].data[:]./1e3) + # push!(deformation_x, chain.coords[1].data[:]./1e3) + # push!(deformation_y, chain.coords[2].data[:]./1e3) @show it += 1 t += dt if plotting @@ -553,7 +566,7 @@ const plotting = true do_vtk = true # set to true to generate VTK files for ParaView # figdir = "Caldera2D_noPguess" figdir = "$(today())_Conduit" -n = 128 +n = 64 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; @@ -575,7 +588,7 @@ else igg end -# main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); +main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); # function plot_particles(particles, pPhases, chain) # p = particles.coords From 3f271439b58e6d6482451b0edb98a63d769ff63a Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 16 Dec 2024 14:10:59 +0100 Subject: [PATCH 62/77] upa miniapp again --- .../stokes2D/Volcano2D/Caldera2D.jl | 35 +++++++++++++++---- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index e76fb3e3b..d1d7b745b 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -104,7 +104,7 @@ function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_ph # if isone(conduit_ratio_ij) || isone(magma_ratio_ij) Ω_T[i+1, j] = Temp[i+1, j] = T_chamber - elseif air_ratio_ij > 0.5 + elseif air_ratio_ij > 0.5 Ω_T[i+1, j] = Temp[i+1, j] = T_air end @@ -114,7 +114,7 @@ function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_ph ni = size(phase_ratios.vertex) @parallel (@idx ni) _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, phase_ratios.vertex, conduit_phase, magma_phase, air_phase) - + @views Ω_T[1, :] .= Ω_T[2, :] @views Ω_T[end, :] .= Ω_T[end-1, :] @views Temp[1, :] .= Temp[2, :] @@ -205,7 +205,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", Ω_T = @zeros(size(thermal.T)...) thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) JustRelax.DirichletBoundaryCondition(Ω_T) - + thermal_bc = TemperatureBoundaryConditions(; no_flux = (; left = true, right = true, top = false, bot = false), dirichlet = (; mask = Ω_T) @@ -486,27 +486,38 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") # ax4 = Axis(fig[2,3], aspect = ar, title = "log10(εII)") ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + ax5 = Axis(fig[3,1], aspect = ar, title = "EII_pl") + ax6 = Axis(fig[3,3], aspect = ar, title = "Melt fraction ϕ") # Plot temperature h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) # Plot particles phase - # h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 3) - # scatter!(ax2, Array(chain_x), Array(chain_y), color=:red, markersize = 3) h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, uconvert.(u"cm/yr",Array(stokes.V.Vy)u"m/s") , colormap=:batlow) + scatter!(ax2, Array(chain_x), Array(chain_y), color=:red, markersize = 3) + # Plot 2nd invariant of strain rate # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε_pl.II)) , colormap=:batlow) h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) # Plot effective viscosity # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:lipari) h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + h5 = heatmap!(ax5, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.EII_pl) , colormap=:batlow) + h6 = heatmap!(ax6, xci[1].*1e-3, xci[2].*1e-3, Array(ϕ_m) , colormap=:lipari) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) + hidexdecorations!(ax4) + hideydecorations!(ax3) + hideydecorations!(ax4) + hideydecorations!(ax6) + Colorbar(fig[1,2], h1) Colorbar(fig[2,2], h2) Colorbar(fig[1,4], h3) Colorbar(fig[2,4], h4) - linkaxes!(ax1, ax2, ax3, ax4) + Colorbar(fig[3,2], h5) + # Colorbar(fig[3,4], h6) + linkaxes!(ax1, ax2, ax3, ax4, ax5) fig save(joinpath(figdir, "$(it).png"), fig) @@ -549,8 +560,18 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", Y./1e3, ) hideydecorations!(ax2) - # save(joinpath(figdir, "thermal_profile_$it.png"), fig) + save(joinpath(figdir, "thermal_profile_$it.png"), fig) + fig + end + + fig1 = let + fig = Figure(; size=(1200, 900)) + ax = Axis(fig[1, 1];title="Drucker Prager") + lines!(ax,[0e6, maximum(stokes.P)]./1e6,[10e6; (maximum(stokes.P)*sind(rheology[1].CompositeRheology[1].elements[end].ϕ.val) - rheology[1].CompositeRheology[1].elements[end].C.val*cosd(rheology[1].CompositeRheology[1].elements[end].ϕ.val))]./1e6, color=:black, linewidth=2) + s1=scatter!(ax, Array(stokes.P./1e6)[:], Array(stokes.τ.II./1e6)[:]; color=Array(stokes.R.RP)[:], colormap=:roma, markersize=3) + Colorbar(fig[1,2], s1) fig + save(joinpath(figdir, "DruckerPrager_$it.png"), fig) end end # ------------------------------ From cc128a0f9e721cbac8cee8bc1debd423174ed1e2 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 16 Dec 2024 15:28:30 +0100 Subject: [PATCH 63/77] up miniapp caldera --- .../stokes2D/Volcano2D/Caldera2D.jl | 35 +++++-------------- .../stokes2D/Volcano2D/Caldera_rheology.jl | 7 ++-- .../stokes2D/Volcano2D/Caldera_setup.jl | 4 +-- 3 files changed, 15 insertions(+), 31 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index d1d7b745b..51b71033f 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -124,7 +124,7 @@ function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_ph end ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) +function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false, extension = 1e-15 * 0) # Physical domain ------------------------------------ ni = nx, ny # number of cells @@ -249,7 +249,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # flow_bcs!(stokes, flow_bcs) # apply boundary conditions # displacement2velocity!(stokes, dt) - εbg = 1e-15 * 1 + εbg = extension apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -583,12 +583,12 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- const plotting = true - +extension = 1e-15 * 1 do_vtk = true # set to true to generate VTK files for ParaView -# figdir = "Caldera2D_noPguess" -figdir = "$(today())_Conduit" +figdir = "Caldera2D" n = 64 nx, ny = n, n >>> 1 + li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; sticky_air = 4e0, @@ -596,11 +596,11 @@ li, origin, phases_GMG, T_GMG = setup2D( flat = false, chimney = true, volcano_size = (3e0, 5e0), - conduit_radius = 0.5e0, + conduit_radius = 2e-1, chamber_T = 900e0, chamber_depth = 7e0, - chamber_radius = 1e0, - aspect_x = 6, + chamber_radius = 1.25, + aspect_x = 2, ) igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -609,21 +609,4 @@ else igg end -main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); - -# function plot_particles(particles, pPhases, chain) -# p = particles.coords -# # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter -# ppx, ppy = p -# # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) -# # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) -# pxv = ppx.data[:] -# pyv = ppy.data[:] -# clr = pPhases.data[:] -# # clr = pϕ.data[:] -# idxv = particles.index.data[:] -# f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma, markersize=1) - -# Colorbar(f[1,2], h) -# f -# end +main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk, extension = extension); diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl index 7bb80e260..6aa020ff9 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl @@ -46,16 +46,17 @@ function init_rheologies(; linear=false) η_reg = 1e18 C = linear ? Inf : 10e6 ϕ = 15 - soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 2) # nonlinear softening law + soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 1e5) # nonlinear softening law # soft_C = NonLinearSoftening() # nonlinear softening law - pl = DruckerPrager_regularised(; C=C, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C=soft_C) + pl = DruckerPrager_regularised(; C=C*MPa, ϕ=ϕ, η_vp=(η_reg)*Pas, Ψ=0.0, softening_C=soft_C) el = ConstantElasticity(; G = 25e9, ν = 0.45) β = 1 / el.Kb.val Cp = 1200.0 magma_visc = ViscosityPartialMelt_Costa_etal_2009(η=LinearMeltViscosity(A = -8.1590, B = 2.4050e+04K, T0 = -430.9606K,η0=1e3Pa*s)) #dislocation laws - disl_top = SetDislocationCreep(Dislocation.dry_olivine_Karato_2003) + disl_top = DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) + # disl_top = SetDislocationCreep(Dislocation.dry_olivine_Karato_2003) # diffusion laws disl_bot = SetDislocationCreep(Dislocation.wet_quartzite_Hirth_2001) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl index d88f10878..d4676c40e 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl @@ -80,8 +80,8 @@ function setup2D( ph = Phases[:,1,:] T = Temp[:,1,:] .+ 273 - V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * 1.0 - printstyled("Magma volume of the initial chamber: $(round(V; digits=3)) km³ \n"; bold=true, color=:red) + V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * (chamber_radius*aspect_x) + printstyled("Magma volume of the initial chamber: $(round(V; digits=3)) km³ \n"; bold=true, color=:red, blink=true) # write_paraview(Grid, "Volcano2D") return li, origin, ph, T, Grid end From 1e9ee4b9807baaf9041b46cf398ea38ae4468b21 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 16 Dec 2024 15:33:35 +0100 Subject: [PATCH 64/77] add cyclicity to T BC's; update test --- .../stokes2D/Volcano2D/Caldera2D.jl | 6 +- test/test_Volcano2D.jl | 138 +++++++++++++----- 2 files changed, 108 insertions(+), 36 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 51b71033f..7981dc03a 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -319,9 +319,9 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end-1, :] .= T_buffer - # if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 - thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) - # end + if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) diff --git a/test/test_Volcano2D.jl b/test/test_Volcano2D.jl index f5654b1ce..f25801df2 100644 --- a/test/test_Volcano2D.jl +++ b/test/test_Volcano2D.jl @@ -82,28 +82,66 @@ function apply_pure_shear(Vx,Vy, εbg, xvi, lx, ly) return nothing end + +function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) + topo_idx = [findfirst(x->x==air_phase, row) - 1 for row in eachrow(phases_GMG)] + yv = xvi[2] + topo_y = yv[topo_idx] + return topo_y +end + +function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_phase, magma_phase, air_phase) + + @parallel_indices (i, j) function _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, vertex_ratio, conduit_phase, magma_phase, air_phase) + # quick escape + conduit_ratio_ij = @index vertex_ratio[conduit_phase, i, j] + magma_ratio_ij = @index vertex_ratio[magma_phase, i, j] + air_ratio_ij = @index vertex_ratio[air_phase, i, j] + + if conduit_ratio_ij > 0.5 || magma_ratio_ij > 0.5 + # if isone(conduit_ratio_ij) || isone(magma_ratio_ij) + Ω_T[i+1, j] = Temp[i+1, j] = T_chamber + + elseif air_ratio_ij > 0.5 + Ω_T[i+1, j] = Temp[i+1, j] = T_air + end + + return nothing + end + + ni = size(phase_ratios.vertex) + + @parallel (@idx ni) _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, phase_ratios.vertex, conduit_phase, magma_phase, air_phase) + + @views Ω_T[1, :] .= Ω_T[2, :] + @views Ω_T[end, :] .= Ω_T[end-1, :] + @views Temp[1, :] .= Temp[2, :] + @views Temp[end, :] .= Temp[end-1, :] + + return nothing +end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) +function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false, extension = 1e-15 * 0) # Physical domain ------------------------------------ ni = nx, ny # number of cells di = @. li / ni # grid steps grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # Physical properties using GeoParams ---------------- rheology = init_rheologies() - dt = 5e2 * 3600 * 24 * 365 # diffusive CFL timestep limiter + dt = 5e2 * 3600 * 24 * 365 # dt = Inf # diffusive CFL timestep limiter # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell = 40 - max_xcell = 60 - min_xcell = 20 + nxcell = 100 + max_xcell = 150 + min_xcell = 75 particles = init_particles( backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni ) @@ -113,24 +151,36 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # material phase & temperature pPhases, pT = init_cell_arrays(particles, Val(2)) - # particle fields for the stress rotation - pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress - # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress - pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity - particle_args = (pT, pPhases, pτ..., pω...) - particle_args_reduced = (pT, pτ..., pω...) - # Assign particles phases anomaly phases_device = PTArray(backend)(phases_GMG) phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); init_phases!(pPhases, phases_device, particles, xvi) + + # Initialize marker chain + nxcell, max_xcell, min_xcell = 100, 150, 75 + initial_elevation = 0e0 + chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); + air_phase = 6 + topo_y = extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) + for _ in 1:3 + @views hn = 0.5 .* (topo_y[1:end-1] .+ topo_y[2:end]) + @views topo_y[2:end-1] .= 0.5 .* (hn[1:end-1] .+ hn[2:end]) + fill_chain_from_vertices!(chain, PTArray(backend)(topo_y)) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + end update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # particle fields for the stress rotation + pτ = StressParticles(particles) + particle_args = (pT, pPhases, unwrap(pτ)...) + particle_args_reduced = (pT, unwrap(pτ)...) + # rock ratios for variational stokes # RockRatios - air_phase = 5 ϕ = RockRatio(backend, ni) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + # update_rock_ratio!(ϕ, phase_ratios, air_phase) + compute_rock_fraction!(ϕ, chain, xvi, di) + # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -142,8 +192,17 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # TEMPERATURE PROFILE -------------------------------- thermal = ThermalArrays(backend, ni) @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) + + # Add thermal anomaly BC's + T_chamber = 1223e0 + T_air = 273e0 + Ω_T = @zeros(size(thermal.T)...) + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + JustRelax.DirichletBoundaryCondition(Ω_T) + thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true, right = true, top = false, bot = false), + no_flux = (; left = true, right = true, top = false, bot = false), + dirichlet = (; mask = Ω_T) ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -153,9 +212,14 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(2)) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + # Melt fraction + ϕ_m = @zeros(ni...) + compute_melt_fraction!( + ϕ_m, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P) + ) # Rheology args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) viscosity_cutoff = (1e18, 1e23) @@ -179,7 +243,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # flow_bcs!(stokes, flow_bcs) # apply boundary conditions # displacement2velocity!(stokes, dt) - εbg = 1e-15 * 1 + εbg = extension apply_pure_shear(@velocity(stokes)..., εbg, xvi, li...) flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -211,11 +275,16 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end-1, :] .= T_buffer + if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) - args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + args = (; ϕ=ϕ_m, T=thermal.Tc, P=stokes.P, dt=Inf) + + stress2grid!(stokes, pτ, xvi, xci, particles) iters = solve_VariationalStokes!( stokes, @@ -236,18 +305,13 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ) ) - center2vertex!(τxx_v, stokes.τ.xx) - center2vertex!(τyy_v, stokes.τ.yy) - centroid2particle!(pτxx , xci, stokes.τ.xx, particles) - centroid2particle!(pτyy , xci, stokes.τ.yy, particles) - grid2particle!(pτxy, xvi, stokes.τ.xy, particles) - grid2particle!(pωxy, xvi, stokes.ω.xy, particles) - rotate_stress_particles!(pτ, pω, particles, dt) + # rotate stresses + rotate_stress!(pτ, stokes, particles, xci, xvi, dt) tensor_invariant!(stokes.ε) tensor_invariant!(stokes.ε_pl) dtmax = 2e3 * 3600 * 24 * 365.25 - dt = compute_dt(stokes, di, dtmax) + dt = compute_dt(stokes, di, dtmax) * 0.5 println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ @@ -264,7 +328,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", kwargs = ( igg = igg, phase = phase_ratios, - iterMax = 50e3, + iterMax = 100e3, nout = 1e2, verbose = true, ) @@ -289,6 +353,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", move_particles!(particles, xvi, particle_args) # check if we need to inject particles # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + center2vertex!(τxx_v, stokes.τ.xx) + center2vertex!(τyy_v, stokes.τ.yy) inject_particles_phase!( particles, pPhases, @@ -297,13 +363,19 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", xvi ) + # advect marker chain + advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + + compute_melt_fraction!( + ϕ_m, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P) + ) + # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) - update_rock_ratio!(ϕ, phase_ratios, air_phase) + # update_rock_ratio!(ϕ, phase_ratios, air_phase) + compute_rock_fraction!(ϕ, chain, xvi, di) - particle2centroid!(stokes.τ.xx, pτxx, xci, particles) - particle2centroid!(stokes.τ.yy, pτyy, xci, particles) - particle2grid!(stokes.τ.xy, pτxy, xvi, particles) tensor_invariant!(stokes.τ) @show it += 1 @@ -311,7 +383,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", end - return iters + return iters end ## END OF MAIN SCRIPT ---------------------------------------------------------------- From 4ace09b9396bbdd645314136d20c100349c4ec3c Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 23 Dec 2024 10:44:09 +0100 Subject: [PATCH 65/77] variational stokes SB MPI --- .../shear_band/ShearBand2D_variational_MPI.jl | 259 ++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl new file mode 100644 index 000000000..19962fc92 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl @@ -0,0 +1,259 @@ +using GeoParams +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO, CairoMakie +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) + +const backend = CPUBackend + +using JustPIC, JustPIC._2D +const backend_JP = JustPIC.CPUBackend + +# HELPER FUNCTIONS ----------------------------------- ---------------------------- +solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) + +# Initialize phases on the particles +function init_phases!(phase_ratios, xci, xvi, radius) + ni = size(phase_ratios.center) + origin = 0.5, 0.5 + + @parallel_indices (i, j) function init_phases!(phases, xc, yc, o_x, o_y, radius) + x, y = xc[i], yc[j] + if ((x-o_x)^2 + (y-o_y)^2) > radius^2 + @index phases[1, i, j] = 1.0 + @index phases[2, i, j] = 0.0 + + else + @index phases[1, i, j] = 0.0 + @index phases[2, i, j] = 1.0 + end + return nothing + end + + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni.+1) init_phases!(phase_ratios.vertex, xvi..., origin..., radius) + return nothing +end + +# MAIN SCRIPT -------------------------------------------------------------------- +function main(igg; nx=64, ny=64, figdir="model_figs") + + # Physical domain ------------------------------------ + ly = 1e0 # domain length in y + lx = ly # 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, 0.0 # origin coordinates + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + dt = Inf + + # Physical properties using GeoParams ---------------- + τ_y = 1.6 # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) + ϕ = 30 # friction angle + C = τ_y # Cohesion + η0 = 1.0 # viscosity + G0 = 1.0 # elastic shear modulus + Gi = G0/(6.0-4.0) # elastic shear modulus perturbation + εbg = 1.0 # background strain-rate + η_reg = 8e-3 # regularisation "viscosity" + dt = η0/G0/4.0 # assumes Maxwell time of 4 + el_bg = ConstantElasticity(; G=G0, Kb=4) + el_inc = ConstantElasticity(; G=Gi, Kb=4) + visc = LinearViscous(; η=η0) + pl = DruckerPrager_regularised(; # non-regularized plasticity + C = C, + ϕ = ϕ, + η_vp = η_reg, + Ψ = 0.0 + ) + + rheology = ( + # Low density phase + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + # CompositeRheology = CompositeRheology((visc, el_bg, )), + CompositeRheology = CompositeRheology((visc, el_bg, pl)), + Elasticity = el_bg, + + ), + # High density phase + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + # CompositeRheology = CompositeRheology((visc, el_inc, )), + CompositeRheology = CompositeRheology((visc, el_inc, pl)), + Elasticity = el_inc, + ), + ) + + # perturbation array for the cohesion + perturbation_C = @rand(ni...) + + # Initialize phase ratios ------------------------------- + radius = 0.1 + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) + init_phases!(phase_ratios, xci, xvi, radius) + air_phase = 0 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.95 / √2.1) + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, Re=3e0, r=0.7, CFL = 0.95 / √2.1) + + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt, perturbation_C = perturbation_C) + + # Rheology + compute_viscosity!( + stokes, phase_ratios, args, rheology, air_phase, (-Inf, Inf) + ) + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + no_slip = (left = false, right = false, top = false, bot=false), + ) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + take(figdir) + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + + + # global array + nx_v = (nx - 2) * igg.dims[1] + ny_v = (ny - 2) * igg.dims[2] + τII_v = zeros(nx_v, ny_v) + η_vep_v = zeros(nx_v, ny_v) + εII_v = zeros(nx_v, ny_v) + τII_nohalo = zeros(nx-2, ny-2) + η_vep_nohalo = zeros(nx-2, ny-2) + εII_nohalo = zeros(nx-2, ny-2) + Vxv_v = zeros(nx_v, ny_v) + Vyv_v = zeros(nx_v, ny_v) + Vx_nohalo = zeros(nx-2, ny-2) + Vy_nohalo = zeros(nx-2, ny-2) + xci_v = LinRange(0, 1, nx_v) , LinRange(0, 1, ny_v) + + local Vx, Vy + Vx = @zeros(ni...) + Vy = @zeros(ni...) + + # Time loop + t, it = 0.0, 0 + tmax = 5 + τII = Float64[] + sol = Float64[] + ttot = Float64[] + + # while t < tmax + for _ in 1:15 + + # Stokes solver ---------------- + iters = solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 50e3, + nout = 2e3, + viscosity_cutoff = (-Inf, Inf), + ) + ) + tensor_invariant!(stokes.ε) + push!(τII, maximum(stokes.τ.xx)) + + it += 1 + t += dt + + push!(sol, solution(εbg, t, G0, η0)) + push!(ttot, t) + + println("it = $it; t = $t \n") + + # visualisation + th = 0:pi/50:3*pi; + xunit = @. radius * cos(th) + 0.5; + yunit = @. radius * sin(th) + 0.5; + + # Gather MPI arrays + velocity2center!(Vx, Vy, @velocity(stokes)...) + @views Vx_nohalo .= Array(Vx[2:end-1, 2:end-1]) # Copy data to CPU removing the halo + @views Vy_nohalo .= Array(Vy[2:end-1, 2:end-1]) # Copy data to CPU removing the halo + gather!(Vx_nohalo, Vxv_v) + gather!(Vy_nohalo, Vyv_v) + + @views τII_nohalo .= Array(stokes.τ.II[2:end-1, 2:end-1]) # Copy data to CPU removing the halo + @views η_vep_nohalo .= Array(stokes.viscosity.η_vep[2:end-1, 2:end-1]) # Copy data to CPU removing the halo + @views εII_nohalo .= Array(stokes.ε.II[2:end-1, 2:end-1]) # Copy data to CPU removing the halo + gather!(τII_nohalo, τII_v) + gather!(η_vep_nohalo, η_vep_v) + gather!(εII_nohalo, εII_v) + + if igg.me == 0 + fig = Figure(size = (1600, 1600), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = 1, title = "τII") + ax2 = Axis(fig[2,1], aspect = 1, title = "η_vep") + ax3 = Axis(fig[1,3], aspect = 1, title = "log10(εII)") + ax4 = Axis(fig[2,3], aspect = 1) + heatmap!(ax1, xci_v..., Array(τII_v) , colormap=:batlow) + heatmap!(ax2, xci_v..., Array(log10.(η_vep_v)), colormap=:batlow) + heatmap!(ax3, xci_v..., Array(log10.(εII_v)) , colormap=:batlow) + lines!(ax2, xunit, yunit, color = :black, linewidth = 5) + lines!(ax4, ttot, τII, color = :black) + lines!(ax4, ttot, sol, color = :red) + hidexdecorations!(ax1) + hidexdecorations!(ax3) + save(joinpath(figdir, "MPI_$(it).png"), fig) + + data_c = (; + τII = τII_v, + εII = εII_v, + η = η_vep_v, + ) + velocity = ( + Array(Vxv_v), + Array(Vyv_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$(it)_$(igg.me)", 6, "0")), + xci_v, + data_c, + velocity; + t=t + ) + + end + end + + return nothing +end + +n = 64 -2 +nx = n *2 +ny = n *2 +figdir = "Variational_ShearBands2D_MPI" +igg = if !(JustRelax.MPI.Initialized()) + IGG(init_global_grid(nx, ny, 1; init_MPI = true)...) +else + igg +end +main(igg; figdir = figdir, nx = nx, ny = ny); From 32554aef3dafc16f361147c5b75fa41bb926699f Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 23 Dec 2024 13:47:23 +0100 Subject: [PATCH 66/77] fix --- .../stokes2D/shear_band/ShearBand2D_variational_MPI.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl index 19962fc92..44a544d00 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_variational_MPI.jl @@ -42,7 +42,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") lx = ly # 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 + di = @. li / (nx_g(),ny_g()) # grid step in x- and -y origin = 0.0, 0.0 # origin coordinates grid = Geometry(ni, li; origin = origin) (; xci, xvi) = grid # nodes at the center and vertices of the cells From dcc8100a5735e11f3abbb546e3e6ce388e6ab6bc Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Fri, 3 Jan 2025 15:04:30 +0100 Subject: [PATCH 67/77] =?UTF-8?q?add=20`get=5F=CE=B1`=20for=20new=20GP=20d?= =?UTF-8?q?ensity=20functions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/JustRelax.jl | 4 ++-- src/rheology/GeoParams.jl | 28 ++++++++++++++++++++++++++++ test/test/Project.toml | 2 -- 3 files changed, 30 insertions(+), 4 deletions(-) delete mode 100644 test/test/Project.toml diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 955419a4a..a23c5e599 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -16,14 +16,14 @@ using Statistics function solve!() end #! format: off function __init__() - println(""" + printstyled(""" _ _ ____ _ _ _ | |_ _ ___| |_| _ \\ ___| | __ ___ __ (_) | _ | | | | / __| __| |_) / _ \\ |/ _` \\ \\/ / | | | | |_| | |_| \\__ \\ |_| _ < __/ | (_| |> < _ | | | \\___/ \\__,_|___/\\__|_| \\_\\___|_|\\__,_/_/\\_(_)/ |_| |__/ - """) + """, bold=true, color=:white) end #! format: on abstract type AbstractBackend end diff --git a/src/rheology/GeoParams.jl b/src/rheology/GeoParams.jl index 01883e301..ca7ec291a 100644 --- a/src/rheology/GeoParams.jl +++ b/src/rheology/GeoParams.jl @@ -22,10 +22,38 @@ function get_α(rho::MeltDependent_Density; ϕ::T=0.0, kwargs...) where {T} return ϕ * αmelt + (1 - ϕ) * αsolid end + +function get_α(rho::BubbleFlow_Density; P = 0e0, kwargs...) + αmelt = get_α(rho.ρmelt, kwargs...) + αgas = get_α(rho.ρgas, kwargs...) + + @unpack_val c0, a = rho + + cutoff = c0^2/a^2 + + if P < cutoff + c = a * sqrt(abs(P)) + else + c = c0 + end + + return inv((c0-c)/αgas + (1-(c0-c))/αmelt) +end + +function get_α(rho::GasPyroclast_Density; kwargs...) + αmelt = get_α(rho.ρmelt, kwargs...) + αgas = get_α(rho.ρgas, kwargs...) + @unpack_val δ, β = rho + + return δ * αgas + (1 - δ) * αmelt +end + @inline get_α(p::MaterialParams) = get_α(p.Density[1]) @inline get_α(p::MaterialParams, args::NamedTuple) = get_α(p.Density[1], args) @inline get_α(p::Union{T_Density,PT_Density}) = GeoParams.get_α(p) @inline get_α(p::Union{T_Density,PT_Density}, ::Any) = GeoParams.get_α(p) @inline get_α(rho::MeltDependent_Density, ::Any) = get_α(rho) +@inline get_α(rho::BubbleFlow_Density, ::Any) = get_α(rho) +@inline get_α(rho::GasPyroclast_Density, ::Any) = get_α(rho) @inline get_α(rho::ConstantDensity, args) = 0 @inline get_α(rho::ConstantDensity) = 0 diff --git a/test/test/Project.toml b/test/test/Project.toml deleted file mode 100644 index ee29853db..000000000 --- a/test/test/Project.toml +++ /dev/null @@ -1,2 +0,0 @@ -[deps] -Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" From efd8047de3af777e7bba7aea333eaedfa971eff9 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Fri, 3 Jan 2025 15:08:49 +0100 Subject: [PATCH 68/77] format --- src/rheology/GeoParams.jl | 11 +++++------ src/variational_stokes/mask.jl | 7 ++++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/rheology/GeoParams.jl b/src/rheology/GeoParams.jl index ca7ec291a..a804946df 100644 --- a/src/rheology/GeoParams.jl +++ b/src/rheology/GeoParams.jl @@ -22,14 +22,13 @@ function get_α(rho::MeltDependent_Density; ϕ::T=0.0, kwargs...) where {T} return ϕ * αmelt + (1 - ϕ) * αsolid end - -function get_α(rho::BubbleFlow_Density; P = 0e0, kwargs...) +function get_α(rho::BubbleFlow_Density; P=0e0, kwargs...) αmelt = get_α(rho.ρmelt, kwargs...) - αgas = get_α(rho.ρgas, kwargs...) + αgas = get_α(rho.ρgas, kwargs...) @unpack_val c0, a = rho - cutoff = c0^2/a^2 + cutoff = c0^2 / a^2 if P < cutoff c = a * sqrt(abs(P)) @@ -37,12 +36,12 @@ function get_α(rho::BubbleFlow_Density; P = 0e0, kwargs...) c = c0 end - return inv((c0-c)/αgas + (1-(c0-c))/αmelt) + return inv((c0 - c) / αgas + (1 - (c0 - c)) / αmelt) end function get_α(rho::GasPyroclast_Density; kwargs...) αmelt = get_α(rho.ρmelt, kwargs...) - αgas = get_α(rho.ρgas, kwargs...) + αgas = get_α(rho.ρgas, kwargs...) @unpack_val δ, β = rho return δ * αgas + (1 - δ) * αmelt diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl index cda4fac9f..46c85bbb9 100644 --- a/src/variational_stokes/mask.jl +++ b/src/variational_stokes/mask.jl @@ -79,7 +79,9 @@ function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,3}, phase_ratios, air_phas return nothing end -@inline function compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} +@inline function compute_rock_ratio( + phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} +) where {N} 1 ≤ air_phase ≤ numphases(phase_ratio) || return 1e0 x = 1 - @index phase_ratio[air_phase, I...] x *= x > 1e-5 @@ -88,7 +90,7 @@ end @inline function compute_air_ratio( phase_ratio::CellArray, air_phase, I::Vararg{Integer,N} -) where {N} +) where {N} 1 ≤ air_phase ≤ numphases(phase_ratio) || return 1e0 return @index phase_ratio[air_phase, I...] end @@ -370,4 +372,3 @@ Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer,N}) where {N} = # cv = c || v # return cv || isvalid(ϕ.Vy, i, j) # end - From 7b75e23b19b3e6be15ae324fd5c238387886b962 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Thu, 16 Jan 2025 11:48:16 +0100 Subject: [PATCH 69/77] free surf stabilization attempt --- src/variational_stokes/Stokes2D.jl | 10 ++- src/variational_stokes/VelocityKernels.jl | 76 +++++++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 135be45e0..f41ff09da 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -31,6 +31,7 @@ function _solve_VS!( nout=500, b_width=(4, 4, 0), verbose=true, + free_surface=false, kwargs..., ) where {T} @@ -69,6 +70,7 @@ function _solve_VS!( θ = deepcopy(stokes.P) λ = @zeros(ni...) λv = @zeros(ni .+ 1...) + Vx_on_Vy = @zeros(size(stokes.V.Vy)) η0 = deepcopy(η) do_visc = true @@ -145,9 +147,14 @@ function _solve_VS!( ) update_halo!(stokes.τ.xy) + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( + Vx_on_Vy, stokes.V.Vx + ) + # @hide_communication b_width begin # communication/computation overlap @parallel (@idx ni .+ 1) compute_V!( @velocity(stokes)..., + Vx_on_Vy, stokes.R.Rx, stokes.R.Ry, stokes.P, @@ -157,6 +164,7 @@ function _solve_VS!( ητ, ϕ, _di..., + dt * free_surface, ) # apply boundary conditions velocity2displacement!(stokes, dt) @@ -184,7 +192,7 @@ function _solve_VS!( push!(err_evo1, err) push!(err_evo2, iter) - if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + if igg.me == 0 && verbose #((verbose && err > ϵ) || iter == iterMax) @printf( "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", iter, diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 4acfd6b2f..547d8897c 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -157,6 +157,82 @@ end return nothing end +@parallel_indices (i, j) function compute_V!( + Vx::AbstractArray{T,2}, + Vy, + Vx_on_Vy, + Rx, + Ry, + P, + τxx, + τyy, + τxy, + ηdτ, + ρgx, + ρgy, + ητ, + ϕ::JustRelax.RockRatio, + _dx, + _dy, + dt, +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A, ϕ) = _av_xa(A, ϕ, i, j) + av_ya(A, ϕ) = _av_ya(A, ϕ, i, j) + + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + if all((i, j) .< size(Vx) .- 1) + if isvalid_vx(ϕ, i + 1, j) + Rx[i, j] = + R_Vx = ( + -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - + av_xa(ρgx, ϕ.center) + ) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + else + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) + end + end + + if all((i, j) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j + 1) + θ = 1.0 + # Interpolated Vx into Vy node (includes density gradient) + Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] + # Vertical velocity + Vyᵢⱼ = Vy[i + 1, j + 1] + # Get necessary buoyancy forces + # i_W, i_E = max(i - 1, 1), min(i + 1, nx) + j_N = min(j + 1, size(ρgy,2)) + ρg_S = ρgy[i, j] * ϕ.center[i, j] + ρg_N = ρgy[i, j_N] * ϕ.center[i, j_N] + # Spatial derivatives + # ∂ρg∂x = (ρg_E - ρg_W) * _dx + ∂ρg∂y = (ρg_N - ρg_S) * _dy + # correction term + ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + Ry[i, j] = + R_Vy = + -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - + av_ya(ρgy, ϕ.center) + ρg_correction + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + else + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + end + end + + return nothing +end + @parallel_indices (i, j, k) function compute_V!( Vx::AbstractArray{T,3}, Vy, From 70590486965bc80e86eb37ea5b797c2d354345ec Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 17 Jan 2025 14:53:54 +0100 Subject: [PATCH 70/77] reverse change for now --- src/phases/topography_correction.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index 3155401e7..0a4b72ab9 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -43,17 +43,17 @@ function _update_phases_given_markerchain_kernel!( above = is_above_chain(xq, yq, chain_coords, cell_vertices) # if the particle is above the surface and the phase is not air, set the phase to air if above && phaseq != air_phase - # @index phase[ip, icell, j] = T(air_phase) - @index coords[1][ip, icell, j] = NaN - @index coords[2][ip, icell, j] = NaN - @index index[ip, icell, j] = false + @index phase[ip, icell, j] = T(air_phase) + # @index coords[1][ip, icell, j] = NaN + # @index coords[2][ip, icell, j] = NaN + # @index index[ip, icell, j] = false end # if the particle is above the surface and the phase is air, set the phase to the closes rock phase if !above && phaseq == air_phase - # @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) - @index coords[1][ip, icell, j] = NaN - @index coords[2][ip, icell, j] = NaN - @index index[ip, icell, j] = false + @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) + # @index coords[1][ip, icell, j] = NaN + # @index coords[2][ip, icell, j] = NaN + # @index index[ip, icell, j] = false end end end From 6d29609c68db322dd4f764bae6fd835655bdee10 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 22 Jan 2025 12:14:05 +0100 Subject: [PATCH 71/77] free surface stabilization --- src/Interpolations.jl | 32 ++++++ src/phases/topography_correction.jl | 4 +- src/variational_stokes/Stokes2D.jl | 65 ++++++++---- src/variational_stokes/VelocityKernels.jl | 124 +++++++++++++++++++++- 4 files changed, 205 insertions(+), 20 deletions(-) diff --git a/src/Interpolations.jl b/src/Interpolations.jl index a468f0444..d6b0f506b 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -32,6 +32,38 @@ end return nothing end +@parallel_indices (i, j) function interp_Vx∂ρ∂x_on_Vy!(Vx_on_Vy, Vx, ρg, ϕ, _dx) + nx, ny = size(ρg) + + iW = clamp(i - 1, 1, nx) + iE = clamp(i + 1, 1, nx) + jS = clamp(j - 1, 1, ny) + jN = clamp(j, 1, ny) + + # OPTION 1 + ρg_L = + 0.25 * ( + ρg[iW, jS] * ϕ.center[iW, jS] + + ρg[i, jS] * ϕ.center[i, jS] + + ρg[iW, jN] * ϕ.center[iW, jN] + + ρg[i, jN] * ϕ.center[i, jN] + ) + ρg_R = + 0.25 * ( + ρg[iE, jS] * ϕ.center[iE, jS] + + ρg[i, jS] * ϕ.center[i, jS] + + ρg[iE, jN] * ϕ.center[iE, jN] + + ρg[i, jN] * ϕ.center[i, jN] + ) + + Vx_on_Vy[i + 1, j] = + (0.25 * (Vx[i, j] + Vx[i + 1, j] + Vx[i, j + 1] + Vx[i + 1, j + 1])) * + (ρg_R - ρg_L) * + _dx + + return nothing +end + # From cell vertices to cell center temperature2center!(thermal) = temperature2center!(backend(thermal), thermal) diff --git a/src/phases/topography_correction.jl b/src/phases/topography_correction.jl index 0a4b72ab9..b251c2bdf 100644 --- a/src/phases/topography_correction.jl +++ b/src/phases/topography_correction.jl @@ -50,7 +50,9 @@ function _update_phases_given_markerchain_kernel!( end # if the particle is above the surface and the phase is air, set the phase to the closes rock phase if !above && phaseq == air_phase - @index phase[ip, icell, j] = closest_phase(coords, (xq, yq), index, ip, phase, air_phase, icell, j) + @index phase[ip, icell, j] = closest_phase( + coords, (xq, yq), index, ip, phase, air_phase, icell, j + ) # @index coords[1][ip, icell, j] = NaN # @index coords[2][ip, icell, j] = NaN # @index index[ip, icell, j] = false diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index f41ff09da..4d07fe8e0 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -147,25 +147,54 @@ function _solve_VS!( ) update_halo!(stokes.τ.xy) - @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( - Vx_on_Vy, stokes.V.Vx - ) + if free_surface + @parallel (@idx ni .+ 1) compute_Vx!( + stokes.V.Vx, + stokes.R.Rx, + stokes.P, + stokes.τ.xx, + stokes.τ.xy, + ηdτ, + ρg[begin], + ητ, + ϕ, + _di..., + ) + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!( + Vx_on_Vy, stokes.V.Vx, last(ρg), ϕ, first(_di) + ) + update_halo!(Vx_on_Vy) - # @hide_communication b_width begin # communication/computation overlap - @parallel (@idx ni .+ 1) compute_V!( - @velocity(stokes)..., - Vx_on_Vy, - stokes.R.Rx, - stokes.R.Ry, - stokes.P, - @stress(stokes)..., - ηdτ, - ρg..., - ητ, - ϕ, - _di..., - dt * free_surface, - ) + @parallel (@idx ni .+ 1) compute_Vy!( + stokes.V.Vy, + Vx_on_Vy, + stokes.R.Ry, + stokes.P, + stokes.τ.yy, + stokes.τ.xy, + ηdτ, + ρg[end], + ητ, + ϕ, + _di..., + dt, + ) + else + # @hide_communication b_width begin # communication/computation overlap + @parallel (@idx ni .+ 1) compute_V!( + @velocity(stokes)..., + Vx_on_Vy, + stokes.R.Rx, + stokes.R.Ry, + stokes.P, + @stress(stokes)..., + ηdτ, + ρg..., + ητ, + ϕ, + _di..., + ) + end # apply boundary conditions velocity2displacement!(stokes, dt) # free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 547d8897c..973083290 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,3 +1,35 @@ +@parallel_indices (i, j) function interp_Vx∂ρ∂x_on_Vy!(Vx_on_Vy, Vx, ρg, ϕ, _dx) + nx, ny = size(ρg) + + iW = clamp(i - 1, 1, nx) + iE = clamp(i + 1, 1, nx) + jS = clamp(j - 1, 1, ny) + jN = clamp(j, 1, ny) + + # OPTION 1 + ρg_L = + 0.25 * ( + ρg[iW, jS] * ϕ.center[iW, jS] + + ρg[i, jS] * ϕ.center[i, jS] + + ρg[iW, jN] * ϕ.center[iW, jN] + + ρg[i, jN] * ϕ.center[i, jN] + ) + ρg_R = + 0.25 * ( + ρg[iE, jS] * ϕ.center[iE, jS] + + ρg[i, jS] * ϕ.center[i, jS] + + ρg[iE, jN] * ϕ.center[iE, jN] + + ρg[i, jN] * ϕ.center[i, jN] + ) + + Vx_on_Vy[i + 1, j] = + (0.25 * (Vx[i, j] + Vx[i + 1, j] + Vx[i, j + 1] + Vx[i + 1, j + 1])) * + (ρg_R - ρg_L) * + _dx + + return nothing +end + @parallel_indices (I...) function compute_∇V!( ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} ) where {T,N} @@ -157,6 +189,96 @@ end return nothing end +@parallel_indices (i, j) function compute_Vx!( + Vx::AbstractArray{T,2}, Rx, P, τxx, τxy, ηdτ, ρgx, ητ, ϕ::JustRelax.RockRatio, _dx, _dy +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A, ϕ) = _av_xa(A, ϕ, i, j) + av_ya(A, ϕ) = _av_ya(A, ϕ, i, j) + + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + if all((i, j) .< size(Vx) .- 1) + if isvalid_vx(ϕ, i + 1, j) + Rx[i, j] = + R_Vx = ( + -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) - + av_xa(ρgx, ϕ.center) + ) + Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ) + else + Rx[i, j] = zero(T) + Vx[i + 1, j + 1] = zero(T) + end + end + + return nothing +end + +@parallel_indices (i, j) function compute_Vy!( + Vy::AbstractArray{T,2}, + Vx_on_Vy, + Ry, + P, + τyy, + τxy, + ηdτ, + ρgy, + ητ, + ϕ::JustRelax.RockRatio, + _dx, + _dy, + dt, +) where {T} + d_xi(A, ϕ) = _d_xi(A, ϕ, _dx, i, j) + d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j) + d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j) + d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j) + av_xa(A, ϕ) = _av_xa(A, ϕ, i, j) + av_ya(A, ϕ) = _av_ya(A, ϕ, i, j) + + av_xa(A) = _av_xa(A, i, j) + av_ya(A) = _av_ya(A, i, j) + harm_xa(A) = _av_xa(A, i, j) + harm_ya(A) = _av_ya(A, i, j) + + if all((i, j) .< size(Vy) .- 1) + if isvalid_vy(ϕ, i, j + 1) + θ = 1.0 + # Interpolated Vx into Vy node (includes density gradient) + Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1] + # Vertical velocity + Vyᵢⱼ = Vy[i + 1, j + 1] + # Get necessary buoyancy forces + # i_W, i_E = max(i - 1, 1), min(i + 1, nx) + j_N = min(j + 1, size(ρgy, 2)) + ρg_S = ρgy[i, j] * ϕ.center[i, j] + ρg_N = ρgy[i, j_N] * ϕ.center[i, j_N] + # Spatial derivatives + # ∂ρg∂x = (ρg_E - ρg_W) * _dx + ∂ρg∂y = (ρg_N - ρg_S) * _dy + # correction term + ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + Ry[i, j] = + R_Vy = + -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - + av_ya(ρgy, ϕ.center) + ρg_correction + Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + else + Ry[i, j] = zero(T) + Vy[i + 1, j + 1] = zero(T) + end + end + + return nothing +end + @parallel_indices (i, j) function compute_V!( Vx::AbstractArray{T,2}, Vy, @@ -211,7 +333,7 @@ end Vyᵢⱼ = Vy[i + 1, j + 1] # Get necessary buoyancy forces # i_W, i_E = max(i - 1, 1), min(i + 1, nx) - j_N = min(j + 1, size(ρgy,2)) + j_N = min(j + 1, size(ρgy, 2)) ρg_S = ρgy[i, j] * ϕ.center[i, j] ρg_N = ρgy[i, j_N] * ϕ.center[i, j_N] # Spatial derivatives From 677dda51e0d72ddf5844463a4a1e21e7ab4b753e Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 22 Jan 2025 14:52:51 +0100 Subject: [PATCH 72/77] typo --- src/variational_stokes/Stokes2D.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl index 4d07fe8e0..fb0766549 100644 --- a/src/variational_stokes/Stokes2D.jl +++ b/src/variational_stokes/Stokes2D.jl @@ -183,7 +183,6 @@ function _solve_VS!( # @hide_communication b_width begin # communication/computation overlap @parallel (@idx ni .+ 1) compute_V!( @velocity(stokes)..., - Vx_on_Vy, stokes.R.Rx, stokes.R.Ry, stokes.P, From 009c5342b4ccdf63dd7873dc5bfaa9b6b5dcf169 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 22 Jan 2025 14:53:10 +0100 Subject: [PATCH 73/77] miniapps --- .../PlumeFreeSurface_VariationalStokes.jl | 14 +- .../RayleighTaylor2D_VariationalStokes.jl | 280 ++++++++++++++++++ 2 files changed, 288 insertions(+), 6 deletions(-) create mode 100644 miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D_VariationalStokes.jl diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl index c414762f8..ed9fe5ad2 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_VariationalStokes.jl @@ -1,5 +1,5 @@ -# const isCUDA = false -const isCUDA = true +const isCUDA = false +# const isCUDA = true @static if isCUDA using CUDA @@ -32,7 +32,8 @@ else end # Load script dependencies -using GeoParams, GLMakie +using GeoParams +using GLMakie import ParallelStencil.INDICES const idx_j = INDICES[2] @@ -198,11 +199,12 @@ function main(igg, nx, ny) igg; kwargs = (; iterMax = 100e3, - nout = 2e3, + nout = 1e3, viscosity_cutoff = (-Inf, Inf), + free_surface = true, ) ) - # dt = compute_dt(stokes, di) / 2 + dt = compute_dt(stokes, di) / 2 # ------------------------------ # Advection -------------------- @@ -261,7 +263,7 @@ end ## END OF MAIN SCRIPT ---------------------------------------------------------------- # (Path)/folder where output data and figures are stored -n = 128 +n = 100 nx = n ny = n igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D_VariationalStokes.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D_VariationalStokes.jl new file mode 100644 index 000000000..472a04f5b --- /dev/null +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D_VariationalStokes.jl @@ -0,0 +1,280 @@ +const isCUDA = false +# const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +using JustPIC, JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend_JP = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +# Load script dependencies +using GeoParams +using GLMakie + +# Velocity helper grids for the particle advection +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + @parallel f_x(A, B) +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 + +function init_phases!(phases, particles, A) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) + + f(x, A, λ) = A * sin(π * x / λ) + + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) + @index phases[ip, i, j] = 2.0 + + if 0e0 ≤ depth ≤ 100e3 + @index phases[ip, i, j] = 1.0 + + elseif depth > (-f(x, A, 500e3) + (200e3 - A)) + @index phases[ip, i, j] = 3.0 + + end + + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, A) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +# (Path)/folder where output data and figures are stored +n = 101 +nx = n +ny = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(igg, nx, ny) + + # Physical domain ------------------------------------ + thick_air = 100e3 # thickness of sticky air layer + ly = 500e3 + thick_air # domain length in y + lx = 500e3 # 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) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = rheology = ( + # Name = "Air", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=1e0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e16),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Crust", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Mantle", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ) + ) + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 125, 175, 75 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + A = 5e3 # Amplitude of the anomaly + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) + init_phases!(pPhases, particles, A) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # RockRatios + air_phase = 1 + ϕ = RockRatio(backend, ni) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + # Initialize marker chain------------------------------- + nxcell, max_xcell, min_xcell = 100, 150, 75 + initial_elevation = -100e3 + chain = init_markerchain(backend_JP, nxcell, min_xcell, max_xcell, xvi[1], initial_elevation); + # ---------------------------------------------------- + + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.98 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend, ni) + # ---------------------------------------------------- + + # Buoyancy forces & rheology + ρg = @zeros(ni...), @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + # @parallel init_P!(stokes.P, ρg[2], xci[2]) + compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, (-Inf, Inf)) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = false), + no_slip = (left = false, right = false, top = false, bot = true), + free_surface = true, + ) + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + figdir = "RayleighTaylor2D" + take(figdir) + + # Time loop + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + dt_max = 3e3 * (3600 * 24 * 365.25) + + while it < 1000 #00 + + if t / (1e3 * 3600 * 24 *365.25) > 500 + dt_max = 25e3 * (3600 * 24 * 365.25) + end + ## variational solver + # Stokes solver ---------------- + solve_VariationalStokes!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + ϕ, + rheology, + args, + dt, + igg; + kwargs = ( + iterMax = 50e3, + iterMin = 1e3, + viscosity_relaxation = 1e-2, + nout = 2e3, + free_surface = true, + viscosity_cutoff = (-Inf, Inf) + ) + ) + dt = compute_dt(stokes, di, dt_max) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_MQS!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + + # advect marker chain + advect_markerchain!(chain, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + update_phases_given_markerchain!(pPhases, chain, particles, origin, di, air_phase) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + update_rock_ratio!(ϕ, phase_ratios, air_phase) + + @show it += 1 + t += dt + + if it == 1 || rem(it, 5) == 0 + px, py = particles.coords + + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + nt = 5 + fig = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs") + # heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array([argmax(p) for p in phase_ratios.vertex]), colormap = :grayC) + scatter!(ax, Array(px.data[:]).*1e-3, Array(py.data[:]).*1e-3, color =Array(pPhases.data[:]), colormap = :grayC) + arrows!( + ax, + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), + color = :red, + ) + fig + save(joinpath(figdir, "$(it).png"), fig) + + end + end + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +main(igg, nx, ny) \ No newline at end of file From 6e5b973ab27029c96bce12e921225487ba8ad4c4 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 22 Jan 2025 16:16:38 +0100 Subject: [PATCH 74/77] up --- src/variational_stokes/VelocityKernels.jl | 42 +++++------------------ 1 file changed, 9 insertions(+), 33 deletions(-) diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl index 973083290..6d553fe35 100644 --- a/src/variational_stokes/VelocityKernels.jl +++ b/src/variational_stokes/VelocityKernels.jl @@ -1,35 +1,3 @@ -@parallel_indices (i, j) function interp_Vx∂ρ∂x_on_Vy!(Vx_on_Vy, Vx, ρg, ϕ, _dx) - nx, ny = size(ρg) - - iW = clamp(i - 1, 1, nx) - iE = clamp(i + 1, 1, nx) - jS = clamp(j - 1, 1, ny) - jN = clamp(j, 1, ny) - - # OPTION 1 - ρg_L = - 0.25 * ( - ρg[iW, jS] * ϕ.center[iW, jS] + - ρg[i, jS] * ϕ.center[i, jS] + - ρg[iW, jN] * ϕ.center[iW, jN] + - ρg[i, jN] * ϕ.center[i, jN] - ) - ρg_R = - 0.25 * ( - ρg[iE, jS] * ϕ.center[iE, jS] + - ρg[i, jS] * ϕ.center[i, jS] + - ρg[iE, jN] * ϕ.center[iE, jN] + - ρg[i, jN] * ϕ.center[i, jN] - ) - - Vx_on_Vy[i + 1, j] = - (0.25 * (Vx[i, j] + Vx[i + 1, j] + Vx[i, j + 1] + Vx[i + 1, j + 1])) * - (ρg_R - ρg_L) * - _dx - - return nothing -end - @parallel_indices (I...) function compute_∇V!( ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N} ) where {T,N} @@ -264,12 +232,20 @@ end # ∂ρg∂x = (ρg_E - ρg_W) * _dx ∂ρg∂y = (ρg_N - ρg_S) * _dy # correction term - ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt + Ry[i, j] = R_Vy = -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - av_ya(ρgy, ϕ.center) + ρg_correction Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ) + + # ρgx_correction = (Vxᵢⱼ) * θ * dt + # ρgy_correction = (Vyᵢⱼ * ∂ρg∂y) * θ * dt + # Ry[i, j] = R_Vy = + # -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) - + # av_ya(ρgy, ϕ.center) + ρgx_correction + # Vy[i + 1, j + 1] += R_Vy * inv(inv(ηdτ / av_ya(ητ)) + ρgy_correction) else Ry[i, j] = zero(T) Vy[i + 1, j + 1] = zero(T) From 6b6080159b562658be66b70afb83752dd47c0e71 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Thu, 23 Jan 2025 13:34:04 +0100 Subject: [PATCH 75/77] add JR versioning to the `_init__()` function --- Project.toml | 2 ++ src/JustRelax.jl | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index dbb1fef8d..6bd0ec8e6 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" [weakdeps] @@ -45,6 +46,7 @@ ParallelStencil = "0.13.6, 0.14" Reexport = "1.2.2" StaticArrays = "1" Statistics = "1" +TOML = "1.0.3" WriteVTK = "1.18, 1.19" julia = "1.9, 1.10" diff --git a/src/JustRelax.jl b/src/JustRelax.jl index a23c5e599..b43ed19a4 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -11,6 +11,7 @@ using HDF5 using CellArrays using StaticArrays using Statistics +using TOML @reexport using JustPIC function solve!() end @@ -23,7 +24,8 @@ function __init__() | |_| | |_| \\__ \\ |_| _ < __/ | (_| |> < _ | | | \\___/ \\__,_|___/\\__|_| \\_\\___|_|\\__,_/_/\\_(_)/ |_| |__/ - """, bold=true, color=:white) + Version: $(TOML.parsefile("Project.toml")["version"]) +""", bold=true, color=:white) end #! format: on abstract type AbstractBackend end From 7284c35987304a8cd44eeeb7e1eddea9b7b3e6a7 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Mon, 27 Jan 2025 13:03:41 +0000 Subject: [PATCH 76/77] =?UTF-8?q?switch=20back=20to=20`F=20=3D=20=CF=84II?= =?UTF-8?q?=5Fij=20-=20C=20-=20max(Pr[I...],=200.0)=20*=20sin=CF=95`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/stokes/StressKernels.jl | 14 +++++++------- src/variational_stokes/StressKernels.jl | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 6828dd71d..b4f073c86 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -597,7 +597,7 @@ end ni = size(Pr) Ic = clamped_indices(ni, I...) - ## yz + ## yz if all(I .≤ size(ε[4])) # interpolate to ith vertex ηv_ij = av_clamped_yz(η, Ic...) @@ -647,7 +647,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[1][I...] = @@ -710,7 +710,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[2][I...] = @@ -774,7 +774,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[3][I...] = @@ -810,7 +810,7 @@ end dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij + 2.0 * ηij * εij) * dτ_r τII_ij = second_invariant(dτij .+ τij) # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ + F = τII_ij - C - max(Pr[I...], 0.0) * sinϕ if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center @@ -898,7 +898,7 @@ end τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2) # yield function @ center - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij,0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[I...] = @@ -933,7 +933,7 @@ end dτij = compute_stress_increment(τij, τij_o, ηij, εij, _Gdt, dτ_r) τII_ij = GeoParams.second_invariant(dτij .+ τij) # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ + F = τII_ij - C - max(Pr[I...],0.0) * sinϕ if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center diff --git a/src/variational_stokes/StressKernels.jl b/src/variational_stokes/StressKernels.jl index 9ac0c3a3a..3f4f11ff9 100644 --- a/src/variational_stokes/StressKernels.jl +++ b/src/variational_stokes/StressKernels.jl @@ -66,7 +66,7 @@ ) # yield function @ center - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[I...] = @@ -106,7 +106,7 @@ dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r τII_ij = GeoParams.second_invariant(dτij .+ τij) # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ + F = τII_ij - C - max(Pr[I...], 0.0) * sinϕ if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center @@ -175,7 +175,7 @@ end ni = size(Pr) Ic = clamped_indices(ni, I...) - ## yz + ## yz if all(I .≤ size(ε[4])) && isvalid_yz(ϕ, I...) # interpolate to ith vertex ηv_ij = av_clamped_yz(η, Ic...) @@ -225,7 +225,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[1][I...] = @@ -288,7 +288,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[2][I...] = @@ -352,7 +352,7 @@ end τIIv_ij = second_invariant(τijv .+ dτijv) # yield function @ vertex - Fv = τIIv_ij - Cv - Pv_ij * sinϕv + Fv = τIIv_ij - Cv - max(Pv_ij, 0.0) * sinϕv if is_pl && !iszero(τIIv_ij) && Fv > 0 # stress correction @ vertex λv[3][I...] = @@ -388,7 +388,7 @@ end dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij + 2.0 * ηij * εij) * dτ_r τII_ij = second_invariant(dτij .+ τij) # yield function @ center - F = τII_ij - C - Pr[I...] * sinϕ + F = τII_ij - C - max(Pr[I...], 0.0) * sinϕ if is_pl && !iszero(τII_ij) && F > 0 # stress correction @ center From b9d2c19a249e29301c388bdb5e9f19d15ddc6d36 Mon Sep 17 00:00:00 2001 From: Pascal Aellig Date: Tue, 28 Jan 2025 13:38:18 +0000 Subject: [PATCH 77/77] update Caldera scripts --- .../stokes2D/Volcano2D/Caldera2D.jl | 199 ++++++++++-------- .../stokes2D/Volcano2D/Caldera_rheology.jl | 85 +++----- .../stokes2D/Volcano2D/Caldera_setup.jl | 27 ++- 3 files changed, 164 insertions(+), 147 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl index 7981dc03a..a4b343e33 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera2D.jl @@ -1,5 +1,5 @@ -const isCUDA = false -# const isCUDA = true +# const isCUDA = false +const isCUDA = true @static if isCUDA using CUDA @@ -32,7 +32,7 @@ else end # Load script dependencies -using GeoParams, GLMakie, CellArrays, Statistics, Dates +using GeoParams, CairoMakie, CellArrays, Statistics, Dates, JLD2 # Load file with all the rheology configurations include("Caldera_setup.jl") @@ -57,7 +57,7 @@ end # Initial pressure profile - not accurate @parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + @all(P) = abs(@all(ρg) * @all_k(z)) #* <(@all_k(z), 0.0) return nothing end @@ -92,18 +92,21 @@ function extract_topo_from_GMG_phases(phases_GMG, xvi, air_phase) return topo_y end -function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_phase, magma_phase, air_phase) +function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_phase, magma_phase, anomaly_phase, air_phase) - @parallel_indices (i, j) function _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, vertex_ratio, conduit_phase, magma_phase, air_phase) + @parallel_indices (i, j) function _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, vertex_ratio, conduit_phase, magma_phase, anomaly_phase, air_phase) # quick escape conduit_ratio_ij = @index vertex_ratio[conduit_phase, i, j] magma_ratio_ij = @index vertex_ratio[magma_phase, i, j] + anomaly_ratio_ij = @index vertex_ratio[anomaly_phase, i, j] air_ratio_ij = @index vertex_ratio[air_phase, i, j] - if conduit_ratio_ij > 0.5 || magma_ratio_ij > 0.5 + # if conduit_ratio_ij > 0.5 || magma_ratio_ij > 0.5 + if conduit_ratio_ij > 0.5 || anomaly_ratio_ij > 0.5 # if isone(conduit_ratio_ij) || isone(magma_ratio_ij) Ω_T[i+1, j] = Temp[i+1, j] = T_chamber - + elseif magma_ratio_ij > 0.5 + Ω_T[i+1, j] = Temp[i+1, j] = T_chamber - 100e0 elseif air_ratio_ij > 0.5 Ω_T[i+1, j] = Temp[i+1, j] = T_air end @@ -113,7 +116,7 @@ function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_ph ni = size(phase_ratios.vertex) - @parallel (@idx ni) _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, phase_ratios.vertex, conduit_phase, magma_phase, air_phase) + @parallel (@idx ni) _thermal_anomaly!(Temp, Ω_T, T_chamber, T_air, phase_ratios.vertex, conduit_phase, magma_phase, anomaly_phase, air_phase) @views Ω_T[1, :] .= Ω_T[2, :] @views Ω_T[end, :] .= Ω_T[end-1, :] @@ -123,8 +126,24 @@ function thermal_anomaly!(Temp, Ω_T, phase_ratios, T_chamber, T_air, conduit_ph return nothing end +function plot_particles(particles, pPhases) + p = particles.coords + # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter + ppx, ppy = p + # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) + # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) + pxv = ppx.data[:] + pyv = ppy.data[:] + clr = pPhases.data[:] + # clr = pϕ.data[:] + idxv = particles.index.data[:] + f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma, markersize=1) + Colorbar(f[1,2], h) + f +end + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false, extension = 1e-15 * 0) +function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false, extension = 1e-15 * 0, cutoff_visc = (1e16, 1e23)) # Physical domain ------------------------------------ ni = nx, ny # number of cells @@ -134,11 +153,16 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # ---------------------------------------------------- # Physical properties using GeoParams ---------------- - rheology = init_rheologies() - dt = 5e2 * 3600 * 24 * 365 - # dt = Inf # diffusive CFL timestep limiter + rheology = init_rheologies(; incompressible=false, magma=false) + rheology_incomp = init_rheologies(; incompressible=true, magma=false) + dt_time = 1e3 * 3600 * 24 * 365 + κ = (4 / (rheology[1].HeatCapacity[1].Cp.val * rheology[1].Density[1].ρ0.val)) # thermal diffusivity # thermal diffusivity + dt_diff = 0.5 * min(di...)^2 / κ / 2.01 + dt = min(dt_time, dt_diff) # ---------------------------------------------------- + # randomize cohesion + perturbation_C = @rand(ni...); # Initialize particles ------------------------------- nxcell = 100 max_xcell = 150 @@ -172,11 +196,6 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) # particle fields for the stress rotation - # pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress - # # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress - # pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity - # particle_args = (pT, pPhases, pτ..., pω...) - # particle_args_reduced = (pT, pτ..., pω...) pτ = StressParticles(particles) particle_args = (pT, pPhases, unwrap(pτ)...) particle_args_reduced = (pT, unwrap(pτ)...) @@ -203,12 +222,12 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", T_chamber = 1223e0 T_air = 273e0 Ω_T = @zeros(size(thermal.T)...) - thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) - JustRelax.DirichletBoundaryCondition(Ω_T) + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, 4, air_phase) + # JustRelax.DirichletBoundaryCondition(Ω_T) thermal_bc = TemperatureBoundaryConditions(; no_flux = (; left = true, right = true, top = false, bot = false), - dirichlet = (; mask = Ω_T) + # dirichlet = (; mask = Ω_T) ) thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) @@ -218,8 +237,11 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(2)) + for _ in 1:5 compute_ρg!(ρg, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + @parallel init_P!(stokes.P, ρg[end], xvi[2]) + end + # stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) # Melt fraction ϕ_m = @zeros(ni...) @@ -227,8 +249,8 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ϕ_m, phase_ratios, rheology, (T=thermal.Tc, P=stokes.P) ) # Rheology - args0 = (; ϕ=ϕ_m,T=thermal.Tc, P=stokes.P, dt = Inf) - viscosity_cutoff = (1e16, 1e23) + args0 = (; ϕ=ϕ_m,T=thermal.Tc, P=stokes.P, dt = Inf, perturbation_C = perturbation_C) + viscosity_cutoff = (1e17, 1e23) compute_viscosity!(stokes, phase_ratios, args0, rheology, air_phase, viscosity_cutoff) # PT coefficients for thermal diffusion @@ -258,12 +280,12 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # IO ------------------------------------------------- # if it does not exist, make folder where figures are stored if plotting - checkpoint = joinpath(figdir, "checkpoint") if do_vtk vtk_dir = joinpath(figdir, "vtk") take(vtk_dir) end take(figdir) + checkpoint = joinpath(figdir, "checkpoint") end # ---------------------------------------------------- @@ -310,23 +332,42 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # Time loop t, it = 0.0, 0 + interval = 0 + iterMax = 150e3 + local iters thermal.Told .= thermal.T - while it < 10 #000 # run only for 5 Myrs + while it < 100 #000 # run only for 5 Myrs + if it >1 && iters.iter > iterMax && iters.err_evo1[end] > pt_stokes.ϵ * 5 + iterMax += 10e3 + iterMax = min(iterMax, 200e3) + println("Increasing maximum pseudo timesteps to $iterMax") + else + iterMax = 150e3 + end # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) @views T_buffer[:, end] .= Ttop @views T_buffer[:, 1] .= Tbot + # clamp!(T_buffer, 273e0, 1223e0) @views thermal.T[2:end-1, :] .= T_buffer - if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=3), 1.5e3) == 0.0 - thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, air_phase) + if it > 1 && rem(it, 5) == 0 + # if mod(round(t/(1e3 * 3600 * 24 *365.25); digits=1), 1e3) == 0.0 + println("Simulation eruption at t = $(round(t/(1e3 * 3600 * 24 *365.25); digits=2)) Kyrs") + thermal_anomaly!(thermal.T, Ω_T, phase_ratios, T_chamber, T_air, 5, 3, 4, air_phase) + interval += 1 + copyinn_x!(T_buffer, thermal.T) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + temperature2center!(thermal) + grid2particle!(pT, xvi, T_buffer, particles) end thermal_bcs!(thermal, thermal_bc) temperature2center!(thermal) - # args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) - args = (; ϕ=ϕ_m, T=thermal.Tc, P=stokes.P, dt=Inf) + args = (;ϕ=ϕ_m, T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc, perturbation_C=perturbation_C) + # args = (; ϕ=ϕ_m, T=thermal.Tc, P=stokes.P, dt=Inf) stress2grid!(stokes, pτ, xvi, xci, particles) @@ -354,10 +395,11 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", println("Stokes solver time ") println(" Total time: $t_stokes s") + println("Extrema T[C]: $(extrema(thermal.T.-273))") tensor_invariant!(stokes.ε) tensor_invariant!(stokes.ε_pl) dtmax = 2e3 * 3600 * 24 * 365.25 - dt = compute_dt(stokes, di, dtmax) * 0.5 + dt = compute_dt(stokes, di, dtmax) println("dt = $(dt/(3600 * 24 *365.25)) years") # ------------------------------ @@ -398,7 +440,6 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", # advect particles in memory move_particles!(particles, xvi, particle_args) # check if we need to inject particles - # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) center2vertex!(τxx_v, stokes.τ.xx) center2vertex!(τyy_v, stokes.τ.yy) inject_particles_phase!( @@ -424,31 +465,43 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", tensor_invariant!(stokes.τ) - # track deformation of free_surface - # push!(deformation_x, chain.coords[1].data[:]./1e3) - # push!(deformation_y, chain.coords[2].data[:]./1e3) @show it += 1 t += dt + + if it == 1 + stokes.EII_pl .= 0.0 + end + if plotting # Data I/O and plotting --------------------- if it == 1 || rem(it, 1) == 0 + ## this is used for plotting for now: + η_eff = @. stokes.τ.II / (2*stokes.ε.II) + ## if igg.me == 0 && it == 1 - metadata(pwd(), checkpoint, basename(@__FILE__), "Caldera_setup.jl", "Caldera_rheology.jl") + metadata(pwd(), checkpoint, basename(@__FILE__), joinpath(@__DIR__, "Caldera_setup.jl"), joinpath(@__DIR__,"Caldera_rheology.jl")) end checkpointing_jld2(checkpoint, stokes, thermal, t, dt, igg) - + checkpointing_particles(checkpoint, particles; phases= pPhases, phase_ratios=phase_ratios, chain=chain, particle_args=particle_args, t=t, dt=dt) (; η_vep, η) = stokes.viscosity if do_vtk velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(T_buffer), - τII = Array(stokes.τ.II), - εII = Array(stokes.ε.II), - εII_pl = Array(stokes.ε_pl.II), + stress_xy = Array(stokes.τ.xy), + strain_rate_xy = Array(stokes.ε.xy), + phase_vertices = [argmax(p) for p in Array(phase_ratios.vertex)], ) data_c = (; P = Array(stokes.P), - η = Array(η_vep), + viscosity = Array(η_eff), + phases = [argmax(p) for p in Array(phase_ratios.center)], + Melt_fraction = Array(ϕ_m), + EII_pl = Array(stokes.EII_pl), + stress_II = Array(stokes.τ.II), + strain_rate_II = Array(stokes.ε.II), + plastic_strain_rate_II = Array(stokes.ε_pl.II), + density = Array(ρg[2]./9.81), ) velocity_v = ( Array(Vx_v), @@ -463,6 +516,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", velocity_v; t = round(t/(1e3 * 3600 * 24 *365.25); digits=3) ) + save_marker_chain(joinpath(vtk_dir, "chain_" * lpad("$it", 6, "0")), xvi[1], Array(chain.h_vertices)) end # Make particles plottable @@ -478,9 +532,9 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", chain_y = chain.coords[2].data[:]./1e3 # Make Makie figure - ar = 2 + ar = DataAspect() fig = Figure(size = (1200, 900), title = "t = $t") - ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(round(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs)") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [C] (t=$(round(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs)") ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [cm/yr]") # ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") ax3 = Axis(fig[1,3], aspect = ar, title = "τII [MPa]") @@ -489,7 +543,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", ax5 = Axis(fig[3,1], aspect = ar, title = "EII_pl") ax6 = Axis(fig[3,3], aspect = ar, title = "Melt fraction ϕ") # Plot temperature - h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:].-273) , colormap=:batlow) # Plot particles phase h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, uconvert.(u"cm/yr",Array(stokes.V.Vy)u"m/s") , colormap=:batlow) @@ -500,7 +554,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II)./1e6 , colormap=:batlow) # Plot effective viscosity # h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:lipari) - h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_eff)), colorrange=log10.(viscosity_cutoff), colormap=:batlow) h5 = heatmap!(ax5, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.EII_pl) , colormap=:batlow) h6 = heatmap!(ax6, xci[1].*1e-3, xci[2].*1e-3, Array(ϕ_m) , colormap=:lipari) hidexdecorations!(ax1) @@ -516,7 +570,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", Colorbar(fig[1,4], h3) Colorbar(fig[2,4], h4) Colorbar(fig[3,2], h5) - # Colorbar(fig[3,4], h6) + Colorbar(fig[3,4], h6) linkaxes!(ax1, ax2, ax3, ax4, ax5) fig save(joinpath(figdir, "$(it).png"), fig) @@ -543,27 +597,7 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", fig end - fig = 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:].-273.15), - Yv./1e3, - ) - lines!( - ax2, - Array(stokes.P[:]./1e6), - Y./1e3, - ) - hideydecorations!(ax2) - save(joinpath(figdir, "thermal_profile_$it.png"), fig) - fig - end - + # Plot Drucker Prager yield surface fig1 = let fig = Figure(; size=(1200, 900)) ax = Axis(fig[1, 1];title="Drucker Prager") @@ -578,29 +612,30 @@ function main(li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", end end - return nothing + return nothing end ## END OF MAIN SCRIPT ---------------------------------------------------------------- const plotting = true -extension = 1e-15 * 1 + do_vtk = true # set to true to generate VTK files for ParaView -figdir = "Caldera2D" -n = 64 + +figdir = "Caldera2D_$(today())" +n = 256 nx, ny = n, n >>> 1 li, origin, phases_GMG, T_GMG = setup2D( nx+1, ny+1; sticky_air = 4e0, - dimensions = (30e0, 20e0), # extent in x and y in km - flat = false, - chimney = true, - volcano_size = (3e0, 5e0), - conduit_radius = 2e-1, - chamber_T = 900e0, - chamber_depth = 7e0, - chamber_radius = 1.25, - aspect_x = 2, + dimensions = (25e0, 20e0), # extent in x and y in km + flat = false, # flat or volcano cone + chimney = true, # conduit or not + volcano_size = (3e0, 6e0), # height, radius + conduit_radius = 4e-1, # radius of the conduit + chamber_T = 900e0, # temperature of the chamber + chamber_depth = 5e0, # depth of the chamber + chamber_radius = 1.25e0, # radius of the chamber + aspect_x = 2.5e0, # aspect ratio of the chamber ) igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -608,5 +643,5 @@ igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid else igg end - -main(li, origin, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk, extension = extension); +extension = 1e-15 +main(li, origin, phases_GMG, T_GMG, igg, ; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk, extension = extension, cutoff_visc = (1e17, 1e23)); diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl index 6aa020ff9..72fa1ad3d 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_rheology.jl @@ -1,64 +1,30 @@ using GeoParams.Dislocation using GeoParams.Diffusion -function init_rheology_nonNewtonian() - #dislocation laws - disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) - # diffusion laws - diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) - - el = ConstantElasticity(; G = 40e9) - - lithosphere_rheology = CompositeRheology( (el, disl_wet_olivine, diff_wet_olivine)) - init_rheologies(lithosphere_rheology) -end +function init_rheologies(; linear=false, incompressible=true, isplastic = true, magma = false) -function init_rheology_nonNewtonian_plastic() - #dislocation laws - disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) - # diffusion laws - diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) - # plasticity - ϕ_wet_olivine = asind(0.1) - C_wet_olivine = 1e6 - η_reg = 1e16 - el = ConstantElasticity(; G = 40e9, ν = 0.45) - lithosphere_rheology = CompositeRheology( - ( - el, - disl_wet_olivine, - diff_wet_olivine, - DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity - ) - ) - init_rheologies(lithosphere_rheology) -end - -function init_rheology_linear() - el = ConstantElasticity(; G = 40e9, ν = 0.49) - # lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), )) - lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), el)) - init_rheologies(lithosphere_rheology) -end - -function init_rheologies(; linear=false) - - η_reg = 1e18 - C = linear ? Inf : 10e6 + η_reg = 1e15 + C = isplastic ? 10e6 : Inf ϕ = 15 + Ψ = 0.0 soft_C = NonLinearSoftening(; ξ₀=C, Δ = C / 1e5) # nonlinear softening law - # soft_C = NonLinearSoftening() # nonlinear softening law - pl = DruckerPrager_regularised(; C=C*MPa, ϕ=ϕ, η_vp=(η_reg)*Pas, Ψ=0.0, softening_C=soft_C) - el = ConstantElasticity(; G = 25e9, ν = 0.45) + pl = DruckerPrager_regularised(; C=C*MPa, ϕ=ϕ, η_vp=(η_reg)*Pas, Ψ=Ψ, softening_C=soft_C) + G0 = 25e9Pa # elastic shear modulus + G_magma = 10e9Pa # elastic shear modulus magma + + el = incompressible ? ConstantElasticity(; G = G0, ν = 0.45) : ConstantElasticity(; G = G0, ν = 0.25) + el_magma= incompressible ? ConstantElasticity(; G = G_magma, ν = 0.45) : ConstantElasticity(; G = G_magma, ν = 0.25) β = 1 / el.Kb.val + β_magma = 1 / el_magma.Kb.val Cp = 1200.0 - magma_visc = ViscosityPartialMelt_Costa_etal_2009(η=LinearMeltViscosity(A = -8.1590, B = 2.4050e+04K, T0 = -430.9606K,η0=1e3Pa*s)) + magma_visc = magma ? ViscosityPartialMelt_Costa_etal_2009(η=LinearMeltViscosity(A = -8.1590, B = 2.4050e+04K, T0 = -430.9606K)) : LinearViscous(η=1e15) + conduit_visc = magma ? ViscosityPartialMelt_Costa_etal_2009(η=LinearMeltViscosity(A = -8.1590, B = 2.4050e+04K, T0 = -430.9606K)) : LinearViscous(η=1e15) #dislocation laws - disl_top = DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) + disl_top = linear ? LinearViscous(η=1e23) : DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) # disl_top = SetDislocationCreep(Dislocation.dry_olivine_Karato_2003) # diffusion laws - disl_bot = SetDislocationCreep(Dislocation.wet_quartzite_Hirth_2001) + disl_bot = linear ? LinearViscous(η=1e21) : SetDislocationCreep(Dislocation.wet_quartzite_Hirth_2001) # Define rheolgy struct rheology = ( @@ -80,43 +46,48 @@ function init_rheologies(; linear=false) HeatCapacity = ConstantHeatCapacity(; Cp = Cp), Conductivity = ConstantConductivity(; k = 2.5), CompositeRheology = CompositeRheology( (disl_bot, el, pl)), - # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e22), el, pl)), - Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve + # CompositeRheology = CompositeRheology( (LinearViscous(; η=1e21), el, pl)), + Melting = MeltingParam_Smooth3rdOrder(a=517.9, b=-1619.0, c=1699.0, d = -597.4), #mafic melting curve Gravity = ConstantGravity(; g=9.81), ), # Name = "magma chamber", SetMaterialParams(; Phase = 3, - Density = T_Density(; ρ0=2.5e3, T0=273.15), + Density = MeltDependent_Density(ρsolid=PT_Density(ρ0=2.4e3, T0=273.15, β=β_magma), ρmelt=T_Density(ρ0=2.2e3, T0=273.15)), + # Density = PT_Density(; ρ0=2.4e3, T0=273.15, β=β_magma), Conductivity = ConstantConductivity(; k = 1.5), # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el)), + CompositeRheology = CompositeRheology( (magma_visc, el_magma)), Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), # Name = "magma chamber - hot anomaly", SetMaterialParams(; Phase = 4, - Density = T_Density(; ρ0=2.5e3, T0=273.15), + # Density = T_Density(; ρ0=2.2e3, T0=273.15), + # Density = BubbleFlow_Density(ρgas=ConstantDensity(ρ=10.0), ρmelt=ConstantDensity(ρ=2.4e3), c0=4e-2), + Density = BubbleFlow_Density(ρgas=ConstantDensity(ρ=10.0), ρmelt=MeltDependent_Density(ρsolid=T_Density(ρ0=2.4e3, T0=273.15), ρmelt= ConstantDensity(ρ=2.2e3)), c0=4e-2), Conductivity = ConstantConductivity(; k = 1.5), # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el, )), + CompositeRheology = CompositeRheology((magma_visc, el_magma, )), Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), # Name = "Conduit", SetMaterialParams(; Phase = 5, - Density = T_Density(; ρ0=1.5e3, T0=273.15), + Density = BubbleFlow_Density(ρgas=ConstantDensity(ρ=10.0), ρmelt=MeltDependent_Density(ρsolid=T_Density(ρ0=2.4e3, T0=273.15), ρmelt= ConstantDensity(ρ=2.2e3)), c0=4e-2), + # Density = BubbleFlow_Density(ρgas=ConstantDensity(ρ=10.0), ρmelt=ConstantDensity(ρ=2.4e3), c0=4e-2), + # Density = T_Density(; ρ0=1.5e3, T0=273.15), Conductivity = ConstantConductivity(; k = 1.5), # HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity()), HeatCapacity = Latent_HeatCapacity(Cp=ConstantHeatCapacity(), Q_L=350e3J/kg), LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - CompositeRheology = CompositeRheology( (LinearViscous(; η=1e18), el, )), + CompositeRheology = CompositeRheology((conduit_visc, el_magma,)), Melting = MeltingParam_Smooth3rdOrder(a=3043.0,b=-10552.0,c=12204.9,d=-4709.0), #felsic melting curve ), # Name = "StickyAir", diff --git a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl index d4676c40e..010bce971 100644 --- a/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl +++ b/miniapps/benchmarks/stokes2D/Volcano2D/Caldera_setup.jl @@ -8,7 +8,7 @@ function setup2D( flat = true, chimney = false, volcano_size = (3e0, 5e0), - conduit_radius = 0.25, + conduit_radius = 0.2, chamber_T = 1e3, chamber_depth = 5e0, chamber_radius = 2e0, @@ -21,10 +21,8 @@ function setup2D( z = range(-dimensions[2], sticky_air, nz); Grid = CartData(xyz_grid(x,y,z)); - # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) + # Allocate Phase and Temp arrays Phases = fill(6, nx, 2, nz); - - # In many (geodynamic) models, one also has to define the temperature, so lets define it as well Temp = fill(0.0, nx, 2, nz); add_box!(Phases, Temp, Grid; @@ -52,6 +50,13 @@ function setup2D( cen = (mean(Grid.x.val), 0, -chamber_depth), axes = (chamber_radius * aspect_x, 2.5, chamber_radius), phase = ConstantPhase(3), + T = ConstantTemp(T=chamber_T-100e0) + ) + + add_ellipsoid!(Phases, Temp, Grid; + cen = (mean(Grid.x.val), 0, -(chamber_depth-(chamber_radius/2))), + axes = ((chamber_radius/1.25) * aspect_x, 2.5, (chamber_radius/2)), + phase = ConstantPhase(4), T = ConstantTemp(T=chamber_T) ) @@ -68,8 +73,6 @@ function setup2D( cap = (mean(Grid.x.val), 0, flat ? 0e0 : volcano_size[1]), radius = conduit_radius, phase = ConstantPhase(5), - # T = LinearTemp(Ttop=20, Tbot=1000), - # T = ConstantTemp(T=800), T = ConstantTemp(T=chamber_T), ) end @@ -80,8 +83,16 @@ function setup2D( ph = Phases[:,1,:] T = Temp[:,1,:] .+ 273 - V = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * (chamber_radius*aspect_x) - printstyled("Magma volume of the initial chamber: $(round(V; digits=3)) km³ \n"; bold=true, color=:red, blink=true) + V_total = 4/3 * π * (chamber_radius*aspect_x) * chamber_radius * (chamber_radius*aspect_x) + V_erupt = 4/3 * π * (chamber_radius/1.25) * aspect_x * (chamber_radius/2) * ((chamber_radius/1.25) * aspect_x) + R = ((chamber_depth-chamber_radius))/(chamber_radius*aspect_x) + chamber_diameter = 2*(chamber_radius*aspect_x) + chamber_erupt = 2*((chamber_radius/1.25) * aspect_x) + printstyled("Magma volume of the initial chamber: $(round(V_total; digits=3)) km³ \n"; bold=true, color=:red, blink=true) + printstyled("Eruptible magma volume: $(round(V_erupt; digits=3)) km³ \n"; bold=true, color=:red, blink=true) + printstyled("Roof ratio (Depth/half-axis width): $R \n"; bold=true, color=:cyan) + printstyled("Chamber diameter: $chamber_diameter km \n"; bold=true, color=:light_yellow) + printstyled("Eruptible chamber diameter: $chamber_erupt km \n"; bold=true, color=:light_yellow) # write_paraview(Grid, "Volcano2D") return li, origin, ph, T, Grid end