Improve and extend detector indexing behaviour

docs/src/api.md

@@ -49,6 +49,8 @@ flush
 PMT
 DetectorModule
 Detector
+getmodule
+modules
 write(::AbstractString, ::Detector)
 write(::IO, ::Detector)
 Hydrophone

src/KM3io.jl

@@ -21,7 +21,7 @@ export ROOTFile
 export H5File, H5CompoundDataset, create_dataset, addmeta
 
 export Direction, Position, UTMPosition, Location, Quaternion, Track, AbstractCalibratedHit
-export Detector, DetectorModule, PMT, Tripod, Hydrophone, center, isbasemodule
+export Detector, DetectorModule, PMT, Tripod, Hydrophone, center, isbasemodule, getmodule, modules
 
 # Acoustics
 export Waveform, AcousticSignal, AcousticsTriggerParameter, piezoenabled, hydrophoneenabled

src/hardware.jl

@@ -266,8 +266,38 @@ function Base.iterate(d::Detector, state=(Int[], 1))
     end
     (d.modules[module_ids[count]], (module_ids, count + 1))
 end
-Base.getindex(d::Detector, module_id) = d.modules[module_id]
+Base.getindex(d::Detector, module_id::Integer) = d.modules[module_id]
+Base.getindex(d::Detector, string::Integer, floor::Integer) = d.locations[string, floor]
+"""
+Return the detector module for a given string and floor.
+"""
+@inline getmodule(d::Detector, string::Integer, floor::Integer) = d[string, floor]
+"""
+Return the detector module for a given string and floor (as `Tuple`).
+"""
+@inline getmodule(d::Detector, loc::Tuple{T, T}) where T<:Integer = d[loc...]
+"""
+Return the detector module for a given location.
+"""
+@inline getmodule(d::Detector, loc::Location) = d[loc.string, loc.floor]
+Base.getindex(d::Detector, string::Int, ::Colon) = sort!(filter(m->m.location.string == string, modules(d)))
+Base.getindex(d::Detector, string::Int, floors::T) where T<:Union{AbstractArray, UnitRange} = [d[string, floor] for floor in sort(floors)]
+Base.getindex(d::Detector, ::Colon, floor::Int) = sort!(filter(m->m.location.floor == floor, modules(d)))
+"""
+Return a vector of detector modules for a given range of floors on all strings.
 
+This can be useful if specific detector module layers of the detector are needed, e.g.
+the base modules (e.g. `detector[:, 0]`) or the top layer (e.g. `detector[:, 18]`).
+"""
+function Base.getindex(d::Detector, ::Colon, floors::UnitRange{T}) where T<:Integer
+    modules = DetectorModule[]
+    for string in d.strings
+        for floor in floors
+            push!(modules, d[string, floor])
+        end
+    end
+    sort!(modules)
+end
 
 """
 

test/hardware.jl

@@ -22,9 +22,11 @@ const SAMPLES_DIR = joinpath(@__DIR__, "samples")
                 # no base modules in DETX version <4
                 @test 342 == length(d)
                 @test 342 == length(d.modules)
+                @test 342 == length(modules(d))
                 @test DetectorModule == eltype(d)
             else
                 @test 361 == length(mods)
+                @test 361 == length(modules(d))
                 @test 106.95 ≈ d.modules[808469291].pos.y  # base module
                 @test 97.3720395 ≈ d.modules[808974928].pos.z  # base module
                 @test 0.0 == d.modules[808469291].t₀  # base module
@@ -35,8 +37,20 @@ const SAMPLES_DIR = joinpath(@__DIR__, "samples")
             if version > 3
                 @test Quaternion(1, 0, 0, 0) ≈ d.modules[808995481].q
                 @test 19 == length(collect(m for m ∈ d if isbasemodule(m)))
+                @test 19 == length(d[:, 0])
+                @test 19 == length(d[30, :])
+                @test 5 == length(d[30, 0:4])
+                for m in d[30, 0:4]
+                    @test m.location.floor in 0:4
+                end
             else
                 @test 0 == length(collect(m for m ∈ d if isbasemodule(m)))
+                @test 0 == length(d[:, 0])
+                @test 18 == length(d[30, :])
+                @test 4 == length(d[30, 1:4])
+                for m in d[30, 1:4]
+                    @test m.location.floor in 1:4
+                end
             end
 
             if version > 4
@@ -51,8 +65,21 @@ const SAMPLES_DIR = joinpath(@__DIR__, "samples")
             end
 
             @test 31 == d.modules[808992603].n_pmts
-            @test 30 ≈ d.modules[817287557].location.string
-            @test 18 ≈ d.modules[817287557].location.floor
+            @test 30 == d.modules[817287557].location.string
+            @test 18 == d.modules[817287557].location.floor
+            @test Location(30, 18) == d[817287557].location
+            @test 817287557 == d[30, 18].id
+            @test 817287557 == getmodule(d, 30, 18).id
+            @test 817287557 == getmodule(d, (30, 18)).id
+            @test 817287557 == getmodule(d, Location(30, 18)).id
+
+            @test 19 == length(d[:, 18])
+            for m in d[:, 17]
+                @test 17 == m.location.floor
+            end
+            for m in d[:, 1:4]
+                @test m.location.floor in 1:4
+            end
 
             @test 478392.31980645156 ≈ d.modules[808992603].t₀