Add repulsion between cylinders/spheres

Repulsion step is skipped for walls. This is based on code from Zhiyu Zheng.

src/geometries/user/random_distribution.jl

 module RandomDistribution
-import StaticArrays: SVector
+import StaticArrays: SVector, MVector
 import Distributions
 
 function draw_radii(box_size::SVector{N, <:Real}, target_density::Real, distribution::Distributions.ContinuousUnivariateDistribution, allowed_range::Tuple{<:Real, <:Real}) where {N}
@@ -28,7 +28,9 @@ function draw_radii(box_size::SVector{N, <:Real}, target_density::Real, distribu
 end
 
 
-function eliminate_overlap(positions::AbstractVector{SVector{N, Float64}}, radii::AbstractVector{Float64}, box_size::SVector{N, Float64}; max_iter=1000) where {N}
+norm_dist(orig_dist::Float64, half_repeat::Float64) = mod(orig_dist + half_repeat, half_repeat * 2) - half_repeat
+
+function eliminate_overlap!(positions::AbstractVector{SVector{N, Float64}}, radii::AbstractVector{Float64}, box_size::SVector{N, Float64}; max_iter=1000) where {N}
     @assert length(positions) == length(radii)
     moved = true
     half_box_size = box_size ./ 2
@@ -37,30 +39,79 @@ function eliminate_overlap(positions::AbstractVector{SVector{N, Float64}}, radii
 
     for _ in 1:max_iter
         moved = false
-        for i in 1:length(positions)
+        for i in eachindex(positions)
             for j in 1:(i-1)
-                displacement = map((dp, hbs) -> mod(dp + hbs, hbs * 2) - hbs, positions[i] - positions[j], half_box_size)
+                displacement = map(norm_dist, positions[i] - positions[j], half_box_size)
                 distsq = sum(p->p*p, displacement)
-                total_radius = radii[i] + radii[j]
-                if distsq < (total_radius * total_radius)
+                min_dist = radii[i] + radii[j]
+                if distsq < (min_dist * min_dist)
                     moved = true
                     dist = sqrt(distsq)
-                    total_movement = max((total_radius - dist) * (1 + step_size), 0.01 * total_radius)
+                    total_movement = max((min_dist - dist) * (1 + step_size), 0.01 * min_dist)
                     # move larger cylinder less
-                    positions[i] = positions[i] .+ displacement .* (total_movement * radii[j] / (total_radius * dist))
-                    positions[j] = positions[j] .- displacement .* (total_movement * radii[i] / (total_radius * dist))
+                    positions[i] = positions[i] .+ displacement .* (total_movement * radii[j] / (min_dist * dist))
+                    positions[j] = positions[j] .- displacement .* (total_movement * radii[i] / (min_dist * dist))
                 end
             end
         end
         step_size *= 0.8
         if !moved
-            return [map((p, hbs) -> mod(p + hbs, hbs * 2) - hbs, pos, half_box_size) for pos in positions]
+            positions[:] = map(pos -> norm_dist.(pos, half_box_size), positions)
+            return
         end
     end
     error("No solution without overlap found after max_iter=$max_iter iterations")
 end
 
 
+function repel_position!(positions::AbstractVector{SVector{N, Float64}}, radii::AbstractVector{Float64}, box_size::SVector{N, Float64}; max_iter=1000, repulsion_strength=1e-3) where {N}
+
+    half_box_size = box_size ./ 2
+
+    function does_repulsed_have_a_collision(proposed_pos, index)
+        for i in eachindex(proposed_pos)
+            if i != index
+                min_dist_sq = (radii[index] + radii[i])^2
+                dist_sq = sum(norm_dist.(proposed_pos .- positions[i], half_box_size) .^ 2)
+                if min_dist_sq >= dist_sq
+                    return true
+                end
+            end
+        end
+        return false
+    end
+    
+    proposed_pos = zero(MVector{N, Float64})
+    # Update circle positions based on repulsion
+    for _ in 1:max_iter
+        success_count=0
+        for i in eachindex(positions)
+            proposed_pos[:] = positions[i]
+            for j in eachindex(positions)
+                if i != j
+                    displacement = norm_dist.(positions[j] .- positions[i], half_box_size)
+
+                    dist_sq = sum(displacement .* displacement)
+
+                    # force = (radius^2*radius^2)*repulsion_strength / (dist * mean_radius^2)
+                    force = (radii[i]^N * radii[j]^N)*repulsion_strength / (dist_sq^(N//2) * (radii[i] + radii[j])^N)
+                    proposed_pos .-= force .* displacement
+                end
+            end
+            # Abandon changes if there's overlap
+            if !does_repulsed_have_a_collision(proposed_pos, i)
+                positions[i] = proposed_pos
+                success_count += 1
+            end
+        end
+        if iszero(success_count)
+            # Stop early when equilibrium is reached and nothing can be moved in a loop
+            break
+        end
+    end
+end
+
+
 """
     random_positions_radii(box_size, target_density, n_dimensions; distribution=Gamma, mean=1., variance=1., max_iter=1000, min_radius=0.1, max_radius=Inf)
 
@@ -74,12 +125,15 @@ Arguments:
 - `mean`: mean of the gamma distribution (ignored if `distribution` explicitly set)
 - `variance`: variance of the gamma distribution (ignored if `distribution` explicitly set)
 - `max_iter`: maximum number of iterations to try to prevent the circles/spheres from overlapping. An error is raised if they still overlap after this number of iterations.
+- `repulsion_strength`: strength of the repulsion in each step (defaults to 0.001).
+- `max_iter_repulse`: maximum number of iterations that circles/spheres will be repulsed from each other
 - `min_radius`: samples from the distribution below this radius will be rejected (in um).
 - `max_radius`: samples from the distribution above this radius will be rejected (in um).
 """
 function random_positions_radii(
     box_size, target_density::Real, ndim::Int;
     distribution=nothing, mean=1., variance=1., max_iter=1000,
+    repulsion_strength=1e-3, max_iter_repulse=1000,
     min_radius=0.1, max_radius=Inf,
 )
     if iszero(variance)
@@ -96,8 +150,12 @@ function random_positions_radii(
     end
     target_density = Float64(target_density)
     radii = draw_radii(box_size, target_density, distribution, (min_radius, max_radius))
-    start_positions = [SVector{ndim, Float64}((rand(ndim) .- 0.5) .* box_size) for _ in 1:length(radii)]
-    return eliminate_overlap(start_positions, radii, box_size; max_iter=max_iter), radii
+    positions = [SVector{ndim, Float64}((rand(ndim) .- 0.5) .* box_size) for _ in 1:length(radii)]
+    eliminate_overlap!(positions, radii, box_size; max_iter=max_iter)
+    if !isone(ndim)
+        repel_position!(positions, radii, box_size; max_iter=max_iter_repulse, repulsion_strength=repulsion_strength)
+    end
+    return positions, radii
 end
 
 end
\ No newline at end of file