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

@@ -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))
+

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()
+

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

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