handle [var][fix] style doubly jagged branch (#188)

* handle [var][fix] style doubly jagged branch

* fix tests for 1.6

Author: Moelf

Project.toml

@@ -1,7 +1,7 @@
 name = "UnROOT"
 uuid = "3cd96dde-e98d-4713-81e9-a4a1b0235ce9"
 authors = ["Tamas Gal", "Jerry Ling", "Johannes Schumann", "Nick Amin"]
-version = "0.8.15"
+version = "0.8.16"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/custom.jl

@@ -43,8 +43,16 @@ Base.iterate(x::FixLenVector) = iterate(x.vec)
 Base.iterate(x::FixLenVector, n) = iterate(x.vec, n)
 Base.getindex(x::FixLenVector, n) = getindex(x.vec, n)
 
+#N.B this is a hack, we deal with ntoh in reinterpret step
+Base.ntoh(x::FixLenVector) = x
 const VorView = Union{Vector{UInt8}, SubArray{UInt8, 1, Vector{UInt8}, Tuple{UnitRange{Int64}}, true}}
 
+function Base.reinterpret(::Type{Vector{FixLenVector{N, T}}}, data::Vector{UInt8}) where {N,T}
+    vs = reinterpret(T, data)
+    @. vs = ntoh(vs)
+    return FixLenVector.(reinterpret(SVector{N, T}, vs))
+end
+
 function Base.reinterpret(::Type{FixLenVector{N, T}}, v::VorView) where {N, T}
     vs = reinterpret(T, v)
     @. vs = ntoh(vs)

src/root.jl

@@ -407,14 +407,31 @@ function auto_T_JaggT(f::ROOTFile, branch; customstructs::Dict{String, Type})
             end
         end
     else
+        # since no classname were found, we now try to determine
+        # type based on leaf information
+        _type, _jaggtype = leaf_jaggtype(leaf, _jaggtype)
+    end
+    return _type, _jaggtype
+end
+
+function leaf_jaggtype(leaf, _jaggtype)
         _type = primitivetype(leaf)
-        if leaf.fLen > 1 # treat NTuple as Nojagg since size is static
-            _type = FixLenVector{Int(leaf.fLen), _type}
-            return _type, Nojagg
+        leafLen = leaf.fLen
+        if leafLen > 1 # treat NTuple as Nojagg since size is static
+            _fTitle = replace(leaf.fTitle, "[$(leafLen)]" => "")
+            # looking for more `[var]`
+            m = match(r"\[\D+\]", _fTitle)
+            _type = FixLenVector{Int(leafLen), _type}
+            if isnothing(m)
+                return _type, Nojagg
+            else
+                #FIXME this only handles [var][fix] case
+                return Vector{_type}, Nooffsetjagg
+            end
         end
         _type = _jaggtype === Nojagg ? _type : Vector{_type}
-    end
-    return _type, _jaggtype
+
+        return _type, _jaggtype
 end
 
 

test/runtests.jl

@@ -318,8 +318,9 @@ end
     @test v2[2].int32_array == data[3].int32_array
 end
 
-@testset "Doubly jagged branches" begin
-    rootfile = ROOTFile(joinpath(SAMPLES_DIR, "tree_with_doubly_jagged.root"))
+@testset "Doubly jagged [var][var] branches" begin
+    # this is vector<vector<blah>>
+    rootfile = UnROOT.samplefile("tree_with_doubly_jagged.root")
     vvi = [[[2], [3, 5]], [[7, 9, 11], [13]], [[17], [19], []], [], [[]]]
     vvf = [[[2.5], [3.5, 5.5]], [[7.5, 9.5, 11.5], [13.5]], [[17.5], [19.5], []], [], [[]]]
     @test UnROOT.array(rootfile, "t1/bi") == vvi
@@ -331,6 +332,33 @@ end
     close(rootfile)
 end
 
+@testset "Doubly jagged [var][fix] branches" begin
+    # issue #187
+    # this is vector<Int[N]>
+    f = UnROOT.samplefile("tree_with_varfix_doubly_jagged.root")
+    tree = LazyTree(f, "outtree")
+    @test tree.nparticles == [4,3,2]
+    @test length.(tree.P) == [4,3,2]
+    @test eltype(tree.P[1]) <: AbstractVector
+    # also compared to uproot
+    @test tree[1].P == [
+                        [0.9411764705882353, 0.8888888888888888, 0.8421052631578947, 0.8],
+                        [1.0, 0.9285714285714286, 0.8666666666666667, 0.8125],
+                        [1.1111111111111112, 1.0, 0.9090909090909091, 0.8333333333333334],
+                        [1.4, 1.1666666666666667, 1.0, 0.875]
+                       ]
+    @test tree[3].P == [
+                        [0.8222222222222222,
+                         0.8043478260869565,
+                         0.7872340425531915,
+                         0.7708333333333334],
+                        [0.8292682926829268,
+                         0.8095238095238095,
+                         0.7906976744186046,
+                         0.7727272727272727]
+                       ]
+end
+
 @testset "NanoAOD" begin
     rootfile = ROOTFile(joinpath(SAMPLES_DIR, "NanoAODv5_sample.root"))
     event = UnROOT.array(rootfile, "Events/event")
@@ -641,8 +669,11 @@ nthreads == 1 && @warn "Running on a single thread. Please re-run the test suite
 
 
     if get(ENV, "CI", "false") == "true"
-        # Make sure CI runs with more than 1 thread
-        @test Threads.nthreads()>1 skip=(nthreads == 1)
+        if nthreads >= 1
+            @test Threads.nthreads()>1 
+        else
+            @warn "CI wasn't run with multi thread"
+        end
     end
     nmus = zeros(Int, Threads.nthreads())
     Threads.@threads for i in 1:length(t)
@@ -671,15 +702,19 @@ t = LazyTree(ROOTFile(joinpath(SAMPLES_DIR, "NanoAODv5_sample.root")), "Events",
     Threads.@threads for evt in t
         nmus[Threads.threadid()] += length(evt.Muon_pt)
     end
-    @test count(>(0), nmus) > 1  skip = (nthreads==1)# test @threads is actually threading
+    if nthreads > 1
+        @test count(>(0), nmus) > 1# test @threads is actually threading
+    end
     @test sum(nmus) == 878
 
 
     nmus .= 0
     Threads.@threads for evt in t
         nmus[Threads.threadid()] += length(evt.Muon_pt)
     end
-    @test count(>(0), nmus) > 1  skip = (nthreads==1)# test @threads is actually threading
+    if nthreads > 1
+        @test count(>(0), nmus) > 1
+    end
     @test sum(nmus) == 878
 
     nmus .= 0

test/samples/tree_with_varfix_doubly_jagged.C

@@ -0,0 +1,27 @@
+#include "TFile.h"
+#include "TTree.h"
+
+int maketree(){
+    TFile f("tree_with_varfix_doubly_jagged.root", "RECREATE", "");
+    TTree tree = TTree("outtree", "outtree");
+    int nparticles{};
+    double P[100][4];
+    tree.Branch("nparticles", &nparticles, "nparticles/I");
+    tree.Branch("P", P, "P[nparticles][4]/D");
+    double counter1 = 1;
+    double counter2 = 1;
+    for (auto ev = 0; ev<3; ++ev){
+        nparticles = 4-ev;
+        for (auto i = nparticles; i>=0; --i){
+            counter1 += 3;
+            for (auto j = 0; j<=3; ++j){
+                P[i][j] = counter1 / (counter2);
+                counter2++;
+            }
+        }
+        tree.Fill();
+    }
+    f.Write();
+    f.Close();
+    return 0;
+}