From fc27e7a42bfc98fcbec34db0455050de1b37ffbc Mon Sep 17 00:00:00 2001
From: Tamas Gal <himself@tamasgal.com>
Date: Mon, 17 Apr 2023 23:58:13 +0200
Subject: [PATCH] Add besttrack and related helper functions

---
 docs/src/api.md             |  2 ++
 src/KM3io.jl                |  7 ++++++-
 src/tools/reconstruction.jl | 20 +++++++++++++++++---
 test/tools.jl               | 30 ++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index dc69f148..40b8c676 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -72,5 +72,7 @@ most_frequent
 
 ### Reconstruction
 ```@docs
+besttrack
 RecStageRange
+hashistory
 ```
diff --git a/src/KM3io.jl b/src/KM3io.jl
index aab59cca..27997b8c 100644
--- a/src/KM3io.jl
+++ b/src/KM3io.jl
@@ -19,7 +19,12 @@ export Waveform, AcousticSignal, AcousticsTriggerParameter, piezoenabled, hydrop
 export Hit, TriggeredHit, Trk, MCHit, MCTrk, EvtHit
 
 export calibrate, floordist, slew
-export RecStageRange, hashistory
+
+export besttrack, RecStageRange, hashistory, hasjppmuonprefit, hasjppmuonsimplex, hasjppmuongandalf,
+       hasjppmuonenergy, hasjppmuonstart, hasjppmuonfit, hasshowerprefit, hasshowerpositionfit,
+       hasshowercompletefit, hasshowerfit, hasaashowerfit, hasreconstructedjppmuon,
+       hasreconstructedjppshower, hasreconstructedaashower
+
 export is3dshower, ismxshower, is3dmuon, isnb, triggered
 export most_frequent, categorize, nthbitset
 
diff --git a/src/tools/reconstruction.jl b/src/tools/reconstruction.jl
index 21a36a27..83b13a67 100644
--- a/src/tools/reconstruction.jl
+++ b/src/tools/reconstruction.jl
@@ -32,7 +32,7 @@ Base.in(rec_stage::T, rsr::RecStageRange) where T<:Integer = rsr.lower <= rec_st
 
 
 """
-Checks if a track with a given `rec_type` contains all the reconstruction stages in `rsr::RecStageRange`.
+Returns `true` if a track with a given `rec_type` contains all the reconstruction stages in `rsr::RecStageRange`.
 """
 function hashistory(t::Trk, rec_type::Integer, rsr::RecStageRange)
     rec_type != t.rec_type && return false
@@ -42,6 +42,9 @@ function hashistory(t::Trk, rec_type::Integer, rsr::RecStageRange)
     true
 end
 
+"""
+Returns `true` if a track with a given `rec_type` contains the `rec_stage`.
+"""
 function hashistory(t::Trk, rec_type::Integer, rec_stage::Integer)
     rec_type != t.rec_type && return false
     rec_stage ∈ t.rec_stages
@@ -62,6 +65,17 @@ hasreconstructedjppmuon(e::Evt) = any(hasjppmuonfit, e.trks)
 hasreconstructedjppshower(e::Evt) = any(hasshowerfit, e.trks)
 hasreconstructedaashower(e::Evt) = any(hasaashowerfit, e.trks)
 
-function besttrack(e::Evt, rsr::RecStageRange)
-    error("not implemented yet")
+
+"""
+
+Return the best reconstructed track for a given reconstruction type and
+reconstruction stage range. If no track could be found, `nothing` is returned.
+
+"""
+function besttrack(e::Evt, rec_type::Integer, rsr::RecStageRange)
+    candidates = filter(e.trks) do t
+        hashistory(t, rec_type, rsr)
+    end
+    length(candidates) == 0 && return nothing
+    sort(candidates; by=c -> (length(c.rec_stages), c.lik)) |> last
 end
diff --git a/test/tools.jl b/test/tools.jl
index 2cbc05d8..31d88a0a 100644
--- a/test/tools.jl
+++ b/test/tools.jl
@@ -1,4 +1,6 @@
+using KM3io
 import KM3io: nthbitset, SnapshotHit, tonumifpossible
+using KM3NeTTestData
 using Test
 
 @testset "tools" begin
@@ -44,3 +46,31 @@ end
     @test 1.1 == tonumifpossible("1.1")
     @test "1.1.1" == tonumifpossible("1.1.1")
 end
+
+@testset "has...() helpers" begin
+    f = ROOTFile(datapath("offline", "km3net_offline.root"))
+    e = f.offline[1]
+    t = e.trks |> first
+    @test hasjppmuonprefit(t)
+    @test !hasjppmuonsimplex(t)
+    @test hasjppmuongandalf(t)
+    @test hasjppmuonenergy(t)
+    @test hasjppmuonstart(t)
+    @test hasjppmuonfit(t)
+    @test !hasshowerprefit(t)
+    @test !hasshowerpositionfit(t)
+    @test !hasshowercompletefit(t)
+    @test !hasshowerfit(t)
+    @test !hasaashowerfit(t)
+    @test hasreconstructedjppmuon(e)
+    @test !hasreconstructedjppshower(e)
+    @test !hasreconstructedaashower(e)
+    close(f)
+end
+
+@testset "besttrack()" begin
+    f = ROOTFile(datapath("offline", "km3net_offline.root"))
+    bt = besttrack(f.offline[1], KM3io.RECONSTRUCTION.JPP_RECONSTRUCTION_TYPE, RecStageRange(KM3io.RECONSTRUCTION.JMUONBEGIN, KM3io.RECONSTRUCTION.JMUONEND))
+    @test 294.6407542676734 ≈ bt.lik
+    close(f)
+end
-- 
GitLab