optimize indices

Author: ucabc46

src/utils.jl

@@ -1,15 +1,52 @@
+using ExactOptimalTransport
 
 function normalized_exp!(weight::AbstractVector)
     weight .-= maximum(weight)
     @. weight = exp(weight)
     weight ./= sum(weight)
 end
 
+function optimized_resample!(resampled_indices::AbstractVector{Int}, nrank::Int)
+    nprt_per_rank = length(resampled_indices) ÷ nrank
+    stock_queue = [Int[] for _ in 1:nrank]
+
+    # Assign each resampled index to its corresponding rank
+    for resampled_idx in resampled_indices
+        rank = div(resampled_idx - 1, nprt_per_rank) + 1
+        push!(stock_queue[rank], resampled_idx)
+    end
+
+    supply_vector = [length(stock_queue[rank]) for rank in 1:nrank]
+    demand_vector = fill(nprt_per_rank, nrank)
+    cost_matrix = ones(Int, nrank, nrank)
+    for i in 1:nrank
+        cost_matrix[i, i] = 0
+    end
+
+    γ = emd(supply_vector, demand_vector, cost_matrix)
+
+    # update resampled_indices
+    for i in 1:nrank
+        idx = 1
+        for j in 1:nrank
+            nmove = Int(γ[j, i])
+            for _ in 1:nmove
+                resampled_indices[(i - 1) * nprt_per_rank + idx] = popfirst!(stock_queue[j])
+                idx += 1
+            end
+        end
+    end
+    return resampled_indices
+end
+
+
 # Resample particles from given weights using Stochastic Universal Sampling
 function resample!(
     resampled_indices::AbstractVector{Int}, 
     weights::AbstractVector{T}, 
-    rng::Random.AbstractRNG=Random.TaskLocalRNG()
+    rng::Random.AbstractRNG=Random.TaskLocalRNG(),
+    optimize_resample::Bool=false,
+    nrank::Int=1,
 ) where T
 
     nprt = length(weights)
@@ -28,6 +65,10 @@ function resample!(
         end
         resampled_indices[ip] = k
     end
+
+    if optimize_resample
+        resampled_indices .= optimized_resample!(resampled_indices, nrank)
+    end
 end
 
 function init_states(model, nprt_per_rank::Int, n_tasks::Int, rng::AbstractRNG)
@@ -174,6 +215,7 @@ function copy_states_dedup!(
         for (id, buffer_indices) in remote_copies
             if length(buffer_indices) > 1
                 source_view = view(buffer, :, buffer_indices[1])
+                # TODO: threading in chunks
                 Threads.@threads for i in 2:length(buffer_indices)
                     k = buffer_indices[i]
                     buffer[:, k] .= source_view
@@ -182,8 +224,12 @@ function copy_states_dedup!(
         end
     end
 
-    @timeit_debug to "write from buffer" particles .= buffer
-
+    @timeit_debug to "write from buffer" begin
+        Threads.@threads for j in 1:size(particles, 2)
+            # @views creates a non-allocating view of the column, which is faster inside a loop
+            @views particles[:, j] .= buffer[:, j]
+        end
+    end
 end
 
 function _determine_sends(resampling_indices::Vector{Int}, my_rank::Int, nprt_per_rank::Int)

test/mpi_copy_states.jl

@@ -39,7 +39,7 @@ my_size = MPI.Comm_size(MPI.COMM_WORLD)
 
 println("Number of threads available: ", Threads.nthreads())
 
-n_particle_per_rank = 1000
+n_particle_per_rank = 100
 n_particle = n_particle_per_rank * my_size
 verbose = "-v" in ARGS || "--verbose" in ARGS
 output_timer = "-t" in ARGS || "--output-timer" in ARGS

test/optimized_resample.jl

@@ -0,0 +1,10 @@
+using ParticleDA
+
+nrank = 5
+n_particle = 10
+
+resampled_indices = sample_indices(n_particle, k=5, p=0.99)
+println("Resampled Indices: ", resampled_indices)
+
+resampled_indices = ParticleDA.optimized_resample!(resampled_indices, nrank)
+println("Optimized Resampled Indices: ", resampled_indices)