From 0b8a21c931b240d4a302ef68ee384ee94d340123 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Tue, 21 Oct 2025 16:46:40 -0400 Subject: [PATCH 1/2] perf: use simpler IMAS trace function to reduce memory allocations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace trace_surfaces with trace_simple_surfaces - Add convenience dispatch for trace_simple_surfaces on equilibrium__time_slice - Update type signatures to use AbstractFluxSurface for better type flexibility Performance improvement (measured with @btime): - Runtime: 8.463 ms → 5.652 ms (1.5x faster) - Allocations: 61,925 → 7,517 (8.2x fewer) - Memory: 6.50 MiB → 2.35 MiB (2.8x less) --- src/PAM.jl | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/PAM.jl b/src/PAM.jl index 95d9276..a6569fa 100644 --- a/src/PAM.jl +++ b/src/PAM.jl @@ -24,7 +24,7 @@ mutable struct AllocationPool{T<:Real} nsource_max_len::Int end -function AllocationPool(::Type{T}, surfaces::Vector{IMAS.FluxSurface}) where {T<:Real} +function AllocationPool(::Type{T}, surfaces::Vector{<:IMAS.AbstractFluxSurface}) where {T<:Real} max_len = maximum(length(surf.r) for surf in surfaces) return AllocationPool{T}(Vector{T}(undef, max_len), max_len) end @@ -58,7 +58,7 @@ function Pellet( eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__profiles_1d, time::Vector{Float64}, - surfaces::Vector{IMAS.FluxSurface}, + surfaces::Vector{<:IMAS.AbstractFluxSurface}, drift_model::Symbol, Bt_dependance::Bool, update_plasma::Bool @@ -555,7 +555,7 @@ function dr_dt!(pelt::Pellet, k::Int) return end -function pellet_density(pelt::Pellet, surface::IMAS.FluxSurface, k::Int) +function pellet_density(pelt::Pellet, surface::IMAS.AbstractFluxSurface, k::Int) # Get buffer view from pool len = length(surface.r) nsource = @view pelt.pool.nsource_buffer[1:len] @@ -583,7 +583,7 @@ function pellet_density(pelt::Pellet, surface::IMAS.FluxSurface, k::Int) return IMAS.flux_surface_avg(nsource, surface) end -function ablate!(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__profiles_1d, pelt::Pellet, surfaces::Vector{IMAS.FluxSurface}) +function ablate!(eqt::IMAS.equilibrium__time_slice, cp1d::IMAS.core_profiles__profiles_1d, pelt::Pellet, surfaces::Vector{<:IMAS.AbstractFluxSurface}) pellet_source = zeros(length(pelt.time), length(surfaces)) for k in 2:length(pelt.time) @@ -659,6 +659,15 @@ end end end +# Convenient dispatch to call IMAS's trace_simple_surfaces +function IMAS.trace_simple_surfaces(eqt::IMAS.equilibrium__time_slice{T}, wall_r::AbstractVector{T}, wall_z::AbstractVector{T}) where {T<:Real} + eqt2d = findfirst(:rectangular, eqt.profiles_2d) + r, z, PSI_interpolant = IMAS.ψ_interpolant(eqt2d) + RA = eqt.global_quantities.magnetic_axis.r + ZA = eqt.global_quantities.magnetic_axis.z + return IMAS.trace_simple_surfaces(eqt.profiles_1d.psi, r, z, eqt2d.psi, PSI_interpolant, RA, ZA, wall_r, wall_z) +end + """ run_PAM(dd::IMAS.dd; t_start::Float64, t_finish::Float64, time_step::Float64, drift_model::Symbol, Bt_dependance::Bool, update_plasma::Bool) @@ -672,7 +681,7 @@ function run_PAM(dd::IMAS.dd; t_start::Float64, t_finish::Float64, time_step::Fl time = collect(range(t_start, t_finish; step=time_step)) # define flux surfaces, will be needed for the pellet source calculations - surfaces = IMAS.trace_surfaces(eqt, IMAS.first_wall(dd.wall)...) + surfaces = IMAS.trace_simple_surfaces(eqt, IMAS.first_wall(dd.wall)...) # initialize the pellet structure pellet = Pellet(dd.pellets.time_slice[].pellet[1], eqt, cp1d, time, surfaces, drift_model, Bt_dependance, update_plasma) From 128b4ff0dc3617a6be1cf2289e4613c6f1747511 Mon Sep 17 00:00:00 2001 From: Min-Gu Yoo Date: Sat, 25 Oct 2025 09:33:09 -0700 Subject: [PATCH 2/2] remove temporary convenience dispatch for `IMAS.trace_simple_surfaces` --- src/PAM.jl | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/PAM.jl b/src/PAM.jl index a6569fa..405e66c 100644 --- a/src/PAM.jl +++ b/src/PAM.jl @@ -659,14 +659,6 @@ end end end -# Convenient dispatch to call IMAS's trace_simple_surfaces -function IMAS.trace_simple_surfaces(eqt::IMAS.equilibrium__time_slice{T}, wall_r::AbstractVector{T}, wall_z::AbstractVector{T}) where {T<:Real} - eqt2d = findfirst(:rectangular, eqt.profiles_2d) - r, z, PSI_interpolant = IMAS.ψ_interpolant(eqt2d) - RA = eqt.global_quantities.magnetic_axis.r - ZA = eqt.global_quantities.magnetic_axis.z - return IMAS.trace_simple_surfaces(eqt.profiles_1d.psi, r, z, eqt2d.psi, PSI_interpolant, RA, ZA, wall_r, wall_z) -end """ run_PAM(dd::IMAS.dd; t_start::Float64, t_finish::Float64, time_step::Float64, drift_model::Symbol, Bt_dependance::Bool, update_plasma::Bool)