Merge pull request #141 from JuliaGeodynamics/adm/new_function

Fuse center and vertex phase ratio functions
JuliaGeodynamics · Sep 25, 2024 · 87fdb05 · 87fdb05
2 parents 5aef76d + bed9db2
commit 87fdb05
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 9 deletions.
diff --git a/ext/JustPICAMDGPUExt.jl b/ext/JustPICAMDGPUExt.jl
@@ -223,6 +223,12 @@ module _2D
 
     # Phase ratio kernels
 
+    function JustPIC._2D.update_phase_ratios!(phase_ratios::JustPIC.PhaseRatios{AMDGPUBackend}, particles, xci, xvi, phases)
+        phase_ratios_center!(phase_ratios, particles, xci, phases)
+        phase_ratios_vertex!(phase_ratios, particles, xvi, phases)
+        return nothing
+    end
+
     function JustPIC._2D.PhaseRatios(
         ::Type{AMDGPUBackend}, nphases::Integer, ni::NTuple{N,Integer}
     ) where {N}
@@ -455,6 +461,12 @@ module _3D
 
     # Phase ratio kernels
 
+    function JustPIC._3D.update_phase_ratios!(phase_ratios::JustPIC.PhaseRatios{AMDGPUBackend}, particles, xci, xvi, phases)
+        phase_ratios_center!(phase_ratios, particles, xci, phases)
+        phase_ratios_vertex!(phase_ratios, particles, xvi, phases)
+        return nothing
+    end
+
     function JustPIC._3D.PhaseRatios(
         ::Type{AMDGPUBackend}, nphases::Integer, ni::NTuple{N,Integer}
     ) where {N}

diff --git a/ext/JustPICCUDAExt.jl b/ext/JustPICCUDAExt.jl
@@ -217,6 +217,12 @@ module _2D
 
     # Phase ratio kernels
 
+    function JustPIC._2D.update_phase_ratios!(phase_ratios::JustPIC.PhaseRatios{CUDABackend}, particles, xci, xvi, phases)
+        phase_ratios_center!(phase_ratios, particles, xci, phases)
+        phase_ratios_vertex!(phase_ratios, particles, xvi, phases)
+        return nothing
+    end
+
     function JustPIC._2D.PhaseRatios(
         ::Type{CUDABackend}, nphases::Integer, ni::NTuple{N,Integer}
     ) where {N}
@@ -450,6 +456,12 @@ module _3D
 
     # Phase ratio kernels
 
+    function JustPIC._3D.update_phase_ratios!(phase_ratios::JustPIC.PhaseRatios{CUDABackend}, particles, xci, xvi, phases)
+        phase_ratios_center!(phase_ratios, particles, xci, phases)
+        phase_ratios_vertex!(phase_ratios, particles, xvi, phases)
+        return nothing
+    end
+
     function JustPIC._3D.PhaseRatios(
         ::Type{CUDABackend}, nphases::Integer, ni::NTuple{N,Integer}
     ) where {N}

diff --git a/src/PhaseRatios/kernels.jl b/src/PhaseRatios/kernels.jl
@@ -1,3 +1,9 @@
+function update_phase_ratios!(phase_ratios::JustPIC.PhaseRatios, particles, xci, xvi, phases)
+    phase_ratios_center!(phase_ratios, particles, xci, phases)
+    phase_ratios_vertex!(phase_ratios, particles, xvi, phases)
+    return nothing
+end
+
 ## Kernels to compute phase ratios at the centers
 
 function phase_ratios_center!(phase_ratios::JustPIC.PhaseRatios, particles, xci, phases)

diff --git a/src/common.jl b/src/common.jl
@@ -93,4 +93,4 @@ include("PhaseRatios/constructors.jl")
 export PhaseRatios
 
 include("PhaseRatios/kernels.jl")
-export phase_ratios_center!, phase_ratios_vertex!
+export update_phase_ratios!, phase_ratios_center!, phase_ratios_vertex!
diff --git a/test/test_CellArrays.jl b/test/test_CellArrays.jl
@@ -1,7 +1,7 @@
 using JustPIC, Test, StaticArrays
 
 @testset "CellArrays - 2D" begin
-    x = 1e0
+    x  = 1e0
     ni = (2, 2)
 
     ## Test a 2x2 grid with 2x1 CellArrays per grid cell
@@ -48,13 +48,13 @@ end
 
     nphases      = 5
     phases,      = JustPIC._2D.init_cell_arrays(particles, Val(1));
-    phases.data .= rand(1:nphases, size(phases.data));
+    T            = typeof(phases.data)
+    phases.data .= T(rand(1:nphases, size(phases.data)));
 
     phase_ratios = JustPIC._2D.PhaseRatios(backend, nphases, ni);
 
-    JustPIC._2D.phase_ratios_vertex!(phase_ratios, particles, xvi, phases) 
-    JustPIC._2D.phase_ratios_center!(phase_ratios, particles, xci, phases) 
-
+    JustPIC._2D.update_phase_ratios!(phase_ratios, particles, xci, xvi, phases)
+
     @test sum(phase_ratios.vertex.data) ≈ prod(ni.+1)
     @test sum(phase_ratios.center.data) ≈ prod(ni)
 end
@@ -108,12 +108,12 @@ end
 
     nphases      = 5
     phases,      = JustPIC._3D.init_cell_arrays(particles, Val(1));
-    phases.data .= rand(1:nphases, size(phases.data));
+    T            = typeof(phases.data)
+    phases.data .= T(rand(1:nphases, size(phases.data)));
 
     phase_ratios = JustPIC._3D.PhaseRatios(backend, nphases, ni);
 
-    JustPIC._3D.phase_ratios_vertex!(phase_ratios, particles, xvi, phases) 
-    JustPIC._3D.phase_ratios_center!(phase_ratios, particles, xci, phases) 
+    JustPIC._3D.update_phase_ratios!(phase_ratios, particles, xci, xvi, phases)
 
     @test sum(phase_ratios.vertex.data) ≈ prod(ni.+1)
     @test sum(phase_ratios.center.data) ≈ prod(ni)