Merge pull request #165 from JuliaGeodynamics/adm/subdiff

Improve interpolatio to/from grid centroids
JuliaGeodynamics · Oct 28, 2024 · 252aa94 · 252aa94
2 parents 7fc97ae + 2e51b34
commit 252aa94
Show file tree

Hide file tree

Showing 8 changed files with 336 additions and 327 deletions.
diff --git a/ext/JustPICAMDGPUExt.jl b/ext/JustPICAMDGPUExt.jl
@@ -97,8 +97,8 @@ module _2D
         return Particles(AMDGPUBackend, coords, index, nxcell, max_xcell, min_xcell, np)
     end
 
-    function JustPIC._2D.SubgridDiffusionCellArrays(particles::Particles{AMDGPUBackend})
-        return SubgridDiffusionCellArrays(particles)
+    function JustPIC._2D.SubgridDiffusionCellArrays(particles::Particles{AMDGPUBackend}; loc::Symbol=:vertex)
+        return SubgridDiffusionCellArrays(particles; loc=loc)
     end
 
     function JustPIC._2D.init_particles(
@@ -217,6 +217,21 @@ module _2D
         return nothing
     end
 
+    function JustPIC._2D.subgrid_diffusion_centroid!(
+        pT,
+        T_grid,
+        ΔT_grid,
+        subgrid_arrays,
+        particles::Particles{AMDGPUBackend},
+        xvi,
+        di,
+        dt;
+        d=1.0,
+    )
+        subgrid_diffusion_centroid!(pT, T_grid, ΔT_grid, subgrid_arrays, particles, xvi, di, dt; d=d)
+        return nothing
+    end
+
     ## MakerChain
 
     function JustPIC._2D.init_markerchain(::Type{AMDGPUBackend}, nxcell, min_xcell, max_xcell, xv, initial_elevation)
@@ -378,8 +393,8 @@ module _3D
         return Particles(AMDGPUBackend, coords, index, nxcell, max_xcell, min_xcell, np)
     end
 
-    function JustPIC._3D.SubgridDiffusionCellArrays(particles::Particles{AMDGPUBackend})
-        return SubgridDiffusionCellArrays(particles)
+    function JustPIC._3D.SubgridDiffusionCellArrays(particles::Particles{AMDGPUBackend}; loc::Symbol=:vertex)
+        return SubgridDiffusionCellArrays(particles; loc=loc)
     end
 
     function JustPIC._3D.init_particles(
@@ -496,6 +511,21 @@ module _3D
         return nothing
     end
 
+    function JustPIC._3D.subgrid_diffusion_centroid!(
+        pT,
+        T_grid,
+        ΔT_grid,
+        subgrid_arrays,
+        particles::Particles{AMDGPUBackend},
+        xvi,
+        di,
+        dt;
+        d=1.0,
+    )
+        subgrid_diffusion_centroid!(pT, T_grid, ΔT_grid, subgrid_arrays, particles, xvi, di, dt; d=d)
+        return nothing
+    end
+
     ## PassiveMarkers
 
     function JustPIC._3D.init_passive_markers(

diff --git a/ext/JustPICCUDAExt.jl b/ext/JustPICCUDAExt.jl
@@ -99,8 +99,8 @@ module _2D
         return Particles(CUDABackend, coords, index, nxcell, max_xcell, min_xcell, np)
     end
 
-    function JustPIC._2D.SubgridDiffusionCellArrays(particles::Particles{CUDABackend})
-        return SubgridDiffusionCellArrays(particles)
+    function JustPIC._2D.SubgridDiffusionCellArrays(particles::Particles{CUDABackend}; loc::Symbol=:vertex)
+        return SubgridDiffusionCellArrays(particles; loc=loc)
     end
 
     function JustPIC._2D.init_particles(
@@ -215,6 +215,21 @@ module _2D
         return nothing
     end
 
+    function JustPIC._2D.subgrid_diffusion_centroid!(
+        pT,
+        T_grid,
+        ΔT_grid,
+        subgrid_arrays,
+        particles::Particles{CUDABackend},
+        xvi,
+        di,
+        dt;
+        d=1.0,
+    )
+        subgrid_diffusion_centroid!(pT, T_grid, ΔT_grid, subgrid_arrays, particles, xvi, di, dt; d=d)
+        return nothing
+    end
+
     ## MakerChain
 
     function JustPIC._2D.init_markerchain(::Type{CUDABackend}, nxcell, min_xcell, max_xcell, xv, initial_elevation)
@@ -377,8 +392,8 @@ module _3D
         return Particles(CUDABackend, coords, index, nxcell, max_xcell, min_xcell, np)
     end
 
-    function JustPIC._3D.SubgridDiffusionCellArrays(particles::Particles{CUDABackend})
-        return SubgridDiffusionCellArrays(particles)
+    function JustPIC._3D.SubgridDiffusionCellArrays(particles::Particles{CUDABackend}; loc::Symbol=:vertex)
+        return SubgridDiffusionCellArrays(particles; loc=loc)
     end
 
     function JustPIC._3D.init_particles(
@@ -495,6 +510,21 @@ module _3D
         return nothing
     end
 
+    function JustPIC._3D.subgrid_diffusion_centroid!(
+        pT,
+        T_grid,
+        ΔT_grid,
+        subgrid_arrays,
+        particles::Particles{CUDABackend},
+        xvi,
+        di,
+        dt;
+        d=1.0,
+    )
+        subgrid_diffusion_centroid!(pT, T_grid, ΔT_grid, subgrid_arrays, particles, xvi, di, dt; d=d)
+        return nothing
+    end
+
     ## PassiveMarkers
 
     function JustPIC._3D.init_passive_markers(

diff --git a/scripts/temperature_advection_centroid.jl b/scripts/temperature_advection_centroid.jl
@@ -0,0 +1,88 @@
+using CUDA
+using JustPIC, JustPIC._2D
+
+# Threads is the default backend, 
+# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA"), 
+# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU")
+# const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+
+using GLMakie
+
+function expand_range(x::AbstractRange)
+    dx = x[2] - x[1]
+    n = length(x)
+    x1, x2 = extrema(x)
+    xI = round(x1-dx; sigdigits=5)
+    xF = round(x2+dx; sigdigits=5)
+    range(xI, xF, length=n+2)
+end
+
+# Analytical flow solution
+vx_stream(x, y) =  250 * sin(π*x) * cos(π*y)
+vy_stream(x, y) = -250 * cos(π*x) * sin(π*y)
+g(x) = Point2f(
+    vx_stream(x[1], x[2]),
+    vy_stream(x[1], x[2])
+)
+
+function main()
+    # Initialize particles -------------------------------
+    nxcell, max_xcell, min_xcell = 24, 30, 12
+    n = 256
+    nx = ny = n-1
+    Lx = Ly = 1.0
+    # nodal vertices
+    xvi = xv, yv = range(0, Lx, length=n), range(0, Ly, length=n)
+    dxi = dx, dy = xv[2] - xv[1], yv[2] - yv[1]
+    # nodal centers
+    xci = xc, yc = range(0+dx/2, Lx-dx/2, length=n-1), range(0+dy/2, Ly-dy/2, length=n-1)
+    # staggered grid velocity nodal locations
+    grid_vx = xv, expand_range(yc)
+    grid_vy = expand_range(xc), yv
+
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi...,
+    )
+
+    # Cell fields -------------------------------
+    Vx = TA(backend)([vx_stream(x, y) for x in grid_vx[1], y in grid_vx[2]]);
+    Vy = TA(backend)([vy_stream(x, y) for x in grid_vy[1], y in grid_vy[2]]);
+    T  = TA(backend)([y for x in xc, y in yc]);
+    V  = Vx, Vy;
+
+    dt = min(dx / maximum(abs.(Array(Vx))),  dy / maximum(abs.(Array(Vy))));
+    dt *= 0.5
+
+    # Advection test
+    particle_args = pT, = init_cell_arrays(particles, Val(1));
+    centroid2particle!(pT, xci, T, particles);
+
+    Tmin, Tmax = 0, 1
+    @assert all(x -> Tmin ≤ x ≤ Tmax, Array(pT.data[particles.index.data]))
+
+    !isdir("figs") && mkdir("figs")
+
+    niter = 250
+    for it in 1:niter
+        @show it 
+        advection!(particles, RungeKutta2(), V, (grid_vx, grid_vy), dt)
+        move_particles!(particles, xvi, particle_args)
+        inject_particles!(particles, (pT, ), xci)
+            @assert all(x -> Tmin ≤ x ≤ Tmax, Array(pT.data[particles.index.data]))
+        particle2centroid!(T, pT, xci, particles)
+            @assert all(x -> Tmin ≤ x ≤ Tmax, Array(T))
+
+        if rem(it, 10) == 0
+            f, ax, = heatmap(xci..., Array(T), colormap=:batlow)
+            streamplot!(ax, g, xvi...)
+            save("figs/test_$(it).png", f)
+            f
+        end
+    end
+
+    @show extrema(T)
+    println("Finished")
+end
+
+main()