Add adiabatic heating (#183)

* add adiabatic heating

* format

* fix indexing

src/rheology/GeoParams.jl

@@ -1,26 +1,20 @@
-function get_bulk_modulus(args...)
+function get_bulk_modulus(args::Vararg{Any,N}) where {N}
     Kb = GeoParams.get_Kb(args...)
     if isnan(Kb) || iszero(Kb)
         return Inf
     end
     return Kb
 end
 
-function get_shear_modulus(args...)
+function get_shear_modulus(args::Vararg{Any,N}) where {N}
     Kb = GeoParams.get_G(args...)
     if isnan(Kb) || iszero(Kb)
         return Inf
     end
     return Kb
 end
 
-function get_thermal_expansion(p)
-    α = get_α(p)
-    if isnan(α) || iszero(α)
-        return 0.0
-    end
-    return α
-end
+get_thermal_expansion(args::Vararg{Any,N}) where {N} = get_α(args...)
 
 @inline get_α(p::MaterialParams) = get_α(p.Density[1])
 @inline get_α(p::ConstantDensity) = 0.0

src/thermal_diffusion/DiffusionPT_GeoParams.jl

@@ -28,11 +28,20 @@ end
     return fn(rheology, phase, args)
 end
 
+@inline function compute_phase(fn::F, rheology, phase::Int) where {F}
+    return fn(rheology, phase, args)
+end
+
 @inline function compute_phase(fn::F, rheology, phase::SVector, args) where {F}
     return fn_ratio(fn, rheology, phase, args)
 end
 
+@inline function compute_phase(fn::F, rheology, phase::SVector) where {F}
+    return fn_ratio(fn, rheology, phase)
+end
+
 @inline compute_phase(fn::F, rheology, ::Nothing, args) where {F} = fn(rheology, args)
+@inline compute_phase(fn::F, rheology, ::Nothing) where {F} = fn(rheology)
 
 @inline Base.@propagate_inbounds function getindex_phase(
     phase::AbstractArray, I::Vararg{Int,N}
@@ -101,3 +110,13 @@ end
 @inline function compute_ρCp(rheology, ρ, phase_ratios::SArray, args)
     return fn_ratio(compute_heatcapacity, rheology, phase_ratios, args) * ρ
 end
+
+# α
+
+function compute_α(rheology, phase::SArray)
+    return fn_ratio(get_α, rheology, phase)
+end
+
+function compute_α(rheology, phase::Union{Int,Nothing})
+    return compute_phase(get_α, rheology, phase)
+end

src/thermal_diffusion/DiffusionPT_kernels.jl

@@ -146,6 +146,7 @@ end
     qTz,
     H,
     shear_heating,
+    adiabatic,
     rheology,
     phase,
     dτ_ρ,
@@ -172,7 +173,8 @@ end
             (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) -
             compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt +
             av(H) +
-            av(shear_heating)
+            av(shear_heating) +
+            adiabatic[i, j, k] * T_ijk
         )
     return nothing
 end
@@ -216,6 +218,7 @@ end
     qTz2,
     H,
     shear_heating,
+    adiabatic,
     rheology,
     phase,
     _dt,
@@ -238,7 +241,8 @@ end
         -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)
+        av(shear_heating) +
+        adiabatic[i, j, k] * T_ijk
 
     return nothing
 end
@@ -402,6 +406,7 @@ end
     qTy,
     H,
     shear_heating,
+    adiabatic,
     rheology,
     phase,
     dτ_ρ,
@@ -436,7 +441,8 @@ end
         av(dτ_ρ) * (
             (-(d_xa(qTx) + d_ya(qTy))) - ρCp * (T_ij - Told[i + 1, j + 1]) * _dt +
             av(H) +
-            av(shear_heating)
+            av(shear_heating) +
+            adiabatic[i, j] * T[i + 1, j + 1]
         )
     return nothing
 end
@@ -475,6 +481,7 @@ end
     qTy2,
     H,
     shear_heating,
+    adiabatic,
     rheology,
     phase,
     _dt,
@@ -507,7 +514,9 @@ end
     ResT[i, j] =
         -ρCp * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) +
         av(H) +
-        av(shear_heating)
+        av(shear_heating) +
+        adiabatic[i, j] * T[i + 1, j + 1]
+
     return nothing
 end
 
@@ -539,6 +548,7 @@ function update_T(
         @qT(thermal)...,
         thermal.H,
         thermal.shear_heating,
+        thermal.adiabatic,
         rheology,
         phase,
         pt_thermal.dτ_ρ,
@@ -547,3 +557,100 @@ function update_T(
         args,
     )
 end
+
+@parallel_indices (i, j, k) function adiabatic_heating(
+    A, Vx, Vy, Vz, P, rheology, phases, _dx, _dy, _dz
+)
+    I = i, j, k
+    I1 = i1, j1, k1 = I .+ 1
+    @inbounds begin
+        α =
+            (
+                compute_α(rheology, getindex_phase(phases, I...)) +
+                compute_α(rheology, getindex_phase(phases, i, j1, k)) +
+                compute_α(rheology, getindex_phase(phases, i1, j, k)) +
+                compute_α(rheology, getindex_phase(phases, i1, j1, k)) +
+                compute_α(rheology, getindex_phase(phases, i, j1, k1)) +
+                compute_α(rheology, getindex_phase(phases, i1, j, k1)) +
+                compute_α(rheology, getindex_phase(phases, i1, j1, k1)) +
+                compute_α(rheology, getindex_phase(phases, I1...))
+            ) * 0.125
+        # cache P around T node
+        P111 = P[I...]
+        P112 = P[i, j, k1]
+        P121 = P[i, j1, k]
+        P122 = P[i, j1, k1]
+        P211 = P[i1, j, k]
+        P212 = P[i1, j, k1]
+        P221 = P[i1, j1, k]
+        P222 = P[i1, j1, k1]
+        # P averages
+        Px_L = (P111 + P121 + P112 + P122) * 0.25
+        Px_R = (P211 + P221 + P212 + P222) * 0.25
+        Py_F = (P111 + P211 + P112 + P212) * 0.25
+        Py_B = (P121 + P221 + P122 + P222) * 0.25
+        Pz_B = (P111 + P211 + P121 + P221) * 0.25
+        Pz_T = (P112 + P212 + P122 + P222) * 0.25
+        # Vx average
+        Vx_av =
+            0.25 *
+            (Vx[I1...] + Vx[i1, j1, k1 + 1] + Vx[i1, j1 + 1, k1] + Vx[i1, j1 + 1, k1 + 1])
+        # Vy average
+        Vy_av =
+            0.25 *
+            (Vy[I1...] + Vy[i1 + 1, j1, k1] + Vy[i1, j1, k1 + 1] + Vy[i1 + 1, j1, k1 + 1])
+        # Vz average
+        Vz_av =
+            0.25 *
+            (Vz[I1...] + Vz[i1 + 1, j1, k1] + Vz[i1, j1 + 1, k1] + Vz[i1 + 1, j1 + 1, k1])
+        dPdx = Vx_av * (Px_R - Px_L) * _dx
+        dPdy = Vy_av * (Py_B - Py_F) * _dy
+        dPdz = Vz_av * (Pz_T - Pz_B) * _dz
+        A[I...] = (dPdx + dPdy + dPdz) * α
+    end
+    return nothing
+end
+
+@parallel_indices (i, j) function adiabatic_heating(
+    A, Vx, Vy, P, rheology, phases, _dx, _dy
+)
+    I = i, j
+    I1 = i1, j1 = I .+ 1
+    @inbounds begin
+        α =
+            (
+                compute_α(rheology, getindex_phase(phases, I...)) +
+                compute_α(rheology, getindex_phase(phases, i, j1)) +
+                compute_α(rheology, getindex_phase(phases, i1, j)) +
+                compute_α(rheology, getindex_phase(phases, I1...))
+            ) * 0.25
+        # cache P around T node
+        P11 = P[I...]
+        P12 = P[i, j1]
+        P21 = P[i1, j]
+        P22 = P[i1, j1]
+        # P averages
+        Px_L = (P11 + P12) * 0.5
+        Px_R = (P21 + P22) * 0.5
+        Py_T = (P12 + P22) * 0.5
+        Py_B = (P11 + P21) * 0.5
+        # Vx average
+        Vx_av = (Vx[I1...] + Vx[i1, j1 + 1]) * 0.5
+        # Vy average
+        Vy_av = (Vy[I1...] + Vy[i1 + 1, j1]) * 0.5
+        dPdx = (Px_R - Px_L) * _dx
+        dPdy = (Py_T - Py_B) * _dy
+        A[i1, j] = (Vx_av * dPdx + Vy_av * dPdy) * α
+    end
+    return nothing
+end
+
+function adiabatic_heating!(thermal, stokes, rheology, phases, di)
+    idx = @idx (size(stokes.P) .- 1)
+    _di = inv.(di)
+    @parallel idx adiabatic_heating(
+        thermal.adiabatic, @velocity(stokes)..., stokes.P, rheology, phases, _di...
+    )
+end
+
+@inline adiabatic_heating!(thermal, ::Nothing, ::Vararg{Any,N}) where {N} = nothing

src/thermal_diffusion/DiffusionPT_solver.jl

@@ -110,6 +110,7 @@ function _heatdiffusion_PT!(
     di;
     igg=nothing,
     phase=nothing,
+    stokes=nothing,
     b_width=(4, 4, 4),
     iterMax=50e3,
     nout=1e3,
@@ -127,6 +128,9 @@ function _heatdiffusion_PT!(
     @copy thermal.Told thermal.T
     !isnothing(phase) && update_pt_thermal_arrays!(pt_thermal, phase, rheology, args, _dt)
 
+    # compute constant part of the adiabatic heating term
+    adiabatic_heating!(thermal, stokes, rheology, phases, di)
+
     # errors
     iter_count = Int64[]
     norm_ResT = Float64[]
@@ -174,6 +178,7 @@ function _heatdiffusion_PT!(
                     @qT2(thermal)...,
                     thermal.H,
                     thermal.shear_heating,
+                    thermal.adiabatic,
                     rheology,
                     phases,
                     _dt,

src/types/constructors/heat_diffusion.jl

@@ -10,6 +10,7 @@ function ThermalArrays(nx::Integer, ny::Integer)
     T = @zeros(nx + 3, ny + 1)
     ΔT = @zeros(nx + 3, ny + 1)
     ΔTc = @zeros(nx, ny)
+    adiabatic = @zeros(nx + 1, ny - 1)
     Told = @zeros(nx + 3, ny + 1)
     Tc = @zeros(nx, ny)
     H = @zeros(nx, ny)
@@ -26,6 +27,7 @@ function ThermalArrays(nx::Integer, ny::Integer)
         Told,
         ΔT,
         ΔTc,
+        adiabatic,
         dT_dt,
         qTx,
         qTy,
@@ -43,6 +45,7 @@ function ThermalArrays(nx::Integer, ny::Integer, nz::Integer)
     T = @zeros(nx + 1, ny + 1, nz + 1)
     ΔT = @zeros(nx + 1, ny + 1, nz + 1)
     ΔTc = @zeros(nx, ny, ny)
+    adiabatic = @zeros(nx - 1, ny - 1, nz - 1)
     Told = @zeros(nx + 1, ny + 1, nz + 1)
     Tc = @zeros(nx, ny, nz)
     H = @zeros(nx, ny, nz)
@@ -56,6 +59,21 @@ function ThermalArrays(nx::Integer, ny::Integer, nz::Integer)
     qTz2 = @zeros(nx - 1, ny - 1, nz)
     ResT = @zeros(nx - 1, ny - 1, nz - 1)
     return JustRelax.ThermalArrays(
-        T, Tc, Told, ΔT, ΔTc, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT
+        T,
+        Tc,
+        Told,
+        ΔT,
+        ΔTc,
+        adiabatic,
+        dT_dt,
+        qTx,
+        qTy,
+        qTz,
+        qTx2,
+        qTy2,
+        qTz2,
+        H,
+        shear_heating,
+        ResT,
     )
 end

src/types/heat_diffusion.jl

@@ -4,6 +4,7 @@ struct ThermalArrays{_T}
     Told::_T
     ΔT::_T
     ΔTc::_T
+    adiabatic::_T # adiabatic term α (u ⋅ ∇P)
     dT_dt::_T
     qTx::_T
     qTy::_T