Add phase ratio calculations on velocity nodes

ext/JustPICAMDGPUExt.jl

@@ -317,6 +317,12 @@ module _2D
         )
         return nothing
     end
+
+    function JustPIC._2D.phase_ratios_midpoint!(phase_midpoint, particles::Particles{AMDGPUBackend}, xci::NTuple{N}, phases, dimension) where N
+        phase_ratios_midpoint!(phase_midpoint, particles, xci, phases, dimension)
+        return nothing
+    end
+
 end
 
 module _3D
@@ -561,6 +567,11 @@ module _3D
         )
         return nothing
     end
+
+    function JustPIC._3D.phase_ratios_midpoint!(phase_midpoint, particles::Particles{AMDGPUBackend}, xci::NTuple{N}, phases, dimension) where N
+        phase_ratios_midpoint!(phase_midpoint, particles, xci, phases, dimension)
+        return nothing
+    end
 end
 
 end # module

ext/JustPICCUDAExt.jl

@@ -314,6 +314,10 @@ module _2D
         return nothing
     end
 
+    function JustPIC._2D.phase_ratios_midpoint!(phase_midpoint, particles::Particles{CUDABackend}, xci::NTuple{N}, phases, dimension) where N
+        phase_ratios_midpoint!(phase_midpoint, particles, xci, phases, dimension)
+        return nothing
+    end
 end
 
 module _3D
@@ -560,6 +564,11 @@ module _3D
         )
         return nothing
     end
+
+    function JustPIC._3D.phase_ratios_midpoint!(phase_midpoint, particles::Particles{CUDABackend}, xci::NTuple{N}, phases, dimension) where N
+        phase_ratios_midpoint!(phase_midpoint, particles, xci, phases, dimension)
+        return nothing
+    end
 end
 
 end # module

src/PhaseRatios/centers.jl

@@ -0,0 +1,29 @@
+## Kernels to compute phase ratios at the centers
+
+function phase_ratios_center!(phase_ratios::JustPIC.PhaseRatios, particles, xci, phases)
+    ni = size(phases)
+    di = compute_dx(xci)
+
+    @parallel (@idx ni) phase_ratios_center_kernel!(
+        phase_ratios.center, particles.coords, xci, di, phases
+    )
+    return nothing
+end
+
+@parallel_indices (I...) function phase_ratios_center_kernel!(
+    ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases
+) where {N,T1,T2,T3}
+
+    # index corresponding to the cell center
+    cell_center = ntuple(i -> xci[i][I[i]], Val(N))
+    # phase ratios weights (∑w = 1.0)
+    w = phase_ratio_weights(
+        getindex.(pxi, I...), phases[I...], cell_center, di, nphases(ratio_centers)
+    )
+    # update phase ratios array
+    for k in 1:numphases(ratio_centers)
+        @index ratio_centers[k, I...] = w[k]
+    end
+
+    return nothing
+end

src/PhaseRatios/midpoints.jl

@@ -0,0 +1,94 @@
+
+function phase_ratios_midpoint!(phase_midpoint, particles, xci::NTuple{N}, phases, dimension) where N
+    ni = size(phases)
+    di = compute_dx(xci)
+
+    offsets = if N == 2
+        offsets = if dimension === :x
+            (1, 0)
+        elseif dimension === :y
+            (0, 1)
+        else
+            throw("Unknown dimensions. Valid dimensions are :x, :y")
+        end
+    elseif N == 3
+        offsets = if dimension === :x
+            (1, 0, 0)
+        elseif dimension === :y
+            (0, 1, 0)
+        elseif dimension === :z
+            (0, 0, 1)
+        else
+            throw("Unknown dimensions. Valid dimensions are :x, :y, :z")
+        end
+    end
+
+    @parallel (@idx ni) phase_ratios_midpoint_kernel!(
+        phase_midpoint, particles.coords, xci, di, phases, offsets
+    )
+    return nothing
+end
+
+@parallel_indices (I...) function phase_ratios_midpoint_kernel!(
+    ratio_midpoints, pxi::NTuple{N}, xci::NTuple{N}, di::NTuple{N,T}, phases, offsets
+) where {N,T}
+
+    # index corresponding to the cell center
+    cell_center   = getindex.(xci, I)
+    cell_midpoint = @. cell_center + di * offsets / 2
+    ni = size(phases)
+    NC = nphases(ratio_midpoints)
+    w = ntuple(_ -> zero(T), NC)
+
+    # general case
+    for offsetsᵢ in (ntuple(_->0, Val(N)), offsets)
+        cell_index = min.(I .+ offsetsᵢ, ni)
+        all(@. 0 < cell_index < ni + 1) || continue
+
+        for ip in cellaxes(phases)
+            p = ntuple(Val(N)) do i
+                @index pxi[i][ip, cell_index...]
+            end
+            any(isnan, p) && continue
+            # check if it's within half cell
+            prod(x -> abs(x[1] - x[2]) ≤ x[3] / 2, zip(p, cell_midpoint, di)) && continue
+            x = @inline bilinear_weight(cell_midpoint, p, di)
+            ph_local = @index phases[ip, cell_index...]
+            # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta
+            w = w .+ x .* ntuple(j -> (ph_local == j), NC)
+        end
+    end
+
+    w = w .* inv(sum(w))
+    for ip in cellaxes(ratio_midpoints)
+        @index ratio_midpoints[ip, (I.+offsets)...] = w[ip]
+    end
+
+    # handle i == 1 or j == 1
+    boundary_offset = @.  1 * offsets * I
+    isboundary = any(isone, boundary_offset)
+    if isboundary
+        # index corresponding to the cell center
+        cell_midpoint = @. cell_center - di * offsets / 2
+        w = ntuple(_ -> zero(T), NC)
+
+        for ip in cellaxes(phases)
+            p = ntuple(Val(N)) do i
+                @index pxi[i][ip, I...]
+            end
+            any(isnan, p) && continue
+            # check if it's within half cell
+            prod(x -> abs(x[1] - x[2]) ≤ x[3] / 2, zip(p, cell_midpoint, di)) && continue
+            x = @inline bilinear_weight(cell_midpoint, p, di)
+            ph_local = @index phases[ip, I...]
+            # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta
+            w = w .+ x .* ntuple(j -> (ph_local == j), NC)
+        end
+        w = w .* inv(sum(w))
+        for ip in cellaxes(ratio_midpoints)
+            @index ratio_midpoints[ip, I...] = w[ip]
+        end
+    end
+
+    return nothing
+end

src/PhaseRatios/utils.jl

@@ -0,0 +1,40 @@
+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
+
+## interpolation kernels
+
+function phase_ratio_weights(
+    pxi::NTuple{NP,C}, ph::SVector{N1,T}, cell_center, di, ::Val{NC}
+) where {N1,NC,NP,T,C}
+
+    # Initiaze phase ratio weights (note: can't use ntuple() here because of the @generated function)
+    w = ntuple(_ -> zero(T), Val(NC))
+    # sumw = zero(T)
+
+    for i in eachindex(ph)
+        p = getindex.(pxi, i)
+        isnan(first(p)) && continue
+        x = @inline bilinear_weight(cell_center, p, di)
+        # sumw += x # reduce
+        ph_local = ph[i]
+        # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta
+        w = w .+ x .* ntuple(j -> (ph_local == j), Val(NC))
+    end
+    w = w .* inv(sum(w))
+    return w
+end
+
+@generated function bilinear_weight(
+    a::NTuple{N,T}, b::NTuple{N,T}, di::NTuple{N,T}
+) where {N,T}
+    quote
+        Base.@_inline_meta
+        val = one($T)
+        Base.Cartesian.@nexprs $N i ->
+            @inbounds val *= muladd(-abs(a[i] - b[i]), inv(di[i]), one($T))
+        return val
+    end
+end

src/PhaseRatios/kernels.jl → src/PhaseRatios/vertices.jl

@@ -1,39 +1,3 @@
-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)
-    ni = size(phases)
-    di = compute_dx(xci)
-
-    @parallel (@idx ni) phase_ratios_center_kernel!(
-        phase_ratios.center, particles.coords, xci, di, phases
-    )
-    return nothing
-end
-
-@parallel_indices (I...) function phase_ratios_center_kernel!(
-    ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases
-) where {N,T1,T2,T3}
-
-    # index corresponding to the cell center
-    cell_center = ntuple(i -> xci[i][I[i]], Val(N))
-    # phase ratios weights (∑w = 1.0)
-    w = phase_ratio_weights(
-        getindex.(pxi, I...), phases[I...], cell_center, di, nphases(ratio_centers)
-    )
-    # update phase ratios array
-    for k in 1:numphases(ratio_centers)
-        @index ratio_centers[k, I...] = w[k]
-    end
-
-    return nothing
-end
-
 ## Kernels to compute phase ratios at the vertices
 
 function phase_ratios_vertex!(phase_ratios::JustPIC.PhaseRatios, particles, xvi, phases)
@@ -123,39 +87,4 @@ end
     end
 
     return nothing
-end
-
-## interpolation kernels
-
-function phase_ratio_weights(
-    pxi::NTuple{NP,C}, ph::SVector{N1,T}, cell_center, di, ::Val{NC}
-) where {N1,NC,NP,T,C}
-
-    # Initiaze phase ratio weights (note: can't use ntuple() here because of the @generated function)
-    w = ntuple(_ -> zero(T), Val(NC))
-    # sumw = zero(T)
-
-    for i in eachindex(ph)
-        p = getindex.(pxi, i)
-        isnan(first(p)) && continue
-        x = @inline bilinear_weight(cell_center, p, di)
-        # sumw += x # reduce
-        ph_local = ph[i]
-        # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta
-        w = w .+ x .* ntuple(j -> (ph_local == j), Val(NC))
-    end
-    w = w .* inv(sum(w))
-    return w
-end
-
-@generated function bilinear_weight(
-    a::NTuple{N,T}, b::NTuple{N,T}, di::NTuple{N,T}
-) where {N,T}
-    quote
-        Base.@_inline_meta
-        val = one($T)
-        Base.Cartesian.@nexprs $N i ->
-            @inbounds val *= muladd(-abs(a[i] - b[i]), inv(di[i]), one($T))
-        return val
-    end
-end
+end

src/common.jl

@@ -92,5 +92,8 @@ export particle2grid!
 include("PhaseRatios/constructors.jl")
 export PhaseRatios
 
-include("PhaseRatios/kernels.jl")
-export update_phase_ratios!, phase_ratios_center!, phase_ratios_vertex!
+include("PhaseRatios/utils.jl")
+include("PhaseRatios/centers.jl")
+include("PhaseRatios/vertices.jl")
+include("PhaseRatios/midpoints.jl")
+export update_phase_ratios!, phase_ratios_center!, phase_ratios_vertex!, phase_ratios_midpoint!