Allow automatic grouping of hits based on a group_distance when creating an Event

Project.toml

@@ -6,6 +6,7 @@ version = "0.10.11"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 Clustering = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Format = "1fa38f19-a742-5d3f-a2b9-30dd87b9d5f8"
@@ -48,6 +49,7 @@ SolidStateDetectorsGeant4Ext = "Geant4"
 Adapt = "3, 4"
 ArraysOfArrays = "0.5, 0.6"
 Clustering = "0.15"
+DataStructures = "0.17, 0.18"
 Distributions = "0.24.5, 0.25"
 FillArrays = "0.9, 0.10, 0.11, 0.12, 0.13, 1"
 Format = "1.2"

src/ChargeClustering/ChargeClustering.jl

@@ -72,3 +72,57 @@ function cluster_detector_hits(table::TypedTables.Table, cluster_radius::RealQua
         )
     ))
 end
+
+
+function _group_points_by_distance(pts::AbstractVector{CartesianPoint{T}}, group_distance::T)::Tuple{Vector{Int}, Vector{Int}} where {T}
+
+    n::Int = length(pts)
+
+    # use BFS to find connected components
+    visited        = falses(n)
+    clustersidx    = similar(eachindex(pts))
+    elem_ptr       = similar(eachindex(pts), n+1)
+    queue          = DataStructures.CircularBuffer{Int}(n)
+
+    counter        = firstindex(pts)
+    Cidx           = firstindex(elem_ptr)
+    elem_ptr[Cidx] = counter
+
+    @inbounds for start in eachindex(pts)
+        if !visited[start]
+            push!(queue, start)
+            visited[start] = true
+            while !isempty(queue)
+                node = popfirst!(queue)
+                clustersidx[counter] = node
+                counter += 1
+                for neighbor in eachindex(pts)
+                    if !visited[neighbor] && distance_squared(pts[node], pts[neighbor]) <= group_distance^2
+                        push!(queue, neighbor)
+                        visited[neighbor] = true
+                    end
+                end
+            end
+            Cidx += 1
+            elem_ptr[Cidx] = counter
+        end
+    end
+    clustersidx, elem_ptr[begin:Cidx]
+end
+
+function group_points_by_distance(pts::AbstractVector{CartesianPoint{T}}, group_distance::T)::Tuple{VectorOfVectors{CartesianPoint{T}}, VectorOfVectors{T}} where {T <: SSDFloat}
+    clustersidx, elem_ptr = _group_points_by_distance(pts, group_distance)
+    VectorOfVectors(pts[clustersidx], elem_ptr)
+end
+
+function group_points_by_distance(pts::AbstractVector{CartesianPoint{T}}, energies::AbstractVector{T}, group_distance::T)::Tuple{VectorOfVectors{CartesianPoint{T}}, VectorOfVectors{T}} where {T <: SSDFloat}
+    @assert eachindex(pts) == eachindex(energies)
+    clustersidx, elem_ptr = _group_points_by_distance(pts, group_distance)
+    VectorOfVectors(pts[clustersidx], elem_ptr), VectorOfVectors(energies[clustersidx], elem_ptr)
+end
+
+function group_points_by_distance(pts::AbstractVector{CartesianPoint{T}}, energies::AbstractVector{T}, radius::AbstractVector{T}, group_distance::T)::Tuple{VectorOfVectors{CartesianPoint{T}}, VectorOfVectors{T}, VectorOfVectors{T}} where {T <: SSDFloat}
+    @assert eachindex(pts) == eachindex(energies) == eachindex(radius)
+    clustersidx, elem_ptr = _group_points_by_distance(pts, group_distance)
+    VectorOfVectors(pts[clustersidx], elem_ptr), VectorOfVectors(energies[clustersidx], elem_ptr), VectorOfVectors(radius[clustersidx], elem_ptr)
+end

src/ChargeDrift/ChargeDrift.jl

@@ -130,7 +130,7 @@ function _add_fieldvector_selfrepulsion!(step_vectors::Vector{CartesianVector{T}
                 if iszero(direction) # if the two charges are at the exact same position
                     continue         # don't let them self-repel each other but treat them as same change
                 end                  # if diffusion is simulated, they will very likely be separated in the next step
-                tmp::T = elementary_charge * inv(4 * pi * ϵ0 * ϵ_r * max(sum(direction.^2), T(1e-10))) # minimum distance is 10µm 
+                tmp::T = elementary_charge * inv(4π * ϵ0 * ϵ_r * max(distance_squared(direction), T(1e-10))) # minimum distance is 10µm 
                 step_vectors[n] += charges[m] * tmp * normalize(direction)
                 step_vectors[m] -= charges[n] * tmp * normalize(direction)
             end

src/ConstructiveSolidGeometry/PointsAndVectors/PointsAndVectors.jl

@@ -3,3 +3,7 @@ abstract type AbstractCoordinateVector{T, S} <: StaticArrays.FieldVector{3, T} e
 
 include("Points.jl")
 include("Vectors.jl")
+
+@inline distance_squared(v::CartesianVector) = v.x * v.x + v.y * v.y + v.z * v.z
+@inline distance_squared(p1::CartesianPoint{T}, p2::CartesianPoint{T}) where {T <: Real} = distance_squared(CartesianVector(p1 .- p2))
+export distance_squared

src/Event/Event.jl

@@ -26,10 +26,16 @@ function Event(location::AbstractCoordinatePoint{T}, energy::RealQuantity = one(
     return evt
 end
 
-function Event(locations::Vector{<:AbstractCoordinatePoint{T}}, energies::Vector{<:RealQuantity} = ones(T, length(locations)))::Event{T} where {T <: SSDFloat}
+function Event(locations::Vector{<:AbstractCoordinatePoint{T}}, energies::Vector{<:RealQuantity} = ones(T, length(locations)); group_distance::Union{<:Real, <:LengthQuantity} = NaN)::Event{T} where {T <: SSDFloat}
+    d::T = T(to_internal_units(group_distance))
+    @assert isnan(d) || d >= 0 "Group distance must be positive or NaN (no grouping), but $(group_distance) was given."
     evt = Event{T}()
-    evt.locations = VectorOfArrays(broadcast(pt -> [CartesianPoint(pt)], locations))
-    evt.energies = VectorOfArrays(broadcast(E -> [T(to_internal_units(E))], energies))
+    if isnan(d) # default: no grouping, the charges from different hits drift independently
+        evt.locations = VectorOfArrays(broadcast(pt -> [CartesianPoint(pt)], locations))
+        evt.energies = VectorOfArrays(broadcast(E -> [T(to_internal_units(E))], energies))
+    else
+        evt.locations, evt.energies = group_points_by_distance(CartesianPoint.(locations), T.(to_internal_units.(energies)), d)
+    end
     return evt
 end
 
@@ -42,6 +48,7 @@ end
 
 function Event(nbcc::NBodyChargeCloud{T})::Event{T} where {T <: SSDFloat}
     evt = Event{T}()
+    # charges in one NBodyChargeCloud should see each other by default
     evt.locations = VectorOfArrays([nbcc.locations])
     evt.energies = VectorOfArrays([nbcc.energies])
     return evt
@@ -55,15 +62,40 @@ function Event(nbccs::Vector{<:NBodyChargeCloud{T}})::Event{T} where {T <: SSDFl
 end
 
 function Event(locations::Vector{<:AbstractCoordinatePoint{T}}, energies::Vector{<:RealQuantity}, N::Int; 
-               particle_type::Type{PT} = Gamma, number_of_shells::Int = 2,
-               radius::Vector{<:RealQuantity} = radius_guess.(T.(to_internal_units.(energies)), particle_type)
-              )::Event{T} where {T <: SSDFloat, PT <: ParticleType}
-
+        particle_type::Type{PT} = Gamma, number_of_shells::Int = 2,
+        radius::Vector{<:Union{<:Real, <:LengthQuantity}} = radius_guess.(T.(to_internal_units.(energies)), particle_type),
+        group_distance::Union{<:Real, <:LengthQuantity} = NaN
+    )::Event{T} where {T <: SSDFloat, PT <: ParticleType}
+
+
     @assert eachindex(locations) == eachindex(energies) == eachindex(radius)
-    return Event(broadcast(i -> 
-                NBodyChargeCloud(locations[i], energies[i], N, particle_type, 
+
+    d::T = T(to_internal_units(group_distance))
+    @assert isnan(d) || d >= 0 "Group distance must be positive or NaN (no grouping), but $(group_distance) was given."
+
+    if isnan(d) # default: no grouping, the charges from different hits drift independently
+        Event(
+            broadcast(i -> 
+                NBodyChargeCloud(locations[i], energies[i], N, 
                 radius = T(to_internal_units(radius[i])), number_of_shells = number_of_shells),
-           eachindex(locations)))
+            eachindex(locations))
+        )
+    else # otherwise: group the CENTERS by distance and compose the NBodyChargeClouds around them
+        loc, ene, rad = group_points_by_distance(CartesianPoint.(locations), T.(to_internal_units.(energies)), T.(to_internal_units.(radius)), d)
+        nbccs = broadcast(i -> 
+            broadcast(j -> 
+                let nbcc = NBodyChargeCloud(loc[i][j], ene[i][j], N, 
+                    radius = rad[i][j], number_of_shells = number_of_shells)
+                    nbcc.locations, nbcc.energies
+                end, 
+            eachindex(loc[i])), 
+        eachindex(loc))
+
+        Event(
+            broadcast(nbcc -> vcat(getindex.(nbcc, 1)...), nbccs),
+            broadcast(nbcc -> vcat(getindex.(nbcc, 2)...), nbccs)
+        )
+    end
 end
 
 function Event(

src/SolidStateDetectors.jl

@@ -48,6 +48,7 @@ import .ConstructiveSolidGeometry: sample, to_internal_units, from_internal_unit
 export CartesianPoint, CartesianVector, CylindricalPoint
 
 import Clustering
+import DataStructures
 import Distributions
 import GPUArrays
 import IntervalSets

test/test_charge_drift_models.jl

@@ -3,7 +3,8 @@
 using Test
 
 using SolidStateDetectors
-using SolidStateDetectors: getVe, getVh, Vl, get_path_to_example_config_files, AbstractChargeDriftModel
+using SolidStateDetectors: getVe, getVh, Vl, get_path_to_example_config_files, AbstractChargeDriftModel, group_points_by_distance, distance_squared
+using ArraysOfArrays
 using InteractiveUtils
 using StaticArrays
 using LinearAlgebra
@@ -324,4 +325,43 @@ end
         @test cdm0.holes == cdmdict.holes
         @test cdm0.crystal_orientation == cdmdict.crystal_orientation
     end
+end
+
+@timed_testset "Test grouping of charges" begin
+    for N in 1:100
+        @timed_testset "Test for length $(N)" begin
+
+            # Generate random data
+            pts = [CartesianPoint{T}(rand(3)...) for _ in Base.OneTo(N)]
+            E = rand(T, N)
+            d = rand(T)
+
+            # Evaluate method
+            ptsg, Eg = group_points_by_distance(pts, E, d)
+
+            # Test correctness
+            s0 = Set(pts)
+
+            # result should be a VectorOfVectors
+            @test ptsg isa VectorOfVectors{<:CartesianPoint{T}}
+
+            # all points should appear in the result
+            @test Set(flatview(ptsg))    == Set(pts)
+            @test length(flatview(ptsg)) == length(pts)
+            @test Set(flatview(Eg))      == Set(E)
+            @test length(flatview(Eg))   == length(E)
+
+            for (i,group) in enumerate(ptsg)
+                @testset "Length $(N), subgroup $(i)" begin
+                    sg = Set(group)
+                    sc = setdiff(s0, group)
+                    for pt in group
+                        swo = setdiff(sg, Set([pt]))
+                        @test isempty(swo) || any(distance_squared.(Ref(pt), swo) .<= d^2)
+                        @test all(distance_squared.(Ref(pt), sc) .> d^2)
+                    end
+                end
+            end
+        end
+    end
 end