Performance Optimization: Memory Pool & Early Termination

src/PAM.jl

@@ -16,7 +16,20 @@ function projected_coordinate(x::AbstractVector{Float64}, y::AbstractVector{Floa
     return r, z
 end
 
-mutable struct Pellet{A,T,N,S,B,X}
+"""
+Preallocated buffer pool to reduce allocations
+"""
+mutable struct AllocationPool{T<:Real}
+    nsource_buffer::Vector{T}
+    nsource_max_len::Int
+end
+
+function AllocationPool(::Type{T}, surfaces::Vector{IMAS.FluxSurface}) where {T<:Real}
+    max_len = maximum(length(surf.r) for surf in surfaces)
+    return AllocationPool{T}(Vector{T}(undef, max_len), max_len)
+end
+
+mutable struct Pellet{A,T,N,S,B,X,P}
     properties::IMAS.pellets__time_slice___pellet
     Btdep::B
     update_plasma::B
@@ -37,6 +50,7 @@ mutable struct Pellet{A,T,N,S,B,X}
     ablation_rate::A
     density_source::N
     temp_drop::A
+    pool::P
 end
 
 function Pellet(
@@ -94,11 +108,16 @@ function Pellet(
 
     radii = cumsum([layer.thickness for layer in pelt.layer])
     @assert pelt.shape.size[1] == radii[end] "The layer's thickness don't add up to the total radius"
-    radius = fill(radii[end], size(time))
+    radius = fill(0.0, size(time))
+    radius[1] = radii[end] # Set initial radius at t=0
     ablation_rate = fill(0.0, size(time))
     temp_drop = fill(0.0, size(time))
     R_drift = fill(0.0, size(time))
     density_source = fill(0.0, (length(time), length(surfaces)))
+
+    # Initialize allocation pool
+    pool = AllocationPool(Float64, surfaces)
+
     return Pellet(
         pelt,
         Bt_dependance,
@@ -119,7 +138,8 @@ function Pellet(
         radius,
         ablation_rate,
         density_source,
-        temp_drop
+        temp_drop,
+        pool
     )
 end
 
@@ -536,7 +556,11 @@ function dr_dt!(pelt::Pellet, k::Int)
 end
 
 function pellet_density(pelt::Pellet, surface::IMAS.FluxSurface, k::Int)
-    # calculations of the pellet cloud size
+    # Get buffer view from pool
+    len = length(surface.r)
+    nsource = @view pelt.pool.nsource_buffer[1:len]
+
+    # Cloud size calculations
     cloudFactor = 5
 
     cloudFactorR = cloudFactor
@@ -545,11 +569,14 @@ function pellet_density(pelt::Pellet, surface::IMAS.FluxSurface, k::Int)
     rcloudR = pelt.radius[k] * cloudFactorR
     rcloudZ = pelt.radius[k] * cloudFactorZ
 
-    # Assume density source as a 2D Gaussian function
-    nsource = exp.(-0.5 .* ((pelt.r[k] .- surface.r .+ 0.5 .* pelt.R_drift[k]) .^ 2 ./ (rcloudR .+ 0.25 .* pelt.R_drift[k]) .^ 2 .+ (pelt.z[k] .- surface.z) .^ 2 ./ rcloudZ .^ 2))
+    # # Compute 2D Gaussian density source (in-place)
+    @. nsource = exp(-0.5 * (
+        (pelt.r[k] - surface.r + 0.5 * pelt.R_drift[k])^2 / (rcloudR + 0.25 * pelt.R_drift[k])^2 +
+        (pelt.z[k] - surface.z)^2 / rcloudZ^2
+    ))
 
     # need to normilize on the surface area under the gaussian shape
-    nsource ./= (2 * π)^2 * (rcloudR + 0.25 * pelt.R_drift[k]) * (pelt.r[k] + 0.5 * pelt.R_drift[k]) * rcloudZ
+    nsource ./= (2π)^2 * (rcloudR + 0.25 * pelt.R_drift[k]) * (pelt.r[k] + 0.5 * pelt.R_drift[k]) * rcloudZ
 
     nsource .*= pelt.ablation_rate[k]
 
@@ -564,23 +591,21 @@ function ablate!(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__pr
 
         if pelt.ρ[k] > 1.0 && pelt.radius[k-1] > 0
             pelt.radius[k] = pelt.radius[k-1]
-
         else
             dr_dt!(pelt, k)
             if isnan(pelt.ablation_rate[k])
-
                 pelt.ablation_rate[k] = 0.0
             end
 
+            if pelt.radius[k] <= 0.0
+                @info "Pellet fully ablated at t = $(pelt.time[k]) s"
+                return 
+            end
+
             drift!(Val(pelt.drift_model), pelt, eqt, cp1d, k)
 
             for (ks, surface) in enumerate(surfaces)
-                tmp = pellet_density(pelt, surface, k)
-                if isnan(tmp)
-                    pellet_source[k, ks] += 0.0
-                else
-                    pellet_source[k, ks] += tmp * dt
-                end
+                pellet_source[k, ks] += pellet_density(pelt, surface, k) * dt
             end
 
             #------calculate energy sink and update plasma -----------

test/runtests.jl

@@ -3,7 +3,10 @@ import IMAS
 using Test
 
 @testset "PAM" begin
-    dd_D3D = IMAS.json2imas(joinpath(@__DIR__, "..", "examples", "template_D3D_1layer_2species.json"))
+    dd_D3D_json = IMAS.json2imas(joinpath(@__DIR__, "..", "examples", "template_D3D_1layer_2species.json"))
+    dd_D3D = IMAS.hdf2imas(joinpath(@__DIR__, "..", "examples", "template_D3D_1layer_2species.h5"))
+
+    @test dd_D3D_json == dd_D3D
 
     dd_D3D.pellets.time_slice[].pellet[1].velocity_initial = 200.0
 
@@ -22,3 +25,41 @@ using Test
     end
 
 end
+
+@testset "Comparison with OMFIT PAM" begin
+    dd_D3D = IMAS.hdf2imas(joinpath(@__DIR__, "..", "examples", "template_D3D_1layer_2species.h5"))
+    dd_D3D.pellets.time_slice[].pellet[1].velocity_initial = 200.0;
+
+    inputs=(
+        t_start = 0.0,
+        t_finish = 0.001,
+        time_step = 0.0001, 
+        drift_model= :HPI2,
+        Bt_dependance=true,
+        update_plasma = false,
+    )
+
+    pellet_PAM = PAM.run_PAM(dd_D3D; inputs...);
+
+    # Reference from OMFIT PAM
+    omfit_radius = 0.01 * [0.1, 0.1, 0.09907301, 0.09514626, 0.08962336, 0.08174432, 0.07028775, 0.05297672, 0.02178495, 0.0, 0.0]
+    omfit_R = [2.3, 2.28, 2.26, 2.24, 2.22, 2.2, 2.18, 2.16, 2.14, 2.12, 2.1]
+    omfit_ablation_D = [0.0, 0.0, 1.15829553e+23, 1.62737315e+23, 2.03252305e+23, 2.43448343e+23, 2.64645468e+23, 2.35555840e+23,
+        9.11641845e+22, 0.0, 0.0]
+
+    @test isapprox(pellet_PAM.radius, omfit_radius, rtol=5e-3)
+    @test isapprox(pellet_PAM.r, omfit_R, rtol=5e-3)
+    @test isapprox(pellet_PAM.ablation_rate, 2 * omfit_ablation_D, rtol=5e-3)
+
+    # Regression test (with previous PAM version)
+    @test isapprox(maximum(pellet_PAM.density_source), 8.79030213181948e19)
+    @test isapprox(sum(pellet_PAM.density_source), 3.9712279520318335e21)
+
+    @test isapprox(pellet_PAM.ablation_rate[4], 3.2596193988579004e23)
+    @test isapprox(pellet_PAM.ablation_rate[9], 1.7870515483957858e23)
+    @test isapprox(pellet_PAM.ne[4], 3.437934198621271e19)
+    @test isapprox(pellet_PAM.Te[4],  517.7166642235971)
+    @test isapprox(pellet_PAM.temp_drop[4],  0.9994708595114168)
+    @test isapprox(pellet_PAM.radius[4],  0.0009515680832308055)
+    @test isapprox(pellet_PAM.x[4], 2.2399999999999998)
+end