clustering detector hits by time as well

Author: sm2909

src/ChargeClustering/ChargeClustering.jl

@@ -2,69 +2,121 @@
 
 # breaking changes in Clustering v0.14 -> v0.15
 @inline _get_clusters(clusters::Clustering.DbscanResult) = clusters.clusters
-@inline _get_clusters(clusters::Vector{Clustering.DbscanCluster}) = clusters
+
+@inline function finalize_group!(time_groups, current_group, new_index)
+    push!(time_groups, current_group)
+    return [new_index]
+end
 
 function cluster_detector_hits(
     detno::AbstractVector{<:Integer},
     edep::AbstractVector{TT},
-    pos::AbstractVector{<:CartesianPoint{PT}},
-    cluster_radius::RealQuantity
-) where {TT<:RealQuantity, PT <: RealQuantity}
+    pos::AbstractVector{CartesianPoint{PT}},
+    thit::AbstractVector{TTT},
+    cluster_radius::RealQuantity,
+    cluster_time::RealQuantity
+) where {TT<:RealQuantity, PT <: RealQuantity, TTT <: RealQuantity}
 
-    pos = isa(first(pos), CartesianPoint) ? [SVector(p.x, p.y, p.z) for p in pos] : pos     #converting CartesianPoint to SVector
+    pos = [SVector(p.x, p.y, p.z) for p in pos]   # converting CartesianPoint to SVector
 
     Table = TypedTables.Table
-    unsorted = Table(detno = detno, edep = edep, pos = pos)
+    unsorted = Table(detno = detno, edep = edep, pos = pos, thit = thit)
     sorting_idxs = sortperm(unsorted.detno)
     sorted = unsorted[sorting_idxs]
     grouped = Table(consgroupedview(sorted.detno, TypedTables.columns(sorted)))
 
     r_detno = similar(detno, 0)
     r_edep = similar(edep, 0)
     r_pos = similar(pos, 0)
+    r_thit = similar(thit, 0)
 
     posunit = unit(PT)
     ustripped_cradius = ustrip(posunit, cluster_radius)
-    
+
+    thitunit = unit(TTT)
+    ustripped_ctime = ustrip(thitunit, cluster_time)
+
     for d_hits_nt in grouped
         d_hits = Table(d_hits_nt)
         d_detno = first(d_hits.detno)
         @assert all(isequal(d_detno), d_hits.detno)
-        if length(d_hits) > 3
-            clusters = Clustering.dbscan(ustrip.(flatview(d_hits.pos)), ustripped_cradius, leafsize = 20, min_neighbors = 1, min_cluster_size = 1)
-            for c in _get_clusters(clusters)
-                idxs = vcat(c.boundary_indices, c.core_indices)
-                @assert length(idxs) == c.size
-                c_hits = view(d_hits, idxs)
-                
-                push!(r_detno, d_detno)
-                esum_u = sum(c_hits.edep)
-                push!(r_edep, esum_u)
-                esum = ustrip(esum_u)
-                if esum ≈ 0
-                    push!(r_pos, mean(c_hits.pos))
-                else
-                    weights = ustrip.(c_hits.edep) .* inv(esum)
-                    push!(r_pos, sum(c_hits.pos .* weights))
+
+        # sort hits by time
+        t_sort_idx = sortperm(d_hits.thit)
+        d_hits = d_hits[t_sort_idx]
+        t_vals = ustrip.(d_hits.thit)
+
+        time_groups = Vector{Vector{Int}}()
+        current_group = [1]
+        start_time = t_vals[1]
+        for i in 2:length(t_vals)
+            if t_vals[i] - start_time ≤ ustripped_ctime
+                push!(current_group, i)
+            else
+                current_group = finalize_group!(time_groups, current_group, i)
+                start_time = t_vals[i]
+            end
+        end
+        current_group = finalize_group!(time_groups, current_group, nothing)
+
+
+        # spatial clustering within each temporal cluster
+        for tg in time_groups
+            t_hits = view(d_hits, tg)
+            if length(t_hits) > 3
+                clusters = Clustering.dbscan(ustrip.(flatview(t_hits.pos)), ustripped_cradius;
+                                              leafsize = 20, min_neighbors = 1, min_cluster_size = 1)
+                for c in _get_clusters(clusters)
+                    idxs = vcat(c.boundary_indices, c.core_indices)
+                    @assert length(idxs) == c.size
+                    c_hits = view(t_hits, idxs)
+
+                    push!(r_detno, d_detno)
+                    esum_u = sum(c_hits.edep)
+                    push!(r_edep, esum_u)
+                    esum = ustrip(esum_u)
+                    if esum ≈ 0
+                        push!(r_pos, mean(c_hits.pos))
+                        push!(r_thit, mean(c_hits.thit))
+                    else
+                        weights = ustrip.(c_hits.edep) .* inv(esum)
+                        push!(r_pos, sum(c_hits.pos .* weights))
+                        push!(r_thit, sum(c_hits.thit .* weights))
+                    end
                 end
+            else
+                append!(r_detno, t_hits.detno)
+                append!(r_edep, t_hits.edep)
+                append!(r_pos, t_hits.pos)
+                append!(r_thit, t_hits.thit)
             end
-        else
-            append!(r_detno, d_hits.detno)
-            append!(r_edep, d_hits.edep)
-            append!(r_pos, d_hits.pos)
         end
     end
 
-    (detno = r_detno, edep = r_edep, pos = r_pos)
+    r_pos = [CartesianPoint(p[1], p[2], p[3]) for p in r_pos]   #converting SVector back to CartesianPoint
+
+    (detno = r_detno, edep = r_edep, pos = r_pos, thit = r_thit)
 end
 
 
-function cluster_detector_hits(table::TypedTables.Table, cluster_radius::RealQuantity)
+function cluster_detector_hits(
+    table::TypedTables.Table, 
+    cluster_radius::PT, 
+    cluster_time::TTT
+) where {PT <: RealQuantity, TTT <: RealQuantity}
     @assert :pos in TypedTables.columnnames(table) "Table has no column `pos`"
+    @assert :thit in TypedTables.columnnames(table) "Table has no column `thit`"
     @assert :edep in TypedTables.columnnames(table) "Table has no column `edep`"
     @assert :detno in TypedTables.columnnames(table) "Table has no column `detno`"
     clustered_nt = map(
-        evt -> cluster_detector_hits(evt.detno, evt.edep, evt.pos, cluster_radius),
+        evt -> cluster_detector_hits(
+            evt.detno,
+            evt.edep,
+            [p isa CartesianPoint ? p : CartesianPoint(p...) for p in evt.pos],
+            evt.thit,
+            cluster_radius,
+            cluster_time
+        ),
         table
     )
     TypedTables.Table(merge(

test/test_geant4.jl

@@ -66,6 +66,23 @@ end
     @test evts isa Table
     @test length(evts) == 100
 
+    # Cluster events by time and radius
+    clustered_evts = @test_nowarn SolidStateDetectors.cluster_detector_hits(evts, 10u"µm", 100u"s")
+    @test length(clustered_evts) == length(evts)
+    @test length(flatview(clustered_evts.pos)) ≤ length(flatview(evts.pos))
+
+    @test all(isa.(flatview(clustered_evts.pos), CartesianPoint{<:Real}))   #checking type consistency for pos
+    @test isapprox(sum(flatview(clustered_evts.edep)), sum(flatview(evts.edep)); rtol=1e-6)   #checking energy conservation after clustering
+    
+    #sanity check that all values in the table are not NaN
+    pos = flatview(clustered_evts.pos)
+    @test all(!isnan.(ustrip.([p.x for p in pos]))) &&
+        all(!isnan.(ustrip.([p.y for p in pos]))) &&
+        all(!isnan.(ustrip.([p.z for p in pos])))
+
+    @test all(!isnan.(ustrip.(flatview(clustered_evts.thit))))
+    @test all(!isnan.(ustrip.(flatview(clustered_evts.edep))))
+
     # Generate waveforms
     simulate!(sim, refinement_limits = [0.2,0.1,0.05,0.03,0.02])
     wf = simulate_waveforms(evts, sim, Δt = 1u"ns", max_nsteps = 2000)