From 8f56ff6b8676c0872ac2ca401680218c458dc160 Mon Sep 17 00:00:00 2001
From: Tamas Gal <himself@tamasgal.com>
Date: Wed, 3 Jul 2024 10:52:20 +0200
Subject: [PATCH] Add orientations readout

---
 src/exports.jl          |  1 +
 src/root/calibration.jl | 59 +++++++++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/exports.jl b/src/exports.jl
index e524c25b..d5c2857d 100644
--- a/src/exports.jl
+++ b/src/exports.jl
@@ -84,6 +84,7 @@ calibratetime,
 combine,
 floordist,
 slew,
+Orientations,
 
 # Reconstruction
 RecStageRange,
diff --git a/src/root/calibration.jl b/src/root/calibration.jl
index c7173232..9cbadeab 100644
--- a/src/root/calibration.jl
+++ b/src/root/calibration.jl
@@ -1,16 +1,59 @@
 struct Orientations
-    times::Vector{Float64}
-    quaternions::Vector{Quaternion}
+    module_ids::Set{Int}
+    times::Dict{Int, Vector{Float64}}
+    quaternions::Dict{Int, Vector{Quaternion}}
+end
+function Base.show(io::IO, o::Orientations)
+    min_t = Inf
+    max_t = -Inf
+    for times in values(o.times)
+        _min_t, _max_t = extrema(times)
+        if _min_t < min_t
+            min_t = _min_t
+        end
+        if _max_t > max_t
+            max_t = _max_t
+        end
+    end
+    t₁ = unix2datetime(min_t)
+    tâ‚‚ = unix2datetime(max_t)
+    print(io, "Orientations of $(length(o.module_ids)) modules ($(t₁) to $(t₂))")
+end
+function (o::Orientations)(module_id::Integer, time::Real)
+    times = o.times[module_id]
+    (time < first(times) || time > last(times)) && error("The requested time is outside of the range of the orientations data.")
+
+    idx2 = searchsortedfirst(times, time)
+    q2 = o.quaternions[module_id][idx2]
+    (idx2 >= length(times) || idx2 == 1) && return q2
+    idx1 = idx2 - 1
+    q1 = o.quaternions[module_id][idx1]
+
+    t1 = times[idx1]
+    t2 = times[idx2]
+    Δt = t2 - t1
+
+    Δt == 0.0 && return q1
+
+    t = (time - t1) / Δt
+
+    slerp(q1, q2, t)
 end
 
 
-function orientations(f::UnROOT.ROOTFile)
-    q_out = Dict{Int, Vector{Quaternion}}()
+function Base.read(filename::AbstractString, T::Type{Orientations})
+    f = UnROOT.ROOTFile(filename)
+    module_ids = Set{Int}()
+    quaternions = Dict{Int, Vector{Quaternion}}()
+    times = Dict{Int, Vector{Float64}}()
     for (module_id, t, a, b, c, d) in zip([UnROOT.LazyBranch(f, "ORIENTATION/ORIENTATION/$(b)") for b in ["id", "t", "JCOMPASS::JQuaternion/a", "JCOMPASS::JQuaternion/b", "JCOMPASS::JQuaternion/c", "JCOMPASS::JQuaternion/d"]]...)
-        if !haskey(q_out, module_id)
-            q_out[module_id] = Quaternion[]
+        if !(module_id ∈ module_ids)
+             push!(module_ids, module_id)
+            quaternions[module_id] = Quaternion[]
+            times[module_id] = Float64[]
         end
-        push!(q_out[module_id], Quaternion(a, b, c, d))
+        push!(quaternions[module_id], Quaternion(a, b, c, d))
+        push!(times[module_id], t)
     end
-    q_out
+    T(module_ids, times, quaternions)
 end
-- 
GitLab