format

PTsolvers · Aug 16, 2024 · de162b4 · de162b4
1 parent 626fd41
commit de162b4
Show file tree

Hide file tree

Showing 6 changed files with 153 additions and 31 deletions.
diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -85,26 +85,51 @@ 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}
+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)
+        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}
+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)
+        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}
+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)
+        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/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -85,26 +85,51 @@ 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}
+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)
+        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}
+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)
+        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}
+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)
+        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/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -84,26 +84,51 @@ 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}
+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)
+        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}
+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)
+        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}
+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)
+        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/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -95,26 +95,51 @@ 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}
+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)
+        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}
+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)
+        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}
+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)
+        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_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl
@@ -122,26 +122,49 @@ function _compute_pt_thermal_arrays!(
     return nothing
 end
 
-function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, phase_ratios, args, dt)
+function update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs, rheology, phase_ratios, args, dt
+)
     ni = size(pt_thermal.dτ_ρ)
     @parallel (@idx ni) 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)
+        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) compute_pt_thermal_arrays!(
-        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+        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)
+function update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs, rheology, ::Nothing, args, dt
+)
     ni = size(pt_thermal.dτ_ρ)
     @parallel (@idx ni) compute_pt_thermal_arrays!(
-        pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt)
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
     )
     return nothing
-end
+end
diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl
@@ -374,7 +374,6 @@ end
     return nothing
 end
 
-
 @parallel_indices (i, j) function update_T!(
     T::AbstractArray{_T,2}, Told, qTx, qTy, H, shear_heating, ρCp, dτ_ρ, _dt, _dx, _dy
 ) where {_T}