diff --git a/src/UnROOT.jl b/src/UnROOT.jl index bf32aab7..2e58975c 100644 --- a/src/UnROOT.jl +++ b/src/UnROOT.jl @@ -11,7 +11,7 @@ import AbstractTrees: children, printnode, print_tree using CodecZlib, CodecLz4, CodecXz, CodecZstd, StaticArrays, LorentzVectors, ArraysOfArrays using Mixers, Parameters, Memoization, LRUCache -import Tables, TypedTables, PrettyTables, DataFrames +import Tables, TypedTables, PrettyTables @static if VERSION < v"1.6" Base.first(a::AbstractVector{S}, n::Integer) where S<: AbstractString = a[1:(length(a) > n ? n : end)] diff --git a/src/custom.jl b/src/custom.jl index c99f9496..7d1c6490 100644 --- a/src/custom.jl +++ b/src/custom.jl @@ -73,7 +73,7 @@ function interped_data(rawdata, rawoffsets, ::Type{Vector{LVF64}}, ::Type{Offset offset .+= 1 VectorOfVectors(real_data, offset) end -function interped_data(rawdata, rawoffsets, ::Type{LVF64}, ::Type{J}) where {T, J <: JaggType} +function interped_data(rawdata, rawoffsets, ::Type{LVF64}, ::Type{Nojagg}) # even with rawoffsets, we know each TLV is destinied to be 64 bytes [ reinterpret(LVF64, x) for x in Base.Iterators.partition(rawdata, 64) @@ -82,48 +82,32 @@ end # TLorentzVector ends # KM3NeT -struct KM3NETDAQHit <: CustomROOTStruct +struct _KM3NETDAQHit <: CustomROOTStruct dom_id::Int32 channel_id::UInt8 tdc::Int32 tot::UInt8 end -function readtype(io::IO, T::Type{KM3NETDAQHit}) +function readtype(io::IO, T::Type{_KM3NETDAQHit}) T(readtype(io, Int32), read(io, UInt8), read(io, Int32), read(io, UInt8)) end -function interped_data(rawdata, rawoffsets, ::Type{Vector{KM3NETDAQHit}}, ::Type{J}) where {T, J <: UnROOT.JaggType} - UnROOT.splitup(rawdata, rawoffsets, KM3NETDAQHit, skipbytes=10) +function interped_data(rawdata, rawoffsets, ::Type{Vector{_KM3NETDAQHit}}, ::Type{Nojagg}) + UnROOT.splitup(rawdata, rawoffsets, _KM3NETDAQHit, skipbytes=10) end - - -# Experimental implementation for maximum performance (using reinterpret) -primitive type DAQHit 80 end -function Base.getproperty(hit::DAQHit, s::Symbol) - r = Ref(hit) - GC.@preserve r begin - if s === :dom_id - return ntoh(unsafe_load(Ptr{Int32}(Base.unsafe_convert(Ptr{Cvoid}, r)))) - elseif s === :channel_id - return unsafe_load(Ptr{UInt8}(Base.unsafe_convert(Ptr{Cvoid}, r)+4)) - elseif s === :tdc - return unsafe_load(Ptr{UInt32}(Base.unsafe_convert(Ptr{Cvoid}, r)+5)) - elseif s === :tot - return unsafe_load(Ptr{UInt8}(Base.unsafe_convert(Ptr{Cvoid}, r)+9)) - end - end - error("unknown field $s of type $(typeof(hit))") +function interped_data(rawdata, rawoffsets, ::Type{Vector{_KM3NETDAQHit}}, ::Type{Offsetjagg}) + UnROOT.splitup(rawdata, rawoffsets, _KM3NETDAQHit, skipbytes=10) end -Base.show(io::IO, h::DAQHit) = print(io, "DAQHit(", h.dom_id, ',', h.channel_id, ',', h.tdc, ',', h.tot, ')') -struct KM3NETDAQTriggeredHit +struct _KM3NETDAQTriggeredHit dom_id::Int32 channel_id::UInt8 tdc::Int32 tot::UInt8 trigger_mask::UInt64 end -function readtype(io::IO, T::Type{KM3NETDAQTriggeredHit}) +packedsizeof(::Type{_KM3NETDAQTriggeredHit}) = 24 # incl. cnt and vers +function readtype(io::IO, T::Type{_KM3NETDAQTriggeredHit}) dom_id = readtype(io, Int32) channel_id = read(io, UInt8) tdc = read(io, Int32) @@ -133,11 +117,14 @@ function readtype(io::IO, T::Type{KM3NETDAQTriggeredHit}) T(dom_id, channel_id, tdc, tot, trigger_mask) end -function UnROOT.interped_data(rawdata, rawoffsets, ::Type{Vector{KM3NETDAQTriggeredHit}}, ::Type{J}) where {T, J <: UnROOT.JaggType} - UnROOT.splitup(rawdata, rawoffsets, KM3NETDAQTriggeredHit, skipbytes=10) +function UnROOT.interped_data(rawdata, rawoffsets, ::Type{Vector{_KM3NETDAQTriggeredHit}}, ::Type{Nojagg}) + UnROOT.splitup(rawdata, rawoffsets, _KM3NETDAQTriggeredHit, skipbytes=10) +end +function UnROOT.interped_data(rawdata, rawoffsets, ::Type{Vector{_KM3NETDAQTriggeredHit}}, ::Type{Offsetjagg}) + UnROOT.splitup(rawdata, rawoffsets, _KM3NETDAQTriggeredHit, skipbytes=10) end -struct KM3NETDAQEventHeader +struct _KM3NETDAQEventHeader detector_id::Int32 run::Int32 frame_index::Int32 @@ -147,9 +134,9 @@ struct KM3NETDAQEventHeader trigger_mask::UInt64 overlays::UInt32 end -packedsizeof(::Type{KM3NETDAQEventHeader}) = 76 +packedsizeof(::Type{_KM3NETDAQEventHeader}) = 76 -function readtype(io::IO, T::Type{KM3NETDAQEventHeader}) +function readtype(io::IO, T::Type{_KM3NETDAQEventHeader}) skip(io, 18) detector_id = readtype(io, Int32) run = readtype(io, Int32) @@ -165,6 +152,6 @@ function readtype(io::IO, T::Type{KM3NETDAQEventHeader}) T(detector_id, run, frame_index, UTC_seconds, UTC_16nanosecondcycles, trigger_counter, trigger_mask, overlays) end -function UnROOT.interped_data(rawdata, rawoffsets, ::Type{KM3NETDAQEventHeader}, ::Type{J}) where {T, J <: UnROOT.JaggType} - UnROOT.splitup(rawdata, rawoffsets, KM3NETDAQEventHeader, jagged=false) +function UnROOT.interped_data(rawdata, rawoffsets, ::Type{_KM3NETDAQEventHeader}, ::Type{Nojagg}) + UnROOT.splitup(rawdata, rawoffsets, _KM3NETDAQEventHeader, jagged=false) end diff --git a/src/displays.jl b/src/displays.jl index 4eff9ac2..e410cbe7 100644 --- a/src/displays.jl +++ b/src/displays.jl @@ -3,17 +3,34 @@ These functions are used to display a ROOTFile is a tree-like fashion by using `AbstractTrees` printing functions. We customize what the children of ROOTFile and a TTree is, and how to print the final `node`. =# +struct TKeyNode + name::AbstractString + classname::AbstractString +end function children(f::ROOTFile) - ch = Vector{TTree}() + # display TTrees recursively + # subsequent TTrees with duplicate fName will be skipped + # since TKey cycle number is guaranteed to be decreasing + # then all TKeys in the file which are not for a TTree + seen = Set{String}() + ch = Vector{Union{TTree,TKeyNode}}() + lock(f) for k in keys(f) - lock(f.fobj) try - push!(ch, f[k]) + obj = f[k] + obj isa TTree || continue + obj.fName ∈ seen && continue + push!(ch, obj) + push!(seen, obj.fName) catch - finally - unlock(f.fobj) end end + for tkey in f.directory.keys + kn = TKeyNode(tkey.fName, tkey.fClassName) + kn.classname == "TTree" && continue + push!(ch, kn) + end + unlock(f) ch end function children(t::TTree) @@ -26,15 +43,16 @@ function children(t::TTree) return ks end end -printnode(io::IO, t::TTree) = print(io, t.fName) +printnode(io::IO, t::TTree) = print(io, "$(t.fName) (TTree)") printnode(io::IO, f::ROOTFile) = print(io, f.filename) +printnode(io::IO, k::TKeyNode) = print(io, "$(k.name) ($(k.classname))") function Base.show(io::IO, tree::LazyTree) _hs = _make_header(tree) _ds = displaysize(io) PrettyTables.pretty_table( io, - tree; + innertable(tree); header=_hs, alignment=:l, vlines=[1], @@ -44,7 +62,7 @@ function Base.show(io::IO, tree::LazyTree) row_number_column_title="Row", show_row_number=true, compact_printing=false, - formatters=(v, i, j) -> _treeformat(v, _ds[2] ÷ min(5, length(_hs[1]))), + formatters=(v, i, j) -> _treeformat(v, _ds[2] ÷ min(8, length(_hs[1]))), display_size=(min(_ds[1], 40), min(_ds[2], 160)), ) end @@ -52,14 +70,15 @@ _symtup2str(symtup, trunc=15) = collect(first.(string.(symtup), trunc)) function _make_header(t) pn = propertynames(t) header = _symtup2str(pn) - subheader = _symtup2str(Tables.columntype.(Ref(t), pn)) + subheader = _symtup2str(Tables.columntype.(Ref(innertable(t)), pn)) (header, subheader) end function _treeformat(val, trunc) - s = if val isa Vector{T} where T<:Integer + s = if val isa AbstractArray{T} where T<:Integer string(Int.(val)) - elseif val isa Vector{T} where T<:AbstractFloat - string(round.(Float64.(val); sigdigits=3)) + elseif val isa AbstractArray{T} where T<:AbstractFloat + T = eltype(val) + replace(string(round.(T.(val); sigdigits=3)), string(T)=>"") else string(val) end diff --git a/src/iteration.jl b/src/iteration.jl index 874420b0..1cfee31b 100644 --- a/src/iteration.jl +++ b/src/iteration.jl @@ -5,7 +5,7 @@ Reads all branches from a tree. """ function arrays(f::ROOTFile, treename) names = keys(f[treename]) - res = Vector{Any}(undef, length(names)) + res = Vector{Vector}(undef, length(names)) Threads.@threads for i in eachindex(names) res[i] = array(f, "$treename/$(names[i])") end @@ -122,12 +122,12 @@ Base.eltype(ba::LazyBranch{T,J,B}) where {T,J,B} = T function Base.show(io::IO, lb::LazyBranch) summary(io, lb) - println(":") - println(" File: $(lb.f.filename)") - println(" Branch: $(lb.b.fName)") - println(" Description: $(lb.b.fTitle)") - println(" NumEntry: $(lb.L)") - print(" Entry Type: $(eltype(lb))") + println(io, ":") + println(io, " File: $(lb.f.filename)") + println(io, " Branch: $(lb.b.fName)") + println(io, " Description: $(lb.b.fTitle)") + println(io, " NumEntry: $(lb.L)") + print(io, " Entry Type: $(eltype(lb))") nothing end @@ -166,27 +166,24 @@ function Base.iterate(ba::LazyBranch{T,J,B}, idx=1) where {T,J,B} return (ba[idx], idx + 1) end -const _LazyTreeType = - TypedTables.Table{<:NamedTuple,1,NamedTuple{S,N}} where {S,N<:Tuple{Vararg{LazyBranch}}} - -struct LazyTree{T} <: DataFrames.AbstractDataFrame +struct LazyTree{T} treetable::T - colidx::DataFrames.Index end + @inline innertable(t::LazyTree) = Core.getfield(t, :treetable) +Base.propertynames(lt::LazyTree) = propertynames(innertable(lt)) +Base.getproperty(lt::LazyTree, s::Symbol) = getproperty(innertable(lt), s) + # a specific branch Base.getindex(lt::LazyTree, row::Int) = innertable(lt)[row] function Base.getindex(lt::LazyTree, rang::UnitRange) - return LazyTree(innertable(lt)[rang], Core.getfield(lt, :colidx)) + return LazyTree(innertable(lt)[rang]) end Base.getindex(lt::LazyTree, ::typeof(!), s::Symbol) = lt[:, s] -Base.getindex(lt::LazyTree, ::Colon, i::Int) = lt[:, propertynames(lt)[i]] -Base.getindex(lt::LazyTree, ::typeof(!), i::Int) = lt[:, propertynames(lt)[i]] Base.getindex(lt::LazyTree, ::Colon, s::Symbol) = getproperty(innertable(lt), s) # the real deal # a specific event -Base.getindex(lt::LazyTree, row::Int, col::Int) = lt[:, col][row] Base.getindex(lt::LazyTree, row::Int, col::Symbol) = lt[:, col][row] Base.getindex(lt::LazyTree, rows::UnitRange, col::Symbol) = lt[:, col][rows] Base.getindex(lt::LazyTree, ::Colon) = lt[1:end] @@ -200,13 +197,9 @@ Base.lastindex(e::Iterators.Enumerate{LazyTree{T}}) where T = lastindex(e.itr) Base.eachindex(e::Iterators.Enumerate{LazyTree{T}}) where T = eachindex(e.itr) Base.getindex(e::Iterators.Enumerate{LazyTree{T}}, row::Int) where T = (row, first(iterate(e.itr, row))) -# interfacing AbstractDataFrame -DataFrames._check_consistency(lt::LazyTree) = nothing #we're read-only +# interfacing Table Base.names(lt::LazyTree) = collect(String.(propertynames(innertable(lt)))) -DataFrames.index(lt::LazyTree) = Core.getfield(lt, :colidx) -DataFrames.ncol(lt::LazyTree) = length(DataFrames.index(lt)) Base.length(lt::LazyTree) = length(innertable(lt)) -DataFrames.nrow(lt::LazyTree) = length(lt) function getbranchnamesrecursive(obj) out = Vector{String}() @@ -223,7 +216,7 @@ end LazyTree(f::ROOTFile, s::AbstractString, branche::Union{AbstractString, Regex}) LazyTree(f::ROOTFile, s::AbstractString, branches::Vector{Union{AbstractString, Regex}}) -Constructor for `LazyTree`, which is close to an `AbstractDataFrame` (interface wise), +Constructor for `LazyTree`, which is close to an `DataFrame` (interface wise), and a lazy `TypedTables.Table` (speed wise). Looping over a `LazyTree` is fast and type stable. Internally, `LazyTree` contains a typed table whose branch are [`LazyBranch`](@ref). This means that at any given time only `N` baskets are cached, where `N` is the number of branches. @@ -251,16 +244,14 @@ function LazyTree(f::ROOTFile, s::AbstractString, branches) @warn "Your tree is quite wide, with $(length(branches)) branches, this will take compiler a moment." end d = Dict{Symbol,LazyBranch}() - d_colidx = Dict{Symbol,Int}() _m(s::AbstractString) = isequal(s) _m(r::Regex) = Base.Fix1(occursin, r) branches = mapreduce(b -> filter(_m(b), getbranchnamesrecursive(tree)), ∪, branches) SB = Symbol.(branches) - for (i, b) in enumerate(SB) + for b in SB d[b] = f["$s/$b"] - d_colidx[b] = i end - return LazyTree(TypedTables.Table(d), DataFrames.Index(d_colidx, SB)) + return LazyTree(TypedTables.Table(d)) end function LazyTree(f::ROOTFile, s::AbstractString) @@ -285,7 +276,7 @@ end function Base.getproperty(evt::LazyEvent, s::Symbol) @inbounds getproperty(Core.getfield(evt, :tree), s)[Core.getfield(evt, :idx)] end -Base.collect(evt::LazyEvent) = Core.getfield(evt, :tree)[Core.getfield(evt, :idx)] +Base.collect(evt::LazyEvent) = @inbounds Core.getfield(evt, :tree)[Core.getfield(evt, :idx)] function Base.iterate(tree::T, idx=1) where {T<:LazyTree} idx > length(tree) && return nothing diff --git a/src/root.jl b/src/root.jl index 8e4d26bc..1335e04a 100644 --- a/src/root.jl +++ b/src/root.jl @@ -193,67 +193,72 @@ on type `T` and jagg type `J`. In order to retrieve data from custom branches, user should defined more speialized method of this function with specific `T` and `J`. See `TLorentzVector` example. """ -function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{J}) where {T, J<:JaggType} - # there are two possibility, one is the leaf is just normal leaf but the title has "[...]" in it - # magic offsets, seems to be common for a lot of types, see auto.py in uproot3 - # only needs when the jaggedness comes from TLeafElements, not needed when - # the jaggedness comes from having "[]" in TLeaf's title - # the other is where we need to auto detector T bsaed on class name - # we want the fundamental type as `reinterpret` will create vector - if J == Nojagg - return ntoh.(reinterpret(T, rawdata)) - elseif J == Offsetjaggjagg # the branch is doubly jagged - jagg_offset = 10 - subT = eltype(eltype(T)) - out = VectorOfVectors(T(), Int32[1]) - @views for i in 1:(length(rawoffsets)-1) - flat = rawdata[(rawoffsets[i]+1+jagg_offset:rawoffsets[i+1])] - row = VectorOfVectors{subT}() - cursor = 1 - while cursor < length(flat) - n = ntoh(reinterpret(Int32, flat[cursor:cursor+sizeof(Int32)-1])[1]) - cursor += sizeof(Int32) - b = ntoh.(reinterpret(subT, flat[cursor:cursor+n*sizeof(subT)-1])) - cursor += n*sizeof(subT) - push!(row, b) - end - push!(out, row) - end - return out - else # the branch is singly jagged - # for each "event", the index range is `offsets[i] + jagg_offset + 1` to `offsets[i+1]` - # this is why we need to append `rawoffsets` in the `readbranchraw()` call - # when you use this range to index `rawdata`, you will get raw bytes belong to each event - # Say your real data is Int32 and you see 8 bytes after indexing, then this event has [num1, num2] as real data - _size = sizeof(eltype(T)) - if J === Offsetjagg - jagg_offset = 10 - dp = 0 # book keeping for copy_to! - lr = length(rawoffsets) - offset = Vector{Int32}(undef, lr) - offset[1] = 0 - @views @inbounds for i in 1:lr-1 - start = rawoffsets[i]+jagg_offset+1 - stop = rawoffsets[i+1] - l = stop-start+1 - if l > 0 - unsafe_copyto!(rawdata, dp+1, rawdata, start, l) - dp += l - offset[i+1] = offset[i] + l - else - # when we have an empty [] in jagged basket - offset[i+1] = offset[i] - end - end - resize!(rawdata, dp) +function interped_data(rawdata, rawoffsets, ::Type{Bool}, ::Type{Nojagg}) + # specialized case to get Vector{Bool} instead of BitVector + return map(ntoh,reinterpret(Bool, rawdata)) +end +function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{Nojagg}) where T + return ntoh.(reinterpret(T, rawdata)) +end +# there are two possibility, one is the leaf is just normal leaf but the title has "[...]" in it +# magic offsets, seems to be common for a lot of types, see auto.py in uproot3 +# only needs when the jaggedness comes from TLeafElements, not needed when +# the jaggedness comes from having "[]" in TLeaf's title +# the other is where we need to auto detector T bsaed on class name +# we want the fundamental type as `reinterpret` will create vector +function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{Nooffsetjagg}) where T + _size = sizeof(eltype(T)) + real_data = ntoh.(reinterpret(T, rawdata)) + rawoffsets .= (rawoffsets .÷ _size) .+ 1 + return VectorOfVectors(real_data, rawoffsets, ArraysOfArrays.no_consistency_checks) +end +function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{Offsetjagg}) where T + # for each "event", the index range is `offsets[i] + jagg_offset + 1` to `offsets[i+1]` + # this is why we need to append `rawoffsets` in the `readbranchraw()` call + # when you use this range to index `rawdata`, you will get raw bytes belong to each event + # Say your real data is Int32 and you see 8 bytes after indexing, then this event has [num1, num2] as real data + _size = sizeof(eltype(T)) + jagg_offset = 10 + dp = 0 # book keeping for copy_to! + lr = length(rawoffsets) + offset = Vector{Int32}(undef, lr) + offset[1] = 0 + @views @inbounds for i in 1:lr-1 + start = rawoffsets[i]+jagg_offset+1 + stop = rawoffsets[i+1] + l = stop-start+1 + if l > 0 + unsafe_copyto!(rawdata, dp+1, rawdata, start, l) + dp += l + offset[i+1] = offset[i] + l else - offset = rawoffsets + # when we have an empty [] in jagged basket + offset[i+1] = offset[i] + end + end + resize!(rawdata, dp) + real_data = ntoh.(reinterpret(T, rawdata)) + offset .= (offset .÷ _size) .+ 1 + return VectorOfVectors(real_data, offset, ArraysOfArrays.no_consistency_checks) +end +function interped_data(rawdata, rawoffsets, ::Type{T}, ::Type{Offsetjaggjagg}) where T + jagg_offset = 10 + subT = eltype(eltype(T)) + out = VectorOfVectors(T(), Int32[1]) + @views for i in 1:(length(rawoffsets)-1) + flat = rawdata[(rawoffsets[i]+1+jagg_offset:rawoffsets[i+1])] + row = VectorOfVectors{subT}() + cursor = 1 + while cursor < length(flat) + n = ntoh(reinterpret(Int32, flat[cursor:cursor+sizeof(Int32)-1])[1]) + cursor += sizeof(Int32) + b = ntoh.(reinterpret(subT, flat[cursor:cursor+n*sizeof(subT)-1])) + cursor += n*sizeof(subT) + push!(row, b) end - real_data = ntoh.(reinterpret(T, rawdata)) - offset .÷= _size - offset .+= 1 - return VectorOfVectors(real_data, offset) + push!(out, row) end + return out end function _normalize_ftype(fType) diff --git a/src/utils.jl b/src/utils.jl index 594c3c31..8c63d395 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -99,3 +99,7 @@ function parseTH(th::Dict{Symbol, Any}) end return counts, edges, sumw2 end + +function samplefile(filename::AbstractString) + return ROOTFile(normpath(joinpath(@__DIR__, "../test/samples", filename))) +end diff --git a/test/runtests.jl b/test/runtests.jl index 6184d52e..a7438143 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -215,6 +215,12 @@ end rootfile = ROOTFile(joinpath(SAMPLES_DIR, "tree_with_large_array.root")) branch = rootfile["t1"]["int32_array"] arr = UnROOT.array(rootfile, branch) + arr2 = UnROOT.arrays(rootfile, "t1")[1] + + @test hash(branch) == hash(rootfile["t1"]["int32_array"]) + @test hash(branch) != hash(rootfile["t1"]["float_array"]) + @test arr == arr2 + table = LazyTree(rootfile, "t1") BA = LazyBranch(rootfile, branch) @test length(arr) == length(BA) @@ -223,6 +229,10 @@ end @test BA[20:30] == arr[20:30] @test BA[1:end] == arr @test table.int32_array[20:30] == BA[20:30] + @test table[:, :int32_array][20:30] == BA[20:30] + @test table[23, :int32_array] == BA[23] + @test table[20:30, :int32_array] == BA[20:30] + @test table[:].int32_array[20:30] == BA[20:30] @test [row.int32_array for row in table[20:30]] == BA[20:30] @test sum(table.int32_array) == sum(row.int32_array for row in table) @test [row.int32_array for row in table] == BA @@ -308,10 +318,11 @@ end @test eltype(HLT_Mu3_PFJet40) == Bool @test HLT_Mu3_PFJet40[1:3] == [false, true, false] tree = LazyTree(rootfile, "Events", [r"Muon_(pt|eta|phi)$", "Muon_charge", "Muon_pt"]) - @test sort(propertynames(tree)) == sort([:Muon_pt, :Muon_eta, :Muon_phi, :Muon_charge]) + @test sort(propertynames(tree) |> collect) == sort([:Muon_pt, :Muon_eta, :Muon_phi, :Muon_charge]) tree = LazyTree(rootfile, "Events", r"Muon_(pt|eta)$") - @test sort(propertynames(tree)) == sort([:Muon_pt, :Muon_eta]) + @test sort(propertynames(tree) |> collect) == sort([:Muon_pt, :Muon_eta]) @test occursin("LazyEvent", repr(first(iterate(tree)))) + @test sum(rootfile["Events/HLT_Mu3_PFJet40"]) == 443 close(rootfile) end @@ -328,15 +339,36 @@ end @test filter_branches(["Muon.pt"]) == Set(["Muon.pt"]) end -@testset "Displaying" begin - files = filter(endswith(".root"), readdir(SAMPLES_DIR)) +@testset "Displaying files" begin + files = filter(x->endswith(x, ".root"), readdir(SAMPLES_DIR)) _io = IOBuffer() for f in files r = ROOTFile(joinpath(SAMPLES_DIR, f)) show(_io, r) close(r) end + + # test that duplicate trees (but different cycle numbers) + # are only displayed once, and that histograms show up + f = UnROOT.samplefile("tree_cycles_hist.root") + @test length(collect(eachmatch(r"Events", repr(f)))) == 1 + @test length(collect(eachmatch(r"myTH2F", repr(f)))) == 1 + close(f) +end + +@testset "Displaying trees" begin + f = UnROOT.samplefile("NanoAODv5_sample.root") + t = LazyTree(f, "Events", ["nMuon","MET_pt","Muon_pt"]) + _io = IOBuffer() + show(_io, t) + show(_io, t[1:10]) + show(_io, t.Muon_pt) + show(_io, t.Muon_pt[1:10]) + s = repr(t[1:10]) + @test length(collect(eachmatch(r"Float32\[", s))) == 0 + close(f) end + # Custom bootstrap things @testset "custom boostrapping" begin @@ -344,18 +376,21 @@ end f_manual = ROOTFile(joinpath(SAMPLES_DIR, "km3net_online.root")) data, offsets = UnROOT.array(f_manual, "KM3NET_EVENT/KM3NET_EVENT/KM3NETDAQ::JDAQEventHeader"; raw=true) - headers_manual = UnROOT.splitup(data, offsets, UnROOT.KM3NETDAQEventHeader; jagged=false) + headers_manual = UnROOT.splitup(data, offsets, UnROOT._KM3NETDAQEventHeader; jagged=false) data, offsets = UnROOT.array(f_manual, "KM3NET_EVENT/KM3NET_EVENT/snapshotHits"; raw=true) - event_hits_manual = UnROOT.splitup(data, offsets, UnROOT.KM3NETDAQHit; skipbytes=10) + event_hits_manual = UnROOT.splitup(data, offsets, UnROOT._KM3NETDAQHit; skipbytes=10) + + data, offsets = UnROOT.array(f_manual, "KM3NET_EVENT/KM3NET_EVENT/triggeredHits"; raw=true) + event_thits_manual = UnROOT.splitup(data, offsets, UnROOT._KM3NETDAQTriggeredHit; skipbytes=10) close(f_manual) # we can close, everything is in memory # automatic interpretation customstructs = Dict( - "KM3NETDAQ::JDAQEvent.snapshotHits" => Vector{UnROOT.KM3NETDAQHit}, - "KM3NETDAQ::JDAQEvent.triggeredHits" => Vector{UnROOT.KM3NETDAQTriggeredHit}, - "KM3NETDAQ::JDAQEvent.KM3NETDAQ::JDAQEventHeader" => UnROOT.KM3NETDAQEventHeader + "KM3NETDAQ::JDAQEvent.snapshotHits" => Vector{UnROOT._KM3NETDAQHit}, + "KM3NETDAQ::JDAQEvent.triggeredHits" => Vector{UnROOT._KM3NETDAQTriggeredHit}, + "KM3NETDAQ::JDAQEvent.KM3NETDAQ::JDAQEventHeader" => UnROOT._KM3NETDAQEventHeader ) f_auto = UnROOT.ROOTFile(joinpath(SAMPLES_DIR, "km3net_online.root"), customstructs=customstructs) headers_auto = f_auto["KM3NET_EVENT/KM3NET_EVENT/KM3NETDAQ::JDAQEventHeader"] @@ -376,6 +411,20 @@ end @test event_hits[3][end].dom_id == 809544061 @test event_hits[3][end].tdc == 63512892 end + for event_thits ∈ [event_thits_manual, event_thits_auto] + @test length(event_thits) == 3 + @test length(event_thits[1]) == 18 + @test length(event_thits[2]) == 53 + @test length(event_thits[3]) == 9 + @test event_thits[1][1].dom_id == 806451572 + @test event_thits[1][1].tdc == 30733918 + @test event_thits[1][end].dom_id == 808972598 + @test event_thits[1][end].tdc == 30733192 + @test event_thits[3][1].dom_id == 808447186 + @test event_thits[3][1].tdc == 63511558 + @test event_thits[3][end].dom_id == 809526097 + @test event_thits[3][end].tdc == 63511708 + end for headers ∈ [headers_manual, headers_auto] @test length(headers) == 3 @@ -530,6 +579,7 @@ end @testset "Parallel and enumerate interface" begin t = LazyTree(ROOTFile(joinpath(SAMPLES_DIR, "NanoAODv5_sample.root")), "Events", ["Muon_pt"]) + @test eachindex(enumerate(t)) == eachindex(t) nmu = 0 for evt in t nmu += length(evt.Muon_pt) diff --git a/test/samples/tree_cycles_hist.py b/test/samples/tree_cycles_hist.py new file mode 100644 index 00000000..4fe7c178 --- /dev/null +++ b/test/samples/tree_cycles_hist.py @@ -0,0 +1,35 @@ +import ROOT as r + +f = r.TFile("tree_cycles_hist.root", "recreate") +t = r.TTree("Events", "") + +obj = r.vector("float")() +t.Branch("Jet_pt", obj) +rows = [[], [27.324586868286133, 24.88954734802246, 20.853023529052734], [], [20.330659866333008], [], []] +for i,row in enumerate(rows): + obj.clear() + for x in row: + obj.push_back(x) + t.Fill() + if i == 3: + t.Write() + + +th1f = r.TH1F("myTH1F", "", 2, -2, 2) +th1d = r.TH1D("myTH1D", "", 2, -2, 2) +th2f = r.TH2F("myTH2F", "", 2, -2, 2, 4, -2, 2) +th2d = r.TH2D("myTH2D", "", 2, -2, 2, 4, -2, 2) + +for x,y,w in [ + [-1.5, -1.5, 20.0], + [+1.5, +1.5, 1.0], + [-1.5, +1.5, 20.0], + [+1.5, -1.5, 1.0], + ]: + th1f.Fill(x, w) + th1d.Fill(x, w) + th2f.Fill(x, y, w) + th2d.Fill(x, y, w) + +f.Write() +f.Close() diff --git a/test/samples/tree_cycles_hist.root b/test/samples/tree_cycles_hist.root new file mode 100644 index 00000000..0c8f4326 Binary files /dev/null and b/test/samples/tree_cycles_hist.root differ