From a318545b5b59ebf65763fb2ce6ec076a0e178e9f Mon Sep 17 00:00:00 2001
From: albert-de-montserrat <albert.de.montserrat@gmail.com>
Date: Fri, 5 Apr 2024 10:51:17 +0200
Subject: [PATCH] thermal stress switch via multiple dispatch

---
 src/stokes/PressureKernels.jl | 31 ++++++++++++++++++++++++++++---
 src/stokes/Stokes2D.jl        |  6 +++---
 src/stokes/Stokes3D.jl        | 15 +++++++--------
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/stokes/PressureKernels.jl b/src/stokes/PressureKernels.jl
index 26275a68..f1fc9b85 100644
--- a/src/stokes/PressureKernels.jl
+++ b/src/stokes/PressureKernels.jl
@@ -31,10 +31,35 @@ Compute the pressure field `P` and the residual `RP` for the compressible case.
 after the implementation of Kiss et al. (2023). The temperature difference `ΔTc` on the cell center is used to compute this
 as well as α as the thermal expansivity.
 """
-@parallel_indices (I...) function compute_P!(
-    P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, args
+function compute_P!(
+    P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, dt, r, θ_dτ, kwargs::NamedTuple
+) where N
+    compute_P!(P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ; kwargs...)
+end
+
+function compute_P!(
+    P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, dt, r, θ_dτ; ΔTc = nothing, kwargs...
+) where N
+    ni = size(P)
+    @parallel (@idx ni) compute_P_kernel!(
+        P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ, ΔTc
+    )
+    return nothing
+end
+
+@parallel_indices (I...) function compute_P_kernel!(
+    P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, ::Nothing
+) where {N,C<:JustRelax.CellArray}
+    K = fn_ratio(get_bulk_modulus, rheology, phase_ratio[I...])
+    RP[I...], P[I...] = _compute_P!(
+        P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ
+    )
+    return nothing
+end
+
+@parallel_indices (I...) function compute_P_kernel!(
+    P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, ΔTc
 ) where {N,C<:JustRelax.CellArray}
-    ΔTc = args.ΔTc
     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!(
diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl
index 43fbf87d..c4ab0b78 100644
--- a/src/stokes/Stokes2D.jl
+++ b/src/stokes/Stokes2D.jl
@@ -561,7 +561,7 @@ function JustRelax.solve!(
 
             @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...)
 
-            @parallel (@idx ni) compute_P!(
+            compute_P!(
                 stokes.P,
                 stokes.P0,
                 stokes.R.RP,
@@ -620,8 +620,8 @@ function JustRelax.solve!(
                 dt,
                 θ_dτ,
             )
-            free_surface_bcs!(stokes)
-            # @views stokes.τ.yy[:, end] .= stokes.P[:, end]
+            free_surface_bcs!(stokes, flow_bcs)
+            
             @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c)
             update_halo!(stokes.τ.xy)
 
diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl
index 84b8153a..5d5506dd 100644
--- a/src/stokes/Stokes3D.jl
+++ b/src/stokes/Stokes3D.jl
@@ -417,19 +417,18 @@ function JustRelax.solve!(
 
     # solver loop
     wtime0 = 0.0
+    ητ = deepcopy(η)
+
     while iter < 2 || (err > ϵ && iter ≤ iterMax)
         wtime0 += @elapsed begin
             # ~preconditioner
-            ητ = deepcopy(η)
-            compute_maxloc!(ητ, η)
-            update_halo!(ητ)
-            # @hide_communication b_width begin # communication/computation overlap
-            #     @parallel compute_maxloc!(ητ, η)
-            #     update_halo!(ητ)
-            # end
+            @hide_communication b_width begin # communication/computation overlap
+                @parallel compute_maxloc!(ητ, η)
+                update_halo!(ητ)
+            end
 
             @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...)
-            @parallel (@idx ni) compute_P!(
+            compute_P!(
                 stokes.P,
                 stokes.P0,
                 stokes.R.RP,