From 32685f67fc2c0428232d960cd3e1eff0ca5f0bd9 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Tue, 19 Nov 2024 13:46:02 +0100 Subject: [PATCH] Internal Dirichlet temperature boundary conditions (#273) * mask prototypes * up * more proto * start moving things into the source code * working inner T Dirichlet BCs * inner T BC miniapp * format * ext * quick patch * formatino * typo * Mask test * comment tests for now * comment Dirichlet test * typos * more typos --- mask.jl | 126 +++++++++++++ .../diffusion/diffusion2D_inner_BCs.jl | 173 ++++++++++++++++++ src/JustRelax.jl | 4 + src/JustRelax_CPU.jl | 8 +- src/boundaryconditions/Dirichlet.jl | 124 +++++++++++++ src/boundaryconditions/types.jl | 10 +- src/common.jl | 4 + src/ext/AMDGPU/2D.jl | 4 +- src/ext/AMDGPU/3D.jl | 4 +- src/ext/CUDA/2D.jl | 4 +- src/ext/CUDA/3D.jl | 4 +- src/mask/constructors.jl | 11 ++ src/mask/mask.jl | 55 ++++++ src/thermal_diffusion/DiffusionPT_kernels.jl | 107 ++++++++--- src/thermal_diffusion/DiffusionPT_solver.jl | 26 ++- test/test_boundary_conditions2D.jl | 78 ++++++++ test/test_mask.jl | 46 +++++ 17 files changed, 750 insertions(+), 38 deletions(-) create mode 100644 mask.jl create mode 100644 miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_inner_BCs.jl create mode 100644 src/boundaryconditions/Dirichlet.jl create mode 100644 src/mask/constructors.jl create mode 100644 src/mask/mask.jl create mode 100644 test/test_mask.jl diff --git a/mask.jl b/mask.jl new file mode 100644 index 000000000..10a9925be --- /dev/null +++ b/mask.jl @@ -0,0 +1,126 @@ +using MuladdMacro +abstract type AbstractMask end + +struct Mask{T} <: AbstractMask + mask::T + + Mask(A::T) where T<:AbstractArray = new{T}(A) +end + +Mask(ni::Vararg{Int,N}) where N = Mask(zeros(ni...)) + +function Mask(nx, ny, inds::Vararg{UnitRange,N}) where N + mask = zeros(nx, ny) + @views mask[inds...] .= 1 + return Mask(mask) +end + +function Mask(nx, ny, nz, inds::Vararg{UnitRange,N}) where N + mask = zeros(nx, ny, nz) + @views mask[inds...] .= 1 + return Mask(mask) +end + +Base.@propagate_inbounds Base.getindex(m::Mask{T}, inds::Vararg{Int,N}) where {T,N} = m.mask[inds...] +Base.@propagate_inbounds Base.getindex(m::Mask{T}, inds::Vararg{UnitRange{Int64},N}) where {T,N} = m.mask[inds...] + +Base.@propagate_inbounds Base.setindex!(m::Mask{T}, val, inds::Vararg{Int,N}) where {T,N} = setindex!(m.mask, val, inds...) +Base.@propagate_inbounds Base.setindex!(m::Mask{T}, val, inds::Vararg{UnitRange,N}) where {T,N} = setindex!(m.mask, val, inds...) + +Base.@propagate_inbounds Base.inv(m::Mask, inds::Vararg{Int,N}) where {N} = 1 - m[inds...] +Base.@propagate_inbounds Base.inv(m::Mask) = 1 .- m.mask + +Base.size(m::Mask) = size(m.mask) +Base.length(m::Mask) = length(m.mask) +Base.axes(m::Mask) = axes(m.mask) +Base.eachindex(m::Mask) = eachindex(m.mask) +Base.all(m::Mask) = all(isone, m.mask) +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) = (A .= inv(m) .* A .+ m.mask .* B) +@inline apply_mask!(::AbstractArray, ::AbstractArray, ::Nothing) = nothing +@inline apply_mask!(A::AbstractArray, B::AbstractArray, 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}) where {N} = nothing + +@inline apply_mask(A::AbstractArray, B::AbstractArray, m::Mask) = inv(m) .* A .+ m.mask .* B +@inline apply_mask(A::AbstractArray, B::AbstractArray, 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, ::AbstractArray, ::Nothing, inds::Vararg{Int,N}) where {N} = A[inds...] + +## +abstract type AbstractDirichletBoundaryCondition{T, M} end +struct DirichletBoundaryCondition{T, M} <: AbstractDirichletBoundaryCondition{T, M} + value::T + mask::M + + function DirichletBoundaryCondition(value::T, mask::M) where {T, M} + return new{T, M}(value, mask) + end + + function DirichletBoundaryCondition() + return new{Nothing, Nothing}(nothing, nothing) + end +end + +function DirichletBoundaryCondition(A::AbstractArray{T}) where T + m = Mask(size(A)...) + copyto!(m.mask, T.(A .!= 0)) + return DirichletBoundaryCondition(A, m) +end + +Base.getindex(x::DirichletBoundaryCondition{Nothing, Nothing}, ::Vararg{Int, N}) where N = 0 +Base.getindex(x::DirichletBoundaryCondition, inds::Vararg{Int, N}) where N = x.value[inds...] * x.mask[inds...] + +struct ConstantArray{T, N} <: AbstractArray{T, N} + val::T + + ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val) +end +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}:") + println(io, " ", A.val) +end + +function Base.show(io::IO, A::ConstantArray{T, N}) where {T, N} + println(io, "ConstantArray{$T,$N}:") + println(io, " ", A.val) +end + +struct ConstantDirichletBoundaryCondition{T, M} <: AbstractDirichletBoundaryCondition{T, M} + value::T + mask::M + + function ConstantDirichletBoundaryCondition(value::N, mask::M) where {N<:Number, M} + v = ConstantArray(value, Val(dims(mask))) + T = typeof(v) + return new{T, M}(v, mask) + end + + function ConstantDirichletBoundaryCondition() + return new{Nothing, Nothing}(nothing, nothing) + end +end + +Base.getindex(x::ConstantDirichletBoundaryCondition, inds::Vararg{Int, N}) where N = x.value * x.mask[inds...] +Base.getindex(x::ConstantDirichletBoundaryCondition{Nothing, Nothing}, ::Vararg{Int, N}) where N = 0 + +@inline apply_dirichlet!(A::AbstractArray, bc::AbstractDirichletBoundaryCondition) = apply_mask!(A, bc.value, bc.mask) +@inline apply_dirichlet!(::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing, Nothing}) = nothing +@inline apply_dirichlet!(A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N}) where {N} = apply_mask!(A, bc.value, bc.mask, inds...) +@inline apply_dirichlet!(::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing, Nothing}, ::Vararg{Int,N}) where {N} = nothing + +@inline apply_dirichlet(A::AbstractArray, bc::AbstractDirichletBoundaryCondition) = apply_mask(A, bc.value, bc.mask) +@inline apply_dirichlet(A::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing, Nothing}) = A +@inline apply_dirichlet(A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N}) where {N} = apply_mask(A, bc.value, bc.mask, inds...) +@inline apply_dirichlet(A::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing, Nothing}, inds::Vararg{Int,N}) where {N} = A[inds...] + +@inline Dirichlet(x::NamedTuple) = Dirichlet(; x...) +@inline Dirichlet(; values = nothing, mask = nothing) = Dirichlet(values, mask) +@inline Dirichlet(::Nothing, mask::Nothing) = DirichletBoundaryCondition() +@inline Dirichlet(::Nothing, mask::AbstractArray) = DirichletBoundaryCondition(mask) +@inline Dirichlet(values::Number, mask::AbstractArray) = ConstantDirichletBoundaryCondition(values, Mask(mask)) + diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_inner_BCs.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_inner_BCs.jl new file mode 100644 index 000000000..39f8f7cee --- /dev/null +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_inner_BCs.jl @@ -0,0 +1,173 @@ +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) + +using GeoParams +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend + +distance(p1, p2) = mapreduce(x->(x[1]-x[2])^2, +, zip(p1, p2)) |> sqrt + +@parallel_indices (i, j) function init_T!(T, z) + if z[j] == maximum(z) + T[i, j] = 300.0 + elseif z[j] == minimum(z) + T[i, j] = 3500.0 + else + T[i, j] = z[j] * (1900.0 - 1600.0) / minimum(z) + 1600.0 + end + return nothing +end + +function elliptical_perturbation!(T, mask, Ω_T, xc, yc, r, xvi) + + @parallel_indices (i, j) function _elliptical_perturbation!(T, mask, Ω_T, xc, yc, r, x, y) + @inbounds if (((x[i]-xc ))^2 + ((y[j] - yc))^2) ≤ r^2 + T[i, j] += Ω_T + mask[i, j] = 1 + end + return nothing + end + + @parallel _elliptical_perturbation!(T, mask, Ω_T, xc, yc, r, xvi...) +end + +function init_phases!(phases, particles, xc, yc, r) + ni = size(phases) + center = xc, yc + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, center, r) + @inbounds for ip in cellaxes(phases) + # quick escape + @index(index[ip, i, j]) == 0 && continue + + x = @index px[ip, i, j] + y = @index py[ip, i, j] + + # plume - rectangular + if (((x - center[1] ))^2 + ((y - center[2]))^2) ≤ r^2 + @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, center, r) +end + +@parallel_indices (I...) function compute_temperature_source_terms!(H, rheology, phase_ratios, args) + + args_ij = ntuple_idx(args, I...) + H[I...] = fn_ratio(compute_radioactive_heat, rheology, phase_ratios[I...], args_ij) + + return nothing +end + +function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) + kyr = 1e3 * 3600 * 24 * 365.25 + Myr = 1e3 * kyr + ttot = 1 * Myr # total simulation time + dt = 50 * kyr # physical time step + + init_mpi = JustRelax.MPI.Initialized() ? false : true + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + + # Physical domain + ni = (nx, ny) + li = (lx, ly) # domain length in x- and y- + di = @. li / ni # grid step in x- and -y + grid = Geometry(ni, li; origin = (0, -ly)) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + + # Define the thermal parameters with GeoParams + rheology = ( + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=3e3, β=0.0, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp0), + Conductivity = ConstantConductivity(; k=K0), + RadioactiveHeat = ConstantRadioactiveHeat(1e-6), + ), + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3.3e3, β=0.0, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp0), + Conductivity = ConstantConductivity(; k=K0), + RadioactiveHeat = ConstantRadioactiveHeat(1e-7), + ), + ) + + # fields needed to compute density on the fly + P = @zeros(ni...) + args = (; P=P) + + ## Allocate arrays needed for every Thermal Diffusion + thermal = ThermalArrays(backend_JR, ni) + @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[2]) + + # Add thermal perturbation + Ω_T = 1050e0 # inner BCs temperature + r = 10e3 # thermal perturbation radius + center_perturbation = lx/2, -ly/2 + mask = zeros(size(thermal.T)...) + elliptical_perturbation!(thermal.T, mask, Ω_T, center_perturbation..., r, xvi) + temperature2center!(thermal) + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 40, 40, 1 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi... + ) + pPhases, = init_cell_arrays(particles, Val(1)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) + init_phases!(pPhases, particles, center_perturbation..., r) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + @parallel (@idx ni) compute_temperature_source_terms!(thermal.H, rheology, phase_ratios.center, args) + + # PT coefficients for thermal diffusion + args = (; P=P, T=thermal.Tc) + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.97 / √2 + ) + + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + dirichlet = (constant = Ω_T, mask=mask) + ) + + # Time loop + t = 0.0 + it = 0 + nt = Int(ceil(ttot / dt)) + while it < nt + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs = (; + phase = phase_ratios, + iterMax = 1e3, + nout = 10, + ) + ) + + it += 1 + t += dt + end + + return (ni=ni, xci=xci, xvi=xvi, li=li, di=di), thermal +end + +thermal = diffusion_2D() + diff --git a/src/JustRelax.jl b/src/JustRelax.jl index d7bbbb1d4..16cae0317 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -32,6 +32,10 @@ include("types/heat_diffusion.jl") include("types/weno.jl") +include("mask/mask.jl") + +include("boundaryconditions/Dirichlet.jl") + include("boundaryconditions/types.jl") export TemperatureBoundaryConditions, DisplacementBoundaryConditions, VelocityBoundaryConditions diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 5022bdb59..5add960f4 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -17,7 +17,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._2D: numphases, nphases @@ -48,7 +50,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._3D: numphases, nphases diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl new file mode 100644 index 000000000..0c8d259fe --- /dev/null +++ b/src/boundaryconditions/Dirichlet.jl @@ -0,0 +1,124 @@ +abstract type AbstractDirichletBoundaryCondition{T,M} end +struct DirichletBoundaryCondition{T,M} <: AbstractDirichletBoundaryCondition{T,M} + value::T + mask::M + + function DirichletBoundaryCondition(value::T, mask::M) where {T,M} + return new{T,M}(value, mask) + end + + function DirichletBoundaryCondition() + return new{Nothing,Nothing}(nothing, nothing) + end +end + +Adapt.@adapt_structure DirichletBoundaryCondition + +function DirichletBoundaryCondition(A::AbstractArray{T}) where {T} + m = Mask(size(A)...) + copyto!(m.mask, T.(A .!= 0)) + return DirichletBoundaryCondition(A, m) +end + +Base.getindex(x::DirichletBoundaryCondition{Nothing,Nothing}, ::Vararg{Int,N}) where {N} = 0 +function Base.getindex(x::DirichletBoundaryCondition, inds::Vararg{Int,N}) where {N} + 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 + +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}:") + return println(io, " ", A.val) +end + +function Base.show(io::IO, A::ConstantArray{T,N}) where {T,N} + println(io, "ConstantArray{$T,$N}:") + return println(io, " ", A.val) +end + +struct ConstantDirichletBoundaryCondition{T,M} <: AbstractDirichletBoundaryCondition{T,M} + value::T + mask::M + + function ConstantDirichletBoundaryCondition(value::N, mask::M) where {N<:Number,M} + v = ConstantArray(value, Val(dims(mask))) + T = typeof(v) + return new{T,M}(v, mask) + end + + function ConstantDirichletBoundaryCondition() + return new{Nothing,Nothing}(nothing, nothing) + end +end + +Adapt.@adapt_structure ConstantDirichletBoundaryCondition + +function Base.getindex(x::ConstantDirichletBoundaryCondition, inds::Vararg{Int,N}) where {N} + return x.value * x.mask[inds...] +end +function Base.getindex( + x::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) +end + +@inline function apply_dirichlet!( + ::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing} +) + return nothing +end + +@inline function apply_dirichlet!( + A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N} +) where {N} + return apply_mask!(A, bc.value, bc.mask, inds...) +end + +@inline function apply_dirichlet!( + ::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing}, ::Vararg{Int,N} +) where {N} + return nothing +end + +@inline function apply_dirichlet(A::AbstractArray, bc::AbstractDirichletBoundaryCondition) + return apply_mask(A, bc.value, bc.mask) +end + +@inline function apply_dirichlet( + A::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing} +) + return A +end + +@inline function apply_dirichlet( + A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N} +) where {N} + return apply_mask(A, bc.value, bc.mask, inds...) +end + +@inline function apply_dirichlet( + A::AbstractArray, + ::AbstractDirichletBoundaryCondition{Nothing,Nothing}, + inds::Vararg{Int,N}, +) where {N} + return A[inds...] +end + +@inline Dirichlet(x::NamedTuple) = Dirichlet(; x...) +@inline Dirichlet(; constant=nothing, mask=nothing) = Dirichlet(constant, mask) +@inline Dirichlet(::Nothing, mask::Nothing) = DirichletBoundaryCondition() +@inline Dirichlet(::Nothing, mask::AbstractArray) = DirichletBoundaryCondition(mask) +@inline Dirichlet(constant::Number, mask::AbstractArray) = + ConstantDirichletBoundaryCondition(constant, Mask(mask)) diff --git a/src/boundaryconditions/types.jl b/src/boundaryconditions/types.jl index fbb2c0a50..fc4cbd54d 100644 --- a/src/boundaryconditions/types.jl +++ b/src/boundaryconditions/types.jl @@ -1,13 +1,15 @@ abstract type AbstractBoundaryConditions end abstract type AbstractFlowBoundaryConditions <: AbstractBoundaryConditions end -struct TemperatureBoundaryConditions{T,nD} <: AbstractBoundaryConditions +struct TemperatureBoundaryConditions{T,D,nD} <: AbstractBoundaryConditions no_flux::T - + dirichlet::D function TemperatureBoundaryConditions(; - no_flux::T=(left=true, right=false, top=false, bot=false) + no_flux::T=(left=true, right=false, top=false, bot=false), + dirichlet=(; constant=nothing, mask=nothing), ) where {T} + D = Dirichlet(dirichlet) nD = length(no_flux) == 4 ? 2 : 3 - return new{T,nD}(no_flux) + return new{T,typeof(D),nD}(no_flux, D) end end diff --git a/src/common.jl b/src/common.jl index 18117e31c..d3ccc3d4d 100644 --- a/src/common.jl +++ b/src/common.jl @@ -1,3 +1,5 @@ +using Adapt, MuladdMacro + include("types/constructors/stokes.jl") export StokesArrays, PTStokesCoeffs @@ -33,6 +35,8 @@ export @allocate, include("types/displacement.jl") export velocity2displacement!, displacement2velocity! +include("mask/constructors.jl") + include("boundaryconditions/BoundaryConditions.jl") export AbstractBoundaryConditions, TemperatureBoundaryConditions, diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 33e999e92..5db87d39c 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -27,7 +27,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._2D: nphases, numphases diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index ddc8432b5..b2e40b70b 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -27,7 +27,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._3D: numphases, nphases diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index f23d81b0a..fd0f2a6fb 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -20,7 +20,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._2D: numphases, nphases diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 9eb49eeec..6cc0605b4 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -19,7 +19,9 @@ import JustRelax: TemperatureBoundaryConditions, AbstractFlowBoundaryConditions, DisplacementBoundaryConditions, - VelocityBoundaryConditions + VelocityBoundaryConditions, + apply_dirichlet, + apply_dirichlet! import JustPIC._3D: numphases, nphases diff --git a/src/mask/constructors.jl b/src/mask/constructors.jl new file mode 100644 index 000000000..27269b2b2 --- /dev/null +++ b/src/mask/constructors.jl @@ -0,0 +1,11 @@ +function Mask(nx, ny, inds::Vararg{UnitRange,2}) + mask = @zeros(nx, ny) + @views mask[inds...] .= 1 + return JustRelax.Mask(mask) +end + +function Mask(nx, ny, nz, inds::Vararg{UnitRange,3}) + mask = @zeros(nx, ny, nz) + @views mask[inds...] .= 1 + return JustRelax.Mask(mask) +end diff --git a/src/mask/mask.jl b/src/mask/mask.jl new file mode 100644 index 000000000..6147cefe9 --- /dev/null +++ b/src/mask/mask.jl @@ -0,0 +1,55 @@ +using MuladdMacro +abstract type AbstractMask end + +struct Mask{T} <: AbstractMask + mask::T + + Mask(A::T) where {T<:AbstractArray} = new{T}(A) +end + +Mask(ni::Vararg{Int,N}) where {N} = Mask(zeros(ni...)) + +Adapt.@adapt_structure Mask + +Base.@propagate_inbounds Base.getindex(m::Mask{T}, inds::Vararg{Int,N}) where {T,N} = + m.mask[inds...] +Base.@propagate_inbounds Base.getindex( + m::Mask{T}, inds::Vararg{UnitRange{Int64},N} +) where {T,N} = m.mask[inds...] + +Base.@propagate_inbounds Base.setindex!(m::Mask{T}, val, inds::Vararg{Int,N}) where {T,N} = + setindex!(m.mask, val, inds...) +Base.@propagate_inbounds Base.setindex!( + m::Mask{T}, val, inds::Vararg{UnitRange,N} +) where {T,N} = setindex!(m.mask, val, inds...) + +Base.@propagate_inbounds Base.inv(m::Mask, inds::Vararg{Int,N}) where {N} = 1 - m[inds...] +Base.@propagate_inbounds Base.inv(m::Mask) = 1 .- m.mask + +Base.size(m::Mask) = size(m.mask) +Base.length(m::Mask) = length(m.mask) +Base.axes(m::Mask) = axes(m.mask) +Base.eachindex(m::Mask) = eachindex(m.mask) +Base.all(m::Mask) = all(isone, m.mask) +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) = + (A .= inv(m) .* A .+ m.mask .* B) +@inline apply_mask!(::AbstractArray, ::AbstractArray, ::Nothing) = nothing +@inline apply_mask!( + A::AbstractArray, B::AbstractArray, 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} +) where {N} = nothing + +@inline apply_mask(A::AbstractArray, B::AbstractArray, m::Mask) = inv(m) .* A .+ m.mask .* B +@inline apply_mask( + A::AbstractArray, B::AbstractArray, 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, ::AbstractArray, ::Nothing, inds::Vararg{Int,N} +) where {N} = A[inds...] diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl index 31a58b3fc..f22ff8e73 100644 --- a/src/thermal_diffusion/DiffusionPT_kernels.jl +++ b/src/thermal_diffusion/DiffusionPT_kernels.jl @@ -1,3 +1,5 @@ +isNotDirichlet(m, inds::Vararg{Int,N}) where {N} = iszero(m[inds...]) +isNotDirichlet(::Nothing, ::Vararg{Int,N}) where {N} = false ## 3D KERNELS @@ -116,6 +118,7 @@ end shear_heating, ρCp, dτ_ρ, + dirichlet, _dt, _dx, _dy, @@ -136,6 +139,8 @@ end av(shear_heating) ) + T[I1...] ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt) + apply_dirichlet!(T, dirichlet, I1...) + return nothing end @@ -151,6 +156,7 @@ end rheology, phase, dτ_ρ, + dirichlet, _dt, _dx, _dy, @@ -162,23 +168,25 @@ end d_ya(A) = _d_ya(A, i, j, k, _dy) d_za(A) = _d_za(A, i, j, k, _dz) - I = i + 1, j + 1, k + 1 + I1 = i + 1, j + 1, k + 1 - T_ijk = T[I...] + T_ijk = T[I1...] args_ijk = (; T=T_ijk, P=av(args.P)) phase_ijk = getindex_phase(phase, i, j, k) ρCp = compute_ρCp(rheology, phase_ijk, args_ijk) - T[I...] = + T[I1...] = ( av(dτ_ρ) * ( -(d_xa(qTx) + d_ya(qTy) + d_za(qTz)) + - Told[I...] * ρCp * _dt + + Told[I1...] * ρCp * _dt + av(H) + av(shear_heating) + adiabatic[i, j, k] * T_ijk ) + T_ijk ) / (one(_T) + av(dτ_ρ) * ρCp * _dt) + apply_dirichlet!(T, dirichlet, I1...) + return nothing end @@ -192,6 +200,7 @@ end H, shear_heating, ρCp, + dirichlet, _dt, _dx, _dy, @@ -202,12 +211,15 @@ end d_za(A) = _d_za(A, i, j, k, _dz) av(A) = _av(A, i, j, k) - I = i + 1, j + 1, k + 1 + I1 = i + 1, j + 1, k + 1 - ResT[i, j, k] = - -av(ρCp) * (T[I...] - Told[I...]) * _dt - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + + ResT[i, j, k] = if isNotDirichlet(dirichlet.mask, I1...) + -av(ρCp) * (T[I1...] - Told[I1...]) * _dt - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + av(H) + av(shear_heating) + else + zero(_T) + end return nothing end @@ -224,6 +236,7 @@ end adiabatic, rheology, phase, + dirichlet, _dt, _dx, _dy, @@ -240,13 +253,15 @@ end args_ijk = (; T=T_ijk, P=av(args.P)) phase_ijk = getindex_phase(phase, i, j, k) - ResT[i, j, k] = + ResT[i, j, k] = if isNotDirichlet(dirichlet.mask, I...) -compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + av(H) + av(shear_heating) + adiabatic[i, j, k] * T_ijk - + else + zero(_T) + end return nothing end @@ -375,7 +390,18 @@ end end @parallel_indices (i, j) function update_T!( - T::AbstractArray{_T,2}, Told, qTx, qTy, H, shear_heating, ρCp, dτ_ρ, _dt, _dx, _dy + T::AbstractArray{_T,2}, + Told, + qTx, + qTy, + H, + shear_heating, + ρCp, + dτ_ρ, + dirichlet, + _dt, + _dx, + _dy, ) where {_T} nx, ny = size(ρCp) @@ -392,16 +418,19 @@ end end #! format: on - T[i + 1, j + 1] = + I1 = i + 1, j + 1 + T[I1...] = ( av(dτ_ρ) * ( -(d_xa(qTx) + d_ya(qTy)) + - Told[i + 1, j + 1] * av(ρCp) * _dt + + Told[I1...] * av(ρCp) * _dt + av(H) + av(shear_heating) - ) + T[i + 1, j + 1] + ) + T[I1...] ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt) + apply_dirichlet!(T, dirichlet, I1...) + return nothing end @@ -416,6 +445,7 @@ end rheology, phase, dτ_ρ, + dirichlet, _dt, _dx, _dy, @@ -443,22 +473,36 @@ end compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) ) * 0.25 - T[i + 1, j + 1] = + I1 = i + 1, j + 1 + T[I1...] = ( av(dτ_ρ) * ( -(d_xa(qTx) + d_ya(qTy)) + - Told[i + 1, j + 1] * ρCp * _dt + + Told[I1...] * ρCp * _dt + av(H) + av(shear_heating) + - adiabatic[i, j] * T[i + 1, j + 1] - ) + T[i + 1, j + 1] + adiabatic[i, j] * T[I1...] + ) + T[I1...] ) / (one(_T) + av(dτ_ρ) * ρCp * _dt) + apply_dirichlet!(T, dirichlet, I1...) + return nothing end @parallel_indices (i, j) function check_res!( - ResT::AbstractArray{_T,2}, T, Told, qTx2, qTy2, H, shear_heating, ρCp, _dt, _dx, _dy + ResT::AbstractArray{_T,2}, + T, + Told, + qTx2, + qTy2, + H, + shear_heating, + ρCp, + dirichlet, + _dt, + _dx, + _dy, ) where {_T} nx, ny = size(ρCp) @@ -475,11 +519,15 @@ end end #! format: on - ResT[i, j] = - -av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - - (d_xa(qTx2) + d_ya(qTy2)) + + I1 = i + 1, j + 1 + ResT[i, j] = if isNotDirichlet(dirichlet.mask, I1...) + -av(ρCp) * (T[I1...] - Told[I1...]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) + av(H) + av(shear_heating) + else + zero(_T) + end + return nothing end @@ -494,6 +542,7 @@ end adiabatic, rheology, phase, + dirichlet, _dt, _dx, _dy, @@ -521,11 +570,15 @@ end compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) ) * 0.25 - ResT[i, j] = - -ρCp * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) + + I1 = i + 1, j + 1 + ResT[i, j] = if isNotDirichlet(dirichlet.mask, I1...) + -ρCp * (T[I1...] - Told[I1...]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) + av(H) + av(shear_heating) + - adiabatic[i, j] * T[i + 1, j + 1] + adiabatic[i, j] * T[I1...] + else + zero(_T) + end return nothing end @@ -535,7 +588,7 @@ end return nothing end -function update_T(::Nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) +function update_T(::Nothing, b_width, thermal, ρCp, pt_thermal, dirichlet, _dt, _di, ni) @parallel update_range(ni...) update_T!( thermal.T, thermal.Told, @@ -544,13 +597,14 @@ function update_T(::Nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) thermal.shear_heating, ρCp, pt_thermal.dτ_ρ, + dirichlet, _dt, _di..., ) end function update_T( - ::Nothing, b_width, thermal, rheology, phase, pt_thermal, _dt, _di, ni, args + ::Nothing, b_width, thermal, rheology, phase, pt_thermal, dirichlet, _dt, _di, ni, args ) @parallel update_range(ni...) update_T!( thermal.T, @@ -562,6 +616,7 @@ function update_T( rheology, phase, pt_thermal.dτ_ρ, + dirichlet, _dt, _di..., args, diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index 5f2bb8862..14c872455 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -53,7 +53,17 @@ function _heatdiffusion_PT!( @parallel flux_range(ni...) compute_flux!( @qT(thermal)..., @qT2(thermal)..., thermal.T, K, pt_thermal.θr_dτ, _di... ) - update_T(nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) + update_T( + nothing, + b_width, + thermal, + ρCp, + pt_thermal, + thermal_bc.dirichlet, + _dt, + _di, + ni, + ) thermal_bcs!(thermal, thermal_bc) update_halo!(thermal.T) end @@ -70,6 +80,7 @@ function _heatdiffusion_PT!( thermal.H, thermal.shear_heating, ρCp, + thermal_bc.dirichlet, _dt, _di..., ) @@ -159,7 +170,17 @@ function _heatdiffusion_PT!( args, ) update_T( - nothing, b_width, thermal, rheology, phases, pt_thermal, _dt, _di, ni, args + nothing, + b_width, + thermal, + rheology, + phases, + pt_thermal, + thermal_bc.dirichlet, + _dt, + _di, + ni, + args, ) thermal_bcs!(thermal, thermal_bc) update_halo!(thermal.T) @@ -182,6 +203,7 @@ function _heatdiffusion_PT!( thermal.adiabatic, rheology, phases, + thermal_bc.dirichlet, _dt, _di..., args, diff --git a/test/test_boundary_conditions2D.jl b/test/test_boundary_conditions2D.jl index 285c8e609..6a8e72138 100644 --- a/test/test_boundary_conditions2D.jl +++ b/test/test_boundary_conditions2D.jl @@ -174,5 +174,83 @@ end @test true === true end end + + # @testset "DirichletBoundaryCondition" begin + # ni = 10, 10 + # A = rand(ni...) + # value = zeros(ni...) + # value[4:7, 4:7] .= 5 + # mask = Float64.(value .== 5) + + # bc = JustRelax.DirichletBoundaryCondition(value, mask) + + # @test all(JustRelax.apply_dirichlet(A, bc)[4:7,4:7] .== 5) + + # apply_mask!(A, bc) + # @test all(A[4:7,4:7] .== 5) + + # A = rand(ni...) + # @test apply_dirichlet(A, bc, 1, 1) == A[1,1] + # @test apply_dirichlet(A, bc, 5, 5) == 5 + + # apply_mask!(A, bc, 1, 1) + # apply_mask!(A, bc, 5, 5) + + # @test A[1, 1] !== 5 + # @test A[5, 5] == 5 + + # bc = DirichletBoundaryCondition() + + # @test all(apply_dirichlet(A, bc) == A) + + # apply_mask!(A, bc) + # @test all(A[4:7,4:7] .!== 5) + + # @test apply_dirichlet(A, bc, 1, 1) == A[1, 1] + # @test apply_dirichlet(A, bc, 5, 5) == A[5, 5] + + # apply_mask!(A, bc, 1, 1) + # apply_mask!(A, bc, 5, 5) + + # @test A[1, 1] !== 5 + # @test A[5, 5] !== 5 + # end + + # @testset "ConstantDirichletBoundaryCondition" begin + + # ni = 10, 10 + # A = rand(ni...) + # value = 5e0 + # mask = Mask(ni..., 4:7, 4:7) + + # bc = ConstantDirichletBoundaryCondition(value, mask) + + # @test apply_dirichlet(A, bc, 1, 1) == A[1,1] + # @test apply_dirichlet(A, bc, 5, 5) == 5 + + # A .= rand(ni...) + # apply_mask!(A, bc, 1, 1) + # apply_mask!(A, bc, 5, 5) + + # @test A[1, 1] !== 5 + # @test A[5, 5] == 5 + + # bc = ConstantDirichletBoundaryCondition() + + # @test all(apply_dirichlet(A, bc) == A) + + # apply_mask!(A, bc) + # @test any(A .== 5) + + # @test apply_dirichlet(A, bc, 1, 1) == A[1, 1] + # @test apply_dirichlet(A, bc, 5, 5) == A[5, 5] + + # A .= rand(ni...) + # apply_mask!(A, bc, 1, 1) + # apply_mask!(A, bc, 5, 5) + + # @test A[1, 1] !== 5 + # @test A[5, 5] !== 5 + # end end end diff --git a/test/test_mask.jl b/test/test_mask.jl new file mode 100644 index 000000000..ecd2f056d --- /dev/null +++ b/test/test_mask.jl @@ -0,0 +1,46 @@ +using Test +using JustRelax, JustRelax.JustRelax2D +import JustRelax: Mask, apply_mask, apply_mask! +import JustRelax.JustRelax2D as JR2 + +@testset "Mask" begin + ni = 10, 10 + + # Test basics + m = Mask(ni...) + + @test size(m) == ni + @test length(m) == prod(ni) + @test axes(m) == (Base.OneTo(ni[1]), Base.OneTo(ni[2])) + @test eachindex(m) == Base.OneTo(prod(ni)) + @test similar(m) isa Mask{Matrix{Float64}} + @test !all(m) + + m = JR2.Mask(ni..., 4:7, 4:7) + @test all(isone, m[4:7, 4:7]) + + # test masking + A = rand(ni...) + B = zeros(ni...) + B[4:7, 4:7] .= 5 + m = JR2.Mask(ni..., 4:7, 4:7) + + C = apply_mask(A, B, m) + @test all(C[4:7, 4:7] .== 5) + + apply_mask!(A, B, m) + @test all(A[4:7, 4:7] .== 5) + + A = rand(ni...) + @test apply_mask(A, B, m, 1, 1) == A[1, 1] + @test apply_mask(A, B, m, 5, 5) == 5 + + apply_mask!(A, B, m, 1, 1) + apply_mask!(A, B, m, 5, 5) + + @test A[1, 1] != 5 + @test A[5, 5] == 5 + + @test isone(inv(m, 1, 1)) + @test iszero(inv(m, 5, 5)) +end \ No newline at end of file