update PT coefs inside iterations

PTsolvers · Aug 16, 2024 · d3b9e85 · d3b9e85
1 parent dcc79be
commit d3b9e85
Show file tree

Hide file tree

Showing 6 changed files with 121 additions and 0 deletions.
diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -85,6 +85,30 @@ function JR2D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, phase_ratios, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, ::Nothing, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR2D.flow_bcs!(
     ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions

diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -85,6 +85,30 @@ function JR3D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, phase_ratios, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, ::Nothing, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR3D.flow_bcs!(
     ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions

diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -84,6 +84,30 @@ function JR2D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR2D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR2D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR2D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR2D.flow_bcs!(
     ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions

diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -95,6 +95,30 @@ function JR3D.update_pt_thermal_arrays!(
     return nothing
 end
 
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR3D.flow_bcs!(
     ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions

diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl
@@ -121,3 +121,27 @@ function _compute_pt_thermal_arrays!(
 
     return nothing
 end
+
+function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, phase_ratios, args, dt)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, args, dt)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
+
+function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, ::Nothing, args, dt)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+    )
+    return nothing
+end
diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl
@@ -147,6 +147,7 @@ function _heatdiffusion_PT!(
 
     while err > ϵ && iter < iterMax
         wtime0 += @elapsed begin
+            update_thermal_coeffs!(pt_thermal, rheology, phase, args, dt)
             @parallel flux_range(ni...) compute_flux!(
                 @qT(thermal)...,
                 @qT2(thermal)...,